AMix-1: A Pathway to Test-Time Scalable
Protein Foundation Model

1Shanghai Artificial Intelligence Laboratory
2Generative Symbolic Intelligence Lab (GenSI), Tsinghua University
3Institute for AI Industry Research (AIR), Tsinghua University
*Full author and affiliation list in Contributions
Paper arXiv Model (Incoming) 📚 Citation

We introduce AMix-1, a powerful protein foundation model built on Bayesian Flow Networks and empowered by a systematic training methodology, encompassing pretraining scaling laws, emergent capability analysis, in-context learning strategy, and test-time scaling algorithm. To guarantee robust scalability, we establish a predictive, noise-aware scaling law and reveals the progressive emergence of structural understanding via loss perspective, culminating in a strong 1.7-billion model. Building on this foundation, we devise a multiple sequence alignment (MSA) based in-context learning strategy to unify protein design into a general framework, where AMix-1 recognizes deep evolutionary signals among MSAs and consistently generates structurally and functionally coherent proteins. This framework enables the successful design of a dramatically improved AmeR variant with an up to 50x activity increase over its wild type. Pushing the boundaries of protein engineering, we further empower AMix-1 with an evolutionary test-time scaling algorithm for \textit{in silico} direction evolution that delivers substantial, scalable performance gains as verification budgets are intensified, laying the groundwork for next-generation lab-in-the-loop protein design.

Scaling Law

We are the first to characterize scaling law and emergent ability for an advanced generative model, i.e. Bayesian Flow Network, and pretrain a family of powerful models, AMix-1, scaling from 8 million to 1.7 billion parameters. Following figures illustrate scaling laws of cross-entropy and training FLOPs across various model sizes and the noise level \(\alpha\). We show the equation between \(\mathcal{L}\) and FLOPs \(F\), indicating that the AMix-1 models have predictive scaling laws.

Emergent Ability

Emergent abilities of AMix-1 models under varying noise levels \(\alpha\). Noise levels of \(\alpha=0.16\) and \(0.32\) exhibit clear emergence. Darker colors indicate models with more parameters. Lower cross-entropy, and consequently stronger downstream performance, emerges as the model size increases .

In-Context Learning

We establish a unified protein design framework leveraging AMix-1's in-context learning ability and evolutionary signals in multiple sequence alignments for novel structure- and function-preserving generation. Following figure illustrate the comparison of standard LLM in-context learning (top) with protein family conditioning via MSA profiles in AMix-1 (bottom). In LLM ICL, few-shot exemplars are token-level demonstrations; in AMix-1, the MSA profile serves as a dense, position-wise family prompt.

We design a high-activity AmeR variant with a \(50\times\) improvement over the wild type, showcasing AMix-1's real-world impact on functional protein engineering. Specifically, we used AMix-1 to generate 40 candidate protein sequences, each containing no more than ten mutations. The functional activity of these variants was measured using a fluorescent reporter assay, where higher DNA-binding affinity leads to stronger repression of a fluorescent protein, quantified as "fold repression." A higher fold repression value indicates better function. AMix-1's performance was compared to three baselines: Wild-Type (WT), a single-mutant library, and EvoAI, a state-of-the-art directed evolution method. The illustration on the right shows the wet-lab verification process guided iteratively by AMix-1.

Test-Time Scaling

With the in-context learning ability of AMix-1, we push this capacity even further: introducing a novel evolutionary test-time scaling algorithm that uses AMix-1 as the proposer and an external plug-and-play verifier to generate candidate sequences from evolutionary constraints, achieving directed protein evolution through arbitrary in-silico verifiers or experimental screening.

Citation

If you find this work useful, please cite our paper:


@article{lv2025amix1,
  title={AMix-1: A Pathway to Test-Time Scalable Protein Foundation Model},
  author={Changze Lv*, Jiang Zhou*, Siyu Long*, Lihao Wang, Jiangtao Feng, Dongyu Xue, Yu Pei, Hao Wang, Zherui Zhang, Yuchen Cai, Zhiqiang Gao, Ziyuan Ma, Jiakai Hu, Chaochen Gao, Jingjing Gong, Yuxuan Song, Shuyi Zhang, Xiaoqing Zheng, Deyi Xiong, Lei Bai, Ya-Qin Zhang, Wei-Ying Ma, Bowen Zhou, Hao Zhou},
  journal={arXiv preprint arXiv:2507.08920},
  year={2025}
}