技術要約：プラトニック・トランスフォーマー：等変性のための確かな選択肢

問題提起

トランスフォーマーは、その柔軟性と効率性によりディープラーニングにおいて広く成功を収めてきましたが、科学的分野（物理学、化学）やコンピュータビジョン（3D点群）に存在する幾何学的対称性に固有の帰納バイアスを欠いています。これらの対称性を尊重する既存の等変性手法は、群畳み込みやクレブシュ・ゴルダン・テンソル積のような複雑な演算に依存することがよくあります。これらのアプローチは、標準的なトランスフォーマーアーキテクチャと比較して、計算上のオーバーヘッドとアーキテクチャの複雑さを増大させ、スケーラビリティと柔軟性を阻害します。中心となる課題は、標準的なトランスフォーマーの速度、柔軟性、および計算効率を犠牲にすることなく、強力な幾何学的帰納バイアスをトランスフォーマーの枠組みに組み込むことです。

手法：プラトニック・トランスフォーマー

著者らは、標準的なトランスフォーマーの注意機構（Attention mechanism）や計算グラフを変更することなく、連続的な並進および離散的な回転反射（プラトオン（正多面体）の対称性）に対する等変性を実現するフレームワークである**プラトニック・トランスフォーマー（Platonic Transformer）**を提案しています。

コアメカニズム

1. 参照フレームのリフティング（Reference Frame Lifting）: 入力特徴量を単一のグローバルな座標系で処理する代わりに、この手法では入力特徴量を離散的な対称群 $G \subset O(n)$（例：正四面体、正八面体、または正二十面体の回転対称群）上の関数へとリフトアップします。特徴量は、$G$ の要素によって定義される一連の参照フレームに関連付けられて表現されます。
2. フレーム依存型RoPE（Frame-Dependent RoPE）: この手法は、回転位置エンコーディング（RoPE）を一般化したものです。標準的なRoPEは並進等変性を保証しますが、プラトニック・トランスフォーマーは各参照フレーム $R \in G$ に対して相対的に位置エンコーディングを再定義します。位置 $p_i$ にあるトークンに対して、各フレームに対する相対的な位置は $p_i(R) = R^{-1}p_i$ となります。その後、これらの相対位置に対してRoPE演算が適用されます。
3. 重み共有スキーム（Weight-Sharing Scheme）: すべての線形層（Query、Key、Value投影、およびMLP）は、群 $G$ に対して等変であるよう制約されます。数学的には、この手法により線形層は**群畳み込み（group convolutions）**へと変換されます。平坦化されたベクトル表現においては、これは重み行列が群上の学習可能なカーネルによって決定されるブロック行列として現れる、構造化された重み共有スキームとして記述されます。これにより、パラメータ数を制約のない層と比較して係数 $|G|$ だけ削減しながら、同一の計算グラフ構造を維持します。
4. 並列アテンション・ヘッド（Parallel Attention Heads）: アテンション機構は、すべての参照フレームにわたって並列に動作します。各フレームは実質的に個別の注意ヘッドとして機能します。生の注意スコアは $s_{ij}(R) = q_i(R)^\top \rho_\Omega(p_j(R) - p_i(R)) k_j(R)$ として計算され（ここで $\rho_\Omega$ はRoPE回転行列）、ソフトマックスはフレームごとに適用されるため、モデルは方向依存の注意パターンを学習することが可能になります。

理論的洞察

• 動的群畳み込み（Dynamic Group Convolution）: 著者らは、RoPEベースのアテンション（ソフトマックスが省略される場合、または線形アテンションの設定において）と動的群畳み込みとの間の形式的な等価性を確立しています。これは、クエリベクトルが疎なフーリエ基底展開の係数として機能し、コンテンツに応じた幾何学的フィルタをオンザフライで合成することを意味します。
• 線形計算量: この畳み込みの観点から、このメカニズムは $O(N)$ の時間（トークン数に対して線形）で計算可能であることが明らかになり、大規模な3Dデータに適したスケーラブルな代替案を提供します。

主な貢献

1. アーキテクチャ: プラトニック固体の対称群を用いた原理的な重み共有スキームを通じて、標準的なトランスフォーマーに幾何学的等変性を統合し、ベースラインと同じアーキテクチャと計算コストを維持するプラトニック・トランスフォーマーの導入。
2. 理論的統一: RoPEベースのアテンションが動的群畳み込みと形式的に等価であることを示し、適応的な幾何学的フィルタを学習するモデルの能力を明らかにしました。
3. スケーラビリティ: 長いシーケンス長にも対応可能な、大規模3Dデータに適した効率的な線形時間型の畳み込みバリアントの提供。
4. 実証的検証: 様々なベンチマークにおける包括的な評価により、幾何学的制約が追加の計算コストなしに性能を向上させることを示しました。

実験結果

論文では、「非等変タスク」（データに標準的な向きがあるもの）と「等変タスク」（データの向きが任意であるもの）の両方についてモデルを評価しています。

• CIFAR-10 (2D Vision): 2D回転対称性（例：$C_6, D_3$）を組み込むことで、並進のみのベースライン（RoPEを用いたViT）よりも精度が向上しました。特筆すべきは、最もパラメータ効率の高い構成（重み共有によりパラメータ数が少ない構成）が最高の性能を達成したことであり、これは利得が容量の増加ではなく、対称性による正則化から来ていることを示唆しています。
• ScanObjectNN (3D Point Clouds): 方向の変化や遮蔽が顕著なタスクにおいて、正四面体群のバリアントはベースラインに対してアテンションで2.43%、畳み込みで1.54%上回りました。線形時間型の畳み込みバリアントは、競争力のある精度を維持しながら大幅な高速化を実現しました。
• QM9 (分子生成および特性予測):
  • 生成: 拡散トランスフォーマー（DiT）内の標準的なトランスフォーマーをプラトニック・トランスフォーマー（DiP）に置き換えることで、ADiT、GeoLDM、EDMなどのベースラインと比較して、より高い妥当性と有意に改善された分子安定性を実現しました。等変バリアント（DiP-Tetrahedron）は、厳格な幾何学的健全性チェック（PoseBusters）において、ほぼ完璧な通過率を達成しました。
  • 特性予測: 標準的なQM9特性予測タスクにおいて、競争力のある結果が得られました。
• OMol25 (大規模な力/エネルギー予測): 400万構造の巨大なデータセットにおいて、プラトニック・トランスフォーマー（正四面体および正八面体バリアント）は、力およびエネルギー予測の精度（MAE）において、最先端のeSEN-smベースラインを大幅に上回り、かつ2.6倍高い学習スループットを達成しました。対称群のサイズを大きくすること（正四面体から正八面体へ）で、性能の低下を最小限に抑えつつパラメータを半分に削減でき、重み共有スキームの効率性がさらに検証されました。
• ProteinMD: 本モデルは、タンパク質のダイナミクスにおける力を予測する上で、従来の最先端手法（G-HyenaやEGNNを含む）を上回る性能を示しました。

意義と主張

本論文は、プラトニック・トランスフォーマーが、対称性の認識とスケーラビリティの間の長年のトレードオフを解決することを主張しています。プラトオンの幾何学的構造を利用して重み共有を誘導することで、本手法は、トランスフォーマーのコアとなる計算グラフを変更したり計算コストを増大させたりすることなく、近似的な $E(n)$ 等変性を達成します。

著者らは、このアプローチが、原理的な等変性と現代的なスケーラビリティは互いに排他的ではないことを示していると強調しています。さらに、動的群畳み込みとの形式的な接続は、適応的な幾何学的フィルタを学習するモデルの能力に理論的根拠を与え、高度にスケーラブルな線形時間型バリアントを可能にします。結果は、物理法則によって支配される科学的分野（等変性が「プラトニックな理想」である領域）において、プラトニック・トランスフォーマーを通じてこれらの制約を課すことが、標準的なアーキテクチャや複雑な等変性代替案と比較して、よりデータ効率が高く、堅牢で、物理的に妥当なモデルをもたらすことを示唆しています。