TR2-D2: Траекторно-ориентированная тонкая настройка с управлением поиском по дереву для дискретной диффузии

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

September 29, 2025

Авторы: Sophia Tang, Yuchen Zhu, Molei Tao, Pranam Chatterjee

cs.AI

Аннотация

Обучение с подкреплением в сочетании со стохастическим оптимальным управлением предлагает перспективную основу для тонкой настройки диффузионных моделей, где предварительно обученная диффузионная модель оптимизируется для генерации траекторий, ведущих к распределению, смещенному в сторону вознаграждения. Хотя такие подходы позволяют проводить оптимизацию без доступа к явным образцам из оптимального распределения, они требуют обучения на траекториях, сгенерированных текущей настроенной моделью, что делает их уязвимыми к закреплению субоптимальных траекторий, дающих низкие вознаграждения. Чтобы преодолеть эту проблему, мы представляем TRee Search Guided TRajectory-Aware Fine-Tuning for Discrete Diffusion (TR2-D2) — новый фреймворк, который оптимизирует траектории дискретной диффузии, направляемые вознаграждением, с использованием поиска по дереву для создания буферов воспроизведения, учитывающих траектории, для тонкой настройки. Эти буферы генерируются с помощью метода поиска по дереву Монте-Карло (MCTS) и затем используются для тонкой настройки предварительно обученной дискретной диффузионной модели в рамках задачи стохастического оптимального управления. Мы проверяем наш фреймворк на задачах одно- и многокритериальной тонкой настройки моделей диффузии биологических последовательностей, демонстрируя общую эффективность TR2-D2 для надежной тонкой настройки, направляемой вознаграждением, в задачах генерации дискретных последовательностей.

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.