기술 요약: MixFlow 학습

문제 정의

본 논문은 확산(diffusion) 및 플로우 매칭(flow matching) 모델에서 흔히 발생하는 학습-테스트 불일치(training-testing discrepancy), 즉 노출 편향(exposure bias) 문제를 다룹니다. 학습 과정에서 예측 네트워크는 각 타임스텝에서 정답 노이즈 데이터(노이즈와 데이터의 보간물)를 입력으로 받습니다. 반면, 테스트(샘플링) 시에는 이전 단계에서 생성된 노이즈가 섞인 데이터를 입력으로 받습니다. 이러한 불일치는 예측 오류, 오류 누적 및 샘플링 드리프트(sampling drift)를 유발합니다. 기존의 해결책들은 샘플링 과정을 수정하거나(예: Time-Shift Samplers, Epsilon Scaling) 학습 중 입력에 섭동(perturbation)을 가하는 방식에 집중해 왔으나, 본 연구는 이 문제를 직접적으로 완화하기 위한 새로운 학습 절차를 제안합니다.

방법론: MixFlow

제안된 MixFlow 접근 방식은 "Slow Flow" 현상에 대한 관찰에서 착안되었습니다. 경험적 분석 결과, 샘플링 타임스텝 $t$에서 생성된 노이즈 데이터와 가장 가까운 정답 노이즈 데이터는 종종 더 높은 노이즈를 가진 타임스텝(즉, "느려진 타임스텝", $m_t \leq t$)에 해당한다는 것이 밝혀졌습니다. 다시 말해, 생성된 데이터가 정답 궤적보다 뒤처지는 현상이 발생합니다.

이러한 관찰을 활용하기 위해 MixFlow는 다음과 같이 학습 절차를 수정합니다:

1. 느려진 보간 혼합(Slowed Interpolation Mixture): 학습 타임스텝 $t$에 대응하는 단일 보간물 $x_t$로 예측 네트워크 $u_\theta$를 학습시키는 대신, MixFlow는 "느려진" 타임스텝들의 범위에서 가져온 보간물들의 혼합(mixture)을 사용합니다.

   • 입력 노이즈 데이터는 $\mathcal{X}_t = \{x_{m_t} \mid m_t \in [(1-\gamma)t, t]\}$ 집합에서 샘플링됩니다. 여기서 $m_t$는 더 높은 노이즈(더 느린) 타임스텝을 나타냅니다.
   • 범위의 크기는 계수 $\gamma$에 의해 제어되며 $t$에 따라 선형적으로 변합니다.

2. 수정된 손실 함수: 학습 목적 함수는 느려진 보간물 $x_{m_t}$를 입력으로 사용하여 예측 네트워크를 최적화하되, 타겟은 느려진 타임스텝 $m_t$에서의 정답 속도(ground-truth velocity)를 유지합니다:
   $$\mathbb{E}_{t, x_0, x_1, m_t} [\| u_\theta(x_{m_t}, t) - u^*(x_{m_t}, m_t) \|_2^2]$$
   결정적으로, 네트워크에 전달되는 입력 타임스텝은 $m_t$가 아니라 학습 타임스텝 $t$입니다.

3. 샘플링 전략:

   • 느려진 타임스텝 ($m_t$): $[(1-\gamma)t, t]$ 범위에서 균등하게(uniformly) 샘플링됩니다.
   • 학습 타임스텝 ($t$): $(x_{m_t}, t)$ 쌍이 가능한 모든 조합에 대해 균등하게 샘플링되도록 하기 위해, 학습 타임스텝 $t$는 균등 분포가 아닌 Beta(2, 1) 분포($p(t) \propto t$)에서 샘플링됩니다.

4. 구현: 이 방법은 매우 가볍며, 기존 학습 파이프라인에 약간의 코드 수정(예: RAE의 경우 5줄)만 필요합니다. 또한 사전 학습된 모델에 적용 가능한 사후 학습(post-training) 전략으로 기능합니다.

주요 기여

• Slow Flow 현상 규명: 본 논문은 생성된 노이즈 데이터가 타임스텝 $t$에서 더 느린(높은 노이즈의) 타임스텝의 정답 데이터와 가장 가깝다는 것을 경험적으로 입증하여, 제안된 학습 수정의 이론적 근거를 제공합니다.
• MixFlow 학습 알고리즘: 학습 분포를 테스트 분포와 더 밀접하게 일치시키기 위해 느려진 보간물의 혼합을 사용하는 새로운 학습 체계를 제안하며, 이를 통해 노출 편향을 효과적으로 완화합니다.
• 최첨단 성능(SOTA): 다양한 모델 아키텍처(SiT, REPA) 및 생성 작업(클래스 조건부 및 텍스트-이미지 생성)에서 유의미한 성능 향상을 달elle 성합니다.

실험 결과

저자들은 클래스 조건부 이미지 생성(ImageNet)과 텍스트-이미지 생성(Stable Diffusion 3.5)에 대해 MixFlow를 평가했습니다.

클래스 조건부 생성 (ImageNet)

• RAE 모델: RAE 모델에 적용된 MixFlow는 새로운 SOTA 결과를 달성했습니다:
  • 256 × 256: 1.43 FID (가이드 없음) 및 1.10 FID (가이드 포함).
  • 512 × 512: 1.55 FID (가이드 없음) 및 1.10 FID (가이드 포함).
• SiT 및 REPA: SiT(B 및 XL 사이즈)와 REPA 모델 전반에서 일관된 개선이 관찰되었습니다. 예를 들어, SiT-B/2(가이드 포함)는 gFID 기준 4.46에서 3.91로 개선되었습니다.
• 샘플링 단계: MixFlow의 이점은 적은 샘플링 단계(예: 20단계)에서 더 두드러졌으며, 이는 모델이 거친 샘플링(coarse sampling)에 내재된 근사 오류를 효과적으로 해결하고 있음을 시사합니다.
• 비교: MixFlow는 입력 섭동(Input Perturbation), Epsilon Scaling, Time-Shift Samplers를 포함한 다른 노출 편향 완화 방법들보다 우수한 성능을 보였습니다.

텍스트-이미지 생성

• Stable Diffusion 3.5에 적용되었을 때, MixFlow는 40단계 및 10단계 샘플링 모두에서 모든 평가 지표(GenEval, DPG-Bench, T2I-CompBench)의 성능을 향상시켰습니다.
• 정성적 결과에서는 개수 파악 정확도, 공간 관계, 객체 형태 준수 능력이 향상되었음을 보여주었습니다.

의의 및 주장

본 논문은 MixFlow가 "Slow Flow" 현상을 고려하여 학습 분포를 수정함으로써 노출 편향 문제에 대한 단순하면서도 효과적인 솔루션을 제공한다고 주장합니다. 저자들은 다음을 강조합니다:

1. 학습 vs 추론: 추론 시점의 조정과 달리, MixFlow는 학습 목적 함수를 변경하여 근본 원인을 해결함으로써 모델이 학습-테스트 간의 분포 변화에 강건하게 만듭니다.
2. 호환성: 이 방법은 다양한 백본 아키텍처(SiT, DiT), 보간 방식(linear, variance-preserving), 그리고 고급 학습 패러다임(REPA, RAE)과 호환됩니다.
3. 효율성: 최소한의 계산 오버헤드만으로 SOTA 결과를 달성하며, 단 몇 줄의 코드 수정만 필요하고 효율적인 사후 학습 단계로 작용합니다.

결론적으로 저자들은 MixFlow가 특히 고해상도 및 저단계 샘플링 시나리오에서 생성 품질을 크게 향상시켜, 확산 기반 이미지 생성의 새로운 벤치마크를 설정했다고 밝히고 있습니다.