Fire Engine

消防士→ITエンジニア→研究者

Equivariant Neural Networksの概論：群論を用いた深層学習の設計の進展

データサイエンス 論文

最近、Equivariant Neural Networksというデータの対称性に着目した深層学習の設計の新しいパラダイムの存在を知り、非常に興味を持っている。Equivariant Neural Networksは、深層学習の更なる汎化性能の向上や学習の効率化、適用分野の拡大などを実現する可能性を秘めている。今回は、Equivariant Neural Networksとは何なのか、これまでの研究事例、どういう分野で応用できるのかなど、最近私が学んだことを概論的にまとめていく。

目次

Equivariant Neural Networksの概要

ニューラルネットワーク（以降、NN）を用いた深層学習は、対象とする問題やデータに応じて最適な設計を行う必要がある。問題やデータが持つ固有の構造や制約は帰納バイアス（inductive bias）と呼ばれ、学習データ以外の有効な事前知識として深層学習の設計に積極的に活用される傾向にある。Equivariant Neural Networksとは、データに潜む対称性を強力な帰納バイアスとして捉え、それをネットワークの設計に取り入れたNNである。

対称性とは、ある対象に特定の操作を適用しても結果が変わらない性質のことであり、身の回りに多く存在する。例えば、手書き数字のデータセットであるMNISTを用いた画像分類問題を例に挙げて考えると、下のように「7」という数字を上下左右に並進移動させても、回転させても、それが「7」であるという分類の結果には変わりがない。これをデータは並進移動や回転に対して対称であるという。

手書き数字画像に対する並進移動と回転の例

Equivariant Neural Networksでは、これらの対称性を不変性（Invariance）を一般化した同変性（Equivariance）として表す。不変性と同変性について詳細な定義は後述する。Equivariant Neural Networksはそれ自体が新しいネットワークの設計を生み出すものであるが、既存のNNの設計を対称性を通して統一的に解釈できるという側面もある。例えば、畳み込みニューラルネットワーク（以降、CNN）の畳み込み層は並進移動に対する同変性を持つし、TransformerのSelf-Attention層は置換に対する同変性を持っている。したがって、Equivariant Neural NetworksはNNの設計の新しい指針を与えるだけでなく、これまでのNNのアーキテクチャに対して俯瞰的な視点を与えてくれる。

不変性と同変性

不変性と同変性について数式を用いて説明する前に、それらがどういった性質を持っているかを図を使って簡単に説明する。関数 f が入力として数字の配列を受け取り、その配列の要素の置換に対する不変性と同変性の例を下に示す。

置換に対する不変性と同変性の例

不変性は、入力配列の順序に関わらず同じ値を返す、つまり置換という操作に対して結果が変わらない。これは例えば、最大値の計算や総和の計算が該当する。次に、同変性を見てみると、入力配列の順序に対応して、出力配列の順序が変わっている。このように入力に応じて「同じように変わる性質」が同変性であるというイメージを持っておくとよい。

ここまでで不変性と同変性のイメージは共有できたと思うので数式を用いた説明に移る。下の図のように入力 xに対してある操作 g（例えば、並進移動や回転など）と変換 \phi（例えば、NNの畳み込み層など）を適用することを考える。

変換 \phiが操作 gに対して同変・不変である場合

以下の式を満たすような g' が存在するとき、「変換 \phi は操作 g に対して同変である」という。

\displaystyle \phi(g(x))=g'(\phi(x))


つまり、入力 x に操作 g を適用した結果 g(x) に対して変換 \phi を適用した結果 \phi(g(x)) が、入力 x に変換 \phi を適用した結果 \phi(x) に対して g に対応する操作 g' を適用した結果 g'(\phi(x)) が一致することを示している。 これは、入力に対して操作した後に変換をしても、変換をした後に操作をしても結果は変わらないことを意味している。

また、以下の式を満たすとき、「変換 \phi は操作 g に対して不変である」という。

\displaystyle \phi(x)=\phi(g(x))


これは、操作 g を適用してもしなくても結果が変わらないことを意味している。 同変の定義における g' を恒等変換とした場合に不変となるため、不変は同変の特殊例であることがわかる。

Equivariant Neural Networksは、ここまでで説明した同変性をNNの設計に活用するわけであるが、なぜ不変性ではなく、同変性を用いるのかについては、Taco Cohen氏の学位論文[1]で以下のように述べられている。

We argue and show empirically that in the context of deep learning it is better to learn equivariant rather than invariant representations, because invariant ones lose information too early on in the network.
(DeepLによる翻訳) 我々は、深層学習の文脈では、不変な表現よりも同変な表現を学習する方が良いことを主張し、経験的に示している。なぜなら、不変な表現はネットワークの早い段階で情報を失いすぎてしまうからである。

つまり、NNの各層の処理が特定の操作に対して不変である場合、例えば画像Aが画像Bに対して、上に並進移動しているとか、時計回りに90度回転しているといった相対的な位置や向きなどの情報が失われてしまう。 一方、同変の場合、このような相対的な位置や向きの情報をNNの中間層で保持することができる。

群論との関係

群論は、「対称性を記述する数学」や「対称性をはかる数学」と言われることからデータの対称性に着目したEquivariant Neural Networksとの繋がりがあることがわかる。ここでは、群の正確な定義は割愛するが、群とは特定の条件を満たす集合である（実際には演算とセットで考えるがここでは深く立ち入らない）。前述の不変性と同変性の説明のところで、ある変換が「操作」に対して不変/同変であることの定義を述べたが、この「操作」を元（集合の要素）として、それら元のなす集合を群として考える。そして、Equivariant Neural Networksでは、特定の群の全ての元に対して、NNの各層の処理が同変となるように設計する。

Geometric Deep Learningとの関係

Equivariant Neural Networksについて調査しているときに、対象領域が非常に近く、よく出てくるキーワードとして「Geometric Deep Learning」という言葉がある。2021年に公開されたGeometric Deep Learningの文献[2]では、Geometric Deep Learningを以下のように説明しており、対称性に着目していることがわかる。

We call this geometrisation attempt "Geometric Deep Learning", and true to the spirit of Felix Klein, propose to derive different inductive biases and network architectures implementing them from first principles of symmetry and invariance.
(DeepLによる翻訳) 我々はこの幾何学的な試みを「Geometric Deep Learning」と呼び、Felix Kleinの精神に忠実に、対称性と不変性の第一原理から様々な帰納的バイアスとそれを実装したネットワークアーキテクチャを導出することを提案する。

また、文献の5.2節「Group-equivariant CNNs」では、群同変なCNNについて解説されている。このように両者は重なる部分が多い研究領域であるが、Geometric Deep Learningは同変性のみにフォーカスしていないことなどからより対象領域が広いと考えられる。

CNNと同変性

同変性をNNの設計に組み込むことで大きな成功を遂げたモデルの代表的な例として、CNNが挙げられる。以下の図で示したようにCNNの畳み込み層は、並進移動に対する同変性を持っている。図は不変性と同変性の定義の説明の際に用いた図と対応している。

畳み込み層と並進移動

図のように、入力画像「7」を畳み込み処理して得られた特徴マップを左上に並進移動させたものは、入力画像を同様に左上に並進移動した画像を畳み込み処理して得られた特徴マップと一致する。 つまり、入力画像を並進移動させても、対応する特徴マップも同様に並進移動する。 これは、畳み込み処理が、画像全体で同じ重みを共有したフィルタを用いて、画像に対してフィルタを並進移動させながら内積をとるといった処理を行っているためである。畳み込み層はこの重みの共有により並進移動に対する同変性と全結合層と比べたパラメータ数の削減を実現している。実際にCNNでは、出力層の前にGlobal Average Pooling（GAP）を適用することが多く、これによりモデル全体としては並進に対して不変性となる。

それでは、画像に対して回転操作を行った場合はどうだろうか。結論は、以下の図のようにCNNの畳み込み層は回転に対して同変とならない。

畳み込み層と回転操作

つまり、画像の畳み込み処理により得られた特徴マップを反時計回りに90度回転したものは、画像を反時計回りに90度回転したものに畳み込み処理をして得られた特徴マップと一致しない。 これは、例に挙げた手書き数字のように画像内の認識対象に上下左右の明確な向きがある場合には問題になりにくい。 しかし、同じ画像でも例えば衛星画像や顕微鏡画像など明確な向きが存在しないものも多く、この場合は回転しても画像の意味は変わらないため、NNの処理が回転に対する同変性や不変性を持つことが望まれる。

この問題を従来のCNNで解決する方法の一つとして、Data Augmentationが挙げられる。つまり、様々な回転操作を施した画像を用意してCNNで学習するという方法である。 しかし、この方法では単純にデータ量が膨大になる。また、モデル自体はある画像とそれを回転させた画像の関係性を捉えることはなく、全く異なる画像を同じラベルに紐付けるように学習する。 このことからも学習効率が良くないことが想像できる。

様々なEquivariant Neural Networks

Equivariant Neural Networksは2016年頃から非常に活発に研究が行われている。 以下のGitHubページには、Equivariant Neural Networksに関連する論文やチュートリアルなどの情報がまとめられている。

github.com

また、E. J. Bekkers氏は、「Group Equivariant Deep Learning」というタイトルの講義をYouTubeで公開しており、既存のEquivariant Neural Networksを数学的な背景から解説してくれている。その講義スライドの中で同変なCNNに関するこれまでの歴史がまとめてある。

出典: Group Equivariant Deep Learning - Lecture 1.5: A Brief History of G-CNNs

ここでは、既存のEquivariant Neural Networksの中からいくつかピックアップして簡単に紹介する。

Group Equivariant Convolution Networks (2016)

Group Equivariant Convolution Networks（G-CNNs）[3]は、Equivariant Neural Networksの先駆け的な存在である。G-CNNsは、並進移動、鏡映、90度の倍数の回転操作を元とするp4m群の作用に対して同変性を持つCNNを実現している。 以下の記事は図が豊富で、G-CNNsの処理の内容がわかりやすく描画されている。

medium.com

本ブログ公開後の2022年5月23日にG-CNNsの論文解説をブログにまとめた。

blog.tsurubee.tech

Steerable CNNs (2017)

前述のG-CNNの一つの課題として、計算量が対象とする群のサイズ（元の個数）に比例して大きくなることが挙げられる。Steerable CNNs [4]は、計算量が群のサイズに依存しない同変なCNNを実現している。これにより、同変なCNNをよりサイズが大きな群にスケールすることが可能となる。論文中では、G-CNNと同様にp4mの群の作用に同変なCNNについて説明や評価が行われている。

Spherical CNNs (2018)

従来のCNNは、2次元平面上のデータを取り扱う問題に対して広く適用されているが、Spherical CNNs [5]では、球面上のデータを取り扱う。これは例えば、ドローンなどから取得される全方位ビジョンや地球気象のモデリングのための球面画像などが挙げられる。Spherical CNNsは、3次の特殊直交群 SO(3)の作用に同変な球面上の畳み込み層を提案している。

towardsdatascience.com

B-Spline CNNs (2020)

B-Spline CNNs [6]は、リー代数上で定義されるB-spline基底関数を用いて畳み込みカーネルを拡張することで、同変な畳み込み層を任意のリー群に一般化するフレームワークを提案している。論文中の評価では、病理画像における癌検出および顔のランドマーク推定の二つのタスクにおいて、それぞれ同変な畳み込み層を設計し、どちらも従来のCNNを大きく上回る性能を達成した。

E(n) Equivariant Graph Neural Networks (2021)

E(n) Equivariant Graph Neural Networks (EGNNs) [7]は、E(n)変換に対して同変なグラフニューラルネットワークである。E(n)とは、n次元ユークリッド空間における等長変換群であり、回転、並進、鏡映、置換から構成される。これまでに紹介したモデルと異なり、グラフ構造に着目しているので、創薬などの分子を対象データとする様々なタスクへの応用が期待される。

応用分野

Equivariant Neural Networksはデータに何らかの対称性を見いだせる場面で活用することで、NNの性能向上が期待できる。ここでは、現在の主な応用分野について紹介する。

医療画像

画像内の認識対象が明確な上下左右の向きを持っていない場合、向きが変更される回転に対して同変性・不変性が求められる。このような場面は、医療画像によく現れる。E. J. Bekkersらは、下図のような病理組織画像における有糸分裂の検出（Mitosis detection）、網膜画像における血管のセグメンテーション（Vessel segmentation）、電子顕微鏡画像における細胞境界のセグメンテーション（Cell boundary segmentation）の三つのタスクにおいて、2次元の特殊ユークリッド群SE(2)に対して同変なCNNを設計し、従来のCNNを上回る性能を達成した [8]。

出典：参考文献[8]のFig. 2

また、P. Müllerらは、MRI画像に3次元の特殊ユークリッド群SE(3)に対して同変なCNNを適用している[9]。 下図は、MRI画像における多発性硬化症病変のセグメンテーションの結果を示している。

出典：参考文献[9]のFigure 8

分子設計・材料設計

分子は3次元構造を持つため、3次元の回転などの操作に対して同変性・不変性が求められる。 また、材料科学の分野において，所望の物性を有する材料を開発することは重要な目標の一つであるため、分子やそれをマクロに捉えた材料の物性を深層学習を用いて予測する研究が盛んに行われている。 この材料物性の予測に対しても、Equivariant Neural Networksが活用され始めている。 T. Leらは、同変なグラフNNであるEQGATを設計し、低分子の量子力学的特性の予測においてSOTA性能を達成した[10]。

また、生成モデルを用いて、所望の分子を生成するアプローチも研究されている。V. G. Satorrasらは、E(n)変換に対して同変なNormalizing FlowsであるE-NFsを提案し、下図のような3次元の分子構造を生成することを可能にした[11]。

出典：参考文献[11]のFigure 1

創薬

医薬品として使われる低分子化合物やタンパク質（抗体）も分子であるため、ここで述べる創薬は分子設計に含めても良いが、背景として創薬にフォーカスしているものを切り出して紹介する。 H. Stärkらは、低分子化合物が特定の標的タンパク質にどのように結合するかを予測するために、SE(3)に対して同変なNNであるEQUIBINDを提案した[12]。EQUIBINDは、受容体の結合位置とリガンドの結合姿勢・向きを予測することができる。 また、O. E. Ganeaらは、タンパク質複合体の立体構造を予測するために、SE(3)に対して同変なNNであるEQUIDOCKを提案した[13]。EQUIDOCKは、結合ポケットの近似やドッキングポーズの予測を行うことができる。

出典：参考文献[12]のFigure 1

さいごに

この記事では、Equivariant Neural Networksとは何かから始まり、その応用事例までを概論的に紹介した。Equivariant Neural Networksは、本文中でも紹介した既存のモデルを深掘りして理解しようとしたときに、群論や微分幾何学などの高度な数学の知識が要求される。現在私は、これらの数学を学びながら理解を進めているため、今後はEquivariant Neural Networksのいくつかの既存のモデルについて、学んだことを各論的にまとめた記事を書いていく予定である。

更新：2022年5月23日公開

blog.tsurubee.tech

参考文献

[1] T. Cohen, "Equivariant Convolutional Networks" (2021).
[2] M. M. Bronstein et al., "Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges" (2021).
[3] T. Cohen et al., "Group Equivariant Convolutional Networks" (2016).
[4] T. Cohen et al., "Steerable CNNs" (2017).
[5] T. Cohen et al., "Spherical CNNs" (2018).
[6] E. J. Bekkers, "B-Spline CNNs on Lie Groups" (2020).
[7] V. G. Satorras et al., "E(n) Equivariant Graph Neural Networks" (2021).
[8] E. J. Bekkers et al., "Roto-Translation Covariant Convolutional Networks for Medical Image Analysis" (2018).
[9] P. Müller et al., "Rotation-Equivariant Deep Learning for Diffusion MRI" (2021).
[10] T. Le et al., "Equivariant Graph Attention Networks for Molecular Property Prediction" (2022).
[11] V. G. Satorras et al., "E(n) Equivariant Normalizing Flows" (2021).
[12] H. Stärk et al., "EQUIBIND- Geometric Deep Learning for Drug Binding Structure Prediction" (2022).
[13] O. E. Ganea et al., "Independent SE(3)-Equivariant Models for End-to-End Rigid Protein Docking" (2022).