TR2-D2: Ajuste Fino Orientado por Busca em Árvore com Consciência de Trajetória para Difusão 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

Resumo

O aprendizado por reforço com controle ótimo estocástico oferece uma estrutura promissora para o ajuste fino de difusão, onde um modelo de difusão pré-treinado é otimizado para gerar caminhos que levam a uma distribuição inclinada por recompensas. Embora essas abordagens permitam a otimização sem acesso a amostras explícitas da distribuição ótima, elas exigem treinamento em execuções sob o modelo ajustado atual, tornando-as suscetíveis ao reforço de trajetórias subótimas que produzem recompensas insatisfatórias. Para superar esse desafio, introduzimos o TRee Search Guided TRajectory-Aware Fine-Tuning for Discrete Diffusion (TR2-D2), uma nova estrutura que otimiza trajetórias de difusão discreta guiadas por recompensas com busca em árvore para construir buffers de replay para ajuste fino consciente da trajetória. Esses buffers são gerados usando a Busca em Árvore de Monte Carlo (MCTS) e subsequentemente usados para ajustar um modelo de difusão discreta pré-treinado sob um objetivo de controle ótimo estocástico. Validamos nossa estrutura no ajuste fino de modelos de difusão de sequências biológicas com objetivos únicos e múltiplos, destacando a eficácia geral do TR2-D2 para ajuste fino confiável guiado por recompensas na geração de sequências 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.