[논문 리뷰] Diffusion Reinforcement Learning via Centered Reward Distillation

Yuanzhi Zhu, Xi Wang외 2명|arXiv (Cornell University)|2026. 03. 14.

Reinforcement Learning in Robotics인용 수 0

한 줄 요약

CRD는 diffusion RL을 중심화된 보상 매칭 문제로 공식화하고, 알려지지 않은 프롬프트 의존 정규화자를 취소하며, 외부 보상으로 더 경쟁력 있고, 더 빠르며, 더 강건한 텍스트-이미지 미세 조정을 달성하기 위한 안정화 기술을 도입한다.

ABSTRACT

Diffusion and flow models achieve State-Of-The-Art (SOTA) generative performance, yet many practically important behaviors such as fine-grained prompt fidelity, compositional correctness, and text rendering are weakly specified by score or flow matching pretraining objectives. Reinforcement Learning (RL) fine-tuning with external, black-box rewards is a natural remedy, but diffusion RL is often brittle. Trajectory-based methods incur high memory cost and high-variance gradient estimates; forward-process approaches converge faster but can suffer from distribution drift, and hence reward hacking. In this work, we present extbf{Centered Reward Distillation (CRD)}, a diffusion RL framework derived from KL-regularized reward maximization built on forward-process-based fine-tuning. The key insight is that the intractable normalizing constant cancels under \emph{within-prompt centering}, yielding a well-posed reward-matching objective. To enable reliable text-to-image fine-tuning, we introduce techniques that explicitly control distribution drift: ( extit{i}) decoupling the sampler from the moving reference to prevent ratio-signal collapse, ( extit{ii}) KL anchoring to a CFG-guided pretrained model to control long-run drift and align with the inference-time semantics of the pre-trained model, and ( extit{iii}) reward-adaptive KL strength to accelerate early learning under large KL regularization while reducing late-stage exploitation of reward-model loopholes. Experiments on text-to-image post-training with exttt{GenEval} and exttt{OCR} rewards show that CRD achieves competitive SOTA reward optimization results with fast convergence and reduced reward hacking, as validated on unseen preference metrics.

연구 동기 및 목표

• 정적 사전 학습 목표를 넘어 외부 보상에 맞춰 확산 모델의 사후 학습 정렬을 촉진한다.
• 역경우 forward-process 확산 RL의 견고한 목표를 개발하되, 계산 불가능한(normalizer)에도 불구하고 잘 정의되어 있도록 유지한다.
• RL 미세 조정 중 보상 남용과 분포 드리프트를 감소시킨다.
• 비분화가능한 보상을 갖는 안정적이고 효율적인 확산 RL을 가능하게 하는 실용적 안정화 기술을 제시한다.

제안 방법

• Center Reward Distillation(CRD)을 KL-정규화 보상 최대화 하에 중심화된 보상 매칭 목표로 도출한다.
• 현재 모델과 움직이는 기준 간의 암묵적 로그-밀도 비를 추정하기 위해 forward diffusion ELBO 대체를 사용한다.
• 각 프롬프트 그룹 내에서 보상을 중앙화하여 프롬프트-의존 정규화자(beta log Z(c))를 제거한다.
• 외부 보상과 온도 매개변수 tau에 의존하는 보상 가중 중심 가중치를 정의한다.
• 비율-신호 붕괴를 방지하기 위해 샘플링을 움직이는 기준으로부터 분리한다.
• CFG-가이드 사전 학습 모델에 KL 고정을 적용하여 장기 드리프트를 제어하고 추론 시 의미론과 정렬한다.
• 초기 학습을 가속화하면서 후기 보상 허점들을 완화하기 위해 보상-적응적 KL 강도를 도입한다.

실험 결과

연구 질문

• RQ1 계산 불가능한 정규화를 모를 때도 잘 정의된 확산 RL 목표를 어떻게 형식화할 수 있는가?
• RQ2 프롬프트 내 중앙화가 프롬프트 의존 정규화자를 제거하고 보상 매칭을 안정화할 수 있는가?
• RQ3 분포 드리프트와 보상 남용에 대해 어떤 실용적 기법이 forward-process 확산 RL을 가장 잘 안정시키는가?
• RQ4 CRD가 Flow-GRPO 스타일 및 forward-process 기반의 기준선과 비교하여 더 빠른 수렴과 감소된 보상 남용으로 경쟁력 있는 보상 최적화를 달성하는가?
• RQ5 CFG 기반 KL 고정이 학습 안정성과 추론 시 의미론과의 정렬에 어떤 영향을 미치는가?

주요 결과

• CRD는 프롬프트 그룹 내에서 보상을 중앙화하여 계산 불가능한 정규화자를 제거함으로써 잘 정의된 보상 매칭 목표를 얻는다.
• CRD는 기준선에 비해 더 빠른 수렴과 향상된 학습 안정성으로 경쟁력 있는 보상 최적화를 달성한다.
• 실용적 안정화 기술—분리된 샘플링, CFG-가이드 사전 학습 모델에 대한 KL 고정, 보상 적응적 KL 강도—은 보상 남용과 분포 드리프트를 완화한다.
• GenEval 및 OCR 보상으로 텍스트-이미지 미세조정 실험은 우호적 품질적·정량적 행동과 감소된 보상 남용을 보여준다.
• CRD는 이전 보상 증류 및 GVPO 스타일 목표를 특수한 경우로 포괄하며, InfoNCA와 관련된 비율 기반 변형을 허용한다.