[論文レビュー] LieTransformer: Equivariant self-attention for Lie Groups

Michael Hutchinson, Charline Le Lan+4 more|arXiv (Cornell University)|Dec 20, 2020

Machine Learning in Materials Science参考文献 58被引用数 35

ひとこと要約

LieTransformerは自己注意をLie群へ拡張することで、入力をG上の関数へリフトし、Lie群とその離散部分群に対して等変なLieSelfAttention層を構築します。形状の数え上げ、QM9分子特性、ハミルトニアン動力学タスクで競争力のある結果を達成します。

ABSTRACT

Group equivariant neural networks are used as building blocks of group invariant neural networks, which have been shown to improve generalisation performance and data efficiency through principled parameter sharing. Such works have mostly focused on group equivariant convolutions, building on the result that group equivariant linear maps are necessarily convolutions. In this work, we extend the scope of the literature to self-attention, that is emerging as a prominent building block of deep learning models. We propose the LieTransformer, an architecture composed of LieSelfAttention layers that are equivariant to arbitrary Lie groups and their discrete subgroups. We demonstrate the generality of our approach by showing experimental results that are competitive to baseline methods on a wide range of tasks: shape counting on point clouds, molecular property regression and modelling particle trajectories under Hamiltonian dynamics.

研究の動機と目的

• Lie群で記述される群の対称性を動機づけ、活用して学習効率と一般化性能を向上させる。
• リフティングベースのフレームワークを用いて自己注意をLie群上で等変に拡張する。
• 回転・平行移動対称性を持つタスク（形状の数え上げ、分子特性予測、ハミルトニアンダイナミクス）で手法を実証する。
• LieTransformerがLie群とその離散部分群を扱えることを、競争力のある性能で示す。

提案手法

• 同次空間Xからリフティング演算子Lを介してLie群G上の関数へ入力データをリフトし、G等変な処理を可能にする。
• リフトされたドメインG上で内容と位置を用いた手がかりを用いてアテンションを行う LieSelfAttention を使用し、G の規則表現の下で等変性を保証する。
• Gまたはリフトされたドメインが連続的な場合（無限のG_fを持つLie群など）、積分を近似するためにモンテカルロサンプリングを用いる。
• LieSelfAttention、LayerNorm、MLPの残差ブロックを適用し、最後に不変なGプーリング層を介してタスク出力を生成する。
• リフティングが等変であることと、LieSelfAttentionが規則表現の下で等変であることを示して等変性の証明を行う。

実験結果

リサーチクエスチョン

• RQ1 自己注意を任意のLie群とその離散部分群に対して等変にすることは可能か？
• RQ2 リフティングベースのLieSelfAttentionアーキテクチャはSE(2)/SE(3)や他のLie群対称性を必要とするタスクで競争力のある性能を発揮するか？
• RQ3 形状の数え上げ、分子特性予測（QM9）、ハミルトニアンダイナミクスのタスクにおいて、非不変ベースラインやLieConvと比較してLieTransformerはどのような性能を示すか？

主な発見

• LieTransformerは形状の数え上げ、QM9、およびハミルトニアンダイナミクスのタスクで強力なベースラインに対して競争力のある性能を達成する。
• LieSelfAttentionはGの規則表現の下で証明可能な等変性を持つ（無限のGにはモンテカルロ近似を適用）。
• SE(3)不変の変種は多くの場合、平行移動のみの変種よりも優れており、回転不変性によって一般化が改善されるが、リフティングのサンプリングによるばらつきがある。
• ハミルトニアンダイナミクスでは、非不変ベースラインよりもデータ効率と一般化が著しく高く、いくつかの領域で1–3オーダーの差が出る。
• 同程度のモデルサイズと群設定の下でLieConvと比較してLieTransformerはしばしばより良い性能を示し、特にT(2)やSE(2)設定で顕著。