用于可扩展等变学习的 SE(3)-Hyena 算子
机器学习
2024-08-14 v2
摘要
在诸多领域如生物学、化学或计算机视觉中，准确预测需要捕获全局几何上下文且保持等变性。这一目标面临计算规模的挑战。现有方法如等变自注意力或基于距离的消息传递，因序列长度呈二次复杂度；而局部方法则牺牲了全局信息。灵感来源于状态空间和长卷积模型的最新成功，本文引入SE(3)-Hyena 算子，这是一种基于Hyena 算子的等变长卷积模型。SE(3)-Hyena 在保持旋转和平移等变性的同时，以亚二次复杂度捕获全局几何上下文。在等变关联记忆和 n 体建模评估中，SE(3)-Hyena 在匹配或超越等变自注意力的同时，需的内存和计算资源显著降低。我们的模型比等变变换器在2万 token 的几何上下文处理速度快3.5倍，同时在相同内存预算下可支持175倍更长的上下文。
关键词
引用
@article{arxiv.2407.01049,
title = {SE(3)-Hyena Operator for Scalable Equivariant Learning},
author = {Artem Moskalev and Mangal Prakash and Rui Liao and Tommaso Mansi},
journal= {arXiv preprint arXiv:2407.01049},
year = {2024}
}