https://arxiv.org/abs/2112.10752
High-Resolution Image Synthesis with Latent Diffusion Models
By decomposing the image formation process into a sequential application of denoising autoencoders, diffusion models (DMs) achieve state-of-the-art synthesis results on image data and beyond. Additionally, their formulation allows for a guiding mechanism t
arxiv.org
해당 논문을 보고 작성했습니다.
Abstract
image를 생성하는 process를 denoising autoencoder의 연속된 application으로 분해함으로써, diffusion models (DMs)이 image data와 그 외의 data에서 가장 좋은 합성 결과를 보여주었습니다. 그리고 retraining 과정 없이 image를 생성하는 과정을 control 하기 위해 guiding mechanism을 사용하였습니다. 하지만 이러한 model들은 일반적으로 pixel space에서 바로 동작하기 때문에 DMs을 최적화하는 것은 수백 개의 GPU가 필요로 하며, inference도 상당히 많은 연산량을 필요로 합니다. DM이 제한된 연산 resource를 가지고 학습하면서도 quality와 flexibility를 유지하기 위해, 저자들은 powerful pretrained autoencoder의 latent space를 사용합니다. 이전 연구들과 다르게, 이러한 representation에서 diffusion model을 학습하면 complexity를 줄이고 detail을 보존할 수 있으며 visual fidelity를 크게 높일 수 있습니다. model architecture로 cross-attention layer를 사용함으로써 저자들은 text나 bounding box와 같은 general conditioning input을 위한 diffusion model을 generator로 사용할 수 있으며, convolutional manner에서 고해상도 합성이 가능해집니다. 저자들의 latent diffusion models (LDMs)은 pixel-based DMs보다 연산량을 상당히 줄였음에도 불구하고, image inpainting과 class-conditional image synthesis에서 가장 좋은 성능을 기록했으며, text-to-image synthesis, unconditional image generation, super-resolution과 같은 다양한 task에서도 비교 가능할 정도의 좋은 성능을 기록했습니다.
Introduction
image 합성은 최근에 가장 급격하게 발전된 computer vision 분야 중 하나이지만 상당히 많은 연산량을 필요로 합니다. 특히 고해상도의 복잡하고 자연스러운 scene을 합성하는 것은 수십억개의 parameter를 가지고 있는 autoregressive (AR) transformer를 확장하여 가능해졌습니다. 반면 GAN의 학습이 어렵고 multi-modal distribution을 modelling 하기 어려워 제한됩니다. 최근에 denoising autoencoder의 계층적 구조를 설계하는 diffusion model은 image 합성과 다양한 task에서 좋은 성능을 보여주고 있습니다. likelihood-based model로써, mode-collapse와 GAN의 학습 불안정성을 경험하지 않으며 parameter를 많이 공유하여 AR model과 같이 수십억 개의 parameter를 포함하지 않고도 자연스러운 image의 매우 복잡한 distribution을 modelling 할 수 있습니다.
Democratizing High-Resolution image Synthesis
DMs은 likelihood-based model의 한 종류로 data에서 인지하기 어려운 detail들을 modelling하는데 과도하게 많은 양의 자원을 사용하기 쉽습니다. reweighted variational objective가 초기 denoising 단계를 undersampling 하여 문제를 해결하려 했지만 DMs은 학습이나 evaluation 과정에서 반복적인 function을 거치기 때문에 여전히 연산량을 많이 사용합니다. 이 model의 접근성을 높이면서 연산량을 줄이기 위해, train과 sample 과정 모두에서 연산 복잡성을 줄여야 합니다. 성능을 저하하지 않고 DMs의 연산 요구량을 줄이는 것이 접근성을 향상시키는 key입니다.
Departure to Latent Space
저자들은 pixel space에서 이미 학습된 diffusion model을 분석하는 것으로 연구를 시작합니다.
위 그림은 학습된 model의 rate-distortion trade-off를 보여줍니다. likelihood-based model의 경우 크게 2 stage로 학습을 나눌 수 있습니다: 먼저 high-frequency detail을 제거하지만 의미적 변화는 학습하는 perceptual compression stage입니다. 그 다음 stage는 generative model가 data의 의미와 개념적 구성을 학습합니다(semantic compression). 그래서 저자들은 지각적으로 동등하지만 연산적으로 적절한 space를 찾는 것을 첫 번째 목표로 하며, 찾은 space를 사용하여 diffusion model가 high-resolution image synthesis를 학습합니다.
일반적인 관행을 따라, 저자들은 학습 과정을 2가지 phase로 나눕니다: 먼저 저자들은 data space와 지각적으로 동등한 lower-dimensional (efficient) representational space를 제공하는 autoencoder를 학습합니다. 이전 연구들과 다르게, 저자들은 공간적 압축에 크게 의존할 필요가 없으며, 저자들은 DMs을 learned latent space에서 학습할 것입니다. 줄어든 복잡성 또한 single network pass로 latent space에서 image를 효율적으로 생성할 수 있도록 만들어줍니다. 이러한 model을 저자들은 Latent Diffusion Models (LDMs)라 부릅니다.
이러한 방식은 universal autoencoding stage를 한번만 학습하면 되고, DM training에 재사용할 수 있다는 장점이 있습니다. 그리고 완전히 다른 task에 대해서도 재사용할 수 있다는 장점이 있습니다. 요약하여 저자들의 연구는 다음과 같은 contribution이 있습니다.
- pure transformer-based 방법과 다르게 저자들의 method는 더 높은 차원의 data로 scaling 될 수 있으며, 이를 통해 이전 연구들보다 더 신뢰할 수 있고 detailed reconstruction을 compression level에서 수행할 수 있고 megapixel image와 같은 고해상도 합성에 효율적으로 사용될 수 있습니다.
- 여러 task에서 저자들은 좋은 성능을 달성했으며, 상당히 적은 연산량을 사용합니다. pixel-based diffusion approach와 비교했을 때 inference cost도 상당히 감소되었습니다.
- encoder/decoder architecture와 score-based prior를 동시에 학습하는 이전 연구와 다르게 저자들은 reconstruction and generative ability에 대한 섬세한 weight를 요구하지 않으며, latent space의 정규화를 거의 요구하지 않습니다.
- 저자들은 super-resolution, inpainting, semantic synthesis와 같은 densely condition이 필요한 task에서 저자들의 model이 convolutional fashion에서 사용될 수 있고 결과로 1024x1024 pixel 크기의 크고 일관성 있는 image를 redering 할 수 있습니다.
- 그리고 저자들은 cross-attention 기반 general-purpose conditioning mechanism을 design 했으며, multi-modal training이 가능합니다.
Method
diffusion model을 고해상도 image 합성을 수행하도록 학습할 때 필요한 연산량을 줄이기 위해, 저자들은 diffusion model이 지각적으로 드러나지 않는 detail들을 무시하도록 loss term을 undersampling 하였지만, 여전히 pixel space의 연산량을 필요로 하고 computation time과 energy resource가 많이 사용됩니다.
저자들은 이러한 문제를 피하기 위해 generative learning phase와 압축 과정을 명시적으로 분리하였습니다. 이를 위해 저자들은 image space와 지각적으로 동등하지만 상당히 줄어든 연산 복잡도를 띄는 space를 학습하는 autoencoding model을 사용했습니다.
이러한 방법은 몇 가지 장점이 있습니다. 1) 고차원 image space를 사용하지 않기 때문에 DMs이 저차원 space에서 sampling을 수행하기 때문에 훨씬 더 적은 연산량을 필요로 합니다. 2) U-Net architecture에서 유도된 DMs의 inducitve bias를 사용합니다. 공간적 구조를 띄는 data에 특히 효율적이고 aggressive quality-reducing compression level의 필요성을 완화합니다. 3) 마지막으로 다양한 generative model을 학습하는데 사용될 수 있는 latent space를 가지는 general-purpose compression model을 얻고, 이를 single-image CLIP-guided synthesis와 같은 downstream application에 사용될 수 있습니다.
Perceptual Image Compression
저자들의 perceptual compression model은 perceptual loss와 patch-based adversarial objective의 combination으로 학습된 autoencoder로 구성됩니다. 이는 local realism을 enforce하여 image manifold에서 reconstruction이 이루어지도록 만들고, pixel-space loss에만 의존하였을 때 생기는 bluriness를 방지했습니다.
구체적으로 말하면, RGB space의 image x ∈ R^{H x W x 3}가 주어졌을 때, encoder ε는 x를 latent representation z로 encode하고, decoder D는 latent로부터 image를 reconstruct 합니다. encoder는 image를 downsample 하며 우리는 downsample 된 결과에 대해 조사를 진행하면 됩니다.
임의의 high-variance latent space를 피하기 위해, 저자들은 2가지 regularization을 실험했습니다. 첫번째는 KL-reg로, VAE처럼 learned latent 가 standard normal이도록 KL-penalty를 부여합니다. VQ-reg는 decoder에 vector quantization layer를 사용합니다. 저자들의 DM은 학습된 latent space z = ε(x)의 2차원 구조에서 동작하도록 design 되었기 때문에, 저자들은 상대적으로 mild compression rate를 사용할 수 있고 매우 좋은 reconstruction 성능을 보여줍니다. 임의의 1D learned space z를 model 하는 이전 연구들은 z의 내재된 구조를 무시하게 됩니다. 그래서 저자들의 compression model은 x의 detail을 더 잘 보존할 수 있습니다.
Latent Diffusion Models
- Diffusion Model
diffusion model은 정규 분포된 variable을 점진적으로 denoising하여 data distribution p(x)를 학습하도록 design 된 확률 model이며, 이는 T 길이의 고정 Markov Chain의 reverse process를 학습하는 것에 해당됩니다. image 합성의 경우, 대부분의 성공적인 model들은 denoising score-matching을 반영하는 p(x)의 reweighted variational lower bound에 의존합니다. 이러한 model들은 denoising autoencoder의 균일한 weighted sequence로 해석될 수 있으며, 이들은 input x_t의 denoised variant를 예측하도록 학습됩니다.
식으로 나타내면 위와 같습니다.
- Generative Modeling of Latent Representation
저자들의 trained perceptual compression model은 encoder ε와 decoder D로 구성됩니다. 이를 통해 이제 high-frequency imperceptible detail들은 요약되는 효율적이고 저차원의 latent space를 사용할 수 있습니다. 고차원 pixel space와 비교했을 때, data의 중요한 부분과 sematic bit에 focus 하고 저차원에서 학습되기 때문에 연산량이 훨씬 효율적인 공간이므로 이 space는 likelihood-based generative model에 더 적합합니다.
autoregressive, attention-based transformer models, discrete latent space에 의존하는 이전 연구들과 다르게, 저자들은 model이 제공하는 image-specific inductive biases의 장점을 사용할 수 있습니다. 2D convolution layer로 구성된 U-Net을 사용할 수 있으며, reweighted bound를 사용하여 지각적으로 가장 관련 있는 요소에 집중할 수 있습니다.
neural backbone ε_θ(o, t)는 time-conditional UNet입니다. forward process가 고정되어 있기 때문에, z_t는 학습과정에서 encoder를 통해 효율적으로 얻을 수 있으며 p(z)에서 sample을 decoder라는 single pass를 통해 image space로 decode할 수 있습니다.
Conditioning Mechanisms
다른 종류의 생성 model과 유사하게 diffusion model은 p(z|y) 형태의 conditional distribution modeling을 수행할 수 있습니다. 이는 conditional denoising autoencoder ε_θ(z_t, t, y)로 구현될 수 있으며, text나 semantic map 또는 다른 image-to-image translation task와 같은 input y를 통해 합성 prcoess를 control 합니다.
하지만 image 합성 관점에서 DM의 생성 성능을 class-label 또는 input image에 대한 blurred variant 이상을 condition으로 하여 결합하는 것은 아직 충분히 연구되지 않은 분야입니다.
저자들은 DM의 U-Net backbone에 cross-attention mechanism을 확장하여 DM을 보다 유연한 conditional image generator로 만들었습니다. 이는 다양한 input modality의 attention based model을 학습하는 데 효과적입니다. 다양한 modality로부터 얻은 y를 전처리하기 위해, 저자들은 y를 중간 representation으로 project 하는 domain specific encoder τ_θ(y)를 사용했습니다. 이 중간 representation은 U-Net의 중간 layer에 있는 cross-attention layer Attention(Q, K, V) = softmax (QK^T/d^(1/2))V를 통해 mapping 되었습니다.
위 식에서 φ_i(z_t) ∈ R^{N x d_ε^i} U-Net 의 중간 flattened representation을 나타내고, W_V와 W_K는 학습 가능한 projection matrix입니다. 그림으로 나타내면 다음과 같습니다.
image-conditioning pair를 기반으로, 저자들은 conditional LDM을 다음 식으로 학습했습니다.
위 식을 통해 τ_θ와 ε_θ가 동시에 최적화됩니다. τ_θ가 domain-specific expert로 parameterize될 수 있기 때문에 conditioning mechanism이 더 유연합니다.
Experiments
LDMs은 다양한 image modality에 대한 유연하고 tractable한 diffusion based image synthesis 방법을 제안합니다. 저자들의 model을 다른 pixel-baseed diffusion model과 비교했을 때 학습 과정과 추론 과정에서의 이점을 분석합니다. 흥미롭게도 저자들은 VQ-regularized latent space에서 학습된 LDM이 더 나은 sample quality를 종종 보인다는 것을 발견했습니다.
On Perceptual Compression Tradeoffs
저자들은 다양한 downsampling factor f ∈ {1, 2, 4, 8, 16, 32}에 대한 LDM의 성능을 분석했습니다. LDM-1은 pixel-based DM과 동일합니다. 실험 model들의 hyperparameter와 reconstruction performance는 다음과 같습니다.
small downsampling factor LDM-{1, 2}는 느린 학습 속도를 보여주지만, factor f가 크다면 몇 step 이후 침체된 fidelity를 보여줍니다. LDM-{4, 8, 16}이 perceptually faithful result와 효율성 사이 balance가 잘 맞으며 pixel-based diffusion보다 FID 성능이 좋은 모습을 보여줍니다.
Conditional Latent Diffusion
- Transformer Encoders for LDMs
cross-attention based conditioning을 LDM에 적용하여 저자들은 다양한 conditioning modality를 사용할 수 있는 diffusion model을 구현했습니다. text-to-image의 image modeling의 경우, BERT-tokenizer와 transformer를 τ로 사용했으며, multi-head cross attention을 통해 U-Net으로 mapping 되는 latent code를 inference 했습니다. semantic layout을 condition으로 사용한 결과는 다음과 같습니다.
Limitations
LDMs는 pixel-based 방식과 비교했을 때 연산량을 상당히 줄였지만, GAN에 비해 연속적인 sampling process는 여전히 느립니다. 그리고 f = 4인 autoencoding model의 경우 image quality의 loss는 매우 작으며 reconstruction 성능은 fine-grained accuracy를 요구하는 task에서는 bottleneck이 됩니다.
Conclusion
저자들은 간단하고 효율적인 방법으로 denoising diffusion model의 학습과 sampling 모두 향상시키는 latent diffusion model을 제안했습니다. 저자들은 실험을 통해 task-specific architecture를 사용하지 않고 광범위한 conditional image synthesis task에서 좋은 성능을 보였습니다.
