Flash Invariant Point Attention

2025-05-16

AI Summary

In structural biology, invariant point attention (IPA) suffers from quadratic computational complexity, hindering its application to long protein/RNA sequences. This work introduces FlashIPA—the first algorithm reformulating IPA into a linear-complexity, hardware-efficient FlashAttention paradigm. By factorizing the attention mechanism, optimizing geometric tensor operations, and incorporating structure-specific embeddings, FlashIPA achieves strict linear scaling in both GPU memory consumption and runtime. Experiments demonstrate that FlashIPA enables end-to-end structure generation for sequences comprising thousands of residues, eliminating length constraints even when retraining generative models. Crucially, it matches or exceeds standard IPA in modeling accuracy while reducing GPU memory usage by 58% and accelerating inference by 2.3×. This work breaks IPA’s fundamental sequence-length bottleneck for the first time, establishing a scalable foundation for geometric modeling of large biomolecular systems.

Technology Category

• Machine Learning: Structured Learning
• Natural Language Processing: Generation
• Humans and AI: Human-Aware Planning and Behavior Prediction

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

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.

Problem

• Reduces quadratic complexity of Invariant Point Attention
• Achieves linear scaling in GPU memory usage
• Enables training on longer biological sequences

Innovation

• Factorized IPA using FlashAttention for efficiency
• Achieves linear scaling in GPU memory usage
• Enables training on longer sequences unrestricted