Research map / Architectures and efficiency / Attention and transformers
Flash Invariant Point Attention
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes FlashIPA, a factorized reformulation of Invariant Point Attention that runs on FlashAttention for linear memory and time scaling, and uses it to retrain FoldFlow and RNA-FrameFlow on unrestricted lengths.
Problem
IPA, the SE(3)-invariant attention used in many protein and RNA structure models, scales quadratically in sequence length, limiting training to short chains and making long biomolecules infeasible.
Approach
Rewrites the IPA softmax argument (including geometric distance terms) as a single inner product by lifting queries/keys/values, and factorizes the pair representation into low-rank pseudo-factors, so the update can run through FlashAttention with O(L) memory and I/O.
Key results
- FlashIPA memory scales linearly with sequence length versus quadratic for original IPA, and FlashIPA allows effective batch size 39 at length 512 where IPA allows 1.
- FoldFlow with FlashIPA gives sc-RMSD consistent with or better than original IPA, and training on the full unrestricted monomer set improves sc-RMSD further.
- RNA-FrameFlow with FlashIPA on a single GPU matches original validity/novelty at roughly a quarter of the compute cost; generation at length 2048 is about 30x faster than original IPA.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace IPA with FlashIPA | RNA-FrameFlow · RNA generated structures, length ≤150 | validity (scTM): 0.42 ± 0.21 -> 0.38 ± 0.20 | No clear effect |
| Replace IPA with FlashIPA | RNA-FrameFlow · RNA generated structures, length ≤150 | diversity (qTM cluster): 0.15 -> 0.14 | No clear effect |
| Replace IPA with FlashIPA | RNA-FrameFlow · RNA generated structures, length ≤150 | novelty (pdbTM): 0.81 ± 0.10 -> 0.82 ± 0.10 | No clear effect |
| Use FlashIPA with batch size 512 on a single GPU | RNA-FrameFlow · RNA generated structures, length ≤150 | validity (scTM): 0.42 ± 0.21 -> 0.41 ± 0.21 | No clear effect |
| Use FlashIPA with batch size 512 on a single GPU | RNA-FrameFlow · RNA generated structures, length ≤150 | diversity (qTM cluster): 0.15 -> 0.08 | Hurt |
| Use FlashIPA with batch size 512 on a single GPU | RNA-FrameFlow · RNA generated structures, length ≤150 | novelty (pdbTM): 0.81 ± 0.10 -> 0.77 ± 0.09 | Helped |
| Train with FlashIPA on all data without maximum length restrictions | RNA-FrameFlow · RNA generated structures, length ≤150 | validity (scTM): 0.42 ± 0.21 -> 0.36 ± 0.14 | No clear effect |
| Train with FlashIPA on all data without maximum length restrictions | RNA-FrameFlow · RNA generated structures, length ≤150 | diversity (qTM cluster): 0.15 -> 0.14 | No clear effect |
| Train with FlashIPA on all data without maximum length restrictions | RNA-FrameFlow · RNA generated structures, length ≤150 | novelty (pdbTM): 0.81 ± 0.10 -> 0.74 ± 0.08 | Helped |
| Replace IPA with FlashIPA | RNA-FrameFlow · RNA generation, length 2048 | runtime speedup: 30 times | Helped |
| Use rank 2 factorization and pair-wise distograms with k=20 (ablation) | FlashIPA · FoldFlow and RNA-FrameFlow training convergence | loss convergence | Helped |
Surprising findings
- Single-GPU FlashIPA training matched the 4-GPU original RNA-FrameFlow on validity and novelty metrics.
Limitations
- Pair representation factorization is an approximation that may not hold for denser pair representations.
- FlashAttention head-dimension limit (256) constrains the factorization rank.
- Compute remains O(L^2) due to softmax even though memory and I/O are linear.
Experimental setup
Datasets: PDB monomers (FoldFlow pipeline) RNASolo2 (BGSU v3.382)
Benchmarks: self-consistency RMSD on generated protein backbones RNA self-consistency TM, diversity and novelty
Metrics: scRMSD scTM qTM cluster diversity pdbTM novelty GPU memory wall-clock time training loss
Baselines: Invariant Point Attention (original IPA) FoldFlow RNA-FrameFlow
Builds on: Invariant Point Attention FlashAttention FoldFlow RNA-FrameFlow
Closest papers by meaning
- FlashBias: Fast Computation of Attention with Bias (NeurIPS 2025 · similarity 0.86)
- FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (ICLR 2024 · similarity 0.84)
- Generating Highly Designable Proteins with Geometric Algebra Flow Matching (NeurIPS 2024 · similarity 0.84)
- Flash3D: Super-scaling Point Transformers through Joint Hardware-Geometry Locality (CVPR 2025 · similarity 0.83)
- Geometric Hyena Networks for Large-scale Equivariant Learning (ICML 2025 · similarity 0.83)
- Proteina: Scaling Flow-based Protein Structure Generative Models (ICLR 2025 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.