Multi-Objective Protein Design via Memory-Aware Test-Time Scaling in Diffusion Models
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 that MoMST exhibits remarkable performance.