mlsanityData to 14 Oct 2026

TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion

arXiv 2509.25171 ↗First posted 29 Sep 2025cs.LG q-bio.BMcited by 14

On the map: Diffusion models → Diffusion and flow policies for robotics

also in 1 more
Abstract (arXiv)

Reinforcement learning with stochastic optimal control offers a promising framework for diffusion fine-tuning, where a pre-trained diffusion model is optimized to generate paths that lead to a reward-tilted distribution. While these approaches enable optimization without access to explicit samples from the optimal distribution, they require training on rollouts under the current fine-tuned model, making them susceptible to reinforcing sub-optimal trajectories that yield poor rewards. To overcome this challenge, we introduce TRee Search Guided TRajectory-Aware Fine-Tuning for Discrete Diffusion (TR2-D2), a novel framework that optimizes reward-guided discrete diffusion trajectories with tree search to construct replay buffers for trajectory-aware fine-tuning. These buffers are generated using Monte Carlo Tree Search (MCTS) and subsequently used to fine-tune a pre-trained discrete diffusion model under a stochastic optimal control objective. We validate our framework on single- and multi-objective fine-tuning of biological sequence diffusion models, highlighting the overall effectiveness of TR2-D2 for reliable reward-guided fine-tuning in discrete sequence generation.

Key facts. Cited by 14 papers in the map (data through 14 Oct 2026). Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models. Builds on 95 of its 144 references that are in the map. First posted on arXiv 29 Sep 2025.

Builds on

95 of 144 references.

Closest ancestors · 15

Same direction, newest first

  1. 2025DiffusionNFT: Online Diffusion Reinforcement with Forward Process19 Sep
  2. 2025Flow-GRPO: Training Flow Matching Models via Online RL8 May
  3. 2025A Simple and Effective Reinforcement Learning Method for Text-to-Image Diffusion Fine-tuning2 Mar
  4. 2025Diffusion-Sharpening: Fine-tuning Diffusion Models with Denoising Trajectory Sharpening17 Feb
  5. 2025Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods3 Feb
  6. 2024Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence24 Dec
  7. 2024Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control13 Sep
  8. 2024MODULI: Unlocking Preference Generalization via Diffusion Models for Offline Multi-Objective Reinforcement Learning28 Aug
  9. 2024Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review18 Jul
  10. 2024Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control23 Feb
  11. 2024Self-Play Fine-Tuning of Diffusion Models for Text-to-Image Generation15 Feb
  12. 2023Diffusion Model Alignment Using Direct Preference Optimization21 Nov
3 more
  1. 2023Directly Fine-Tuning Diffusion Models on Differentiable Rewards29 Sep
  2. 2023DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models25 May
  3. 2023Training Diffusion Models with Reinforcement Learning22 May

Other references · 80

  1. 2020Denoising Diffusion Probabilistic Modelscited by 21,076
  2. 2020Denoising Diffusion Implicit Modelscited by 9,781
  3. 2024DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Modelscited by 9,379
  4. 2020Score-Based Generative Modeling through Stochastic Differential Equationscited by 9,122
  5. 2015Deep Unsupervised Learning using Nonequilibrium Thermodynamicscited by 7,058
  6. 2022Classifier-Free Diffusion Guidancecited by 5,847
  7. 2016XGBoost: A Scalable Tree Boosting Systemcited by 5,166
  8. 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,703
  9. 2024Scaling Rectified Flow Transformers for High-Resolution Image Synthesiscited by 4,181
  10. 2022DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generationcited by 3,300
  11. 2021Structured Denoising Diffusion Models in Discrete State-Spacescited by 1,719
  12. 2025Large Language Diffusion Modelscited by 912
38 more
  1. 2024Simple and Effective Masked Diffusion Language Modelscited by 848
  2. 2019Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learningcited by 782
  3. 2023Discrete Diffusion Modeling by Estimating the Ratios of the Data Distributioncited by 745
  4. 2024Simplified and Generalized Masked Diffusion for Discrete Datacited by 626
  5. 2022A Continuous Time Framework for Discrete Denoising Modelscited by 451
  6. 2023Universal Guidance for Diffusion Modelscited by 437
  7. 2025Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Modelscited by 409
  8. 2021Flow Network based Generative Models for Non-Iterative Diverse Candidate Generationcited by 392
  9. 2024Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Datacited by 381
  10. 2025Mercury: Ultra-Fast Language Models Based on Diffusioncited by 235
  11. 2023Practical and Asymptotically Exact Conditional Sampling in Diffusion Modelscited by 232
  12. 2025A General Framework for Inference-time Scaling and Steering of Diffusion Modelscited by 224
  13. 2024Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Samplingcited by 208
  14. 2017Max-value Entropy Search for Efficient Bayesian Optimizationcited by 206
  15. 2025DiffuCoder: Understanding and Improving Masked Diffusion Models for Code Generationcited by 187
  16. 2025d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learningcited by 186
  17. 2022Sketch-Guided Text-to-Image Diffusion Modelscited by 180
  18. 2020Differentiable Expected Hypervolume Improvement for Parallel Multi-Objective Bayesian Optimizationcited by 179
  19. 2025Seed Diffusion: A Large-Scale Diffusion Language Model with High-Speed Inferencecited by 178
  20. 2023Protein Design with Guided Discrete Diffusioncited by 157
  21. 2024Simple Guidance Mechanisms for Discrete Diffusion Modelscited by 155
  22. 2023Unexpected Improvements to Expected Improvement for Bayesian Optimizationcited by 149
  23. 2024Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decodingcited by 136
  24. 2024Unlocking Guidance for Discrete State-Space Diffusion and Flow Modelscited by 129
  25. 2021MARS: Markov Molecular Sampling for Multi-objective Drug Discoverycited by 125
  26. 2024Dirichlet Flow Matching with Applications to DNA Sequence Designcited by 115
  27. 2014Bayesian Optimization with Unknown Constraintscited by 109
  28. 2022Multi-Objective GFlowNetscited by 101
  29. 2024Diffusion Language Models Are Versatile Protein Learnerscited by 98
  30. 2025Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Expertscited by 96
  31. 2015Predictive Entropy Search for Multi-objective Bayesian Optimizationcited by 88
  32. 2021Multi-Objective Bayesian Optimization over High-Dimensional Search Spacescited by 87
  33. 2020Solving high-dimensional Hamilton-Jacobi-Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path spacecited by 82
  34. 2024Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Designcited by 79
  35. 2023Dirichlet Diffusion Score Model for Biological Sequence Generationcited by 75
  36. 2019Multi-objective Bayesian Optimization using Pareto-frontier Entropycited by 57
  37. 2024Meissonic: Revitalizing Masked Generative Transformers for Efficient High-Resolution Text-to-Image Synthesiscited by 56
  38. 2025Muddit: Liberating Generation Beyond Text-to-Image with a Unified Discrete Diffusion Modelcited by 55

50 most cited of 80 All 80 (JSON)

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 2; 2026: 12.

14 papers.

Sort:

By area:

  1. 2025Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimizationcited by 20
  2. 2026RAVEN: Real-time Autoregressive Video Extrapolation with Consistency-model GRPOcited by 7
  3. 2026Entropy Aware Reward Guidance for Diffusion Language Model Alignmentcited by 3
  4. 2026Advances in GRPO for Generation Models: A Surveycited by 2
  5. 2026A2D2: Fine-Tuning Any-Length Discrete Diffusion for Adaptive Decodingcited by 2
  6. 2025Entangled Schr"odinger Bridge Matchingcited by 1
  7. 2026TD3B: Transition-Directed Discrete Diffusion for Allosteric Binder Generationcited by 1
  8. 2026mRNAutilus: Multi-Objective-Guided Discrete Generation of mRNA with Optimized Therapeutic Propertiescited by 1
  9. 2026AlloGen: Conformation-Selective Binder Generation with Differential State Scoringcited by 1
  10. 2026FluxLite: Inference-Time Proposal Control for Discrete Diffusion Modelscited by 1
  11. 2026Minimal-Action Discrete Schr"odinger Bridge Matching for Peptide Sequence Designcited by 0
  12. 2026On the Design Space of Discrete Diffusion Online Adaptation for Molecular Optimizationcited by 0
2 more
  1. 2026pCoMole: Pareto-Constrained Molecule Editing with Discrete Flowscited by 0
  2. 2026Bridging the Evidence-to-Execution Gap:A Reflective Agent for Multi-Objective Peptide Designcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2026A Continuous-Time Reinforcement Learning Framework for Fine-Tuning Discrete Diffusion Models
  2. 2026Scaling Reinforcement Learning for Diffusion Models via Velocity Matching
  3. 2024Avoiding mode collapse in diffusion models fine-tuned with reinforcement learning
  4. 2025Maximize Your Diffusion: A Study into Reward Maximization and Alignment for Diffusion-based Control
  5. 2025Tree-Guided Diffusion Planner
  6. 2026Beyond Imitation: Reinforcement Learning Fine-Tuning for Adaptive Diffusion Navigation Policies
  7. 2026Diffusion Controller: Framework, Algorithms and Parameterization
  8. 2025VARD: Efficient and Dense Fine-Tuning for Diffusion Models with Value-based RL