[論文レビュー] Frame Averaging for Invariant and Equivariant Network Design

Omri Puny, Matan Atzmon+5 more|arXiv (Cornell University)|Oct 7, 2021

Machine Learning in Materials Science参考文献 62被引用数 10

ひとこと要約

この論文は、計算的に扱いにくい群平均を、データに依存する小さな群要素の集合（フレームと呼ばれる）による平均に置き換える、体系的なフレームワークであるフレーム平均化（FA）を導入する。これにより、ニューラルネットワークにおける正確な不変性または等長性を実現できる。FAはバックボーンアーキテクチャの表現力も保ち、ポイントクラウド法線推定、グラフ分離、n体ダイナミクス予測の分野で最先端の性能を達成する。

ABSTRACT

Many machine learning tasks involve learning functions that are known to be invariant or equivariant to certain symmetries of the input data. However, it is often challenging to design neural network architectures that respect these symmetries while being expressive and computationally efficient. For example, Euclidean motion invariant/equivariant graph or point cloud neural networks. We introduce Frame Averaging (FA), a general purpose and systematic framework for adapting known (backbone) architectures to become invariant or equivariant to new symmetry types. Our framework builds on the well known group averaging operator that guarantees invariance or equivariance but is intractable. In contrast, we observe that for many important classes of symmetries, this operator can be replaced with an averaging operator over a small subset of the group elements, called a frame. We show that averaging over a frame guarantees exact invariance or equivariance while often being much simpler to compute than averaging over the entire group. Furthermore, we prove that FA-based models have maximal expressive power in a broad setting and in general preserve the expressive power of their backbone architectures. Using frame averaging, we propose a new class of universal Graph Neural Networks (GNNs), universal Euclidean motion invariant point cloud networks, and Euclidean motion invariant Message Passing (MP) GNNs. We demonstrate the practical effectiveness of FA on several applications including point cloud normal estimation, beyond $2$-WL graph separation, and $n$-body dynamics prediction, achieving state-of-the-art results in all of these benchmarks.

研究の動機と目的

• 置換やオイラー変換などの複雑な対称性に対して正確に不変または等長となるニューラルネットワークを設計する課題に取り組み、表現力と計算効率を維持すること。
• 大規模または連続的な対称性群における完全な群平均の非可算性を克服し、群要素の小さな、効率的に計算可能な部分集合による平均に置き換えること。
• GNN、ポイントクラウドネットワーク、MLPなど多様なアーキテクチャに適用可能な汎用的で体系的なフレームワークを提供すること。
• FAに基づくモデルがバックボーンアーキテクチャの最大表現力を保持することを証明し、重要な設定において普遍近似が可能であることを示すこと。
• 複数のベンチマークにおいてFAを実証的に検証し、近似平均化や完全群平均化のベースラインと比較して優れた性能と不変性の安定性を示すこと。

提案手法

• フレーム平均化は、対称性群Gのすべての要素に対する非可算な群平均を、入力Xに依存する小さな有限部分集合（フレームF(X)と呼ばれる）による平均に置き換える。
• フレームF(X)は、群作用における集合的等長性を満たすように構築され、これにより完全な群作用下での正確な不変性または等長性が保たれる。
• このフレームワークは、MLP（置換不変性）、GNN（グラフレベルの不変性）、メッセージパッシングGNN（オイラー変換不変性／等長性）など、さまざまなアーキテクチャに適用される。
• オイラー変換対称性（E(d)）の場合、フレームは安定化部分群から導き出し、幾何的不変量を活用して効率的な計算を可能にする。
• 特徴量が各層でフレーム全体にわたって平均化されるように変更されたメッセージパッシング機構が使用され、等長性を保ちつつ表現力を維持する。
• 理論的分析により、FAベースのモデルがバックボーンが普遍的である場合に普遍近似を達成することが証明され、フレームの選択が表現能力を低下させないことが示された。

実験結果

リサーチクエスチョン

• RQ1 完全な群平均の計算コストを負担せずに、正確な不変性または等長性をニューラルネットワークに強制する体系的なフレームワークを開発できるか？
• RQ2 小さな、データに依存する群要素の部分集合（フレーム）が完全な群平均に置き換え可能である条件は何か？ ただし、正確な対称性特性が保持されるものとする。
• RQ3 フレーム平均化は、対称性に配慮した学習タスクにおいて、下位のバックボーンアーキテクチャの表現力を維持または向上させるか？
• RQ4 近似平均化法（例：モンテカルロ法）や完全群平均化と比較して、FAは不変性の正確さとモデル性能の面でどのように差をつけるか？
• RQ5 FAは置換やオイラー変換などの多様な対称性群に一般化可能であり、GNN やポイントクラウドネットワークを含む幅広いアーキテクチャに適用可能か？

主な発見

• フレーム平均化は、計算的に非現実的な完全群平均を避けるために、入力に依存する小さなフレームによる平均化により、正確な不変性と等長性を達成する。
• FAベースのモデルは、バックボーンアーキテクチャの最大表現力を維持し、グラフおよびポイントクラウド学習タスクにおいて普遍近似が可能である。
• n体ダイナミクス予測タスクにおいて、FA-GNNはテストMSEが0.0057を達成し、パrameter数が同等のSOTAであるEGNN（0.0071）を20%以上上回った。
• 2-WLを超えるグラフ分離ベンチマークでは、FA-MLPおよびFA-GIN+IDが完全な分離を達成し、普遍的な表現力の証明となった。
• k=1のフレームサンプルを用いた近似FAは、k=1の完全群平均化（GA）よりも顕著に低い不変性誤差を示し、優れた安定性と一般化性能を示した。
• 実験的結果から、FAは低サンプル状態や高対称性環境において、近似平均化および完全群平均化の両方と比較して、より不変的かつ効果的であることが示された。