TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete
  Diffusion
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.
68dc27df4159d1f2418f98ee
Sophia Tang
68dc27df4159d1f2418f98ef
Yuchen Zhu
68dc27df4159d1f2418f98f0
Molei Tao
68dc27df4159d1f2418f98f1
Pranam Chatterjee
TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete
  Diffusion
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.
A novel framework, TR2-D2, uses tree search to optimize reward-guided discrete diffusion trajectories for fine-tuning biological sequence models.
reinforcement learning stochastic optimal control diffusion fine-tuning reward-tilted distribution rollouts sub-optimal trajectories TRee Search Guided TRajectory-Aware Fine-Tuning for Discrete Diffusion TR2-D2 discrete diffusion trajectories Monte Carlo Tree Search MCTS trajectory-aware fine-tuning biological sequence diffusion models single-objective fine-tuning multi-objective fine-tuning