最適化喫茶１：機械学習における陰関数微分の利用（2レベル最適化・ハイパーパラメータ最適化）

前置き

数学の記事を試しに書いてみたいという動機だけで書き進めたので、内容は大雑把で抽象的なのはご容赦（必要あれば修正します。質問も気兼ねなくどうぞ）。

最近、機械学習の分野で陰関数微分を使った経験残差の最小化が注目を浴びている。具体的にどういう問題を考えているのかというと、ハイパーパラメータチューニングで使われるメタ学習と呼ばれる類のモデルである。このモデルは学習を上位レベルと下位レベルに分け、上位はハイパーパラメータを最適化する問題を解き、下位はハイパーパラメータを含む学習タスク（残差最小化）を行う。

この構造を持ったモデルは数理最適化では２レベル最適化と呼ばれており、以下のような形で定義される問題のことを指す。

$$
\begin{aligned}
\min_{x\in\mathcal{X}} \quad & F(x,y(x)) \\
\text{s.t.} \quad & y(x)\in \arg\min_{y\in\mathcal{Y}} G(x,y)
\end{aligned}
$$

特徴は上位の最適化問題の制約条件に下位の最適化問題を内包している構造を持っている点である。2レベル最適化問題を解くにはどんな手法を用いたらよいだろうか？

もっとも容易に思い浮かぶのは、上位と下位をそれぞれ交互に解くようなアルゴリズムだろうか。つまり、先に $${(x^0, y^0)}$$ という初期点を与えた後、まず$${y^0}$$ を固定して $${x}$$ について解き、$${x^1}$$ が得られた後、それを固定して $${y}$$ について解く… というのを繰り返して最適なハイパーパラメータと学習を行う方法である。

しかし、この方法によって得られる上位・下位の最適解のペア $${(x^*,y^*)}$$ は上記の2レベル最適化を解いておらず、下位と上位が同等な関係にあるような同時手番の最適化、つまり、非協力ゲームにおけるナッシュ均衡を求めているに過ぎない。そのため、2レベル最適化は解法の構築からして複雑であり、この20～30年で様々な解法が提案されてきた。

本記事では、下位の最適解（最適応答）を上位の変数に関する陰関数としてみなし、その微分情報を用いる陰関数微分アプローチによって2レベル最適化問題を解く手法を紹介する。

陰関数微分

仮定

1. 関数 $${G\colon\mathcal{X}\times\mathcal{Y}\to\Bbb R}$$ は2回連続的微分可能

2. 関数 $${G(x,\cdot)}$$ は任意に与えられた $${x}$$ に対して強凸

3. $${\mathcal{Y}=\Bbb R^m}$$ （すなわち無制約最適化）

4. $${F\colon\mathcal{X}\times\mathcal{Y}\to\Bbb R}$$ は $${(x,y)}$$ について連続的微分可能

凸性の下での下位レベル最適化問題の特徴・陰関数定理

下位レベルの最適化問題に関して、強凸性から任意に与えた $${x}$$ に対して最適解 $${y^*}$$ が一意に存在することが保証されている（これは最適化理論でよく知られた結果を利用）。また、凸な無制約最適化問題を考えているので、$${\nabla_y G(x,y)=0}$$ を満たす $${y}$$ は最適解となる。すなわち、 方程式 $${\nabla_y G(x,y)=0}$$ の $${y}$$ に関する解＝下位にとって、与えられた上位の変数 $${x}$$ に対する最適な応答を表す。

陰関数定理：ここで、この方程式 $${\nabla_y G(x,y^*)=0}$$ について、$${(x,y)}$$ 近傍で$${y}$$ について $${\nabla_y G}$$ をさらにもう一回微分したヘッセ行列 $${\nabla^2_{yy} G(x,y^*)}$$ が正則（逆行列を持つ）ならば、$${G(x,y^*(x))=0}$$ を満たす$${x}$$ について微分可能な関数 $${y^*(x)}$$ が存在することが知られており、これを陰関数定理とよぶ。ここで、$${G}$$ は $${y}$$ について強凸であるから、ヘッセ行列 $${\nabla_{yy} G(x,y^*)}$$ は正定値となるため、正則であるから、陰関数定理が成り立つ。

陰関数微分

陰関数定理を用いることで、$${y^*(x)}$$ は閉形式として陽に書けないが、その $${x}$$ に関する微分（ヤコビ行列） $${\nabla_x y^*(x)}$$ は以下のように書ける：

$$
\begin{aligned}
& \nabla_x G(x,y^*(x)) = 0 \\
\Leftrightarrow\ & \nabla_{xy} G(x,y^*) + \nabla^2_{yy}G(x,y^*) \nabla_x y^*(x) = 0 \\
\Leftrightarrow\ & \nabla_x y^*(x) = -[\nabla_{yy} G(x,y^*)]^{-1}\nabla_{xy} G(x,y^*)
\end{aligned}
$$

陰関数微分を用いた上位レベルの最適化

次に、上位の最適化問題を考える。ここでは簡単のため、上位も無制約最適化を解くものとする。すなわち、$${\mathcal{X}=\Bbb R^n}$$ とする。無制約最適化においてよく知られた、最適解を求める最も基本的なアルゴリズムとして最急降下法があるので紹介する（一般に最急降下法自体は最適解までの収束スピードが遅いので、加速化手法（Nesterovの加速化など）がしばしば用いられる）。詳しくは数理最適化に関する書籍に当たることをおすすめする。

準備

下位の変数を含まない以下の無制約凸最適化問題を考える。

$$
\min_x\quad f(x)
$$

この最適解を得るには、凸性より $${\nabla_x f(x) = 0}$$ となる解を求めればよい。最急降下法は $${f(x)}$$ を $${x^k}$$ 付近で一次近似に近接項を加えた局所的に定義される関数

$$
f(x^k)+\nabla_x f(x^k)^\top (x-x^k)+\frac{1}{2\gamma_k}\|x-x^k\|^2
$$

を最小にする点 $${x^*}$$ を次の点（$${x^{k+1}}$$）とし

$$
x^{k+1}:=x^k-\gamma_k\nabla_x f(x^k)
$$

によって逐次的に関数を小さくする点に更新する手法である。ここで、$${\gamma_k>0}$$ とし、これはステップ幅とか、機械学習の分野では学習率と呼ばれる正数である。
※近接項の意味：一次近似 $${f(x^k)+\nabla_x f(x^k)^\top (x-x^k)}$$（一次式）だけでは、無制約の場合、どこまでも小さくできてしまうため、$${x^k}$$ から離れるほど罰則が科されるような二次の項を加えていることに注意する。これを近接化、あるいは正則化と呼んだりする。

最急降下法を用いた上位レベルの最適化

さて、もとの上位の最適化問題の最小化に戻る。陰関数表現された目的関数 $${F(x,y^*(x))}$$ はもはや $${x}$$ のみを変数とする関数に他ならないので、$${f(x)=F(x,y^*(x))}$$ として先ほどの問題に適用して考えてみよう。

合成関数の勾配は連鎖律を用いればよいので、$${\nabla_x F(x,y^*(x)) = \nabla_x F(x,y^*) + \nabla_x y^*(x)^\top \nabla_y F(x,y^*)}$$ となる。ここで、$${\nabla_x y^*(x)}$$ は先ほど陰関数定理を用いて導いたものを代入してあげれば良いので、

$$
\nabla_x y^*(x)^\top \nabla_y F(x,y^*) = -\nabla_{xy} G(x,y^*)^\top[\nabla_{yy} G(x,y^*)]^{-\top} \nabla_y F(x,y^*)
$$

となる。ここで、$${-\top}$$ は逆行列に転置をとったものを表す。

したがって、上位の最適化における最急降下法の更新式は

$$
x^{k+1}=x^k-\gamma_k \{\nabla_x F(x^k,y^*) -\nabla_{xy} G(x,y^*)^\top[\nabla_{yy} G(x,y^*)]^{-\top} \nabla_y F(x,y^*)\}
$$

ここで、$${y^*}$$ は $${x^k}$$ を与えたときの下位の最適解となることに注意する（※ 上の式では陰関数の情報を使うか否かはっきりさせたかったので、あえて陽に $${y^*}$$ を用いる際は $${y^*(x^k)}$$ とは記さなかった）。

おわりに：陰関数微分を用いた最適化手法の長所・短所、陽関数微分アプローチとの比較

本記事で紹介してきた手法は下位の陰関数微分を用いることで、2レベル最適化問題を解く方法を紹介した。しかし、陰関数微分は（一般には疎でない）大規模なヘッセ行列の逆行列を必要とするため、メモリ消費量が大きいことから、1反復当たりの計算コストが高くなることが欠点である。しかし、行列計算法の観点から様々な計算負荷を軽くする手法が提案されているので、問題（インスタンス）によってはその短所をうまくカバーできるかもしれない。また、実用的には下位レベルの最適応答 $${y^*}$$ はアルゴリズムにおいて、初めのうちは厳密に求めなくても、段々精度を上げていく形で求めれば十分であることが知られており、計算の負荷を下げるための一つの工夫へとつながる。

本記事では陽関数微分な方法は紹介しなかったが、こちらも活発に研究されている。こちらの手法は下位の更新式

$$
y^{t+1} := y^t - \eta_t \nabla_y G(x^k,y^t)
$$

を一つの動的システムの状態方程式 $${y(t+1)=f_{\eta_t}(y(t);x^k)}$$ の平衡状態 $${t\to\infty}$$ を求めるために差分方程式を解くことで近似的に$${\nabla_x y(x^k)}$$ を求める方法で、陰関数微分手法よりも制御工学的なアイデアが含まれている。

陽関数微分アプローチは一般に陰関数微分アプローチよりも収束性を保証するために多くの仮定を必要としないが、$${\nabla_x y(x^k)}$$ の近似値を計算するためにいくらかの行列の積を必要とし、精度を求めるとその回数が増大するため、時間・空間計算量ともに陰関数微分アプローチよりも多く必要とする。