Technical Summary: AEGPO – Adaptive Entropy-Guided Policy Optimization for Diffusion Models

1. Problem Statement

Reinforcement Learning from Human Feedback (RLHF), particularly Group Relative Policy Optimization (GRPO), has become a standard paradigm for aligning diffusion and flow-matching models with human preferences. However, existing GRPO variants suffer from two fundamental inefficiencies due to their reliance on static heuristics:

1. Inefficient Sample Allocation: Standard methods assign a uniform number of rollouts to every prompt, ignoring the significant variation in "learning value" across different samples. Some prompts induce minimal policy changes, while others drive substantial improvements.
2. Static Timestep Exploration: Exploration strategies (e.g., branching in tree-based rollouts) typically follow fixed, predefined schedules. These fail to adapt to the dynamic nature of the denoising process, where critical moments for exploration vary significantly between prompts and training stages.

These limitations result in wasted computational resources on low-value samples and suboptimal exploration at non-critical timesteps, slowing convergence and limiting final alignment performance.

2. Methodology: AEGPO

The authors propose Adaptive Entropy-Guided Policy Optimization (AEGPO), a plug-and-play framework that leverages intrinsic attention signals to dynamically guide both sample selection and exploration. The core insight is that Attention Entropy serves as a powerful dual-signal proxy.

2.1 Core Signals

The method defines two intrinsic metrics derived from the model's attention maps:

• $\Delta$Entropy (Sample-Level Signal): The relative change in attention dispersion between the current policy and the base policy.
  • Definition: $\Delta\text{Entropy} = \frac{1}{T} \sum_{t=1}^{T} |\text{Entropy}_{\theta}(t) - \text{Entropy}_{\theta_{base}}(t)|$.
  • Function: It acts as a robust proxy for sample learning value. High $\Delta$Entropy indicates that a prompt forces the model to significantly reorganize its attention strategy, implying high learning potential.
• Entropy($t$) (Timestep-Level Signal): The absolute attention dispersion at a specific denoising step $t$.
  • Definition: The mean Shannon entropy of attention weights across text tokens for image features at step $t$.
  • Function: Peaks in Entropy($t$) identify critical exploration timesteps. These peaks correspond to moments of high uncertainty or broad semantic integration (often occurring in early structural and late refinement stages), where exploration yields the highest reward variance and diversity.

2.2 Dual-Level Adaptive Strategy

AEGPO utilizes these signals to reallocate computation without increasing the total rollout budget:

1. Global Adaptive Allocation:

   • The batch of prompts is split into "high-value" and "low-value" tiers based on the median $\Delta$Entropy.
   • High-value prompts (high $\Delta$Entropy) are assigned a larger rollout budget ($r_{high}$), while low-value prompts receive fewer rollouts ($r_{low}$).
   • Constraint: $r_{high} + r_{low} = 2r_{avg}$, ensuring total compute remains constant.

2. Local Adaptive Exploration:

   • Instead of fixed branching schedules, the method dynamically identifies the top-$K$ timesteps with the highest Entropy($t$) peaks for each rollout.
   • Branching (exploration) is performed exclusively at these identified high-dispersion timesteps.
   • This ensures exploration is concentrated on the most informative moments of the denoising trajectory.

3. Key Contributions

• Dual-Signal Discovery: The paper establishes Attention Entropy as a dual-signal intrinsic proxy, where relative change ($\Delta$Entropy) indicates sample learning value and absolute peaks (Entropy($t$)) identify critical exploration timesteps.
• AEGPO Framework: Introduction of a dual-level adaptive framework that improves sample selection and timestep exploration using lightweight, entropy-guided mechanisms.
• Generalizability: Demonstration that AEGPO is model-agnostic and framework-agnostic, successfully enhancing diverse policy optimization methods (DanceGRPO, BranchGRPO, FlowGRPO, DiffusionNFT) and base models (FLUX.1-dev, SD3.5-M).

4. Experimental Results

Experiments were conducted on the HPDv2.1 and Pick-a-Pic datasets using FLUX.1-dev and SD3.5-M as base models.

• Convergence Speed: AEGPO significantly accelerates policy optimization. It achieves 2$\times$ faster convergence on DanceGRPO and 5$\times$ faster convergence on DiffusionNFT compared to standard variants, reaching superior final rewards in fewer steps.
• Performance Metrics:
  • On FLUX.1-dev, BranchGRPO + AEGPO achieved a new state-of-the-art HPS-v2.1 score of 0.374 (up from 0.363) and Pick Score of 0.232 (up from 0.229).
  • On SD3.5-M, DiffusionNFT + AEGPO improved HPS-v2.1 to 0.343 and ImageReward to 1.57.
• Reward-KL Trade-off: AEGPO models consistently dominate the baseline on the Reward-KL Pareto frontier, achieving significantly higher alignment rewards for any given level of policy deviation (KL divergence).
• Ablation Studies: Both Global Allocation (G) and Local Exploration (L) provide independent benefits, with the combined approach (G+L) yielding the highest performance.
• Computational Overhead: The method introduces a modest overhead (approx. 11.1% increase in step time and negligible VRAM increase) due to extracting intermediate attention maps, which is justified by the dramatic reduction in total training steps required to reach peak performance.

5. Significance and Claims

The paper claims that AEGPO addresses the core inefficiencies of static sampling in diffusion model alignment. By shifting from heuristic-based allocation to intrinsic signal-based adaptation, the method ensures that computational resources are concentrated on the most informative samples and the most critical moments of the generation process.

The authors emphasize that AEGPO is a generalizable, plug-and-play module that does not require new loss functions or complex reward modeling. Instead, it leverages the model's own internal attention dynamics to guide optimization, offering a more efficient path to high-quality alignment across diverse diffusion architectures and RLHF frameworks. The work suggests that intrinsic signals like attention entropy are underutilized resources for improving the efficiency of policy optimization in generative models.