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

• URL: http://arxiv.org/abs/2605.10983v2
• Date: Wed, 13 May 2026 08:00:49 GMT
• Title: TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment
• Abstract summary: We propose Trajectory Matching Policy Balance Optimization (TMPO), which replaces reward with human-level reward distribution matching.<n>TMPO improves generative diversity over state-of-the-art methods by 9.1%, and achieves competitive performance in all downstream and efficiency metrics.<n>To further reduce multi-trajectory training time on large-scale flow- prefixes, TMPO incorporates Dynamic Tree Sampling models, where trajectories share denoising and branch at dynamically scheduled steps.
• Score: 52.570581883709345
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Scaling Reinforcement Learning for Diffusion Models via Velocity Matching [66.34541945134916]
  RewardBench fine-tuning is an important tool for adapting diffusion models to human preferences and task-specific objectives.<n>Unlike autoregressive models, diffusion models do not provide tractable likelihoods for generated samples.<n>We propose reward-based velocity matching (RVM), a simple trajectory-free update that acts directly on the velocity field.
  arXiv  Detail & Related papers  (2026-08-24T17:13:43Z)
• RSPO: Reward-Swap Policy Optimization for Multi-Turn LLM Agents [24.215311212134882]
  We propose Reward-Swap Policy Optimization (RSPO) to leverage the rich information from dense process rewards to facilitate training with outcome rewards.<n>RSPO ensures the diversity of sampled trajectories while guaranteeing consistency between the optimization objective and the true outcome rewards.
  arXiv  Detail & Related papers  (2026-07-06T06:32:39Z)
• DiffusionOPD: A Unified Perspective of On-Policy Distillation in Diffusion Models [55.01951088768769]
  DiffusionOPD is a new multi-task training paradigm for diffusion models based on Online Policy Distillation (OPD)<n>We show that DiffusionOPD consistently surpasses both multi-reward RL and cascade RL baselines in training efficiency and final performance.
  arXiv  Detail & Related papers  (2026-05-14T16:49:09Z)
• Pareto-Guided Optimal Transport for Multi-Reward Alignment [59.776579791864805]
  We show how to optimize against a unified global target under heterogeneous reward upper bounds.<n>We develop both online and offline optimization strategies tailored to diverse reward signal characteristics.<n>Our approach outperforms strong baselines with an 11% gain in Joint Domination Rate (JDR) and achieves a near 80% win rate in human evaluations.
  arXiv  Detail & Related papers  (2026-05-13T08:19:48Z)
• MARBLE: Multi-Aspect Reward Balance for Diffusion RL [71.6241143519038]
  Reinforcement learning has become the dominant approach for aligning diffusion models with human preferences.<n>Existing practice deal with multiple rewards by training one specialist model per reward.<n>We propose MARBLE, a gradient-space optimization framework that maintains independent advantage estimators for each reward.
  arXiv  Detail & Related papers  (2026-05-07T16:20:42Z)
• HAD: Combining Hierarchical Diffusion with Metric-Decoupled RL for End-to-End Driving [51.268878540511054]
  We propose HAD, an end-to-end planning framework with a Hierarchical Diffusion Policy.<n>We show that HAD achieves new state-of-the-art performance on both NAVSIM and HUGSIM.
  arXiv  Detail & Related papers  (2026-04-04T04:12:47Z)
• SetPO: Set-Level Policy Optimization for Diversity-Preserving LLM Reasoning [50.93295951454092]
  We introduce a set level diversity objective defined over sampled trajectories using kernelized similarity.<n>Our approach derives a leave-one-out marginal contribution for each sampled trajectory and integrates this objective as a plug-in advantage shaping term for policy optimization.<n>Experiments across a range of model scales demonstrate the effectiveness of our proposed algorithm, consistently outperforming strong baselines in both Pass@1 and Pass@K across various benchmarks.
  arXiv  Detail & Related papers  (2026-02-01T07:13:20Z)
• GARDO: Reinforcing Diffusion Models without Reward Hacking [54.841464430913476]
  Fine-tuning diffusion models via online reinforcement learning (RL) has shown great potential for enhancing text-to-image alignment.<n>The mismatch often leads to reward hacking, where proxy scores increase while real image quality deteriorates and generation diversity collapses.<n>We propose Gated and Adaptive Regularization with Diversity-aware Optimization (GARDO) to address the competing demands of sample efficiency, effective exploration, and mitigation of reward hacking.
  arXiv  Detail & Related papers  (2025-12-30T10:55:45Z)
• VARD: Efficient and Dense Fine-Tuning for Diffusion Models with Value-based RL [28.95582264086289]
  VAlue-based Reinforced Diffusion (VARD) is a novel approach that first learns a value function predicting expection of rewards from intermediate states.<n>Our method maintains proximity to the pretrained model while enabling effective and stable training via backpropagation.
  arXiv  Detail & Related papers  (2025-05-21T17:44:37Z)
• Test-time Alignment of Diffusion Models without Reward Over-optimization [8.981605934618349]
  Diffusion models excel in generative tasks, but aligning them with specific objectives remains challenging.<n>We propose a training-free, test-time method based on Sequential Monte Carlo (SMC) to sample from the reward-aligned target distribution.<n>We demonstrate its effectiveness in single-reward optimization, multi-objective scenarios, and online black-box optimization.
  arXiv  Detail & Related papers  (2025-01-10T09:10:30Z)
• Mimicking Better by Matching the Approximate Action Distribution [48.95048003354255]
  We introduce MAAD, a novel, sample-efficient on-policy algorithm for Imitation Learning from Observations. We show that it requires considerable fewer interactions to achieve expert performance, outperforming current state-of-the-art on-policy methods.
  arXiv  Detail & Related papers  (2023-06-16T12:43:47Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.