[論文レビュー] e3nn: Euclidean Neural Networks

Mario Geiger, Tess Smidt|arXiv (Cornell University)|Jul 18, 2022

Medical Imaging and Analysis被引用数 106

ひとこと要約

tldr: irreps、tensor products、球面調和関数を用いて E(3) 等変換ニューラルネットワークを構築し、3D 幾何を扱うための Python フレームワーク。

ABSTRACT

We present e3nn, a generalized framework for creating E(3) equivariant trainable functions, also known as Euclidean neural networks. e3nn naturally operates on geometry and geometric tensors that describe systems in 3D and transform predictably under a change of coordinate system. The core of e3nn are equivariant operations such as the TensorProduct class or the spherical harmonics functions that can be composed to create more complex modules such as convolutions and attention mechanisms. These core operations of e3nn can be used to efficiently articulate Tensor Field Networks, 3D Steerable CNNs, Clebsch-Gordan Networks, SE(3) Transformers and other E(3) equivariant networks.

研究の動機と目的

• 3Dデータに対して対称性を意識したモデルの利用を提唱し、過度のデータ拡張を回避し、一般化性能を向上させる。
• 3D ユークリッド対称性 (E(3)) を尊重する等変ニューラルネットワークを構築するための柔軟で統一されたフレームワークを提供する。
• 数学的プリミティブ（irreps、tensor products、球面調和関数）を説明し、それらを畳み込みやトランスフォーマーなどのモジュールに組み込む方法。
• これらのビルディングブロックが Tensor Field Networks、3D Steerable CNNs、Clebsch-Gordan Networks、SE(3) Transformers などのアーキテクチャをどのように実現するかを示す。

提案手法

• データ変換をモデル化するために O(3) の irreps とパリティを導入する。
• Clebsch–Gordan 係数を用いて双線形入力を irreps に分解する一般的な TensorProduct 演算を開発する。
• テンソル積における複数の「経路」を定義して、全ての双線形等変演算を捉え、加重結合を可能にする。
• 索引対称性を伴う irreps への相互作用の分解を提供する ReducedTensorProduct ユーティリティ。
• 安定した訓練と等変性を維持するための初期化方式と成分正規化を説明する。

実験結果

リサーチクエスチョン

• RQ1 入力 irreps 間の全ての双線形等変演算を、出力 irreps に写像する形でどのように体系的に構築できるか？
• RQ2 統一された TensorProduct フレームワークを用いて、等変モジュール（畳み込み、アテンション など）を効率的に実装・組み合わせるにはどうすればよいか？
• RQ3 データ効率性と一般化の観点で、3D ジオメトリックデータの学習に対してE(3) 等変性を課すことの利点は何か？

主な発見

• このフレームワークは、単一の TensorProduct 原子でさまざまな E(3) 等変アーキテクチャを統合する。
• 球面調和関数と Clebsch–Gordan 係数を用いて、実数値基底でテンソル積を irreps に分解する。
• 複数経路に基づくテンソル積構成は、学習可能な重みを伴う完全な双線形等変演算を可能にする。
• ReducedTensorProduct は、インデックス対称性を尊重しつつ高階テンソルを irreps に体系的に約化する。
• 本論文は、等変モデルの初期化とデータ効率性の利点を、学習曲線の傾きへの影響を含めて論じる。