Linear Memory SE(2) Invariant Attention

著者 Ethan Pronovost, Neha Boloor, Peter Schleede, Noureldin Hendy, Andres Morales, Nicholas Roy
所属 Zoox
投稿日 2025年07月26日
カテゴリ cs.LG
arXiv 2507.18597v1 ↗

Linear Memory SE(2) Invariant Attention

基本情報

  • arXiv ID: 2507.18597v1 (https://arxiv.org/abs/2507.18597)
  • 著者: Ethan Pronovost, Neha Boloor, Peter Schleede, Noureldin Hendy, Andres Morales, Nicholas Roy
  • 所属: Zoox
  • 投稿日: 2025年07月26日
  • カテゴリ: cs.LG

簡単に説明すると

この論文は、自動運転におけるマルチエージェント（車、歩行者など）の動きを予測するための新しいAIアーキテクチャを提案しています。

自動運転では、周囲の車や歩行者の位置と向き（SE(2)＝2次元空間での位置と回転）を考慮することが重要です。従来の手法では、全てのエージェント間の相対的な位置関係を計算するため、エージェント数の2乗に比例したメモリが必要でした。これは、エージェント数が多い場合に大きな問題となります。

本研究では、フーリエ級数展開を使用した新しいアプローチを提案し、エージェント数に対して線形のメモリで済むようにしました。これにより、大規模言語モデル（LLM）で使われている高効率なアテンション機構（Flash Attentionなど）を活用でき、平行移動と回転に対して不変なモデルを効率的に実現できます。

実験の結果、特に車が曲がる場面での予測精度が従来手法よりも向上したことが示されました。論文ではGitHubやHuggingFaceのリンクは提供されていませんが、自動運転企業Zooxの研究成果です。

1. 研究概要

1.1 背景と動機

自動運転システムにおいて、動的なエージェント（車両、歩行者など）の存在下での計画や制御には、各エージェントの空間的幾何学的特性をモデル化することが不可欠です。特に、SE(2)（Special Euclidean group in 2D）における位置と向きの情報は重要です。

近年、マルチエージェントの動作予測、計画、シミュレーションなどのタスクにおいて、SE(2)不変なネットワークアーキテクチャが最先端の性能を達成しています。これらの手法は、オブジェクト間の相対的なポーズ（位置と向き）のみを考慮することで、座標系の選択に依存しないモデルを実現しています。

しかし、既存のSE(2)不変モデルアーキテクチャは、全ての要素ペア間の相対ポーズを明示的に計算するため、GPUの高帯域幅メモリ（HBM）をシーン要素数の2乗に比例して必要とします。これは、大規模なシーンでは重大なボトルネックとなります。

一方、Transformerモデルでも当初は同様の二次メモリ問題がありましたが、Flash Attentionなどの技術により線形メモリでの計算が可能になりました。これにより、「SE(2)不変なTransformerアーキテクチャで線形HBMを実現できるか？」という自然な疑問が生まれました。

本研究は、この疑問に対して肯定的な答えを提供し、実用的な解法を提示します。

1.2 主要な貢献

本研究は、SE(2)不変なアテンション機構の効率化に関して以下の重要な貢献を行っています：

  • 線形メモリで動作する新しいSE(2)不変スケールド・ドットプロダクト・アテンション機構の提案
  • フーリエ級数展開を用いた相対ポーズ表現の近似手法の開発
  • 近似誤差が16ビット浮動小数点精度に匹敵するレベルであることの理論的・実験的証明
  • 既存のFlash Attention実装を変更せずに使用可能なアルゴリズムの提示
  • 実際のマルチエージェントシミュレーションタスクでの有効性の実証

2. 提案手法

2.1 手法の概要

提案手法の核心は、SE(2)相対ポーズを効率的にエンコードする新しいアテンション機構です。基本的なアイデアは、ロータリポジションエンベディング（RoPE）をSE(2)群に拡張することです。

一般的に、相対アテンションでは、クエリとキーの位置から相対位置を介した関数φを定義します。本手法の重要な洞察は、この関数を因数分解することです：
φ(p_{n→m}) = φ_q(p_n) φ_k(p_m)

この因数分解により、クエリとキーを事前に変換し、標準のスケールド・ドットプロダクト・アテンション（Flash Attentionなど）を使用し、その後で変換することが可能になります。これにより、二次のテンソルを明示的に生成することなく、線形メモリで計算を完了できます。

具体的には、クエリとキーベクトルにブロック対角の2D回転行列を適用します。これらの回転角度は、クエリとキー間の相対的なx位置、y位置、または向き（yaw）に基づきます。

2.2 技術的詳細

フーリエ近似の詳細：

SE(2)における相対ポーズの表現において、最大の挑戦はxとyの相対位置の処理です。例えば、相対的なx位置は次のように表されます：

x_{n→m} = (x_m - x_n)コサインθ_n + (y_m - y_n)サインθ_n

この式はx_mコサインθ_nのような項を含み、クエリとキーの両方のポーズを含むため、直接因数分解できません。

しかし、フーリエ級数を使用して近似的に因数分解することが可能です。具体的には、コサインとサイン関数の基底を使用して、相対位置に依存する項を展開します。

アルゴリズムの流れ：

  1. クエリポーズ p_n とクエリベクトル q_n が与えられたとき、変換 q̃_n = φ_q(p_n)^T q_n を適用
  2. キーポーズ p_m とキーベクトル k_m に対して、変換 k̃_m = φ_k(p_m) k_m を適用
  3. 標準のスケールド・ドットプロダクト・アテンション（SDPA）を変換されたベクトルに適用
  4. 出力に対して逆変換 o_n = φ_q(p_n) õ_n を適用

近似精度：

フーリエ展開の項数Fを調整することで、近似精度を制御できます。位置の大きさが4以下の場合、わずか18項でfloat16の精度に匹敵する近似を達成できます。

2.3 新規性

本研究の新規性は、以下の点にあります：

1. 線形メモリでのSE(2)不変性の実現
従来のSE(2)不変モデルは全て二次のメモリを必要としていましたが、本手法は初めて線形メモリで完全なSE(2)不変性（平行移動と回転の両方）を達成しました。

2. フーリエ級数による巧妙な因数分解
相対位置の非線形関数をフーリエ基底で展開し、クエリとキーに依存する部分を分離するというアプローチは、SE(2)群のような非アーベル群への新しい適用例です。

3. 既存インフラとの互換性
Flash Attentionのような既存の最適化されたGPUカーネルを変更せずに使用できるため、実装が容易で保守性が高いです。

4. 既存手法との比較

  • 2D RoPE: 平行移動不変ですが回転不変ではありません
  • SE(2)表現法: 完全なSE(2)不変ですが、位置のスケーリングが必要で訓練が不安定
  • 本手法: 完全なSE(2)不変で、サイン・コサインエンコーディングにより安定した訓練が可能

3. 実験結果

3.1 実験設定

提案手法の有効性を評価するため、以下の設定で実験を行いました：

データセット：

  • 3300万シナリオを含むプライベートデータセット
  • エージェントシミュレーションタスク：エージェントの履歴、道路地図、交通信号に基づいて次のアクションを予測

モデル設定：

  • 次トークン予測モデルを使用
  • エージェント（車両、歩行者）と地図要素（車線、横断歩道）をトークン化
  • Transformerで処理し、離散アクションのカテゴリ分布を予測
  • 位置は大きさが4以下になるようダウンスケール

比較手法：

  • 絶対位置埋め込み
  • 2D RoPE
  • SE(2)表現法
  • SE(2)フーリエ（提案手法）

評価指標：

  • 負の対数尤度（NLL）：地上真実アクションの予測確率
  • 最小平均変位誤差（minADE）：16個の6秒軌道をサンプリングし、最も良い予測の誤差
  • 軌道タイプ別評価：静止、直進、旋回

3.2 主要な結果

実験結果は、提案手法の有効性を明確に示しています：

全体的な性能：

  • NLL: 0.190（最良、2D RoPEと同等）
  • 静止軌道minADE: 0.23（最良、他手法と同等）
  • 直進軌道minADE: 1.79（2D RoPEの1.78に次ぐ）
  • 旋回軌道minADE: 2.60（最良、最も難しいカテゴリで大幅改善）

重要な発見：

  1. 2D RoPEとSE(2)フーリエは、絶対位置埋め込みやSE(2)表現法よりも優れた性能を示しました
  2. 2D RoPEは直進軌道でわずかに優れていましたが、SE(2)フーリエは旋回軌道で大幅に優れていました
  3. これは、回転不変性が旋回動作の予測に重要であることを示しています

近似精度の検証：

  • キー位置の大きさが2、4、8のとき、それぞれ基底サイズ12〔18〔28で16ビット浮動小数点精度に匹敵する平均近似誤差を達成
  • キー位置の大きさを2倍にすると、スペクトルノルム近似誤差、10^-3程度を維持するためには基底サイズを約50%増やす必要がある

3.3 既存手法との比較

提案手法を4つのアプローチと比較した結果から、以下の洞察が得られました：

絶対位置埋め込み：

  • 最も基本的なアプローチですが、全ての指標で最下位の性能
  • SE(2)変換に対して不変ではないため、一般化性能が低い

2D RoPE：

  • 平行移動不変であるが回転不変ではない
  • 直進軌道で最良の性能を示したが、旋回軌道ではSE(2)フーリエに劣る
  • 線形メモリで動作し、実装が簡単

SE(2)表現法：

  • 完全なSE(2)不変性を持つが、位置座標を直接使用するため訓練が不安定
  • 位置のダウンスケーリングが必要
  • 性能は2D RoPEとSE(2)フーリエの中間

SE(2)フーリエ（提案手法）：

  • 完全なSE(2)不変性を持ち、サイン・コサインエンコーディングにより安定した訓練
  • 線形メモリで動作し、Flash Attentionと互換
  • 特に旋回軌道で優れた性能を示し、実用的な価値が高い

4. 実用性評価

4.1 実装の容易性

提案手法は、実装の容易性において優れた特性を持ちます：

プラス面：

  • 既存のFlash Attention実装を変更せずに使用可能
  • 専用GPUカーネルの開発が不要
  • アルゴリズムがシンプルで理解しやすい
  • 大規模言語モデルでの成功事例を活用可能

実装上の考慮点：

  • フーリエ係数の事前計算が必要（数値積分を2F点で実施）
  • 位置の大きさに応じた基底サイズの選択が必要
  • Flash Attentionがサポートする特徴次元（最低256）内で設計が必要

全体的に、既存の深層学習フレームワークへの統合が容易で、特に自動運転分野での実用化が期待できます。

4.2 計算効率

システムの計算効率に関して、以下の特徴があります：

メモリ効率：

  • GPU HBM使用量：O(N)（Nはシーン要素数）
  • 従来手法のO(N^2)から大幅に改善
  • 大規模シーンでもスケーラブル

計算速度：

  • フーリエ係数の事前計算は高速（オフラインで実施可能）
  • 推論時はFlash Attentionの最適化された実装を活用
  • 基底サイズによる速度への影響は軽微

スケーラビリティ：

  • エージェント数に対して線形にスケール
  • バッチサイズの増加にも柔軟に対応
  • 近年の大規模言語モデルと同様のスケーリング特性

実用上、計算効率の向上は、より大規模なシーンやより多くのエージェントを扱うことを可能にします。

4.3 応用可能性

提案手法は、幅広い応用可能性を持ちます：

直接的な応用：

  • マルチエージェント動作予測
  • 自動運転における経路計画
  • エージェントシミュレーション
  • シナリオ生成

拡張可能な領域：

  • 他の群（SE(3)など）への拡張
  • ロボティクスにおける空間認識
  • コンピュータビジョンの空間的タスク
  • 物理シミュレーション

産業へのインパクト：

  • 自動運転車の安全性向上
  • リアルタイム処理の実現
  • 計算リソースの節約
  • より複雑なシーンの処理

特に、自動運転分野での実用化は直近の課題であり、Zooxを含む多くの企業がこの技術に注目しています。

5. まとめと所感

5.1 論文の意義

本研究は、自動運転分野における空間認識タスクの効率化に重要な貢献をしました。

学術的意義：
非アーベル群にRoPEスタイルのアプローチを適用するという挑戦に対し、フーリエ級数展開を用いたエレガントな解法を提示しました。これは、群論と機械学習の交差点における重要な進歩です。

実用的意義：
GPUメモリの制約は、実用的な自動運転システムにおいて大きなボトルネックでした。本手法により、より大規模なシーンやより多くのエージェントを扱うことが可能になり、安全性と性能の向上につながります。

技術的洞察：
回転不変性が特に旋回軌道の予測において重要であることが実験的に示されました。これは、空間的タスクにおける不変性の設計の重要性を改めて示しています。

限界と課題：
論文でも言及されているように、現時点では一つのデータセットでの評価に留まっています。また、データ拡張や二次メモリSE(2)不変アテンションとの詳細な比較が今後の課題です。

5.2 今後の展望

本研究は以下のような将来の研究方向を示唆しています：

手法の拡張：

  • SE(3)や他のリー群への拡張
  • より効率的な基底関数の探索
  • 適応的な基底サイズ選択の自動化
  • 近似誤差のさらなる削減

実験的検証：

  • より幅広い運転タスクとデータセットでの評価
  • データ拡張との比較
  • 二次メモリSE(2)不変アテンションとの詳細な比較
  • 計算速度とメモリ使用量の詳細な分析

応用領域の拡大：

  • コンピュータビジョンタスクへの適用
  • 物理シミュレーションへの応用
  • 分子動力学や材料科学への展開

長期的ビジョン：
本研究は、空間的タスクにおける効率的な不変性の実現という重要な方向性を示しました。将来的には、このようなアプローチが、より安全で効率的な自動運転システムの実現に貢献し、社会実装を加速させることが期待されます。