用树搜索优化扩散模型生成路径，提升序列设计可靠性
TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion
- 引入树搜索构建回放缓冲区，避免劣质路径被强化
- 在生物序列生成中实现单/多目标优化，性能显著提升
- 适合需要可靠奖励导向生成的序列设计任务
基于随机最优控制的强化学习为扩散模型微调提供了有前景的框架，通过优化生成路径以达到奖励倾斜的分布。然而，现有方法需在当前微调模型上进行轨迹训练，易强化低质量路径导致性能下降。为此，我们提出针对离散扩散模型的树搜索引导轨迹感知微调框架TR2-D2：利用蒙特卡洛树搜索（MCTS）构建回放缓冲区，并在此基础上以随机最优控制目标微调预训练离散扩散模型。我们在生物序列扩散模型的单目标与多目标微调任务上验证了该方法的有效性，表明TR2-D2能实现更可靠的奖励导向序列生成。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。