旋转位置嵌入的高效矩阵实现
机器学习
2026-04-14 v1 计算机视觉与模式识别
摘要
旋转位置嵌入 (Rotary Position Embedding, RoPE) 已成为现代Transformer架构中处理语言、视觉和3D领域的核心组件。然而，现有实现依赖向量级别的分割与合并操作，带来显著的计算开销，在注意力优化中常被忽视。多维情境（如2D和3D RoPE）下的问题更为严重，额外的向量操作和不均匀特征划分会损害硬件利用率。为克服这些限制，我们提出RoME (Rotary Matrix position Embedding)，一种在数学等价性下以统一矩阵变换取代向量操作的高效表述。RoME消除维度特定操作，简化实现，使其能够在现代NPU的Cube和Vector单元上实现融合并行执行。实验表明，RoME在算子和完整模型层面均实现显著加速。其实现已公开：https://gitcode.com/cann/ops-transformer/blob/master/experimental/posembedding/rope_matrix/README.md。
关键词
引用
@article{arxiv.2604.09742,
title = {Efficient Matrix Implementation for Rotary Position Embedding},
author = {Chen Minqi and Zhongqi Yue and Shihao Zhang and Yun Xu and Peng Wu and kaixiang Xu and Zeyi Huang and Hanwang Zhang},
journal= {arXiv preprint arXiv:2604.09742},
year = {2026}
}