[論文レビュー] Fast, Expressive SE$(n)$ Equivariant Networks through Weight-Sharing in Position-Orientation Space

Erik J. Bekkers, Sharvaree Vadgama+3 more|arXiv (Cornell University)|Oct 4, 2023

Advanced Image and Video Retrieval Techniques被引用数 5

ひとこと要約

本稿では、位置・方位空間（R³×S²）における重み共有を用いた高速で表現力のある SE(n) 同変ネットワークを提案する。重み共有は、SE(n) 変換における点対の同値類を形式化することで実現される。メッセージ関数をこれらの同値類を一意に特定する幾何的属性に条件づけることで、原子間エネルギー予測、N体軌道予測、拡散モデルを用いた分子生成を含む 3D ポイントクラウドベンチマークにおいて、最先端の精度と速度を達成する。

ABSTRACT

Based on the theory of homogeneous spaces we derive geometrically optimal edge attributes to be used within the flexible message-passing framework. We formalize the notion of weight sharing in convolutional networks as the sharing of message functions over point-pairs that should be treated equally. We define equivalence classes of point-pairs that are identical up to a transformation in the group and derive attributes that uniquely identify these classes. Weight sharing is then obtained by conditioning message functions on these attributes. As an application of the theory, we develop an efficient equivariant group convolutional network for processing 3D point clouds. The theory of homogeneous spaces tells us how to do group convolutions with feature maps over the homogeneous space of positions $\mathbb{R}^3$, position and orientations $\mathbb{R}^3 { imes} S^2$, and the group $SE(3)$ itself. Among these, $\mathbb{R}^3 { imes} S^2$ is an optimal choice due to the ability to represent directional information, which $\mathbb{R}^3$ methods cannot, and it significantly enhances computational efficiency compared to indexing features on the full $SE(3)$ group. We support this claim with state-of-the-art results -- in accuracy and speed -- on five different benchmarks in 2D and 3D, including interatomic potential energy prediction, trajectory forecasting in N-body systems, and generating molecules via equivariant diffusion models.

研究の動機と目的

• 従来の畳み込み設定を超えて、メッセージ伝達ネットワークにおける重み共有を形式化すること。
• 同次空間理論を用いて、SE(n) 同変メッセージ関数の幾何的に最適な属性を導出すること。
• SE(n) 同変性を維持したまま、3D ポイントクラウド処理に適した効率的かつ表現力豊かなアーキテクチャを開発すること。
• 分子系や物理系を含む多様な 3D 学習ベンチマークにおいて、最先端の性能を実証すること。

提案手法

• SE(n) 変換における点対の同値類を一意に特定する属性にメッセージ関数を条件づけることで、重み共有を形式化する。
• 同次空間理論を用いて幾何的属性を導出し、方向情報と効率性の観点から、R³×S² が方向情報の最適特徴空間であることを同定する。
• R³×S² 上に分離可能な群畳み込みを構築し、方位とチャネルにわたる並列化を可能にすることで計算効率を向上させる。
• 球面畳み込みとバッチ処理されたチャネル混合を用いて、推論と学習の高速化を図る。
• 相対的位置、相対的な方位、相互作用特徴（例：電荷積）などの不変属性にメッセージ関数を条件づける。
• ConvNeXt ブロックと学習可能な球面埋め込みを用いた完全畳み込み型アーキテクチャを実装し、入出力特徴を処理する。

実験結果

リサーチクエスチョン

• RQ1 メッセージ伝達ネットワークにおける重み共有を、並進同変畳み込みを超えて、完全な SE(n) 同変性に一般化するにはどうすればよいか？
• RQ2 SE(n) 変換における点対の同値類を一意に特定する幾何的属性は何か？
• RQ3 なぜ、R³ や完全な SE(3) よりも、位置・方位空間（R³×S²）が SE(3) 同変学習においてより効率的で表現力に優れているのか？
• RQ4 完全に畳み込み型の SE(n) 同変アーキテクチャが、多様な 3D ポイントクラウドベンチマークで最先端の性能を達成できるか？
• RQ5 導出された属性にメッセージ関数を条件づけることで、同変メッセージ伝達における表現力と効率性がどのように向上するのか？

主な発見

• 提案手法は、原子間ポテンシャルエネルギー予測において最先端の性能を達成し、精度と推論速度の両面で先行手法を上回った。
• N体軌道予測において、ベースラインの同変モデルと比較して顕著に短い推論時間を達成し、最先端の結果を得た。
• 拡散モデルを用いた分子生成において、高品質で物理的に妥当な分子を生成し、サンプリング効率が向上した。
• 位置・方位空間（R³×S²）の使用により、完全な SE(3) インデクシングと比較して計算コストを削減しながら、方向感度を保持し、性能を向上させた。
• 理論的枠組みにより、幾何的属性と共有メッセージ関数のみを用いて、SE(n) 同変関数の普遍的近似が可能であることが示された。
• 実験的結果から、導出された属性にメッセージ関数を条件づけることで、全ベンチマークで一般化性能が向上し、収束が速くなった。