TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment

    We propose TMPO, a trajectory-level reward distribution matching framework for diffusion and flow model alignment. It replaces scalar reward maximization with a Softmax Trajectory Balance objective, preserves coverage over acceptable generation trajectories, and accelerates multi-trajectory training with Dynamic Stochastic Tree Sampling; across human preference, compositional generation, and text rendering tasks, TMPO improves generative diversity by 9.1% while maintaining competitive reward and efficiency.
    🌟This work extends my exploration of unified intelligence from structured visual manipulation to distribution-aware diffusion post-training—aligning generative agents with preference signals while preserving diverse, plausible visual worlds.

BibTeX

@article{li2026tmpo,
  title = {TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment},
  author = {Li, Jiaming^{1,2} and Zhu, Chenyu^{1} and Yi, Nanxi^1 and Bao, Youjun^2 and Sun, Li^2 and Lv, Quanying^2 and Fang, Xiang^3 and Liu, Daizong^4 and Li, Jianjun^1 and He, Kun^1 and Zhou, Bowen^5 and Ma, Zhiyuan^{1\dag}},
  journal = {arXiv:2605.10983},
  year = {2026},
  month = {may},
  url = {https://arxiv.org/abs/2605.10983}
}