arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2407.01049cs.LG
收藏

面向可扩展等变学习的SE(3)-Hyena算子

SE(3)-Hyena Operator for Scalable Equivariant Learning

Johnson and Johnson Innovative Medicine(强生创新医药)

Artem Moskalev, Mangal Prakash, Rui Liao, Tommaso Mansi

AI总结:

针对现有等变学习方法复杂度高或丢失全局信息的问题,提出基于Hyena算子的SE(3)-Hyena等变长卷积模型,以次二次复杂度实现全局几何上下文建模,在相关任务上性能优异且计算效率更高、内存占用更低。

摘要:

在生物学、化学或计算机视觉等诸多领域中,在保持等变性的同时对全局几何上下文进行建模对实现准确预测至关重要。然而,由于大规模处理高维数据的计算需求,这一目标颇具挑战性。现有方法(如等变自注意力或基于距离的消息传递)存在与序列长度相关的二次复杂度问题,而局部化方法则会牺牲全局信息。受近期状态空间模型和长卷积模型取得的成功启发,本研究提出了SE(3)-Hyena算子——一种基于Hyena算子的等变长卷积模型。SE(3)-Hyena以次二次复杂度捕获全局几何上下文,同时保持对旋转和平移的等变性。在等变联想回忆和n体建模任务上的评估结果显示,SE(3)-Hyena的性能与等变自注意力相当或更优,且在处理长序列时所需的内存和计算资源显著更少。我们的模型处理2万个token的几何上下文的速度是等变Transformer的3.5倍,并且在相同内存预算下可支持的上下文长度是其175倍。

英文摘要:

Modeling global geometric context while maintaining equivariance is crucial for accurate predictions in many fields such as biology, chemistry, or vision. Yet, this is challenging due to the computational demands of processing high-dimensional data at scale. Existing approaches such as equivariant self-attention or distance-based message passing, suffer from quadratic complexity with respect to sequence length, while localized methods sacrifice global information. Inspired by the recent success of state-space and long-convolutional models, in this work, we introduce SE(3)-Hyena operator, an equivariant long-convolutional model based on the Hyena operator. The SE(3)-Hyena captures global geometric context at sub-quadratic complexity while maintaining equivariance to rotations and translations. Evaluated on equivariant associative recall and n-body modeling, SE(3)-Hyena matches or outperforms equivariant self-attention while requiring significantly less memory and computational resources for long sequences. Our model processes the geometric context of 20k tokens x3.5 times faster than the equivariant transformer and allows x175 longer a context within the same memory budget.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑