用于大规模等变学习的几何 Hyena 网络
机器学习
2025-05-29 v1
摘要
在对生物、化学和物理系统进行建模时，处理全局几何上下文同时保持等变性至关重要。然而，由于等变性和大规模全局上下文的计算需求，这极具挑战性。标准方法如等变自注意力存在二次复杂度问题，而局部方法如基于距离的消息传递则牺牲了全局信息。受状态空间模型和长卷积模型最近成功的启发，我们引入了 Geometric Hyena，这是第一个用于几何系统的等变长卷积模型。Geometric Hyena 以次二次复杂度捕获全局几何上下文，同时保持对旋转和平移的等变性。在大型 RNA 分子的全原子性质预测和完整蛋白质分子动力学评估中，Geometric Hyena 优于现有的等变模型，同时所需的内存和计算量远低于等变自注意力。值得注意的是，我们的模型处理 30k token 的几何上下文速度比等变 Transformer 快 20 倍，并且在相同预算下允许 72 倍长的上下文。
引用
@article{arxiv.2505.22560,
title = {Geometric Hyena Networks for Large-scale Equivariant Learning},
author = {Artem Moskalev and Mangal Prakash and Junjie Xu and Tianyu Cui and Rui Liao and Tommaso Mansi},
journal= {arXiv preprint arXiv:2505.22560},
year = {2025}
}