PI-Mamba: Linear-Time Protein Backbone Generation via Spectrally Initialized Flow Matching

Tianyu Wu; Lin Zhu

q-bio.BMMar 17, 2026

PI-Mamba: Linear-Time Protein Backbone Generation via Spectrally Initialized Flow Matching

Authors: Tianyu Wu, Lin Zhu

Organizations: Center for Biophysics and Quantitative Biology, University of Illinois Urbana-Champaign, 600 S Mathews Ave., 61801, IL, USA · School of Information Science, University of Illinois Urbana-Champaign, 501 E Daniel St., 61820, IL, USA

Abstract

Motivation: Generative models for protein backbone design have to simultaneously ensure geometric validity, sampling efficiency, and scalability to long sequences. However, most existing approaches rely on iterative refinement, quadratic attention mechanisms, or post-hoc geometry correction, leading to a persistent trade-off between computational efficiency and structural fidelity. Results: We present Physics-Informed Mamba (PI-Mamba), a generative model that enforces exact local covalent geometry by construction while enabling linear-time inference. PI-Mamba integrates a differentiable constraint-enforcement operator into a flow-matching framework and couples it with a Mamba-based state-space architecture. To improve optimisation stability and backbone realism, we introduce a spectral initialization derived from the Rouse polymer model and an auxiliary cis-proline awareness head. Across benchmark tasks, PI-Mamba achieves 0.0% local geometry violations and high designability (scTM = $0.91\pm 0.03$ , n = 100), while scaling to proteins exceeding 2,000 residues on a single A5000 GPU (24 GB).

Figures & tables

Figure 1: Architecture of PI-Mamba for scalable protein design. The framework processes a sequence length L and a noise prior on SE(3)L through a unified generator optimized for O(L) inference. The pipeline consists of three integrated lanes: (1) Physics-informed spectral initialization , which uses the Rouse spectrum $\lambda_{p}=4sin^{2}(p\pi/2L)$ to initialize SSM decay rates $A_{p}=exp(-\lambda_{p}\Delta t/\tau(x))$ , ensuring global topology is captured by low modes and local details by high modes; (2) a Bidirectional Mamba backbone that predicts tangent velocities $\xi_{\theta}(t,T)\in SE(3)^{L}$ for flow matching, enabling iterative updates $T_{t+\Delta t}=T_{t}\cdot exp(\Delta t\xi_{\theta})$ and backbone reconstruction via NeRF; and (3) Constraint enforcement , where periodic $C_{\alpha}$ retraction $\Pi_{K}$ maintains ideal 3.80Åbond lengths and peptide planarity.

Figure 2: Training dynamics of PI-Mamba. (A) Loss components (log scale): total loss, flow matching, hydrogen-bond, and Ramachandran losses all converge stably. (B) Rouse mode concentration $\mathcal{C}_{10}$ evaluated periodically during training, measuring the fraction of hidden-state variance captured by the lowest 10 Rouse modes.

Model	Raw Viols $\downarrow$	Final Viols $\downarrow$	scTM $\uparrow$	Div $\uparrow$	Total Time $\downarrow$
FrameDiff	12.1%	0.2%	0.58	12.8 Å	72.0s
FrameFlow	3.2%	0.0%	0.70	12.2 Å	22.2s
Genie2	2.1%	0.0%	0.93	14.4 Å	60.0s
Chroma	4.5%	0.0%	0.60	17.1 Å	73.0s
RFdiffusion	8.2%	0.0%	0.97	8.9 Å	37.2s
Proteus	5.4%	0.0%	0.86	11.2 Å	13.5s

Table 1: Benchmark comparison at

L=100

with

n=100

samples. Raw Viols : % bond violations before projection/refinement. Final Viols : % after projection (PI-Mamba) or refinement (baselines). scTM : self-consistency TM-score (mean

\pm

SD when available). Div : mean pairwise RMSD (definition in Supplemental File). Total Time : end-to-end runtime including projection or refinement.

Metric	Euclidean baseline	PI-Mamba (Final)
Bond violations ( $\lvert d-3.8\rvert>0.5\,\text{\AA }$ )	76.3%	0.0%
Angle violations ( $>10^{\circ}$ )	3.2%	0.0%
Cis-proline recall ( $\omega$ agreement)	35.0%	99.2%
Trans-proline recall ( $\omega$ agreement)	92.1%	99.8%

Table 2: Geometric validity breakdown (rounded). Rates computed over

n=100

generated backbones at

L=100

Panel of Figure 3 — Figure 3: Diversity and validity of generated backbones ( $L=200$ ). (Left) Overlay of $6$ random samples rendered in PyMOL PyMOL (white/black spheres: N-/C-termini). (Right) $C_{\alpha}$ – $C_{\alpha}$ bond-length distribution. PI-Mamba exhibits a narrow peak near $3.8$ Å after projection.

Figure 4: Computational efficiency comparison on a single NVIDIA RTX A5000 (24 GB). (A) Inference time per sample as a function of sequence length. PI-Mamba scales near-linearly, while diffusion baselines exhibit superlinear growth. (B) Peak GPU memory usage. PI-Mamba remains below 0.5 GB up to $L{=}1000$ , whereas most baselines exceed 10 GB and encounter OOM failures beyond $L{\approx}500$ . Endpoints indicate the maximum length each method can handle before OOM on this hardware: FrameDiff, Genie2, and Proteus fail at $L{>}500$ ; Proteina at $L{>}300$ ; only FrameFlow and RFdiffusion reach $L{=}1000$ .

Topology	scTM	scRMSD (Å)	Recovery
All- $\alpha$	0.94	1.2	45%
All- $\beta$	0.89	1.5	39%
Mixed ( $\alpha/\beta$ )	0.88	1.3	36%

Table 3: Performance by CATH Topology Class (

L=100

n=100

Figure 5: Physics-aligned statistics emerge without explicit supervision. (A) Rouse mode amplitude spectrum of PI-Mamba hidden states follows the theoretical $1/p^{2}$ decay, confirming alignment with polymer normal modes. (B) Radius of gyration scales as $R_{g}\propto L^{0.48}$ , consistent with Rouse theory ( $\nu=0.5$ ). (C) End-to-end distance distribution at $L{=}200$ matches the Gaussian chain prediction.

Method	Representation	Seq. Recovery
RFdiffusion	Full backbone	33.4%
Proteina	Full backbone	32.1%
PI-Mamba ( $C_{\alpha}$ only)	$C_{\alpha}$ trace	23.9%
PI-Mamba + Ideal Recon.	Full backbone	31.2%

Table 4: Sequence recovery with full backbone reconstruction.

Projection frequency	scTM $\uparrow$	Comment
$k=1$ (every step)	0.61	over-constrained
$k=10$ (default)	0.91	best trade-off
$k=\infty$ (final only)	0.65	drift before projection

Table 5: Projection frequency ablation (

L=100

S=200

). Full results are reported in Appendix 15 , Table 8 .

Panel of Figure 6 — Figure 6: Gallery of Failure Modes (PyMOL Renders). (Left) Loop Collapse : The N-terminus (green) forms multiple tight loops before extending into an unstructured tail. (Right) Extended Chain : A non-compact, elongated backbone lacking any tertiary fold. Both cases arise from poor long-range coherence in the Mamba encoder.

Figure 7: Sequence recovery distribution. Mean 23.9% (red) vs random 5% (gray).

Mamba Backbone
Parameter	Value
Number of Layers	16
Model Dimension ( $d_{model}$ )	512
State Dimension ( $d_{state}$ )	32
Conv Kernel Size	4
Expansion Factor	2

Table 6: Model Hyperparameters

Symbol	Description
$L$	Sequence length (number of residues)
$T_{i}\in SE(3)$	Rigid body frame of residue $i$ (rotation $R_{i}$ and translation $\mathbf{t}_{i}$ )
$\mathcal{K}$	Constraint manifold of valid protein backbones (fixed bond lengths/angles)
$\Pi_{\mathcal{K}}$	Kinematic projection operator onto $\mathcal{K}$
$\tau_{\text{relax}}$	Rouse relaxation time constant (physics-informed)
$\tau_{\text{step}}$	Integration step size $\Delta t$ in the ODE solver

Table 7: Summary of Mathematical Notation

$k$ (project every)	$\omega$ dev ( ∘ )	CA-CA (Å)	scTM	Time (s)
1 (every step)	0.0	3.80	0.61	2.1
5	0.0	3.80	0.68	1.8
10 (default)	0.0	3.80	0.91	1.6
20	0.0	3.80	0.88	1.5
$\infty$ (final only)	0.0	3.80	0.65	1.4
No projection	85.2	3.45	0.65	1.4

Table 8: Effect of Projection Frequency on Generated Structures (

L=100

S=200

steps)

Explore similar work

SPID: Distilled Protein Backbone Generation
Oct 3, 2025Liyang Xie, Haoran Zhang, Zhendong Wang +2Diffusion Model Distillation Few-Step Diffusion Sampling
The University of Texas at Austin · The Memorial Sloan Kettering Cancer Center
Proteo-R1: Reasoning Foundation Models for De Novo Protein Design
May 1, 2026Fang Wu, Weihao Xuan, Heli Qi +26Generative Modeling 3D Molecular Generation
Stanford University
ProHiFlo: Hierarchical Flow Matching with Functional Guidance for De Novo Protein Generation
Jun 3, 2026Chuanzhen Wang, Meade Cleti, Pete JanoFlow Matching 3D Molecular Generation
3Tongji University · 1Arizona State University · University of Wisconsin-Madison

Oct 3, 2025J/K move · Enter open · S save

SPID: Distilled Protein Backbone Generation

Liyang Xie, Haoran Zhang, Zhendong Wang, Wesley Tansey+1

The University of Texas at Austin · The Memorial Sloan Kettering Cancer Center

Diffusion Model Distillation Few-Step Diffusion Sampling Molecular Generation Protein Design

Open paper AskPDF arXiv

Diffusion- and flow-based generative models have recently demonstrated strong performance in protein backbone generation tasks, offering unprecedented capabilities for de novo protein design. However, despite their generation quality, these models are constrained by slow sampling, often requiring hundreds of iterative steps. This computational bottleneck limits their practical utility in large-scale protein discovery, where thousands to millions of candidate structures are needed. To address this challenge, we explore the techniques of score distillation, which has shown great success in reducing the number of sampling steps in the vision domain while maintaining high generation quality. However, a straightforward adaptation of these methods results in unacceptably low designability. We introduce Score Protein identity Distillation (SPID), which resolves this incompatibility by combining few-step generation with inference-time noise scaling. SPID adapts the Score identity Distillation (SiD) framework to both diffusion- and flow-based models without requiring access to pretraining data. Applied to the Proteina flow-matching model, our 16-step generator achieves 94.4% designability, matching the 400-step teacher, while delivering more than a 20-fold reduction in effective backbone-generation time and maintaining comparable diversity and novelty. SPID generalizes across unconditional generation, fold-class conditional generation, and motif scaffolding, and extends to equivariant diffusion architectures, achieving significant reduction in generation time with comparable generation quality to the teacher in all tasks. The resulting reduction in inference cost could facilitate large-scale in silico protein design, thereby advancing diffusion-based models toward real-world protein engineering applications. The PyTorch implementation is available at https://github.com/LY-Xie/SiD_Protein

Abstract

Explore similar work

SPID: Distilled Protein Backbone Generation

Proteo-R1: Reasoning Foundation Models for De Novo Protein Design

ProHiFlo: Hierarchical Flow Matching with Functional Guidance for De Novo Protein Generation