SE(3)不変とタンパク質立体構造予測③：球面調和関数とSE(3)同変性

前回

cake-by-the-river.hatenablog.jp

前回に引き続き、SE(3)不変/同変のための数学的な基礎を説明します。前回よりもっと難しいので、不明点はコメント等で教えてもらえるとありがたいです。

球面調和関数
- 球面調和関数による展開
- 球面調和関数とSO(3)同変性
Clebsch-Gordan係数
まとめ
補足
- 関数のヒルベルト空間と球面調和関数の完全直交性

球面調和関数

前回「軌道量子数」を紹介した時、なぜ量子？と思われた方がいると思います。実は、この軌道量子数は電子の軌道と対応していることが知られています。

量子力学において、電子は原子核の周囲を「回っている」と考えられています。つまり、何かしらの軌道に沿って動いていることになりますが、粒子がそのような運動（角運動）をしているならば、対応するエネルギーを持っているはずです。量子の世界では、エネルギーは連続的ではなく離散的にしか取れないため（これは、量子が「波」であり、特定のエネルギー以外は打ち消しあってしまうからです）、電子の軌道もそのエネルギーに対応して離散的に分布しています。

この離散的な電子軌道には、s軌道やp軌道といった名前がついています。化学の本で次のような画像を見たことがある人は多いと思います。

これらは、球面調和関数と呼ばれる関数です。球面調和関数は 1個, 3個, 5個, ...と種類が分かれていますが、この数値はSO(3)の各既約表現の次元と一致しています。ここでは証明などを省いてしまいますが、軌道量子数 $l$ で表される $2l+1$ 次元の既約表現に対し、対応する $2l+1$ 個の球面調和関数が得られます。

例えば、回転行列などの3次元の既約表現に対して、p軌道に対応する $l=1$ の3つの球面調和関数は

$Y^{(l)}_{m} : Y^{(1)}_{-1}, Y^{(1)}_{0}, Y^{(1)}_{1}$

と書かれます。2つ目の添え字 $m$ がマイナスから始まるのは、導出や量子力学における意味合いなどの都合によるものですが、 $m=-l, -l+1, \dots, l-1, l$ で $2l+1$ 個が表せる便利な記法と言えます。

球面調和関数による展開

さて、ここがポイントなのですが、ある既約表現に対応した表現空間へ出力する写像 $f(x)$ は、対応する軌道量子数の球面調和関数を「基底」として展開することが可能です。そのため、3次元ベクトルを入力とするSO(3)不変なあらゆる写像は、対応する各既約表現の球面調和関数の線形和（定数倍と足し算）によって表現できることになります。

恐らく今の説明はだいぶ飛躍しているように見えるでしょう。ここからは、この文章の意味をそれぞれ分解してゆっくり考えます。

まず、今考えている関数は3次元ベクトルを入力とし、何かしらの次元のベクトルを出力するものに限定しています。例えば、3次元空間上の任意の点を入力としたベクトル場（電場など）やテンソル場（応力場など）です。

ここで、3次元の入力ベクトル $(x,y,z)$ を ”球座標” すなわち原点からの距離 $r$ （動径方向）と対応する球面上の点 $(r,\theta, \phi)$ として考えます。

3次元の極座標。半径rと二つの角度θ,φによって位置が指定できる。Wikipediaより引用。

今、考えている関数 $f$ はこの3次元の点の変化によって値が変わるもの $f(r,\theta,\phi)$ です。そのため、角度を固定し半径 $r$ の動径方向の変化のみを考えた関数 $f(r, \theta=固定, \phi=固定)$ と、半径を固定し球面上を動かしたときの変化 $f(r=固定, \theta, \phi)$ の二つを考えることができます。

このとき、SO(3)の同変性を満たすような関数は、球面上の関数の変化（出力ベクトルの値の変化）と動径方向の関数の変化が独立しているはずです。もしこれらが独立していないとしたら、対象を回転することにより半径に応じた関数の変化が生じます。これの意味するところは「カメラを回転したらリンゴの大きさや色が変化する」のようなもので、明らかに同変性が崩れることになります。したがって、動径方向の関数と球面方向の関数の二つに分離して関数を書くことができるはずです。

$f(r,\theta,\phi) = d(r)Y(\theta,\phi)$

ここで、 $d(r)$ は半径のみに依存する回転不変な関数です。そのため、球面を入力とした関数 $Y(\theta, \phi)$ のみが（回転を表す）SO(3)の表現に対応する部分だと分かります。

この球面上の関数が、球面調和関数によって書けるはずだということです。詳しくは補足に示しますが、ある既約表現に対応した $2l+1$ 次元の表現空間の中では、無数にある球面上の関数は $2l+1$ 個の球面調和関数の線形和に分解して記述することが可能です。これは、「基底関数による展開」と呼ばれています。

さらっと言ってしまいましたが、これは結構凄そうに聞こえます。球面上のあらゆる点を入力とする関数は、例えば地球表面の気圧のようなもので、想像もつかないほどのパターンを含むはずです。その一見してコンピュータで扱いにくそうな無数の球面上の関数が、 $2l+1$ 個の球面調和関数を考えることで、その線形和の係数を並べた $2l+1$ 次元のベクトルに置き換えられる、ということです。例えば、3次元の既約表現（回転行列など）に対応する関数の場合、 $l=1$ の3つの球面調和関数 $Y^{(1)}_{-1}, Y^{(1)}_{0}, Y^{(1)}_{1}$ の線形和

$f(x) = a_{1,-1} Y^{(1)}_{-1} + a_{1,0} Y^{(1)}_{0} + a_{1,1} Y^{(1)}_{1}$

で書けます。そのため、コンピュータの内部では

$\left( \begin{array}{c} a_{1,-1} \\ a_{1,0} \\ a_{1,1} \end{array} \right)$

という3次元のベクトルを持っておけば、あとは球面調和関数をライブラリなどで計算して処理すれば事足ります。

これは余談ですが、同じように「関数を展開」することができる手段は他にも色々と存在します。例えば、フーリエ変換では、関数を各周波数に対応した波（モード）で分解し、周波数を入力とした関数に置き換えられます。大縄跳びのようにヒモを振り回したときに出来るサインカーブで、あらゆるヒモの形式を表現する、といったイメージです。各周波数ごとのモードに分解した後は、対応する周波数での数値が得られるため、それらをベクトルのように並べていくことも出来ます。このとき、周波数に制限を加える（例えば、デジタル音楽は44,100Hzのサンプル数に限定されています）ことで、有限個の周波数のみで関数を表すことになり、対応するベクトルも有限次元に収まります。

球面調和関数による展開は、これとほとんど同じことをしています。ただ、周波数の制限の代わりに、既約表現の表現空間という分かりにくい制限を採用しているだけです（もっと言えば、実は球面上でフーリエ変換を行うことと球面調和関数で展開することは同じことであり、球面上の"モード"が球面調和関数に対応していることが分かっています）。

一部の人は「なぜ想像もつかないほどたくさんある球面上の関数が高々 $2l + 1$ 個のベクトルに収まるのか？」という風に考えるかもしれません。そしてその考えは正しいですし、上の説明は誇張したところがあります。つまり、 $2l + 1$ 個のベクトルで表せる関数は、 $2l + 1$ 次元の既約表現で扱える程度の関数でしかないからです。フーリエ変換の例も、ローパスフィルターが低周波数のうちに表現できる関数が対応していましたが、同じように低軌道量子数の関数しか対応していないのです。しかし、実際に考える関数においては、それでも十分な場合があります。人間も44,100Hz程度で感じ取れる周波数が限界を迎えるためにそれ以上の周波数は切ってしまってよいのです。いずれにせよ、球面調和関数を用いてSO(3)不変/同変な関数を考えるのが、計算科学など人間が取り扱う観点で非常に理にかなっていると言えます。

ところで、球面調和関数は球面を入力とする関数だと言いましたが、先ほどの電子軌道は全く球面の上にあるようには見えません。それは、球面調和関数が本来は複素関数であり、電子軌道として紹介されるものは実数のみの関数となるようにそれらを組み合わせたものだからです。私たちが実際に取り扱う際も、実数化した関数を用いたところで本質的には何の問題がないため、都合の良い実数による表記を採用しようと思います。

球面調和関数とSO(3)同変性

球面調和関数 $Y^{(l)}_{m}$ は、入出力空間の既約表現に対応したSO(3)に同変な関数となっています。

今考えている関数は3次元を入力とするため、入力ベクトル（3次元）の既約表現は回転行列 $T_g = R_g$ に対応しています。また、SO(3)の既約表現はWignerのD行列 $S_g = D^{l}_{m,m'} \in (2l+1)\times(2l+1)行列$ を用いて表すことができるのでした。これらと球面調和関数 $Y^{(l)}_{m}$ の間には、同変の定義式

$Y^{(l)}_{m} (T_g[x]) = S_g [ Y^{(l)}_{m}(x)]$

が成り立ちます。既約表現に具体的に代入し展開した形としては、

$\displaystyle Y^{(l)}_{m} (R_g x) = \sum_{m'}{D^{(l)}_{m,m'} Y^{(l)}_{m'}(x)}$

が成立します。そのため、球面調和関数の線形和によって記述された関数もSO(3)に同変と言えます。これは、球面調和関数に基づいて展開した特徴量ベクトルをいくら弄っても、最終的に球面調和関数との線形和で関数を記述している分には、SO(3)同変性が崩れないことを意味しています。また、既約表現の直和で表される表現空間での関数も、これらの関数の組み合わせ（対応する空間での関数の結合）としてSO(3)同変性を保つことができます。

Clebsch-Gordan係数

さて、SO(3)に同変な関数 $f(x)$ の出力ベクトルについて、その各既約表現に対応する表現空間へ出力するパートは、球面調和関数を用いることで機械学習でよく用いる特徴ベクトルの形式で記述可能であることを前節で述べました。しかし、私たちが最終目的としている不変性・同変性を保つ機械学習では、例えばレイヤーを複数重ねたり、畳み込んだりするなど、特徴ベクトルをどんどんと変化させる必要があります（これは関数の合成に対応しています）。特に、畳み込みのような計算は特徴ベクトルの次元（すなわち表現空間の次元）を変化させることが多々あります。こうした操作を行った結果、次元の変化に応じて表現空間とその既約表現は毎回様変わりすることが予想されます。

どのような表現空間であっても、SO(3)同変であれば既約分解を通じて球面調和関数で書けることは明らかとなっていますが、関数の合成を行うことで既約表現の組み合わせが変化するかはまだ説明していません。つまり、どのように合成後の関数を機械的に（コンピュータ的に）既約分解すればよいのかが分かっていません（前回挙げた複数の例は全て私たちの直感を使って解いていました）。

この問題は、角運動量の合成という分野でよく研究された内容であり、実際にはClebsch-Gordan係数（CG係数）と呼ばれる数値によって対応付けられることが知られています。

まずは、今問題となっていることをより数学的に示しましょう。SO(3)同変なある関数 $f(x)$ について、別な関数 $g$ を合成してまた新しい関数 $h(x)$ を作ります。

$h(x) = g(f(x))$

この関数 $h(x)$ がSO(3)同変であるとするとき、関数 $g$ が満たすべき条件は何でしょうか？ $h(x)$ の出力ベクトルは $f(x)$ の出力ベクトルと異なる次元を取りうるため、作用する表現空間は異なる可能性が高いです。

ここで、次元が変化するような関数の合成の中では、2つのSO(3)の既約表現を『合成』した空間を相手にしている、ということに注目します。ここでいう『合成』は、厳密には外積やテンソル積と呼ばれるものです。ただ、そこまで深く考えなくてもよく、「2つの空間のペア」という空間に変化した、と思えば大丈夫です。例えば1次元の空間（数直線）と2次元の空間（平面）のペアは3次元の空間に相当しますが、それと類似したものとして考えるので結構です（勿論色々と異なる点はあります）。そして、合成後の関数もSO(3)同変を仮定するので、これも既約分解が可能なはずです。

結論から言うと、CG係数とは合成された空間を改めて既約表現の直和で表す際に必要な係数、ということです。数式で表すならば、2つの既約表現 $D^{(k)}, D^{(l)}$ のテンソル積（ $\otimes$ で表します）からなる空間は、CG係数 $Q_{lk}$ を用いて

$\displaystyle D^{(k)} \otimes D^{(l)} = Q_{lk}^{T} \left[\bigoplus^{k+l}_{j=|k-l|} D^{(j)}\right] Q_{lk}$

として表されます。この式の形は、相似変換であり、CG係数も直交行列の一種であることが分かります。前回説明したように、あらゆるSO(3)の表現はこうした直交行列の相似変換で書けるため、CG係数はそのような計算処理に必要な直交行列のことであると言えます。

上の例のCG係数。行が合成後の既約表現、列が合成前の既約表現のペアを表す。色がついていないところはゼロ。

CG係数を具体的に計算する理論は面白いですが、ここでは手に負えないほどの量子力学の知見を必要とするため、省略します。重要なポイントは、ある程度の次元の既約表現についてはすでに計算されたCG係数が用意されており、それらを利用すれば計算上は手軽に『合成』することができる、ということです。

これは関数の合成をする際にはかなり心強いことです。つまり、関数 $g$ がどのようなものであっても、SO(3)の同変性を保とうとするならば、必ず $f(x)$ と $h(x)$ の既約表現間をCG係数によって結びつけるようなものである、ということが保証されます。裏を返せば、そのような関数に限定しておけば、むしろどのような関数を設定してもよく、それを学習対象とすることが可能であるということです。これが、SE(3)同変な機械学習を考える際の根本的なアイデアとなります。

余談ですが、「角運動量の合成」という問題は、量子力学に由来するものです。電子を始めとした量子は、軌道に沿った運動（角運動）に基づき角運動量を持ちます。こうした回転運動は当然座標を回転しても不変であるため、角運動量も回転に不変であり、SO(3)の既約表現の組み合わせで表されます。例えば、二つの電子の角運動量を合成したものは、（量子の世界ではエネルギーが離散的なことを踏まえると）別のエネルギー状態に対応した新しい角運動量（と軌道）の形態をとるはずであり、それもまた回転に不変です。

また、ここまでに紹介した角運動量は原子核に対する「公転」のようなものでしたが、電子の「自転」に対応する角運動量（スピンと呼ばれています）も存在します。この「公転」と「自転」の角運動量の合成（LS結合）も、原子核物理・物性物理・量子化学・分子軌道論などありとあらゆる分野の基本に通じています。

いずれにせよ、これらの理論の基本は二つの角運動量の（テンソル積による）合成であり、CG係数はそのような研究の過程で発見されたものです。そもそも、SO(3)を始めとしたLie群論の多くも、こうした量子力学による長年の研究の積み重ねで発展した重厚な（少なくとも数学科でも一大テーマとなるような）理論分野であり、これらの研究分野の基にタンパク質立体構造予測という真新しい問題が解決されているところは、学際的な面白さがあります。

まとめ

２回分のまとめとして、前回の最初に挙げた問題とその答えを示しましょう。

入力空間での操作 $T_g$ 、および対応する出力空間での操作 $S_g$ としてどのようなものを選ぶか

今回の入力空間は3次元ベクトルであり、その操作は回転行列により表現されていました。一方、どのような出力空間であっても、何かしらの軌道量子数を持つ既約表現の直和を用いることで、SO(3)の表現となりました。したがって、これらを選ぶのが適切だと分かります。

どう写像 $f$ を構築すればよいか

既約表現の表現空間の中では、球面調和関数を用いている限り、SO(3)の同変性を担保できていました。したがって、実際に連続で無限個存在する関数を考える必要はなく、球面調和関数のもとに展開、すなわちベクトルとして関数を取り扱って良いことになります。また、それらの関数の合成等、次元の異なる表現空間のベクトル間を結ぶ際には、Clebsch-Gordan係数を用いてスムーズに演算可能でした。これらを機械学習モジュールの計算の基礎に据えることで、SO(3)の同変性を保つ計算が可能であることが分かりました。

「SO(3)の同変性は分かったけど、やりたいのはSE(3)の同変性では？」と思った方がいるでしょう。実は、SE(3)の同変性はSO(3)の同変性に「相対ベクトルで考える」ということを追加で意識するだけで解決します。SE(3)は回転と平行移動の組み合わせからなりますが、お互いは独立した動作にできます。より群論的に言うと、SE(3)はSO(3)と「平行移動群」の直和で表現できます。

したがって、SE(3)の同変性は、SO(3)に同変かつ平行移動群に同変、という条件でクリアできます。平行移動で同変性を保つ方法は、単に絶対座標系ではなく相対座標系を用いるだけで解決することは容易に想像できます。したがって、そのように入力を制限してしまえば、ほとんどの場合にクリアできることが分かります（実際に満たすかどうかはもちろん証明した方が良いですが）。

こうして、SE(3)同変な関数・機械学習を考える基礎ができました。次回は、これらの知見を実際にニューラルネットワークの形で導入した Tensor Field Network, SE(3)-Transformers などを解説していこうと思います。

cake-by-the-river.hatenablog.jp

補足

関数のヒルベルト空間と球面調和関数の完全直交性

後で書きます。