TR2-D2: 이산 확산 모델을 위한 트리 탐색 기반 궤적 인식 미세 조정

초록

확률적 최적 제어를 활용한 강화 학습은 확산 미세 조정을 위한 유망한 프레임워크를 제공하며, 여기서 사전 훈련된 확산 모델은 보상 기울어진 분포로 이어지는 경로를 생성하도록 최적화됩니다. 이러한 접근 방식은 최적 분포로부터의 명시적 샘플에 접근하지 않고도 최적화를 가능하게 하지만, 현재 미세 조정된 모델 하에서의 롤아웃에 대한 훈련이 필요하여 낮은 보상을 초래하는 차선의 궤적을 강화할 가능성이 있습니다. 이러한 문제를 극복하기 위해, 우리는 이산 확산을 위한 트리 탐색 기반 궤적 인식 미세 조정 프레임워크인 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.

TR2-D2: 이산 확산 모델을 위한 트리 탐색 기반 궤적 인식 미세 조정

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

초록

Support