TR2-D2: Baumsuchgesteuerte, trajektorienbewusste Feinabstimmung für diskrete Diffusion

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

Reinforcement Learning mit stochastischer optimaler Steuerung bietet einen vielversprechenden Rahmen für die Feinabstimmung von Diffusionsmodellen, bei dem ein vortrainiertes Diffusionsmodell optimiert wird, um Pfade zu erzeugen, die zu einer belohnungsorientierten Verteilung führen. Während diese Ansätze eine Optimierung ohne Zugriff auf explizite Stichproben aus der optimalen Verteilung ermöglichen, erfordern sie das Training auf Rollouts unter dem aktuell feinabgestimmten Modell, was sie anfällig für die Verstärkung suboptimaler Trajektorien macht, die geringe Belohnungen liefern. Um diese Herausforderung zu bewältigen, führen wir TRee Search Guided TRajectory-Aware Fine-Tuning for Discrete Diffusion (TR2-D2) ein, ein neuartiges Framework, das belohnungsgesteuerte diskrete Diffusionstrajektorien mit Baumsuche optimiert, um Replay-Puffer für trajektorienbewusste Feinabstimmung zu konstruieren. Diese Puffer werden mithilfe von Monte Carlo Tree Search (MCTS) erzeugt und anschließend verwendet, um ein vortrainiertes diskretes Diffusionsmodell unter einem stochastischen optimalen Steuerungsziel feinabzustimmen. Wir validieren unser Framework anhand der Einzel- und Mehrziel-Feinabstimmung von biologischen Sequenzdiffusionsmodellen und unterstreichen die allgemeine Effektivität von TR2-D2 für zuverlässige belohnungsgesteuerte Feinabstimmung in der diskreten Sequenzerzeugung.

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.