TR2-D2: Boomzoek-geleide trajectbewuste fine-tuning voor discrete diffusie

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

September 29, 2025

Auteurs: Sophia Tang, Yuchen Zhu, Molei Tao, Pranam Chatterjee

cs.AI

Samenvatting

Reinforcement learning met stochastische optimale controle biedt een veelbelovend raamwerk voor diffusie-finetuning, waarbij een vooraf getraind diffusiemodel wordt geoptimaliseerd om paden te genereren die leiden naar een beloning-georiënteerde verdeling. Hoewel deze benaderingen optimalisatie mogelijk maken zonder expliciete voorbeelden van de optimale verdeling, vereisen ze training op rollouts onder het huidige gefinetunede model, waardoor ze vatbaar zijn voor het versterken van suboptimale trajecten die slechte beloningen opleveren. Om deze uitdaging te overwinnen, introduceren we TRee Search Guided TRajectory-Aware Fine-Tuning for Discrete Diffusion (TR2-D2), een nieuw raamwerk dat beloning-gestuurde discrete diffusietrajecten optimaliseert met behulp van boomzoeken om replay-buffers te construeren voor trajectbewuste finetuning. Deze buffers worden gegenereerd met behulp van Monte Carlo Tree Search (MCTS) en vervolgens gebruikt om een vooraf getraind discreet diffusiemodel te finetunen onder een stochastisch optimaal controle-doel. We valideren ons raamwerk op enkelvoudige en meervoudige doelstellingen voor finetuning van biologische sequentiediffusiemodellen, waarbij de algehele effectiviteit van TR2-D2 voor betrouwbare beloning-gestuurde finetuning in discrete sequentiegeneratie wordt benadrukt.

English

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.