[Paper Review] Proteina: Scaling Flow-based Protein Structure Generative Models

Tomas Geffner, Kieran Didi+9 more|ArXiv.org|Mar 2, 2025

Genetics, Bioinformatics, and Biomedical Research3 citations

TL;DR

The paper introduces ourmodel, a large-scale flow-based protein backbone generator conditioned on hierarchical fold class labels, achieving state-of-the-art performance on backbones up to 800 residues and enabling designable, diverse structures with scalable training data and novel conditioning/guidance mechanisms.

ABSTRACT

Recently, diffusion- and flow-based generative models of protein structures have emerged as a powerful tool for de novo protein design. Here, we develop Proteina, a new large-scale flow-based protein backbone generator that utilizes hierarchical fold class labels for conditioning and relies on a tailored scalable transformer architecture with up to 5x as many parameters as previous models. To meaningfully quantify performance, we introduce a new set of metrics that directly measure the distributional similarity of generated proteins with reference sets, complementing existing metrics. We further explore scaling training data to millions of synthetic protein structures and explore improved training and sampling recipes adapted to protein backbone generation. This includes fine-tuning strategies like LoRA for protein backbones, new guidance methods like classifier-free guidance and autoguidance for protein backbones, and new adjusted training objectives. Proteina achieves state-of-the-art performance on de novo protein backbone design and produces diverse and designable proteins at unprecedented length, up to 800 residues. The hierarchical conditioning offers novel control, enabling high-level secondary-structure guidance as well as low-level fold-specific generation.

Motivation & Objective

• Scale protein backbone generation with a flow-based model inspired by diffusion/flow matching.
• Incorporate hierarchical fold-class conditioning (CATH-based) for controllability over secondary structure and fold-level generation.
• Develop a scalable transformer architecture capable of handling long protein backbones (up to 800 residues).
• Introduce new distribution-level metrics to quantify how generated proteins match reference distributions (FPSD, fJSD, fS).
• Demonstrate data scaling to tens of millions of synthetic protein structures and explore fine-tuning/guidance strategies (LoRA, CFG, autoguidance) for improved designability and control.

Proposed method

• Use flow-matching to learn a vector-field v_t that transforms Gaussian noise into realistic C_alpha coordinates of protein backbones.
• Condition the model on hierarchical fold-class labels (C, A, T levels) using learnable embeddings and hierarchical dropout for unconditional and conditional generation (classifier-free guidance and autoguidance).
• Adopt a scalable non-equivariant transformer architecture to process residue and pair representations, with optional triangle layers and QK normalization enhancements.
• In sampling, derive a score s_t from v_t to form an SDE that enables stochastic sampling with a tunable noise scale gamma.
• Introduce three probabilistic metrics (FPSD, fJSD, fS) based on a fold-class predictor p_phi to assess distributional similarity between generated and reference structures.
• Train on two datasets: a 2nd-order Foldseek AFDB clustered set (F_FS) and a 21M high-quality AFDB subset (F_21M), enabling up to ~21M structures and ~35x data scale over prior work.
• Explore training strategies including LoRA-based fine-tuning, distogram auxiliary loss, self-conditioning, and fold-class conditioning dropout.

Experimental results

Research questions

• RQ1 Can a large-scale, non-equivariant flow-based model generate high-quality, designable protein backbones up to 800 residues?
• RQ2 Does hierarchical fold-class conditioning provide controllable generation over secondary structure content and fold-level features?
• RQ3 How does scaling training data to tens of millions of synthetic protein structures impact designability, diversity, and distributional similarity to reference data?
• RQ4 What are effective training and sampling recipes (e.g., LoRA, CFG, autoguidance, t-sampling) for protein backbone generation?
• RQ5 Can new distribution-level metrics (FPSD, fJSD, fS) meaningfully quantify how closely generated structures match reference distributions across fold classes?

Key findings

• uModel ourmodel achieves state-of-the-art performance in unconditional and fold-class conditioned backbone generation, including long chains up to 800 residues.
• uHierarchical fold-class conditioning enables novel control over generated structures, including fold-specific synthesis and enhanced beta-sheet content.
• uTraining on up to 21 million high-quality synthetic structures (35x larger than Foldseek AFDB subset) is feasible and improves distributional realism.
• uNew metrics FPSD, fJSD, and fS quantify distribution-level similarity and diversity/novelty across fold classes, complementing existing designability metrics.
• uThe model scales to more than 400M parameters in the transformer and demonstrates effective fine-tuning with LoRA and guidance strategies (CFG and autoguidance) for designable outputs.
• uSampling can be performed via an ODE or SDE, with a tunable noise scale gamma affecting designability/diversity trade-offs.