提出新方法提升扩散模型生成多样性与效率,避免奖励欺骗导致的模式崩溃。
TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment

- 用轨迹分布匹配替代奖励最大化,避免集中于少数高奖励路径。
- 在多种任务上生成多样性提升9.1%,同时保持优秀下游性能。
- 适合追求高质量、多样化生成结果的研究者与应用开发者。
强化学习在对齐扩散模型与下游任务方面展现出巨大潜力,但多数方法仍存在严重的奖励欺骗问题,导致生成多样性下降和质量退化,表现为视觉模式崩溃并放大不可靠奖励。我们识别出根本原因在于现有方法具有模式聚焦特性:仅最大化期望奖励,而未有效约束可接受轨迹的概率分布,导致概率集中在少数高奖励路径。为此,我们提出轨迹匹配策略优化(TMPO),将标量奖励最大化替换为轨迹级奖励分布匹配。具体而言,TMPO引入软最大轨迹平衡(Softmax-TB)目标,使策略对K条轨迹的概率匹配由奖励诱导的玻尔兹曼分布。我们证明该目标继承前向KL散度的覆盖特性,在优化奖励的同时保留所有可接受轨迹的覆盖。为降低大规模流匹配模型上的多轨迹训练时间,TMPO采用动态随机树采样,轨迹共享去噪前缀,并在动态调度步骤分叉,减少冗余计算,提升训练效率。大量实验表明,无论是在人类偏好、组合生成还是文本渲染等多样任务中,TMPO均比当前最优方法生成多样性提升9.1%,并在所有下游指标与效率指标上表现优异,实现了奖励与多样性的最佳平衡。
原文摘要 · Abstract (English)
Reinforcement learning (RL) has shown extraordinary potential in aligning diffusion models to downstream tasks, yet most of them still suffer from significant reward hacking, which degrades generative diversity and quality by inducing visual mode collapse and amplifying unreliable rewards. We identify the root cause as the mode-seeking nature of these methods, which maximize expected reward without effectively constraining probability distribution over acceptable trajectories, causing concentration on a few high-reward paths. In contrast, we propose Trajectory Matching Policy Optimization (TMPO), which replaces scalar reward maximization with trajectory-level reward distribution matching. Specifically, TMPO introduces a Softmax Trajectory Balance (Softmax-TB) objective to match the policy probabilities of K trajectories to a reward-induced Boltzmann distribution. We prove that this objective inherits the mode-covering property of forward KL divergence, preserving coverage over all acceptable trajectories while optimizing reward. To further reduce multi-trajectory training time on large-scale flow-matching models, TMPO incorporates Dynamic Stochastic Tree Sampling, where trajectories share denoising prefixes and branch at dynamically scheduled steps, reducing redundant computation while improving training effectiveness. Extensive results across diverse alignment tasks such as human preference, compositional generation and text rendering show that TMPO improves generative diversity over state-of-the-art methods by 9.1%, and achieves competitive performance in all downstream and efficiency metrics, attaining the optimal trade-off between reward and diversity.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。