Technical Summary: CoCA – Step-Level Reward for Free in RL-Finetuned T2I Diffusion Models

1. Problem Statement

State-of-the-art text-to-image (T2I) diffusion models increasingly rely on Reinforcement Learning (RL) for fine-tuning to align with human preferences (e.g., aesthetic appeal, semantic correctness). However, these methods face a critical limitation: reward sparsity.

In standard RL formulations for diffusion, the process is modeled as a Markov Decision Process (MDP) where the agent receives a single terminal reward only after the final denoising step ($t=0$). Existing policy-gradient methods (e.g., REINFORCE, DDPO) distribute this scalar reward uniformly across all $T$ denoising steps. This creates a mismatch:

• Unequal Impact: Early timesteps are decisive for establishing the global structure of the image, while later timesteps primarily refine fine-grained details.
• Inefficient Credit Assignment: Uniform weighting fails to identify which specific steps contributed meaningfully to the final quality, leading to suboptimal sample efficiency and limited precision in attributing rewards to generative actions.
• Limitations of Existing Solutions: Previous attempts to address this include fixed temporal discounting (which relies on a priori assumptions that fail to generalize across diverse prompts) or training auxiliary networks (e.g., step-wise critics) to predict intermediate rewards, which introduces significant computational overhead and requires additional training.

2. Methodology: CoCA (Contribution-based Credit Assignment)

The authors propose CoCA, a framework that redistributes the terminal reward across the denoising trajectory to generate dense, step-level signals without training auxiliary networks.

Core Mechanism

CoCA estimates the contribution of each denoising step by measuring the latent progression toward the realized clean image.

1. Latent Progress Proxy: Instead of assuming fixed importance, CoCA tracks the similarity between the intermediate latent state ($x_{T-t}$) and the realized terminal latent ($x_0$) in the embedding space.
   • It computes the cosine similarity: $Sim_t = \frac{\langle x_{T-t}, x_0 \rangle}{\|x_{T-t}\| \cdot \|x_0\|}$.
   • The contribution of a step is derived from the change in this similarity ($\Delta Sim_t$).
2. Window Smoothing: To reduce fluctuations in per-step similarity, the trajectory is segmented into non-overlapping windows of fixed size $W$. The average similarity within a window is computed, and the progress score is the difference between consecutive windows ($\Delta \bar{Sim}_i$).
3. Reward Redistribution: The terminal reward $r(x_0, c)$ is redistributed to each step $t$ based on the normalized progress scores:
   $$\hat{R}(s_t, a_t) = w_t \cdot r(x_0, c)$$
   where $w_t$ is proportional to the estimated contribution of that step.

Stabilization and Theoretical Guarantees

• Reward Normalization: To stabilize training, CoCA applies normalization twice:
  1. Per-prompt normalization before redistribution to maintain ranking consistency across prompts.
  2. Per-prompt per-timestep normalization after redistribution to ensure consistent reward scales across timesteps and reduce variance.
• Policy-Gradient Consistency: The authors prove that if the redistributed step rewards sum to the original terminal reward for every trajectory ($\sum \hat{R}_t = r(x_0, c)$), the expected policy gradient remains unbiased. Thus, CoCA preserves the stationary points of the original objective while providing denser learning signals.

3. Key Contributions

1. Novel Framework: Introduction of CoCA, a contribution-based credit assignment mechanism that addresses the mismatch between equal policy updates and varying step importance in RL-finetuned diffusion models.
2. No Auxiliary Training: Unlike prior step-level methods, CoCA requires no additional networks (critics) or reward annotations; it operates purely as a reward redistribution mechanism conditioned on realized trajectories.
3. Empirical Observation: The paper demonstrates that timestep-level latent progress patterns vary substantially across different prompts and trajectories, explaining why fixed temporal weighting schedules fail to generalize and motivating the need for adaptive assignment.
4. General Applicability: The framework is algorithm-agnostic, successfully integrated with both REINFORCE-based (DDPO) and GRPO-based (DanceGRPO) optimization pipelines.

4. Experimental Results

Experiments were conducted across four human preference reward functions (Aesthetic, ImageReward, PickScore, HPSv2) and two optimization settings (DDPO and DanceGRPO).

• Sample Efficiency: CoCA improves optimization efficiency by 25% to 100% compared to trajectory-level baselines (DDPO) and existing step-level baselines (TDPO, UCA). It achieves steeper learning curves and faster convergence.
• Performance: In the REINFORCE setting, CoCA achieved the highest scores across all metrics, outperforming both trajectory-level and step-level baselines. In the GRPO setting, it showed superior results, particularly in multi-reward optimization scenarios.
• Generalization: CoCA demonstrated improved generalization on unseen prompts and unseen reward functions (e.g., improving PickScore when trained on HPSv2), suggesting the improvements stem from genuine image quality enhancements rather than reward hacking.
• Ablation Studies:
  • Window Size: A moderate window size ($W=5$) yielded the best performance, balancing stability and granularity.
  • Distance Metric: Cosine similarity outperformed $\ell_2$ distance for measuring latent progress.
  • Normalization: Combining both pre- and post-redistribution normalization provided the most stable training.

5. Significance and Claims

The paper claims that CoCA offers a "free" improvement in RL-finetuning by solving the credit assignment problem inherent in sparse reward settings. Its significance lies in:

• Efficiency: Drastically reducing the number of reward queries (and thus compute/data) required to align generative models with human preferences.
• Simplicity: Providing a general, plug-and-play solution that does not increase model complexity or require auxiliary training.
• Theoretical Soundness: Ensuring that the redistribution of rewards does not bias the optimization objective, maintaining the integrity of the original policy gradient.

The authors conclude that by providing denser, outcome-aligned learning signals, CoCA lowers the barrier for beneficial applications in content creation and assistive design, while acknowledging that improved efficiency could also facilitate misuse or amplify biases if not carefully managed.