Flash不变点注意力
机器学习
2025-05-20 v1 人工智能
生物大分子
摘要
不变点注意力（IPA）是结构生物学中几何感知建模的关键算法，对许多蛋白质和RNA模型至关重要。然而，其二次复杂度限制了输入序列长度。我们引入了FlashIPA，这是IPA的一种因式分解重构，它利用硬件高效的FlashAttention，在GPU内存和实际运行时间上实现了随序列长度的线性缩放。FlashIPA达到或超越了标准IPA的性能，同时大幅降低了计算成本。FlashIPA将训练扩展到以前无法达到的长度，我们通过在没有长度限制的情况下重新训练生成模型并生成数千个残基的结构来证明了这一点。FlashIPA可在https://github.com/flagshippioneering/flash_ipa获取。
引用
@article{arxiv.2505.11580,
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 A Ban and Olivia Viessmann},
journal= {arXiv preprint arXiv:2505.11580},
year = {2025}
}