Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Tianyu Wu, Lin Zhu

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).

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

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 = , n = 100), while scaling to proteins exceeding 2,000 residues on a single A5000 GPU (24 GB).

Paper Structure

This paper contains 58 sections, 1 theorem, 20 equations, 7 figures, 8 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Materials and Methods
  4. Overview
  5. Notation
  6. Geometric flow matching on $SE(3)^L$
  7. Physics-informed spectral initialisation of the Mamba backbone
  8. Constraint enforcement during generation
  9. Training objective and auxiliary geometric supervision
  10. Training data, distillation, and length curriculum
  11. Sampling procedure
  12. Results
  13. Benchmark comparison on CATH backbones
  14. Geometric validity and stereochemical consistency
  15. Comparison of Computational efficiency and scalability
  16. ...and 43 more sections

Key Result

Proposition 1

The output of Algorithm alg:kinematic satisfies the following claims: Claim A (Exact Local Covalent Geometry). NeRF reconstruction yields idealized bond lengths ($L_{N-CA}, L_{CA-C}, L_{C-N}$), bond angles, and peptide planarity $\omega=180^\circ$ for backbone atoms by construction. Claim B (Exact C

Figures (7)

  • 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.
  • 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$.
  • 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.
  • ...and 2 more figures

Theorems & Definitions (1)

  • Proposition 1: Exact Local Covalent Geometry