提出新型旋转等变注意力机制,实现高效全局建模且支持高阶特征。
Clebsch-Gordan Transformer: Fast and Global Equivariant Attention
- 基于SO(3)不可约表示的克莱布施-戈登卷积,实现等变全局注意力。
- 计算复杂度降至O(N log N),在多个数据集上提升速度与精度。
- 支持任意阶等变特征,适合物理、化学及机器人等几何建模任务。
全局注意力是Transformer成功的关键,但其计算成本随标记数呈二次增长。而等变模型虽在物理、生化、计算机视觉和机器人任务中表现更优,却因计算开销大,仅支持低阶等变特征和局部上下文窗口,限制了表达能力。本文提出克莱布施-戈登变换器(Clebsch-Gordan Transformer),通过在SO(3)不可约表示上设计新型克莱布施-戈登卷积,实现高效全局注意力。该方法可在保持所有阶等变特征建模能力的同时,将输入标记复杂度降至O(N log N)。此外,利用克莱布施-戈登矩阵的稀疏性,有效扩展至高阶不可约表示。还通过权重共享或数据增强引入可选的标记置换等变性。在多体模拟、QM9、ModelNet点云分类及机器人抓取数据集上的实验表明,相比现有等变Transformer,本方法在GPU内存占用、运行速度和准确率方面均有显著提升。
原文摘要 · Abstract (English)
The global attention mechanism is one of the keys to the success of transformer architecture, but it incurs quadratic computational costs in relation to the number of tokens. On the other hand, equivariant models, which leverage the underlying geometric structures of problem instance, often achieve superior accuracy in physical, biochemical, computer vision, and robotic tasks, at the cost of additional compute requirements. As a result, existing equivariant transformers only support low-order equivariant features and local context windows, limiting their expressiveness and performance. This work proposes Clebsch-Gordan Transformer, achieving efficient global attention by a novel Clebsch-Gordon Convolution on $\SO(3)$ irreducible representations. Our method enables equivariant modeling of features at all orders while achieving ${O}(N \log N)$ input token complexity. Additionally, the proposed method scales well with high-order irreducible features, by exploiting the sparsity of the Clebsch-Gordon matrix. Lastly, we also incorporate optional token permutation equivariance through either weight sharing or data augmentation. We benchmark our method on a diverse set of benchmarks including n-body simulation, QM9, ModelNet point cloud classification and a robotic grasping dataset, showing clear gains over existing equivariant transformers in GPU memory size, speed, and accuracy.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。