TR2-D2 : Ajustement fin guidé par recherche arborescente avec prise en compte des trajectoires pour la diffusion discrète

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

September 29, 2025

papers.authors: Sophia Tang, Yuchen Zhu, Molei Tao, Pranam Chatterjee

cs.AI

papers.abstract

L'apprentissage par renforcement avec contrôle optimal stochastique offre un cadre prometteur pour le réglage fin de modèles de diffusion, où un modèle de diffusion pré-entraîné est optimisé pour générer des trajectoires menant à une distribution biaisée par une récompense. Bien que ces approches permettent une optimisation sans accès à des échantillons explicites de la distribution optimale, elles nécessitent un entraînement sur des déploiements sous le modèle actuellement réglé, les rendant vulnérables au renforcement de trajectoires sous-optimales produisant de faibles récompenses. Pour surmonter ce défi, nous introduisons TRee Search Guided TRajectory-Aware Fine-Tuning for Discrete Diffusion (TR2-D2), un nouveau cadre qui optimise les trajectoires de diffusion discrète guidées par récompense à l'aide de la recherche arborescente pour construire des mémoires tampons de rejeu destinées au réglage fin prenant en compte les trajectoires. Ces mémoires tampons sont générées à l'aide de la recherche arborescente de Monte Carlo (MCTS) et utilisées ensuite pour régler finement un modèle de diffusion discrète pré-entraîné sous un objectif de contrôle optimal stochastique. Nous validons notre cadre sur le réglage fin mono- et multi-objectif de modèles de diffusion de séquences biologiques, mettant en évidence l'efficacité globale de TR2-D2 pour un réglage fin fiable guidé par récompense dans la génération de séquences discrètes.

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.