陰関数微分：勾配計算で計算グラフをワープする

本記事はロボットとAI技術の専門誌『日経Robotics』のデジタル版です

［画像のクリックで別ページへ］

　多くの機械学習は学習問題を最適化問題として定式化し、勾配降下法を使って最適化する。ニューラルネットワーク（NN）は誤差逆伝播法（後ろ向き自動微分）を使うことで複雑な関数でも勾配を効率的に計算できている。

著者の岡野原大輔氏

　勾配計算のもう1つの重要なアルゴリズムとして陰関数微分がある。深層学習（ディープラーニング）の分野での陰関数微分の利用は限定的であったが、最近は急速に適用範囲が広がっている。今回はこれについて紹介する。

陰関数微分とは

　始めに、陰関数微分の基本的な考えを説明する。次のような二変数関数から成る方程式を考える。

\[ F(x, y) = 0 \]

このとき、ある条件を満たした場合、$y = \phi(x)$が存在し、

\[ F(x, \phi(x)) = 0 \]

を満たせることが分かっている（陰関数定理）。つまり$F(x, y)=0$という方程式から$y =\phi(x)$という関数が導出されたことになる。この$y=\phi(x)$を陰関数と呼ぶ。

　方程式が単純でも、陰関数は非常に複雑になる場合があり、むしろこのような場合に陰関数は多く使われる。この方程式の両辺を$x$で微分すると

\[ 0 = \frac{\mathrm{d}}{\mathrm{d}x}F(x, \phi(x)) = F_x + F_y \frac{\mathrm{d} \phi(x)}{\mathrm{d}x} \]

となる。ここで

\[ F_x = \frac{\partial}{\partial x}F(x, \phi(x)) \]

\[ F_y = \frac{\partial}{\partial y}F(x, \phi(x)) \]

はそれぞれ$F$の$x$と$y$についての偏微分である。$F_y$は正則であることを仮定（これは陰関数定理の条件の1つである）した上でさきほどの式を変形すると

\[ \frac{\mathrm{d} \phi(x)}{\mathrm{d}x} = - F_y^{-1} F_x \]

が得られる。これは陰関数の入力についての微分である。このように陰関数がどれだけ複雑な形をしていても、その微分は方程式の偏微分の組み合わせで表現することができる。

Deep Equilibrium Model

　それでは実際に陰関数微分を利用したモデルとしてDeep Equilibrium Model^1）をみていく。このモデルでは入力$\mathbf{x}$と適当に初期化した隠れ状態$\mathbf{z}^{[0]}$に対し、繰り返し同じ関数を適用していくことを考える。

\[ \mathbf{z}^{[i+1]} = f_{\theta}(\mathbf{z}^{[i]}, \mathbf{x}) \]

通常のNNとは違って、常に入力が引数として与えられている（入力からスキップ接続がある）、各層のパラメータが共有されているようなモデルである。RNNで入力が時刻によらず固定の場合と考えても良い。このようなモデルとしてUniversal Transformer^2）、Trellis Network^3）などが知られており、それらは多くの系列問題の最高精度を達成している。

　この関数を無限回適用した場合、$f_{\theta}$が特定の条件を満たすならば、隠れ状態はある値（不動点）に収束する。

\[ \lim_{i \to \infty} \mathbf{z}^{[i]} = \lim_{i \to \infty}f_{\theta}(\mathbf{z}^{[i]}; \mathbf{x}) \equiv f_{\theta}(\mathbf{z}^{*}; \mathbf{x}) = \mathbf{z}^{*} \]

これは$f_{\theta}$が定義するダイナミクス上で隠れ状態が均衡（Equilibrium）状態に到達したといえる。この入力$\mathbf{x}$から、不動点$\mathbf{z}^*$を返す一連の操作を1つの関数とみなす。

　さきほどの$f_{\theta}$を繰り返し適用して不動点を求める操作は不動点反復法と呼ばれ、条件を満たせば線形収束する。このほかにも不動点を求める方法であれば、より高速に求められる方法がある。具体的には

\[ g_{\theta}(\mathbf{z}; \mathbf{x}) = f_{\theta}(\mathbf{z};\mathbf{x}) - \mathbf{z} \]

という関数を定義し、この関数が$0$をとるような$\mathbf{z}$を求める問題を考え、（準）ニュートン法を適用する。

\[ \mathbf{z}^{[i+1]} = \mathbf{z}^{[i]} - \alpha B g_{\theta}(\mathbf{z}^{[i]}; \mathbf{x}) \]

ただし$B$は$g_{\theta}$の$\mathbf{z}^{[i]}$におけるヤコビアンの逆行列である。この$g_{\theta}(\mathbf{z}; \mathbf{x}) =0$となる解を求めるプロセス全体を$\mathbf{z}^* = \mathtt{RootFind}(g_{\theta};\mathbf{x})$というブラックボックスの関数としよう。この関数は上にあるような準ニュートン法などを使うかもしれないし、不動点反復法を収束するまで繰り返したものかもしれない。

　このように得られた解$\mathbf{z}^*$に損失関数を適用した結果を目的関数とする。

\[ l(\mathbf{z}^*) = l(\mathtt{RootFind}(g_{\theta}; \mathbf{x})) \]

学習ではこの目的関数の$\theta$や$\mathbf{x}$についての勾配が必要となる。ここではそれらをまとめて$\frac{\partial l}{\partial (\cdot)}$と表し、$(\cdot)$に$\theta$や$\mathbf{x}$が入ることとしよう。ここで陰関数微分を使うことで、勾配計算時に$\mathtt{RootFind}$を使わずに勾配が求められる。

　まず、$\mathbf{z}^* = f_{\theta}(\mathbf{z}^*;\mathbf{x})$という条件を利用する。この両辺を$(\cdot)$について微分すると

\[ \frac{\mathrm{d} \mathbf{z}^*}{\mathrm{d} (\cdot)} = \frac{\mathrm{d} f_{\theta}(\mathbf{z}^*; \mathbf{x})}{\mathrm{d} (\cdot)} =\frac{\partial f_{\theta}(\mathbf{z}^*; \mathbf{x})}{\partial (\cdot)} + \frac{\partial f_{\theta}(\mathbf{z}^*;\mathbf{x})}{\partial \mathbf{z}^*}\frac{\mathrm{d}\mathbf{z}^*}{\mathrm{d} (\cdot)} \]

これを式移動すると次が得られる。

\[ \left ( I - \frac{\partial f_{\theta}(\mathbf{z}^*;\mathbf{x})}{\partial \mathbf{z}^*} \right ) \frac{\mathrm{d}\mathbf{z}^*}{\mathrm{d} (\cdot)} =\frac{\partial f_{\theta}(\mathbf{z}^*; \mathbf{x})}{\partial (\cdot)} \]

また、$g_{\theta}(\mathbf{z}^*; \mathbf{x}) = f_{\theta}(\mathbf{z}^*;\mathbf{x}) - \mathbf{z}^*$に対し、$\mathbf{z}^*$について微分をとると

\[ J_{g_{\theta}}|_{\mathbf{z}^*} = - \left ( I - \frac{\partial f_{\theta}(\mathbf{z}^*; \mathbf{x})}{\partial \mathbf{z}^*} \right ) \]

が得られる。ここで$J_{g_{\theta}}|_{\mathbf{z}^*}$は$g_{\theta}$の$\mathbf{z^*}$におけるヤコビアンである。これらを組み合わせると、損失関数の勾配は

\[ \frac{\partial l}{\partial (\cdot)} = \frac{\partial l}{\partial\mathbf{z}^*}\frac{\mathrm{d} \mathbf{z}^*}{\mathrm{d} (\cdot)} = -\frac{\partial l}{\partial \mathbf{z}^*} \left (J^{-1}_{g_{\theta}}|_{\mathbf{z}^*} \right )\frac{\partial f_{\theta}(\mathbf{z}^*; \mathbf{x})}{\partial (\cdot)} \]

と求められる。これはさきほどの陰関数微分の定理で$J^{-1}_{g_{\theta}}|_{\mathbf{z}^*}$が$F_y^{-1}$、$\frac{\partial f_{\theta}(\mathbf{z}^*;\mathbf{x})}{\partial (\cdot)}$が$F_x$に対応していると思ってもらってよい。

　この$\frac{\partial l}{\partial \mathbf{z}^*}J^{-1}_{g_{\theta}}|_{\mathbf{z}^*} $を求めるのは困難そうに見えるが、代わりに

\[ \left (J^T_{g_{\theta}}|_{\mathbf{z}^*} \right ) \mathbf{x}^T +\left ( \frac{\partial l}{\partial \mathbf{z}^*} \right ) ^T =\mathbf{0} \]

という線形方程式の解として求めることができる。第一項のベクトル・ヤコビアン積（vjp）は自動微分を使うことでヤコビアンを直接求めなくても解くことができる（この点は2019年6月号の本欄のNeural ODEの回でも説明した）。また線形方程式も準ニュートン法などを用いて高速に解くことができる。Deep Equilibrium Modelはこの陰関数微分を使うことで従来モデルに比べて1/10程度のメモリ使用量で学習できている（表1）。

表1　各モデルで言語モデルを学習させた場合のネットワークサイズやメモリ容量

（文献^1）より引用）

[画像のクリックで拡大表示]

メタ学習

　陰関数微分はメタ学習でも利用されている^4）。メタ学習では似たようなタスクがたくさんある場合に、一部のタスクで学習し、新しいタスクをより少ない学習データで効率よく学習できるようにすることが目的である。ここでは各タスクの学習（内側ループ）と、タスク集合に対するメタ学習用のパラメータの最適化（外側ループ）の2つがある。

　これまでは内側ループの最適化プロセス上で誤差逆伝搬法を走らせる必要があるため、各タスクの更新回数は数回と限られていた。これに対し、各タスクの最適化プロセスは陰関数で表すことができるため、内側ループでの更新回数によらず、陰関数微分を使うことで外側ループで必要なパラメータの勾配を効率よく求めることができる。

Implicit Deep Learning

　今回紹介したように、方程式を満たすような変数間の関係を利用して表現したい関数を表すニューラルネットワークを米University of California BerkeleyのLaurent El Ghaoni氏らはImplicit Deep Learning^5）と呼んでいる。このようなモデルは達成してほしい制約や目標を直接モデルに組み込むことができるメリットがある。さらに、入力から均衡点を返すという関数は、入力の変化に対し滑らかであり安定的であることが示唆されており、遷移関数を工夫することでノルムを変えず、勾配発散/消失が起きないRNNを作れることが分かっている^6）。

本記事はロボットとAI技術の専門誌『日経Robotics』のデジタル版です

［画像のクリックで別ページへ］

　一方、SGDによる最適化と同様に、均衡状態に到達する経路や初期値が重要になる可能性もある。例えば、オンライン学習などでは前のパラメータに近いという制約を入れることで新しいデータに対応しつつも以前のデータを記憶することを達成できている。同様に均衡状態に収束する前にあえて止めることで前の記憶を残すことができる。究極的には入力や拘束条件が次々と変化していく環境下で均衡状態に収束する前に次の均衡状態に向かうというダイナミクスを持ったモデルを考える必要があるだろう。

岡野原 大輔（おかのはら・だいすけ）
Preferred Networks 代表取締役副社長

2006年にPreferred Infrastructureを共同創業。2010年、東京大学大学院博士課程修了。博士（情報理工学）。未踏ソフト創造事業スーパークリエータ認定。東京大学総長賞。