Flash Invariant Point Attention

TL;DR

Invariant Point Attention (IPA) enables geometry-aware modeling of proteins and RNAs but suffers from quadratic $O(L^2)$ scaling that limits input length. FlashIPA provides a factorized reformulation that leverages FlashAttention to achieve linear memory and wall-clock scaling with sequence length, while preserving SE(3) invariance and, in many cases, improving performance over standard IPA. Key innovations include a factorized pair representation $z_{ij}=z_i^{1 op} z_j^2$, a local nearest-neighbor distogram bias, and lifting IPA components into regular attention to enable efficient computation. Empirically, FlashIPA matches or exceeds IPA in validation and enables training and generation for much longer biomolecular structures (thousands of residues and RNAs up to thousands of nucleotides) with substantial reductions in compute cost, making long-context, geometry-aware modeling more accessible and scalable. The work provides an open-source implementation at https://github.com/flagshippioneering/flash_ipa and demonstrates practical impact on backbone generative models like FoldFlow and RNA-FrameFlow.

Abstract

Invariant Point Attention (IPA) is a key algorithm for geometry-aware modeling in structural biology, central to many protein and RNA models. However, its quadratic complexity limits the input sequence length. We introduce FlashIPA, a factorized reformulation of IPA that leverages hardware-efficient FlashAttention to achieve linear scaling in GPU memory and wall-clock time with sequence length. FlashIPA matches or exceeds standard IPA performance while substantially reducing computational costs. FlashIPA extends training to previously unattainable lengths, and we demonstrate this by re-training generative models without length restrictions and generating structures of thousands of residues. FlashIPA is available at https://github.com/flagshippioneering/flash_ipa.

Figures (9)

• Figure 1: The protein backbone frame: The $C_\alpha$ is the centre. The vectors spanning $C_{\alpha}-N$ and $C\alpha-C$ define the third axis via Gram–Schmidt. The oxygen atom $O$ is parameterized via the torsion angle $\psi$ around the $C_\alpha - C$ axis.
• Figure 2: Scaling as a function of input sequence length on a single-sample batch forward pass. [A] GPU memory usage in GB. Original IPA scaled approximately quadratically with sequence length ($y\,\text{[MB]} = 2.4\times 10^{-3}\cdot L^2 + 1.4\cdot10^{-2} \cdot L$), FlashIPA follows a linear trend ($y\,\text{[MB]} = -7\cdot 10^{-12}\cdot L^2 + 7.5\cdot10^{-2} \cdot L$). [B] Wall-clock time in seconds.
• Figure 3: FoldFlow self-consistency validation after 200k optimization steps. A) The sc-RMSD of the FlashIPA (red, green) models is consistent or better than the original IPA model (blue). Extending training to larger structures with FlashIPA further improves sc-RMSD. B) Three exemplar generated backbone structures with FoldFlow FlashIPA. *ESMFold gets out of memory for lengths beyond 500+ residues and sc-RMSD could not be assessed.
• Figure 4: Scaling of FlashIPA versus IPA for RNA generation using RNA-FrameFlow model, with a number of diffusion timestep $N_T$ = 50. [A] Impact of the generated sequence length on the generation runtime, using a batch size of $1$. [B] Impact of the generated batch size on the generation runtime, for generated sequence of length $128$.
• Figure 5: RNA-FrameFlow generated RNAs results. [A] Comparison of the scTM score depending on the generated RNA sequence length and the IPA module used. The RNA-FrameFlow and RNA-FrameFlow + FlashIPA models are both trained on only short sequences $\leq 150$, while the All data model has been trained without maximum sequence lenghth limit. We don't observe a significant difference between the different trained RNA-FrameFlow models. [B] Example of short generated RNA structures with our FlashIPA RNA-FrameFlow model trained on short sequences. [C] Example of long generated RNA structured with our FlashIPA RNA-FrameFlow model trained on the full dataset.