将蛋白结构建模中的IPA算法加速至线性复杂度,支持超长序列生成。
Flash Invariant Point Attention
- 用分治法重构IPA,结合FlashAttention实现硬件高效计算
- 序列长度达数千残基时仍保持高精度,内存与时间开销线性增长
- 适合需要长序列建模的蛋白质生成与结构预测研究者
不变点注意力(IPA)是结构生物学中几何感知建模的关键算法,广泛应用于蛋白质和RNA模型。然而其二次复杂度限制了输入序列长度。我们提出FlashIPA,一种基于分块重写重构的IPA方法,利用硬件高效的FlashAttention实现GPU内存和实际运行时间随序列长度呈线性增长。FlashIPA在性能上达到或超过标准IPA,同时显著降低计算成本。该方法使训练扩展至此前无法实现的长序列,我们通过无长度限制的生成模型重新训练,成功生成了含数千残基的结构。FlashIPA代码已开源:https://github.com/flagshippioneering/flash_ipa。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。