Research map / Reinforcement learning / RL from feedback
Confronting Reward Overoptimization for Diffusion Models: A Perspective of Inductive and Primacy Biases
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes TDPO, a per-timestep RL framework with a temporal critic for aligning diffusion models, and TDPO-R, which adds periodic resets of critic active neurons to mitigate reward overoptimization.
Problem
Reward overoptimization in reward-driven diffusion model alignment, where finetuning on learned rewards degrades ground-truth quality and cross-reward generalization, with a trade-off against sample efficiency.
Approach
Assign timestep-dependent rewards to intermediate denoising states via a temporal critic that reuses the reward model's encoder and predicts a residual from the final reward; update policy and critic per timestep to match temporal inductive bias; and periodically reset critic active neurons to counter primacy bias, since dormant neurons act as implicit regularization against overoptimization.
Key results
- TDPO and TDPO-R improve sample efficiency over DDPO-2 and DDPO-100 on Aesthetic Score, PickScore, and HPSv2 according to reward-vs-queries curves.
- In cross-reward generalization to ImageReward when finetuned on Aesthetic Score, TDPO and TDPO-R maintain higher out-of-domain scores than DDPO and AlignProp as training reward rises.
- Resetting critic active neurons mitigates overoptimization, whereas resetting dormant neurons or all neurons exacerbates it; KL regularization is reported as less effective than the neuron reset.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use timestep-dependent temporal rewards and per-timestep gradient updates instead of final-image rewards with per-batch updates | TDPO · Aesthetic Score, PickScore, and HPSv2 | sample efficiency | Helped |
| use timestep-dependent temporal rewards and per-timestep gradient updates instead of final-image rewards with per-batch updates | TDPO · Cross-reward evaluation, including unseen prompts | cross-reward generalization | Helped |
| periodically reset active neurons in the critic (ablation) | TDPO · Aesthetic Score to ImageReward cross-reward evaluation | cross-reward generalization | Helped |
| periodically reset dormant neurons in the critic (ablation) | TDPO · Aesthetic Score to ImageReward cross-reward evaluation | cross-reward generalization | Hurt |
| periodically reset all neurons in the critic (ablation) | TDPO · Aesthetic Score to ImageReward cross-reward evaluation | cross-reward generalization | Hurt |
| periodically reset dormant neurons in the policy LoRA layers (ablation) | TDPO · Aesthetic Score cross-reward evaluation | cross-reward generalization | No clear effect |
| periodically reset active neurons in the policy LoRA layers (ablation) | TDPO · Aesthetic Score | sample efficiency | Hurt |
| add KL regularization instead of using critic active-neuron reset (ablation) | TDPO · Aesthetic Score cross-reward evaluation | cross-reward generalization | Hurt |
| replace the HPSv2 temporal critic encoder with an Aesthetic Score encoder (ablation) | TDPO · HPSv2 reward optimization | reward optimization performance | Hurt |
| replace standard ReLU with leaky ReLU (ablation) | TDPO-R · Aesthetic Score to PickScore cross-reward evaluation | cross-reward generalization | Helped |
| increase DDPO update frequency from 2 to 100 gradient updates per epoch | DDPO · Diffusion model reward optimization | sample efficiency | Helped |
| increase DDPO update frequency from 2 to 100 gradient updates per epoch | DDPO · Cross-reward generalization | cross-reward generalization | Hurt |
Surprising findings
- Resetting dormant neurons in the critic exacerbates reward overoptimization, contrary to prior deep RL findings that dormant neurons should be reset.
- Dormant neurons in the critic persist across training (100% overlap between iterations) and act as an implicit regularizer against overoptimization.
Limitations
- Policy neuron resets with LoRA cause catastrophic forgetting and hinder learning, so neuron reset is only applied to the critic.
- Experiments use LoRA finetuning rather than full model finetuning, so internal neuron analysis of the diffusion model is not comprehensive.
- Multi-reward learning is unexplored.
Experimental setup
Datasets: HPSv2 prompts (Human Preference Dataset v2) animal name prompts
Benchmarks: ImageReward PickScore HPSv2 Aesthetic Score
Metrics: average reward vs reward queries (sample efficiency) cross-reward generalization score
Baselines: Stable Diffusion v1.4 DDPO-2 DDPO-100 AlignProp KL regularization
Builds on: DDPO AlignProp DDPM LoRA primacy bias
Closest papers by meaning
- Reward Score Matching: Unifying Reward-based Fine-tuning for Flow and Diffusion Models (NeurIPS 2026 · similarity 0.89)
- Directly Fine-Tuning Diffusion Models on Differentiable Rewards (ICLR 2024 · similarity 0.89)
- Test-time Alignment of Diffusion Models without Reward Over-optimization (ICLR 2025 · similarity 0.88)
- Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control (ICLR 2025 · similarity 0.87)
- Using Human Feedback to Fine-tune Diffusion Models without Any Reward Model (CVPR 2024 · similarity 0.87)
- Training Diffusion Models with Reinforcement Learning (ICLR 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.