🤖 AI

E2Former-V2: On-the-Fly Equivariant Attention with Linear Activation Memory

E2Former-V2は、等変軸並列スパース化（Equivariant Axis-Aligned Sparsification）とオンザフライ等変アテンション（On-the-Fly Equivariant Attention）メカニズムを導入することで、主流のモデルにおける計算上のボトルネックを克服し、3D原子系における予測性能を維持しつつ20倍のTFLOPS向上を実現する、スケーラブルな等変グラフニューラルネットワークアーキテクチャである。

原著者： Lin Huang, Chengxiang Huang, Ziang Wang, Yiyue Du, Chu Wang, Haocheng Lu, Yunyang Li, Xiaoli Liu, Arthur Jiang, Jia Zhang

公開日 2026-06-08

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Lin Huang, Chengxiang Huang, Ziang Wang, Yiyue Du, Chu Wang, Haocheng Lu, Yunyang Li, Xiaoli Liu, Arthur Jiang, Jia Zhang

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、膨大な群衆（原子）が部屋の中でどのように動き、相互作用するかを予測しようとしていると想像してください。化学や物理学の世界では、これを「原子論的モデリング（atomistic modeling）」と呼びます。長い間、科学者たちは「等変グラフニューラルネットワーク（Equivariant Graph Neural Network: EGNN）」と呼ばれる特別な種類のAIを使用してきました。これらのネットワークは、部屋全体を回転させても、その中の物理現象は変わらないという性質を理解しているため、非常に優れています。

しかし、現在のバージョンのネットワークには大きな問題があります。それは、群衆が大きくなると不器用で動作が遅くなることです。

以下は、この論文の解決策であるE2Former-V2を、日常的な比喩を用いて分かりやすく解説したものです。

問題点：「エッジ中心」の交通渋滞

パーティーでは、誰もが隣の人と話をしたがっています。

旧来の手法（従来のEGNN）： このシステムは、パーティーで行われているあらゆる会話の巨大な物理リストを作成します。もし1,000人がいて、それぞれが50人の隣人と会話する場合、システムは計算を行う前に、膨大なホワイトボード（メモリ）に50,000個もの個別の会話メモを書き留めなければなりません。
ボトルネック： パーティーの規模が大きくなる（原子が増える）につれ、このホワイトボードは巨大になり、コンピュータの容量が足りなくなります。コンピュータは、実際に思考するよりも、これらのメモを移動させることだけに時間を費やしてしまいます。これが、古いモデルがタンパク質のような大きな分子をシミュレーションしようとすると、クラッシュしたり、動作が極端に遅くなったりする理由です。

解決策：E2Former-V2

著者らは、すべての会話メモを書き留めるのをやめ、プロセスを電光石火のように速くするための2つの巧妙なトリックを用いた新しいシステムを構築しました。

トリック1：「魔法のコンパス」（等変軸方向スパース化）

旧システムでは、原子がどのように相互作用するかを計算することは、あらゆる方向に回転できるパズルのピースを解くようなもので、非常に煩雑で多くの計算を必要としました。

比喩： 複雑な3Dパズルを持っていると想像してください。パズルが回転している間に解こうとする代わりに、「魔法のコンパス」を使って、特定のピースが常に真上を向くようにパズル全体を回転させます。
結果： パズルが整列されると、ほとんどのピースは空（ゼロ）になります。もはや空の部分に対して計算を行う必要はありません。論文ではこれをEAASと呼んでいます。これは、乱雑で密な計算を、単純な「再インデックス（ピースの並べ替え）」ゲームへと変え、これにより計算が大幅に高速化されます。

トリック2：「ストリーミング・シェフ」（オンザフライの等変アテンション）

これが最大のブレイクスルーです。

旧来の方法（ビュッフェ形式）： 旧システムは、全員が食べ始める前に、すべてのビュッフェ（5万個の会話メモすべて）を用意してテーブルに並べます。もしテーブルが小さすぎると、料理が溢れ出し、パーティーは中断してしまいます。
新しい方法（ストリーミング・シェフ）： 新しいシステムであるE2Former-V2は、顧客の目の前で一度に一皿ずつ料理を作り、すぐに提供するシェフのように振る舞います。
- メモを書き留めません。
- 全体のリストを保存しません。
- 2つの要素間の相互作用を計算し、その結果を最終スコアに加算した後、次のペアのために即座に詳細を忘れます。
結果： これは「オンザフライ（On-the-Fly）」処理と呼ばれます。巨大なホワイトボードを作成しないため、メモリをほとんど消費しません。計算をコンピュータの超高速な内部キャッシュ（SRAM）に収めることができるため、旧来の手法よりも20倍高速になります。

彼らが実際に達成したこと

著者らは、この新システムを主に2つの対象でテストしました。

低分子（SPICE）： 薬物のような分子を用いてテストを行いました。新システムは、既存の最高水準のモデルと同等の精度を維持しながら、より高速に動作しました。
巨大分子（OMol25）： 巨大な化学構造を用いてテストを行いました。
- 限界： 旧来のモデルは、10,000個以上の原子を含むシステムをシミュレートしようとすると、メモリ不足でクラッシュしました。
- 勝利： E2Former-V2は、クラッシュすることなく100,000個の原子を含むシステムをシミュレートすることに成功しました。

言及されている実世界のテスト

著者らは単に数値を動かしただけでなく、実際の物理現象をシミュレートしました。

液体の水： 水分子のボックスをシミュレートしました。新モデルは、水分子がどのように配置されるか（動径分布関数）をほぼ完璧に予測し、従来のモデルよりも実世界の実験結果に近く、正確に一致しました。
タンパク質： 水中に浮遊する大きなタンパク質（約30,000原子）をシミュレートしました。モデルはタンパク質の安定性を維持し、実験データと一致する形状への折り畳みを助け、複雑な生物学的構造を扱えることを証明しました。

まとめ

E2Former-V2は、巨大で遅いホワイトボードにすべての会話を書き留めるシステムから、スマートなシェフが会話を一つずつ即座に調理して提供するシステムへとアップグレードしたようなものです。これにより、科学者は標準的なコンピュータを使用しながら、メモリ不足に陥ることなく、巨大なタンパク質のような大規模な分子システムを、物理的な正確さを保ったままシミュレートできるようになります。

テクニカルサマリー: E2Former-V2: 線形活性化メモリを用いたオンザフライ等変アテンション

1. 問題提起

等変グラフニューラルネットワーク（EGNN）は、回転および並進対称性を保持する能力により、3次元原子系をモデリングするための標準となっています。しかし、主流のアーキテクチャ（例：MACE、EquiformerV2、eSCN）は、その**エッジ中心（edge-centric）**の設計に起因する決定的な効率性のボトルネックに苦しんでいます。これらのモデルは、幾何学的特徴を明示的に構築するか、すべてのエッジに対して高密度なテンソル積を実行するため、以下の問題を引き起こします：

メモリ複雑性: $O(|E|)$ の活性化メモリ。ここで $|E| \approx kN$ （ $k$ は平均近傍数）です。大規模なシステムでは、これによりGPU容量を超える膨大なメモリフットプリントが発生します。
ハードウェアの非効率性: エッジテンソルの明示的な実体化は、高帯域幅メモリ（HBM）とオンチップSRAM間のデータ移動を強制し、スループットを制限するメモリウォールを生み出します。
スケーラビリティの限界: システムサイズ（ $N$ ）が増加するにつれ、レイテンシが悪化するため、大規模な分子動力学（MD）シミュレーションを汎用的なハードウェアで行うことは計算上極めて困難になります。

近年のTransformer（例：FlashAttention）における進歩は、ストリーミングやタイルベースの実行を通じて、シーケンスデータにおける同様のボトルネックを解決してきましたが、これらの原理はSO(3)等変アーキテクチャにおいては未開拓のままでした。

2. 手法

E2Former-V2は、代数的疎性化（algebraic sparsification）とハードウェアを意識したノード中心の実行モデルを統合することで、これらの問題に対処します。このアーキテクチャは、2つの核心的なイノベーションで構成されています。

2.1. 等変軸方向疎性化 (EAAS)

標準的なSO(3)テンソル縮退は、クレブシュ・ゴーダン係数に関する高密度な総和を伴い、 $O(L^6)$ の複雑さと不規則なメモリアクセスパターンをもたらします。E2Former-V2は、これらの高密度演算を疎な置換ベースの演算へと変換するために、EAASを導入しています。

SO(3) $\to$ SO(2) 基底変換: この手法は、グローバル座標系をローカルベクトル $\vec{r}$ （2つの原子を結ぶ軸）に整合させます。この整合されたフレーム内では、実球調和関数は極限まで疎になり（補題 4.1）、 $m=0$ 成分のみが非ゼロとなります。
再インデックス演算子: 回転をテンソル積を通じて交換することで、高密度な縮退は決定論的な再インデックス規則（式 13）へと崩壊します。この演算子は、磁気指数の高密度な総和を排除し、疎なブロック単位の線形写像を介して入力特徴量を出力特徴量へとマッピングします。
結果: これにより、畳み込みステージの複雑さが軽減され、標準的な実装に対して約1.5倍の高速化を実現します。

2.2. オンザフライ等変アテンション

$O(|E|)$ のメモリボトルネックを排除するため、著者らはカスタムTritonカーネルによって実装された、完全に融合されたノード中心のアテンションメカニズムを提案しています。

ノード中心の因子分解: エッジ特徴量を実体化する代わりに、アテンションメカニズムは計算を分離します。クエリ（Query）およびキー（Key）ベクトルはSO(3)不変なスカラー投影（式 6）を用いて構築され、等変な幾何学的情報は値（Value）のパスを通じてのみ伝播されます。
ストリーミング削減: カーネルは近傍をストリーミング方式で処理します。近傍 $j \in \mathcal{N}(i)$ を反復処理し、アテンションスコアを計算し、オンザフライでSRAM内で値の更新を累積します。
エッジ実体化の排除: $N \times K \times H \times D$ の形状を持つ中間テンソルは、決してHBMに書き込まれません。カーネルは、オンチップで実行統計（最大値、正規化、および値の累積器）を維持し、活性化メモリの複雑性を $O(|E|)$ から $O(|V|)$ へと低減させます。
オンラインソフトマックス: ストリーミング中の数値的安定性を確保するため、累積器を逐次的に更新するオンラインソフトマックス定式化が使用されます。

3. 主な貢献

等変軸方向疎性化 (EAAS): Wigner-6j再結合とSO(2)軸整合を統一し、高密度なSO(3)テンソル縮退を疎な再インデックス演算へと変換する新しい代数的技術。
オンザフライ等変アテンション: エッジの実体化を回避し、SRAMの再利用を最大化して、標準的な実装と比較して最大20倍の高いTFLOPSを達成する、完全に融合されたTritonカーネル。
スケーラブルな性能: 単一のGPU上で10万個の原子を扱う能力を実証。これは、従来の等変Transformerがメモリ不足（OOM）エラーにより失敗する領域です。
物理的忠実度: 液体の水の構造の正確な再現およびAlphaFold予測タンパク質構造（約3万原子）の安定した精緻化における検証。

4. 実験結果

著者らは、SPICE（創薬化学）およびOMol25（基盤モデル規模）のデータセット、ならびにMDシミュレーションにおいてE2Former-V2を評価しました。

予測精度:
- SPICEにおいて、E2Former-V2はほとんどのサブセットで最低のエネルギーおよび力の平均絶対誤差（MAE）を達成し、DimersにおけるEnergy MAEにおいてMACE-Largeを48%上回りました。
- OMol25において、保守的（Conservative）なバリアントは1.27 meV/atomの集計Energy MAEを達成し、eSEN-smallに匹敵し、UMA-S（3.62 meV/atom）を大幅に上回りました。
- OC20 S2EF-2Mにおいて、EquiformerV2やeSCNのような強力なベースラインと同等の精度を達成しました。
効率性とスケーラビリティ:
- スループット: 融合カーネルは、ナイーブなPyTorch実装と比較して、平均20倍のTFLOPS高速化を達成しました。
- メモリ: ピークメモリ使用量は、ベースラインと比較して16.9倍から19.5倍減少しました。
- スケーラビリティ: 従来のモデル（例：EquiformerV2、MACE-Large）は1万〜5万原子を超えるとOOMエラーに直面しましたが、E2Former-V2は10万原子まで確実にスケールしました。
- 推論速度: 10万原子において、E2Former-V2 (Conservative) は0.12 steps/sに達し、UMA-Sよりも約3倍高速でした。Direct設定では、小規模スケールにおいてAllegroやEquiformerV2よりも1桁高速でした。
分子動力学:
- 周期境界条件を持つバルク水（216分子）のシミュレーションにおいて、E2Former-V2は実験的な動径分布関数（RDF）に密接に一致し、MACE-OFFを上回りました。
- 溶媒和タンパク質（約3万原子）の2 nsシミュレーションは、ネイティブ構造に対する高いTMスコアを維持し、安定したダイナミクスを示しました。

5. 意義と主張

本論文は、E2Former-V2が汎用GPU上での大規模等変Transformerの実用的な実現可能性を示していると主張しています。実行パラダイムをエッジ中心からノード中心へと移行し、ハードウェアを意識したストリーミングを活用することで、等変モデルのスケールを制限してきた根本的なメモリボトルネックを解決しています。

著者らは、これは単なるアルゴリズムの改善ではなく、理論的な因子分解（複雑性を $O(|E|)$ から $O(|V|)$ へ削減）とハードウェア効率の間のギャップを埋める、システムレベルの実装であると強調しています。この研究は、これまで計算上不可能であった大規模な分子基盤モデルや長時間スケールのMDシミュレーションが、標準的なGPUインフラを使用してアクセス可能になったことを示唆しています。コミュニティでのさらなる採用を促進するため、コードは公開されています。