Research map / Reinforcement learning / RL from feedback

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

NeurIPS 2026 · Reinforcement learning › RL from feedback · approach: RL fine-tuning of generative models · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes TMPO, a trajectory-level reward distribution matching method (Softmax Trajectory Balance with prefix-sharing tree rollouts) that aligns FLUX.1-dev while preserving diversity and cutting training time.

Problem

RL alignment of diffusion/flow text-to-image models via reward maximization suffers from reward hacking, causing visual mode collapse (reduced diversity) and amplified unreliable proxy rewards.

Approach

Replace reward maximization with within-group trajectory-level distribution matching: a Softmax-TB advantage log(q_i/p_i) matches policy trajectory probabilities to a Boltzmann target exp(beta R), cancelling the partition function by normalizing within K trajectories sharing a prefix; rollouts use a dynamic stochastic tree with curriculum-scheduled branch points.

Key results

• Under GenEval-only training, TMPO reaches 0.949 GenEval vs 0.943 for Flow-GRPO, with LGMD 0.131 vs -0.092 for Flow-GRPO.
• TMPO reduces iteration time to 68.3 s vs 108.9 s for Flow-GRPO and 79.1 s for TreeGRPO under PickScore-only training.
• Under PickScore-only training TMPO obtains PickScore 24.277 (best) and LGMD 0.204 vs -0.107 for Flow-GRPO; paper reports 9.1% average diversity gain and up to 27% training time reduction.

Surprising findings

• Reward maximization baselines drive LGMD negative (mode collapse) even while raising reward scores.
• Step-wise (per-step) detailed balance causes severe loss oscillation and near-zero LGMD, unlike trajectory-level balance.

Limitations

• Boltzmann target quality is upper-bounded by the underlying reward model.
• Evaluated only on text-to-image generation; video, 3D, and robotics are untested.

Experimental setup

Datasets: HPDv2 GenEval prompts (task-specific constructed) OCR prompts (task-specific constructed) Pick-a-Pic

Benchmarks: GenEval OCR accuracy DrawBench HPS-v2.1 ImageReward PickScore

Metrics: GenEval accuracy OCR accuracy (1-NED) HPS-v2.1 ImageReward PickScore LGMD Cosine Diversity iteration time

Baselines: Flow-GRPO MixGRPO TreeGRPO GARDO

Builds on: GFlowNet trajectory balance Flow-GRPO TreeGRPO GRPO-Guard CPS

Closest papers by meaning

• Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View (NeurIPS 2026 · similarity 0.92)
• Composite Flow Matching for Reinforcement Learning with Shifted-Dynamics Data (NeurIPS 2025 · similarity 0.89)
• GenPO: Generative Diffusion Models Meet On-Policy Reinforcement Learning (NeurIPS 2025 · similarity 0.89)
• Aligning Diffusion Behaviors with Q-functions for Efficient Continuous Control (NeurIPS 2024 · similarity 0.89)
• Beyond the Dirac Delta: Mitigating Diversity Collapse in Reinforcement Fine-Tuning for Image Generation (NeurIPS 2026 · similarity 0.89)
• Flow-GRPO: Training Flow Matching Models via Online RL (NeurIPS 2025 · similarity 0.89)