2024. 6. 6. 19:31ㆍDevelopers 공간 [Basic]/Vision & Audio
이번 글에서는 StabilityAI에서 공개한 SD(Stable Diffusion)이 LDM기반으로 등장한 이후에, 개선된 다양한 모델의 히스토리를 살펴보고자합니다.
이런 다양한 SD들을 활용하기 위해서는 보통 Diffusers라는 라이브러리를 활용해 구현합니다.
따라서 다양한 모델들의 버전별 큰 특징들을 먼저 살펴보고, Diffusers라는 라이브러리를 활용해 구현할 때 성능 개선을 위해 사용하는 다양한 옵션에 대해서도 논문과 함께 살펴보고자 합니다.
<구성>
1. Core Models
a. Stable Diffusion 2.0
b. Stable Diffusion XL
c. Stable Diffusion 3.0
d. Stable Diffusion Turbo
2. Optional Improvements
a. EMA(Exponential Moving Average)
b. X-formers
c. DPM-Solver++
d. freeU
글효과 분류1 : 코드
글효과 분류2 : 폴더/파일
글효과 분류3 : 용어설명
글효과 분류4 : 글 내 참조
글효과 분류5 : 글 내 참조2
글효과 분류6 : 글 내 참조3
1. Core Models
StabilityAI에서 2022년 8월에 Stable Diffusion 1.0이 등장했습니다.
** Blog : https://stability.ai/news/stable-diffusion-public-release
** Weight : https://huggingface.co/CompVis/stable-diffusion
이후에 v1-1, v1-2, v1-3, v1-4, v1-5으로 개선된 Weight들이 등장하지만, 모델 구조에 대한 변화보다는 데이터를 통한 Weight 개선으로 나아갑니다.
** Weight-v1.5 : https://huggingface.co/runwayml/stable-diffusion-v1-5
하지만, 그 이후 이미지 생성 외 다양한 분야의 모델도 등장했을 뿐아니라, 이미지 모델도 다양한 버전으로 아래와 같이 지속적으로 발전해 나갑니다.
| SD 1.5 | SD 2.0 | SDXL 1.0 | SDXL Turbo |
SD 3.0 | SA 1.0 | ||
| 출시 |
22.10 | 22.11 | 23.07 | 23.11 | 24.02 | 24.06 | |
| 파라미터 |
860M | 865M | 2.6B | 3.1B | 2B | 1.2B | |
| 이미지 사이즈 |
512x512 | 768x768 | 1024x1024 | 1024x1024 | 1024x1024 | . | |
| Control |
Text | Text | Text x 2 Size Crop |
Text x 2 Size Crop |
Text x 3 Noisy latent |
Text | |
| Sampler | Scheduler | PNDM Scheduler |
DDIM Scheduler |
Euler Discrete Scheduler |
EulerAncestral Discrete Scheduler |
FlowMatchEuler Discrete Scheduler |
CosineDPMSolver Multistep Scheduler |
| Objective | epsilon | v_prediction | epsilon | epsilon | Rectified Flow | v_prediction | |
| Core Model? |
X | X | X | O | O | O | |
** Sampler는 다른 방법으로 표현할 수도 있지만, Diffusers에 있는 Object(https://huggingface.co/docs/diffusers/v0.30.1/en/api/schedulers/overview)를 기준으로 표기했습니다.
위 Core Model 이란, StabilityAI에서 지정한 모델들을 의미합니다. 이런 Core 모델들은 License의 대상이 되므로 파악할 필요가 있어 표시해두었습니다.
** https://stability.ai/core-models
** https://stability.ai/license
이제 초기 모델 이후 개선된 모델들의 큰 특징에 대해 자세히 살펴보고자 합니다.
a. Stable Diffusion 2.0
Blog : https://stability.ai/news/stable-diffusion-v2-release
Weight : https://huggingface.co/stabilityai/stable-diffusion-2
SD 2.0은 v-objective(v-prediction, v-diffusion)을 활용해 학습한 것이 가장 큰 특징입니다. 따라서 v-diffusion이 등장한 논문을 살펴보려고합니다.
** Progressive Distillation for Fast Sampling of Diffusion Models(arxiv'22)
기존 DM의 단점은 sampling time이 느리다는 것(Nx100~Nx1000 steps)이었습니다. 따라서 본논문은 이를 개선하기 위해 아래 두 가지의 contribution을 제안합니다.
- 적은 sampling step으로도 stability를 높일 수 있는 DM의 새로운 parametrization 제안
- 학습된 deterministic diffusion sampler를 distill하는 방법 제안
구현에 앞서, 본 논문에서는 Diffusion의 forward process를 SNR(Signal-To-Noise Ratio)이 image(t=0.0)에서 noise(t=1.0)으로 갈수록 단조롭게(monotonically) 감소하도록 noise를 scheduling합니다.
** 이 Scheduling 방법을 제안한 VDM에 대해 궁금하신 분들은 더보기를 참조하세요
$$SNR(t)=\frac{\alpha_t^2}{\sigma_t^2}$$
----------------------------------------------------------------
<VDM>
** Variational Diffusion Models(NIPS’21)
기존에 데이터 x와 noisy한 latent z의 정의는 아래와 같습니다.
- data : $x\sim p(x)$
- latents : $z=\{z_t|t\in [0,1]\}$
이 때 diffusion process라고 불리는 forward process를 아래와 같이 정의합니다.
$$\begin{aligned}
q(z_t|x)&=\mathcal{N}(z_t;&\alpha_t\boldsymbol{x}, &\sigma^2_t\boldsymbol{I}) \\
q(z_t|z_s)&=\mathcal{N}(z_t;&(\alpha_t/\alpha_s)\boldsymbol{z_s}, &\sigma^2_{t|s}\boldsymbol{I})\\
&&&\sigma^2_{t|s}=(1-e^{\lambda_t-\lambda_s})\sigma^2_t
\end{aligned}$$
즉, 위 식에서 $0\leq s< t\leq 1$일 때, 미분가능한 함수인 $\alpha_t,\sigma_t$에 대한 noise schedule을 통해 수식을 완성하는 것이 diffusion(forward) process인데, 본 논문에서는 아래의 SNR(Signal-To-Noise Ratio)이 image(t=0.0)에서 noise(t=1.0)으로 갈수록 단조롭게(monotonically) 감소하도록 noise를 scheduling합니다.
$$\lambda_t=log\begin{bmatrix}
\alpha^2_t/\sigma^2_t
\end{bmatrix} $$
이 때, 분산 $\sigma_t$에 대해서는 fixed-form을 사용하지 않고, 아래 식과 같이 NN $\gamma_\eta$을 통한 parametrization으로 예측하도록 했으며,
$$\sigma_t^2=sigmoid(\gamma_\eta(t))$$
만약 VP 모델에서 활용하는 $\alpha_t=\sqrt{1-\sigma^2_t}$의 형태를 활용한다면, 아래와 같이 수식을 나타낼 수도 있습니다.
** sigmoid(logistic) 수식 :$\sigma(z)=\frac{1}{1+e^{-z}}$
$$\begin{aligned}
SNR(t)&=\frac{\alpha_t^2}{\sigma_t^2}\\
&=exp(-\gamma_\eta(t))
\end{aligned}$$
VDM의 forward process에 대해서 추가적으로 아래와 같이 정리됩니다.
- $\alpha_t$와 분산 $\sigma_t$가 각각 diffusion process에 영향을 주는 것이 아니라, $\alpha_t$에 따른 rescaling만 있다면 SNR형태로 영향을 미친다는 것을 수식으로 증명해 보였습니다.
** 추가적으로 $SNR_{max},SNR_{min}$이 같다면 loss도 같아집니다. - 결과적으로 forward process를 score-based에서 사용한 $\mathrm{d}x_t=f(t)x_t\mathrm{d}t+g(t)\mathrm{d}\omega_t$ 형태로 표현했을 때, f(t)와 g(t)는 각각 아래와 같습니다.
$$\begin{aligned}
f(t)&=\frac{\mathrm{d}log\alpha_t}{\mathrm{d}t}\\
g^2(t)&=\frac{\mathrm{d}\sigma^2_t}{\mathrm{d}t}-2\frac{\mathrm{d}log\alpha_t}{\mathrm{d}t}\sigma^2_t\end{aligned}$$
이렇게 forward proces가 정의되면, 우리의 목적은 reverse process를 위해 noise prediction model $\hat{\epsilon}_\theta(z_t;t)$를 학습하는 것이 목표입니다.
하지만, data x 로부터 만들어낸 latent variable $z_t$를 denoise한 denoising model $\hat{x}_\theta(x_t;t)$은 아래와 같이 data prediction model이며, 내부 식 중 노이즈를 예측하는 $\epsilon$-space에서 noise prediction model $\hat{\epsilon}_\theta(z_t;t)$만이 네트워크 입니다.
$$\hat{x}_\theta(z_t;t)=\frac{1}{\alpha_t}(z_t-\sigma_t\hat{\epsilon_\theta}(z_t;t))$$
또한 결과적으로 reverse process SDE는 score-based에서 활용한 식과 동일하겠지만, 기존 정의에서의 $\epsilon_\theta(x_t,t)$는 $\triangledown_xlogq_t(x_t)$가 아닌 scaled score fuunction인 $-\sigma_t\triangledown_xlogq_t(x_t)$를 예측하는 것이 목표이므로 아래와 같이 표현할 수도 있습니다.
\begin{matrix}
dx_t&=[f(t)x_t-g^2(t)&\triangledown_xlogq_t(x_t)&]dt+g(t)d\bar{\omega}_t\\
&=[f(t)x_t-g^2(t)&\frac{1}{\sigma_t}\epsilon_\theta(x_t,t)&]dt+g(t)d\bar{\omega}_t
\end{matrix}
이에 상응하는 reverse process ODE는 아래와 같이 표현할 수도 있겠습니다.
$$dx_t=f(t)x_t+\frac{g^2(t)}{2\sigma_t}\epsilon_\theta(x_t,t)$$
----------------------------------------------------------------
본 논문에서 제안하는 progressive distillation의 방법은 아래와 같은 순서로 진행됩니다.
- 먼저 teacher DM을 기존의 학습법과 같이 얻습니다.
- student DM을 아래와 같은 순서로 학습합니다.
- teacher DM의 파라미터와 모델로 초기화합니다.
- data $x\sim p(x)$로부터 sample을 얻어 noise가 섞인 $z_t$를 얻어냅니다.
- Teacher Model을 통한 2번의 DDIM step을 이용해 target $\tilde{x}(z_t)=[z_t\rightarrow z_{t-(1/N)}]$를 얻어냅니다.
** N을 student의 1번의 sampling step이며 결과 $\tilde{x}(z_t)$를 “target”이라 부릅니다.
** $z_{t}$는 teacher model의 현재 t step에서의 noisy data이며, $z_{t’’}$는 teacher model의 t-(1/N)step에서의 noisy data입니다.
$$\tilde{x}=\frac{z_{t''}-(\sigma_{t''}/\sigma_t)z_t}{\alpha_{t''}-(\sigma_{t''}/\sigma_t)\alpha_t}$$ - Student Model은 target인 $\tilde{x}$를 보고 $\hat{x}_\theta$를 학습해냅니다.
기존의 모델에서는 서로 다른 data point x 가 주어졌을 때, 같은 noisy data $z_t$가 될수도 있으므로, noisy data $z_t$가 주어졌을 때 예측하는 data point x는 blurry한 예측이 될 수 있었습니다.
하지만 이와 반대로 본 논문의 학습법으로는 noisy data $z_t$가 주어졌을 때 학습된 student DM의 target인 $\tilde{x}(z_t)$는 완전히 결정되므로, 이전보다 sharp한 예측이 될 수 있다고 합니다.
이후에 N개의 step으로 학습한 student DM을 teacher로 활용해, N/2 step을 위한 student DM으로 또 distillation할 수 있습니다.
** 이 때의 초기 noise인 $z_1$은 DDPM과 다르게 $\mathcal{N}(0,\boldsymbol{I})$, 즉 SNR 0 인 곳에서 시작하니 기존 모델을 학습할 때와 distillation 모두에서 성능이 더 좋았다고 합니다.
마지막으로, 보통 아래 식과 같이 $\epsilon$-space에서 모델을 학습해 활용하는데 이와 같은 방법은 기존에는 넓은 범위의 SNR에서 학습하는데는 유용하지만 distillation에서는 noise에 가까운 낮은 SNR에서 더 많이 학습하므로 좋지 않다고 합니다.
$$\hat{x}_\theta(z_t;t)=\frac{1}{\alpha_t}(z_t-\sigma_t\hat{\epsilon_\theta}(z_t;t))$$
즉, 위와 같은 상황에서 학습하면 $\epsilon$-space에서의 결과의 작은 변화는 $\alpha_t\rightarrow 0$로 나눠지므로, x-space에서의 denoising model $\hat{x}_\theta(z_t;t)$은 증폭되어 둘 사이의 link가 무너져 distillation 과정에서는 적합하지 않습니다.
따라서 아래와 같이 $\hat{x}_\theta(z_t)$를 다양한 parametrization을 통해 기존의 DM을 학습해보았습니다.
- (sample prediction) $\hat{x}_\theta(z_t)$를 직접 예측하는 방법
- (data prediction) $\tilde{x}_\theta(z_t)$와 $\tilde{\epsilon}_\theta(z_t)$을 각각 예측해 아래와 같은 수식을 통해 $\hat{x}$를 만들어내는 방법
$$\hat{x}=\sigma^2_t\tilde{x}_\theta(z_t)+\alpha_t(z_t-\sigma_t\tilde{\epsilon}_\theta(z_t))$$ - (v-prediction) velocity $v \equiv \alpha_t\epsilon-\sigma_t x$를 예측해 아래와 같은 수식을 통해 $\hat{x}$를 만들어내는 방법
** velocity에 대해 궁금하신 분들은 아래 더보기를 참조하세요
$$\hat{x}=\alpha_tz_t-\sigma_t\hat{v}_\theta(z_t)$$
----------------------------------------------------------------
<Velocity의 정의>
$\sigma_t$와 $\alpha_t$에 대한 parameterization을 SNR대신에 아래와 같이 각도 $\phi_t$로 표현할 수 있습니다.
$$\begin{aligned}
\phi_t&=arctan(\sigma_t/\alpha_t) \\
\sigma_t&=sin(\phi) \\
\alpha_t&=cos(\phi)\\
z_\phi&=cos(\phi)x+sin(\phi)\epsilon
\end{aligned}$$
이제 $z_\phi$의 velocity인 $v_\phi$는 아래와 같습니다.
$$\begin{aligned}
v_\phi&=\frac{dz_\phi}{d\phi}\\
&=cos(\phi)\epsilon-sin(\phi)x\\
\epsilon&=sin(\phi)z_\phi+cos(\phi)v_\phi\\
x&=cos(\phi)z_\phi-sin(\phi)v_\phi
\end{aligned}$$
이 때, 위 그림과 같이 z가 $\phi_1$에서 $\phi_2$로 이동하면, velocity $v_\phi$를 포함한 수식으로 아래와 같이 정리할 수 있고, 결과적으로 velocity를 예측하면 이동하고 난 위치의 z를 찾아낼 수 있다는 것을 알 수 있습니다.
$$z_{\phi_{2}}=cos(\phi_{2}-\phi_{1})z_{\phi_{1}}+sin(\phi_{2}-\phi_{1})\hat{v}_\theta(z_{\phi_{1}})$$
따라서 위와 같은 정의에 따라, 결과적으로 denoising model $\hat{x}_\theta(x_t;t)$에 대한 표현으로 나타내면 아래와 같이 velocity를 통해 denoised 결과를 예측해낼 수 있습니다.
$$\hat{x}=\alpha_tz_t-\sigma_t\hat{v}_\theta(z_t)$$
----------------------------------------------------------------
위와 같은 다양한 parametrization으로 학습한 결과는 아래와 같으며, 모두 좋은 성능을 냈지만 v-prediction는 DDIM의 step size가 정의상 SNR과 독립적인 특징이 있으므로 가장 안정적이었다고 합니다.
또한, reconstruction loss를 아래와 같이 weighting function $w(\lambda_t)$와 함께 학습했는데, 0 SNR을 사용하는 경우 0 weight를 줄 수 있는 가능성이 있어 대안으로 truncated SNR weighting 기법과 SNR+1 weighting기법을 제안하기도 했습니다.
$$\mathbb{E}_{\epsilon, t}=\left[w(\lambda_t)\left\|\hat{x}_\theta(z_t)-x\right\|^2_2\right]$$
실험은 CIFAR-10, ImageNet, LSUN에서 진행했으며, 총 8192steps에서 시작해 perceptual quality를 잃지 않고(FID : 3.0 on CIFAR-10) 4 step까지 줄일 수 있었다고 합니다.
b. Stable Diffusion XL
SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis(arxiv'23)
Blog : https://stability.ai/news/stable-diffusion-sdxl-1-announcement
Weight : https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0
기존의 T2I SD모델에서 개선된 성능의 SD-XL이라는 모델을 선보입니다. 각각의 contribution에 대해 보이기 전에, 먼저 전체 구조는 아래와 같습니다.
two-stage 파이프라인으로 되어있으며 LDM구조를 가지고, latent noise(128x128)로부터 "base 모델"을 통과하고 "refinement 모델"을 활용해 최종 latent를 생성해냅니다.
본 논문의 Contribution은 아래와 같습니다.
- Model Improvement
- Multiple Conditioning
- Multiple Aspect Ratio
- Refinement Model
그럼 하나씩 살펴보겠습니다.
1. Model Improvement
SD-XL는 LDM구조를 따라가, 128x128의 latent space에서 UNet을 통과합니다.
기존에는 Convolutional UNet이 이미지 합성에 dominant하게 사용되었고, 이후에 self attention이 추가되거나 개선된 upscaling 레이어를 활용하는 등 다양한 발전이 이루어졌습니다.
이 중 transformer 연산을 low-level feature에서 활용하는 UNet을 아래 논문을 참조해, UNet내에서 transformer block의 heterogeneous 분포를 활용합니다.
** Simple Diffusion: End-to-end diffusion for high resolution images (arxiv’23)
또한 기존의 LDM의 autoencoder의 locality와 high-frequency detail을 향상시키기 위해 같은 구조를 활용하되,
- 더 큰 batch-size에서(9→256)
- EMA도 활용해
학습한 후 사용합니다.
2. Multiple Conditioning
- Text Conditioning
Text conditioning을 위해 OpenCLIP ViT-bigG와 CLIP-ViT-L text encoder 두가지 활용하며, 마지막 output을 concat해 cross-attention을 활용해 들어가도록 합니다.
- Size Conditioning
기존의 LDM은 학습할 때 최소 이미지 사이즈를 고려해 학습을 진행하므로, 최소 resolution이 안되는 이미지를 버리거나 upscale해서 사용합니다. 하지만 아래와 같은 단점이 존재합니다.
- 이미지를 버리는 경우, 학습 데이터중 39%의 데이터를 버리게 되었습니다.
- 이미지를 upscaling해 사용하는 경우, 마지막 결과가 blurry하게 학습하게 됩니다.
따라서 본 논문에서는 대신에, resize하는 경우 $c_{size}=(h_{orig}, w_{orig})$형태로 기존 이미지의 resolution을 제공하는 size-conditioning을 제공해 학습합니다.
** Fourier feature encoding한뒤, timestep embedding과 더해서 제공합니다.
- Crop Conditioning
기존에, height와 width가 다른 경우 tensor 사이즈를 동일하게 맞추기 위해 사이즈가 작은쪽으로 resize하거나, 사이즈 긴쪽에 대해 random cropping했었는데, 이는 결과적으로 학습된 모델이 생성한 object가 cropped된 것처럼 생성되는 단점이 있습니다.
본 논문에서는 대신에, crop하는 경우, 각각 $c_{top},c_{left}$형태의 coordinates를 embedding해 concat한 뒤 추가적인 파라미터로 넣어줍니다.
** Fourier feature encoding한뒤, timestep embedding과 더해서 제공합니다.
3. Multiple Aspect Ratio
실제 학습에 사용되는 이미지들은 다양한 aspect ratio(landscape 16:9, portrait, 6:16)와 다양한 size를 가지고 있지만 T2I모델에는 주로 512x512 혹은 1024x1024같이 sqaure의 픽셀을 사용하게 됩니다.
따라서 본 논문은 Data Bucketing이라는 방법을 활용합니다.
** https://blog.novelai.net/novelai-improvements-on-stable-diffusion-e10d38db82ac
즉, 데이터를 서로 다른 aspect ratio의 bucket으로 나누어, 학습시 하나의 batch 내에서는 같은 bucket의 이미지끼리 학습이 진행되도록 하며, 매 step 서로 다른 bucket에서 학습이 되도록 합니다.
** bucket을 정할때는, 64배수의 height width를 가지고, 1024x1024의 픽셀을 최대한 유지하도록 합니다.
이 때, bucket마다의 size는 target size $c_{ar}=(h_{tgt},w_{tgt})$형태로 condition을 제공합니다.
** Fourier feature encoding한뒤, timestep embedding과 더해서 제공합니다.
단, pretraining 단계에서는 고정된 aspect ratio & resolution으로 학습하고, 이후 finetuning 단계에서 앞서 제안된 Size Conditioning & Crop Conditioning과 함께 학습합니다.
4. Refinement Model
결과적으로 3단계까지의 학습방법을 포함해 Base Model을 학습하는 단계는 아래와 같습니다.
- Step1. 먼저 Base Model을 내부 데이터셋으로 pretraining합니다.
- Resolution : 256x256 pix
- Batch Size : 2048
- Training : 600,000 steps
- Schedule : Discrete-Time Diffusion Schedule(1000steps)
- +Size-conditioning, Crop-Conditioning
- Step2. 다른 사이즈로 Pretraining을 진행합니다.
- Resolution : 512x512 pix
- Training : 200,000 steps
- Step3. Multi-aspect training을 함께 진행합니다.
- Resolution : ~1024x1024 pix
- + 0.05-level의 offset noise활용
** Offset Noise : $\epsilon\sim\mathcal{N}(\mathbf{0},\mathbf{I})$대신 $\epsilon\sim\mathcal{N}(0.1\delta_c,\mathbf{I})$를 활용하는 것
** Common Diffusion Noise Schedules and Sample Steps are Flawed (WACV’24)
하지만 이렇게 학습을해도 low-local quality의 샘플을 생성하는 문제가 있어, 이를 위해 Refinement Stage를 추가합니다.
즉, latent space에서 SDEdit의 noising-denoising process를 활용해 high-quality, high-resolution의 데이터를 만들어낼 수 있는 분리된 LDM을 학습해 base SDXL에서 생성된 latent를 refine할 수 있도록 하는 것입니다.
Refine단계는 optional이겠지만, 본 논문에서는 꽤나 큰 quality상승을 확인할 수 있었다고 합니다.
c. Stable Diffusion 3.0
Scaling Rectified Flow Transformers for High-Resolution Image Synthesis(ICML'24)
Blog1 : https://stability.ai/news/stable-diffusion-3
Blog2 : https://stability.ai/news/stable-diffusion-3-research-paper
Weight : https://huggingface.co/stabilityai/stable-diffusion-3-medium
Weight : https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers
데이터와 noise를 straight한 선으로 연결하는 모델 formulation을 제안한 Rectified Flow를 도입해 sampling을 개선하고, transformer-based T2I모델을 통해 scaling이 가능하며 text 정보를 더 잘 이해하도록 Stable Diffusion 3.0을 제안합니다.
이 두가지를 하나씩 살펴보도록 하겠습니다.
1. Rectified Flow
기존의 미분가능한 ODE를 통해 풀어나가는 것은 computationally expensive하기 때문에 "직접 vector field를 학습시키는 방법"을 활용하는 Rectified Flow가 제안되었습니다.
** Rectified Flow에 대해 궁금하시면 아래 더보기를 참조하세요
------------------------------------------------------------------
<Rectified Flow>
** Flow Straight And Fast : Learning to generate and transfer data with Rectified Flow(arxiv'22)
$X_0$에서 $X_1$으로 이동하는 ODE를 drift force(velocity) $v$로 표현해 학습하는 방법입니다.
$$\begin{aligned}
\mathrm{d}Z_t&=v(Z_t,t)dt\\
&\underset{v}{min}\int^1_0\mathbb{E}[\left\| (X_1-X_0)-v(X_t,t)\right\|^2]\mathrm{d}t
\end{aligned}$$
이외에도 flow의 crossing한 부분을 줄이기 위해 rewire하거나,
transport cost를 줄이기 위해 straight하게 학습하도록 하는 reflow방법을 추가해 step수를 줄일 수 있음을 보였으며,
더 나아가 k-th rectified flow $Z^k$를 distillation을 통해 더 단축시킬 수 있음을 보였습니다.
Rectified Flow 알고리즘은 아래와 같이 정리할 수 있습니다.
------------------------------------------------------------------
------------------------------------------------------------------
<다양한 formalization >
Rectified Flow를 포함한 다양한 Model Formulation이 존재합니다. 간단히 아래와 같이 정리해보았습니다.
1. Rectified Flow
forward process를 데이터 distribution과 normal distribution사이의 straight한 경로로 정의해 표현합니다
$$z_t=(1-t)x_0+t\epsilon$$
2. EDM
** Elucidating the design space of diffusion-based generative models (arxiv’22)
forward process를 자체적으로 정의한 $b_t$를 활용해 표현했으며, Loss는 자체적으로 정의한 F-prediction을 활용합니다.
$$z_t=x_0+b_t\epsilon$$
3. Cosine
** Improved denoising diffusion probabilistic models (ICML’21)
forward process를 cosine과 sine을 활용해 표현했으며, Loss는 $\epsilon$- 혹은 v-prediction을 활용합니다.
$$z_t=cos(\frac{\pi}{2}t)x_0+sin(\frac{\pi}{2}t)\epsilon$$
4. Linear (LDM)
기존 DDPM VP schedule($\beta_t=\beta_0+\frac{t}{T-1}(\beta_{T-1}-\beta_0)$)을 modify해 $\beta_t=\left(\sqrt{\beta_0}+\frac{t}{T-1}(\sqrt{\beta_{T-1}}-\sqrt{\beta_0})\right)^2$와 같이 schedule합니다.
$$z_t=\sqrt{\bar{\alpha}_t}x_0+\sqrt{1-\bar{\alpha}_t}\epsilon$$
------------------------------------------------------------------
본 논문에서는 Rectified Flow에서 사용하는 vector field $u_t$를 표현하기 위해 아래와 같이 정의합니다
** 아래 두번째 식에서 $\psi'_t(x_0|\epsilon)=a'_tx_0+b'_t\epsilon$과 $\psi^{-1}_t(z|\epsilon)=(z-b_t\epsilon)/a_t$를 활용합니다.
** 아래에서 $\lambda_t:=log{(a^2_t/b^2_t)}$이므로 $\lambda'_t:=2((a'_t/a_t)-(b'_t/b_t))$입니다.
$$\begin{aligned}
&\psi_t(\cdot|\epsilon)&:x_0\rightarrow a_tx_0+b_t\epsilon&&\\
z'_t=&\boldsymbol{u_t}(z|\epsilon)&:=\psi'_t(\psi_t^{-1}(z|\epsilon)|\epsilon)&=\frac{a'_t}{a_t}z_t- b_t(\frac{a'_t}{a_t}-\frac{b'_t}{b_t})\epsilon&=\frac{a'_t}{a_t}z_t-\frac{b_t}{2}\lambda'_t\epsilon
\end{aligned}$$
결과적으로 tractable하기 위해 위 conditional vector field $u_t(z|\epsilon)$를 활용한 CFM(Conditional Flow Matching) Loss는 아래와 같이 정리할 수 있습니다.
$$\begin{aligned}
L_{CFM}&=\mathbb{E}_{t,p_t,p}\left\|v_\theta(z,t)-u_t(z|\epsilon)\right\|^2_2\\
&=\mathbb{E}_{t,p_t,p}(-\frac{b_t}{2}\lambda'_t)^2\left\|\epsilon_\theta(z,t)-\epsilon\right\|^2_2\\
&when\ \epsilon_\theta:=\frac{-2}{\lambda'_tb_t}(v_\theta-\frac{a'_t}{a_t}z)
\end{aligned}$$
또한 기존의 Rectified Flow는 velocity를 예측하기 위한 모델 $v_\theta$를 [0,1]의 timestep에서 uniform하게 학습하지만, 사실 $\epsilon-x_0$ 결과는 [0,1]의 양쪽인 t=0과 t=1에 비해 중간지점에서 더 예측하기 어렵습니다.
따라서 일반적으로는 아래와 같이 density 함수인 $\pi(t)$를 활용해 time distribution을 바꿔줌으로써, 중간 timestep에 좀더 weight를 주어 더 자주 학습될 수 있도록 합니다.
$$w^\pi_t=\frac{t}{1-t}\pi(t)$$
위 density함수 $\pi(t)$는 다양하게 줄 수 있는 방법이 있습니다.
a. Logit-Normal Sampling
중간 step에 힘을 주기 위한 logit-normal distribution으로, location 파라미터 m과 scale파라미터 s로 이루어져 있습니다.
** Logistic-normal distributions: Some properties and uses (’80)
$$\begin{aligned}
\pi_{LN}(t;m,s)&=\frac{1}{s\sqrt{2\pi}(t\cdot (1-t))}exp(-\frac{(logit(t)-m)^2}{2s^2})\\
&when\ logit(t)=log(t/(1-t))
\end{aligned}$$
location 파라미터는 학습 timestep을 데이터 distribution인 $p_0$에 가깝게 하거나(m<0) 노이즈 distribution인 $p_1$에 가깝게 (m>0) 할 수 있습니다.
b. Mode Sampling with Heavy Tails
위 Logit-Normal Sampling은 endpoint인 0, 1에서 density가 없어지는 문제가 있습니다. 따라서 먼저 scale 파라미터 s에 따른 아래와 같은 monotonic한 함수를 정의하고,
$$\begin{aligned}
f_{mode}(u;s)&=1-u-s\cdot(cos^2)(\frac{\pi}{2}u)-1+u\\
&when\ -1\leq s\leq \frac{2}{\pi-2}
\end{aligned}$$
다음과 같은 density함수를 정의해 활용하면 scale 파라미터가 midpoint에 가깝게 학습하거나(s>0) endpoint에 가깝게 학습하는(s<0) sampling을 조절할 수 있습니다.
** s=0일 때는 기존과 같은 Uniform weighting입니다.
$$\pi_{mode}(t;s)=|\frac{d}{dt}f^{-1}_{mode}(t)|$$
c. CosMap
cosine schedule을 RF에서 활용하기 위해 uniform한 input u([0,1])에 의한 timestep t는 아래와 같이 정의됩니다.
$$t=f(u)=1-\frac{1}{tan(\frac{\pi}{2}u)+1}$$
이를 weight로 나타내면 아래와 같습니다.
$$\pi_{CosMap}(t)=|\frac{d}{dt}f^{-1}(t)|=\frac{2}{\pi-2\pi t+2\pi t^2}$$
2. DiT (Diffusion model with Transformers)
SD 3.0은 일반적으로 T2I LDM 구조를 따르지만, UNet이 아닌 DiT기반으로 구성되어 있습니다.
이 DiT의 Text Condition의 경우, 기존 DiT구조와 비슷하게, class Label $c_{vec}$과 timestep을 합친 pooled text representation을 condition으로 modulation을 진행합니다.
하지만 위 정보만으로는 coarse-grained 정보만 처리할 수 있으므로, SDXL과 같이 여러개(본 논문은 3개)의 text-encoder를 통해 sequence representation $c_{ctxt}$ 정보를 만들고, 이도 활용합니다.
위 sequence 정보를 활용할 때는 text 정보에 input image(noisy data)의 embedding을 섞어주는데, 구체적으로는 아래 순서로 진행합니다.
- Step1. Image $x\in\mathbb{R}^{h\times w\times c}$를 encoding해 patch embedding을 만들어줍니다.
- 길이 $\frac{1}{2}\cdot h\cdot\frac{1}{2}\cdot w$의 2x2 사이즈의 patch들을 만들어내고
- 각각의 patching를 embedding화 합니다.
- Step2. Text embedding의 sequence $c_{ctxt}$와 위 patch embedding의 차원을 맞춰줍니다.
- Step3. 두 sequence를 concat해 modulated attention과 MLP를 적용합니다.
** 각각의 modality에 대해서는 구분된 weight를 활용합니다.
그림으로 나타내면 위 3번 과정은 아래 그림중 오른쪽에 느낌표로 표현된 $\bigodot$과 같습니다.
** 그림에서 $\bigodot$는 concatenation을 의미하며 $\ast $는 element-wise multiplication을 의미합니다.
이와 같이 고정된 text representation을 사용하는 것보다 이미지와 text에 대해 학습가능한 각각의 stream을 도입하는 것이 text 이해력을 높일 수 있었다고 합니다.
차원의 경우 "모델의 depth", 즉 "attention block의 개수"를 d라고 했을때, 실험적으로 attention의 hidden size를 64xd로 정했으며, attention의 head의 개수는 d로 정했습니다.
d. Stable Diffusion Turbo
[SD Turbo, SD-XL Turbo] Adversarial Diffusion Distillation (arxiv'23)
Blog : https://stability.ai/research/adversarial-diffusion-distillation
Weight : https://huggingface.co/stabilityai/sd-turbo
Weight-XL : https://huggingface.co/stabilityai/sdxl-turbo
[SD3.0 Turbo] Fast High-Resolution Image Synthesis with Latent Adversarial Diffusion Distillation (arxiv'24)
Weight : https://huggingface.co/spaces/prodia/sd3-turbo
기존의 Consistency Model에서는 ODE trajectory에 대한 consistency regularization과 few-shot setting을 통해 기존의 성능을 유지하면서 단 4 sampling step만으로 생성이 가능한 performance를 보였고,
** Consistency models (ICML’23)
** [LCM] Latent consistency models: Synthesizing high-resolution images with few-step inference (arxiv’23)
Rectified Flow라는 모델이 등장하기도 했지만, 합성된 샘플은 여전히 blurry한 경향이 있었습니다.
이때, SDS(Score Distillation Sampling, Score Jacobian Chaining)는 distillation을 활용해 2D에서 3D Object를 만들어내기 위해 제안되었는데,
** DREAMFUSION: TEXT-TO-3D USING 2D DIFFUSION (arxiv’22)
본 논문은 이 SDS를 fidelity를 위한 adversarial loss와 합쳐 ADD(Adversarial Diffusion Distillation)라는 방법을 제안했습니다.
이를 통해 SDXL의 fidelity를 유지하면서도, 적은 step만을 사용하는 GAN이나 LCM같이 단 1~4step안에 large-scale의 이미지를 생성할 수 있는 SDXL-Turbo를 보였습니다.
먼저 objective의 식과 구조에 대한 그림은 아래와 같습니다.
$$L=L^G_{adv}(\hat{x}_\theta(x_s,s),\phi)+\lambda L_{distill}(\hat{x}_\theta(x_s,s),\psi)$$
학습단계에 대해 이해하는 것이 가장 중요하므로, 아래와 같이 순서대로 살펴보겠습니다.
- 모델을 준비합니다.
- $\theta$의 파라미터를 가지는 student DM, $\hat{x}_\theta(x_s,s)$(pretrained UNet-DM으로 초기화)
- $\phi$의 파라미터를 가지는 discriminator
- $\psi$의 파라미터를 가지는 Teacher DM, $\hat{x}_\psi(\hat{x}_{\theta,t},t)$
- real 이미지 $x_0$를 활용해 noised data를 찾아냅니다.
** student DM의 $\alpha_s$와 $\sigma_s$는 Teacher DM의 coefficient를 활용합니다.
** s는 zero-terminal SNR로 schedule된 $T_{student}=\left\{\tau_1,\dots,\tau_n\right\}, \tau_n=1000$에서 N=4개의 timestep을 uniform하게 뽑아 활용합니다.
** Common Diffusion Noise Schedules and Sample Steps are Flawed(WACV’24)
$$x_s=\alpha_xx_0+\sigma_s\epsilon$$ - discriminator는 $x_0$와 student DM이 만들어낸 $\hat{x}_\theta$를 구분해내고, adversarial loss를 만들어냅니다.
- Student DM이 만들어낸 $\hat{x}_\theta$에서 Teacher DM의 forward process를 통해 $\hat{x}_{\theta,t}$를 만들어냅니다.
- Teacher DM은 $\hat{x}_\psi$를 만들어내고, Student DM이 만들어낸 $\hat{x}_\theta$과 비교해 distillation loss를 만들어냅니다.
** Teacher와 Student가 같은 LDM구조를 활용한다면, distillation loss를 latent space에서 측정할 수도 있지만 pixel space에서 하는 것이 더욱 안정된 gradients를 만들어낼 수 있다고 합니다.
그럼 위 과정에서 등장한 Adversarial Loss와 Distillation Loss를 조금 더 깊게 보겠습니다.
a. Adversarial Loss
다른 distillation방법에서 발견되는 것과 같은 blurry함이나 특이한 artifact 없이 적은 step으로 생성할 수 있도록 하기 위해 Discriminator 모델을 활용한 Adversarial Loss를 활용합니다. 이 때 사용하는 두 모델은 아래와 같이 정리됩니다.
- Discriminator 모델 : 아래와 같이 구성됩니다.
- Pretrained feature network $F$ : ViT (DINOv2, CLIP)
- Discriminator heads $D_{\phi,k}$ (Trainable)
** feature network의 서로 다른 layer feature에 적용 - Condition(Projection) : text embedding(CLIP-ViT-g-14), given image embedding(DINOv2 ViT-L)
- Adversarial Objective Function : hinge loss
** R1 : R1 gradient penalty of discriminator
$$\begin{aligned}
L^D_{adv}(\hat{x}_\theta(x_s,s),\phi)&=\mathbb{E}_{x_0}[\sum_kmax(0,1-D_{\phi,k}(F_k(x_0)))+\gamma R1(\phi)]\\
&+\mathbb{E}_{{\color{red}\hat{x}_\theta}}[\sum_kmax(0,1+D_{\phi,k}({\color{red}F_k(\hat{x}_\theta)}))]
\end{aligned}$$
- Student DM 모델 : 아래와 같이 구성됩니다.
- Adversarial Objective Function : hinge loss
$$L^D_{adv}(\hat{x}_\theta(x_s,s),\phi)=-\mathbb{E}_{s,\epsilon,x_0}[\sum_kD_{\phi,k}({\color{red}F_k(\hat{x}_\theta(x_s,s))})]$$
- Adversarial Objective Function : hinge loss
b. Distillation Loss
또한 pretrained DM을 Teacher로 사용했을 때, compositionality를 보존하고 생성 이미지의 quality를 보존하기 위해 Distillation Loss를 활용합니다.
SDS Loss라는 것을 활용하는데, 식은 아래와 같습니다.
** d() : distance metric, 여기서는 L2 norm을 활용
** c() : weighting function, 여기서는 (exponential weighting, SDS weighting) 중 선택
$$L_{distill}(\hat{x}_\theta(x_s,s),\psi)=\mathbb{E}_{t,e'}[c(t)d({\color{red}\hat{x}_\theta}, \hat{x}_{\psi};t)]$$
위에 추가로 NFSD(Noise-Free Score Distillation) 또한 사용합니다.
** Noise-free score distillation (arxiv’23)
이와 같은 ADD를 SD-XL에 적용해 SD-XL Turbo가 등장합니다.
이후에는 distillation 과정에서 image space로 바꾸기 위한 decoding을 제거한 LADD(Latent ADD)도 등장해 SD3.0에 적용되어 SD3.0-Turbo가 만들어집니다.
즉, 아래 그림과 같이 pixel space에서 distillation loss와 adversarial loss(w/ DINOv2 feature)를 활용하던 ADD와 달리, latent space에서 학습을 진행합니다.
학습하는 과정을 정리하면 아래와 같습니다.
- Teacher 모델을 활용해 latent를 생성해냅니다.
- 위 결과를 time step $\hat{t}$에 대해 renoise합니다.
- 2번의 결과를 Student모델을 활용해 적용해 $\hat{x}_\theta(x_t,t)$를 얻어냅니다.
- 2번의 결과를 Teacher모델을 활용해 다시 적용하면서, 각 attention 블록 뒤의 token sequence들을 모아둡니다.
- 3, 4번의 결과 중, token sequence 각각에 대해 독립적인 discriminator head를 적용합니다.
** discriminator의 condition은 noise-level과 CLIP embedding을 적용했습니다.
** discriminator의 feature network는 MAR(Multi-Aspect Ratio) setting을 위해 1D conv대신 token sequence를 기존의 spatial layout으로 다시 바꾸어 2D conv를 적용했습니다.
위와 같은 과정을 적용하면 아래와 같은 이점이 있습니다.
- 장점1. Discriminator와 Teacher Model의 Unification : discriminative(classification, self-supervised objective)의 pretraining task로 학습된 feature network 대신 generative feature를 활용하면 아래와 같은 장점이 있습니다.
- image space로 decoding하지 않아도 돼서 메모리와 구조가 효율적입니다.
- diffusion training에서 noise에 따른 loss weighting을 하는데 있어, noise에 대한 파라미터를 적용해 학습할 수 있습니다.
- Teacher Model이 MAR(Multi-Aspect Ratio)에서 학습된 만큼, 관련있는 셋팅이 가능합니다.
- 장점2. Synthetic data를 학습에 활용 : 적은 step에서 CFG를 활용해 condition을 주는 경우, text에 align되기보다는 over saturation되는 현상이 발생합니다.
하지만 학습시에 고정된 CFG를 활용해 Teacher Model이 생성한 synthetic data를 활용하면, image-text align을 고려해 distilling이 가능하다고합니다.
** 추가적으로, 위 같은 경우를 방지하기 위해 1-shot보다는 multi-step이 좋았으며, dynamic thresholding도 이런 현상을 완화할 수 있다고 합니다.
결과적으로 LADD를 SD3에 적용해 4 sampling step만에 높은 퀄리티의 multi-aspect megapixel 이미지를 생성가능한 SD3-turbo를 얻었습니다.
** 추가적으로 preference alignment를 활용하는 DPO(Direct Preference Optimization)를 LoRA를 활용해 학습하니 성능이 더 좋았다고 합니다.
2. Optional Improvements
위에서 살펴본 것과 같이 모델의 큰 특징적인 발전외에도 Optional하게 더 좋은 결과를 위한 다양한 논문들이 등장했습니다.
Diffusers 패키지를 활용해 구현할 때 활용할 수 있는 다양한 옵션들을 논문과 함께 살펴보겠습니다.
a. EMA(Exponential Moving Average)
EMA는 기존 모델 weight의 moving average을 유지하는 weight 및 방법으로, Stable Diffusion 2.0에서 사용되었습니다.
각각의 time step에서 EMA 모델을 업데이트하는 경우를 예로 들어보겠습니다.
decay factor(smoothing) 0.99인 경우, (0.99의 current EMA 모델 + 0.01의 새로운 weight)로 EMA모델을 업데이트 해 줍니다.
$$W^{t+1}_{ema\ model}=smoothing\cdot w^t_{ema\ model}+(1-smoothing)\cdot w^t_{model}$$
하지만 위와 같이 모든 weight를 Read/Write해야하므로 많은 메모리 operation 때문에 느려질 위험이 있기도 합니다.
따라서, 이런 EMA모델을 활용하는 방법에서 최종 결과에 영향을 미치는건 사실 최종 몇 step뿐이므로, 학습이 끝나기 전 몇 %만(~3.5%~) EMA를 적용하기도 합니다.
EMA모델은 아래와 같은 두가지 장점을 가집니다.
- (학습시) 모델의 Generalization을 향상시키며, 학습이 안정화됩니다.
- (생성시) 역시나 Generalization이 더 잘되어있고 안정적이며, 일반적인 파라미터보다 성능이 조금 더 좋은 것으로 알려져있습니다.
이렇게 저장된 결과 weight를 EMA weight라고 부르며 아래와 같이 일반적인 weight와 다릅니다.
- sd-v1-4.ckpt : 파라미터가 저장되어있으며, fine-tuning 등 학습시에 필요합니다.
- sd-v1-4-full-ema.ckpt : 파라미터의 running average가 저장되어있으며, inference시에 안정적인 결과를 냅니다.
EMA는 ema_pytorch를 활용하면 쉽게 구현할 수 있는데, 아래 "pytorch lightning에서 training wrapper을 만들 때"를 예로 보여 ema model을 어떻게 선언하고 관리하는지를 살펴볼 수 있습니다.
from ema_pytorch import EMA
from safetensors.torch import save_file
import torch
class TrainingWrapper(pl.LightningModule):
def __init__(
self,
model: MyModel,
use_ema: bool = True,
):
super().__init__()
self.my_model = model
if use_ema:
self.my_model_ema = EMA(
self.my_model.model,
beta=0.9999,
power=3/4,
update_every=1,
update_after_step=1,
include_online_model=False
)
else:
self.my_model_ema = None
def configure_optimizers(self):
pass;
def training_step(self, batch, batch_idx):
pass;
def on_before_zero_grad(self, *args, **kwargs):
if self.my_model_ema is not None:
self.my_model_ema.update()
def export_model(self, path, use_safetensors=False):
if self.my_model_ema is not None:
self.my_model.model = self.my_model_ema.ema_model
if use_safetensors:
save_file(self.my_model.state_dict(), path)
else:
torch.save({"state_dict": self.my_model.state_dict()}, path)b. X-formers
https://github.com/facebookresearch/xformers
X-formers는 PyTorch-based 파이썬 패키지로, Transformers의 다양한 block들(특히 Attention Block)이 더 빠른 속도와 적은 메모리를 소비하도록 구현되었습니다.
Geforce 1000이상의 GPU시리즈에서 활용이 가능하며 생성환경에서 활용하면 속도를 높여 생성시간을 단축시킬 수 있습니다.
다양한 컴포넌트들이 CUDA kernel로 구현되어 있는데, 이 중 대표적으로 활용되는 것이 바로 Flash Attention입니다.
** FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (NIPS'22)
** FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (arxiv'23)
** FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (arxiv'24)
** https://github.com/Dao-AILab/flash-attention
그럼 예를 들어 Flash Attention만 한번 깊게 살펴보겠습니다.
기존에 아래와 같은 많은 approximation 방법들을 통해 attention의 compute & memory 요구사항을 줄이고자 노력했지만, wall-clock speedup은 보이지 못해왔습니다.
** wall-clock time : 일을 수행하는 데 걸리는 시간
** cpu execution time : 일을 수행하는 데 걸린 cpu만의 시간
- approximation 방법1. sparse-approximation
- approximation 방법2. low-rank approximation
이는 대부분의 방법이 wall-clock time과 무관한 FLOP(FLoating point OPeration)s를 줄이는 것에 집중하며, memory access IO를 무시했기 때문이라고 합니다.
따라서 아래와 같은 서로 다른 메모리에 대한 read/write를 고려한 IO-Aware attention을 제안했습니다.
- (fast) GPU on-chip SRAM
- (slower) GPU HBM(High Bandwidth Memory)
이를 위해서 아래와 같은 요구조건을 만족시킴으로써 attention matrix을 활용하기 위한 HBM read/write를 줄였고, 9배의 적은 HBM memory access를 통해 학습을 빠르게 할 뿐 아니라, 이를 통해 context를 길게 함으로써 quality도 증가시킬 수 있다고 합니다.
** 결과적으로 기존 $\Omega(Nd+N^2)$ HBM access를 $\Omega(N^2d^2M^{-1})$까지 줄였습니다. 이때 d는 head dimension, N는 sequence length, M은 SRAM 사이즈입니다.
- 조건1. softmax reduction을 전체 input에 대한 접근 없이 할 수 있어야 합니다.
** Softmax 함수는 output이 0~1사이의 값으로 모두 정규화되며, output의 총합은 항상 1이 되는 특성을 가지는 함수입니다.
$$f(z)_i=\frac{e^{z_i}}{\sum^K_{j=1}e^{z_j}}$$ - 조건2. 큰 크기의 attention matrix는 backward pass에서 저장해두면 안됩니다.
위 조건을 만족시키기 위해 구체적으로는 아래와 같은 방법을 사용합니다.
- Tiling : input을 block으로 나누어 softmax reduction연산을 효율적으로 진행합니다.
- Recomputing : sotmax의 normalization factor를 forward pass에 저장해두고, backward pass에서 attention을 on-chip에서 recompute합니다.
이런 방법은 중간 attention matrix를 HBM으로 읽어오는 것보다 FLOPs는 증가할 수 있지만 더 빠릅니다. (GPT2에서 7.6배 정도 빠릅니다.)
즉, 아래 그림을 보면..
- NxN의 attention matrix가 있을 때, outer loop(빨간선)에 따라 tiling을 하고 (K,V) matrix block 각각을 SRAM(빠름)에 올립니다.
- 이후에 각각의 block에서는 inner loop(파란선)에 따라 Q를 SRAM으로 load하며 결과를 HBM을 통해 바로 입력해줍니다.
Flash Attention은 패키지가 구현되어 있어 직접 PyTorch 모델 모듈에 포함시켜 아래와 같이 구현할 수 있습니다
pip3 install flash-attnfrom flash_attn import flash_attn_func, flash_attn_kvpacked_func
# Flash Attention 2 requires FP16 inputs
fa_dtype_in = q.dtype
q, k, v = map(lambda t: rearrange(t, 'b h n d -> b n h d').to(torch.float16), (q, k, v))
out = flash_attn_func(q, k, v, causal = causal)
out = rearrange(out.to(fa_dtype_in), 'b n h d -> b h n d')
그리고 Diffusers 패키지에서 활용하기 위해서는 아래와 같이 Xformers를 활성화해 사용할 수 있습니다.
from diffusers import DiffusionPipeline
import torch
pipe = DiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")
pipe.enable_xformers_memory_efficient_attention()
with torch.inference_mode():
sample = pipe("a small cat")
# optional: You can disable it via
# pipe.disable_xformers_memory_efficient_attention()c. DPM-Solver++
DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models (arxiv'22)
기존의 DDIM과 같은 1st order ODE solver는 100~250step의 적은 step으로 guidance sampling에 많이 활용됩니다.
이후에 제안된 아래 논문의 high-order diffusion ODE solver는 더 적은 sampling으로 속도를 개선할 수 있음이 보여졌지만 guidance sampling을 위해 아래와 같은 몇가지 challenge가 있었습니다.
** Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps (arxiv’22)
** [DEIS] Fast sampling of diffusion models with exponential integrator (arxiv’22)
- 아래 그림과 같이 guidance scale이 크면 instable해집니다.
- train-test mismatch 문제가 발생합니다.
즉, 원래라면 정해진 bound 안에 존재하는 image data([-1,1])가 gudiance scale로 인해 true noise에서 멀어진 noise를 통해 정해진 bound밖으로 옮겨져 saturated 혹은 unnatural한 결과를 야기합니다.
** DPM Solver논문에 대해 궁금하시면 아래 더보기를 참조하세요
----------------------------------------------------------------
<DPM Solver>
Dpm-solver : A fast ode solver for diffusion probabilistic model sampling in around 10 steps (arxiv'22)
scaled score function을 반영한 reverse ODE의 식은 아래와 같았는데, 아래 식 중 $f(t)x_t$는 $x_t$에 대해 linear하지만 $\frac{g^2(t)}{2\sigma_t}\epsilon_\theta(x_t,t)$는 $x_t$에 대해 non-linear합니다.
$$dx_t=f(t)x_t+\frac{g^2(t)}{2\sigma_t}\epsilon_\theta(x_t,t)$$
따라서 이런 ODE를 semi-linear ODE라고 할 수 있는데, 기존의 score-based논문에서처럼 non-linear함을 무시하고 ODE를 적용하면 discretization error가 발생한다고 합니다.
따라서 위 ODE대신 아래와 같은 "variation of constants"를 활용한 semi-linear ODE를 활용해 나타내고, VDM에서의 f(t)와 g(t) 정의를 log-SNR을 활용해 정리하면 아래와 같이 나타낼 수 있습니다.
** Numerical solution of ordinary differential equations (Oxford'11)
$$\begin{aligned}
x_t&=e^{\int^t_sf(\tau)d\tau}x_s+\int^t_s\left(e^{\int^t_\tau f(r)dr}\frac{g^2(\tau)}{2\sigma_\tau}\epsilon_\theta(x_\tau,\tau)\right)d\tau \\
&=\frac{\alpha_t}{\alpha_s}x_s-\alpha_t\int^{\lambda_t}_{\lambda_s}e^{-\lambda}\hat{\epsilon}_\theta(\hat{x}_\lambda,\lambda)d\lambda\\
&when\ \lambda_t:=log(\frac{\alpha_t}{\sigma_t})
\end{aligned}$$
즉, $\hat{\epsilon}_\theta$의 exponentially weighted integral을 활용해 diffusion ODE의 solution으로 사용하면, linear term의 aprroximation error를 cancel하기 때문에, 결과적으로 truncation error를 줄여 더욱 빠르고 정확한 solver를 제안할 수 있습니다.
이제 위 수식과 "exponential integrator" 방법을 활용해 high-order solver를 제안합니다.
** Exponential integrators (Acta Numer'10)
즉, $ti(=t_{i})$에서 $ti+(=t_{i+1})$로 이동한다고 할 때, $\hat{\epsilon}^{(n)}_\theta(\hat{x}_\lambda,\lambda)$의 $\lambda_{ti}$에서 $\lambda$에 대한 (k-1)-th order Taylor Expansion을 활용해 위 식을 정리하면 아래와 같습니다.
** $\hat{\epsilon}^{(n)}_\theta$는 $\hat{\epsilon}_\theta$를 $\lambda$로 n번 미분한 값을 의미합니다.
**$O(h_i^{k+1})$는 kth order에 대한 error term
** $h_i$는 $\lambda_{ti+}-\lambda_{ti}$를 의미
$$x_{ti\rightarrow ti+}=\frac{\alpha_{ti+}}{\alpha_{ti}}\tilde{x}_{ti}-\alpha_{ti+}\sum^{k-1}_{n=0}\underbrace{\hat{\epsilon}^{(n)}_\theta(\hat{x}_{\lambda_{ti}},\lambda_{ti})}_{\text{estimated}}\underbrace{\int^{\lambda_{ti+}}_{\lambda_{ti}}e^{-\lambda}\frac{(\lambda-\lambda_{ti})^n}{n!}d\lambda}_{\text{analytically computed}}+\underbrace{O(h_i^{k+1})}_{\text{omitted}}$$
이제, 위 식에서 $O(h_i^{k+1})$를 drop하고 stiff order condition을 활용해 $\hat{\epsilon}^{(n)}_\theta\ (when\ n\leq k-1)$만 approximate하면 k-th order ODE solver(DPM Solver-k)를 만들 수 있습니다.
위 식의 상황에서 예를 들어 k=1인 경우 아래와 같고, 이를 DPM-Solver-1이라고 부릅니다.
$$\begin{aligned}
x_{ti\rightarrow ti+}&=\frac{\alpha_{ti+}}{\alpha_{ti}}\tilde{x}_{ti}-\alpha_{ti+}\epsilon_\theta(\hat{x}_{\lambda_{ti}},\lambda_{ti})\int^{\lambda_{ti+}}_{\lambda_{ti}}e^{-\lambda}d\lambda+O(h_i^{2})\\
&=\frac{\alpha_{ti+}}{\alpha_{ti}}\tilde{x}_{ti}-\sigma_{ti+}(e^{h_i}-1)\epsilon_\theta(\hat{x}_{\lambda_{ti}},\lambda_{ti})+O(h_i^{2})
\end{aligned}$$
다음 예로 k=2인 경우 아래와 같고, 이를 DPM-Solver-2라고 부릅니다. 결국 $\bar{u}$라는 중간값을 구한 뒤, $\bar{x}_t$로 향하는 것이 특징입니다.
** $t_\lambda()$는 $\lambda(t)$함수의 inverse function입니다.
$$\begin{aligned}
&0<t<s<T, h:=\lambda_t-\lambda_s,r_1=0.5\\
s_1&=t_\lambda(\lambda_s+r_1h)\\
\bar{u}&=\frac{\alpha_{s1}}{\alpha_s}\boldsymbol{x}_s-\sigma_{s1}(e^{r_1h}-1)\epsilon_\theta(\boldsymbol{x}_s,s)\\
\boldsymbol{D}&=(1-\frac{1}{2r_1})\epsilon_\theta(\boldsymbol{x}_s,s)+\frac{1}{2r_1}\epsilon_\theta(\bar{u},s_1)\\
\bar{\boldsymbol{x}}_t&=\frac{\alpha_{t}}{\alpha_s}\boldsymbol{x}_s-\sigma_t(e^h-1)\boldsymbol{D}
\end{aligned}$$
본 논문에서는, 기존에 제안된 DDIM은 DPM Solver-1과 동일하며, RK(Runge-Kutta) method의 경우는 non-linearity를 반영한 DPM Solver의 semi-linear ODE와 비교해보면 step size가 클수록 discretization error가 많이 발생할 수 밖에 없고, unstable하다고 합니다.
** RK(Runge-Kutta) method는 two-stage 2nd order method입니다.
결과적으로 DPM solver는 학습없이도 discrete & continuous time에 모두에 적합하며, high quality의 sample을 10~20 step만으로 생성이 가능하다고 합니다.
----------------------------------------------------------------
본 논문에서는 직관적으로 위의 원인이 guidance scale이 커질 수록 output과 $\hat{\epsilon}_\theta$의 미분값 모두 증폭되기 때문이라고 합니다.
$$\begin{aligned}
\tilde{\epsilon}_\theta(x_t,t,c)_{CG}&:=&\epsilon_\theta(x_t,t)&-s\cdot \sigma_t\bigtriangledown x_tlogp_\phi(c|x_t,t)\\
\tilde{\epsilon}_\theta(x_t,t,c)_{CFG}&:=s\cdot &\epsilon_\theta(x_t,t,c)&-(1-s)\cdot \epsilon_\theta(x_t,t,\oslash )
\end{aligned}$$
즉, derivative order가 증가할수록 convergence radius가 좁아져 더 작은 step size를 요구하게 되거나, ODE solver가 수렴해갈 convergence range가 기존의 데이터 data distribution을 향하는 것이 아니게 되는 것입니다.
따라서 본 논문은 training-free한 ODE를
- Data Prediction Model 활용
- Multi-Step 활용
- Dynamic Thresholding 활용
의 방법을 통해 위 문제를 해결해 DPM-Solver++를 제안하고, 이는 guided sampling에서도 high quality sample을 15step으로 만들어내 속도를 빠르게 했습니다.
1. Data Prediction Model
기존의 DPM-Solver는 exponential integrator를 기반으로한 ODE solver로, 아래와 같은 수식을 소개했었습니다.
$$\begin{aligned}
x_t=\frac{\alpha_t}{\alpha_s}x_s-\alpha_t\int^{\lambda_t}_{\lambda_s}e^{-\lambda}\hat{\epsilon}_\theta(\hat{x}_\lambda,\lambda)d\lambda\\
when\ \lambda_t:=log(\frac{\alpha_t}{\sigma_t})
\end{aligned} $$
하지만 본 논문은 Data Prediction Model을 활용하기 위해 위 식을 아래와 같이 수정합니다.
$$\begin{aligned}
x_t=\frac{\sigma_t}{\sigma_s}x_s-\sigma_t\int^{\lambda_t}_{\lambda_s}e^{-\lambda}\hat{x}_\theta(\hat{x}_\lambda,\lambda)d\lambda\\
when\ \lambda_t:=log(\frac{\alpha_t}{\sigma_t})
\end{aligned} $$
** Data Prediction Model에 대한 정보는 아래 더보기를 참조하세요
----------------------------------------------------------------
<Prediction Model Parameterization>
기본적으로 DPM은 noisy input $x_T$에서 데이터 $x_0$를 recover하기 위한 denoising을 학습하는 방법입니다.
하지만 목적은 같더라도 아래와 같이 "모델이 어떤 것을 예측하는 것이 목표인지"에 따라 다양한 Parametrization이 가능합니다.
1. Noise(Epsilon, Score) Prediction Model : DDPM, DDIM, Score-based
가장 기본적인 모델로 noisy 데이터 $x_t$에 포함된 noise를 예측하기 위해 $\epsilon_\theta(x_t,t)$를 학습합니다.
$$\underset{\theta}{min}\mathbb{E}_{x_0,\mathbf{\epsilon},t}[\omega(t)||\epsilon_\theta(x_t,t)-\epsilon||^2_2]$$
diffusion ODE를 활용한 sampling은 아래와 같은 ODE를 discretization함으로써 진행합니다.
$$\frac{dx_t}{dt}=f(t)x_t+\frac{g^2(t)}{2\sigma_t}\epsilon_\theta(x_t,t),x_T\sim N(0,\tilde{\sigma}^2I)$$
2-1. Data Prediction Model : VDM
위와 다르게 $x_t$에 포함된 noise가 제거된 $x_0$를 직접적으로 예측하기 위한 denoiser function $x_\theta(x_t,t)$를 학습합니다.
실제로는 위와 같이 noise를 학습하는 것이지만, 예를 들어 VDM의 경우 아래와 같이 residual $x_t$와 normalization기법들이 포함됩니다.
$$x_\theta(x_t,t):=\frac{1}{\alpha_t}(x_t-\sigma_t\epsilon_\theta(x_t,t)) $$
diffusion ODE를 활용한 sampling은 아래와 같은 ODE를 discretization함으로써 진행합니다. (위 식과 사실상 같은 식)
$$\frac{dx_t}{dt}=(f(t)+\frac{g^2(t)}{2\sigma_t})x_t+\alpha_t\frac{g^2(t)}{2\sigma_t}x_\theta(x_t,t),x_T\sim N(0,\tilde{\sigma}^2I)$$
2-2. F-Prediction Model : EDM
** Elucidating the Design Space of Diffusion-Based Generative Models (NIPS'22)
$$\begin{align}
D_\theta(x,\sigma) &= x - \sigma F_\theta(x,\sigma)\\
&\simeq c_{\text {skip }}(\sigma) x+c_{\text {out }}(\sigma) F_{\theta}\left(c_{\text {in }}(\sigma) x ; c_{\text {noise }}(\sigma)\right)\\
F_{\text {target }}(\boldsymbol{y}, \boldsymbol{n} ; \sigma)&=\frac{1}{c_{\text {out }}(\sigma)}\left(\boldsymbol{y}-c_{\text {skip }}(\sigma)(\boldsymbol{y}+\boldsymbol{n})\right)(\because x=y+n)
\end{align}$$
3-1. V-Prediction Model : V-Diffusion
** Progressive Distillation for Fast Sampling of Diffusion Models (arxiv'22)
$$\hat{x}=\alpha_tz_t-\sigma_t\hat{v}_\theta(z_t)$$
4. Rectified Flow : Rectified Flow
** Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow(arxiv'22)
$$z_t=(1-t)x_0+t\epsilon$$
5. Sample Prediction Model : $x_0$를 직접 예측하는 모델 $x_\theta(x_t,t)$를 활용하는 것
결국 5번을 제외하고는 이론적으로는 동일한 식이지만 실험적인 결과에 영향을 미칩니다.
----------------------------------------------------------------
이제 DPM Solver에서와 같이 $ti(=t_{i})$에서 $ti+(=t_{i+1})$로 이동한다고 할 때, $\hat{\epsilon}^{(n)}_\theta(\hat{x}_\lambda,\lambda)$의 $\lambda_{ti}$에서 $\lambda$에 대한 (k-1)-th order Taylor Expansion을 활용해 위 식을 정리하면 아래와 같습니다.
** $\hat{\boldsymbol{x}}^{(n)}_\theta$는 $\hat{\boldsymbol{x}}_\theta$를 $\lambda$로 n번 미분한 값을 의미합니다.
** $O(h_i^{k+1})$는 kth order에 대한 error term
$$\tilde{x}_{ti+}={\color{red}\frac{\sigma_{ti+}}{\sigma_{ti}}}\tilde{x}_{ti}+{\color{red}\sigma_{ti+}}\sum^{k-1}_{n=0}\underbrace{{\color{red}x_\theta^{(n)}}(\hat{x}_{\lambda_{ti}},\lambda_{ti})}_{\text{estimated}}\underbrace{\int^{\lambda_{ti+}}_{\lambda_{ti}}e^\lambda\frac{(\lambda-\lambda_{ti})^n}{n!}d\lambda}_{\text{analytically computed}}+\underbrace{O(h_i^{k+1})}_{\text{omitted}}$$
역시나 DPM solver와 같이 $O(h_i^{k+1})$를 drop하고 stiff order condition을 활용해 $\hat{x}^{(n)}_\theta\ (when\ n\leq k-1)$만 approximate하면 k-th order ODE solver (DPM Solver++k)를 만들 수 있습니다.
$k=1$일 때, DPM solver처럼 DDIM과 같으며
$k=2$일 때의 DPM-Solver(2S)는 아래와 같고, 기존 DPM-Solver보다 더욱 constant한 경향이 적어 discretization error가 적었다고 합니다.
$$\begin{aligned}
&0<t_{i-1}<t_i<T, h_i:=\lambda_{t_i}-\lambda_{t_{t-1}}\\
r_i&=\frac{\lambda_{s_i}-\lambda_{t_{i-1}}}{h_i}\\
\boldsymbol{u}_i&=\frac{\sigma_{s_i}}{\sigma_{t_{i-1}}}\tilde{\boldsymbol{x}}_{t_{i-1}}-\alpha_{s_i}(e^{-r_ih_i}-1)\boldsymbol{x}_\theta(\tilde{\boldsymbol{x}}_{t_{i-1}},t_{i-1})\\
\boldsymbol{D}_i&=(1-\frac{1}{2r_i})\boldsymbol{x}_\theta(\tilde{\boldsymbol{x}}_{t_{i-1}},t_{i-1})+\frac{1}{2r_i}\boldsymbol{x}_\theta(\boldsymbol{u}_i,s_i)\\
\tilde{\boldsymbol{x}}_{t_i}&=\frac{\sigma_{t_i}}{\sigma_{t_{i-1}}}\tilde{\boldsymbol{x}}_{t_{i-1}}-\alpha_{t_i}(e^{-h_i}-1)\boldsymbol{D}_i
\end{aligned}$$
마지막으로 $k\geq 3$일 때의 high-order solver는 큰 guidance scale에 적합하지 않으므로 본 논문에서는 사용하지 않았다고 합니다.
2. Multi-Step
위 DPM-Solver(2S)에서는 $t_{i-1}$에서 $t_i$로 이동하는 step에서는 제안된 step solver가 두개의 연속적인 함수가 필요하며, 이는 ODE solver에서 일반적인 single step방법이라고 불립니다.
** Numerical solution of ordinary differential equations(’11)
하지만, 이 과정에서의 중간 값인 $\boldsymbol{u}_i$가 한번만 사용되기 때문에, 이전의 정보가 충분히 활용되지 않아 효율적이지 않다고 합니다.
따라서 이전 정보를 잘 활용하기 위해 기존의 ODE solution에서 활용되는 또다른 mainstream 접근인 multistep방법을 활용하면 아래와 같습니다.
** Numerical solution of ordinary differential equations(’11)
$$\begin{aligned}
&0<t_{i-1}<t_i<T, h_i:=\lambda_{t_i}-\lambda_{t_{t-1}}\\
r_i&=\frac{h_{i-1}}{h_i}\\
\boldsymbol{D}_i&=(1-\frac{1}{2r_i})\boldsymbol{x}_\theta(\tilde{\boldsymbol{x}}_{t_{i-1}}, t_{i-1})-\frac{1}{2r_i}\boldsymbol{x}_\theta({\color{red}\tilde{\boldsymbol{x}}_{t_{i-2}}}, t_{i-2})\\
\tilde{\boldsymbol{x}}_{t_i}&=\frac{\sigma_{t_i}}{\sigma_{t_{i-1}}}\tilde{\boldsymbol{x}}_{t_{i-1}}-\alpha_{t_i}(e^{-h_i}-1)\boldsymbol{D}_i
\end{aligned}$$
즉, 이전의 정보인 $\tilde{\boldsymbol{x}}_{t_{i-1}}$에 중간정보 $\boldsymbol{u}_i$ 대신 $\tilde{\boldsymbol{x}}_{t_{i-2}}$를 활용하는 것입니다.
이런식으로 하면, 전체 step이 N이라고 했을 때 single step에서는 실질적으로 $N*k$step을 진행했던 반면 multi step에서는 $N$step만 진행하면 되므로 step수가 적어지고, error term인 $O(h_i^k)$도 줄어들 수 있습니다.
3. Dynamic Thresholding
높은 guidance scale 때문에 sample이 out-of-bound가 되는 것을 방지하기 위해서 기존에 DDPM과 Imagen에서 활용했던 Thresholding방법을 활용합니다.
즉, 아래 식을 활용해 data를 예측했을 때, 결과적인 $\hat{\boldsymbol{x}}_\theta(\boldsymbol{x}_t,t,c)$는 element-wise하게 clipping을 해주었다는 것입니다.
$$x_\theta(x_t,t):=\frac{1}{\alpha_t}(x_t-\sigma_t\epsilon_\theta(x_t,t)) $$
여기까지 reverse ODE의 최적화를 위해 진행했다면, 이를 SDE로 확장하기도 합니다. 기존의 score-based에서 활용하는 reverse SDE의 수식이 아래와 같았다면,
$$\mathrm{d}\boldsymbol{x}_t=[f(t)\boldsymbol{x}_t+\frac{g^2(t)}{\sigma_t}\boldsymbol{\epsilon}_\theta(\boldsymbol{x}_t,t)]\mathrm{d}t+g(t)\mathrm{d}\boldsymbol{\bar{\omega}}_t$$
앞서 진행했던 variation-of-constants 수식을 활용해 $0<t<s<T$일 때, 아래와 같이 얻어 낼 수 있습니다.
$$\boldsymbol{x}_t=\frac{\sigma_t}{\sigma_s}e^{-(\lambda_t-\lambda_s)}\boldsymbol{x}_s+2\alpha_t\int^{\lambda_t}_{\lambda_s}e^{-2(\lambda_t-\lambda)}\boldsymbol{x}_\theta(\boldsymbol{x}_\lambda,\lambda)\mathrm{d}\lambda+\sqrt{2}\sigma_t\int^{\lambda_t}_{\lambda_s}e^{-(\lambda_t-\lambda)}\mathrm{d}\boldsymbol{\omega}_\lambda$$
보통의 training-free sampling 방법들은 모두 diffusion SDE를 discretization 혹은 diffusion ODE를 discretization함으로써 구현하는데, 이 중 fast sampling 방법인 diffusion ODE solver들과 비교하면 아래와 같습니다.
** 이후에 등장한 DPM Solver v3가 궁금하시면 아래 더보기를 참조하세요
----------------------------------------------------------------
<DPM Solver-v3>
DPM-Solver-v3: Improved Diffusion ODE Solver with Empirical Model Statistics (NIPS'23)
이후에 DPM-Solver-v3가 등장하는데, 아래 그림과 같이 새로운 Parametrization을 통해 기존의 방법들을 Generalization해 최적화를 진행합니다.
간단히 수식을 소개하면, 기존의 semi-linear ODE를 log-SNR을 활용해 표현하면 아래와 같았습니다.
$$\begin{aligned}
\frac{\mathrm{d}\boldsymbol{x}_\lambda}{\mathrm{d}\lambda}&=\frac{\dot{\alpha}_\lambda}{\alpha_\lambda}\boldsymbol{x}_\lambda-\sigma_\lambda\boldsymbol{\epsilon}_\theta(\boldsymbol{x}_\lambda,\lambda)\\
&\dot{\alpha}_\lambda:=\frac{\mathrm{d}\alpha_\lambda}{\mathrm{d}\lambda}\ and\ \lambda_t:=log(\alpha_t/\sigma_t)
\end{aligned}$$
본 논문은, 이전과 다르게 linear part를 Rosenbrock-type exponential integrator를 활용해 표현하면 아래와 같이 변형됩니다.
** Exponential rosenbrock-type methods (SINUM'09)
$$\frac{\mathrm{d}\boldsymbol{x}_\lambda}{\mathrm{d}\lambda}=\underbrace{\left(\frac{\dot{\alpha}_\lambda}{\alpha_\lambda}{\color{red}-\boldsymbol{l}_\lambda}\right)\boldsymbol{x}_\lambda}_{\text{linear part}}-\underbrace{\left(\sigma_\lambda\boldsymbol{\epsilon}_\theta(\boldsymbol{x}_\lambda,\lambda){\color{red}-\boldsymbol{l}_\lambda\boldsymbol{x}_\lambda}\right)}_{\text{non-linear part}}$$
위에서 $\boldsymbol{l}_\lambda$는 $\lambda$-dependent한 D-차원의 coefficient입니다. 아래와 같이 $\boldsymbol{x}$에 대한 non-linear part의 Frobenius norm을 제약하는데 사용합니다.
** $p^\theta_\lambda$는 $\lambda$에 대한 ODE trajectory GT sample의 분포를 의미합니다.
$$\boldsymbol{l}^*_\lambda=\underset{\boldsymbol{l}_\lambda}{argmin}\mathbb{E}_{p^\theta_\lambda(\boldsymbol{x}_\lambda)}\left\| \nabla_{\boldsymbol{x}}\mathcal{N}_\theta(\boldsymbol{x}_\lambda,\lambda) \right\|^2_F$$
결과적으로 diffusion ODE는 아래와 같이 정리됩니다.
$$\boldsymbol{x}_{\lambda_t}=\alpha_{\lambda_t}\underbrace{e^{-\int^{\lambda_t}_{\lambda_s}\boldsymbol{l}_\lambda\mathrm{d}\lambda}}_{\text{linear coefficient}}\left(\frac{\boldsymbol{x}_{\lambda_s}}{\alpha_{\lambda_s}}-\int^{\lambda_t}_{\lambda_s}\underbrace{e^{\int^\lambda_{\lambda_s}(\boldsymbol{l}_\tau+\boldsymbol{s}_\tau)\mathrm{d}\tau}}_{\text{scaling coefficient}}\left(\boldsymbol{g}_\theta(\boldsymbol{x}_\lambda,\lambda)+\underbrace{\int^{\lambda}_{\lambda_s}e^{-\int^r_{\lambda_s}\boldsymbol{s}_\tau\mathrm{d}\tau}\boldsymbol{b}_r\mathrm{d}r}_{\text{bias coefficient}}\right)\mathrm{d}\lambda\right)$$
역시나 $\boldsymbol{g}_\theta$의 $\lambda$에 대한 n-th order Taylor Expansion을 활용해 discretization을 진행하면 아래와 같이 high-order solver에 대한 식으로 나타낼 수 있습니다.
$$\boldsymbol{x}_{t}=\alpha_{t}\underbrace{\boldsymbol{A}(\lambda_s,\lambda_t)}_{\text{linear coefficient}}\left(\frac{\boldsymbol{x}_{s}}{\alpha_{s}}-\int^{\lambda_t}_{\lambda_s}\underbrace{\boldsymbol{E}_{\lambda_s}(\lambda)}_{\text{scaling coefficient}}\underbrace{\boldsymbol{B}_{\lambda_s}(\lambda)}_{\text{bias coefficient}}\mathrm{d}\lambda-\sum^n_{k=0}\boldsymbol{g}^{(k)}_s\int^{\lambda_t}_{\lambda_s}\underbrace{\boldsymbol{E}_{\lambda_s}(\lambda)}_{\text{scaling coefficient}}\frac{(\lambda-\lambda_s)^k}{k!}\mathrm{d}\lambda\right)+O(h^{n+2})$$
위 식을 활용해 time s에서 time t로 이동해 $\hat{\boldsymbol{x}}_t$를 얻기 위해서는
- n개의 $\boldsymbol{x}_{\lambda_{i_k}}$와
- t에 해당하는 $\boldsymbol{g}_{i_k}$가
필요하고, 진행되는 과정을 아래 그림과 같이 표현할 수 있습니다.
----------------------------------------------------------------
이와 같은 DPM Solver++는 아래와 같이 Diffusers의 Scheduler를 활용해 구현할 수 있습니다.
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
model_id = "stabilityai/stable-diffusion-2-1"
# Use the DPMSolverMultistepScheduler (DPM-Solver++) scheduler here
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")
prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]d. freeU
FreeU: Free Lunch in Diffusion U-Net(CVPR'24)
FreeU에서는 main feature map인 $F_b$가 denoising에 ciritical한 반면, skip-connection의 feature $F_s$는 high-frequency 정보를 주로 다룬다는 것을 확인했습니다.
하지만 이런 차이가 있는 정보를 단순히 concat한다면, main feature의 무시할 수 없는 semantic content를 해칠 수 있어 channel-independent하며, emprically-tuned된 파라미터를 활용합니다.
이를 통해 main feature를 증진시키면서도 skip-feature의 low-frequency요소를 감소시키는 방법으로 FreeU를 제안했습니다.
Diffusers에서 freeU를 활성화하는 방법은 아래와 같습니다.
import torch
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, safety_checker=None
).to("cuda")
pipeline.enable_freeu(s1=0.9, s2=0.2, b1=1.5, b2=1.6)
generator = torch.Generator(device="cpu").manual_seed(33)
prompt = ""
image = pipeline(prompt, generator=generator).images[0]DPM
https://ml.cs.tsinghua.edu.cn/dpmv3/static/DPMv3.pdf
'Developers 공간 [Basic] > Vision & Audio' 카테고리의 다른 글
| [NLP] GPT 기초 정리 (0) | 2024.06.06 |
|---|---|
| [Calibration] Camera와 좌표계 변환 기초 (4) | 2022.12.21 |
| [Dataset] Waymo Open Dataset 개요 (0) | 2022.12.21 |