The authors propose TMPO (Trajectory Matching Policy Optimization) to align diffusion models using reinforcement learning without suffering from reward hacking — the problem where models collapse onto a few high-reward outputs and lose diversity. Instead of maximizing expected reward (which is "mode-seeking"), TMPO matches the model's probability distribution over entire generation trajectories to a reward-induced Boltzmann distribution. This "mode-covering" approach preserves diversity over all acceptable outputs while still optimizing reward. They also introduce a tree-sampling trick to share computation across multiple trajectories during training, speeding up large flow-matching models.
Main takeaways:
- Standard RL fine-tuning of diffusion models is mode-seeking: it concentrates probability on a few high-reward paths, causing visual mode collapse and reward hacking.
- TMPO uses trajectory-level distribution matching (Softmax Trajectory Balance objective) to cover all acceptable trajectories, not just maximize reward.
- Improves generative diversity by 9.1% over state-of-the-art methods while maintaining competitive reward and efficiency.
- Dynamic Stochastic Tree Sampling shares denoising prefixes across K trajectories, reducing redundant computation during multi-trajectory training.
- Effective across human preference alignment, compositional generation, and text rendering tasks.