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

2025-07-11

AI Summary

Current protein foundation models lack test-time scalability and unified capability across diverse protein design tasks. Method: We propose AMix-1, a Bayesian flow network-based model integrating MSA-driven in-context learning with an evolutionary test-time scaling algorithm, optimized according to pretraining scaling laws. Contribution/Results: AMix-1 is the first model to demonstrate progressive emergence of structural understanding capabilities *during inference*, enabling efficient generation of functional proteins directly from sequence inputs. Experimentally, it designed an AmeR variant with 50× enhanced activity. Moreover, its performance consistently improves with increased computational budget at test time, confirming strong test-time scalability. These advances establish a closed-loop paradigm for *in silico* directed evolution followed by wet-lab validation, significantly enhancing protein engineering capacity.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Search and Optimization: Mixed Discrete/Continuous Search
• Computer Vision: Diffusion Models for Vision

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Incentives in network design for Web infrastructures and ecosystems

Abstract

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 mechanism, and test-time scaling algorithm. To guarantee robust scalability, we establish a predictive scaling law and reveal 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 $50 imes$ 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 in silico directed evolution that delivers substantial, scalable performance gains as verification budgets are intensified, laying the groundwork for next-generation lab-in-the-loop protein design.

Problem

• Develops scalable protein foundation model AMix-1
• Unifies protein design via MSA-based learning framework
• Enables test-time scaling for in silico directed evolution

Innovation

• Bayesian Flow Networks for protein modeling
• MSA-based in-context learning strategy
• Evolutionary test-time scaling algorithm