Flash Invariant Point Attention
Andrew Liu, Axel Elaldi, Nicholas T Franklin, Nathan Russell, Gurinder S. Atwal, Yih-En Andrew Ban, Olivia Viessmann
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.
BibTeX
@inproceedings{
liu2025flash,
title={Flash Invariant Point Attention},
author={Andrew Liu and Axel Elaldi and Nicholas T Franklin and Nathan Russell and Gurinder S. Atwal and Yih-En Andrew Ban and Olivia Viessmann},
booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
year={2025},
url={https://openreview.net/forum?id=gKsG5qR3Bt}
}