要旨: 回転位置埋め込み（RoPE）は、言語、視覚、3D領域にまたがる現代のTransformerアーキテクチャの中核となる構成要素になっています。しかし、既存の実装はベクトル単位での分割・結合（split/merge）操作に依存しており、注意（attention）の最適化において見落とされがちな、無視できない計算オーバーヘッドが生じます。この問題は、さらに多次元の設定（例：2Dおよび3D RoPE）において増幅されます。そこでは、追加のベクトル操作や不均一な特徴分割がハードウェア利用率を低下させます。これらの制限を克服するために、我々はRoME（Rotary Matrix position Embedding）を提案します。RoMEは、RoPEと数学的に同等でありながら、ベクトル操作を統一された行列変換で置き換えることにより計算効率を高めた再定式化です。RoMEは次元固有の操作を排除し、実装を単純化し、さらに最新のNPUにおいてCubeユニットとVectorユニットをまたいだ融合（fused）並列実行を可能にします。実験の結果、RoMEはオペレータ・レベルおよびフルモデル・レベルの両方で大幅な高速化をもたらすことが示されました。実装はhttps://gitcode.com/cann/ops-transformer/blob/master/experimental/posembedding/rope_matrix/README.mdで利用できます。
回転位置埋め込み（RoPE）のための効率的な行列実装
arXiv cs.LG / 2026/4/14
📰 ニュースDeveloper Stack & InfrastructureIdeas & Deep AnalysisModels & Research
要点
- 本論文は、一般的な回転位置埋め込み（Rotary Position Embedding: RoPE）の実装が、ベクトル単位の分割／統合（split/merge）操作によって避けられるオーバーヘッドを負っていることを指摘している。特に2D/3Dの多次元設定では、ハードウェアの利用効率が低下する。
- RoME（Rotary Matrix position Embedding）を提案し、RoPEと数学的に同等な再定式化によって、ベクトル演算を統一的な行列変換に置き換える。
- 次元固有の操作を取り除くことで、RoMEは実装を単純化し、現代のNPUにおける（CubeおよびVectorユニット間での）より効率的な融合並列実行を支援する。
- 実験では、個々のオペレータレベルと、完全なTransformerモデル全体の双方で速度向上が報告されており、マイクロ最適化を超えた実運用上の性能メリットが示される。
- 著者らは、評価や統合のための実装リファレンスを公開リポジトリのリンクとして提供している。
