PaperlayerView exact revision on arXiv

Source-linked AI summary

Proteina: Scaling Flow-based Protein Structure Generative Models

Tomas Geffner, Kieran Didi, Zuobai Zhang, Danny Reidenbach, Zhonglin Cao, Jason Yim, Mario Geiger, Christian Dallago, Emine Kucukbenli, Arash Vahdat, Karsten Kreis

arXiv:2503.00710v1cs.LG

TL;DR

Protein structure generators face limited scale, controllability, and distribution-level evaluation. Proteina addresses these gaps with a scalable flow-based backbone model, hierarchical fold conditioning, expanded synthetic training data, and adapted evaluation, training, and sampling methods. It achieves state-of-the-art designable and diverse backbone generation, including structures up to 800 residues, while enabling secondary-structure and fold-specific control.

  • Problem

    Previous protein structure generators used limited datasets, offered little synthesis-time control, and were commonly evaluated without rigorous distribution-level metrics.

  • Method

    Proteina combines flow matching, a scalable non-equivariant transformer, hierarchical fold-class conditioning, large synthetic training data, distribution-level metrics, and adapted fine-tuning and guidance methods.

  • Results

    Proteina achieves state-of-the-art unconditional, class-conditional, and motif-scaffolding backbone generation, producing diverse and designable structures up to 800 residues.

  • Takeaways & Limitations

    Hierarchical conditioning provides control over secondary-structure content and specific folds, while scalable modeling supports designable generation at unprecedented protein lengths.

  • Takeaways & Limitations

    The proposed probabilistic metrics complement rather than replace designability, diversity, and novelty, and are not intended as model-ranking optimization targets.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Proteina addresses limited scale and controllability in protein structure generation with a large flow-based backbone model conditioned on hierarchical fold labels. It introduces distribution-level evaluation and achieves state-of-the-art generation, including diverse, designable backbones up to 800 residues.

  • Motivation: Proteina targets the limited training scale and weak synthesis control of previous unconditional protein structure generators.Previous models used at most half a million structures and generally lacked synthesis-time control.
  • Model: Proteina uses flow matching with a scalable transformer architecture and hierarchical CATH fold-class labels for controllable backbone generation.The labels support both high-level secondary-structure control and low-level fold-specific guidance.
  • Scaling: Training scales to almost 21 million high-quality synthetic protein structures, a 35× increase over previous work.The model is also scaled beyond 400M parameters.
  • Evaluation: Proteina introduces distribution-level metrics that compare generated samples with reference distributions in a nonlinear feature space and assess fold-class distributions.These metrics complement conventional diversity, novelty, and designability measures.
  • Results: Proteina achieves state-of-the-art protein backbone generation and produces diverse, designable structures up to 800 residues.The paper also explores adjusted objectives, stage-wise training, LoRA fine-tuning, classifier-free guidance, and autoguidance.

2 BACKGROUND AND RELATED WORK

The background describes flow matching as learning a vector field that transports Gaussian noise into protein-structure samples, and situates Proteina among earlier diffusion and flow-based backbone generators.

  • Flow matching: Flow matching models a probability-density path from an analytically tractable noise distribution to a data distribution over time.Proteina uses an ODE whose learned vector field transports samples from noise toward protein structures.
  • Conditional flow matching: Conditional flow matching regresses a neural vector field against an analytically tractable conditional target using intermediate samples and Monte Carlo marginalization.In expectation, this provides the same gradients as regression against the intractable marginal ground-truth field.
  • Rectified flow: Rectified flow constructs intermediate states with the linear interpolant x_t = t x_1 + (1 − t)x_0 and uses target velocity x_1 − x_0.This formulation connects noise and data samples directly.
  • Related work: Prior protein backbone generators include Chroma, RFDiffusion, FrameDiff, FoldFlow, FrameFlow, Genie, and related diffusion- or flow-based models.These approaches differ in their use of Euclidean coordinates, residue frames, or Riemannian manifold formulations.

3 PROTE´INA

Proteina scales protein backbone generation through larger datasets, hierarchical fold-class conditioning, a scalable transformer, and protein-specific training and sampling methods.

  • Scaling protein structure training data with fold classes: 21 million structures make D21M 35× larger than DFS, extending protein structure training beyond prior datasets.D21M contains 20,874,485 filtered AFDB structures, while DFS contains 588,318 cluster representatives.
  • Scaling protein structure training data with fold classes: Hierarchical CATH labels provide both secondary-structure-level control and guidance toward specific fold classes.Proteina uses the C, A, and T levels and focuses on the mainly α, mainly β, and mixed α/β classes.
  • Training objective: Proteina models Cα coordinates with flow matching, optionally adding a 64-bin pairwise-distance distogram loss and self-conditioning.The auxiliary distogram loss is attached to the pair representation and is used when that representation is updated.
  • A scalable protein structure transformer architecture: Its streamlined non-equivariant transformer builds sequence and pair representations, then processes them with conditioned, pair-biased attention layers.The architecture can optionally update the pair representation and inject fold labels and time through adaptive layer normalization.
  • Sampling: Sampling uses stochastic Langevin-style dynamics with a tunable noise scale, classifier-free guidance, and autoguidance for fold-conditioned generation.The paper reports CFG and autoguidance as new guidance schemes for protein structure generation.
  • Probabilistic metrics for protein structure generative models: The paper introduces probabilistic distribution-level metrics because standard designability, diversity, and novelty measures do not rigorously evaluate learned distributions.These metrics provide additional insights but are not intended as optimization targets for ranking models.

4 EXPERIMENTS

Proteina’s experiments evaluate unconditional, conditional, distribution-level, guided, and long-chain protein backbone generation. Across these settings, the models achieve strong designability, diversity, distribution matching, controllability, and scalability.

  • Unconditional generation: MFS outperforms all baselines in designability and diversity, while remaining competitive on novelty across reference sets.Inference-time γ provides different designability, diversity, and novelty trade-offs.
  • PDB-LoRA: LoRA fine-tuning on a small dataset of designable PDB proteins improves designability, moves distributions toward PDB, and doubles β-sheet content.Diversity decreases, while distributional distances to PDB decrease and those to AFDB increase.
  • Fold-class conditional generation and new metrics: Fold-class conditioning achieves state-of-the-art TM-Score diversity and the best FPSD, fS, and fJSD scores while maintaining very high designability.Classifier-free guidance with ω ≥1.0 enhances these effects, although novelty does not improve.
  • Autoguidance: Autoguidance boosts M21M designability by up to 70% in conditional full-distribution generation.The experiment uses an early training checkpoint as the “bad” guidance checkpoint.
  • Long chain generation: Proteina significantly outperforms all baselines beyond 300 residues and produces mostly designable, diverse backbones at length 800.Previous generators had not successfully trained on proteins up to this length.
  • C-level class guidance: Fold-class guidance dramatically increases β-sheet content, especially for “mainly β” classes, while samples remain designable and diverse.Restricting generation to specific classes causes a slight expected decrease in diversity.

5 CONCLUSIONS

The paper presents Proteina as a scalable foundation model for controllable protein backbone generation. It combines large synthetic-data training, a scalable transformer, guidance and fine-tuning methods, and distributional metrics, while emphasizing reproducibility and biosecurity caution.

  • Contributions: Proteina combines fold-class conditioning, a scalable non-equivariant transformer, and new distributional metrics for protein backbone generation.The model is scaled beyond 400M parameters and generates designable, diverse backbones up to 800 residues.
  • Results: Proteina achieves state-of-the-art performance in unconditional, class-conditional, and motif-scaffolding benchmarks.The conclusions also report classifier-free guidance, autoguidance, and LoRA-based fine-tuning in protein structure flow models.
  • Reproducibility: The authors report reproducible data processing, architecture, inference, evaluation, and baseline-comparison procedures with extensive appendix details.Dataset construction includes hierarchical CATH fold-class labels, and the processing and evaluation tools are publicly available.
  • Release: The authors provide model and code release through Proteina’s GitHub repository and project page.The cited release locations are given directly in the paper.
  • Ethics: De novo protein generation has promising applications but also potential biosecurity risks, requiring cautious application.The stated application areas include therapeutics, vaccines, antibodies, enzymes, biomaterials, and protein-biology research.

A ADDITIONAL PROTE´INA SAMPLE VISUALIZATIONS

The appendix presents additional Proteina backbone visualizations across chain lengths, conditioning modes, and motif-scaffolding examples. The shown samples are reported as designable under the paper’s definition.

  • Additional visualizations: Additional unconditional Proteina samples cover the model’s entire chain-length spectrum, while other figures show fold-class-conditioned samples.The appendix also visualizes successful motif-scaffolding examples.
  • Designability: All samples shown in the figures are designable according to the paper’s definition of designability.This statement applies to the displayed samples across the visualizations.
  • Motif scaffolding: Proteina’s motif-scaffolding capability is evaluated on the RFDiffusion benchmark using motif structure and motif-mask embeddings.Inference and centering are based on motif coordinates, and training adds a motif-structure auxiliary loss.

B.2 MOTIF-SCAFFOLDING RESULTS

Proteína is evaluated on motif-scaffolding tasks using sequence design, structure prediction, threshold-based success criteria, and clustering into unique successes. Across the benchmark, it leads overall, while the hardest individual tasks remain challenging.

  • Evaluation protocol: A backbone counts as successful when one designed sequence meets scRMSD ≤2 Å, motifRMSD ≤1 Å, pLDDT ≥70, and pAE ≤5.Eight ProteinMPNN sequences are generated per backbone and evaluated with ESMFold before applying these criteria.
  • Evaluation protocol: All successes are clustered by single-linkage hierarchical clustering at a TM-score threshold of 0.6 to obtain unique successes.
  • Benchmark-wide results: Proteína generated 2094 unique successes across the benchmark, compared with 1445 for Genie2.Each method generated 1000 backbones; the benchmark totals use clustered successes.
  • Benchmark-wide results: Proteína was the sole best method in 8 tasks, compared with 5 tasks for Genie2.
  • Task-level results: Proteína outperforms mostly on easy and medium tasks, whereas the hardest tasks with 1 or 0 successes remain challenging.

C.2 MODEL PARAMETERS, SAMPLING SPEED AND MEMORY CONSUMPTION

Proteína’s scalable architecture supports larger models while preserving practical inference efficiency. Triangle-free models provide especially large batches and fast batched generation, while the fold-conditioning analysis shows stronger performance for common classes than rare ones.

  • Model size: The smaller MsmallFS model has around 60M parameters and still outperforms most baselines, including RFDiffusion, though larger models show better diversity and novelty.
  • Sampling speed: Proteína models without triangle layers are substantially faster for single-protein generation, while triangle-layer models remain faster than RFDiffusion and Genie2.Triangle-layer models are slower than FrameFlow at all lengths and slower than Chroma at longer lengths.
  • Memory and batching: Proteína’s triangle-free models support batches up to 1.6k proteins of length 100.The models can fit very large batches despite the larger parameter counts of Proteína relative to baselines.
  • Sampling speed: Proteína benefits strongly from batched inference, enabling less than 1 second per chain for short chain lengths.The benefit is especially pronounced for models without triangle layers and for shorter sequences.
  • Scaling trade-offs: Model-size scaling does not reduce inference efficiency: Proteína supports batches as large as or larger than baselines and samples as fast as or faster.
  • Fold-conditioned generation: Fold re-classification is generally better for common classes than rare classes, and accuracy increases with classifier-free guidance weight.The increasingly fine A- and T-level class spectrum makes conditioning more challenging.

D.1 RE-CLASSIFICATION ANALYSIS OF FOLD CLASS-CONDITIONAL CHROMA SAMPLING

The Chroma comparison emphasizes that fold-conditioned generation must be evaluated against the requested labels, not only by aggregate properties of generated backbones. Re-classification analysis is the fine-grained test used for this purpose.

  • Re-classification comparison: Chroma generally performs poorly compared with Proteína in fold-class re-classification.
  • Evaluation limitation: Designability, diversity, and novelty metrics score the overall backbone set and do not test whether generated structures match their conditioning labels.
  • Evaluation limitation: Re-classification analysis specifically evaluates fold-class conditioning capabilities in a fine-grained manner.

E EQUIVARIANCE ANALYSIS

The equivariance analysis tests how Proteína’s predictions respond to rotations of noisy protein backbones. The model is approximately, but not perfectly, rotationally equivariant across noise levels.

  • Analysis setup: The analysis compares clean-sample predictions for noisy backbones and rotated versions of those backbones.Three metrics are used to characterize rotational behavior.
  • Equivariance metrics: For a perfectly equivariant model, Er(t)=0; larger Er(t) indicates less equivariance.
  • Equivariance metrics: For approximately equivariant models, Eu(t) is low and Eu(t)≤Er(t), with a small gap when the optimal alignment rotation is close to the input rotation.
  • Results: The invariant-model metric E(t) is considerably higher than the other two metrics, as expected for an approximately equivariant model.
  • Results: Proteína achieves Eu(t)≈Er(t)<0.5 Å for all t, confirming approximate rather than perfect equivariance.The authors attribute this behavior to random rotation augmentations applied to clean samples during training.
  • Relation to designability: The observed rotational-prediction RMSDs are below the 2 Å similarity threshold used to measure designability.

F ESTABLISHED METRICS: DESIGNABILITY, DIVERSITY, NOVELTY & SECONDARY STRUCTURE

Proteina’s evaluation combines established sample-level metrics with new distribution-level measures for assessing generated protein backbones. The framework covers designability, diversity, novelty, secondary-structure composition, and similarity to reference distributions.

  • Established metrics: Designability tests whether generated backbones can be matched to sequences that fold into those structures.Each backbone is paired with eight ProteinMPNN sequences, structurally predicted with ESMFold, and evaluated using RMSD.
  • Established metrics: Diversity is measured through average pairwise TM-score and cluster-based diversity, with lower TM-scores indicating more distinct designable samples.Cluster-based diversity divides the number of designable clusters by the number of designable samples using a TM-score threshold of 0.5.
  • Established metrics: Novelty measures the average maximum TM-score between each designable generated structure and a reference set, where lower values indicate greater distinctness.The reference sets are the PDB and the DFS AlphaFold DB subset.
  • Established metrics: Secondary-structure content reports normalized proportions of alpha helices, beta sheets, and coils in designable backbones.The proportions are computed as α/(α + β + c), β/(α + β + c), and c/(α + β + c).
  • Distribution-level metrics: Existing metrics do not rigorously evaluate how closely a generator matches a target distribution, motivating distribution-level evaluation.They assess individual-sample properties such as designability, diversity, and novelty, but can overlook resemblance to realistic proteins.
  • Distribution-level metrics: FPSD compares generated and reference distributions in nonlinear feature space, while fS measures fold-class quality and diversity and fJSD measures fold-class distribution similarity.FPSD uses feature distributions, fS uses a trained fold classifier, and fJSD compares predicted fold classes across distributions.

G.4 METRIC VALIDATION

Validation experiments show that the proposed distribution-level metrics respond appropriately to fold-distribution differences, class balance, noise, and unseen structures. Their behavior supports using them to assess realism, diversity, and distributional similarity.

  • fS validation: The Fold Class-Balanced Subset receives the highest fS, the Homogeneous Subset the lowest, and the Fold Class-Imbalanced Subset an intermediate score.These outcomes match the intended sensitivity to realistic fold diversity and class balance.
  • fS validation: Adding noise gradually decreases the Fold Score as declining structure quality reduces classifier confidence.The metric is evaluated at the C, A, and T hierarchy levels.
  • FPSD and fJSD validation: 452.44 FPSD and 4.21 fJSD distinguish fold class-disjoint PDB splits, whereas a random split yields approximately 10 FPSD and 0 fJSD.The class-disjoint split contains no overlap at A-level classes, while the random split is expected to follow the same distribution.
  • Validation summary: FPSD, fJSD, and fS remain robust to unseen samples while detecting degradation in noisy protein structures.The authors summarize the metrics as measuring distribution similarity, realism, diversity, and balance.

H.1 FINE-GRAINED DIVERSITY EVALUATIONS

Proteina’s fine-grained evaluations examine diversity, reference-distribution behavior, sampling dynamics, and guidance. They show competitive distribution modeling, strong long-range generation, and sensitivity to inference-time stochasticity and training choices.

  • Diversity evaluations: Proteina’s pairwise TM-score distributions remain reasonable across lengths, with no observed signs of mode collapse.The comparison includes Proteina variants, Genie2, RFDiffusion, and FrameFlow.
  • Reference comparisons: Reference datasets display higher diversity than all evaluated models and baselines, leaving room to improve diversity optimization.The reference metrics use representative PDB and AFDB subsets.
  • Full distribution modeling: Full-distribution sampling gives MFS performance on par with or better than Genie2, while M21M and LoRA fine-tuning substantially improve designability.The comparison evaluates sampling without the temperature or noise reductions commonly used to increase designability.
  • Time sampling: Sampling more heavily at large t is designed to emphasize accurate local details, while mixing in uniform sampling avoids zero density as t approaches 0.The proposed distribution uses a Beta component and is compared with uniform and logit-normal alternatives.
  • Sampling dynamics: The stochasticity schedule strongly affects final performance, with g(t) outperforming gtan(t) and g1−t(t) under numerical sampling with γ < 1.For γ = 1, the schedules have the same marginal distributions in principle, but numerical SDE simulation differentiates them in practice.

L.3 QK LAYER NORM, REGISTERS AND ROPE EMBEDDINGS

Architecture and data-processing ablations clarify the trade-off between pair-update performance and scalability, while the training pipeline uses filtered, clustered datasets and hierarchical fold labels. The resulting design supports long-backbone generation and controllable structural conditioning.

  • Architecture ablations: Pair updates with triangular multiplicative layers improve performance but reduce scalability.The Mno-triFS model remains competitive and is significantly more computationally efficient.
  • Architecture ablations: The computational efficiency of Mno-triFS enables scaling to protein backbones of up to 800 residues.This establishes scalability as the practical benefit of omitting triangular pair updates.
  • Attention components: Registers and QK layer normalization slightly improve training losses, whereas RoPE slightly increases training loss and is excluded from the final architecture.These comparisons are made using small models with directly comparable training losses.
  • Data processing: The filtered PDB dataset contains 114,076 single protein chains, and the designable subset contains 90,423 chains, or 79.26% of the original.Designability filtering uses ProteinMPNN and ESMFold with an scRMSD threshold below 2 Å.
  • Data processing: The high-quality AFDB subset contains approximately 21 million structures after length, confidence, secondary-structure, and radius-of-gyration filtering.The resulting dataset contains 20,874,485 structures and is clustered by 50% sequence similarity.
  • Hierarchical conditioning: Hierarchical CATH labels encode class, architecture, topology, and homologous superfamily, while Proteina focuses on the CAT levels for structural guidance.The H level is discarded, and the main conditioning uses the three major C classes.

O.2 UNCONDITIONAL GENERATION EXPERIMENTS

The unconditional-generation experiments compare Proteina models across sampling choices, noise scales, self-conditioning, and training checkpoints. Results include distributional metrics, designability trade-offs, and scalability to long proteins.

  • Sampling configurations: Mno-triFS remained competitive at γ = 0.45 despite omitting pair-representation updates, improving scalability.The reported run used an earlier 80k-step checkpoint because it offered better designability-diversity trade-offs than the 360k-step model.
  • Sampling configurations: M21M used γ ∈ {0.3, 0.6} to separately target extremely high designability or improved trade-offs across metrics.The lower-noise run used self-conditioning and a large dataset filtered for high-quality structures; the alternative omitted self-conditioning.
  • Sampling configurations: ODE sampling with self-conditioning produced significantly better FPSD, fS, and fJSD values for the MFS and M21M models.The authors attribute this to noise-scale effects that modify the sampled distribution.
  • Long-length generation: Long-length generation fine-tuned Mno-triFS on filtered AFDB representatives and sampled 100 proteins at each length from 300 to 800 residues.The final model used γ = 0.35 and 400 sampling steps after training with maximum lengths of 512 and then 768.
  • Long-length generation: Class-specific guidance was combined with long-length generation to control secondary-structure content in large proteins.Guidance used labels at the C level of the CATH hierarchy.
Loading 2503.00710v1…