DiffusionOPD: Diffusion 모델에서 On-Policy Distillation의 통합적 관점

다중 목표를 동시에 만족해야 하는 실제 사용에서, 단일 보상 최적화 RL은 태스크 간 간섭으로 해결이 어렵다. DiffusionOPD는 단일 태스크 탐색과 다중 태스크 능력 통합을 분리하고, 각 태스크별 교사를 통해 학생으로의 지식을 온폴리시 롤아웃에 맞춰 증류한다. 이를 통해 교사 훈련의 간섭을 피하고, 샘플링 다이어그램의 차원에서 KL 기반 지도 학습의 분산을 줄여 다중 도메인에서의 최종 성능과 학습 효율을 동시에 향상시킨다.

출발점: 다중 태스크 RL은 서로 다른 보상 신호 간의 충돌과 학습 난이도 불균형 등의 문제로 인해 학습이 비효율적이고 성능이 특정 태스크에 편향될 수 있다. 기존의 cascade RL은 여러 단계 학습으로 인한 비용과 망각 문제를 야기한다. 해결 원리: DiffusionOPD는 학습을 두 단계로 나누어, 1) 각 태스크에 대해 독립적 교사를 학습하고 2) 학생이 자신의 롤아웃 경로를 따라 교사들의 지식을 증류하도록 한다. diffusion 도메인에 맞춰 OPD를 확장하고, per-step KL을 닫힌 형태로 유도하여 PPO 스타일의 대조적 경사보다 낮은 분산을 달성한다. 달라지는 점: SDE와 ODE 샘플링 모두에 대해 동일한 공분산을 공유하는 Gaussian 전이에서 KL이 유도되는 점을 활용해, 확률적/결정적 샘플링 간 일관된 지도 학습이 가능해지며, 학습 효율과 성능의 상향이 동시에 가능하다.

단계별 구성: (1) 두 단계 학습 파이프라인의 제시, (2) DiffusionOPD의 수학적 프레임워크 도출, (3) 학습 전략 및 구현 세부사항. (1) Stage 1: M개 태스크 각각에 대해 task-specific teacher v(m)ϕm를 off-the-shelf diffusion RL 알고리즘으로 학습한다. (2) Stage 2: Stage 1에서 학습된 교사들을 이용해 unified student vθ를 온폴리시 롤아웃에 따라 증류한다. 각 태스크 c에서 학생의 롤아웃 {xtj}에 대해 교사 v(m)ϕm를 이용해 per-step KL을 추정하고 Eq. (11) 또는 Eq. (12)로 손실을 계산한다. (3) 디테일: 알고리즘 1에 기반한 라운드 로빈 식으로 학습하며, 각 라운드에서 모든 태스크를 순회한 뒤 한 번의 역전파로 업데이트한다. 학습 손실은 G=M의 누적 증가로 안정화한다. 수학적 기반: pS(·|xtj)와 pT(·|xtj) 사이의 KL은 동일 공분산 Σ = σ¯²j Id를 공유하는 두 Gaussians 간의 차이의 제곱으로 표현되며, KL = ∥µS(xtj; θ) − µT(xtj)∥² / (2σ¯²j)로 계산된다. Deterministic 샘플러의 경우 µS(xtj; θ)와 µT(xtj)의 차이를 직접 제곱해 L2 손실로 활용한다. PPO와의 관계: PPO 스타일의 경사는 두 가지 구성요소(경로적 항, 스코어-함수 항)로 나뉘며, πθold = πθ일 때 경로적 항은 0에 수렴해도 KL 기반 목표의 기대 gradient는 동일하게 유지된다. 분산 측면에서 닫힌 KL은 샘플 노이즈에 의존하지 않는 pathwise gradient를 제공한다. 구현 세부: Algorithm 1의 두 단계 구조, 라운드-로빈 업데이트, 노이즈 수준 a와 스케줄 tj, ∆tj에 따른 샘플러 설정; Stage 1의 교사 학습은 각 태스크의 특성에 맞는 RL 알고리즘을 사용한다.

주요 벤치마크에서의 성능: DiffusionOPD는 Multi-Task GRPO-Guard 및 Cascade NFT를 상회하며, Average 점수 0.929로 최상위를 차지한다(GenEval 0.96, OCR 0.94, PickScore 23.99, ClipScore 0.297, HPSv2.1 0.342, Aesthetic 6.15, ImgRwd 1.50, UniRwd 3.50). 훈련 시간은 85.75+11.26 GPU 시간으로 보고되며, Multi-Task RL-baseline 대비 학습 효율이 크게 개선된다. Figure 1(a)에서 DiffusionOPD의 수렴 속도는 다른 다중태스크 RL 방법들보다 빠르고 성능 상한이 더 높다. Figure 1(b)에서 GenEval, OCR, aesthetics 등 다양한 도메인에서의 일관된 우수 성능이 확인된다. Ablation 연구에서 Distillation 방법(DMD, TDM, SFT) 비교 시 DiffusionOPD가 가장 빠르게 수렴하고 최종 성능 상한도 가장 높다. Loss formulation 비교에서 closed-form KL이 PPO 기반 대비 더 빠르게 성능을 올리고, 샘플링 노이즈 감소 시 수렴 속도와 성능이 증가한다. 노이즈 레벨 0에 가까울수록 실험에서 더 빠른 수렴과 더 높은 점수를 얻는다.