Title: Flash Invariant Point Attention
Venue: NeurIPS 2025
Accepted: yes
Keywords: geometric deep learning

structural biology

invariant point attention
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents FlashIPA, a computationally efficient reformulation of Invariant Point Attention (IPA), which has been widely adopted in protein and RNA structure modeling. FlashIPA replaces the original quadratic attention mechanism with a FlashAttention-compatible variant that uses low-rank factorization of the pairwise attention bias term z_{ij}. This enables SE(3)-invariant attention to scale linearly in memory with sequence length, allowing modeling of long biomolecular sequences previously limited by IPA’s memory footprint.\n\nThe authors demonstrate that FlashIPA preserves geometric invariance, matches or exceeds the performance of IPA in protein and RNA generative models, and enables training and inference on sequences with thousands of residues.\n\nSTRENGTHS AND WEAKNESSES:\nSTRENGTHS\n1. The memory and compute limitations of IPA are well-known, and the need to scale structure models to longer sequences is urgent.\n2. The key technical insight (recasting IPA as a form of dot-product attention compatible with FlashAttention) is sound and well-executed. \n3. The paper includes thorough experiments on FoldFlow and RNA-FrameFlow, showing comparable or improved performance at a fraction of the compute cost.\n4. Providing a plug-and-play implementation and benchmarking setup improves reproducibility and impact.\n\nWEAKNESSES\n1. In practice, all attention terms are collapsed into a single residue embedding s_i, and the pair embedding z_{ij} is not retained across layers. The “modularity” argument for keeping pairwise and nodewise paths separate appears more architectural than essential, given that the final representation is merged anyway.\n2. It would be useful to evaluate a version of the model where the pairwise bias (z¹/z²) is simply folded into the q/k projection. This would test whether the factorized form offers any real gain over just using a single expressive attention term.\n3. Only ranks r = 1, 2 are tested. It’s unclear whether larger ranks would provide meaningful gains, or if performance saturates early. A deeper ablation could clarify this.\n4. While the paper seems to pose FlashIPA as a drop-in replacement for standard IPA, this only holds for models that do not maintain an explicit z_{ij} representation. In models like AlphaFold, where z_{ij} is evolved through the Evoformer via TriangleAttention and other operations, FlashIPA’s assumption of factorized z¹/z² vectors is not directly compatible. The paper does not address how to approximate or decompose an existing z_{ij}, nor does it evaluate the impact of such an approximation on downstream structure quality. This undermines the generalizability of the method and should be considered a notable practical limitation.\n\nQUESTIONS:\n1. Did you try collapsing the pair bias term into the q/k vectors directly and training with a single unified attention score? How does that compare in performance?\n2. Are there examples or tasks (e.g., complex assemblies or RNA tertiary structures) where evolving the pairwise state would be essential? Could your framework support that in future work?\n3. Could you provide more intuition or theoretical grounding for why the factorized z^1_i{}^\\top z^2_j is a good approximation of the original pair tensor z_{ij}? Under what assumptions does this hold?\n\nLIMITATIONS:\nYes!", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nInvariant point attention (IPA) is a variant of the attention mechanism that is used in SOTA protein and RNA structure prediction methods, allowing to apply the mechanism to 3D conformations. Unfortunately, its complexity is $O(L^2)$. The authors adapt an existing variant of attention, FlashAttention, which is $O(L)$, to IPA, resulting in a method called Flash IPA (FIPA) that exhibits good performance and reduced computational cost, at the expense of introducing some approximations and apparantly a complexity that is still $O(L^2)$.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\nIPA is an important algorithm in SOTA 3D protein and RNA structure prediction. The authors address a an important computational bottleneck that hinders training on large macromolecules, and demonstrate its efficacy using drop-in modified versions of FoldFlow (for proteins) and RNA-FrameFlow (for RNA).\n\nWeaknesses \n\nThere is some confusion wrt the final complexity of the method. In the final section, “Limitations”, the authors write “the the underlying compute cost is still $O(L^2)$ due to the softmax”. This should be clearly reflected in the abstract and introduction, as these now give the impression of a linear complexity.\n\nThe authors test FIPA on only one protein and one RNA prediction method.\n\nQUESTIONS:\nIt would benefit the article from adding one additional drop-in test case for both RNA and proteins. If this is not possible, why not?\n\nWhat are the current practical implications of the $O(L^2)$ complexity due to the softmax? \n\nHow significant is the 256 head dimension limit in practice? \n\nPage 4, line 129: What are these “inductive biases” that are being relaxed? Could you clarify the effect of the low-rank factorization? \n\nPage 5, line 164: Rank 2 is quite low. Could you elaborate on why this works well?\n\nPage 4: Could you elaborate on the use of the kNN algorithm? What algorithm was used, what is its complexity, and what are the practical and theoretical  implications?\n\nLIMITATIONS:\nLimitations are addressed in the final section (but see my first item under “Weaknesses”).", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 4, "ORIGINALITY": 2, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces FlashIPA, a novel reformulation of Invariant Point Attention (IPA) designed to overcome IPA's quadratic computational complexity, which limits its application to long protein and RNA sequences. FlashIPA achieves linear scaling in GPU memory and wall-clock time with sequence length by factorizing the IPA algorithm and leveraging hardware-efficient FlashAttention. The proposed method aims to make geometry-aware modeling more accessible and scalable for structural biology. Experimental results demonstrate that FlashIPA matches or surpasses the performance of standard IPA, significantly reduces computational costs, and enables training and generation of biomolecular structures with thousands of residues, a feat previously difficult to achieve due to memory constraints. The authors provide FlashIPA as a readily usable package.\n\nSTRENGTHS AND WEAKNESSES:\n### Strength\n\n- Integrating FlashAttention into IPA is the biggest contribution. The core strength of FlashIPA is its ability to reduce the computational complexity of IPA from quadratic (O(L2)) to effectively linear (O(L)) in terms of GPU memory and wall-clock time with sequence length. This is a major breakthrough for handling long biomolecular sequences.\n- FlashIPA not only achieves significant efficiency gains but also matches or even exceeds the performance of standard IPA in Protein/RNA de novo benchmark.\n- The linear scaling of FlashIPA allows training and generation of structures with thousands of residues, extending to lengths previously unattainable due to memory limitations of standard IPA. It provides integration tests with established protein (FoldFlow) and RNA (RNA-FrameFlow) generative models.\n\n### Weakness\n\nThis paper exhibits no major weaknesses. However, a significant omission is the lack of testing FlashIPA within AlphaFold2, which is the most prevalent and advanced platform for long-protein modeling. While *de novo* design is important, it typically does not necessitate the generation of long protein/RNA structures, and the resulting designs predominantly feature alpha-helical conformations.\n\nQUESTIONS:\nNA\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces flash invariant point attention (FlashIPA), which reformulates IPA to leverate FlashAttention and thus achieve linear scaling in GPU memory wrt sequence length (the sequences being AA sequences here, for instance). The paper is well-written, with a clear structure and addresses a clear need in the structural biology x AI community, and the package is also well-structured and documented. I am not able to well-review the software engineering aspects of the work but I can at least comment on its significance which is high.\n\nEdit: I have updated my score following the thoughtful response and review from the authors.\n\nSTRENGTHS AND WEAKNESSES:\nMajor Points\n- The \"anonymous\" linked repo was not, in fact, anonymous, but contained a bibtex for the arxiv version of this work.\n- Is it not alarming that none of the generated protein backbones are very self-consistent (i.e., sc-RMSD>4 Å), even with FlashIPA? I guess it is not clear to me why this task was chosen as the baseline for FlashIPA, because we do not have a fully accurate ground truth, but I guess it can be thought of as a lower bound? Some clarification in the paper would be good.\n- In Table 1, it would be good to indicate the cost along with the metrics, because the metrics just show worse performance without emphasizing that with FlashIPA you got those results more quickly (but, how much more quickly?).\n- Is TM score here the same as template modeling score? I would clarify this in an article for this venue.\n- It is not clear to me why in Figure 5 there is such a big difference between the scTM scores at some lengths (e.g., 120). Is not what you expect a similar scTM score across all methods (just knowing that it was obtained more quickly with scTM)? So what is the reason for the discrepancy?\n- Again, it is unclear to me in the RNA experiments whether the number of RNAs generated \"per length\" correspond to the same sequence or different sequences.\n\nMinor Points\n- Inconsistent cross-referencing to figures/tables in text (sometimes \"Figure\" sometimes \"Fig.\")\n\nQUESTIONS:\n- Are the reference/generated structures aligned for computing the sc-RMSD (e.g, experiments in Fig. 3, Fig. 5)?\n- What is the reason for the big variance in sc-RMSD (e.g., outliers), even for the same sequence length? Is it specific structural motifs (or flexible protein regions) that make some systems very inconsistent? And are the \"outliers\" for the original IPA and FlashIPA experiments the same (that is, is it the same structures that are very inconsistent and/or reliably consistent)?\n- Related to the above point, it is also unclear to me if for each sequence length, does this indicate just a single structure with that length (but many times), or are there many structures with a given length that were evaluated?\n\nLIMITATIONS:\nSee above, some things are unclear in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}