arXiv:2505.11157cs.LGcs.AI2025-05NeurIPS被引 10

提出球面注意力机制,让Transformer更好处理球形数据。

Attention on the Sphere

  • 将数值积分权重融入注意力,实现几何保真球面计算。
  • 在三个任务中均优于平面模型,尤其在球面物理模拟上优势明显。
  • 适合需要保留球对称性的气象、宇宙学和机器人领域。

我们提出一种适用于球面域的广义注意力机制,使Transformer架构能原生处理定义在二维球面上的数据——这在大气物理学、宇宙学和机器人学等领域至关重要,因保持球面对称性和拓扑结构对物理准确性不可或缺。通过将数值积分权重集成到注意力机制中,我们获得几何忠实的球面注意力,近似旋转等变性,提供强归纳偏置,性能优于笛卡尔方法。为进一步提升可扩展性和模型表现,我们提出球面邻域注意力,将交互限制在测地邻域内,降低计算复杂度并引入局部性归纳偏置,同时保持方法的对称性。我们提供了优化的CUDA核函数和内存高效实现以确保实际应用。该方法在三个不同任务上得到验证:旋转球面上浅水方程模拟、球面图像分割和球面深度估计。在所有任务中,我们的球面Transformer持续优于其平面对应模型,凸显了几何先验在球面域学习中的优势。

原文摘要 · Abstract (English)

We introduce a generalized attention mechanism for spherical domains, enabling Transformer architectures to natively process data defined on the two-dimensional sphere - a critical need in fields such as atmospheric physics, cosmology, and robotics, where preserving spherical symmetries and topology is essential for physical accuracy. By integrating numerical quadrature weights into the attention mechanism, we obtain a geometrically faithful spherical attention that is approximately rotationally equivariant, providing strong inductive biases and leading to better performance than Cartesian approaches. To further enhance both scalability and model performance, we propose neighborhood attention on the sphere, which confines interactions to geodesic neighborhoods. This approach reduces computational complexity and introduces the additional inductive bias for locality, while retaining the symmetry properties of our method. We provide optimized CUDA kernels and memory-efficient implementations to ensure practical applicability. The method is validated on three diverse tasks: simulating shallow water equations on the rotating sphere, spherical image segmentation, and spherical depth estimation. Across all tasks, our spherical Transformers consistently outperform their planar counterparts, highlighting the advantage of geometric priors for learning on spherical domains.

球面注意力几何先验Transformer物理建模

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。