Technical Summary: TMPO for Diverse and Efficient Diffusion Alignment

1. Problem Statement

Reinforcement Learning (RL) has demonstrated significant potential in aligning diffusion and flow-matching models with downstream tasks. However, existing methods predominantly rely on scalar reward maximization, which the authors identify as inherently mode-seeking. This optimization objective causes models to concentrate probability mass on a few high-reward trajectories, leading to two critical failure modes:

1. Degradation of Generative Diversity: The model collapses into a narrow set of compositions and styles (visual mode collapse), ignoring the anisotropic distribution of plausible outputs learned during pre-training.
2. Amplification of Unreliable Rewards: By over-reinforcing specific features that proxy rewards (e.g., OCR accuracy or specific aesthetic scores) favor, the model amplifies spurious attributes, often at the expense of overall generation quality.

Current mitigation strategies, such as KL regularization, fail to alter the fundamental mode-seeking nature of the objective, while existing Generative Flow Network (GFlowNet) approaches often rely on intractable state-flow estimations or local constraints that accumulate errors over long denoising horizons.

2. Methodology: Trajectory Matching Policy Optimization (TMPO)

TMPO fundamentally shifts the optimization paradigm from scalar reward maximization to trajectory-level reward distribution matching.

2.1 Softmax Trajectory Balance (Softmax-TB)

Instead of maximizing expected reward, TMPO aims to match the policy's trajectory probabilities to a target distribution induced by rewards.

• Objective: For a group of $K$ trajectories generated from the same prompt, the method matches the policy probabilities $P_\theta(\tau)$ to a Boltzmann distribution $q(\tau) \propto \exp(\beta R(\tau))$.
• Partition-Free Formulation: By normalizing probabilities and rewards within the group of $K$ trajectories, the intractable global partition function is eliminated. The method defines a Softmax-TB advantage:
  $$A_i = \log q_i - \log p_i = \log \frac{\exp(\beta R_i)}{\sum_j \exp(\beta R_j)} - \log \frac{P_\theta(\tau_i)}{\sum_j P_\theta(\tau_j)}$$
• Theoretical Guarantee: The authors prove that minimizing this objective is equivalent to minimizing the forward KL divergence ($D_{KL}(q \parallel p)$) within the group. Unlike reverse KL (used in standard policy gradients), forward KL is mode-covering: it imposes an infinite penalty if the policy probability $p_i$ approaches zero for any trajectory $i$ that has positive target weight $q_i$. This theoretically guarantees the preservation of diversity without requiring auxiliary KL regularization.

2.2 Dynamic Stochastic Tree Sampling

To address the computational cost of sampling multiple independent trajectories for large-scale flow-matching models, TMPO introduces a tree-structured sampling strategy:

• Prefix Sharing: Trajectories share a deterministic ODE prefix (common denoising steps) and branch at dynamically scheduled stochastic points.
• Dynamic Branching: Branch points are not fixed; they follow a curriculum that shifts from early high-noise regions (coarse-grained diversity) to later low-noise regions (fine-grained variations) as training progresses. Branch positions are stochastically perturbed using a Beta distribution to prevent overfitting to a fixed tree geometry.
• Efficiency: This approach reduces redundant computation by sharing the forward pass for the shared prefix, significantly lowering training time while maintaining effective exploration.

3. Key Contributions

1. Root Cause Identification: The paper identifies the mode-seeking nature of reward maximization as the primary cause of reward hacking and diversity degradation in diffusion RL, proposing a shift to distribution matching.
2. Softmax-TB Objective: Introduction of a partition-free, trajectory-level distribution matching objective that inherits the mode-covering properties of forward KL divergence, naturally preserving generative diversity.
3. Dynamic Stochastic Tree Sampling: A scalable sampling mechanism that reduces training time for large-scale models by sharing denoising prefixes and dynamically adjusting branching schedules.

4. Experimental Results

The authors evaluated TMPO on the FLUX.1-dev backbone across three alignment tasks: compositional image generation (GenEval), visual text rendering (OCR), and human preference alignment (PickScore).

• Diversity: TMPO improved generative diversity by 9.1% on average compared to state-of-the-art methods (Flow-GRPO, MixGRPO, TreeGRPO, GARDO). It was the only method to consistently maintain positive Log Geometric Mean Distance (LGMD) scores, indicating healthy diversity, whereas baselines often exhibited negative LGMD (mode collapse).
• Performance: TMPO achieved competitive or superior performance on task-specific metrics (GenEval accuracy, OCR accuracy) and human preference scores (HPS-v2.1, ImageReward, PickScore).
• Efficiency: By utilizing prefix sharing and tree sampling, TMPO reduced iteration training time by up to 27% compared to existing methods.
• Trade-offs: Pareto analysis demonstrated that TMPO achieves the optimal trade-off between reward maximization, diversity preservation, and training efficiency.

5. Significance and Claims

The paper claims that TMPO offers a fundamental solution to the reward hacking dilemma in generative AI alignment. By reformulating the objective to match reward-induced distributions rather than maximizing scalar rewards, TMPO structurally enforces diversity preservation without relying on external regularization or heuristic reweighting. The method demonstrates that it is possible to achieve high alignment scores while maintaining the rich, multimodal output distribution learned during pre-training, all while reducing the computational overhead typically associated with multi-trajectory RL methods. The authors position TMPO as a scalable framework suitable for large-scale flow-matching models, with potential extensions to video generation and robotic action generation.