논문리뷰

robotics : 논문리뷰 : Enhancing Continuous Diffusion Models with Discrete Latents (DisCo-Diff)

AI바라기 2025. 6. 20. 12:15

쉬운 설명

이 논문의 핵심 아이디어는 복잡한 그림을 그릴 때, 먼저 "무엇을 그릴지" 큰 주제를 정하고 세부 묘사에 들어가는 것과 같습니다.

기존 Diffusion Model은 마치 백지(noise)에서 곧바로 아주 복잡한 유화(데이터)를 한 번에 그리려는 화가와 같습니다. 과정이 매우 복잡하고 어렵습니다.

DisCo-Diff는 화가에게 먼저 "이번엔 허스키를 그릴 거야" 혹은 "이번엔 불독을 그릴 거야" 와 같이 **큰 주제(discrete latent)**를 정해주고 시작하게 합니다. 어떤 종류의 개를 그릴지 주제가 정해지면, 그 개의 구체적인 털 질감, 자세, 표정 등 세부 묘사(continuous generation)에만 집중하면 되므로 그림을 그리는 과정(Generative ODE)이 훨씬 단순하고 명확해집니다.

 

 

DisCo-Diff 학습 노트

용어 설명 (Terminology)

  • DisCo-Diff (Discrete-Continuous Latent Variable Diffusion Models): 이 논문에서 제안하는 모델의 이름. 이름 그대로 Discrete latent와 Continuous latent (Diffusion Model의 가우시안 노이즈)를 함께 사용하는 모델.
  • Generative ODE (Ordinary Differential Equation): Continuous-time diffusion model에서, 가우시안 노이즈(latent)로부터 실제 데이터로 변환되는 결정론적(deterministic) 생성 경로.
  • ODE Curvature: Generative ODE 경로가 얼마나 복잡하게 휘어져 있는지를 나타내는 척도. 이 논문에서는 기존 DM의 ODE curvature가 높아 학습이 어렵다고 주장하며, DisCo-Diff가 이 curvature를 낮추는 것을 핵심 장점으로 내세움.
  • Gumbel-Softmax: Discrete 변수를 sampling하는 과정에서 gradient가 흐를 수 있도록 해주는 연속적인 근사(continuous relaxation) 기법. Encoder와 Denoiser를 end-to-end로 학습시키기 위해 필수적임.
  • Group Hierarchical DisCo-Diff: 제안된 모델의 변형 구조. U-Net의 각기 다른 resolution level에 서로 다른 discrete latent 그룹을 주입하여, latents가 계층적인(hierarchical) feature(예: 저해상도에서는 형태, 고해상도에서는 질감)를 학습하도록 유도하는 방식.
  • Universally Applicable: 특정 domain(예: 이미지)에만 국한된 pre-trained encoder에 의존하지 않아, 분자 구조 등 다양한 데이터에 보편적으로 적용 가능하다는 의미.

Purpose of the Paper

이 논문은 기존 Diffusion Model(DM)이 가진 근본적인 학습의 어려움을 해결하고자 합니다.

  • 기존 DM의 한계: 복잡하고 multimodal한 실제 데이터 분포 전체를, 단 하나의 단순한 unimodal 가우시안 분포로 매핑하는 것은 매우 어려운(unnecessarily challenging) 학습 문제입니다.
  • 문제점: 이로 인해 noise로부터 데이터를 생성하는 경로, 즉 Generative ODE의 curvature가 매우 높아져 학습이 비효율적이고 불안정해집니다.
  • 제안하는 접근 방식: 이 문제를 완화하기 위해, DM의 continuous latent(가우시안 노이즈)를 보완하는 소수의 discrete latent variables를 추가합니다.
  • 차별점:
    1. End-to-end 학습: CLIP이나 MoCo 같은 pre-trained encoder에 의존하지 않습니다. 대신, denoiser의 손실(loss)을 줄이는 데 가장 도움이 되는 정보를 담도록 encoder와 discrete latents를 DM과 함께 end-to-end로 학습시킵니다. 이 덕분에 이미지뿐만 아니라 분자(molecule) 구조와 같이 pre-trained encoder가 없는 domain에도 적용 가능한 universality를 가집니다.
    2. 역할 분담: 데이터의 전체 정보를 discrete token으로 압축하는 VQ-VAE와 달리, DisCo-Diff는 소수의 global discrete latents로 데이터 분포의 **거시적인 mode (예: 스타일, 색상 패턴, 큰 구조)**를 포착하고, DM의 continuous latents는 각 mode 내의 미세하고 연속적인 변화를 모델링하는 역할 분담 구조를 가집니다.

Key Contributions

  • Discrete-Continuous Latent Framework 제안: Diffusion Model을 보강하기 위해 소수의 learnable discrete latents를 도입하는 DisCo-Diff라는 새로운 범용 프레임워크를 제안했습니다.
  • 학습 과정 단순화 및 ODE Curvature 감소 검증: Discrete latents가 데이터의 multimodal 구조를 분해하여 DM의 denoising task를 단순화하고, 결과적으로 Generative ODE의 curvature를 크게 낮추어 학습을 용이하게 함을 실험적으로 증명했습니다. (특히 noise level이 높은, 불확실성이 큰 구간에서 training loss가 크게 감소)
  • State-of-the-Art 성능 달성: 제안된 방법론을 통해 class-conditioned ImageNet-64/128 데이터셋에서 ODE sampler 기준 SOTA FID 점수를 달성했습니다. 이는 더 곧은(straighter) ODE 경로가 실제 생성 품질 향상으로 이어진다는 것을 보여줍니다.
  • 프레임워크의 보편성(Universality) 입증: 이미지 생성뿐만 아니라, 2D toy data, 그리고 molecular docking task에서도 성능 향상을 보여주어 DisCo-Diff가 특정 domain에 국한되지 않는 범용적인 프레임워크임을 입증했습니다.

Novelty

  • 기존 연구들이 pre-trained encoder의 고정된 representation을 활용한 것과 달리, denoiser 성능 향상이라는 목표 함수에 맞춰 discrete latents를 직접 end-to-end로 학습했다는 점이 가장 큰 독창성입니다. 이는 latents가 denoising에 가장 유용한 high-level 정보를 스스로 학습하도록 만듭니다.
  • Auxiliary latents로 continuous variable 대신 discrete variable을 사용해야 하는 이유를 명확히 제시했습니다. Ablation study를 통해 continuous latents는 KLD 가중치 조절이 어려워 encoder collapse나 overfitting을 유발하기 쉬운 반면, 제안된 discrete latents는 안정적으로 유의미한 mode를 학습함을 보였습니다.

Experimental Highlights

  • ImageNet SOTA (ODE Sampler):
    • ImageNet-64: FID 1.65 달성 (baseline EDM: 2.36)
    • ImageNet-128: FID 1.98 달성 (baseline VDM++: 2.29)
  • Toy 2D Gaussian Mixture 실험:
    • 시각적으로 기존 DM의 구부러진 ODE trajectory가 DisCo-Diff에서는 거의 직선에 가깝게 펴지는 것을 확인 (Fig. 3).
    • 정량적으로 ODE curvature와 denoiser network의 Jacobian norm이 전 구간에 걸쳐 크게 감소함을 보여줌 (Fig. 4). 이는 모델의 핵심 가설을 직접적으로 증명하는 결과입니다.
  • Molecular Docking 성능 향상:
    • Full test set 성공률: 32.9% → 35.4%
    • 더 어려운 'unseen proteins' subset 성공률: 13.9% → 18.5%. 이는 프레임워크의 universality를 강력하게 뒷받침합니다.
  • Ablation Study:
    • Classifier-free guidance (CFG)를 discrete latents에 적용했을 때 FID가 1.81에서 1.65로 향상되어, latents가 데이터 분포의 의미 있는 mode를 학습했음을 보여주었습니다.

Limitations and Future Work

  • Limitation (Stochastic Sampler 문제): ImageNet-128 실험에서, DDPM 같은 stochastic sampler를 사용했을 때 오히려 성능이 저하되었습니다. 논문은 discrete latents가 global, low-frequency 정보에 과도하게 집중한 나머지, sampling 후반부에 필요한 high-frequency 디테일을 모델이 놓치게 만들었을 것이라 추측합니다.
  • Future Work (Stochastic Sampler 문제 해결): Text-to-image 모델처럼 null token을 함께 사용하여, sampling 후반부(작은 timestep)에는 모델이 discrete latents의 영향을 무시하도록 학습시키는 방안을 제안합니다.
  • Future Work (적용 분야 확장):
    • Text-to-image 생성에 적용하여 text 조건부 정보와 상호보완적인 역할을 하는지 검증.
    • Diffusion model 외에 Flow-matching이나 Rectified Flow 같은 다른 continuous flow 모델에도 DisCo-Diff 아이디어를 적용.
  • Future Work (학습 프로세스 통합): 현재는 (1) Denoiser/Encoder 학습, (2) Autoregressive model 학습의 2단계로 나뉘어 있는데, 이를 하나의 seamless end-to-end 학습 파이프라인으로 통합하는 연구를 제안합니다.

Overall Summary

DisCo-Diff는 기존 Diffusion Model에 소수의 learnable discrete latent variables를 추가하여 학습의 근본적인 복잡성을 해결하는 새로운 프레임워크입니다. 이 discrete latents는 데이터의 거시적인 mode를 포착하여 DM이 세부적인 변화에만 집중하도록 만들고, 이를 통해 Generative ODE의 curvature를 낮추어 학습을 단순화합니다. 이 접근법은 ImageNet에서 SOTA 성능을 달성하고 molecular docking과 같은 다른 domain에서도 효과를 입증하여, pre-trained 모델 없이도 보편적으로 적용 가능한 강력한 생성 모델링 방법론의 가능성을 제시합니다.

 

Abstract
Diffusion models (DMs)은 generative learning에 혁명을 일으켰습니다. DMs는 diffusion process를 활용하여 data를 단순한 Gaussian distribution으로 encode 합니다. 그러나 복잡하고 잠재적으로 multimodal인 data distribution을 단일 continuous Gaussian distribution으로 encoding하는 것은 불필요하게 어려운 learning 문제가 될 수 있습니다.

우리는 보완적인 discrete latent variables를 도입하여 이 task를 단순화하는 Discrete-Continuous Latent Variable Diffusion Models (DisCo-Diff)를 제안합니다. 우리는 DMs를 learnable discrete latents로 augment하고, 이는 encoder로 inferred되며, DM과 encoder를 end-to-end로 train합니다. DisCo-Diff는 pre-trained networks에 의존하지 않으므로 framework가 보편적으로 적용 가능합니다.

discrete latents는 DM의 generative ODE의 curvature를 줄임으로써 DM의 복잡한 noise-to-data mapping 학습을 크게 단순화합니다. 추가적인 autoregressive transformer는 discrete latents의 distribution을 models합니다. DisCo-Diff는 작은 codebooks를 가진 몇 개의 discrete variables만 필요로 하므로 이는 간단한 단계입니다. 우리는 toy data, 여러 image synthesis tasks 및 molecular docking에 대해 DisCo-Diff를 validate했으며, discrete latents를 도입하는 것이 일관되게 model performance를 향상시킨다는 것을 발견했습니다. 예를 들어, DisCo-Diff는 ODE sampler를 사용하여 class-conditioned ImageNet-64/128 datasets에서 state-of-the-art FID score를 달성합니다.

 

 

더보기

Figure 1. Discrete-Continuous Latent Variable Diffusion Models (DisCo-Diff)

이 그림은 DisCo-Diff model이 어떻게 작동하는지 설명합니다. DisCo-Diff는 DMs를 추가적인 discrete latent variables로 augment하여 global appearance patterns를 capture합니다. 여기서는 허스키 images를 예로 보여줍니다.

(a) Training 및 Sampling 과정

이 다이어그램은 model의 Training과 Sampling 두 가지 주요 과정을 보여줍니다.

  • Training 과정:
    1. 원본 허스키 image가 Encoder (예: Vision Transformer)에 입력됩니다.
    2. Encoder는 image로부터 Discrete Latent Variables를 inferred (추론)합니다. 이 과정에서 Backpropagation이 가능하도록 Gumbel-Softmax distribution을 통한 continuous relaxation 기법이 사용됩니다.
    3. 이 discrete latents는 Cross-Attention을 통해 Denoiser Neural Network에 전달됩니다.
    4. 동시에, noise가 추가된 image가 Denoiser Network에 입력됩니다.
    5. Denoiser Network는 discrete latents 정보를 활용하여 noise를 제거하고 원본에 가까운 image를 복원하도록 train됩니다.
  • Sampling 과정:
    1. 새로운 image를 generate하기 위해, Autoregressive Model (Transformer)이 학습된 distribution에 따라 새로운 Discrete Latent Variables를 생성합니다.
    2. 생성된 discrete latents는 Training 때와 동일한 Denoiser Network에 입력되어, random noise로부터 새로운 허스키 image를 최종적으로 generate합니다.

(b) 서로 다른 Discrete Latent Variables에 대한 Generative Process

이 부분은 generative denoising diffusion trajectories의 schematic visualization을 보여줍니다.

  • 그래프: 그래프의 각기 다른 색상 선(trajectories)은 서로 다른 discrete latent variables를 나타냅니다. 각 discrete latent는 noise 상태에서 시작하여 최종 image로 향하는 경로(trajectory)를 서로 다른 mode로 이끌어갑니다. 이는 discrete latents가 생성 과정의 방향을 제어함을 의미합니다.
  • 생성된 이미지: 그래프 오른쪽에는 생성된 허스키 image들이 있습니다. 각 그룹은 그래프의 특정 색상 trajectory에 해당하며, discrete latent variable에 따라 털 색깔이나 배경(눈, 풀밭 등)과 같은 global style이 일관되게 나타나는 것을 볼 수 있습니다. 이는 model이 discrete latents를 통해 image의 전반적인 특징을 효과적으로 제어하고 있음을 보여줍니다.

 

Introduction

더보기

Diffusion models (DMs)은 최근 다양한 domains에서 generative modeling의 획기적인 발전을 이끌었습니다. 예를 들어, 이 model들은 표현력이 풍부한 high-resolution imagery를 합성할 수 있거나 정확한 molecular structures를 generate할 수 있습니다. DMs는 training data를 간단한 unimodal Gaussian prior distribution에 효과적으로 encodes하는 forward diffusion process를 활용합니다. Generation은 stochastic process 또는 더 편리하게는 Gaussian prior에서 random noise를 input으로 받아 generative ordinary differential equation (ODE)를 통해 data로 transforms하는 deterministic process로 공식화될 수 있습니다. Gaussian prior는 DM의 continuous latent variables에 해당하며, 여기서 data는 ODE-defined mapping을 통해 uniquely encoded됩니다.

그러나 실제 data distributions는 일반적으로 high-dimensional하고 complex하며 종종 multimodal입니다. 이러한 data를 단일 unimodal Gaussian distribution으로 직접 encoding하고 해당 reverse noise-to-data mapping을 learning하는 것은 어려운 과제입니다. mapping, 즉 generative ODE는 필연적으로 강한 curvature를 가진 매우 complex해야 하며, 전체 data distribution을 단일 Gaussian distribution에 mapping하는 것은 부자연스럽다고 간주될 수 있습니다. 실제로는 class labels나 text prompts와 같은 conditioning information이 DM의 denoiser에 더 정확한 denoising을 위한 추가적인 cues를 제공함으로써 복잡한 mapping을 단순화하는 데 종종 도움이 됩니다. 그러나 이러한 conditioning information은 일반적으로 semantic한 성격을 가지며, class나 text prompt가 주어져도 mapping은 매우 complex하게 유지됩니다. 예를 들어, images의 경우, 하나의 class 내에서도 pixel space에서 큰 거리에 해당하는 매우 다른 styles과 color patterns을 가진 images를 찾을 수 있습니다.

여기서 우리는 Discrete-Continuous Latent Variable Diffusion Models (DisCo-Diff)를 제안합니다. 이는 data에 대한 추가적인 high-level information을 encode하고 main DM이 denoising task를 단순화하는 데 사용할 수 있는 추가적인 discrete latent variables로 augmented된 DMs입니다. 이러한 discrete latents는 encoder network를 통해 inferred되며 DM과 함께 end-to-end로 learnt됩니다. 이를 통해 discrete latents는 DM의 score matching objective를 줄이고, simple noise를 complex data로 mapping하는 DM의 어려운 task를 더 쉽게 만드는 데 유용한 information을 직접 encode하도록 learn합니다. 실제로, 우리는 이들이 DM의 generative ODE의 curvature를 상당히 줄이고, 특히 denoising이 가장 ambiguous하고 어려운 large diffusion times에서 DM training loss를 줄인다는 것을 발견했습니다. previous work와 달리, 우리는 domain-specific pre-trained encoder networks에 의존하지 않아 우리의 framework를 general하고 보편적으로 적용할 수 있도록 만듭니다. inference 중 discrete latent variables의 sampling을 용이하게 하기 위해, 우리는 두 번째 단계에서 discrete latents에 대한 autoregressive model을 learn합니다. 우리는 비교적 작은 codebooks를 가진 소수의 discrete latents만을 사용하므로, autoregressive model의 추가적인 training이 쉬워집니다. 우리는 특히 continuous latents 대신 auxiliary discrete의 사용을 지지합니다; Sec. 3.2를 참조하십시오.

previous works는 images를 model하기 위해 fully discrete latent variable-based approaches를 사용하지만, 이는 일반적으로 큰 codebooks를 가진 spatially arranged latents의 큰 집합을 필요로 하여 그 distribution을 learning하는 것을 어렵게 만듭니다. 반면에 DisCo-Diff는 discrete latents를 DM의 continuous latents (Gaussian prior)와 신중하게 결합하여 data 내의 discrete 및 continuous variations의 modeling을 효과적으로 분리합니다. 이는 단지 몇 개의 discrete latents만을 필요로 합니다.

그 보편성을 입증하기 위해, 우리는 여러 다른 tasks에서 DisCo-Diff framework를 validate합니다. 동기 부여 예시로서, 우리는 2D toy distributions을 연구하는데, 여기서 discrete latents는 sampling 중에 더 작은 curvature를 가진 different modes를 capture하도록 learn합니다. 그런 다음 우리는 image synthesis를 다루는데, 여기서 discrete latents는 종종 global style 및 color patterns과 관련된 large-scale appearance를 learn합니다. 이를 통해 semantic conditioning information에 상호 보완적인 이점을 제공합니다. 정량적으로, DisCo-Diff는 보편적으로 output quality를 boosts하고 여러 ImageNet generation benchmarks에서 state-of-the-art performance를 달성합니다. 또한, 우리는 실험적으로 auxiliary discrete latents가 우리의 setup에서 continuous latents보다 superior하다는 것을 확인하고, discrete latents를 DM network에 injecting하기 위한 다양한 network architectures를 연구합니다. 신중한 hierarchical design은 generative adversarial networks에 대한 문헌의 관찰과 유사하게, 다른 discrete latents가 shape 대 color와 같은 다른 image characteristics를 encode하도록 장려할 수 있습니다. 우리는 또한 drug discovery의 중요한 task인 molecular docking에 DisCo-Diff를 적용하는데, 여기서 discrete latents는 interaction에서 critical atoms를 indicate하도록 learning함으로써 performance를 다시 향상시키고, 이런 방식으로 각 pose의 continuous variability로부터 다른 가능한 poses에 의해 주어지는 multimodal uncertainty를 deconvolving합니다. 더욱이, 우리는 Poisson Flow Generative Models를 discrete latent variables로 augment하여 이 framework가 regular DMs 이외의 다른 "iterative" generative models에도 적용될 수 있음을 showcase하고 유사한 이점을 관찰합니다.

Contributions. (i) 우리는 DMs에서 discrete 및 continuous latent variables를 보편적인 방식으로 combining하기 위한 새로운 framework인 DisCo-Diff를 제안합니다. (ii) 우리는 DisCo-Diff를 광범위하게 validate하여 모든 experiments에서 model quality를 크게 boosting하고 여러 image synthesis tasks에서 state-of-the-art performance를 달성합니다. (iii) 우리는 discrete latent variables의 unique benefits와 그것들이 main denoiser network에 어떻게 공급될 수 있는지를 보여주는 상세한 analyses 및 ablation 및 architecture design studies를 제시합니다. (iv) 전반적으로, 우리는 performant generative models를 designing하기 위한 insights를 제공합니다. 우리는 real-world data가 discrete 및 continuous latents를 모두 leverage하는 generative frameworks로 가장 잘 modeled된다는 것을 보여줌으로써 discrete latents의 타당성을 주장합니다. 우리는 community에 널리 적용 가능한 modeling approach를 제공하기 위해 pre-trained encoders에 의존하지 않는 simple하고 universal한 framework를 의도적으로 개발했습니다.

 

DisCo-Diff Introduction 정리노트 (for AI Researchers)

1. 문제 제기 (Problem Definition)

  • 기존 Diffusion Models (DMs)의 근본적인 한계: 복잡하고 multimodal한 실제 data distribution 전체를 단 하나의 unimodal Gaussian distribution으로 mapping하는 것은 매우 어려운 learning 문제임.
  • 이로 인해 generative ODE의 curvature가 매우 높아져 noise-to-data mapping 학습이 비효율적임.
  • Class label이나 text prompt 같은 conditioning 정보가 주어져도, class 내에 존재하는 다양한 style, color pattern 등의 variation 때문에 mapping의 복잡성은 여전히 높게 유지됨.

2. 핵심 아이디어 (Core Idea: DisCo-Diff)

  • 기존 DM의 continuous latent space (Gaussian prior)에 추가적으로 소수의 discrete latent variables를 도입하여 문제를 해결.
  • 이 discrete latent는 data의 global style, color와 같은 high-level의 semantic 정보를 capture하는 역할을 함.
  • DM은 이 discrete latent 정보를 denoising 과정에 활용하여, noise로부터 data를 generation하는 task의 난이도를 낮춤.

3. 주요 메커니즘 (Key Mechanism)

  • End-to-End 학습: Encoder가 data로부터 discrete latents를 infer하고, 이 encoder와 주 DM을 end-to-end로 함께 train함. discrete latent는 DM의 score matching objective를 가장 효과적으로 줄일 수 있는 방향으로 학습됨.
  • Framework의 일반성 (Universality): 특정 domain에 종속적인 pre-trained encoder에 의존하지 않아, 별도의 사전 학습 없이 어떤 data에도 범용적으로 적용 가능함.
  • Sampling: Inference 시 새로운 discrete latents를 sampling하기 위해, 이 latents의 distribution을 학습하는 간단한 autoregressive model을 추가로 사용. latent의 개수와 codebook 크기가 작아 학습이 용이함.

4. 차별점 (Key Differentiators)

  • Hybrid Approach: Latent space 전체를 discrete로 대체하는 방식들과 달리, DM의 continuous latents와 소수의 discrete latents를 결합하여 modeling의 효율성을 극대화함. Discrete와 continuous variation의 분리를 유도.
  • Auxiliary Discrete Latents 선호: 보조적인 latent variable로서 continuous가 아닌 discrete를 사용하는 것이 더 우월함을 주장 (Sec 3.2에서 상세히 다룰 예정).
  • 확장성: 제안하는 framework는 Poisson Flow Generative Models와 같은 다른 iterative generative models에도 적용 가능함을 보임.

5. 핵심 성과 (Claimed Performance)

  • 2D toy distribution, image synthesis, molecular docking 등 다양한 task에서 일관된 성능 향상을 보임.
  • ImageNet generation benchmark에서 state-of-the-art (SOTA) 성능을 달성.
  • Ablation study와 architecture 분석을 통해 discrete latents의 명확한 이점과 효과적인 주입 방법을 제시.

쉬운 설명 :

Diffusion model을 아주 뛰어난 화가에 비유할 수 있습니다. 이 화가에게 "허스키를 그려줘"라고 요청하면, 그림 실력은 뛰어나지만 어떤 허스키를 그릴지는 전적으로 화가의 상상에 달려있습니다. 눈밭에 있는 검은 허스키를 그릴 수도, 숲속에 있는 갈색 허스키를 그릴 수도 있죠. 즉, 결과물이 너무 다양해서 제어하기가 어렵습니다. 이처럼 너무 많은 가능성을 한 번에 다루는 것이 기존 diffusion model의 어려운 점입니다.

DisCo-Diff는 이 문제를 해결하기 위해 화가에게 아주 간단한 "스타일 카드" 몇 장을 추가로 주는 것과 같습니다. 예를 들어, 1번 카드는 '눈 배경에 흑백 털', 2번 카드는 '숲 배경에 갈색 털' 같이 그림의 전반적인 분위기를 결정하는 high-level 정보가 담겨 있습니다.

화가(DM)는 "허스키를 그려줘"라는 기본 요청과 함께 "1번 스타일 카드로!"라는 구체적인 가이드를 받습니다. 덕분에 화가는 그림의 큰 틀(배경, 색감)을 고민할 필요 없이 세부적인 묘사에만 집중할 수 있게 되어 훨씬 쉽고 안정적으로 고품질의 그림을 완성할 수 있습니다.

여기서 "스타일 카드"가 바로 이 논문이 제안하는 discrete latent variable입니다. 어떤 스타일 카드를 사용할지는 encoder라는 보조 모델이 원본 이미지들을 보고 자동으로 학습하여 정해줍니다. 이 간단한 아이디어를 통해 모델의 학습 부담을 크게 줄이고, 결과물의 quality를 획기적으로 높인 것이 이 논문의 핵심입니다.

 

 

 

2. Background

더보기

DisCo-Diff는 continuous-time DMs를 기반으로 하며, 우리는 EDM framework를 따릅니다. DMs는 고정된 forward process에서 인 clean data를 -variance Gaussian noise를 사용하여 perturb합니다. 여기서 이고 는 diffusion process 상의 time을 나타냅니다. 결과로 나오는 distribution은 로 표기되며, 입니다. 충분히 큰 에 대해, 이 distribution은 pure random Gaussian noise와 거의 동일합니다. DMs는 이 관찰을 활용하여 를 sample한 다음, 개의 점진적으로 감소하는 noise levels  ()의 sequence를 통해 sample을 iteratively denoise합니다. 여기서 이고 입니다. 는 continuous  function의 discretization에 해당합니다. 만약 이라면, final 은 data distribution을 따릅니다.

Sampling은 deterministic 또는 stochastic differential equation을 simulating하는 것에 해당합니다: 

 여기서 는 standard Wiener process이고 는 diffused distribution 의 score function입니다. 위 식의 첫 번째 항은 Probability Flow ODE로, samples를 높은 noise levels에서 낮은 noise levels로 이동시킵니다. 두 번째 항은 Langevin Diffusion SDE로, 다른 noise levels 에 대한 equilibrium sampler이며, synthesis 중 오류를 수정하는 데 도움을 줄 수 있습니다. 이 구성 요소는 time-dependent parameter 에 의해 조절될 수 있습니다. 으로 설정하면 pure ODE-based synthesis가 됩니다. 일반적으로, generative ODE/SDE를 풀기 위해 다양한 sampling methods가 사용될 수 있습니다.

DM을 Training하는 것은 다루기 힘든(intractable) score function 을 approximate하기 위해 model을 learning하는 것에 해당합니다. EDM framework를 따라, 우리는 로 parametrize합니다. 여기서 는 noisy inputs으로부터 clean data를 predict하도록 trained된 learnable denoiser neural network이며, noise level 에 conditioned됩니다. 이는 denoising score matching을 사용하여 trained될 수 있으며, 다음을 minimizing합니다:  여기서 는 diffusion times 에 대한 distribution 로부터 나오며, 이고, 는 다른 noise levels에 대해 objective에 다른 weight를 부여하는 function입니다.

이 연구에서는 별도의 언급이 없는 한, 를 사용하고 EDM work의 configuration을 따릅니다. 또한, 우리는 discrete latent variables에 conditioning할 때 DisCo-Diff에서 classifier-free guidance를 활용합니다. Classifier-free guidance는 unconditional diffusion model과 conditional diffusion model의 score functions를 combines하여 conditioning을 amplify합니다.

 

 

DisCo-Diff Background 정리노트 (for AI Researchers)

1. 기반 Framework

  • 이 논문은 Song et al.의 continuous-time Diffusion Models (DMs) 방법론을 기반으로 하며, 구체적인 구현과 수식 전개는 Karras et al.의 EDM (Elucidating the Design Space of Diffusion Models) framework를 따름.

2. 핵심 프로세스 (EDM Framework 기준)

  • Forward Process: Clean data 를 noise variance 를 갖는 Gaussian noise로 점진적으로 perturb함.
  • Score Parametrization: 다루기 힘든 score function 을 직접 근사하는 대신, denoiser network 를 학습시켜 clean data 를 예측하게 하고, 이를 이용해 와 같이 score를 계산함.
  • Training Objective: Denoising score matching loss를 objective function으로 사용하여 를 학습.

3. Generation (Sampling) 방식

  • ODE/SDE 공식 사용: Generation은 Probability Flow ODE (deterministic)와 Langevin Diffusion SDE (stochastic)가 결합된 미분 방정식을 푸는 과정으로 정의됨.
    • Probability Flow ODE: sample을 noise에서 data 분포로 이끄는 주된 흐름을 형성.
    • Langevin Diffusion SDE: synthesis 과정 중 발생하는 오류를 보정하는 sampler 역할을 수행.
  • 구체적 설정: 별도 언급이 없으면 noise schedule은 EDM 논문과 동일하게 를 사용.

4. 본 논문과의 연결점 (Key for DisCo-Diff)

  • Classifier-Free Guidance 활용: 본 논문에서 제안하는 discrete latent variables에 대한 conditioning을 구현하기 위해 classifier-free guidance를 채택함. 이는 unconditional model과 conditional model의 score function을 조합하여 conditioning 효과를 증폭시키는 기법으로, DisCo-Diff의 핵심 구현 방식 중 하나임.

쉬운 설명 :

이 섹션은 DisCo-Diff라는 새로운 모델을 만들 때 어떤 기술들을 기본 도구로 사용했는지 설명하는 부분입니다.

핵심 비유: 흐릿한 사진을 선명하게 복원하기

  1. 학습 과정 (사진 일부러 흐리게 만들기): 컴퓨터에게 사진 복원 기술을 가르치기 위해, 먼저 수많은 선명한 원본 사진을 일부러 점차 흐리게 만듭니다. 사진이 아주 조금 흐릿한 상태부터 완전히 알아볼 수 없는 noise(지직거리는 화면) 상태가 될 때까지 모든 단계를 컴퓨터가 지켜보게 합니다. 이 과정을 통해 컴퓨터는 각 noise 단계에서 원본 사진이 어땠을지 학습합니다. 이것이 Diffusion Model의 기본 학습 원리입니다.
  2. 복원 과정 (새로운 사진 만들기): 이제 컴퓨터는 완전히 noise만 있는 상태에서 출발하여, 학습했던 과정을 거꾸로 거슬러 올라가며 noise를 단계적으로 제거합니다. 그러면 세상에 없던 새로운 사진이 마법처럼 나타납니다. 이 복원 과정이 바로 Generation 또는 Sampling입니다.
  3. 복원 전문가의 두뇌 (): 이 논문에서 사용하는 모델()은 복원 전문가의 두뇌와 같습니다. 전문가는 아무리 흐릿한 사진을 보더라도, 그 사진의 원본이 어땠을지를 예측할 수 있습니다. 모델은 noise 낀 이미지를 보고 원본 이미지를 예측하도록 학습하며, 이 예측 능력이 전체 복원 과정을 이끌어갑니다.
  4. 복원을 위한 힌트 (Classifier-Free Guidance): 만약 "강아지 사진을 복원해줘"라는 힌트를 준다면 어떨까요? 전문가는 그냥 사진을 복원하는 동시에, 결과물이 '강아지'처럼 보이도록 더 신경 써서 복원할 것입니다. 이처럼 힌트(conditioning)를 주었을 때 그 효과를 극대화하는 기술이 Classifier-Free Guidance입니다. 이 논문은 나중에 자신들이 제안하는 discrete latent라는 새로운 종류의 힌트를 모델에게 줄 때 바로 이 기술을 사용하겠다고 밝히고 있습니다.

 

3. DisCo-Diff

더보기

Sec. 3.1에서는 먼저 DisCo-Diff의 generative model과 training framework를 공식적으로 정의하고, Sec. 3.2에서 우리의 approach를 상세히 논의하고 신중하게 동기를 부여합니다. Sec. 3.3에서는 중요한 architecture 고려 사항을 강조합니다.

3.1. Generative Model and Training Objective

우리의 DisCo-Diff framework (Fig. 1)에서는 DM의 learning process를 -dimensional discrete latent 로 augment합니다. 여기서 각 dimension은 codebook size 의 categorical distribution으로부터 나온 random variable입니다. 세 가지 learnable components가 있습니다:

  1. denoiser neural network : DisCo-Diff의 DM에 해당하며, diffusion time 와 discrete latent 에 conditioned된 denoised images를 predict합니다.
  2. encoder : clean images 가 주어졌을 때 discrete latents를 infer하는 데 사용됩니다. 각 discrete latent에 대해 개의 categories에 대한 categorical distribution을 outputs합니다.
  3. post-hoc auto-regressive model : 학습된 discrete latents 의 distribution을 로 approximates합니다.

DisCo-Diff의 training process는 두 단계로 나뉩니다. 첫 번째 단계에서는 denoiser 와 encoder 가 end-to-end 방식으로 함께 co-optimized됩니다. 이는 denoising score matching objective (Eq. 2에 표현된)를 각 data 와 연관된 learnable discrete latents 를 포함하도록 확장함으로써 달성됩니다:  score 의 reparameterization 학습에 초점을 맞추는 Eq. 2의 standard objective와 달리, 우리 approach의 denoiser는 본질적으로 conditional score 의 reparameterization을 learning하고 있습니다. 여기서 probability density functions의 convolution은  입니다. 이 conditional score는 DM을 encoder 에 의해 inferred된 discrete latents 에 conditioning함으로써 비롯됩니다. 만약 각 sub-distribution 에 대한 score가 단순화된다면, denoiser network 는 time-dependent score를 더 잘 capture할 수 있습니다 (즉, 감소된 loss를 달성). 따라서, clean input data에 접근할 수 있는 encoder 는 유용한 information을 discrete latents로 encode하여 denoiser가 data를 더 정확하게 reconstruct하도록 돕도록 장려됩니다. discrete latent variables 의 sampling을 통해 encoder로 gradients를 naively backpropagating하는 것은 불가능합니다. 따라서 training 중에는 Gumbel-Softmax distribution에 기반한 continuous relaxation에 의존합니다 (자세한 내용은 App. D 참조).

denoiser network를 training할 때, 우리는 0.1의 확률로 discrete latent variables를 정보가 없는 null-embedding으로 무작위로 대체합니다. 이를 통해 DM은 discrete latent variable-conditioned score와 regular, unconditional score를 모두 learn합니다. sampling 중에는, model 자신의 discrete latents에 대해 classifier-free guidance를 위해 이 score들을 combine하고, 그 conditioning effect를 amplify할 수 있습니다 (자세한 내용은 App. B 참조).

우리는 DisCo-Diff를 discrete latents와 DM을 decoder로 사용하는 variational autoencoder (VAE)로 해석할 수 있습니다. VAE는 종종 latents에 regularization을 사용합니다. 우리는 매우 low-dimensional latent variables (예: ImageNet experiments에서 10개)와 상대적으로 작은 codebooks를 사용하기 때문에 이것이 필요하지 않다고 판단했습니다. 더욱이, 우리는 Gumbel-Softmax relaxation에서 엄격하게 0이 아닌 temperature를 사용하여 stochasticity를 장려합니다.

두 번째 단계에서는, train된 encoder를 통해 clean data를 통과시켜 정의된 discrete latent variables의 distribution 를 capture하기 위해 autoregressive model 를 train합니다. 우리는 다음과 같이 maximum likelihood objective를 사용합니다:  우리는 을 상대적으로 작은 수로 설정하기 때문에, model이 이러한 짧은 discrete vectors를 처리하기가 매우 쉬워져 이 second-stage training이 효율적이 됩니다. 또한 이 autoregressive component로 인한 DM 상의 추가적인 sampling overhead는 무시할 수 있게 됩니다. inference time에, DisCo-Diff를 사용하여 novel samples를 generate할 때, 우리는 먼저 autoregressive model에서 discrete latent variable을 sample한 다음, ODE 또는 SDE solver로 DM을 sample합니다. training 및 sampling을 위한 algorithm pseudocode는 Appendix C에서 제공합니다.

3.2. Motivation and Related Work

이제 우리는 우리의 design choices를 비판적으로 논의하고 동기를 부여하며, 가장 관련 있는 related works에 대해서도 논의할 것입니다. related work에 대한 확장된 논의는 App. A를 참조하십시오.

The curvature of diffusion models. DMs는 더 간단한 ODE-based formulation( in Eq. (1))에서 복잡한 noise-to-data mapping을 learn합니다. noise는 해석적으로 다루기 쉬운 unimodal Gaussian distribution에서 추출됩니다. data가 이 distribution에 encoded되므로, 우리는 이 high-dimensional Gaussian distribution을 DM의 continuous latent variables로 간주할 수 있습니다 (DMs는 일반적으로 deep latent variable models로 볼 수 있음). 그러나, 구조화되지 않은 noise에서 다양하고 일반적으로 multimodal한 data distribution으로의 mapping은 필연적으로 매우 complex해야 합니다. 이는 강한 curvature를 가진 고도로 non-linear한 generative ODE에 해당하며, 이는 learn하기 어렵고 미세한 discretization을 요구하여 synthesis를 느리게 만듭니다. 이 점을 설명하기 위해, 우리는 간단한 2D mixture of Gaussians에서 DM을 train했고, data 근처에서 구부러진 ODE trajectories를 관찰했습니다 (Fig. 3, 중간). 이 효과는 high dimensions에서 훨씬 더 강합니다.

A simpler mapping with discrete latent variables. DisCo-Diff에서 discrete latents의 역할은 이 complexity를 줄이고 DM의 learning task를 더 쉽게 만드는 것입니다. 단일 noise-to-data mapping은 효과적으로 더 간단한 mappings의 set으로 partitioned되며, 각 mapping은 generative ODE에서 더 적은 curvature를 가집니다. 우리는 전체 multimodal complex data distribution을 단일 continuous Gaussian distribution에 mapping하는 것이 부자연스럽다고 주장합니다. 대신, 우리는 이상적인 generative model이 discrete와 continuous latent variables를 모두 combine해야 한다고 믿습니다. 여기서 discrete latents는 global multimodal structure를 capture하고 continuous latents는 local continuous variability를 model합니다. 이를 염두에 두고, 우리는 작은 codebooks를 가진 적당한 수의 discrete latents만 사용할 것을 제안합니다. 한편으로, 몇 개의 latents만으로도 DM의 learning task를 상당히 단순화할 수 있습니다. 다른 한편으로, 작은 codebooks를 가진 몇 개의 latents만 있다면, discrete latent variable distribution 자체에 대한 generative model(우리의 경우 autoregressive one)을 training하는 것이 간단해질 것입니다 (이는 우리가 관찰한 바이며, Sec. 4 참조).

Validation in 2D. 우리의 reasoning을 validate하기 위해, toy 2D mixture of Gaussians를 다시 살펴보겠습니다. Fig. 3, 오른쪽에서 우리는 DisCo-Diff model의 synthesized data를 보여줍니다. discrete latents는 different modes를 capture하도록 learn하고, DisCo-Diff의 DM component는 individual modes를 models합니다. 다른 latents에 대한 DM의 ODE trajectories는 이제 거의 완벽하게 직선이며, 이는 간단한 conditional score function을 나타냅니다. Fig. 4, 왼쪽에서 우리는 전체 diffusion time 에 걸쳐 강하게 감소된 curvature를 정량적으로 보여줍니다. Fig. 4, 오른쪽에서, network complexity의 척도로서 우리는 사용된 denoiser networks의 Jacobians의 norms도 보여줍니다. 우리는 모든 에 대해 DisCo-Diff에서 상당히 감소된 norms를 보며, 이는 denoiser의 task가 실제로 강하게 단순화되고 더 적은 network capacity가 필요함을 시사합니다.

Using few, global latents with relatively small codebooks is important. DisCo-Diff는 discrete latent variables를 사용하는 대부분의 contemporary generative modeling frameworks와 근본적으로 다릅니다. 이러한 연구들은 autoencoders를 사용하여 images를 inputs의 spatially-arranged, downsampled representations으로 전체적으로 encode하며, 다양한 training data modes를 capturing하는 것보다 image fidelity를 preserving하는 데 더 중점을 둡니다. 그러나 이것 또한 부자연스럽습니다: images의 smooth pose, shape, 또는 color variations와 같은 continuous variability를 discrete latents로 encoding하려면 매우 큰 codebooks를 사용해야 하며, 게다가 이러한 models는 일반적으로 매우 high-dimensional spatial grids of discrete latents에 의존합니다 (예: codebooks >1,000인  latents, 반면 우리는 main image models에서 codebook size가 100인 10개의 latents만 사용). 이로 인해 이러한 유형의 models에서는 discrete latents에 대한 distribution을 learning하는 것이 매우 어려운 반면, DisCo-Diff에서는 DM을 보완할 뿐이므로 간단합니다. DisCo-Diff에서는, 단지 몇 개의 global latents를 사용하여 continuous와 discrete latent variables 모두에서 최상의 것을 얻습니다.

End-to-end training is essential. DisCo-Diff의 discrete latents는 정신적으로 non-learnt conditioning information을 활용하는 것과 유사합니다. 지적된 바와 같이, 이는 strong class-conditional 또는 text-to-image DMs와 같은 high-performance generative models의 training을 용이하게 하는 데 결정적이었습니다. 그러나 DisCo-Diff는 주어진 conditioning data에 의존하기보다는 근본적으로 문제를 해결하는 것을 목표로 합니다. 더욱이, data는 예를 들어 class label이 주어져도 일반적으로 상당한 variability를 가집니다. 우리의 discrete latents는 complexity를 더욱 줄일 수 있습니다 (관찰된 바와 같이, Sec. 4 참조). 그러나, CLIP이나 다른 pre-trained encoder networks를 사용하여 condition할 encodings를 생성하고, 그 distribution을 두 번째 단계에서 model할 수 있을까요? 이는 previous works에서 탐구되었지만 중요한 단점이 있습니다: (i) 가장 중요한 단점은 이러한 encoders가 보편적으로 사용 가능하지 않고, 일반적으로 images에만 해당된다는 것입니다. 그러나 우리는 보편적으로 적용 가능한 framework를 개발하고자 합니다. 예를 들어, 우리는 DisCo-Diff를 적합한 pre-trained networks가 없는 molecular docking에도 적용합니다 (Sec. 4.2 참조). (ii) DisCo-Diff에서 discrete latents의 역할은 DM의 denoising task를 더 쉽게 만드는 것이며, 이는 large noise levels에서 특히 ambiguous합니다 (실제로, 우리는 latents가 특히 이러한 high noise levels에서 loss를 줄이는 데 도움이 된다는 것을 발견했습니다, Fig. 7 참조). 이를 위해 latents가 data에 대한 어떤 information을 가장 잘 encode해야 하는지는 명확하지 않습니다. DM objective 자체와 jointly하게 learning함으로써, latents는 DM이 더 나은 denoisers와 lower curvature generative ODEs를 learn하도록 돕기 위해 직접적으로 trained됩니다. (iii) 두 번째 단계에서는 encodings에 대한 generative model이 trained되어야 합니다. DisCo-Diff에서는, DM의 denoising task를 단순화하는 동시에 두 번째 단계에서 autoregressive model의 쉬운 learning을 용이하게 하기 위해 적절한 수의 latents와 codebook size를 자유롭게 선택할 수 있습니다. pre-trained encoders를 사용할 때는, generative modeling을 위해 개발되지 않은 이러한 방법들의 encodings를 사용해야 합니다. 우리는 DisCo-Diff의 강력한 generation performance를 end-to-end learning 덕분이라고 생각합니다.

The latent variables must be discrete. auxiliary continuous latent variables를 사용할 수도 있을까요? continuous latents에 대한 generative models는 거의 항상 uni-modal Gaussian distribution에서 latents의 distribution으로의 mappings에 기반합니다. 따라서, 만약 이러한 continuous latents가 DisCo-Diff의 discrete latents처럼 main DM의 denoising task를 단순화하기 위해 data의 multimodal structure를 learnt다면, 두 번째 단계에서 그것들에 대한 distribution을 learning하는 것은 다시 Gaussian noise에서 multimodal encodings로의 고도로 non-linear하고 learn하기 어려운 mapping을 필요로 할 것입니다. 이것이 바로 DisCo-Diff가 처음부터 해결하고자 하는 문제입니다. 몇몇 연구는 DMs를 non-spatial continuous latent variables로 augment하지만, 그들은 semantic face image manipulation에만 초점을 맞춥니다. InfoDiffusion은 DMs를 discrete latent variables에 conditions하지만, disentangled representations을 learning하는 데 초점을 맞추며, 주로 low-resolution face synthesis를 위해 mutual information-based objective를 사용합니다. DisCo-Diff와는 반대로, 이들 중 어느 연구도 어렵고 다양한 datasets에 대한 high-quality synthesis를 다루지 않습니다.

우리의 ablation studies (Sec. 4.1)에서, 우리는 여기서 제시한 design choices와 motivations를 더욱 validate합니다.

3.3. Architecture

논의된 바와 같이, DisCo-Diff는 data의 global underlying discrete structure를 capture하기 위한 learnable discrete latent variables를 통합하여 continuous DMs의 training을 향상시킵니다. DisCo-Diff가 의도한 대로 작동하도록 보장하기 위해서는 적합한 network architectures가 필요합니다. 아래에서는 image synthesis를 위한 DisCo-Diff에 초점을 맞춰 우리의 design choices를 요약합니다. 그러나, 이 framework는 general하며, clean input data에서 discrete latents를 infer하기 위한 encoder와 이러한 discrete latents를 denoiser network에 integrates하는 conditioning mechanism만 필요합니다. 실제로, 우리는 우리 model을 2D toy data와 molecular docking에도 적용합니다.

Encoder. image modeling을 위해, 우리는 encoder의 backbone으로 ViT를 활용합니다. 우리는 ViTs의 classification mechanism을 확장하고, 각 discrete token을 다른 classification token으로 취급합니다. 구체적으로, 우리는 image patches의 sequence에 개의 추가 classification tokens를 추가합니다. 이 architectural design은 각 discrete latent가 data classification을 수행하는 것과 유사하게 images의 global characteristic을 효과적으로 capture하도록 자연스럽게 허용합니다.

Discrete latent variable conditioning. image experiments의 경우, DisCo-Diff의 denoisers는 DMs에 널리 사용되는 U-Nets입니다. discrete latent variable conditioning을 위해, 우리는 cross-attention을 활용합니다. text-to-image generation에서 영감을 받아, DisCo-Diff의 discrete latents는 text와 유사하게 기능하여 denoiser의 output에 global influence를 미칩니다. 구체적으로, image features는 queries 역할을 하고 discrete latents는 cross-attention layer에서 keys와 values가 되어, discrete latents가 image features를 global하게 형성할 수 있도록 합니다. 우리는 U-Net 내의 각 self-attention layer 뒤에 cross-attention layer를 추가합니다. 우리의 main models에서는 모든 discrete latents가 모든 cross-attention layers에 제공됩니다.

Group hierarchical models. discrete latents의 interpretability를 향상시키기 위해, 우리는 또한 U-Net architecture에 내재된 inductive bias를 탐구하고, Fig. 5에 표시된 것처럼 U-Net의 up-sampling branch에 있는 다양한 resolution features에 서로 다른 latent groups을 feed합니다. 이 approach는 StyleGAN에서 영감을 얻었으며, 여기서 서로 다른 latents가 다른 resolutions에서 도입되어 각각이 neural network의 inductive bias에 의해 다른 image characteristics를 capture할 수 있도록 합니다. 이 design은 group hierarchy를 조성하며, higher-resolution features와 연관된 groups은 lower-resolution features와 관련된 groups에 conditioned된 supplementary information을 제공합니다. 우리는 이 정제된 model을 group hierarchical DisCo-Diff라고 부릅니다.

molecular docking task에서, 기존 denoisers는 molecular structures를 나타내는 3D point clouds에 대해 permutation equivariant 방식으로 message passing을 통해 작동합니다. 우리는 이 property와 architectural bias를 latent variables에 직접 구축하여, point cloud의 한 node를 나타내는 값을 갖도록 합니다 (따라서 모든 point cloud에 대해 codebook size는 nodes의 수와 같습니다). 이 latent design choice는 encoder가 structure에서 key roles를 하는 atoms를 결정한다는 직관과 일치하며, latents가 모든 node에 대한 additional features를 단순히 represent하는 score model의 minimal modification을 허용합니다. encoder는 또한 유사한 equivariant message passing, e3nn, network로 구성되며, 여기서 각 node에 대해 latent당 하나의 logit이 predicted됩니다. molecular docking task를 위한 architecture에 대한 자세한 내용은 App. E.4에서 찾을 수 있습니다.

discrete latents의 distribution에 대한 auto-regressive model은 image experiments에서 standard Transformer decoder를 사용하여 구현됩니다. molecular docking의 경우, 이는 protein structure와 molecular graph의 conditioning information을 feed받는 e3nn network를 다시 사용합니다. 일반적으로, DisCo-Diff는 denoiser와 auto-regressive model에 inputs으로 추가될 수 있는 class labels와 같은 다른 conditional inputs와 호환됩니다. 우리는 단순성을 위해 auto-regressive model을 사용하며, DisCo-Diff의 second stage가 discrete state diffusion models와 같은 다른 discrete data generative models와도 동일하게 잘 작동할 것으로 기대합니다. 2D toy data experiments를 포함한 architecture details는 App. F에 있습니다.

architecture design을 둘러싼 중요한 질문은 적절한 수의 latents와 codebook size를 선택하는 방법입니다. 직관적으로, latents의 수와 codebook size를 늘리는 것이 더 복잡한 data structures를 capture하여 reconstruction error를 더욱 줄이는 간단한 방법처럼 보일 수 있지만, 이 approach는 추가적인 complexity도 도입합니다. 특히, 더 큰 latents set이나 확장된 codebook size는 auto-regressive model의 task를 복잡하게 만들어 잠재적으로 increased errors로 이어질 수 있습니다. 따라서, 더 상세한 discrete structures를 통해 model performance를 향상시키는 것과 auto-regressive model에 대한 관리 가능한 modeling complexity를 유지하는 것 사이의 balance를 찾는 것이 중요합니다. 우리는 현재 diffusion models에 대해 적당한 수의 latent(예: 10-30)와 codebook size(예: 50-100)를 사용할 것을 권장하며, optimal hyper-parameters의 탐색은 future works에 맡깁니다. 예를 들어, 우리의 image generation experiments에서, 우리는 100의 codebook size를 가진 10개의 latents 구성이 복잡한 ImageNet dataset에서 performance를 크게 향상시킨다는 것을 발견했습니다. 우리는 computational constraints로 인해 이 hyper-parameter를 tune하지 않았습니다. 우리는 latents의 정확한 수와 codebook size에 대한 더 신중한 hyper-parameter optimization이 DisCo-Diff의 performance를 더욱 boost할 것이라고 믿습니다.

DisCo-Diff (Section 3) 정리노트 (for AI Researchers)

1. Generative Model & Training (Sec 3.1)

  • 핵심 구성요소 (3 Learnable Components):
    1. Denoiser (): 표준 DM denoiser에 discrete latent 를 추가적인 condition으로 입력받음.
    2. Encoder (): Clean data 로부터 discrete latent 를 infer (추론)하는 역할을 수행.
    3. Autoregressive Model (): Training 완료 후, sampling 시 latent 의 분포 를 model하기 위해 post-hoc으로 학습되는 모델.
  • 2단계 학습 프로세스 (Two-Stage Training):
    • Stage 1: 와 의 End-to-End 최적화
      • Objective: 표준 denoising score matching loss를 encoder가 추론한 에 대한 기댓값으로 확장.
      • 핵심 원리: 는 conditional score 를 학습. 동시에 는 이 conditional score 학습을 가장 쉽게 만들어주는 (즉, loss를 최소화하는) latent 를 생성하도록 학습됨.
      • 기술적 해결: discrete한 의 sampling 과정에 gradient를 흘려보내기 위해 Gumbel-Softmax trick을 사용.
      • Guidance 기법: Training 시 10% 확률로 를 null-embedding으로 치환하여, inference 시 모델 자신의 latent에 대한 classifier-free guidance가 가능하도록 학습.
    • Stage 2: (Autoregressive Model) 학습
      • 학습된 를 고정(freeze)하고, 전체 dataset에 대해 latent 를 추출하여 latent dataset을 구축.
      • 이 latent dataset을 이용해 의 분포 를 maximum likelihood로 학습. 는 저차원 discrete sequence이므로 이 과정은 매우 효율적임.

2. 설계 동기 및 핵심 주장 (Sec 3.2)

  • "Curvature" 문제 해결: 표준 DM은 복잡한 multimodal 분포를 단일 unimodal Gaussian에 mapping해야 하므로, generative ODE의 curvature가 매우 높아 학습이 어렵고 synthesis가 느림. DisCo-Diff는 latent 를 통해 이 mapping을 여러 개의 단순한 conditional mapping으로 분할하여, 각 mapping의 curvature를 낮춤.
  • End-to-End의 필요성: Pre-trained encoder(e.g., CLIP)를 사용하는 대신 end-to-end 학습을 고수하는 이유.
    1. 범용성 (Universality): Pre-trained encoder는 image 외의 domain(e.g., molecular)에는 존재하지 않음.
    2. 과업 최적화 (Task-Alignment): Latent는 DM의 denoising task를 가장 쉽게 만들도록 직접적으로 최적화됨. Pre-trained encoder의 feature는 이 목적에 최적이 아닐 수 있음.
    3. 단순성 (Simplicity): Latent의 차원()과 codebook 크기()를 Stage 2의 autoregressive model이 학습하기 쉽게 설계 가능.
  • "Discrete" Latent의 필요성: 왜 auxiliary latent가 continuous가 아닌 discrete여야 하는가?
    • 만약 auxiliary latent가 continuous이면서 multimodal 분포를 학습한다면, 이 latent를 sampling하기 위한 또 다른 generative model이 필요해짐. 이는 결국 해결하려던 원점의 "Curvature" 문제로 되돌아가는 것과 같음. Discrete latent는 이 순환을 끊어줌.

3. 아키텍처 (Sec 3.3)

  • Image Domain Architecture:
    • Encoder: ViT 기반. Image patch sequence에 개의 추가적인 classification token을 두어 개의 latent를 각각 추론.
    • Denoiser: 표준 U-Net 사용.
    • Conditioning 방식: U-Net의 각 self-attention 블록 뒤에 cross-attention 레이어를 추가하여 의 정보를 주입.
  • 계층적 모델 (Group Hierarchical Model): StyleGAN에서 영감을 받아, U-Net의 resolution 레벨별로 다른 latent group을 주입하여 latent의 interpretability를 향상시키는 변형 모델도 제안.
  • Domain 확장성: Molecular docking task에서는 permutation equivariant 특성을 가진 e3nn 네트워크를 모든 component에 적용하는 등, domain 특성에 맞게 architecture를 유연하게 변경 가능함을 보임.
  • Hyperparameter (): Latent 수()와 codebook 크기()는 trade-off 관계. 너무 크면 reconstruction은 유리하지만 Stage 2의 autoregressive 모델링이 복잡해짐. 논문에서는 ImageNet에 을 사용했으며, 추가적인 tuning으로 성능 향상 여지가 있다고 언급.

쉬운 설명 :

이 섹션은 DisCo-Diff 모델의 구체적인 설계도와 작동 원리, 그리고 왜 그렇게 설계했는지를 설명하는 가장 핵심적인 부분입니다.

핵심 비유: '스타일 분석가'와 '트렌드 예측가'를 고용하여 화가를 돕는 법

  1. 모델의 구성과 학습 (어떻게 팀을 꾸리고 훈련시키는가?)
    • 팀 구성: 실력이 뛰어난 메인 화가(Denoiser)를 돕기 위해 두 명의 전문가를 추가로 고용합니다.
      1. 스타일 분석가 (Encoder): 수많은 사진을 보고 "눈밭 스타일", "숲속 스타일" 등 그림의 전체적인 분위기를 결정하는 몇 개의 핵심 '스타일 카드'(discrete latent)를 자동으로 만들어내는 전문가입니다.
      2. 트렌드 예측가 (Autoregressive Model): 스타일 분석가가 만든 카드들을 보고, 어떤 스타일이 가장 인기 있는지, 어떤 스타일들이 서로 잘 어울리는지 학습하는 전문가입니다.
    • 훈련 방식 (2단계 훈련):
      • 1단계: '화가'와 '스타일 분석가'가 한 팀이 되어 훈련합니다. 분석가는 화가가 그림을 가장 그리기 쉽게 도와주는 최적의 스타일 카드를 만들려고 노력하고, 화가는 그 카드를 사용해 그림 그리는 연습을 반복합니다. 이 end-to-end 협업을 통해 환상의 호흡을 맞춥니다.
      • 2단계: 훈련이 끝나면, '트렌드 예측가'가 분석가가 만든 모든 스타일 카드들을 학습하여 트렌드를 파악합니다.
  2. 설계 이유 (왜 이런 방식을 선택했는가?)
    • 화가의 부담을 덜어주기 위해: 화가에게 복잡한 그림 전체를 알아서 그리라고 하는 대신, "눈밭 스타일"이라는 카드 한 장을 주면 화가는 스타일에 대한 고민 없이 디테일에만 집중할 수 있습니다. 즉, 복잡한 한 문제를 간단한 여러 문제로 나눈 것입니다.
    • 맞춤형 전문가가 최고: 유명 잡지의 스타일 분류(기존에 pre-trained된 encoder)를 빌려오는 것보다, 우리 화가를 전담하는 '스타일 분석가'를 고용하는 편이 훨씬 효율적입니다. 우리 팀의 목표에 꼭 맞는 맞춤형 도움을 줄 수 있기 때문입니다.
    • 선택은 단순하게: 스타일을 "눈밭 73%, 숲속 27%"처럼 애매한 연속적인 값(continuous)으로 주는 것보다, "1번: 눈밭", "2번: 숲속" 처럼 명확하게 구분되는 카드(discrete) 중에서 고르게 하는 것이 훨씬 쉽습니다. 애매한 값을 또다시 고민하는 것은 원래 문제와 다를 바 없기 때문입니다.
  3. 작업실 환경 (Architecture)
    • 화가(Denoiser)는 U자 모양의 특수 캔버스(U-Net)를 사용합니다.
    • 스타일 분석가(Encoder)는 사진을 분석하기 위해 ViT라는 최첨단 분석 도구를 사용합니다.
    • 스타일 카드의 내용은 cross-attention이라는 특별한 방식으로 그림에 반영됩니다.

요약하자면, DisCo-Diff는 똑똑한 보조 모델(Encoder, Autoregressive Model)들을 도입하고 이들을 메인 모델(Denoiser)과 유기적으로 훈련시켜, 근본적으로 어려운 generation 문제를 훨씬 쉽고 효율적인 문제로 바꾸는 영리한 framework라고 할 수 있습니다.

 

 
 
 
 
 
 
더보기

 주인장 이해

 

1. 가우시안 노이즈부터 쭉 목표 이미지까지 생성하는건 사실 매우 복잡함.

2. 그걸 해결하기 위해 learnable한 latent variables라는 개념을 저장하는 어떤 잠재변수를 만들어서 복잡성을 해결하고자 한다.

3. 2번에서 학습시킨 개념은 거시적인 mode 정보를 담고 있어서 디퓨전 모델은 노이즈 제거 과정에서 매 스텝마다 attention을 통해 가장 관련 있는 잠재 변수를 참고하면서 노이즈를 제거하게 된다.




학습 알고리즘

1. 인코더 -> z를 만듬
2. z와 c를 받아서 디노이저는 깨끗한 이미지를 만듬
3. 재구성 loss를 통해 인코더와 디노이저 학습




근데 오토리그레시브 모델도 학습해야함.

클래스 레이블을 주고 z1~z30을 맞추게함. 이걸 인퍼런스에 활용