TR2-D2: Ottimizzazione Fine Guidata da Ricerca ad Albero con Consapevolezza della Traiettoria per Diffusione Discreta

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

September 29, 2025

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

cs.AI

Abstract

L'apprendimento per rinforzo con controllo ottimo stocastico offre un quadro promettente per l'affinamento di modelli di diffusione, in cui un modello di diffusione pre-addestrato viene ottimizzato per generare percorsi che conducono a una distribuzione inclinata verso una ricompensa. Sebbene questi approcci consentano l'ottimizzazione senza accesso a campioni espliciti della distribuzione ottimale, richiedono l'addestramento su rollout sotto il modello attualmente affinato, rendendoli suscettibili al rafforzamento di traiettorie sub-ottimali che producono ricompense scarse. Per superare questa sfida, introduciamo TRee Search Guided TRajectory-Aware Fine-Tuning for Discrete Diffusion (TR2-D2), un nuovo framework che ottimizza le traiettorie di diffusione discreta guidate dalla ricompensa utilizzando la ricerca ad albero per costruire buffer di replay per l'affinamento consapevole delle traiettorie. Questi buffer sono generati utilizzando la Monte Carlo Tree Search (MCTS) e successivamente utilizzati per affinare un modello di diffusione discreta pre-addestrato sotto un obiettivo di controllo ottimo stocastico. Validiamo il nostro framework sull'affinamento mono e multi-obiettivo di modelli di diffusione di sequenze biologiche, evidenziando l'efficacia complessiva di TR2-D2 per un affidabile affinamento guidato dalla ricompensa nella generazione di sequenze discrete.

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.