Multi-Objective Protein Design via Memory-Aware Test-Time Scaling in Diffusion Models

Ming Yang, Xin Zheng, Yi Li, Yizhen Zheng, Huan Yee Koh, Yanqing Guo, Xiaofeng Cao, Shirui Pan
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:147313-147333, 2026.

Abstract

Multi-objective protein design is essential for meeting the complex demands of synthetic biology. To adapt to shifting multi-functional targets without the prohibitive cost of retraining, test-time scaling has emerged as a flexible, training-free alternative. However, current test-time diffusion methods face critical challenges: i) ineffective learning from interaction history leading to repetitive design errors, ii) over-reliance on successful cases as the reward signal, and iii) difficulties in balancing multi-objective functional trade-offs . To address these limitations, we propose MoMST, a framework for Multi-objective protein design via Memory-aware Self-contrastive learning with Test-time scaling in diffusion models. At test time, we develop a memory bank to extract generalizable reasoning experience from historical iterations. Building on this powerful experience learner, we derive rich residue-level relative preference signals from both successful and failed cases via self-contrastive learning for guiding protein generation. To ensure balance among competing multi-objective functions, we present an inference-time Pareto alignment strategy to resolve objective conflicts. Evaluations on both single-objective and complex multi-objective tasks demonstrate the remarkable performance of MoMST, with code available at https://github.com/MingYangi/MoMST.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-yang26du, title = {Multi-Objective Protein Design via Memory-Aware Test-Time Scaling in Diffusion Models}, author = {Yang, Ming and Zheng, Xin and Li, Yi and Zheng, Yizhen and Koh, Huan Yee and Guo, Yanqing and Cao, Xiaofeng and Pan, Shirui}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {147313--147333}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/yang26du/yang26du.pdf}, url = {https://proceedings.mlr.press/v306/yang26du.html}, abstract = {Multi-objective protein design is essential for meeting the complex demands of synthetic biology. To adapt to shifting multi-functional targets without the prohibitive cost of retraining, test-time scaling has emerged as a flexible, training-free alternative. However, current test-time diffusion methods face critical challenges: i) ineffective learning from interaction history leading to repetitive design errors, ii) over-reliance on successful cases as the reward signal, and iii) difficulties in balancing multi-objective functional trade-offs . To address these limitations, we propose MoMST, a framework for Multi-objective protein design via Memory-aware Self-contrastive learning with Test-time scaling in diffusion models. At test time, we develop a memory bank to extract generalizable reasoning experience from historical iterations. Building on this powerful experience learner, we derive rich residue-level relative preference signals from both successful and failed cases via self-contrastive learning for guiding protein generation. To ensure balance among competing multi-objective functions, we present an inference-time Pareto alignment strategy to resolve objective conflicts. Evaluations on both single-objective and complex multi-objective tasks demonstrate the remarkable performance of MoMST, with code available at https://github.com/MingYangi/MoMST.} }
Endnote
%0 Conference Paper %T Multi-Objective Protein Design via Memory-Aware Test-Time Scaling in Diffusion Models %A Ming Yang %A Xin Zheng %A Yi Li %A Yizhen Zheng %A Huan Yee Koh %A Yanqing Guo %A Xiaofeng Cao %A Shirui Pan %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-yang26du %I PMLR %P 147313--147333 %U https://proceedings.mlr.press/v306/yang26du.html %V 306 %X Multi-objective protein design is essential for meeting the complex demands of synthetic biology. To adapt to shifting multi-functional targets without the prohibitive cost of retraining, test-time scaling has emerged as a flexible, training-free alternative. However, current test-time diffusion methods face critical challenges: i) ineffective learning from interaction history leading to repetitive design errors, ii) over-reliance on successful cases as the reward signal, and iii) difficulties in balancing multi-objective functional trade-offs . To address these limitations, we propose MoMST, a framework for Multi-objective protein design via Memory-aware Self-contrastive learning with Test-time scaling in diffusion models. At test time, we develop a memory bank to extract generalizable reasoning experience from historical iterations. Building on this powerful experience learner, we derive rich residue-level relative preference signals from both successful and failed cases via self-contrastive learning for guiding protein generation. To ensure balance among competing multi-objective functions, we present an inference-time Pareto alignment strategy to resolve objective conflicts. Evaluations on both single-objective and complex multi-objective tasks demonstrate the remarkable performance of MoMST, with code available at https://github.com/MingYangi/MoMST.
APA
Yang, M., Zheng, X., Li, Y., Zheng, Y., Koh, H.Y., Guo, Y., Cao, X. & Pan, S.. (2026). Multi-Objective Protein Design via Memory-Aware Test-Time Scaling in Diffusion Models. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:147313-147333 Available from https://proceedings.mlr.press/v306/yang26du.html.

Related Material