🤖 Off-policy RL enables learning from diffusion trajectories from the non-gradient tracking policy model by storing samples in a replay buffer for repeated use.
🌳 Tree search efficiently explores high-dimensional discrete sequence spaces to find the (often sparse) subspace of high-reward sequences and leverages the structural similarities of optimal sequences to exploit optimal sampling paths in the next iteration.
We use this framework to develop an efficient discrete diffusion fine-tuning strategy that leverages Monte-Carlo Tree Search (MCTS) to curate a replay buffer of optimal trajectories combined with an off-policy control-based RL algorithm grounded in stochastic optimal control theory, yielding theoretically guaranteed convergence to the optimal distribution. 🌟
In this experiment, we fine-tune the pre-trained
unconditional peptide SMILES MDM from PepTune (
Tang et al. 2024) to optimize
multiple therapeutic properties, including target protein binding affinity, solubility, non-hemolysis, non-fouling, and permeability. We show that one-shot generation from the fine-tuned policy outperforms inference-time multi-objective guidance, marking a significant advance over prior fine-tuning methods. Code and instructions to reproduce our results are provided in
/tr2d2-pep.
If you find this repository helpful for your publications, please consider citing our paper:
1@article{tang2024tr2d2,
2 title={TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion},
3 author={Sophia Tang and Yuchen Zhu and Molei Tao and Pranam Chatterjee},
4 journal={arXiv preprint arXiv:2509.25171},
5 year={2025}
6}
To use this repository, you agree to abide by the
PepTune License.
