TR2-D2: Ajuste Fino Guiado por Búsqueda en Árbol con Conciencia de Trayectoria para Difusión Discreta

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

September 29, 2025

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

cs.AI

Resumen

El aprendizaje por refuerzo con control óptimo estocástico ofrece un marco prometedor para el ajuste fino de difusión, donde un modelo de difusión preentrenado se optimiza para generar trayectorias que conduzcan a una distribución inclinada por recompensas. Si bien estos enfoques permiten la optimización sin acceso a muestras explícitas de la distribución óptima, requieren entrenamiento en simulaciones bajo el modelo ajustado actual, lo que los hace susceptibles a reforzar trayectorias subóptimas que producen recompensas deficientes. Para superar este desafío, presentamos TRee Search Guided TRajectory-Aware Fine-Tuning for Discrete Diffusion (TR2-D2), un marco novedoso que optimiza trayectorias de difusión discreta guiadas por recompensas mediante búsqueda en árbol para construir buffers de repetición para el ajuste fino consciente de trayectorias. Estos buffers se generan utilizando Búsqueda de Monte Carlo en Árbol (MCTS) y posteriormente se utilizan para ajustar un modelo de difusión discreta preentrenado bajo un objetivo de control óptimo estocástico. Validamos nuestro marco en el ajuste fino de objetivos únicos y múltiples en modelos de difusión de secuencias biológicas, destacando la efectividad general de TR2-D2 para un ajuste fino confiable guiado por recompensas en la generación de secuencias discretas.

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.