Normalizing Flows with Iterative Denoising

논문 요약: Normalizing Flows with Iterative Denoising (iTARFlow)

1. 연구 배경 및 문제 정의 (Problem)

최근 이미지 생성 분야에서 **확산 모델 (Diffusion Models)**과 **자기회귀 모델 (Autoregressive Models)**이 주류를 이루고 있습니다. 확산 모델은 높은 화질을 제공하지만, 수백 단계의 반복적 디노이징 과정으로 인해 추론 비용이 크고 확장성 (Scaling) 예측이 어렵다는 단점이 있습니다. 반면, 자기회귀 모델은 빠른 샘플링과 명확한 확장성을 가지지만, 이산 토큰으로의 양자화 (Quantization) 과정에서 정보 손실이 발생하거나, 연속 픽셀 공간에서 작동할 경우 시퀀스 길이가 너무 길어 학습 및 샘플링이 비효율적입니다.

이러한 맥락에서 **정규화 흐름 (Normalizing Flows, NF)**은 확률 밀도 기반의 정확한 가능도 (Likelihood) 학습과 연속 공간 표현의 장점을 결합할 수 있는 대안으로 재조명받고 있습니다. 특히 최근 TARFlow 는 NF 가 이미지 모델링에서 경쟁력 있는 성능을 보임을 입증했습니다.

그러나 기존 NF 기반 모델들은 **노이즈 딜레마 (Noise Dilemma)**라는 근본적인 문제에 직면해 있습니다:

• 노이즈가 너무 작을 때: 모델은 국소적인 텍스처는 풍부하지만 전역적인 구조가 엉망인 이미지를 생성합니다. 또한, 작은 노이즈 하에서의 가능도 학습은 불안정해집니다.
• 노이즈가 너무 클 때: 모델은 전역 구조는 잘 잡지만, 트위디 (Tweedie) 보조정리에 따라 디노이징 과정에서 이미지가 과도하게 흐려지고 세부 사항이 손실됩니다.

기존의 TARFlow 는 단일 노이즈 레벨에서 학습하거나 중간 정도의 노이즈를 사용했으나, 이 딜레마를 해결하여 고해상도에서 구조적 일관성과 세부 텍스처를 모두 보존하는 방법은 부재했습니다.

2. 제안 방법: iTARFlow (Methodology)

저자들은 **반복적 TARFlow (iterative TARFlow, iTARFlow)**를 제안하여 위 문제를 해결합니다. 핵심 아이디어는 NF 를 다양한 노이즈 레벨에서 학습시키고, 추론 시 확산 모델과 유사한 반복적 디노이징 (Iterative Denoising) 과정을 도입하는 것입니다.

• 다중 노이즈 레벨 학습 (Multi-noise Training):

  • 모델은 $t \in [t_{min}, t_{max}]$ 범위의 연속적인 노이즈 레벨에서 학습됩니다.
  • 학습 목적함수는 다양한 노이즈 스케일 $t$에 대한 로그 가능도 $\log p_\theta(x_t, t)$를 최적화합니다.
  • 이를 통해 모델은 다양한 노이즈 강도에서 데이터 분포를 학습하게 되어, 노이즈 딜레마를 우회할 수 있습니다.

• 샘플링 프로세스 (Two-Phase Sampling):

  1. 자기회귀 생성 (Autoregressive Generation): TARFlow 의 특성을 활용하여, 가우시안 사전분포에서 시작해 $t_{max}$ 노이즈 레벨의 이미지를 패치 (Patch) 단위로 순차적으로 생성합니다. 이 단계에서는 Classifier-Free Guidance (CFG) 를 적용하여 구조적 일관성을 높입니다.
  2. 병렬 반복 디노이징 (Parallel Iterative Denoising): 생성된 노이즈가 많은 이미지 $x_{t_{max}}$를 기반으로, 학습된 모델의 로그 가능도 기울기 (Score function, $\nabla_x \log p_\theta$) 를 자동 미분 (Automatic Differentiation) 을 통해 계산합니다.
     • 이 Score 를 이용해 확률적 ODE (Ordinary Differential Equation) 를 풀며 $t_{max}$에서 $t_{min}$까지 반복적으로 이미지를 정제합니다.
     • 확산 모델과 달리 이 단계는 완전 병렬화가 가능하며, 추가적인 디코더 (VAE 등) 없이 NF 모델 자체만으로 수행됩니다.

3. 주요 기여 (Key Contributions)

1. 노이즈 딜레마의 규명 및 해결: NF 학습에서 노이즈 레벨 선택이 전역 구조와 국소 디테일 간의 트레이드오프를 결정한다는 점을 체계적으로 분석하고, 이를 해결하기 위해 다중 노이즈 학습과 반복적 디노이징을 결합한 새로운 프레임워크를 제시했습니다.
2. 성능의 획기적 개선: 기존 TARFlow 및 STARFlow 와 비교하여 ImageNet(64, 128, 256 픽셀) 에서 FID(Fréchet Inception Distance) 지표를 크게 개선했습니다. 특히 확산 모델과의 성능 격차를 크게 좁혔습니다.
3. 유연한 하이브리드 아키텍처 가능성: iTARFlow 는 사전 학습된 DiT(Diffusion Transformer) 기반 디노이저와도 자연스럽게 결합될 수 있음을 입증했습니다. 이는 NF 와 확산 모델의 장점을 통합한 하이브리드 프레임워크로의 진전을 시사합니다.

4. 실험 결과 (Results)

• ImageNet-64: 770M 파라미터 규모의 iTARFlow-B 모델은 FID 1.68 을 기록하여, 300M 파라미터의 확산 모델 (EDM-SDE, FID 1.55) 과 유사한 성능을 달성했습니다. 기존 TARFlow(2.66) 대비 파라미터를 절반 이하로 줄이면서 성능을 대폭 향상시켰습니다.
• ImageNet-128: 1.2B 파라미터 모델 (iTARFlow-L) 은 FID 2.44 를 기록하여, 2B 파라미터의 Simple Diffusion 모델 (FID 1.94) 에 근접하는 성능을 보였습니다.
• ImageNet-256: 픽셀 공간에서 FID 3.32 를 달성하여, 기존 TARFlow(5.56) 와 STARFlow(4.69) 를 압도했습니다. 또한, 잠재 공간 (Latent Space) 실험에서도 경쟁력 있는 성능을 보였습니다.
• 효율성: 확산 모델에 비해 학습이 더 효율적이며 (600 에포크 vs 1400 에포크), 5~10 단계의 반복 디노이징만으로도 고화질 이미지를 생성할 수 있습니다.

5. 의의 및 한계 (Significance & Limitations)

• 의의:

  • NF 기반 생성 모델이 확산 모델과 자기회귀 모델의 장점을 모두 취할 수 있음을 증명했습니다.
  • NF 의 확장성 (Scalability) 과 확산 모델의 표현력 (Expressiveness) 을 결합하여, 고해상도 이미지 생성의 새로운 패러다임을 제시했습니다.
  • NF 와 DiT 와 같은 확산 기반 모델 간의 호환성을 입증하여, 향후 하이브리드 모델 연구의 토대를 마련했습니다.

• 한계 및 향후 과제:

  • 계산 비용: 반복적 디노이징 과정에서 전체 네트워크에 대한 자동 미분을 수행해야 하므로, GPU 메모리 사용량이 많다는 것이 주요 계산적 병목입니다.
  • 생성 오류 (Artifacts): 고해상도 (256x256) 에서 배경이 검게 붕괴되거나 (Collapsed Background), 첫 번째 자기회귀 토큰의 불일치로 인한 흐릿함 (Blur) 이 발생하는 사례가 관찰되었습니다. 이는 향후 CFG 스케일 조정이나 초기 토큰 학습 전략 개선으로 해결할 수 있는 과제로 남았습니다.
  • 이론적 최적화: 노이즈 범위 ($t_{min}, t_{max}$) 와 가중치 ($\gamma_t$) 등의 최적 하이퍼파라미터를 결정하는 이론적 근거는 아직 부족하여 경험적 탐색에 의존하고 있습니다.

결론적으로, iTARFlow 는 정규화 흐름 (NF) 의 한계를 극복하고 확산 모델에 필적하는 생성 성능을 달성한 획기적인 연구로, 생성 모델 분야의 새로운 방향성을 제시합니다.