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

2025/09/29 by Shaoqing Tang, Yuchen Zhu, Tang, Sophia +5 · 2 citations
Computer Science · Engineering · #Biomolecules (q-bio.BM) #FOS: Biological sciences #FOS: Computer and information sciences #Machine Learning (cs.LG) #Music Technology and Sound Studies #Music and Audio Processing #VLSI and FPGA Design Techniques

paper · pdf · doi:10.48550/arxiv.2509.25171

openalex publication_date 2025/09/29 · openalex created_date 2025/10/19 · openalex updated_date 2026/07/28

Abstract

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.

Citations

• DiffusionNFT: Online Diffusion Reinforcement with Forward Process
• Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference
• MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control
• Seed Diffusion: A Large-Scale Diffusion Language Model with High-Speed Inference
• Theory-Informed Improvements to Classifier-Free Guidance for Discrete Diffusion Models
• Iterative Distillation for Reward-Guided Fine-Tuning of Diffusion Models in Biomolecular Design
• DiffuCoder: Understanding and Improving Masked Diffusion Models for Code Generation
• Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models
• Provable Maximum Entropy Manifold Exploration via Diffusion Models
• Mercury: Ultra-Fast Language Models Based on Diffusion
• Diffuse Everything: Multimodal Diffusion Models on Arbitrary State Spaces
• Solving Inverse Problems via Diffusion-Based Priors: An Approximation-Free Ensemble Sampling Approach
• Inference-time Scaling of Diffusion Models through Classical Search
• Muddit: Liberating Generation Beyond Text-to-Image with a Unified Discrete Diffusion Model
• Multi-Objective-Guided Discrete Flow Matching for Controllable Biological Sequence Design
• Flow-GRPO: Training Flow Matching Models via Online RL
• d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning
• Preference-Guided Diffusion for Multi-Objective Offline Optimization
• Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models
• Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Experts
• A Simple and Effective Reinforcement Learning Method for Text-to-Image Diffusion Fine-tuning
• Direct Discriminative Optimization: Your Likelihood-Based Visual Generative Model is Secretly a GAN Discriminator
• Diffusion-Sharpening: Fine-tuning Diffusion Models with Denoising Trajectory Sharpening
• Large Language Diffusion Models
• Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods
• Diffusion Generative Modeling for Spatially Resolved Gene Expression Inference from Histology Images
• A General Framework for Inference-time Scaling and Steering of Diffusion Models
• Stochastic Control for Fine-tuning Diffusion Models: Optimality, Regularity, and Convergence
• PepTune: De Novo Generation of Therapeutic Peptides with Multi-Objective-Guided Discrete Diffusion
• Simple Guidance Mechanisms for Discrete Diffusion Models
• ParetoFlow: Guided Flows in Multi-Objective Optimization
• Constrained Multi-objective Bayesian Optimization through Optimistic Constraints Estimation
• Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design
• COS-DPO: Conditioned One-Shot Multi-Objective Fine-Tuning Framework
• Steering Masked Discrete Diffusion Models via Discrete Denoising Posterior Prediction
• Meissonic: Revitalizing Masked Generative Transformers for Efficient High-Resolution Text-to-Image Synthesis
• Plug-and-Play Controllable Generation for Discrete Masked Models
• Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control
• Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling
• MODULI: Unlocking Preference Generalization via Diffusion Models for Offline Multi-Objective Reinforcement Learning
• Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding
• Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review
• Simple and Effective Masked Diffusion Language Models
• Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data
• Simplified and Generalized Masked Diffusion for Discrete Data
• Unlocking Guidance for Discrete State-Space Diffusion and Flow Models
• Trivialized Momentum Facilitates Diffusion Generative Modeling on Lie Groups
• Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond
• Scaling Rectified Flow Transformers for High-Resolution Image Synthesis
• Diffusion Language Models Are Versatile Protein Learners
• Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control
• Self-Play Fine-Tuning of Diffusion Models for Text-to-Image Generation
• Dirichlet Flow Matching with Applications to DNA Sequence Design
• BioXP-0.5B: Explainable Medical-AI via RL-GRPO
• Multi-Objective Optimization via Wasserstein-Fisher-Rao Gradient Flow
• Diffusion Model Alignment Using Direct Preference Optimization
• Unexpected Improvements to Expected Improvement for Bayesian Optimization
• Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution
• Directly Fine-Tuning Diffusion Models on Differentiable Rewards
• PeptideBERT: A Language Model based on Transformers for Peptide Property Prediction
• BOtied: Multi-objective Bayesian optimization with tied multivariate ranks
• DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models
• Training Diffusion Models with Reinforcement Learning
• Dirichlet Diffusion Score Model for Biological Sequence Generation
• Universal Guidance for Diffusion Models
• Sample-efficient Multi-objective Molecular Optimization with GFlowNets
• Evolutionary-scale prediction of atomic-level protein structure with a language model
• Sketch-Guided Text-to-Image Diffusion Models
• Multi-Objective GFlowNets
• A Pareto-optimal compositional energy-based model for sampling and optimization of protein sequences
• DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation
• DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation
• Classifier-Free Diffusion Guidance
• A Continuous Time Framework for Discrete Denoising Models
• Multi-Objective Bayesian Optimization over High-Dimensional Search Spaces
• Structured Denoising Diffusion Models in Discrete State-Spaces
• AutoDock Vina 1.2.0: New Docking Methods, Expanded Force Field, and Python Bindings
• Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation
• RoFormer: Enhanced Transformer with Rotary Position Embedding
• Effective gene expression prediction from sequence by integrating long-range interactions
• MARS: Markov Molecular Sampling for Multi-objective Drug Discovery
• Score-Based Generative Modeling through Stochastic Differential Equations
• Denoising Diffusion Implicit Models
• Denoising Diffusion Probabilistic Models
• Solving high-dimensional Hamilton-Jacobi-Bellman PDEs using neural\n networks: perspectives from the theory of controlled diffusions and measures\n on path space
• Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning
• Multi-objective Bayesian Optimization using Pareto-frontier Entropy
• Efficient Computation of Expected Hypervolume Improvement Using Box Decomposition Algorithms
• Max-value Entropy Search for Efficient Bayesian Optimization
• Predictive Entropy Search for Multi-objective Bayesian Optimization
• The sample size required in importance sampling
• Deep Unsupervised Learning using Nonequilibrium Thermodynamics
• Bayesian Optimization with Unknown Constraints
• ChEMBL: a large-scale bioactivity database for drug discovery
• AutoDock4 and AutoDockTools4: Automated docking with selective receptor flexibility

Cited by

• Entangled Schrödinger Bridge Matching
• Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization
• MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control

Related

• Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review
• A Continuous-Time Reinforcement Learning Framework for Fine-Tuning Discrete Diffusion Models
• Efficient Controllable Diffusion via Optimal Classifier Guidance
• Diffusion-ES: Gradient-free Planning with Diffusion for Autonomous Driving and Zero-Shot Instruction Following
• Maximize Your Diffusion: A Study into Reward Maximization and Alignment for Diffusion-based Control