【論文紹介】LieGG: Studying Learned Lie Group Generators

この記事は BrainPad Advent Calender 2022 の2日目の記事となります。

今回もまた社内の論文紹介で話した内容の公開記事です。またまた機械学習と対称性の話ですが、今回はデータを学習させたモデルから対称性を見つけ出す、という取り組みがちょっと興味深く、あまり広く知られていないテーマでもあると思いましたので読んでみました。

概要

データが何らかの(例えば回転や並進などの)対称性を持つとき、それらの変換に対しての対称性をモデルに課すことでモデルの精度を高めていこう、という取り組みは色々と行われています(例えば以前紹介した E(n) Equivariant Graph Neural Networks もその一つ)。

pseudo-theory-of-everything.hatenablog.com

この論文では、モデルに対称性を課すのではなく、データが何らかの対称性を持つときにそれを学習したモデルから、対称性を抜き出す(！？)、という手法を提案しています。
対称性について簡単に述べておくと、ここではリー群の作用に対する不変性や同変性(共変性)のことを指しています。対称性を抜き出す、というのは学習したモデルからリー群の生成子を見つけ出す作業に対応しています。

データ分析系の論文紹介でリー群・リー代数を既知とします、というのは流石に優しくないと思うので、イメージを掴める程度にそれらを紹介し、その上で論文で提案された手法の説明に入ります。

リー群と生成子

リー群

数学的な正しさに拘らず、具体例を使ってリー群、生成子について説明します。
次の \(2 \times 2\) 行列
\begin{align}
\begin{pmatrix}
\cos\theta & - \sin\theta \\
\sin\theta & \cos\theta
\end{pmatrix}
\end{align}
はパラメータ \(\theta\) を与えると、その角度だけ2次元平面の回転として作用する行列です。異なる \(\theta\) によって異なる角度の回転を表しますが、この連続的なパラメータ \(\theta\) でラベルされる要素全体の成す群*1を特殊直交群 \(SO(2)\) と呼びます。

この他にも、3次元空間の回転に対応する \(SO(3)\) 、複素数 \(z\) の位相回転に対応する \(U(1)\) などがあり、これらのように連続的なパラメータで要素がラベルされるような群、をリー群とイメージしていただければよいかと思います。
(この手の記事にはつきものになってしまいますが、イメージを優先して正しさを犠牲にしていますので、きちんと理解したい場合は教科書を参考にしてください)

生成子

リー群 \(G\) の元は連続パラメータでラベルされるので、無限個の元がありますが、
これらの元は有限個の生成子 \(\boldsymbol{\mathfrak{h}} = (\mathfrak{h}_1, \ldots, \mathfrak{h}_n)\) から、
\begin{align}
e^{\boldsymbol{t} \cdot \boldsymbol{\mathfrak{h}}} \in G
\end{align}
で与えられることが知られています。ただし、 \(\boldsymbol{t} = (t_1, \ldots, t_n) \in \mathbb{R}^n\) はパラメータです。

ここでも \(SO(2)\) を例に具体的な生成子の例を紹介します。天下り的ではありますが、

\begin{align}
H =
\begin{pmatrix}
0 & -1 \\
1 & 0
\end{pmatrix}
\end{align}
という行列を考えてみます*2。このとき、
\begin{align}
H^2 =
\begin{pmatrix}
-1 & 0 \\
0 & -1
\end{pmatrix}, \quad
H^3 =
\begin{pmatrix}
0 & 1 \\
-1 & 0
\end{pmatrix}, \quad
H^4 =
\begin{pmatrix}
1 & 0 \\
0 & 1
\end{pmatrix}, \quad
H^5 = H
\end{align}
が成り立つので、パラメータを \(\theta\) として、この行列を指数関数の肩に乗せた \(e^{\theta H}\) を計算してみると、

\begin{align*}
e^{\theta H}
&= \sum_{n=0}^\infty \frac{1}{n!}(\theta H)^n \\ \nonumber
&=
\begin{pmatrix}
1 & 0 \\
0 & 1
\end{pmatrix}
+ t
\begin{pmatrix}
0 & -1 \\
1 & 0
\end{pmatrix}
+ \frac{t^2}{2!}
\begin{pmatrix}
-1 & 0 \\
0 & -1
\end{pmatrix}
+ \frac{t^3}{3!}
\begin{pmatrix}
0 & 1 \\
-1 & 0
\end{pmatrix}
+ \frac{t^4}{4!}
\begin{pmatrix}
1 & 0 \\
0 & 1
\end{pmatrix}
+ \cdots \\
&=
\begin{pmatrix}
\sum_{n=0}^\infty (-1)^n\frac{t^{2n}}{(2n)!} & -\sum_{n=0}^\infty (-1)^n\frac{t^{2n+1}}{(2n+1)!} \\
\sum_{n=0}^\infty (-1)^n\frac{t^{2n+1}}{(2n+1)!} & \sum_{n=0}^\infty (-1)^n\frac{t^{2n}}{(2n)!}
\end{pmatrix} \\
&=
\begin{pmatrix}
\cos\theta & - \sin\theta \\
\sin\theta & \cos\theta
\end{pmatrix}
\end{align*}

となり、\(SO(2)\) の元が得られることがわかります。つまり、上で導入した行列 \(H\) は \(SO(2)\) の生成子でした。

LieGGのベースとなる考え

以上の背景をもとに、この論文で提案するリー群生成子を求める手法、LieGG(Lie Group Generators) 基本的な考え方を説明します。LieGG は学習済みのニューラルネットからリー群の生成子を抽出する、という手法になります。

データセットの識別関数 \(f : \mathbb{R}^n \rightarrow \mathbb{R}\) を、データセット \({\cal D}\) に含まれる \(\boldsymbol{x} \in {\cal D}\) に対してのみ \(f(\boldsymbol{x}=0)\) を満たす関数として定義します。識別関数はダウンストリームタスクの対称データにフィッティングしたモデルを考えることで自然に導入できます。ここで、次の定理*3を使うことで、生成子を具体的に求めるための方程式を求めることができるようになります。

定理

連結リー群 \(G\) がn次元多様体 \({\cal X}\) に線形に作用する場合を考える。写像 \(F : {\cal X} \rightarrow \mathbb{R}^l, l\leq n\) が連立代数方程式

\begin{align}
F_\nu(\boldsymbol{x}) = 0, \quad \nu = 1, \ldots, l
\end{align}

を与え、またこの連立方程式は最大ランクを持つ(全ての解に対してヤコビアンのランクが \(l\) )であると仮定する。このとき、群 \(G\) がこの方程式系の対称群であるのは、以下が成り立つときに限る：
\begin{align}
\sum_{i=1}^{n}\sum_{j=1}^{l} \frac{\partial F_\nu}{\partial x_i} \cdot \mathfrak{h}_{ij} \cdot x_j = 0, \quad
\mbox{whenever} \quad F_\nu (\boldsymbol{x}) = 0
\end{align}

ただし、 \(\mathfrak{h}_{ij}\) は \(G\) の生成子を行列表示したときの \(i,j\) 成分。

難しいのでこちらも具体例を考えます。
わかりやすくデータセットが \(n - 1\)次元の球面上に分布している場合を考えます :
\begin{align}
{\cal D} = \{\boldsymbol{x} \in \mathbb{R}^n : x_1^2 + \cdots + x_n^2 = 1\}
\end{align}

球面上に分布しているデータなので、このデータは \(n\) 次元回転に対する対称性を持っていることがわかります。またデータの識別関数は \(f = x_1^2 + \cdots + x_n^2 - 1\) であり、\(\frac{\partial f}{\partial x_i} = 2x_i\) となるので、求める生成子 \(H\) の \(i,j\) 成分を \(h_{ij}\) と表すと、定理より
\begin{align}
\sum_{ij} x_i \cdot h_{ij} \cdot x_j = 0
\end{align}
が得られます。この方程式を満たす \(H\) は \(H + H^T = 0\) すなわち反対称行列になります。実は反対称行列は直交行列のなす群 \(O(n)\) の生成子となっています*4。
\begin{align*}
O &= e^{tH}, \quad O^T = (e^{tH})^T = e^{tH^T} = e^{-tH}, \quad \therefore OO^T = 1
\end{align*}

LieGG

Figure 1 (論文より引用)

識別関数 \(f\) が与えられたとき、方程式

\begin{align}
\sum_{i,j}^n \frac{\partial f}{\partial x_i} \cdot \mathfrak{h}_{ij} \cdot x_j = 0 ,
\quad \mbox{for}~ x_i \in {\cal D} \label{sys_eq}
\end{align}

を解くことで生成子 \(\mathfrak{h}_{ij}\) を求める、というのが基本的なLieGGの考え方となります。
では、具体的にこれを解くための手順を説明していきます。求める変数は \(n^2\) 個の \(\mathfrak{h}_{ij}\) なので、これを成分とした縦ベクトルを \(\bar{\mathfrak{h}}\) と表すと、方程式(\(\ref{sys_eq}\)) を行列の固有方程式の形に書き直すことができます。
\begin{align}
{\cal E} \bar{\mathfrak{h}} = 0
\end{align}

この論文では \({\cal E}\) を network polarization matrix と呼んでいます(ので、この記事ではこれを分極行列と呼ぶことにします)。分極行列 \({\cal E}\) はデータセットの数 \(|{\cal D}|\) を行数、 \(n\) を列数とした \(|{\cal D}| \times n\) 行列なので、この方程式の解を求めるために特異値分解
\begin{align}
{\cal E} = USV^T
\end{align}
を行い、ゼロ特異値に対応する \(V\) の特異ベクトルを求めます。

\(\mathbb{R}^2\) に作用する群の生成子

画像分類のようなタスクに対して LieGG を使って対称性を見つけ出すことを考えてみましょう。このとき、2次元平面に並進や回転のような群 \(G\) の作用が画像を変化させる、ということを考えることになります。
方程式を立てるために、きちっと「画像」や「モデル」を数学的な言葉に焼き直してみます。まず、画像は平面上のコンパクトな部分空間 \(\Omega \subset \mathbb{R}^2\) 上に値を取る関数 \(f\) とします。具体的にはには2次元平面上でピクセルの位置 \(\boldsymbol{x}\) を与えると数値が返ってくる、という関数です。画像全体の空間は \({\cal I} := \{f, f:\Omega \rightarrow \mathbb{R}\}\) と表すことができるので、画像タスクのモデルは \(F : {\cal I} \rightarrow \mathbb{R}\) という関数だと考えることができます。ピクセルのグリッド \( (\boldsymbol{x}_1,\ldots,\boldsymbol{x}_n) \) を考えると、画像は \(\mathbb{R}^n\) 上の点 \( (f(\boldsymbol{x}_1),\ldots,f(\boldsymbol{x}_n)) \) で与えられ、モデルはn個の引数を取る \(F(f(\boldsymbol{x}_1),\ldots,f(\boldsymbol{x}_n))\) となるので、このセッティングに対して方程式(\(\ref{sys_eq}\)) を立ててみます。

\begin{align}
 &F(f(e^{t\mathfrak{h}}\boldsymbol{x}_1),\ldots,f(e^{t\mathfrak{h}}\boldsymbol{x}_n)) = 0 \nonumber \\
\Leftrightarrow &
\sum_{p=1}^n\sum_{k,j=1}^2
\frac{\partial F}{\partial f(\boldsymbol{x_p})}
\frac{{\partial f(\boldsymbol{x_p})}}{\partial x^k_p}
\mathfrak{h}_{kj} x^j_p = 0 \label{image_eq}
\end{align}

実際にこれを解くためには以下のステップを行います。

1. グリッドをガウシアンフィルタで平滑化する(元の理論が連続値に対してのみ定義されているのでこれをするらしい)
2. ニューラルネットワークを学習する
3. 方程式(\(\ref{image_eq}\))から分極行列を計算する
4. (近似的な)ゼロ特異値に対応する特異ベクトルを求める

対称性バリアンスとバイアス

ここまでLieGGを使って学習済みモデルからデータの持つ対称性を抽出する方法を説明してきました。実際のデータではなかなかきれいな対称性が見えることも少ないと考えられるので、指標として次のようなものを考えます。

1. 学習したモデルがどのくらい変換に対して不変であるか
2. 仮にデータの持つ真の対称性がわかっているとき、得られた対称性はどれほど元の対称性を再現しているか

1つ目を測る指標を対称性バリアンス(symmetry variance)、2つ目を対称性バイアス(symmetry bias)と呼びます。

まず、変換に対するモデルの不変性を測る指標を次のように定義します。
\begin{align}
  R(g) := \mathbb{E}(F(g\boldsymbol{x}) - F(\boldsymbol{x}))^2
  \sim \frac{1}{|{\cal{D}}|}\sum_{\boldsymbol{x}_i\in{\cal D}}
  (F(g\boldsymbol{x}) - F(\boldsymbol{x}))^2
  =: \tilde{R}(g), \quad g \in G
\end{align}

ここで、差分 \(F(g\boldsymbol{x}) - F(\boldsymbol{x})\) をパラメータ \(t\) で展開することで、分極行列を使った次の表記を得ることができます。
\begin{align}
\tilde{R}(g) = \frac{1}{|{\cal{D}}|}(t^2 ||{\cal E}\bar{\mathfrak{h}}||^2_2 + {\cal O}(t^3))
\end{align}

さらに特異値分解 \({\cal E} = USV^T\) を行うことで、指標の近似値として次の形が得られます。
\begin{align}
\tilde{R}(g) \sim \frac{1}{|{\cal{D}}|}||{\cal E}\bar{\mathfrak{h}}||^2_2
= \frac{1}{|{\cal{D}}|} \sum_l s^2_l (V^T \bar{\mathfrak{h}})^2_l \label{inv_r}
\end{align}

ここで、 \(s_l\) は降順に並べた \(l\) 番目の特異値です。

モデルの変換に対する不変性が特異値で展開できる事がわかったので、モデルの不変性を測る対称性バリアンス \({\cal V}({\cal E})\) を特異値の最小値を使った次の形で定義します。
\begin{align}
{\cal V}({\cal E}) &= \frac{s_{\mbox{min}}({\cal E})^2}{|{\cal D}|}
\end{align}

例えばゼロ特異値があれば対称性バリアンスはゼロになり、対応する特異ベクトルから対称性を再現する生成子が作れますが、ゼロでない場合はどのくらい不変である状況に近いか、という量になるイメージですね。

もう一つ、今度は (\(\ref{inv_r}\)) の特異ベクトルに着目してみます。今度は仮にデータがリー群で表現できる対称性を持っていることがわかっているときに、学習したモデルがきちんと対称性を再現できているかを測る対称性バイアス \({\cal B}_i({\cal E},{\cal G})\) を
\begin{align}
{\cal B}_i({\cal E},{\cal G}) &= ||\bar{v}_i({\cal E}) - {\cal G}||_F
\end{align}
と定義します。\(\bar{v}_i({\cal E})\) は \(i\) 番目の特異ベクトルを行列に戻したもので*5、\({\cal G}\) は真の対称性の生成子です(複数個ありますがノルムが一番小さくなるものを採用するようです)。

普通はモデルの学習で元のデータが対称性を持つかどうかなどあまり気にしないと思うのですが、データに対称性があるかわからなくても、対称性バリアンス自体は計算できて、対称性に関する情報を引き出せるのは面白いな、と思います。

実験

論文では考案された LieGG について擬似データや MNIST で実験を行い、実際に対称性を再現できるか、といったことや、モデルの構成と対称性バリアンス・バイアスの関係などを調べています。各実験について簡単に紹介します。実験のセットアップなど、詳細については原論文を確認していただければと思います。

人工的なデータから対称性を再現

Finzi, Welling, Wilson*6で提案された \(O(5)\) の対称性を持つ回帰問題
\begin{align}
f(\boldsymbol{x}_1, \boldsymbol{x}_2) =
\sin(||\boldsymbol{x}_1||) - 0.5||\boldsymbol{x}_2||^3
+ \frac{\boldsymbol{x}_1^T \boldsymbol{x}_2}{||\boldsymbol{x}_1||||\boldsymbol{x}_2||}
\end{align}

を考えます。この関数から学習用データを生成し、多層パーセプトロン(MLP)でモデルを作成します。目的は \(O(5)\) 対称性をこのモデルから抽出できるかを確認することで、対称性バイアス、対称性バリアンスを計算し確かめます。

Figure 2 (論文より引用)

結果は論文の Figure 2 のグラフのようになります。
特に注目したいのは対称性バイアスで、15番目以降の特異ベクトルではバイアスがゼロ、つまりよく対称性を再現できていることがわかります。

Rotation MNIST を使った対称性の再現

Rotation MNIST*7はその名の通り MNIST の画像を回転させたものを含むデータセットのようで、平面の回転なので \(SO(2)\) の対称性があるはずです。こちらのデータを2層と6層のパーセプトロン(パラメータ数はそれぞれ8,24)で学習し、 LieGG で生成子を作成と対称性バイアスを計算します。この実験では、MLPの層の数の違いで対称性の再現に差が出るか、ということの確認も目的となります。

Figure 3 (論文より引用)

結果は論文の Figure 3 のようになります。
(a) の図を見ると対称性バイアスについて、2層より6層のパーセプトロンの方がバイアスが小さいことが示されています。実際にそれぞれのモデルから生成子を作り、円と MNIST のデータに作用させた結果が (b) となります。2層(上の図)では円の輪郭がぼんやりするくらいには変換にずれが発生していることがわかり、MNIST のデータも回転角が大きくなるにつれて縮小しているように思えます。一方で6層の方(下の図)では比較的きれいに回転の変換ができているように見えます。

ネットワークの構成と対称性

こちらの実験ではより詳しく、層の数の違い、パラメータの違いの影響を確認します。
また、モデルの精度と対称性の再現の関係についても調べます。モデルはパラメータ数を20,000ずつ変化させ、40,000〜200,000パラメータを、隠れ層が1〜5のMLPで作成しています。

Figure 4 (論文より引用)

結果は論文の Figure 4 です。
(a) の図からはパラメータが多くなるほど対称性の再現が良いこと、一方で層が深くなればよく再現できるというわけではない、ということが言えそうです。

また、「モデルの精度が悪ければ対称性の再現なんてできないのでは？」という疑問に対する実験結果が図の (b) です。散布図の色はパラメータの数を表しています。こちらは想像通りというか、Accuracy が高い方が対称性バリアンスが小さい傾向が見えます。
ただし、同じ Accuracy でもバリアンス・バイアスに差が出ていることもわかり、同一精度のモデルでも対称性を上手く抽出できるかに違いが出ることもありそうです。

終わりに

モデルに制約を課すタイプの対称性の話はよく見かけていたのですが、今回のように特に制約なく学習させたモデルからデータの持つ対称性を見出す、という話はあまり知らなかったので、とても興味深かったです。現状実用的な技術の紹介ではないのですが、この辺りの機械学習と対称性や幾何学を絡めた話はとても面白く思えます。いつかどこかで参考になれば幸いです。ありがとうございました。

*1:群というのは集合の元同士に掛け算を導入したもの、くらいのイメージをしてください。ちゃんと知りたい方はググればすぐわかります！

*2:一般的な定義とは符号が逆になっている可能性が高いです

*3:Peter J Olver. Applications of Lie groups to differential equations, volume 107. Springer Science & Business Media, 2000.

*4: \(SO\) じゃないのかい！？ということについてざっくり説明すると、鏡映の対称性も持っているからです。詳しくは教科書を......

*5: \(\bar{h}\) を作ったのと逆の操作で戻すのだと思います

*6:Marc Finzi, Max Welling, and Andrew Gordon Wilson. A practical method for constructing equivariant multilayer perceptrons for arbitrary matrix groups. In International Conference on Machine Learning (ICML), 2021.

*7:Hugo Larochelle, Dumitru Erhan, Aaron Courville, James Bergstra, and Yoshua Bengio. An empirical evaluation of deep architectures on problems with many factors of variation. In Proceedings of the 24th International Conference on Machine Learning (ICML), 2007.