Confronting Reward Overoptimization for Diffusion Models: A Perspective of Inductive and Primacy Biases

TL;DR

This work tackles reward overoptimization in diffusion-model alignment by framing it through temporal inductive bias and primacy bias. It introduces Temporal Diffusion Policy Optimization (TDPO) to exploit the multi-step denoising timeline via timestepped rewards and a temporal critic, achieving improved sample efficiency. Building on TDPO, it proposes TDPO-R, which periodically resets active neurons in the critic to counter primacy bias, yielding superior cross-reward generalization and robustness to out-of-domain rewards. The results demonstrate that attention to temporal structure and neuron-state dynamics can substantially mitigate overoptimization, advancing practical diffusion-model alignment for human-preference-driven generation, with code available at the authors’ GitHub repository.

Abstract

Bridging the gap between diffusion models and human preferences is crucial for their integration into practical generative workflows. While optimizing downstream reward models has emerged as a promising alignment strategy, concerns arise regarding the risk of excessive optimization with learned reward models, which potentially compromises ground-truth performance. In this work, we confront the reward overoptimization problem in diffusion model alignment through the lenses of both inductive and primacy biases. We first identify a mismatch between current methods and the temporal inductive bias inherent in the multi-step denoising process of diffusion models, as a potential source of reward overoptimization. Then, we surprisingly discover that dormant neurons in our critic model act as a regularization against reward overoptimization while active neurons reflect primacy bias. Motivated by these observations, we propose Temporal Diffusion Policy Optimization with critic active neuron Reset (TDPO-R), a policy gradient algorithm that exploits the temporal inductive bias of diffusion models and mitigates the primacy bias stemming from active neurons. Empirical results demonstrate the superior efficacy of our methods in mitigating reward overoptimization. Code is avaliable at https://github.com/ZiyiZhang27/tdpo.

Figures (14)

• Figure 1: TDPO-R first samples trajectories $(x_T, x_{T-1}, ... , x_0)$ from the denoising process of a fixed diffusion model parameterized by $\theta_\mathrm{old}$ for each epoch. At each timestep $t$, it performs a one-step denoising using the current diffusion model parameterized by $\theta$, estimates a temporal reward $\mathcal{T}_{\phi}(x_t)$ using a temporal critic parameterized by $\phi$, and updates the gradients for both $\theta$ and $\phi$ according to the corresponding objective functions. Additionally, TDPO-R resets active neurons of $\phi$ at the end of every $F$ epochs.
• Figure 2: Image generation results sampled from models that are either pre-trained or further finetuned on Aesthetic Score via AlignProp, DDPO-2, DDPO-100, as well as our TDPO and TDPO-R. For a fair comparison, all images are generated using a fixed random seed of 42. Additionally, for the fine-tuned models, the aesthetic scores of the generated images achieve similar values around 7 $\pm$ 0.1.
• Figure 3: Out-of-domain evaluation results via cross-reward generalization against ImageReward (left), PickScore (middle), and HPSv2 (right) when finetuning the diffusion model on Aesthetic Score (left), HPSv2 (middle), and PickScore (right), respectively.
• Figure 4: Quantitative evaluation results for the efficacy of our TDPO and TDPO-R in improving sample efficiency when finetuning the diffusion model on the reward functions of PickScore (left), HPSv2 (middle), and Aesthetic Score (right), compared to DDPO with the update frequencies of 2 (DDPO-2) and 100 (DDPO-100) per epoch.
• Figure 5: Cross-reward generalization results evaluated on a text prompt set of unseen animals when finetuning the diffusion model on Aesthetic Score.