기술 요약: NoiseRater – 확산 모델 훈련을 위한 메타 학습된 노이즈 가치 평가

1. 문제 제기

확산 모델은 점진적인 노이즈 추가 과정을 역으로 학습함으로써 생성 작업에서 최첨단 성능을 달성해 왔습니다. 그러나 표준 훈련 패러다임은 주입된 노이즈를 균일하게 유익한 것으로 간주합니다. 최근 연구는 추론 시간에 노이즈를 최적화하는 것 (예: 가이드 스케일링, 노이즈 리샘플링) 이나 시간 단계 수준에서 노이즈를 재가중치하는 것 (예: P2, Min-SNR) 에 초점을 맞추고 있지만, 훈련 중 인스턴스 수준의 노이즈 변이의 역할은 아직 충분히 탐구되지 않았습니다.

저자들은 근본적인 질문을 제기합니다: 모든 노이즈 실현이 학습에 동등하게 유용한가요? 동일한 시간 단계와 동일한 데이터 샘플에서도 서로 다른 노이즈 인스턴스는 다양한 수준의 학습 신호를 전달할 수 있습니다. 일부는 노이즈 제거를 위한 명확한 감독을 제공할 수 있는 반면, 다른 일부는 모호하거나 중복적이거나 해로울 수 있습니다. 이러한 변이를 균일하게 취급하면 비최적의 학습 역학으로 이어질 수 있습니다. 본 연구는 테스트 시간 노이즈 조작에서 훈련 시간 노이즈 가치 평가로 초점을 이동함으로써 이러한 격차를 해소합니다.

2. 방법론: NoiseRater

제안된 프레임워크인 NoiseRater는 개별 노이즈 실현에 중요도 점수를 부여하는 메타 학습 접근법을 도입하여 훈련 목적 함수의 적응형 재가중치를 가능하게 합니다.

2.1 매개변수 노이즈 평가기 (Parametric Noise Rater)

핵심 구성 요소는 "노이즈 평가기"로 작용하는 매개변수 함수 $\phi_\eta$입니다. 깨끗한 이미지 $x_0$, 조건 $c$(드롭될 수 있음), 시간 단계 $t$, 그리고 특정 노이즈 실현 $\epsilon$이 주어지면, 평가기는 해당 노이즈 인스턴스가 학습 신호에 가지는 상대적 중요도를 나타내는 스칼라 점수를 출력합니다.

평가기가 이미지 콘텐츠가 아닌 노이즈를 평가하도록 보장하기 위해, 본 방법은 그룹화된 노이즈 샘플링을 사용합니다:

• 미니배치의 각 데이터 샘플 $(x_0, c, t)$에 대해 $K$개의 독립적인 노이즈 샘플 $\{\epsilon^{(k)}\}_{k=1}^K$를 추출합니다.
• 이를 통해 $K$개의 노이즈가 추가된 입력 $\{x_t^{(k)}\}_{k=1}^K$를 생성합니다.
• 평가기 점수는 그룹 단위로 소프트맥스 함수를 사용하여 정규화됩니다:
  $$w^{(k)}(x_0) = \frac{\exp(\phi_\eta(\epsilon^{(k)}, t, x_0, \tilde{c}))}{\sum_{k'=1}^K \exp(\phi_\eta(\epsilon^{(k')}, t, x_0, \tilde{c}))}$$
  이는 고정된 조건 하에서 노이즈 인스턴스의 상대적 유용성을 반영하도록 가중치를 보장하며, 모델이 이미지 콘텐츠의 변이를 악용하는 것을 방지합니다.

2.2 이중 최적화 프레임워크

노이즈 평가기는 하류 일반화 성능을 직접 최적화하기 위해 이중 최적화 프레임워크를 통해 훈련됩니다:

• 내부 루프 (모델 훈련): 확산 모델 매개변수 $\theta$는 평가기의 점수 $w^{(k)}$로 재가중치된 가중 손실 $\mathcal{L}_{inner}$를 사용하여 업데이트됩니다. 여기서 표준 MSE 손실이 재가중치됩니다.
  $$\mathcal{L}_{inner}(\theta; \eta) = \mathbb{E}_{x_0 \sim \mathcal{B}} \left[ \sum_{k=1}^K w^{(k)}(x_0) \cdot \|\epsilon^{(k)} - \epsilon_\theta(x_t^{(k)}, t, \tilde{c})\|^2 \right]$$
• 외부 루프 (평가기 훈련): 평가기 매개변수 $\eta$는 내부 루프 업데이트 후 확산 모델의 검증 손실 $\mathcal{L}_{val}$을 최소화하도록 업데이트됩니다.
  $$\eta \leftarrow \arg\min_\eta \mathcal{L}_{val}(\theta^*(\eta))$$
  여기서 $\theta^*(\eta)$는 $S$개의 내부 루프 경사 단계 후의 모델 매개변수를 나타냅니다. 이 메타 학습 설정을 통해 평가기는 휴리스틱 기준에 의존하는 대신, 어떤 노이즈 인스턴스가 더 나은 일반화로 이어지는지 학습할 수 있습니다.

2.3 분리된 2 단계 파이프라인

효율적인 배포를 가능하게 하고 표준 훈련 중 여러 노이즈 후보를 유지하는 계산 오버헤드를 피하기 위해, 저자들은 2 단계 파이프라인을 제안합니다:

1. 메타 훈련 단계: 노이즈 평가기는 소프트 가중치(모든 $K$개의 후보 사용) 를 통한 이중 최적화를 사용하여 훈련됩니다.
2. 표준 훈련 단계: 훈련된 평가기는 고정됩니다. 이후 확산 훈련 중 각 데이터 샘플에 대해 $K'$개의 후보 노이즈가 샘플링되고, Top-1(최고 점수) 노이즈가 훈련 업데이트에 선택됩니다. 이 "하드 선택" 전략은 노이즈 평가를 훈련 과정과 분리하여 인스턴스 수준 가치 평가의 이점을 유지하면서도 표준 훈련 효율성을 보존합니다.

3. 주요 기여

본 논문은 네 가지 주요 기여를 제시합니다:

1. 문제 형식화: 기존 시간 단계 수준 노이즈 설계 및 테스트 시간 제어 방법과 보완적인 차원으로서 훈련 시간 노이즈 가치 평가를 식별하고 형식화합니다.
2. 메타 학습된 평가기: 일반화에 대한 기여도에 기반하여 노이즈 샘플을 적응형으로 가중치하는 메타 학습을 통해 훈련된 매개변수 노이즈 평가기를 제안합니다.
3. 효율적인 파이프라인: 소프트 가중치 (메타 훈련) 에서 하드 노이즈 선택 (표준 훈련) 으로 전환되는 분리된 2 단계 훈련 파이프라인을 도입하여 확장 가능한 배포를 가능하게 합니다.
4. 경험적 검증: 모든 노이즈 실현이 동등하게 유용하지 않으며, 유익한 노이즈를 선택적으로 강조하는 것이 훈련 효율성과 생성 품질 모두를 향상시킨다는 것을 입증합니다.

4. 실험 결과

본 방법은 Diffusion Transformer (DiT) 백본을 사용하여 FFHQ(256×256) 및 ImageNet(256×256) 에서 평가되었습니다.

• 성능 향상: 표준 훈련 (Vanilla) 과 단순 노이즈 선택 전략 (최소/최대 노름 노이즈 선택) 과 비교하여 NoiseRater 는 일관되게 더 낮은 FID-50K 점수를 달성했습니다. ImageNet(무조건부) 에서 이 방법은 +80k 스텝 시 FID 를 61.56 에서 59.95 로 감소시켰습니다 (상대적으로 2.6% 감소).

• 애블레이션 연구:

  • 입력 조건: 평가기는 최적의 성능을 내기 위해 모든 입력 (노이즈, 이미지 $x_0$, 시간 단계 $t$, 클래스 레이블) 이 필요합니다. 시간 단계 임베딩을 제거하면 가장 큰 성능 저하가 발생했는데, 이는 최적의 노이즈가 확산 단계마다 크게 달라진다는 것을 나타냅니다.
  • 후보 풀 크기 ($k$): $k=4$의 풀 크기가 품질 향상과 계산 비용 사이의 가장 좋은 균형을 제공했습니다. 더 큰 풀 ($k=8$) 은 체감 수익을 감소시키고 가우시안 사전 분포에서 편향을 증가시켰습니다.
  • 일반화: 더 작은 백본 (DiT-S/2) 에서 훈련된 평가기는 재훈련 없이 더 큰 백본 (DiT-B/2 및 DiT-L/2) 으로 성공적으로 이전되었으며, 이는 평가기가 노이즈 품질의 백본 무관한 속성을 학습함을 시사합니다.
  • 훈련 단계: 평가기는 훈련의 다양한 단계 (100k, 400k, 500k 스텝) 에 적용될 때 일관된 개선을 제공했으며, 가장 큰 이득은 약 400k 체크포인트 부근에 적용되었을 때 관찰되었습니다.

• 해석 가능성: 평가기 행동에 대한 분석은 다음과 같은 사실을 밝혔습니다:

  • 초기 훈련 후 노이즈 노름과의 상관관계가 약 0 이 되어 단순 노이즈 노름 필터로 수렴하지 않습니다.
  • 확산 모델의 신뢰도가 가장 낮고 학습 신호가 가장 중요한 중간 노이즈 수준 ($t \in [0.6, 0.7]$) 에서 가장 차별적입니다.
  • 그 차별 능력은 확산 모델 훈련 약 60k 스텝 후 안정화됩니다.

5. 중요성 및 주장

저자들은 NoiseRater 가 확산 모델 훈련 개선을 위한 이전에 충분히 탐구되지 않았던 축으로서 노이즈 가치 평가를 확립한다고 주장합니다. 개별 노이즈 샘플의 중요성을 명시적으로 모델링함으로써, 이 방법은 노이즈가 균일하게 유익하다는 가정을 넘어섭니다.

이 연구의 중요성은 다음과 같습니다:

• 상호 보완성: 기존 시간 단계 수준 스케줄링 및 테스트 시간 제어 방법과 보완되는 새로운 최적화 방향을 제공합니다.
• 효율성: 확산 모델의 아키텍처 변경이나 추론 절차 수정 없이 훈련 효율성과 생성 품질을 향상시킵니다.
• 메타 학습 적용: 주입된 노이즈를 가치 평가하는 특정 문제에 이중 최적화를 성공적으로 적용하여, 가치 평가 객체를 데이터 포인트에서 노이즈 실현으로 전환합니다.

본 논문은 유익한 노이즈를 우선시하는 것이 확산 모델을 강화하기 위한 실현 가능한 전략이며, 제안된 프레임워크가 이를 달성하기 위한 유연하고 효과적인 메커니즘을 제공한다고 결론지었습니다.