TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion
This addresses the challenge of reliable reward-guided fine-tuning in discrete diffusion models for applications like biological sequence generation, representing an incremental improvement over existing methods.
The paper tackled the problem of fine-tuning pre-trained diffusion models for reward-guided discrete sequence generation, which often reinforces sub-optimal trajectories, by introducing TR2-D2, a framework that uses tree search to construct replay buffers for trajectory-aware fine-tuning, resulting in effective optimization for biological sequence generation tasks.
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.