Flash Invariant Point Attention
Authors: Andrew Liu
Axel Elaldi
Nicholas T Franklin
Nathan Russell
Gurinder S. Atwal
Yih-En Andrew Ban
Olivia Viessmann
Venue: NeurIPS 2025 spotlight ()
Keywords: geometric deep learning
structural biology
invariant point attention
Primary area: machine_learning_for_sciences
TL;DR: We reduce IPA, one of the most popular geometric deep learning methods in structural biology, from quadratic to linear scaling without performance loss, and train on previously non-trainable data.

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.

Flash Invariant Point Attention
Andrew Liu, Axel Elaldi, Nicholas T Franklin, Nathan Russell
Gurinder Atwal, Yih-En Andrew Ban, Olivia Viessmann
Flagship Pioneering
Cambridge, MA, United States
Janliu,aeladd, infrnnnl in, nrussell,
matwal, aratalll aban, , oviessmann] @flaeaaannnn@llgagaaaaa ort.eessaannlllggggggppppeeeeeen
Abstract
Invariant Point Attention (IPA) is 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, factorized reformulation of IPA that leverages hardware-efficient FlashAttention 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- re-training generative models without length restrictions and generating structures of thousands of residues. FlashIPA is available at hhttts:// / gighu....... com httm//lllahippoooornnnn.. Flashiipa
Introduction
Invariant Point Attention, IPA, a geometry-aware attention operation that has been the workhorse of generative structural design models for proteins and RNA, initially popularized by Alphafold2 [ and AlphaFold-Multime [2, and subsequently widely adopted across structural biology modeling.
Amongst them structure prediction models like OpenFold and ESMFold for proteins [
RhoFold for RNA [5]. generative protein backbone models such as FrameDiff( FrameFlow
the FoldFlow family 99. [0), FrameDiPT (11), Proteus [12), FADiff [13], Genie 14 IgDiff [5
GAFL [16. P2DFlow 17 RNA generative models such as RNA-Frameelow
18 and scoring
models like lociPARSE
19 A list of models that rely IPA for their structure modeling provided
in Appendix A. The advantage of IPA is its roto-trans oototrranslatioall (SE(3) invariant representation of the molecular structure, enforeing the idea that rotations and translations of molecule results in an equivalent structure prediction. This inductive geometric bias accelerates training and improves performance in limited data settings, is the case for structurally resolved biomolecules. IPA's quadratic scaling (O(L²)) limits its scalability, rapidly exhausting GPU memory when modeling longer biomolecules. As of May 2025, 42% of structures within the PDB have more than 512 residues, and 33% with more than 756 and 23% with more than 1024 (see Appendix A.2). Most trainings across the literature reduce their data to chains and cropped structures below 512 residues. Despite this, they still commonly run on costly multi-GPU setups with trainings that span from days up to month 55
Such engineering compromises, namely truncating biologically relevant lengths, limiting dataset sizes, and relying on expensive computational infrastructure, restrict the progress of the field. In this work, we propose recasting of the original IPA algorithm to a simple attention form to leverage off-the shelf I/O reduction methods like FlashAttention, which replace the quadratic ((2)) with an effective linear O(L) scaling behavior. We show empirically that FlashIPA exceeds the validation performance of IPA benchmarking models and datasets. We then demonstrate the memory and
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

compute time efficiency by retraining benchmarking models more efficiently and without length restrictions on the data.
We provide FlashIPA as an importable uy package at https://github. com/ [laasshippooneering hatg/flash_ipppp with an API similar to existing repositories using IPA to facilitate drop-in usage.
2 Preliminaries
Introduced by AlphaFold2 M1 IPA is a specialized attention mechanism designed to preserve 3D geometric relationships directly in transformer-based models for structural biology applications. IPA enforces invariance under 3D rotations and translations by utilizing learnable local coordinate frame representation for each residue (point) within proteins and RNAs.
Formally, a function f x + y is invariant to transformations in the abstract group G if f(g z) fa): for all G. This differs from equivariance to G, where the transformations of the group commute with f. or f(g > f [x) for all E G. Here, we are concerned with the transformations T (referred to "frames") in the special Euclidean group SE(3), continuous group of rigid transformations that includes rotations and translations, but not reflections. Intuitively, the output of an SEB3) function does change under rotations and translations, which is desirable for biomolecular structures, that do not change function under those transformations.
2.1 Frame representations for proteins and RNA
Meaningful frames for biomolecules are commonly defined on the backbone structures. proteins, each residue frame is typically determined by three of the four atoms from the backbone, namely the alpha carbon (C,), nitrogen (N), and carbon (C), with C. at the origin. The frame transformation T maps the standard positions of these backbone atoms into their actual positions in global coordinates: IN,C Ca. 0] = T [N', C',, C, O'I. The frame has a rotation and translation component T = (R, t). The rotation component of the transformation. R, is computed using Fram-Schmide ortho- normalization on vectors formed by C.
and Ca - C' bonds, while the position of the oxygen (O) is captured by the torsion angle around the Ca C' bond.
For RNA, the backbone structure involves 13 atoms per nucleotide, leading to more complex flexibility. RNA-FrameFlow [18], define RNA frames using the atoms C3', C'',, and O4', identified as the least variable backbone positions. The rest of the RNA backbone geometry can then be parameterized efficiently through set of eight torsion angles.
2.2 IPA algorithm
The IPA algorithm is 1D Sequence-to-sequee transformation. It transforms an input sequence representation s; local frames T,, and pair representations 25j (with sequence indices i, 1, ..., L) into an output sequence 8, Each attention head, indexed by / E 1,.., H. involves linear transforma- tions of the input sequence to produce scalar queries, keys, values (q) k4, vi R9) and Euclidean (geometric) counterparts (h hhp vhp E R²)) Additionally, the pair representation Ziy E glxLxd, contributes to bias term by through a linear projection:
Single rep: q 1
hhp kh. The v v.a <- Linear (8:)
Pair rep. (bias term): by + Linear (zis)
(2)
at = softmax, wL (
IPA update:
(Voat Tkg + bij
Σ
T o 0t - T, K.ip
2
F
O(12)
(3)
Attention aggregation: o%
Σ al of =
===
hhp
T 0 Σ a (T hpp
0

---

Concat, project, return: S
Linear
concata.p
o
| (2))
(5)
The proof of SE(3) invariance of the IPA transformation can be found in the appendix of AlphaFold2 [. This form of IPA incurs (TT)) I/O and memory complexity due to the explicit materialization of the quadratic attention matrix in equation
2.3 FlashAttention and 1/O reduction
The quadratic scaling of the attention operation is a well-known challenge in deep learning. Operations are generally 1/O-bound on the GPU, with performance primarily limited by data transfers between the GPU's high-bandwidth memory (HBM) and its static random access memory (SRAM). I/O reduction is typically achieved through kernel fusion, which reduces memory traffic by combining multiple operations into a single CUDA kernel. FlashAttention [20] 21 achieves kernel fusion via an online and tiled computation of the softmax, which, instead of materializing the full quadratic M softmax (qk)) matrix, performs an equivalent computation by accumulating partial contributions to the output Y MV one tile at time. This technique improves scalability and performance for transformer-bant based models. We provide the Flash-Attention pseudo-code in appendix A.3 Building on Fltentnttn FlashAttention-22 further reduced the number of non-matmul FLOPs, increased parallelism across thread blocks, and distributed work between warps to reduce communication through shared memory (21].
2.4 FlashIPA: combining geometry-awarenes and efficiency
Our efficiency gains come from expressing the entire IPA update in form that is amenable to FlashAttention (i.e. a factorized version of form qTk). which then computes the attention update without materializing any quadratic attention matrix. To do so, rewrite the softmax argument in the IPA update (eq. (3) as a single inner product. We also need to parameterize the pair representation (bias term eq.
(2) Zi E RLxLxd, in a factorized form Zij = z Zz]. where z/, z, E R 1xxxdd Here, can be interpreted as the "rank" of the factorization, and is the dimension on which we perform the contraction. We then expand the sum of squared norms in the third term of eq. (3) and collect terms, resulting in equivalent update rule, given in Algorithm [ Note that factorizing Zij allows not only for combining the softmax terms, but also allows us to
compute the pair representation part of the attention as 0t = E, KiyZiy = z.!T [E, aizzz) avoiding materializing any quadratic object. As result, the attention components from eq. (1) are lifted to
q k e 5.geSr v E 1e++ trd, Regular attention is then computed on the lifted components.
With the update in the form of regular attention, then leverage FlashAttention perform the computation. This leads to significant speedups in wall-clock time and since the online computation of FlashAttention avoids materializing the quadratic attention matrix, we only need store the queries, keys, and values in the GPU's HBM, incurring only (L) memory.
2.5 Factorizing the pair representation
General forms of pair representations Zij typically involve quadratic memory complexity and are not inherently decomposable. However, common implementations of IPA often compute these quadratic representations from components that intrinsically scale linearly, implying redundancy in the resulting quadratic tensor Zıy (eq. (2)) This suggests that pair representations can be efficiently approximated through low-rank factorization in latent space. In particular, common pairwise tensors in structural biology modeling, such as distance matrices and contact maps, are either low-rank or sparse due to the smooth, local nature of physical interactions like electrostatics and steric constraints. This motivates a simplified factorization strategy for initializing and updating Ziy- Instead of explicitly constructing and storing full quadratic tensor of shape (B, L, L, d), we represent each feature (e.g.. positional encodings, distograms, diffusion time embeddings in generative models, a.s.o.) using two lower-dimensional tensors (pseudo-factors) each of shape [B, I., r, d). obtained via linear projections. We show empirically that this parametrization, despite relaxing certain inductive biases, recovers downstream performance with substantially lower memory consumption.
3

---

Algorithm FlashIPA with factorized pair representations
Given input sequence s E Rixdim, factorized pair representation Znj z 22 (where z x? E Rlxxxd, and
R
Lxxxx,, frames T, (R,, t,), Vi,j (1,...,.]
I: q khp
q: k/ k.y
vh, v.a Linear (s;)
2: bbl Linear
3: by² Linear
z
4: q < concat
qf [T; o kpp
Nows
{ [i qha
Noua
Nouny buu
p=1
2 p=1
5: k + concat
ve k (MMWwwCT kipp {pp
Newy
wy 2
we
[i o k. hp 2
Now
p=1
Nous
2
p=1
Nuass
6: v + concat
(Tr T p=1 2
7: 6 + FlashAttention( Rh, Vh);;
8: o O o + s
hp
9: o
T: T
10: 0 f z; to
11: 6, < Linear concata. of : oh 0. ))
Return 5
In many structure models zıj is updated multiple times by a non-linear Edge EdgeTransitton module that is not part of IPA. Because of that each factor is usually highly non-linear with respect to the input, despite the low-dimensional contraction. Thus, we still retain substantial representation power with this factorization. The benefits of highly nonlinear low-dimensional factors is well motivated in the literature: Instead of large head dimensions and linear factors, i.c. q, k Linear (X), as is the case for transformers, many state space models, such as Mamba [22] use factors that are non- linear transformations of the input (e.g. g. Swish (Conv ld (X)), and much smaller contraction dimensions (often r = 16).
We also note that most implementations of IPA rely pair representations Z, that depend on distograms computed from pairwise distances. To simultaneously ensure invariance and avoid computing activations the full distogram, we only keep the distances of the & nearest neighbors (B, I,, k, Whins) instead of (B, L, D. In practice, we compute these using a full L L distance matrix during pre-processing. which does not impact the runtime complexity of FlashIPA. We keep track of neighbor identities by applying positional encodings on their indices. Emphasizing nearest neighbor distances is also motivated by the fact that local geometry (c.g. chemical bonds) is more important than global geometry when it comes to validity of biomolecular structures. We surmise that in future work it may be beneficial compress all (LL1)) distances into factors without materializing the pairwise matrix, in a manner similar to FlashAttention.
3 Experiments
3.1 FlashIPA is SE(3) invariant and scales linearly in memory and wall-clock time
We tested the SEvaaaiie by applying random roto-translations input Gaussian point clouds and measuring the output deviation after applying single layer of the original IPA and FlashIPA. Original IPA output deviation was < 10-6 and FlashIPA 10-3, We evaluate the scaling behavior in GPU memory and wall-clock time by passing ingle-sample batches of increasing lengths L through original and FlashIPA, see Fig. [ We use a polynomial (green and red dotted lines), and find an approximate GPU memory scaling for FlashIPA of [MB] = -7 10 12 L + 7.5 10 2 L versus original IPA [MB] 2.4 10-3 3 1,2 + 1.4 10 2. L. The computational complexity of attention still remains O(I)) ) due to the softmax operation. FlashAttention reduces the memory complexity to OLL). Practically I/O, rather than computation, often dominates runtime on GPU hardware, and this is reflected in the observed linear wall-clock scaling in Fig. 2A

---

[A] Flash IPA memory usage
[B] Flash IPA runtime
100
Auseoon lrst
10
OrigiA
GPU Memory [GB]
GPU run time [s]
0.1
FlashIPA OriginA
0.001
120
512
2048
0197
12x
1314
524k
128
512
2048
8192
32k
131k
524k
Sequence length [L]
Sequence length [L]
Figure 1: Scaling as 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 [MB]= 2.4
10- 3 L 1.4 10 L), FlashIPA follows linear trend (y [MB]
-7 10 12 L 7.5 10 2 L). [B] Wall-clock time in seconds.
3.2
Integration test with external repositories
We selected two recent model approaches where code and data were available for retraining of the original model and allowed for FlashIPA insertion. We chose FoldFlow [10) for proteins and RNA- FrameFlow [ for RNAs, both are flow-matching generative backbone models. All experiments were run on L40S GPU instances with 48 GB HBM memory. FlashIPA hyperparameters were matched to the IPA parameters chosen by the original authors (embedding sizes, hidden dimensions, number of heads, etc.). For the factorization of the pair representation we tested rank E 21, and found rank and pair-wise distograms with = 20 to match and partially surpass the loss convergence of original IPA.
3.3 FlashIPA improves performance and extends to larger proteins
We retrained the FoldFlow Base model in its original form with IPA and with FlashIPA. We reran the PDB data pre-processing pipeline by the authors, which yielded a total of 40,492 single-chain protein monomers for training. The original training used maximum length cut-off of 512 residues, which results in a 10% reduction of the training data to 36,600 structures. We match the training strategy of the original authors and train on this reduced dataset on 4 GPUs with DDP. We kept model and train parameters identical to the published values of the authors, however for the FlashIPA version we had to make two adjustments: The original repository runs 4 blocks of IPA with a hidden dimension of 256. FlashAttention becomes incompatible with DDP at that dimension, so we reduced FlashIPA hidden dimension to 128 and used blocks (instead of 4) to match model parameter sizes (17.4M versus 17.1M, theirs versus ours). To keep memory consumption at an efficient constant level, the original implementation heuristically chose an effective batch size according to the quadratic rule eff,,- max [round 5500.00 neeris) ) ,]]. which kept GPU memory at approximately 90% throughout training. For FlashIPA we were able to achieve linear effective batch size effe, max [round (20.000 <RGPU> 1] that resulted in similar memory consumption. For example IPA only has batch size of 1 for a protein of length 512, whereas FlashIPA can batch together 39 samples. We found that FlashIPA converged slightly faster than IPA as expected with larger batch size, but we also compared partial loss terms, and found that particularly the number of steric clashes reduced faster for FlashIPA (we provide loss curves in the Appendix A.4). This suggests that our local k-nearest neighbour distogram emphasis is more efficient at capturing local geometry. To demonstrate the immediate enablement of linear scaling we ran one additional training with FlashIPA on the entire monomer dataset without length restrictions. The largest chain in the dataset 8.8k residues. We follow the original authors' validation test that generates structures varying lengths, subsequently inverse folds with Protein-MPNN 23 and forward-folds with ESMFold [4] Generated and forward-folded structures are then aligned via the Kabsch algorithm and the resulting self-consistency RMSD Sc-RMDS) is proxy for generation fidelity. We perform this test for all three models: original IPA and FlashIPA trained up to length 512, and FlashIPA on data. Similar to the original paper we sampled 50 structures per length from 100 to 500, and sample 8 inverse folded sequences to forward fold per strucutre. All models are evaluated on checkpoint number 200,000. Fig. 2 A) shows that the Sc-RMSD is lower for FlashIPA compared to the original IPA
N
5

---

[A] Self consistency RMSD of generated protein backhones
[B] Geuerated protein backbous structures with FlashIPA
oripnna FlashIPA, data
acRMSD N.L."
Eength 1on <cOSSD 19
152
250 203 Sequence length
Length a 000n
Figure 2: FoldFlow self-consistency validation after 200k optimization steps. A) The sc-RMSD of the FlashIPA (red, green) models is consistent 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.
trained on the same length-restricted data. Sc-RMSD further improved when FlashIPA training is extended to the full dataset, as expected. We emphasize that FlashIPA is not designed to improve generation quality per se, but to enable training on larger and more complex structures. We suspect that the observed Sc-RMSD values reflect accumulated model errors across three stages of validation: FoldFlow generation, Protein-MPNN inverse folding, and ESMFold forward folding. We could not assess Sc-RMSF for larger structures, as the ESMFold also depends IPA and ran out of memory. Overall we conclude that FlashIPA is more performant and more efficient than IPA. Fig. 2 B) shows example structures generated with the full data trained FlashIPA version.
3.4 FlashIPA trains more efficiently and extends to larger RNAS
We retrained the RNA MMA-FrameFloow model using the code provided by the original authors. We ingested the same BGSU version 3.382 of the RNASolo2 dataset (24), comprising a total of 14.995 structures (see Appendix A.5) First we re-trained RNA ANA-FrameFlow with original and FlashIPA on all structures within 40 150 residue (6,030 structures total), as proposed by the authors. We kept all hyperparameters consistent with the original model. We matched the authors' training setup and models were trained on GPUs with DDP and an effective batch size of 4 28. 4 GPUs are necessary for effective training with IPA, but with FlashIPA it becomes feasible to run comparable training on a single GPU instance. To demonstrate the cost efficiency we also performed an additional FlashIPA training run on the same data with a batch size of 512 on single GPU. All models were trained for fixed compute time of hours, which resulted in approximately 156k iterations with original IPA, roughly 230k with FlashIPA on 4GPUs and 194k with FlashIPA single GPU. During training we found models to converge comparably with overlapping loss curves, see Appendix A.6 For an poples-to-apples comparison we validated the models using identical validity, novelty, and diversity metrics as defined in the original paper, which uses gRNAde 25 for unverse-ooddin of the generated structure and RhoFold (26] for horward-folding For validity, report the average self-consic self-consistency template modeling score (sc- TM). To do so, we generated 50 structures per length ranging from 40 to 150 residues, in increments of 10. The results are presented in Table [ and Fig. 4 We observe comparable scores between all models. In particular, the single-GPU training matches performance quarter of the compute cost than the original.
To highlight the advantages of using FlashIPA in the GPU runtime required to generate RNA backbones based on RNA sequence length and batch size. The results are in Fig. 3 We observed improvements with FlashIPA compared to the original IPA model, especially for longer sequences and larger batch sizes. For example it takes 30 times longer to generate an RNA sequence of length 2048 with the original IPA than with FlashIPA.
model, we analyzed the
We also re- re-trained ANA-FFameFlow with FlashIPA on the full RNASolo dataset without maximum length restrictions (the longest structures being 4,417 nucleotides) and filtering out structures shorter than 40 nucleotides. We trained the model on 4 GPUs for 48 hours, with batch size of 4 x 28, resulting in approximately 310k iterations. In Fig. 4 we present samples of large generated RNA structures with lengths of 2000 and 4000 nucleotides. In comparison, the original ANA-FrameFloy can not be

---

Model
Validity (1)
Diversity (1)
Novelty (U)
Checkpoint
Cost
(x<TM i std)
(qTM cluster)
(pdbTM sid)
(#steps)
(# GPU hours)
RNA-FrameFlow
0.42 1 0.21
0.15
0.81 + 0.10
156k
80
FlashIPA (ours)
0.38 0.20
0.14
0.82 + 0.10
230k
80
FlashIPA single GPU ** (ours)
0.41 0.21
0.08
0.77 0.09
194k
+ FlashIPA data (ours)
0.36 = 0.14
0.14
0.74 + 0.08
310k
192
Table 1:
Average validity, diversity, and novelty scores for 600 generated RNAs of length - 150 with with and without FlashIPA. FlashIPA provides competitive results at fraction of the cost an original IPA layer. ** indicates that batch size has been increased to match GPU memory capacities.
[A] RNA generation Sequence length effect runtime
[B] RNA generation
Batch size effect
124
criae
Orignn
GPU run time [s]
GPU run time
Sequence length
Batch size
Figure 3: Scaling of FlashIPA versus IPA for RNA generation using ANA-FrameFlow model, with a number of diffusion timestep N7 = 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.
trained or generate such long sequences. The inverse-forwarrd-olldngg consistency test cannot be run for structures of thousands of residues the validation models (gRNAde 25 for inverse-folding and RhoFold [26] for forward-foldee folding) run out of memory at such lengths. However, we qualitatively observe that the generated structures resemble RNA, we don't necessary suspect those be valid, as the amount of training data at such length is still comparably small. We note that sc-TM scores vary across sequence lengths, with some fluctuations (c.g.. at 120 residues). This pattern is consistent with the original RNA- ANA-FrameFlow paper, which attributes it to length imbalance in the training distribution and biases introduced by the RhoFold structure predictor used during evaluation (see [18 Appendix A.3).
Discussion
FlashIPA provides memory and wall-clock efficient SE(3) invariant training on biomolecules and allows training on structures of thousands of residues. This approach opens up the possibility of multi-chain complex modeling in situations constrained by the original IPA module. Our work permit usage of IPA in contexts where it may have been ruled out due to its scaling limitations, including representations beyond polymeric backbone that include more atomic detail.
Efficiency and scalability long contexts has often been neglected aspect of geometric deep learning research and many modules have quadratic to cubic complexity. An example is triangular attention used popular models like AF-3 Chai- [ and Boltz-1 1 229 Recently, TriFast [30] was released, which uses fused triangular attention kernels to reduce I/O complexity from cubic to quadratic. We assume sub-quadratic models tend not to incorporate geometric inductive biases. Our work can be seen as first step at bridging this divide, enabling models that respect invariances, are cost-effective, and can scale to larger more fine-gras finengrained biomolecular systems.
factorized version might even achieve a linear scaling. Most efficient,

---

[A] Self-consistency TM of generated RNA
[B] Example of small generated RNA
function of sequence length
XNA FrandFlow
self-eonsisteney TM
single
(oun)
Length aTTM
Length sclM cun..
Lengh scidd -
Sequence length [1.]
[C] Example of large generated RNA
Length
Length 4000
Figure 4: 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 only short sequences 150, while the All data model has been trained without maximum sequence lenghth limit. We don't 't observe a significant difference between the different trained RNA-FrameFlow models. [B] Example of short generated RNA structures with our FlashIPA RNA-FrameFF model trained on short sequences. [C] Example of long generated RNA structured with FlashIPA RNA-FrameFlow model trained on the full dataset.
4.1 Limitations
FlashIPA relies on factorization of the pair representation, which ultimately is approximation. Here we did not find a decrease in performance, but this may not be guaranteed for more general applications to other forms of more dense pair representations. Exploring alternative schemes specialized for certain pair representations (e.g. cross-comcatenatioon, distograms, pair index differences) may further improve downstream performance.
Despite substantial memory and I/O savings, current FlashAttention implementations (e.g.. FlashAt- tention2) impose maximum head dimension of 256, due to the kernel's 's internal design, which expects stacked query/key/value vectors. Since our method relies on augmenting the attention components, this means that we must have max(c + 5N query + rdz, + 3.N Here, we achieved competitive results using head dimensions around 200, so this was not a practical limitation. However, the Triton implementation of FlashAttention2 based AMD CDNA (MI200, MI300) and RDNA GPUs allows for extending the head dimension to arbitrary size. We expect these improvements to apply more broadly to wider class of GPU architectures in the months to come. Unlocking larger head dimension would allow us to increase our factorization rank, thereby better approximating dense pair representations and closing potential gaps against quadratic IPA.
rdz) < 256.
Furthermore, we note that despite achieving O(L) in I/O and memory, the underlying compute cost is still ( due to the softmax. Removing the softmax and using linear attention variants such as Mamba would allow us achieve (T.) cost both compute and memory, and is thus promising avenue of future work.

---

References
[1] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool. Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583 2021.
[2] Richard Evans, Michael O'Neill, Alexander Pritzel, Natasha Antropova, Andrew Senior, Tim Green, Augustin Zidek, Russ Bates, Sam Blackwell, Jason Yim, et al. Protein complex prediction with alphafold multimer. biorxiv, pages 2021-10, 2021.
[3] Gustaf Ahdritz, Nazim Bouatta, Christina Floristean, Sachin Kadyan, Qinghui Xia, William Gerecke, Timothy J O' O'Donnell, Daniel Berenberg, Ian Fisk, Niccolò Zanichelli, al. Open- fold: Retraining alphafold2 yields new insights into its learning mechanisms and capacity for generalization. Nature Methods, 21(8):1514-1222 2024.
[4] Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Allan dos Santos Costa, Maryam Fazel-Zarandi, Tom Sercu, Sal Candido, et al. Language models of protein sequences at the scale evolution enable accurate structure prediction. BioRxiv, 2022:500902, 2022.
[5] Tao Shen, Zhihang Hu, Siqi Sun, Di Liu, Felix Wong, Jiuming Wang, Jiayang Chen, Yixuan Wang, Liang Hong, Jin Xiao, et al. Accurate ma 3d structure prediction using a language model-based deep learning approach. Nature Methods, pages 1-12, 2024.
[6] Jason Yim, Brian L. Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. Se (3) diffusion model with application to protein backbone generation. arXiv preprint arXiv:2302.02277 2023.
[7] Jason Yim, Andrew Campbell, Emile Mathieu, Andrew Y. K. Foong, Michael Gastegger, Jose Jimenez-Luna, Sarah Lewis, Victor Garcia Satorras, Bastiaan S. Veeling, Frank Noe, Regina Barzilay, and Tommi Jaakkola. Improved notif-scaffoldin with SE(3) flow matching. Transactions Machine Learning Research, 2024.
[8] Jason Yim, Andrew Campbell, Andrew YK Foong, Michael Gastegger, José Jiménez-Luna Sarah Lewis, Victor Garcia Satorras, Bastiaan S Veeling, Regina Barzilay, Tommi Jaakkola, et al. Fast protein backbone generation with (3) flow matching. arXiv preprint arXiv:2310.05299 2023.
[9] Guillaume Huguet, James Vuckovic, Kilian Fatras, Eric Thibodeau-Laufer, Pablo Lemos, Riashat Islam, Cheng-Hao Liu, Jarrid Rector-Brooks Tara Akhound-Sadegue Michael Bronstein, et al. Sequence-augmented se (3)-flow matching for conditional protein backbone generation. Advances in neural information processing systems, 2024.
Jarrid [10] Joey Bose, Tara Schound-Sadegh Guillaume Huguet, Kilian FATRAS, Jarrid Rector-Brooks, Cheng-Hao Liu, Andrei Cristian Nica, Maksym Korablyov, Michael M Bronstein, and Alexander Tong. Se (3)-stcchastic flow matching for protein backbone generation. In The Twelfth International Conference on Learning Representations, 2024.
[11] Cheng Zhang, Adam Leach, Thomas Makkink, Miguel Arbesú, Ibtissem Kadri, Daniel Luo, Liron Mizrahi, Sabrine Krichen, Maren Lang, Andrey Tovchigrechko, Nicolas Lopez Carranza, Uğur Şahin, Karim Beguir, Michael Rooney, and Yunguan Fu. Framedipt: Se(3) diffusion model for protein structure inpainting. bioRxiv, 2023.
[12] Chentong Wang, Yannan Qu, Zhangzhi Peng. Yukai Wang, Hongli Zhu, Dachuan Chen, and Longxing Cao. Proteus: exploring protein structure generation for enhanced designability and efficiency. bioRxiv, pages 2024 02, 2024.
[13] Ke Liu, Weian Mao, Shuaike Shen, Xiaoran Jiao, Zheng Sun. Hao Chen, and Chunhua Shen. Floating anchor diffusion model for multi-motif scaffolding. 2024.
[14] Yeqing Lin and Mohammed AlQuraishi. Generating novel, designable, and diverse protein structures by equivariantly diffusing oriented residue clouds, arXiv preprint arXiv:2301. /2485, 2023.

---

[15] Daniel Cutting, Frédéric A Dreyer. David Errington. Constantin Schneider, and Charlotte M Deane. De novo antibody design with se (3) diffusion. Journal of Computational Biology, 2024.
[16] Simon Wagner, Leif Scute, Vsevolod Viliuga, Nicolas Wolf, Frauke Gräter, and Jan Stühmer. Generating highly designable proteins with geometric algebra flow matching. In Thirty-eighth Conference on Neural Information Processing Systems, 2024.
[17] Yaowei Jin, Qi Huang. Ziyang Song, Mingyue Zheng, Dan Teng, and Qian Shi. P2dflow: A protein ensemble generative model with se(3) flow matching. Journal of Chemical Theory and Computation, 2025.
[18] Rishabh Anand, Chaitanya K. Joshi, Alex Morehead, Arian R. Jamasb, Charles Harris, Simon Mathis, Kieran Didi, Bryan Hooi, and Pietro Liò. Rna-frameflow: Flow matching for de novo 3d ma backbone design. arXiv preprint 1rXiv:2406.18833 2024.
[19] Sumit Tarafder and Debswapna Bhattacharya. lociparse: a locality- invariant point attention model for scoring ma 3d structures. Journal of Chemical Information and Modeling. 64(22):8655 8664, 2024.
[20] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory--fficient exact attention with io-awareness. Advances in neural information processing systems, 35:16344-16359, 2022.
[21] Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. International Conference Learning Representations, 2024.
[22] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00772, 2023.
[23] Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert Ragotte, Lukas Milles, Basile IM Wicky, Alexis Courbet, Rob J de Haas, Neville Bethel, et al. Robust deep learning- based protein sequence design using proteinmpnn. Science, 378(6615):49 2022.
[24] Bartosz Adamczyk, Maciej Antezak, and Marta Szachniuk. Rnasolo: repository of cleaned pdb-derived ma 3d structures. Bioinformatics, 38 (6(1))3668-33666 -6700, 2022.
[25] Chaitanya K Joshi, Arian Jamasb, Ramon Viñas, Charles Harris, Simon V Mathis, Alex Morehead, Rishabh Anand, and Pietro Liò. gRNAde: Geometric deep learning for 3d ma inverse design. In International Conference on Learning Representations (ICLR), 2025.
[26] Tao Shen, Zhihang Hu, Zhangzhi Peng, Jiayang Chen, Peng Xiong, Liang Hong, Liangzhen Zheng, Yixuan Wang, Irwin King, Sheng Wang, et al. E2efold-3d: End-to-end deep learning method for accurate de novo ma 3d structure prediction. arXiv preprint #rXiv:2207.0188, 2022.
[27] Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630(8016):4933-555) 490-5000 2024.
[28] Chai Discovery team, Jacques Boitreaud, Jack Dent, Matthew McPartion, Joshua Meier, Vinicius Reis, Alex Rogozhonikov. and Kevin Wu. Chai-1: Decoding the molecular interactions of life. BioRxiv, pages 2024-10, 2024.
[29] Jeremy Wohlwend, Gabriele Corso, Saro Passaro, Mateo Reveiz, Ken Leidal, Wojtek Swiderski, Tally Portnoi, Itamar Chinn, Jacob Silterra, Tommi Jaakkola, et al. Boltz-1: Democratizing biomolecular interaction modeling. bioRxiv, pages 2024 11, 2024.
[30] Liam Atkinson. Trifast. htts://gghuu..co//aakkn////
2025.
[31] Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Gener- ative flows on discrete state-spaces, Enabling multimodal flows with applications to protein co-design. arXiv preprint arXiv:2402.04999, 2024.
[32] Aleksa Gordic. Eli5: Flashattention. http://getttiiles..mcmmuummmmmms www.://eass.mmedimooooooo// 5c44017022ad, 2023.
10

---

A Appendix
A.1 List of biomolecular design models with IPA structure modules
Table 2: A list of models that based IPA structure modeling. Models identified via GitHub advanced search, searching for commonalities in code implementations.
Model name
Modeling approach
Training data
Train length cut-offs
OpenFold
protein structure prodiction
130k single protein chains (PDB)
<256, fine-tuning <384<5122<5110
FrameDiff X
diffusion for backhone generation
20k monomers (PDB)
60-512
PrameDiPT M
diffusion for backbone inpainting
32k monomers (PDB)
60-512
Proseus
diffusion for backhone generation
single chains from oligo-
60-512
and monomers (PDB)
FADiI
diffusion multi motif hackhone generation
to FrameDiff'
60-512
Genie
diffusion for backbone generation
monomers (PDB)
60-512
Protein
DiDi
Frame finetuning for antibody backbone generation
Synthetic antibody structures from folding
Not reported
FrameFlow
X
flow-matching backbone generation and motif scallolding
Ideatical to FrameDitf
60-512
FoldFlow 10
Stochastic matching generation
22. 22.2 monomers (PDR)
60-512
16
Geometric algehra matchhing hackhone generation
to FrameDiff'
60.512
Multillow
continuous flow matching for joint sequence backbome generation flow matching backbone ensemble generation
monomers
60-384
P2DFlow
simolation
Not reported
trajectories from ATLAS
Rho-Fold
3
RNA stracture prediction
5.5k chains (PDB)
<1024
RNA FrameFlow
flow matching for backhone generation
5.3k original stractures 1.lk
40-150
RNA
B
cropped structures (PDB)
lcciPARSE
19
Locality aware IPA structure scoring
52k synthetic structures from
200
1.4k sequence targets
A.2
Residue counts of macromolecules in the PDB
Distribution macromolcculare residue counts the PDB
30000
42.2%
Number of entries
25000
20000
32.7%
15000
23.4%
10000
5000
250
350
450
500
BI
nd
1sso
1,00
1200
1ND0
AUD
1500
150
1700
1600
1910
185
289
380
409
300
880
1995
1100
1200
1800
1900
You
Residue count
Figure 5: Distribution of protein residue counts of the proteins resolved in the PDB (Figure reproduced and adapted from the data at RCSB PDB Statistics: Sequence length distribution; http://www...sborrrttadddddiiiis cetpsssrrbbbbioossssuuuoooooo accessed May 2025).
A.3 FlashAttention Algorithm
Pseudo-code for the FlashAttention kernel, an online tiled computation of the softmax, which, instead of materializing the full quadratic M = softmax (qkT) matrix, performs an equivalent computation
by accumulating partial contributions to the output Y MV one tile at time 20 21 introduction with line-by-ling explanations can also be found here (3)..
good
11

---

Algorithm 2 Flash Attention-2 via online softmax (Dao 2024)
Given Q,K, V E Rxdd Te. T. row blocks of size B,, Te column blocks of size B.-
1: for 1 to T, do
2:
Load Q from HBM SRAM
3:
Og = 0B, xd. tg = B, mo = -00B.
4:
for j = to T. do
5:
Load K,. v, from HBM to SRAM.
On chip, compute ( Q.Ky
6:
(1-1)
rowmaxx (s! pp()) = exp (s - -m
7:
On chip, compute m; max
m
1
em -m (0-1) + rowsum
(p!i)
On chip, compute O = e
(1-1) -mi (5)
0.--11 + .
8:
diag
9: end for
10: On chip, compute O diag (1.4.))
0
11:
Write O, to HBM as ith block of O
12:
end for
13: Return O
A.4
FoldFlow training convergence
We provide loss curves for the training of the FoldFlow base model with original and FlashIPA below. FlashIPA converged slightly faster for the same number of optimization steps, which is expected from the bigger effective batch size, but we also found in particular local loss terms, such as the steric clash loss to decrease noticably more efficient.
Total Loss
Steric clash mean loss
Backbone atom loss
orignal IPA
FlashIPA
Figure Loss behaviour for FoldFlow model training.
A.5
Nucleotide residue counts of RNA structures in the RNASolo2 dataset.
Distribution of RNA nucleotides counts in the RNASolo2 dataset Only for sequence with more than 40 nucleotides
Number of entries
3000
30%
RNA Flow training cut-off L<150
100
19
200
400
1801
and
so0 1000
2100
xted Aeta
1000 2,00
1300 1488
14le-u Aidd
1600 Fipe
1000
100
200
200
s0s
neb
603
1800
1sep
198
Nucleotide count
Figure 7: Distribution of nucleotide residue counts of the RNA in the RNASolo2 dataset, filtering out the short structures of length < 40 nucleotides. The training cut-off of RNA Flow discards all sequences of length 150, accounting for 30% of the dataset.
12

---

A.6 RNA-FrameFlow training convergence
We provide loss curves for the training of the RNA-FrameFlow base model with original and FlashIPA with and without FlashIPA
below. At similar hyperparameters and hardware, behaves similarly.
RNA-FrameFlow Training loss
RNA PrameFlow
Training loss
HiA-FrrmeHoow FisshIPA (ours)
12k
20c
12k
524
60
78
Relative time step
Figure Loss behaviour for RNA-FrameFlow model training.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction summarize the scaling laws and downstream performance that we demonstrate in the main text. We provide experiments and benchmarks to validate the claims. We perform invariance tests and run a scaling study to demonstrate our claim of linear scaling with FlashIPA. We apply the proposed method popular models in generative biomoelcular design, namely a protein and an RNA backbone generative model, to demonstrate the impact.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We clearly
any assumptions, for example clearly state that the factorization approximate, we provide experiments on the downstream performance to assess if there any effects of the approximations made. We include comprehensive Limitations section where we address limitations of FlashIPA, including restrictions from FlashAttention and opportunities for improvement.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We provide a thorough walk-through our proposed algorithm and state all assumptions. The proof of invariance is referred from the AlphaFold2 paper, which provides full proof in their appendix.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yss]
Justification: We describe our experiments, including hyperparameters and hardware, as well the data used for training. We provide the FlashIPA code via anonymized github link. The repository provides a test script as well.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide an anonymised link github repository with the FlashIPA code and test script. The data and code we used for re-training is reverenced, as it is adopted from other author's repositories. The data we used for training is provided by other sources, which referene as well.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Training hyperparameters and hardware are described in the experiments section. In many cases we adopted hyperparameter choices from external authors' models, and reference their publication.
14

---

7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We compare invariance enrors between IPA and FlashIPA. For all other results, there are no statistical error bars does not apply to the nature of the measure.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We state the instance type and GPU specifications and not how many GPUs were used for how many hours/days.
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
No human subjects were used the study. Training was performed on publicly available academic datasets. The entire manuscript is anonymised, with no link reference to the authors. The github link is anonymised too.
Justification:
Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA].
There are no anticipated negative or positive societal impacts.
Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
There no risk of misuse our knowledge.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All references are cited throughout the text. We did not use any work with restricted licenses.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA].
No new assest are introduced.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well details about compensation (if any)?
Answer: [NA].
15

---

No crowdsoureing with human subjects was used.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country institution) were obtained?
Answer:[Na NA]
No human subjects involved in this research.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality the research, declaration not required.
Answer: [NA]
Justification: LLMs only used for editing.
16


Reviews:
[{"review_id": "AV3Rv9uF24", "reviewer": "Reviewer_szby", "summary": "This 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.", "questions": "- 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?", "limitations": "See above, some things are unclear in the paper.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Major 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.\")", "was_revised": true, "final_justification": "I have updated my score following the thoughtful response and review from the authors from a 4 to a 5.", "rebuttal": "### Major points\n1. The linked repo contained a bibtex for an arxiv version.  \nApologies, we removed the bibtex information in the linked repo.\n2. ***a) Is it not alarming that none of the generated protein backbones are very self-consistent (i.e., sc-RMSD>4 Å), even with FlashIPA?***  \nFlashIPA does indeed produce self-consistent structures with sc-RMSDs in the range of 1.5–2.5 Å for lengths 100–150 (see Fig. 3A), well below the 4 Å threshold mentioned. But the reviewer is correct, this level does not hold throughout all length ranges, especially once lengths increase (less training samples).   \n***b) Why was this task chosen?***  \nThis task was chosen to match the benchmark used by FoldFlow, which is a widely adopted protein backbone generation model. Our goal was to ensure a direct, apples-to-apples comparison. We have amended the manuscript to clarify this with the following *“For an apples-to-apples comparison we validated the models using identical validity, novelty, and diversity metrics as defined in the original paper, which uses gRNAde [25] for inverse-folding of the generated structure and RhoFold [26] for forward-folding.”*  \n***c) There is no ground truth for generation, is sc-RMSD in Figure 3A an indicator of accuracy lower bound?***  \nFlashIPA is not a generative model itself, but an efficiency framework that enables scaling up training—both in data volume and structure length. This scaling translates into improved sc-RMSD performance when applied within the FoldFlow → protein-MPNN → ESMFold validation pipeline: from ~6–22 Å with the original IPA training set to 1.5–13 Å with the full dataset enabled by FlashIPA (see Fig. 3A, purple vs. green box plots). We believe the observed performance reflects accumulated errors/variability across the three evaluation model stages rather than a lower bound on achievable accuracy.  \n***d) ”Some clarification in the paper would be good.”***  \nWe take the point for more clarity and added this paragraph:\n*\"We emphasize that FlashIPA is not designed to improve generation quality per se, but to enable training on larger and more complex structures. We suspect that the observed sc-RMSD values reflect accumulated model errors across three stages of validation: FoldFlow generation, Protein-MPNN inverse folding, and ESMFold forward folding.”*  \n\n3. ***In Table 1, it would be good to indicate the cost along with the metrics***  \nVery useful feedback, we added two columns to Table 1 indicating the training step at evaluation and the number of GPU training hours (cost). Now it is clear how similar results are obtained at a quarter of the cost.  \n\n4. ***Clarify TM-score***   \nWe added details to the manuscript    \n\n5. ***What is the reason for the scTM discrepancy across models in Figure 5?***   \n RNA-FramFlow is a sequence-agnostic generative model and samples a backbone for a given length. We follow the evaluation protocol of the original paper, which samples 50 structures per length. Additional variability in generated structures across models is expected because the models were not trained with matched optimization steps, but under fixed 20-hour wall-clock budgets to reflect equivalent computational cost. As a result, FlashIPA models underwent more training iterations (156k for original IPA vs. 230k and 194k for FlashIPA variants), likely affecting generalization behavior.  \nThe original RNA-FrameFlow paper also reported length variability in TM scores due to length distribution imbalances in the training set and length-dependent biases in the RhoFold structure predictor:  \n*”We believe the over representation of certain lengths in the training distribution causes the fluctuation of TM-scores. We can also partially attribute this to the inherent length bias of RhoFold; see Appendix A.3. With better structure predictors, we expect more samples to be valid.”  \n***b) what explains the discrepancy across models at some lengths, e.g. 120?***  \nOur model trained with FlashIPA and all data (Figure 5A, green box) sees a more diverse training distribution than the original model that had to exclude data above 150 residues. We assume that the FlashIPA model has reduced overfitting effects seen at certain lengths (e.g., 120 residues). \nTo clarify and avoid confusion we added this paragraph:  \n*”We note that sc-TM scores vary across sequence lengths, with some fluctuations (e.g., at 120 residues). This pattern is consistent with the original RNA-FrameFlow paper, which attributes it to a length imbalance in the training distribution and biases introduced by the RhoFold structure predictor used during evaluation (see [18] Appendix A.3).”*\n\n6. ***It is unclear in the RNA experiments whether the number of RNAs generated per length correspond to the same sequence or different sequences***  \nThe generated RNAs per length do not correspond to the same sequence. RNA-FrameFlow is a generative, sequence-agnostic model that samples diverse backbone structures conditioned only on length. For evaluation, sequences are inferred post hoc via gRNAde (inverse folding), and structures are then forward-folded with RhoFold, described in the manuscript in section 3.4.\n\n### Minor points\n1. ***sometimes you use Figures and sometimes Fig.***  \nThanks, fixed.\n\n### Questions\n1. ***Are the reference/generated structures aligned for computing the sc-RMSD?***  \nYes, we added this detail to the manuscript:\n”Generated and forward-folded structures are then aligned via the Kabsch algorithm…”\n\n2. ***What is the reason for the variance in sc-RMSD even for the same sequence length? Are the outliers for the original IPA and FlashIPA the same?***   \nPlease see our answer to the Major Points, #5b. We have not explicitly investigated whether outlier structures are consistent across IPA and FlashIPA, but this would be an interesting avenue for future analysis. However, we’d suspect that this will likely be the case, as it is highly likely an artifact of the training data distribution and not the algorithmic approach to the method. \n\n3. ***It is also unclear to me if for each sequence length, does this indicate just a single structure with that length (but many times)?***  \nPlease see our previous answers to Major Points #6. Both FoldFlow and RNAFrameFlow are generative, sequence-agnostic methods. For each target sequence length, we generate 50 distinct backbone structures, each sampled independenly."}, {"review_id": "YjqfrZoJgX", "reviewer": "Reviewer_Hu1L", "summary": "This 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.", "questions": "NA", "limitations": "NA", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### 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.", "was_revised": true, "final_justification": "I appreciate auhtors' responses and raise my score accordingly.", "rebuttal": "### Summary \nThank you for the positive feedback. Besides the AF2 retraining, given the comment ***“this paper has no major weaknesses”*** and your enthusiasm for the significant cost reduction and readily usable package — could you clarify what improvements would raise this review from “borderline accept” to “accept” or “strong accept”?\n\n### Weaknesses  \n1. ***A significant omission is the lack of testing FlashIPA within AlphaFold2, the most prevalent platform for long-protein modeling***  \nWe completely agree—evaluating FlashIPA within AF2 would have been ideal. Unfortunately, the AF2 training code is not publicly available. We found that the repo for OpenFold, an open source replication of AF2, has several problems we were unable to resolve, including: outdated PyTorch Lightning dependencies, import errors, and difficulties around template generation. Despite opening issues on the OF GitHub, we were unable to resolve them. For apples-to-apples comparison on loss convergence behaviour we would have also needed to retrain AF2 for one to two epochs at least. In practice this requires massive compute resources (≥128 GPUs), which further makes such experiments infeasible with the original IPA implementation. Our hope is that FlashIPA will make it possible for more researchers to train these models on affordable hardware in the future.  \n\n### Paper Formatting Concerns\n1.  ***The font deviates from the NeurIPS templates***  \nThanks, we used an outdated template and have updated it, which fixed the font :)"}, {"review_id": "RWfyJfMK8g", "reviewer": "Reviewer_wkTi", "summary": "Invariant 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)$.", "questions": "It 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?", "limitations": "Limitations are addressed in the final section (but see my first item under “Weaknesses”).", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 4, "originality": 2, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "The authors provided satisfactory answers to all my concerns, notably about the complexity of the final method.", "rebuttal": "### Questions\n1. ***Can you add additional drop-in test cases for both RNA and proteins?***\nOur original plan was to test FlashIPA with OpenFold, but we ran into several issues with the repository and were unable to run training (see response to Reviewer Hu1L). We thoroughly evaluated nearly every model listed in Appendix Table 2 for retraining feasibility, which required not only complete training code and configuration files but also access to the exact training and validation data. FoldFlow and RNA-FrameFlow were the models we found to be fully reproducible end-to-end (we note that FoldFlow2 was not reproducible in our hands due to inconsistent code). If more repositories provided fully reproducible setups—including original datasets and exact training configs—we would have included a third benchmark. We’ve made FlashIPA publicly available and have already seen strong community interest and early integration efforts, so we’re optimistic that additional validations will follow from the community.\n\n2. ***There is some confusion wrt the final complexity of the method. What are the practical implications of the $O(L^2)$ complexity due to the softmax? Isn't the model's true complexity still $O(L^2)$***   \nGreat point, it's well taken—we agree that the distinction between **computational** and **I/O** complexity could have been made more clear in the text. There are different notions of complexity: compute, I/O, and memory. Due to the softmax, the **computational** complexity of attention remains fundamentally **$O(L^2)$**. However, FlashAttention reduces the **I/O and memory** complexity from $O(L^2)$ to **$O(L)$** through kernel fusion. I/O is usually the dominant bottleneck in practice, and leads to effective wall-clock time improvements even though computation remains quadratic.\nWe’ve reviewed the manuscript to clarify this distinction. The abstract already correctly states that FlashIPA achieves “linear scaling in GPU memory and wall-clock time,\". This is supported both experimentally in Figure 2 and explicitly discussed in Section 3.1.\nWe clarify these points and amended this paragraph:  \n*“We use a polynomial fit (green and red dotted lines), and find an approximate GPU memory scaling for FlashIPA of $y [MB] = −7 · 10^{−12} · L^2 + 7.5 · 10^{−2} · L $ versus original IPA $y [MB] = 2.4 × 10^{−3} · L^2 + 1.4 · 10^{−2} · L$. The computational complexity of attention still remains $O(L^2)$ due to the softmax operation, but FlashAttention reduces the memory complexity to $O(L)$. Practically I/O, rather than computation, often dominates runtime on GPU hardware, and this is reflected in the observed linear wall-clock scaling in Fig. 2A”*  \nBy contrast, linear attention variants (which eliminate the softmax) can achieve true $O(L)$ complexity across compute, memory, and I/O by exploiting matrix multiplication associativity. Extending these ideas to support IPA-style geometric invariance would be a promising direction for future work.\n\n3. ***How significant is the 256 head dimension limit in practice?***  \nIn our experiments, we obtained competitive results with a head dimension of ~200, so the 256 limit was not a practical constraint. This limitation arose only because we used the existing FlashAttention-2 kernel as-is, which required stacking multiple vectors. In principle, this stacking isn’t essential—linear I/O and memory complexity can still be achieved as long as the pair representation is factorized. A custom kernel that avoids stacking could remove this limit entirely, and we may explore that in future. We clarify this point in the manuscript and added:  \n*”Despite substantial memory and I/O savings, current FlashAttention implementations (e.g., FlashAttention2) impose a maximum head dimension of 256, due to the kernel's internal design, which expects stacked query/key/value vectors. Since our method relies on augmenting the attention components, this means that we must have $\\max(c+5N_{query}+rd_z, c+3N_{value}+rd_z)\\leq 256$.  Here, we achieved competitive results using head dimensions around 200, so this was not a practical limitation.”*\n\n4. ***Page 4, line 129: What are these “inductive biases” that are being relaxed? Could you clarify the effect of the low-rank factorization?***  \nWe call the factorization of the pair representation a relaxation of inductive bias because the original 2D matrix corresponds to highly interpretable pairwise interactions. Representing them as compressed factors loses this interpretability. \nWe aren’t exactly directly factorizing the 2D matrix, in the sense that the two 1D sequences, when contracted, aren’t intended to numerically approximate the original 2D matrix. Instead, we are representing/compressing the 2D matrix using two low-dimensional 1D sequences. Such compression is justified when there is substantial redundancy or sparsity in the original 2D matrix, which is often the case.\n\n5. ***Page 5, line 164: Rank 2 is quite low. Why does this work well?***  \nFor a tensor of size $(B,L,r,d)$, you can think of the “effective inner dimension” as really being $r\\times d$. Thus, even with small $r$, the effective inner dimension can still be quite large. We really only introduced the new $r$ dimension so that we can contract the tensors once and still be left with another $d$ dimension to work with.\n\n6. ***Page 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?***  \nGreat question. In our implementation we used the simplest form of kNN where we actually materialized the $L\\times L$ distance matrix on-the-fly. This is done only once as a preprocessing step and is thus considered separate from the FlashIPA algorithm complexity. Also, if you precomputed these knn indices and stored them in the dataset you would have a memory complexity of $O(L\\cdot k)$.\nYou could also compute the $k$ nearest indices in a block-wise streaming fashion. This should be very easy to add (we might just add this in the code).\nAlternatively, you can also try to use approximate kNN methods, such as those that use hashing (LSH) or FAISS, which will achieve on-the-fly memory complexity of $O(L\\cdot k)$. But again, since the k-nearest neighbors can be precomputed beforehand to achieve $O(L\\cdot k)$ memory, the exact method used for computing them isn’t very relevant to our algorithm.  \nWe amended the manuscript to clarify this point:  \n*”…, we only keep the distances of the $k$ nearest neighbors $(B,L,k,n_{bins})$ instead of $(B,L,L,n_{bins})$. In practice, we compute these using a full $L \\times L$ distance matrix during pre-processing, which does not impact the runtime complexity of FlashIPA.”*\n\n### Paper Formatting Concers\n1. ***You only use 8 pages, move A3 (FlashAttention background) to the main text***   \nFair point. We moved A3 FlashAttention explanations up into the main text, we left the pseudo-code for the online tiled computation of the softmax in the appendix as it not our original work and not strictly necessary for the reader to fully grasp."}, {"review_id": "l2WYWco7zs", "reviewer": "Reviewer_N3zQ", "summary": "This 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.", "questions": "1. 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?", "limitations": "Yes!", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "STRENGTHS\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.", "was_revised": true, "final_justification": "In light of the authors’ clarifications, I have revised my score to an acceptance.", "rebuttal": "### Weaknesses/Questions\n1. ***Only ranks $r=1,2$ are tested. It's unclear whether larger ranks would provide meaningful gains.***  \nLower ranks $(r = 1, 2)$ have lower memory usage, and as wHigher ranks start to reduce the memory efficiency gains. We were able to recover original IPA convergence with $r=2$ (see Appendix A4 and A6), we do not expect further computational gains with higher rank approximations. Nonetheless, we will add look into running a full $rank = 3$ training so we can add it to the convergence plots in the appendix for completeness. \n\n2. ***FlashIPA does not apply to 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^1/z^2$ vectors is not 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.***  \nWe agree with the reviewer — these questions are of particular interest and are the subject of our ongoing research. Alphafold switched from using IPA in the AF-2 Evoformer to using Triangular Attention and Multiplication in AF-3. TriangleAttention and the TriangularMultiplication algorithm are distinct algorithms to IPA and they require a different set of kernel fusion approaches, that cannot be covered by a FlashAttention drop-in. Thus, we consider triangular attention out of scope for this paper, but hope to provide an answer to this question in a future paper.  \nHowever, we disagree that this undermines the generalizability of FlashIPA, because IPA and triangular attention and multiplication are two very different modeling approaches ranging from their invariance properties to the explicit algorithm. We are optimistic that lessons from FlashIPA will be useful improving the efficiency of Triangle attention.\n\n3. ***Did you try collapsing the pair bias term into the q/k vectors directly and training with a single unified attention score?***  \nIt's an interesting comment, we didn’t test this. We do expect it to perform similarly since it follows the same intuition of leveraging the compressibility of the pair representation.\n\n4. ***Are there examples or tasks (e.g., complex assemblies or RNA tertiary structures) where evolving the pairwise state would be essential? Will your framework support this?***  \nThat's a good question. In structural biology most interactions can be effectively modeled with sparse or low-rank pairwise representations. In large macromolecular complexes or self-assembling systems where long-range interactions in sequence space can be spatially proximal. Examples could be viral capsids, ribosomes, chromosomal assemblies, and RNA tertiary structures with complex motifs like pseudoknots. For these phenomena we'd still expect the distance/contact maps to be reasonably sparse, but that is a guesstimate and truly needs to be tested. Outside structural biology, many problems, eg in social or communication networks may require dense pairwise modeling as well, though they typically lack $SE(3)$ symmetry and fall outside the scope of our method.\n\n5. ***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?***  \nThe factorized form $\\mathbf{z}_{ij} \\approx \\mathbf{z}^1_i{}^\\top \\mathbf{z}^2_j$ is a good approximation when the original pairwise tensor is compressible, i.e., sparse, low-rank, or redundant due to being derived from underlying 1D sequences. More specifically, this factorized form is the optimal rank-k linear approximation (Eckart-Young-Mirsky theorem) and underlies a number of common techniques in scientific settings. In machine learning, a similar approximation underlies effective techniques such as LORA and matrix factorization in recommender systems.  \nProtein and RNA structures are governed by local and smoothly varying interacting forces (sterics and electrostatics) that are intrinsically low dimensional and make the pairwise representations highly compressible and motivate the use of low-rank approximations.   \nWe amended the text to clarify:\n*”This suggests that pair representations can be efficiently approximated through low-rank factorization in latent space. In particular, common pairwise tensors in structural biology modeling, such as distance matrices and contact maps, are either low-rank or sparse due to the smooth, local nature of physical interactions like electrostatics and steric constraints. This motivates a simplified factorization strategy… “*"}]

Author rebuttal:

Meta-review: The submission presents FlashIPA, a factorized reformulation of Invariant Point Attention (IPA) that leverages hardware-efficient FlashAttention to achieve linear scaling in GPU memory and wall-clock time with sequence length. The IPA layers offer invariance under 3D rotations and translations and have been widely used in AlphaFold2 and many other protein neural networks. The proposed FlashAttention allows training models on longer proteins and RNAs more efficiently. The reviewers have unanimously voted for the acceptance of the paper. Given the significance of the contributions, I am happy to support the recommendation.
Decision: Accept (spotlight)