Flash Invariant Point Attention

TL;DR

FlashIPA is a new, efficient reformulation of Invariant Point Attention that uses FlashAttention to achieve linear scaling, enabling training on longer sequences and generating larger biological structures.

Contribution

We introduce FlashIPA, a hardware-efficient reformulation of IPA that reduces computational costs and enables training on longer sequences in structural biology models.

Findings

FlashIPA achieves linear scaling in GPU memory and time.

It matches or exceeds standard IPA performance.

It enables training on sequences with thousands of residues.

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.