기술 요약: 보상 모멘트 매칭 증류를 이용한 확산 미세 조정 (Diffusion Fine-tuning with Rewarded Moment Matching Distillation)

문제 정의

확산 모델(Diffusion models)은 안정적인 학습 목적 함수 덕분에 고충실도 이미지 합성 및 과학적 예측 분야의 표준이 되었습니다. 그러나 실제 배포 시에는 수십에서 수백 단계의 디노이징(denoising) 단계가 필요하여 추론 비용이 높다는 문제가 있습니다. 증류(Distillation) 기술은 이러한 모델을 몇 단계의 생성기로 압축하며, 강화 학습(RL) 미세 조정은 이를 특정 보상 함수(예: 인간의 선호도 또는 과학적 지표)에 정렬시킵니다. 하지만 이 두 단계를 결합하는 것은 여전히 까로운 과제입니다.

증류와 보상 최적화를 병합하려는 기존 방식들은 다음과 같은 중대한 한계에 직면해 있습니다:

1. 결합 학습의 취약성: 증류와 보상 미세 조정을 하나의 단계로 병합하면, 생성된 샘플이 교사 모델(teacher model)의 입력 분포에서 벗어나 증류 신호를 무효화하는 현상이 발생합니다.
2. 메모리 및 편향 문제: 다단계 역전파(예: DRaFT)를 통해 이미 증류된 모델을 미세 조정하는 것은 메모리 집약적입니다. 역전파를 중단하는 방식(예: ReFL)은 노이즈가 섞인 잠재 변수(latents)에서 보상을 평가함으로써 편향을 유발합니다.
3. 구조적 제약: HyperNoise와 같은 방법은 초기 노이즈를 섭동(perturbing)함으로써 체인 역전파를 피하지만, 저주파 이미지 변화에 국한되어 복잡한 고주파 구조적 보상을 최적화하는 데 한계가 있습니다.

방법론: 보상 모멘트 매칭 증류 (RMMD)

저자들은 증류와 보상 미세 조정을 원칙적으로 연결하면서, 고급 증류의 고충실도 "자연스러움(naturalness)"을 보존하는 두 단계 프레임워크인 RMMD를 제안합니다.

1단계: 모멘트 매칭 증류 (MMD)

과정은 **모멘트 매칭 증류(Moment-Matching Distillation, MMD)**를 사용하여 기본 교사 모델을 증류하는 것으로 시작됩니다. 이 방법은 샘플링 궤적을 따라 중간 디노이징 분포를 매칭하여, 교사의 주변 분포(marginals)를 밀접하게 추적하는 다단계 학생 모델(student model)을 생성합니다. 이 학생 모델은 이후 안정적인 분포 참조 역할을 하기 위해 동결됩니다.

2단계: 온-폴리시(On-Policy) 보상 미세 조정

두 번째 단계에서는 학생 모델이 분포적 충실도를 유지하면서 보상 함수를 최대화하도록 미세 조정됩니다. 주요 혁신 사항은 다음과 같습니다:

1. 온-폴리시 샘플링: 오프-폴리시(off-policy) 데이터 포인트를 사용하는 대신, 현재 학생의 정책으로부터 노이즈가 섞인 상태 $x^{pol}_t$를 샘플링합니다(학생을 $K_{student}$ 단계 동안 실행하고, 출력을 재노이즈화한 뒤 그래디언트를 중단하는 방식). 이는 보상 목적 함수가 학생의 기대 보상을 직접 근사하도록 보장하며, 분포 드리프트가 학습 신호를 무효화하는 것을 방지합니다.
2. 재활용된 정규화(Recycled Regularization): 핵심 기여는 모멘트 매칭 손실을 미세 조정 중에 정규화 도구로 재사용하는 것입니다. 이 손실은 단순한 휴리스틱 페널티가 아니라, 학생의 중간 디노이징 분포와 동결된 참조 모델 간의 불일치를 명시적으로 처벌합니다. 이를 통해 보상 최대화와 분포적 충실도 사이의 절충을 조절할 수 있는 해석 가능한 노브(weight)를 제공합니다.
3. 단일 단계 그래디언트: 다단계 언롤링(unrolling) 없이도 전체 주파수 범위를 포착하기 위해, 오염된 온-폴리시 샘플에 대해 단일 디노이징 단계를 수행하고 보상을 평가합니다.
4. L2 정규화: 초기 MMD 증류 모델의 예측으로부터의 유의미한 이탈을 방지하기 위해 추가적인 L2 항을 적용합니다.

전체 목적 함수는 온-폴리시 보상 그래디언트와 모멘트 매칭 정규화 및 L2 항을 결합합니다:
$$\nabla_\theta \mathcal{L}_{RMMD} = \mathbb{E} \left[ \left(\frac{\partial \hat{x}_0}{\partial \theta}\right)^\top \left( m_{student} - m_{teacher} - \lambda \nabla_{\hat{x}_0} R(\hat{x}_0) \right) \right] + \lambda_{reg} \mathcal{L}_{L2}$$

주요 기여

1. RMMD 절차: 정규화를 위해 온-폴리시 버전의 모멘트 매칭 손실을 활용하여, 보상 함수를 동시에 최대화하는 새로운 확산 증류 방법을 제시했습니다.
2. 경험적 우월성: 다양한 보상 함수(픽셀 수준 메트릭부터 의미적 정렬인 CLIP까지)에 대해 RMMD가 단일 단계 베이스라인(DI++) 및 다단계 경쟁 모델(DRaFT, HyperNoise)보다 FID-Reward 파레토 프런트(Pareto front)에서 일관되게 우수한 성능을 보임을 입증했습니다.
3. 과학적 응용: RMMD를 최첨단 확산 기반 기상 예측 모델인 GenCast에 성공적으로 적용하여 연속 순위 확률 손실(CRPS)을 최적화했습니다.

결과

ImageNet 실험

• 다단계의 이점: 8단계 MMD 증류된 학생을 시작점으로 사용했을 때 단일 단계 베이스라인(DI++)보다 성능이 크게 향상되었습니다. 8단계 RMMD는 우수한 FID-Reward 파레토 프런트 트레이드오프를 달성했습니다.
• 비교: RMMD는 다양한 해상도(64x64 및 512x512)와 보상 유형(CLIP 정렬, Inception Score, 매끄러움)에 대해 DRaFT 및 HyperNoise보다 더 나은 파레토 프런트를 생성했습니다.
• 질적 분석: 시각화 결과, DRaFT-1은 콘텐츠는 유지하지만 적대적 아티팩트(adversarial artifacts)를 유발했고, DRaFT-2는 보상 해킹(reward hacking, 급격한 분포 변화) 문제를 겪은 반면, RMMD는 원래의 데이터 분포에서 크게 벗어나지 않으면서 보상을 높이기 위한 미세한 수정을 성공적으로 통합했습니다.

GenCast 기상 예측

• 속도 향 향상: 증류된 모델은 교사 모델에 비해 7.5배의 속도 향상(8단계 vs 59회 함수 평가)을 달성했습니다.
• 정확도: CRPS를 최적화한 RMMD 모델은 타겟 기상 변수의 **93%**에서 교사 모델보다 우수한 성능을 보였습니다(최고의 MMD 전용 베이스라인은 75%).
• 보정(Calibration): RMMD는 MMD 증류에서 흔히 발생하는 과소 분산(under-dispersion) 문제를 해결했습니다. 온-폴리시 변형이 오프-폴리시 변형보다 약간 낮은 분산 지표를 보였으나, 전반적인 CRPS 개선 측면에서는 가장 우수했습니다.
• 일반화: 12시간 선행 예측에서의 CRPS 최적화는 최대 7일까지의 예측에서 일관된 개선을 보여주었으며, 이는 이 방법이 단순히 즉각적인 보상에 과적합되는 것이 아니라 근본적인 분포 모델링을 개선함을 시사합니다.

의의 및 주장

본 논문은 RMMD가 결합 증류 및 미세 조정 과정에서 발생하는 "드리프트(drift)" 문제에 대한 원칙적인 해결책을 제공한다고 주장합니다. 모멘트 매칭 손실을 정규화제로 재활용하고 온-폴리시 샘플링을 활용함으로써, RMMD는 생성 품질이나 증류된 모델의 안정성을 희생하지 않으면서도 복잡하고 미분 가능한 목적 함수(기상 예측의 CRPS와 같은)를 최적화할 수 있게 합니다.

저자들은 이 접근 방식이 복잡하고 고차원적인 과학 영역으로 확장 가능하다는 점을 강조합니다. GenCast에 대한 성공적인 적용은 확산 모델이 확률적 예측을 유의미하게 가속화하는 동시에 개선할 수 있음을 보여주며, 미분 가능한 목적 함수가 존재하는 과학 계산 분야에서 더 효율적이고 정확한 확산 기반 모델을 위한 길을 열어줍니다. 이 연구는 모멘트 매칭을 통해 안정적인 분포 참조를 유지하는 것이 확산 모델의 효과적인 보상 미세 조정에 필수적임을 시사합니다.