Improving Efficiency of Diffusion Models via Multi-Stage Framework and Tailored Multi-Decoder Architectures

1. 문제 정의 (Problem)

확산 모델 (Diffusion Models) 은 이미지 생성, 역문제 해결 등 다양한 분야에서 뛰어난 성능을 보이지만, 훈련 및 샘플링 속도가 매우 느리다는 치명적인 단점을 가지고 있습니다. 이 비효율성의 주요 원인은 다음과 같습니다.

• 파라미터 중복성 (Parameter Redundancy): 기존 확산 모델은 모든 시간 단계 (timestep) 에 대해 단일의 통합된 아키텍처를 사용합니다. 그러나 저자들은 고노이즈 수준 (시간 단계가 초기) 에서는 적은 파라미터로도 스코어 함수 (score function) 를 학습할 수 있지만, 저노이즈 수준 (시간 단계가 후기) 에서는 더 많은 파라미터가 필요하다는 것을 발견했습니다. 즉, 통합 아키텍처는 특정 구간에서는 과소적합 (underfitting) 을, 다른 구간에서는 과다파라미터화 (overparameterization) 를 초래합니다.
• 기울기 불일치 (Gradient Dissimilarity): 서로 다른 노이즈 수준에서의 데이터 분포 모양이 달라서 발생하는 기울기 (gradient) 의 불일치로 인해, 경사 하강법 기반의 훈련 과정이 비효율적으로 진행됩니다.
• 기존 접근법의 한계:
  • 통합 아키텍처: 모든 단계에 동일한 모델을 사용하여 자원 낭비가 발생합니다.
  • 분리된 아키텍처 (Separate Architectures): 각 구간마다 별도의 모델을 사용하면 과적합 (overfitting) 문제가 발생하고, 단계 간 간섭 (interference) 이 심화될 수 있습니다.

2. 제안된 방법론 (Methodology)

저자들은 위 문제들을 해결하기 위해 다단계 프레임워크 (Multi-Stage Framework) 를 제안하며, 이는 두 가지 핵심 요소로 구성됩니다.

A. 맞춤형 멀티 디코더 U-Net 아키텍처 (Tailored Multi-Decoder U-Net)

기존의 통합 아키텍처와 완전히 분리된 아키텍처의 장점을 결합한 하이브리드 구조를 설계했습니다.

• 공유 인코더 (Shared Encoder): 모든 시간 구간 (interval) 에서 공통적으로 사용되는 인코더를 도입합니다. 이는 모든 시간 단계에 걸쳐 정보를 공유하여 과적합을 방지하고 훈련 안정성을 높입니다.
• 구간별 전용 디코더 (Stage-specific Decoders): 각 시간 구간마다 별도의 디코더를 할당합니다.
  • 고노이즈 구간 (초기 단계): 학습이 상대적으로 쉬우므로 파라미터 수가 적은 디코더를 사용합니다.
  • 저노이즈 구간 (후기 단계): 세부적인 이미지 복원이 필요하므로 파라미터 수가 많은 디코더를 사용합니다.
• 효과: 계산 자원을 각 단계의 난이도에 맞게 최적화하여 할당하고, 구간 간 기울기 간섭을 줄입니다.

B. 최적 디노이저 기반 시간 단계 클러스터링 (Optimal Denoiser-based Timestep Clustering)

전체 시간 구간 $[0, 1]$을 몇 개의 구간으로 나누는 최적의 분할점 ($t_1, t_2, \dots$) 을 결정하는 새로운 알고리즘을 제안했습니다.

• 원리: 각 구간 내에서 회귀 작업 (regression task) 을 가능한 한 균일하게 (homogenize) 만듭니다.
• 방법: Proposition 1 에서 유도된 최적 디노이저 (Optimal Denoiser) 의 함수적 거리 (functional distance) 를 기반으로 시간 단계들을 클러스터링합니다.
  • 특정 구간 $[0, t_1)$ 내에서 최적 디노이저와 초기 상태 ($t=0$) 간의 평균 함수적 유사도가 임계값 이상인 최대 $t_1$을 찾습니다.
  • 구간 $[t_2, 1]$ 내에서 최적 디노이저와 최종 상태 ($t=1$) 간의 유사도가 임계값 이상인 최소 $t_2$를 찾습니다.
• 이 방식은 기존의 단순 시간 기반, SNR 기반, 또는 기울기 기반 클러스터링보다 더 정교하게 구간을 분할합니다.

3. 주요 기여 (Key Contributions)

1. 비효율성의 근원 규명: 확산 모델 훈련의 비효율성이 '모델 용량 요구 사항의 큰 차이'와 '기울기 불일치'에서 비롯됨을 실증적으로 증명했습니다.
2. 새로운 다단계 아키텍처: 공유 인코더와 구간별 맞춤형 디코더를 결합한 U-Net 구조를 설계하여, 기존 통합/분리 방식의 단점을 모두 보완했습니다.
3. 성능 및 효율성 극대화:
   • 훈련 효율성: CelebA 데이터셋에서 Latent Diffusion Model (LDM) 의 훈련 계산량을 30% 로, EDM 의 경우 82% 로 줄이면서도 유사한 생성 품질을 달성했습니다.
   • 샘플링 품질: CIFAR-10 에서 DPM-Solver 의 FID 점수를 2.84 에서 2.37로, EDM 의 FID 를 2.05 에서 1.96으로 개선했습니다.
   • 샘플링 효율성: 20 번의 함수 평가 (NFE) 만으로 PFGM 이 147 번의 NFE 로 달성한 것과 동등한 FID 점수를 기록했습니다.

4. 실험 결과 (Results)

• 데이터셋: CIFAR-10, CelebA (32x32, 256x256).
• 비교 대상: DDPM, Score SDE, DPM-Solver, EDM, LDM 등 최신 확산 모델.
• 주요 성과:
  • 생성 품질 (FID): 모든 비교 모델 (DPM-Solver, EDM, LDM) 에서 제안된 다단계 아키텍처가 더 낮은 FID 점수를 기록했습니다.
  • 훈련 효율 (PFLOPs): 동일한 생성 품질을 달성하는 데 필요한 총 훈련 계산량 (PFLOPs) 이 기존 모델 대비 대폭 감소했습니다 (예: LDM 기준 30% 수준).
  • 아키텍처 비교 실험: 통합 아키텍처, 분리 아키텍처 (조기 종료 및 파라미터 맞춤 적용 포함), 제안된 다단계 아키텍처를 비교한 결과, 제안된 방식이 가장 우수한 FID (2.35) 를 보였습니다. 이는 공유 인코더가 과적합을 방지하고 수렴을 돕는 동시에, 전용 디코더가 각 구간의 특성을 잘 포착함을 의미합니다.
  • 클러스터링 비교: 제안된 '최적 디노이저 기반 클러스터링'이 기존 방법들 (시간, SNR, 기울기 기반) 보다 더 낮은 FID 를 달성하여 구간 분할의 유효성을 입증했습니다.

5. 의의 및 결론 (Significance)

이 논문은 확산 모델의 가장 큰 병목 현상인 계산 비용과 훈련/샘플링 속도 문제를 해결하기 위한 체계적인 접근법을 제시합니다.

• 자원 최적화: 시간 단계별로 필요한 모델 용량을 정밀하게 분석하여, 불필요한 파라미터를 줄이고 중요한 단계에 자원을 집중시킴으로써 효율성을 극대화했습니다.
• 확장성: 제안된 프레임워크는 unconditional diffusion 모델뿐만 아니라, 조건부 확산 모델 (conditional diffusion) 이나 역문제 해결 모델 등 다양한 응용 분야에 적용 가능하며, 대규모 스테이블 디퓨전 (Stable Diffusion) 같은 모델의 훈련 비용을 획기적으로 낮출 수 있는 잠재력을 가집니다.
• 실용성: 동일한 하드웨어 환경에서 더 빠른 훈련과 더 높은 품질의 생성을 가능하게 하여, 실시간 생성이 필요한 응용 분야에서의 확산 모델 활용도를 높일 것으로 기대됩니다.

요약하자면, 이 연구는 **"공유된 지능 (인코더)"**과 **"구체화된 전문성 (디코더)"**을 결합하고, 이를 이론적으로 최적화된 시간 구간 분할과 연결함으로써 확산 모델의 효율성과 성능을 동시에 비약적으로 향상시켰습니다.