Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Step-level Reward for Free in RL-based T2I Diffusion Model Fine-tuning

Xinyao Liao, Wei Wei, Xiaoye Qu, Yu Cheng

TL;DR

This work tackles reward sparsity in RL-based fine-tuning of text-to-image diffusion models by introducing Contribution-based Credit Assignment (CoCA). CoCA estimates each denoising step's contribution to the final image using cosine similarity between intermediate latent representations and the final latent, then redistributes the terminal reward across steps with fixed-window smoothing and two-stage normalization to produce dense, informative step-level rewards. The authors prove that this reward shaping is potential-based and invariant with respect to the optimal policy, ensuring theoretical consistency with the original objective. Empirically, CoCA achieves 1.25x–2x improvements in sample efficiency across four human-preference reward functions and generalizes better to unseen rewards and prompts, without adding auxiliary networks. These results suggest that adaptive, contribution-aware credit assignment can substantially enhance RL-guided T2I diffusion fine-tuning and facilitate more efficient, controllable image synthesis.

Abstract

Recent advances in text-to-image (T2I) diffusion model fine-tuning leverage reinforcement learning (RL) to align generated images with learnable reward functions. The existing approaches reformulate denoising as a Markov decision process for RL-driven optimization. However, they suffer from reward sparsity, receiving only a single delayed reward per generated trajectory. This flaw hinders precise step-level attribution of denoising actions, undermines training efficiency. To address this, we propose a simple yet effective credit assignment framework that dynamically distributes dense rewards across denoising steps. Specifically, we track changes in cosine similarity between intermediate and final images to quantify each step's contribution on progressively reducing the distance to the final image. Our approach avoids additional auxiliary neural networks for step-level preference modeling and instead uses reward shaping to highlight denoising phases that have a greater impact on image quality. Our method achieves 1.25 to 2 times higher sample efficiency and better generalization across four human preference reward functions, without compromising the original optimal policy.

Step-level Reward for Free in RL-based T2I Diffusion Model Fine-tuning

TL;DR

This work tackles reward sparsity in RL-based fine-tuning of text-to-image diffusion models by introducing Contribution-based Credit Assignment (CoCA). CoCA estimates each denoising step's contribution to the final image using cosine similarity between intermediate latent representations and the final latent, then redistributes the terminal reward across steps with fixed-window smoothing and two-stage normalization to produce dense, informative step-level rewards. The authors prove that this reward shaping is potential-based and invariant with respect to the optimal policy, ensuring theoretical consistency with the original objective. Empirically, CoCA achieves 1.25x–2x improvements in sample efficiency across four human-preference reward functions and generalizes better to unseen rewards and prompts, without adding auxiliary networks. These results suggest that adaptive, contribution-aware credit assignment can substantially enhance RL-guided T2I diffusion fine-tuning and facilitate more efficient, controllable image synthesis.

Abstract

Recent advances in text-to-image (T2I) diffusion model fine-tuning leverage reinforcement learning (RL) to align generated images with learnable reward functions. The existing approaches reformulate denoising as a Markov decision process for RL-driven optimization. However, they suffer from reward sparsity, receiving only a single delayed reward per generated trajectory. This flaw hinders precise step-level attribution of denoising actions, undermines training efficiency. To address this, we propose a simple yet effective credit assignment framework that dynamically distributes dense rewards across denoising steps. Specifically, we track changes in cosine similarity between intermediate and final images to quantify each step's contribution on progressively reducing the distance to the final image. Our approach avoids additional auxiliary neural networks for step-level preference modeling and instead uses reward shaping to highlight denoising phases that have a greater impact on image quality. Our method achieves 1.25 to 2 times higher sample efficiency and better generalization across four human preference reward functions, without compromising the original optimal policy.

Paper Structure

This paper contains 47 sections, 3 theorems, 18 equations, 12 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Reward Fine-tuned T2I Diffusion Models
  4. Dense Reward for RL Fine-tuned Models
  5. Prelimilaries
  6. Diffusion Models
  7. Reinforcement Learning
  8. Method
  9. Trajectory-level Reward
  10. Denoising as a MDP with the Trajectory-level Reward
  11. Policy Gradient of the Trajectory-level Reward
  12. Contribution-based Credit Assignment
  13. Timestep Contribution Estimation
  14. Step-level Reward Redistribution
  15. Two-Stage Reward Normalization
  16. ...and 32 more sections

Key Result

Lemma 1

(Follwing Lemma 4.1 in dpok2023) The policy gradient of $\nabla_\theta \mathcal{J}_{\text{TR}}(\pi_\theta)$ is:

Figures (12)

  • Figure 1: (I) By illustrating the evolution of rewards (HPSv2) and reward gains ($\Delta$ HPSv2) over timesteps, we observe that reward gains diminish as timesteps progress, whereas trajectory-level sparse rewards assign equal importance to all timesteps. (II) Our step-level credit assignment provides dense rewards based on the actual contribution of each timestep by weighting the trajectory-level reward, enabling faster emergence of coherent global structures.
  • Figure 2: Learning curves by sample efficiency. Reward functions (From left to right: (a) Aesthetic Score, (b) ImageReward Score, (c) PickScore, (d) HPSv2 Score) are evaluated to compare DDPO, TDPO, UCA, and our method.
  • Figure 4: Cross-reward generalization results of methods trained on PickScore.
  • Figure 6: Qualitative comparison of samples of seen and unseen prompts in the HPSv2 dateset generated by DDPO, TDPO, UCA, and our methods trained on HPSv2 and evaluated on.
  • Figure 7: Qualitative comparison of unseen animals generated by SD-v1.5, DDPO, TDPO, UCA, CoCA trained on ImageReward.
  • ...and 7 more figures

Theorems & Definitions (6)

  • Lemma 1
  • Lemma 2
  • proof
  • Lemma 3
  • proof
  • proof