论文

TMPO:轨迹匹配策略优化,实现多样化、高效的扩散对准

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

模型训练强化学习

摘要

强化学习(RL)在将扩散模型与下游任务结合起来方面表现出了非凡的潜力,但其中大多数仍然受到严重的 奖励作弊 的影响,它通过诱导视觉模式崩溃和放大不可靠的奖励来降低生成多样性和质量。我们将根本原因确定为这些方法的模式寻求性质,这些方法最大化预期奖励,而没有有效限制可接受轨迹上的概率分布,导致集中在一些高奖励路径上。相比之下,我们提出轨迹匹配策略优化(TMPO),它用轨迹级奖励分配匹配代替标量奖励最大化。具体来说,TMPO 引入了 Softmax 轨迹平衡 (Softmax-TB) 目标,将 K 轨迹的策略概率与奖励诱导的玻尔兹曼分布相匹配。我们证明这个目标继承了前向 KL 散度的模式覆盖特性,在优化奖励的同时保留了对所有可接受轨迹的覆盖。为了进一步减少大规模流匹配模型上的多轨迹训练时间,TMPO 结合了动态随机树采样,其中轨迹在动态调度的步骤中共享去噪前缀和分支,减少冗余计算,同时提高训练效果。人类偏好、构图生成和文本渲染等不同对齐任务的广泛结果表明,TMPO 比最先进的方法将生成多样性提高了 9.1%,并在所有下游和效率指标中实现了有竞争力的性能,实现了奖励和多样性之间的最佳权衡。