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

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.

Current protein design methods usually optimize for a single function and struggle to balance multiple conflicting goals, such as stability and solubility. They tend to repeat the same design mistakes across iterations, and adapting to new goals requires costly retraining. We propose MOMST, which learns from past successes and failures through a memory mechanism to avoid repeating mistakes. Crucially, it can flexibly combine and balance multiple functional goals during generation without any retraining. This makes protein design faster and more adaptable, lowering the barrier for applications in synthetic biology and drug discovery.