Denoising, Fast and Slow: Difficulty-Aware Adaptive Sampling for Image Generation

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

현대적인 확산 (Diffusion) 및 흐름 매칭 (Flow Matching) 기반 이미지 생성 모델들은 계산 자원을 공간적으로 균일하게 할당합니다. 즉, 모든 이미지 패치 (patch) 가 동일한 시간 단계 (timestep) 와 동일한 함수 평가 횟수 (NFE) 로 업데이트됩니다.

• 핵심 문제: 자연 이미지는 이질적입니다. 배경이나 단순한 영역은 탈노이즈 (denoising) 가 쉽지만, 객체 경계, 작은 텍스트, 복잡한 질감 등은 후기 단계까지 모호함을 유지하며 더 많은 정제와 문맥이 필요합니다.
• 기존 방식의 한계: 모든 영역을 동일하게 처리함으로써, 쉬운 영역에는 불필요한 계산 자원이 낭비되고, 어려운 영역에는 필요한 추가 정제와 문맥이 부족해집니다.
• 기존 시도 (Diffusion Forcing) 의 결함: 최근 'Diffusion Forcing'과 같이 패치별로 다른 노이즈 스케일을 적용하는 시도가 있었으나, 단순히 패치별 시간 단계를 무작위로 샘플링하면 훈련 시 모델이 추론 시점에는 존재하지 않는 과도하게 정제된 (clean) 문맥을 접하게 되어 **훈련 - 추론 간극 (Train-Test Gap)**이 발생합니다. 이로 인해 성능이 저하됩니다.

2. 제안 방법: Patch Forcing (PF) (Methodology)

저자들은 이미지 생성을 위해 **패치별 난이도 (Patch Difficulty)**를 인식하고 계산 자원을 동적으로 할당하는 Patch Forcing (PF) 프레임워크를 제안합니다.

2.1. 핵심 아이디어

• 이질적인 탈노이즈 궤적: 이미지 내의 각 패치가 서로 다른 노이즈 궤적을 따르도록 합니다.
• 문맥 제공: 탈노이즈가 쉬운 (신뢰도가 높은) 영역을 먼저 빠르게 진행시켜, 어려운 (불확실성이 높은) 영역에 '미래'의 문맥 (context) 을 제공합니다.

2.2. 주요 구성 요소

A. 훈련 시간 타임스텝 샘플링 전략 (LTG Sampler)

• 문제 해결: 기존 SRM(Spatial Reasoning Models) 방식은 평균 시간 단계를 균일하게 샘플링하지만, 최대 시간 단계 ($t_{max}$) 가 1 에 가까워 훈련 시 과도한 정보가 유출되는 문제가 있었습니다.
• 해결책: 저자들은 Logit-Normal Truncated Gaussian (LTG) 샘플러를 제안합니다.
  1. 먼저 각 샘플의 최대 시간 단계 $t_{max}$를 Logit-Normal 분포에서 샘플링합니다.
  2. 각 패치별 시간 단계 $t_i$는 $t_{max}$를 기준으로 하반부 (lower half) 의 가우스 분포에서 조건부 샘플링합니다.
  3. 이는 훈련 시 모델이 추론과 유사한 (순수한 노이즈에서 시작하는) 상태를 경험하도록 하여 훈련 - 추론 간극을 해소하고 성능을 극대화합니다.

B. 패치 난이도 예측 헤더 (Uncertainty Head)

• 모델의 속도 추정 (velocity estimate) 에 대한 표준 편차를 예측하는 경량 헤더를 추가합니다.
• 이 예측값을 '불확실성 (Uncertainty)'이 아닌 **패치 난이도 (Patch Difficulty)**로 해석합니다.
• 높은 난이도 = 더 많은 계산과 문맥이 필요함.
• 실험 결과, 예측된 불확실성 지도는 패치별 생성 난이도와 높은 상관관계를 보이며, 추가적인 공간 문맥이 제공되면 불확실성이 감소함이 확인되었습니다.

C. 적응형 샘플링 전략 (Adaptive Samplers)
예측된 난이도 지도를 기반으로 추론 시 계산 자원을 동적으로 할당하는 두 가지 전략을 도입합니다.

1. Dual-Loop (이중 루프):
   • 낮은 불확실성 (쉬운) 패치는 큰 시간 단계로 빠르게 진행 (Outer Loop).
   • 높은 불확실성 (어려운) 패치는 진행된 쉬운 패치들의 상태를 문맥으로 활용하여 작은 단계로 정제 (Inner Loop).
   • 두 서브셋이 같은 시간 단계에 도달하면 불확실성을 재평가하고 반복합니다.
2. Look-Ahead (선제적 탐색):
   • 신뢰도 높은 패치들을 현재 시간 $t$에서 미래 시간 $t_{ctx}$로 직접 전진시켜 "미래" 상태를 생성합니다.
   • 이 전진된 상태를 어려운 패치들의 탈노이즈 조건 (conditioning) 으로 사용하여 모호성을 줄입니다.

3. 주요 기여 및 발견 (Key Contributions & Findings)

1. 훈련 - 추론 간극 해소: 패치별 타임스텝을 무작위로 샘플링하는 기존 방식의 문제를 해결하기 위해, 최대 정보량 ($t_{max}$) 을 제어하는 LTG 샘플러를 도입하여 표준 SiT(Scalable Interpolant Transformers) 베이스라인보다 우수한 성능을 달성했습니다.
2. 난이도 인식 적응형 샘플링: 예측된 불확실성 지도를 활용하여 계산 자원을 어려운 영역에 집중시키는 새로운 샘플링 전략 (Dual-Loop, Look-Ahead) 을 제안했습니다.
3. 세 가지 핵심 발견:
   • 추가적인 문맥은 생성 품질을 향상시킵니다.
   • 예측된 불확실성은 패치 난이도와 강하게 상관관계가 있습니다.
   • 추가적인 공간 문맥 (쉬운 영역의 진행 상태) 은 어려운 영역의 불확실성을 감소시킵니다.
4. 직교성 (Orthogonality): 제안된 방법은 기존 Classifier-Free Guidance (CFG) 나 표현 정렬 (Representation Alignment, REPA) 방법과 직교하며, 이를 결합했을 때 추가적인 성능 향상을 보입니다.

4. 실험 결과 (Results)

• ImageNet (Class-conditional):
  • SiT-B/2, L/2, XL/2 모델에서 Patch Forcing Transformer (PFT) 는 표준 Euler 샘플링만으로도 SiT 베이스라인을 능가합니다.
  • Look-Ahead 샘플러를 적용 시 SiT-XL/2 의 FID-50k 가 17.2 에서 9.8로 획기적으로 개선되었습니다 (REPA 와 결합 시 6.7 까지 향상).
  • 고정된 NFE(함수 평가 횟수) 하에서 모든 스케일에서 일관된 성능 향상을 보였습니다.
• Text-to-Image Synthesis:
  • 12 억 파라미터 규모의 PFT 를 COYO 데이터셋으로 학습시켰습니다.
  • T2I-CompBench++ 및 GenEval 벤치마크에서 경쟁력 있는 결과를 보였습니다.
  • 텍스트 렌더링 향상: 기존 Flow Matching 모델 대비 Patch Forcing 모델이 텍스트 생성의 정확도와 가독성 (OCR 평가 기준) 에서 현저히 우수한 성능을 보였습니다. 이는 이질적인 노이즈 스케일 훈련이 텍스트와 같은 정밀한 구조 생성에 유리함을 시사합니다.

5. 의의 및 결론 (Significance)

이 논문은 이미지 생성 모델이 **공간적 이질성 (Spatial Heterogeneity)**을 인식하고, 계산 자원을 가장 필요한 곳에 동적으로 할당할 수 있음을 증명했습니다.

• 패치 수준의 적응성: 모든 픽셀을 동일하게 처리하는 패러다임에서 벗어나, 이미지의 난이도에 따라 "빠르게" 또는 "느리게" 탈노이즈하는 새로운 접근법을 제시했습니다.
• 효율성과 품질의 동시 달성: 동일한 계산 비용 (NFE) 으로 더 높은 품질의 이미지를 생성하거나, 동일한 품질을 더 적은 비용으로 달성할 수 있는 가능성을 열었습니다.
• 미래 방향: 이 방법은 증류 (Distillation) 모델이나 소수 단계 (Few-step) 생성 모델로 확장될 수 있으며, 적응형 생성 (Adaptive Generation) 의 새로운 기반이 될 것으로 기대됩니다.

요약하자면, Patch Forcing은 이미지 생성 과정에서 "어떤 부분이 어려운지"를 학습하고, 쉬운 부분을 먼저 해결하여 어려운 부분을 돕는 지능적이고 적응적인 탈노이즈 전략을 성공적으로 구현한 연구입니다.