ニューラルタンジェントカーネル(NTK)をわかりやすく:無限幅で深層学習が線形になる理論

ニューラルネットワークの学習には、よく考えると不思議なことがたくさんあります。損失関数はパラメータについて激しく非凸で、谷も尾根も無数にあるはずです。それなのに、ただの勾配降下法で訓練すると、訓練誤差はほぼゼロまで落ちます。しかも、パラメータ数がデータ数よりはるかに多い「過剰パラメータ」の状態でも、過学習で崩壊するどころか、新しいデータにそこそこうまく当たります。なぜこんなに都合よくいくのでしょうか。

この謎に1つの明快な答えを与えたのが、ニューラルタンジェントカーネル(Neural Tangent Kernel, NTK) です(Jacot, Gabriel, Hongler, 2018)。主張はとても大胆です。ネットワークの幅を無限に広げると、勾配降下で訓練されるニューラルネットは「ある固定されたカーネルを使った線形モデル」と区別がつかなくなる、というのです。線形モデルなら学習は凸最適化に帰着します。だから大域最適に届くのは当たり前。汎化もカーネルの性質で決まる——こうして冒頭の謎が一気に見通せます。

この考え方が活きる場面は、少なくとも2つあります。1つは深層学習の理論的理解です。「なぜ勾配降下が効くのか」「なぜ過剰パラメータでも汎化するのか」という問いに、カーネル法という枯れた道具で答えられます。もう1つはカーネル回帰としての予測です。NTKを実際に計算してしまえば、ネットワークを訓練しなくても、カーネル回帰の閉じた式で予測値が出せます。ニューラルネットの「気持ち」を、行列演算1発で再現できるわけです。

本記事の内容

  • なぜネットワークが「初期点まわりの線形モデル」に化けるのか(1次テイラー展開)
  • パラメータ勾配が特徴写像になり、その内積がNTKになること
  • 無限幅でNTKが決定的な値に収束し、訓練中も不変になること(lazy training)
  • 2乗損失の訓練ダイナミクスが線形ODEで解け、予測がカーネル回帰の閉形式になること
  • NTKの含意(大域収束・汎化)と、限界(特徴学習が起きないlazy regime)
  • PyTorchの自動微分で経験的NTKを計算し、理論を実測で確かめる

前提知識

この記事を読む前に、以下の記事を読んでおくと理解がぐっと深まります。

NTKのアイデアを一言でいうと

まず、ゴールのイメージを先につかんでおきましょう。ニューラルネット $f(x;\theta)$ は、入力 $x$ とパラメータ $\theta$ の両方の関数です。学習とは、$\theta$ を動かして $f$ を望ましい形に近づける作業です。ところが $f$ は $\theta$ について恐ろしく複雑な非線形関数です。これがそのままだと手に負えません。

NTKの発想はこうです。幅が十分広いと、訓練の間じゅう $\theta$ はほとんど初期値 $\theta_0$ の近くにとどまる。ならば、$\theta_0$ のまわりで $f$ を1次のテイラー展開で近似してしまえばよい。すると $f$ は $\theta$ について線形な関数になります。線形なら、勾配降下のダイナミクスは線形微分方程式で解けます。そして、その線形モデルの「内積構造」を決めているのが、まさにNTKというカーネルなのです。

次の図は、この発想の核心を1枚にまとめたものです。

NTKの出発点:初期点まわりでネットワークを線形化する

青い曲線が「本物のネットワーク」、つまり $\theta$ について非線形な $f(\theta)$ です。オレンジの破線が初期点 $\theta_0$ での1次近似(接線)です。一般には、$\theta$ が $\theta_0$ から離れれば、両者は大きくずれます。しかし幅が広いと、訓練中は $\theta$ が薄く塗った帯の中($\theta_0$ のすぐ近く)からほとんど出ないので、ネットワークと線形近似が事実上一致してしまう、というのがNTKのカラクリです。

ここまでで「なぜ線形化が許されるのか」という直感の入口に立てました。では、その線形近似を式で書き下し、そこからカーネルがどう現れるのかを見ていきましょう。

1次テイラー展開:パラメータについて線形化する

ネットワークを初期パラメータ $\theta_0$ のまわりで、$\theta$ について1次テイラー展開します。ここで微分する相手は入力 $x$ ではなくパラメータ $\theta$ であることに注意してください。

$$ \begin{equation} f(x;\theta) \approx f(x;\theta_0) + \nabla_\theta f(x;\theta_0)^\top (\theta – \theta_0) \end{equation} $$

右辺を眺めると、構造がはっきり見えます。第1項 $f(x;\theta_0)$ は初期出力で、定数です。第2項は $\nabla_\theta f(x;\theta_0)$ という「パラメータについての勾配ベクトル」と、パラメータの変位 $\theta – \theta_0$ との内積です。

ここで決定的に大事な視点があります。$\theta – \theta_0$ を「重み」、$\nabla_\theta f(x;\theta_0)$ を「特徴ベクトル」と読み替えてみてください。すると式 (1) は、

$$ f(x;\theta) – f(x;\theta_0) \approx \phi(x)^\top w, \qquad \phi(x) := \nabla_\theta f(x;\theta_0),\quad w := \theta – \theta_0 $$

という、ただの線形モデルになります。特徴写像 $\phi(x)$ が固定されていて、その上で重み $w$ を学習する——これは線形回帰そのものです。深層学習が、初期点のまわりで線形回帰に化けたわけです。

ここで自然な疑問が湧きます。特徴写像 $\phi(x)=\nabla_\theta f(x;\theta_0)$ とは、いったいどんなベクトルなのでしょうか。

NTKの定義:パラメータ勾配の内積

線形モデルでは、特徴ベクトル同士の内積がすべてを決めます(カーネルトリックの心臓部です)。そこで、2つの入力 $x, x’$ に対する特徴ベクトルの内積を取り、これに名前を付けます。これがニューラルタンジェントカーネルです。

$$ \begin{equation} \Theta(x, x’) := \nabla_\theta f(x;\theta_0)^\top \nabla_\theta f(x’;\theta_0) = \phi(x)^\top \phi(x’) \end{equation} $$

言葉にすれば、「$x$ で測ったパラメータ勾配」と「$x’$ で測ったパラメータ勾配」の内積です。2つの入力が、パラメータをいじったときに似た方向に出力を動かすなら $\Theta(x,x’)$ は大きく、関係がなければ小さくなります。これはまさにカーネル(似ているほど大きい類似度)の役割です。

次の図は、この流れを模式化したものです。

パラメータ勾配が特徴写像になりその内積がNTK

入力 $x$ にパラメータ勾配 $\nabla_\theta f(x;\theta_0)$ を作用させると、特徴ベクトル $\phi(x)$ が得られます。$\phi(x)$ と $\phi(x’)$ の内積を取れば、それがカーネル値 $\Theta(x,x’)$ です。特徴空間の次元はパラメータ数 $P$ なので、過剰パラメータのネットワークほど超高次元の特徴空間で物事を見ていることになります。

ここで前提記事との接続を1つ。ヤコビアンの記事では、ヤコビアンには「掛け算すると学習」「行列式を取ると確率」「ノルムを測ると感度」という3つの顔がある、と紹介しました。NTKはその第1の顔の親戚です。ただし、あの記事のヤコビアンは入力 $x$ についての微分でした。NTKで使うのはパラメータ $\theta$ についての勾配(パラメータ・ヤコビアン)です。サンプルごとのパラメータ勾配を1行に並べたヤコビアン行列を $J$ とすると、NTKのグラム行列はちょうど $\Theta = J J^\top$ という「ヤコビアンの積」で書けます。学習の話にヤコビアンの積が顔を出す、という構図はここでも生きています。

理屈の上では $\Theta$ は初期化 $\theta_0$ に依存するランダムな量です。初期化を引き直せば値も変わるはずです。それなのに「カーネルが固定されている」と言えるのはなぜでしょうか。鍵は「無限幅」にあります。

無限幅極限:カーネルが決定的になる

ここがNTKの最も驚くべき部分です。ネットワークの幅 $m$(隠れ層のユニット数)を大きくしていくと、ランダム初期化で計算した $\Theta(x,x’)$ の値が初期化ごとにばらつかなくなり、ある決定的な値に収束します。幅が無限大の極限では、$\Theta$ はもはや乱数ではなく、入力 $x, x’$ だけで決まる固定された関数になるのです。

直感的には大数の法則です。$\Theta$ は隠れユニットごとの寄与の和で、適切なスケーリング(出力を $1/\sqrt{m}$ で割る等)のもとで、ユニット数 $m$ が増えるほど和は期待値に集中します。サイコロを1個振ると目はバラバラですが、1万個振って平均すれば必ず約3.5になる——あの現象が、カーネルの値で起きていると思えば十分です。

実際にPyTorchで経験的NTKを計算し、幅を変えてばらつきがどう縮むか見てみましょう。

幅を増やすとNTKが決定的な値に収束する

横軸が幅 $m$(対数)です。各幅で30回、初期化を引き直して $\Theta(0.5, -0.7)$ を計算し、その平均±標準偏差を青で、標準偏差そのものをオレンジで描いています。幅10では標準偏差が約0.17もあり、初期化のたびに値が大きくぶれます。ところが幅3000では標準偏差が約0.009まで縮み、ほぼ一定値に張り付いています。幅を広げるほどNTKが決定的になることが、実測で確認できました。

カーネルが入力だけで決まる関数になったので、その全体像を行列で可視化してみましょう。

NTKグラム行列のヒートマップ

これは入力を $-2$ から $2$ まで25点並べ、すべてのペアで $\Theta(x,x’)$ を計算したグラム行列です。対角($x=x’$)に近いほど値が大きく、離れるほど小さくなる、なめらかな構造が見えます。つまりNTKは「近い入力ほど似ている」という、ガウスカーネルにも通じる素直な類似度になっています。この行列こそが、これから訓練ダイナミクスを支配する主役です。

無限幅でカーネルが固定された、というのは「初期化の話」でした。次は「訓練が進んでもカーネルが変わらない」という、もう一段強い性質を確認します。

lazy training:訓練中もカーネルが動かない

線形近似の式 (1) が訓練の最後まで成り立つには、$\theta$ が初期値の近くにとどまり、かつ $\nabla_\theta f$(=特徴写像)が訓練中ほとんど変化しないことが必要です。広い幅では、これが実際に起こります。1つ1つのパラメータが出力に与える影響は $1/\sqrt{m}$ で小さく、目標を達成するのに各パラメータがほんの少しだけ動けば足りるからです。結果として、パラメータは初期値からほとんど離れず、勾配 $\nabla_\theta f$ も初期のまま。これを lazy training(怠惰な訓練) と呼びます。

「カーネルが動かない」を、訓練前後でグラム行列がどれだけ変わるかで測ってみましょう。

幅が広いほど訓練中もNTKがほぼ不変

縦軸は、訓練前のカーネル $\Theta_0$ と訓練後のカーネル $\Theta_{\text{end}}$ の相対変化率 $\|\Theta_{\text{end}}-\Theta_0\| / \|\Theta_0\|$ です。幅10では約190%も変化していて、線形近似はまったく当てになりません。ところが幅を広げると変化率は急速に落ち、幅4000ではわずか0.3%です。幅が広いほど、訓練を回してもNTKはほぼ凍ったままであり、線形近似が訓練の最後まで通用することがわかります。これがlazy trainingの正体です。

カーネルが初期化でも訓練中でも固定だとわかれば、あとは「固定カーネルの線形モデルが勾配降下でどう動くか」を解くだけです。これは線形ODEなので、きれいに解けます。

訓練ダイナミクス:線形ODEとカーネル勾配流

2乗損失で訓練する場合を考えます。訓練入力を $X$、目標を $y$、ネットワークの出力ベクトルを $f_t = f(X;\theta_t)$ とします(添字 $t$ は訓練の時刻)。損失は $L(\theta) = \tfrac{1}{2}\|f_t – y\|^2$ です。

連続時間の勾配降下(勾配流)では、パラメータは損失の勾配の逆向きに動きます。

$$ \frac{d\theta_t}{dt} = -\nabla_\theta L(\theta_t) = -\nabla_\theta f_t^\top (f_t – y) $$

ここで知りたいのは $\theta$ ではなく、出力 $f_t$ がどう動くかです。連鎖律で $f_t$ の時間変化を書き下します。$f_t$ の変化は「パラメータの変化」を「出力への感度 $\nabla_\theta f_t$」で受けたものなので、

$$ \frac{df_t}{dt} = \nabla_\theta f_t \, \frac{d\theta_t}{dt} $$

先ほどの $d\theta_t/dt$ を代入すると、

$$ \frac{df_t}{dt} = -\,\nabla_\theta f_t \, \nabla_\theta f_t^\top \,(f_t – y) $$

となります。ここで右辺に現れた $\nabla_\theta f_t \, \nabla_\theta f_t^\top$ こそ、訓練点同士のNTKグラム行列 $\Theta$ にほかなりません。lazy trainingのおかげで $\Theta$ は時間によらず一定 $\Theta_0$ とみなせるので、

$$ \begin{equation} \frac{df_t}{dt} = -\,\Theta_0\,(f_t – y) \end{equation} $$

という、$f_t$ についての線形微分方程式が得られました。非線形だったネットワークの学習が、固定行列 $\Theta_0$ を係数とする線形ODEに化けた瞬間です。残差 $r_t := f_t – y$ について書き直すと $\dfrac{dr_t}{dt} = -\Theta_0 r_t$ で、これは指数減衰の方程式です。解は行列指数を使って、

$$ \begin{equation} f_t – y = e^{-\Theta_0 t}\,(f_0 – y) \end{equation} $$

と閉じた形で書けます。$f_0$ は初期出力です。$\Theta_0$ が(半)正定値なので、$t\to\infty$ で $e^{-\Theta_0 t}\to 0$、つまり残差はゼロに収束し、訓練誤差は消えます。非凸に見えた学習が、線形ODEの安定収束に帰着したのです。

この収束の「速さ」が何で決まるかを見るために、$\Theta_0$ を固有値分解しましょう。固有値を $\lambda_k$、固有ベクトルを $v_k$ とし、初期残差を $r_0 = \sum_k c_k v_k$ と分解します。これを式 (4) に入れると、各固有モードは独立に、

$$ r_t = \sum_k c_k\, e^{-\lambda_k t}\, v_k $$

と減衰します。つまり大きい固有値 $\lambda_k$ に対応する成分は速く、小さい固有値の成分はゆっくり消えるわけです。これは「スペクトルバイアス」と呼ばれる現象で、ニューラルネットがまず大局的(低周波)な成分を学び、細かい成分を後から学ぶ理由を説明します。図で確かめましょう。

損失は固有値ごとに指数的に減衰する

左はNTKグラム行列の固有値スペクトルで、桁の違う固有値が並んでいます。右は規格化した二乗残差の時間変化です。最大固有値に対応する「最速モード」(オレンジ)はあっという間に消え、小さい固有値の「遅いモード」(グレー)はゆっくり減衰します。全体の損失(赤)は、序盤に速いモードのおかげで急落したあと、遅いモードに引きずられてだらだら落ちていきます。収束の速さが固有値で決まるという式どおりの挙動です。

ODEを解いて訓練の最終形がわかったので、$t\to\infty$ での予測を任意の入力 $x$ に対して書き下しましょう。これがカーネル回帰の閉形式です。

予測はカーネル回帰の閉形式になる

訓練が収束した(残差ゼロの)後、新しいテスト点 $x$ に対する予測 $f_\infty(x)$ は、固定カーネルを使ったカーネルリッジ回帰(リッジ係数ゼロの極限) の式と完全に一致します。

$$ \begin{equation} f_\infty(x) = \Theta(x, X)\,\Theta(X, X)^{-1}\, y \end{equation} $$

ここで $\Theta(X,X)$ は訓練点同士のグラム行列、$\Theta(x,X)$ はテスト点と全訓練点の間のカーネルベクトルです。記号を見たことがある人もいるはずで、これはガウス過程回帰の予測平均とまったく同じ形です。実際、無限幅ネットワークはガウス過程と深く結びついており、NTKはそのカーネルを与えます。

式 (5) の意味は明快です。テスト点 $x$ の予測は、各訓練点の目標値 $y_i$ を、$x$ との類似度 $\Theta(x,x_i)$ で重み付けして混ぜたものです($\Theta(X,X)^{-1}$ は訓練点同士の重なりを補正する係数)。ネットワークを1ステップも訓練せずに、カーネルを計算するだけで予測が出せる——これがNTKの実用的な側面です。

実際にこの式で1次元回帰をやってみましょう。

NTKによるカーネル回帰の閉形式予測

赤点が訓練データ、点線が真の関数 $\sin(2x)$、青の実線が式 (5) で計算したNTKカーネル回帰の予測です。訓練データの範囲($-2$ から $2$)では真の関数をほぼ完璧になぞっています。両端の外挿域でわずかにずれるのは、データが無いところでカーネルの推定が不確かになるためで、これも素直な振る舞いです。ネットワークを訓練せず、カーネル回帰の閉形式だけで予測が再現できたわけです。

ここまでは「無限幅ならこうなる」という理論の話でした。本当に、実際に勾配降下で訓練した有限幅のネットワークが、このカーネル回帰の予測に一致するのでしょうか。直接ぶつけて確かめます。

理論と実物を突き合わせる

無限幅の理論が正しいなら、十分広いネットワークを普通に勾配降下で訓練した結果は、式 (5) のカーネル回帰の予測とほぼ重なるはずです。幅4000のネットワークを実際に訓練し、同じ初期化から計算したNTKカーネル回帰の予測と並べてみます。

NTK理論の予測と訓練した広いNNの一致

青がNTK理論の予測(カーネル回帰)、オレンジの破線が実際に勾配降下で訓練した幅4000のネットワークの予測です。2本の線は、目で見て区別がつかないほど重なっています。両者の食い違いはRMSEで $0.0004$ 程度でした。「無限幅では訓練したネットがカーネル回帰に化ける」という理論が、有限でも十分広ければ実際に成り立つことが確認できました。

この一致は幅を広げるほど良くなるはずです。幅を変えながら「訓練後のネットワーク」と「そのNTKによるカーネル回帰予測」の食い違いを測ってみましょう。

幅を広げると訓練後NNがNTK予測に近づく

横軸が幅、縦軸が両者のRMSE(どちらも対数)です。幅10では食い違いが約0.05ありますが、幅を広げるにつれて単調に縮み、幅3000では約0.002になります。幅を広げるほど、実際に訓練したネットワークがNTK理論の予測に収束していく様子が、はっきり見て取れます。理論が「極限の話」ではなく、有限幅への近似としても効いていることの証拠です。

理論と実物が一致することがわかったので、ここから何が言えるのか、その含意を整理しましょう。

NTKから何が言えるか(含意)

NTKの枠組みに乗せると、冒頭の謎たちがきれいに片付きます。

第一に、大域収束が保証されます。無限幅では学習が線形ODE (3) に帰着し、損失は関数空間で凸になります。NTKグラム行列 $\Theta_0$ が正定値であれば、勾配降下は必ず訓練誤差ゼロの解にたどり着きます。「非凸なのに勾配降下が効く」のは、十分広ければ実質的に凸問題を解いているからだ、という説明です。

第二に、汎化がカーネルの性質で決まります。予測は式 (5) のカーネル回帰なので、どんな関数をうまく当てられるか(汎化性能)は、NTKという固定カーネルのなめらかさ・スペクトルで決まります。スペクトルバイアスにより、低周波(なめらか)な成分が優先的に学ばれるので、過剰パラメータでも素直な関数を選びやすく、極端な過学習に陥りにくい——という見立てが立ちます。

第三に、深層学習とカーネル法・ガウス過程が橋渡しされます。「ニューラルネットは謎の黒魔術」ではなく、少なくとも無限幅の入口では、よく理解されたカーネル法の言葉で語れます。これは理論を進める足場として非常に強力です。

とはいえ、NTKは深層学習のすべてを説明する万能理論ではありません。むしろ、その「効きすぎる」性質が限界も浮かび上がらせます。

NTKの限界:特徴学習が起きない(lazy vs rich)

NTKの世界では、特徴写像 $\phi(x)=\nabla_\theta f(x;\theta_0)$ が初期値のまま固定でした。これはlazy trainingの裏返しで、「ネットワークが特徴を学習しない」ことを意味します。しかし、実際の深層学習の強さの源は、まさに特徴学習(feature learning) にあると考えられています。画像なら、層を重ねるうちにエッジ→部品→物体と、データに合わせて表現そのものを作り変えていく。これがNTKの捉えられない部分です。

この対比を図で見ましょう。

lazy regimeとrich regimeの特徴の動き

左の lazy regime(無限幅NTK)では、各データ点の特徴ベクトル(青)が訓練後(オレンジ)もほとんど動きません。固定された特徴の上で重みだけを調整しています。右の rich regime(有限幅・特徴学習)では、特徴ベクトルが大きく移動し、データに合わせて表現が再編成されています。NTKは左の世界。実用的な深層学習の旨味の多くは右の世界にある、というのが正直な評価です。

つまりNTKは、深層学習を理解するための強力な「出発点」であって、終着点ではありません。無限幅にすると確かに解析できるが、その代償として「特徴を学ぶ」という深層学習の本質が抜け落ちる。lazy regimeとrich regimeの境目、初期化スケールやパラメータ化(標準パラメータ化 vs 平均場パラメータ化)の選び方が、どちらの世界に入るかを左右する——このあたりが、NTK以降の活発な研究テーマになっています。

ここまでで理論と限界を一通り見ました。最後に、これらを自分の手で確かめられるよう、経験的NTKを計算するコードをまとめておきます。

Pythonでの実装:経験的NTKを計算する

理論の心臓は式 (2)、すなわち「パラメータ勾配の内積」です。PyTorchの自動微分を使えば、これはそのまま計算できます。ヤコビアンの記事で学んだ torch.autograd.grad の応用です。まず、簡単な2層ネットワークを定義します。出力を $1/\sqrt{m}$ でスケールするのがNTKパラメータ化のポイントです。

import numpy as np
import torch

torch.manual_seed(0)
np.random.seed(0)

class TwoLayerNet(torch.nn.Module):
    def __init__(self, width, din=1):
        super().__init__()
        self.width = width
        self.W = torch.nn.Parameter(torch.randn(width, din))
        self.b = torch.nn.Parameter(torch.randn(width))
        self.v = torch.nn.Parameter(torch.randn(width))

    def forward(self, x):
        pre = x @ self.W.t() + self.b           # (N, width)
        h = torch.tanh(pre)
        return h @ self.v / np.sqrt(self.width)  # 1/sqrt(m) スケール

このネットワークは入力1次元・隠れ層 width 個の素朴な作りです。forward の最後で幅の平方根で割っているのが、無限幅でNTKが決定的に収束するための鍵となるスケーリングです。

次に、サンプルごとにパラメータ勾配を取り、その内積でNTKグラム行列を組み立てる関数です。

def empirical_ntk(net, x1, x2):
    """経験的NTK Theta(x1,x2) = <grad_theta f(x1), grad_theta f(x2)>"""
    params = list(net.parameters())

    def jac(x):
        rows = []
        for i in range(x.shape[0]):
            net.zero_grad()
            y = net(x[i:i+1])
            g = torch.autograd.grad(y.sum(), params, retain_graph=True)
            rows.append(torch.cat([gi.reshape(-1) for gi in g]))  # 平坦化
        return torch.stack(rows)   # (n, P): 各行が phi(x_i)

    J1 = jac(x1)   # (n1, P)
    J2 = jac(x2)   # (n2, P)
    return (J1 @ J2.t()).detach()  # (n1, n2) = J1 J2^T

jac は各サンプル $x_i$ について $\phi(x_i)=\nabla_\theta f(x_i)$ を計算し、全パラメータの勾配を平坦化して1本のベクトルにしています。それを縦に積んだ行列 $J$ が、サンプル×パラメータのヤコビアン行列です。最後の J1 @ J2.t() がまさに $\Theta = J_1 J_2^\top$、式 (2) の内積をまとめて計算したものです。これだけで経験的NTKが手に入ります。

このNTKを使って、式 (5) のカーネル回帰で予測してみましょう。

torch.manual_seed(7)
Xtr = torch.linspace(-2, 2, 9).reshape(-1, 1)
ytr = torch.sin(2 * Xtr).reshape(-1).numpy()
grid = torch.linspace(-2.5, 2.5, 120).reshape(-1, 1)

net = TwoLayerNet(width=3000)
KXX = empirical_ntk(net, Xtr, Xtr).numpy()       # Theta(X, X)
KgX = empirical_ntk(net, grid, Xtr).numpy()       # Theta(x, X)

reg = 1e-6 * np.trace(KXX) / KXX.shape[0]          # 数値安定化のごく小さなリッジ
alpha = np.linalg.solve(KXX + reg * np.eye(len(Xtr)), ytr)
pred = KgX @ alpha                                  # f_inf(x) = Theta(x,X) alpha

alpha は $\Theta(X,X)^{-1} y$ を安定に解いたもので、pred がテスト格子上の予測 $f_\infty(x)=\Theta(x,X)\,\Theta(X,X)^{-1}y$ です。前掲の「カーネル回帰の閉形式予測」の図は、まさにこのコードの出力です。訓練ループは一切回していないのに、ネットワークの予測を再現できているのがポイントです。

最後に、本記事の主張「広いネットを訓練するとNTKカーネル回帰に一致する」を直接検証するコードを置きます。

torch.manual_seed(21)
Xtr = torch.linspace(-2, 2, 9).reshape(-1, 1)
ytr = torch.sin(2 * Xtr).reshape(-1)
grid = torch.linspace(-2.5, 2.5, 120).reshape(-1, 1)

net = TwoLayerNet(width=4000)
f0_g  = net(grid).detach().numpy()    # 初期出力(線形化の定数項)
f0_tr = net(Xtr).detach().numpy()

# 初期化由来の出力を差し引いた残差にカーネル回帰
KXX = empirical_ntk(net, Xtr, Xtr).numpy()
KgX = empirical_ntk(net, grid, Xtr).numpy()
reg = 1e-6 * np.trace(KXX) / KXX.shape[0]
alpha = np.linalg.solve(KXX + reg*np.eye(len(Xtr)), ytr.numpy() - f0_tr)
ntk_pred = f0_g + KgX @ alpha

# 実際に勾配降下で訓練
opt = torch.optim.SGD(net.parameters(), lr=0.2)
for _ in range(3000):
    opt.zero_grad()
    loss = ((net(Xtr) - ytr) ** 2).sum()
    loss.backward(); opt.step()
trained_pred = net(grid).detach().numpy()

rmse = np.sqrt(np.mean((ntk_pred - trained_pred) ** 2))
print("NTK予測と訓練NNの食い違い RMSE =", rmse)   # => 約 0.0004

ここでは初期出力 $f_0$ を線形化の定数項として正しく扱うため、残差 $y – f_0$ に対してカーネル回帰しています。出力の食い違いはRMSEで約 $0.0004$ となり、実際に訓練したネットワークとNTKカーネル回帰の予測がほとんど一致することが、数値で確かめられました。本記事の図はすべてこのseed固定のコードから生成しており、本文の数値(幅3000で標準偏差約0.009、幅4000でカーネル変化率0.3%、食い違いRMSE約0.0004)と一致しています。

補遺:導出を全行で追う

この補遺は読み飛ばしても大丈夫です。 ここから先は、本文で要所だけ示した導出を一行ずつ追う部分です。前半までで「無限幅で固定カーネルの線形モデルになる」というイメージは押さえられているので、結論を使えれば十分という方は次の節へ進んでください。導出を自分の手で追いたい方、論文を読めるようになりたい方のための部分です。

本文では流れを優先して要所だけを示しました。ここでは、論文(Jacot, Gabriel, Hongler, 2018; Lee et al., 2019; Arora et al., 2019)の核心となる4つの導出を、できるだけ1行ずつ追えるように書き下します。記号を最初にそろえておきます。

  • $\nabla_\theta$(ナブラ・シータ)は、パラメータベクトル $\theta\in\mathbb{R}^P$ についての勾配です。スカラー出力 $f(x;\theta)\in\mathbb{R}$ に対して $\nabla_\theta f \in\mathbb{R}^P$ は縦ベクトル($P$ 次元)です。
  • $\mathbb{E}[\,\cdot\,]$(期待値)は、初期化のランダムさ、すなわちパラメータ $\theta_0$ を引く確率分布についての平均です。
  • $\otimes$ はクロネッカー積(テンソル積)で、$u\otimes v$ は外積に相当する $\,uv^\top$ を意味します。
  • $\lambda_k$ は行列の固有値、$v_k$ はそれに対応する固有ベクトルです。
  • $X=(x_1,\dots,x_n)$ は訓練入力、$y\in\mathbb{R}^n$ は目標、$f_t=f(X;\theta_t)\in\mathbb{R}^n$ は時刻 $t$ での出力ベクトルです。

補遺1:1次テイラー展開と線形化モデル(全行)

多変数のテイラー展開を、パラメータ $\theta$ について初期点 $\theta_0$ のまわりで行います。一般に、スカラー値関数 $g(\theta)$ を $\theta_0$ のまわりで展開すると

$$ g(\theta) = g(\theta_0) + \nabla_\theta g(\theta_0)^\top (\theta-\theta_0) + \tfrac{1}{2}(\theta-\theta_0)^\top H(\theta_0)(\theta-\theta_0) + \cdots $$

です。ここで $H(\theta_0)=\nabla_\theta^2 g(\theta_0)$ はヘッセ行列(2階微分)です。$g(\theta)=f(x;\theta)$ と置き、2次以降を落とすと

$$ \begin{equation} f(x;\theta) \approx f(x;\theta_0) + \nabla_\theta f(x;\theta_0)^\top (\theta – \theta_0) \tag{A1} \end{equation} $$

を得ます。これが本文の式 (1) です。2次項を落としてよい根拠は補遺4で幅依存性として定量化します(広い幅では $H$ の寄与が $O(1/\sqrt{m})$ に潰れます)。

ここで $\phi(x):=\nabla_\theta f(x;\theta_0)\in\mathbb{R}^P$、$w:=\theta-\theta_0\in\mathbb{R}^P$ と定義します。$\phi(x)$ は $\theta$ には依らない定ベクトル($\theta_0$ で1回測ったきり固定)であることが肝心です。すると (A1) は

$$ f(x;\theta)-f(x;\theta_0) \approx \phi(x)^\top w $$

となり、右辺は重み $w$ について完全に線形です。$\phi(x)$ を固定した特徴写像、$w$ を学習対象の係数とみなせば、これは特徴 $\phi$ 上の線形回帰そのものです。つまり「非線形に見えたネットワークは、$\theta_0$ のまわりでは特徴 $\phi(x)=\nabla_\theta f(x;\theta_0)$ を持つ線形モデルに化ける」ことが、テイラー展開の1行から従います。重要なのは、非線形性はすべて固定された特徴 $\phi$ の中に押し込まれ、学習する量 $w$ の側は線形になっている、という構造です。

補遺2:訓練ダイナミクスの線形ODEとカーネル回帰(全行)

2乗損失を

$$ L(\theta) = \tfrac{1}{2}\|f(X;\theta) – y\|^2 = \tfrac{1}{2}\sum_{i=1}^{n}\big(f(x_i;\theta)-y_i\big)^2 $$

とします。まず勾配 $\nabla_\theta L$ を成分計算します。$L$ の各成分への連鎖律を使うと、$\partial L/\partial\theta_p = \sum_i (f(x_i;\theta)-y_i)\,\partial f(x_i;\theta)/\partial\theta_p$ です。これをベクトルにまとめると

$$ \nabla_\theta L(\theta) = \sum_{i=1}^{n}\big(f(x_i;\theta)-y_i\big)\,\nabla_\theta f(x_i;\theta) = J(\theta)^\top\,(f_t – y) $$

となります。ここで $J(\theta)\in\mathbb{R}^{n\times P}$ は 第 $i$ 行が $\nabla_\theta f(x_i;\theta)^\top$ のヤコビアン行列です($J_{ip}=\partial f(x_i;\theta)/\partial\theta_p$)。本文の $\nabla_\theta f_t^\top(f_t-y)$ はこの $J(\theta)^\top(f_t-y)$ の略記です。

連続時間の勾配降下(勾配流)は、損失の勾配の逆向きにパラメータを流します。

$$ \begin{equation} \frac{d\theta_t}{dt} = -\nabla_\theta L(\theta_t) = -\,J(\theta_t)^\top\,(f_t – y) \tag{A2} \end{equation} $$

次に、知りたいのは出力 $f_t=f(X;\theta_t)$ の時間発展です。$f_t$ は $\theta_t$ を通じてのみ $t$ に依存するので、合成関数の微分(連鎖律)により

$$ \frac{df_t}{dt} = \frac{\partial f(X;\theta_t)}{\partial\theta}\,\frac{d\theta_t}{dt} = J(\theta_t)\,\frac{d\theta_t}{dt} $$

です。ここでヤコビアン $J(\theta_t)$ が1回目として現れました。これに (A2) を代入すると、

$$ \frac{df_t}{dt} = J(\theta_t)\,\big(-J(\theta_t)^\top(f_t-y)\big) = -\,J(\theta_t)\,J(\theta_t)^\top\,(f_t-y) $$

となり、ヤコビアンが2回目として現れ、$J J^\top$ という積を作ります。この $n\times n$ 行列こそ訓練点上のNTKグラム行列です。

$$ \Theta_t := J(\theta_t)\,J(\theta_t)^\top, \qquad (\Theta_t)_{ij} = \nabla_\theta f(x_i;\theta_t)^\top \nabla_\theta f(x_j;\theta_t) $$

成分を見れば、$(\Theta_t)_{ij}$ は本文の式 (2) のカーネル $\Theta(x_i,x_j)$ そのものです($JJ^\top$ の $(i,j)$ 成分は $i$ 行と $j$ 行の内積だから)。lazy training(補遺4)により、広い幅では $\Theta_t\approx\Theta_0$(時間によらず一定)とみなせるので、

$$ \begin{equation} \frac{df_t}{dt} = -\,\Theta_0\,(f_t – y) \tag{A3} \end{equation} $$

が得られます。本文の式 (3) です。残差 $r_t:=f_t-y$ で書くと、$y$ は定数なので $dr_t/dt=df_t/dt$ となり

$$ \frac{dr_t}{dt} = -\,\Theta_0\,r_t $$

という、行列係数の線形・斉次・1階の常微分方程式になります。スカラーの $\dot r=-\lambda r$ の解 $r=r_0 e^{-\lambda t}$ を行列に持ち上げると、行列指数 $e^{-\Theta_0 t}:=\sum_{k\ge0}\frac{(-\Theta_0 t)^k}{k!}$ を使って

$$ r_t = e^{-\Theta_0 t}\,r_0 \quad\Longleftrightarrow\quad \begin{equation} f_t – y = e^{-\Theta_0 t}\,(f_0 – y) \tag{A4} \end{equation} $$

と閉じた形で解けます(本文の式 (4))。これが解であることは、両辺を $t$ で微分して $\frac{d}{dt}e^{-\Theta_0 t}=-\Theta_0 e^{-\Theta_0 t}$ を使えば直ちに確認できます。$\Theta_0=JJ^\top$ は半正定値(任意の $u$ に対し $u^\top JJ^\top u=\|J^\top u\|^2\ge0$)なので、固有値はすべて $\lambda_k\ge0$、よって $e^{-\Theta_0 t}\to0$($t\to\infty$)となり、残差は消え、訓練誤差はゼロに収束します。

カーネル回帰の閉形式への接続。 任意のテスト点 $x$ での出力 $f_t(x)=f(x;\theta_t)$ も同じ流れで追えます。テスト点と訓練点の間のカーネルベクトルを $\Theta(x,X)=\big(\Theta(x,x_1),\dots,\Theta(x,x_n)\big)\in\mathbb{R}^{1\times n}$ とすると、テスト点出力の時間発展は

$$ \frac{d f_t(x)}{dt} = \nabla_\theta f(x;\theta_t)^\top\frac{d\theta_t}{dt} = -\,\nabla_\theta f(x;\theta_t)^\top J(\theta_t)^\top (f_t-y) = -\,\Theta(x,X)\,(f_t-y) $$

です(最後の等号で $\nabla_\theta f(x)^\top J^\top$ の第 $i$ 成分が $\nabla_\theta f(x)^\top\nabla_\theta f(x_i)=\Theta(x,x_i)$ になることを使いました)。$(f_t-y)$ に (A4) を代入して時間で積分すると、

$$ f_t(x) = f_0(x) – \Theta(x,X)\int_0^t e^{-\Theta_0 s}\,ds\,(f_0-y) $$

となります。行列積分は $\int_0^t e^{-\Theta_0 s}ds = \Theta_0^{-1}\big(I-e^{-\Theta_0 t}\big)$($\Theta_0$ が正則のとき)なので、$t\to\infty$ で $e^{-\Theta_0 t}\to0$ を使い、$\Theta_0=\Theta(X,X)$ と書き直すと

$$ f_\infty(x) = f_0(x) – \Theta(x,X)\,\Theta(X,X)^{-1}\,(f_0 – y) $$

を得ます。これが線形化モデルでの正確な極限予測です。初期出力を無視できる(または前掲コードのように残差 $y-f_0$ に回帰する)標準的な設定では、$f_0\approx0$ とみなせて

$$ \begin{equation} f_\infty(x) = \Theta(x, X)\,\Theta(X, X)^{-1}\, y \tag{A5} \end{equation} $$

という本文の式 (5)、すなわちカーネルリッジ回帰(リッジ係数ゼロ極限)の式に一致します。本文のPython実装が「残差 $y-f_0$ にカーネル回帰し、最後に $f_0$ を足し戻す」処理をしていたのは、まさにこの $f_0$ 項を正しく扱うためです。

補遺3:固有値分解と収束レート(スペクトルバイアスの全行)

$\Theta_0$ は対称半正定値なので、スペクトル分解(固有値分解)

$$ \Theta_0 = U\Lambda U^\top = \sum_{k=1}^{n}\lambda_k\,v_k v_k^\top, \qquad \lambda_1\ge\lambda_2\ge\dots\ge\lambda_n\ge0 $$

ができます。ここで $U=[v_1,\dots,v_n]$ は直交行列($U^\top U=I$、$v_k$ は正規直交固有ベクトル)、$\Lambda=\mathrm{diag}(\lambda_1,\dots,\lambda_n)$ です。行列指数はこの基底で対角化されます。$\Theta_0^m = U\Lambda^m U^\top$ より

$$ e^{-\Theta_0 t} = \sum_{m\ge0}\frac{(-t)^m}{m!}\Theta_0^m = U\Big(\sum_{m\ge0}\frac{(-t)^m}{m!}\Lambda^m\Big)U^\top = U\,e^{-\Lambda t}\,U^\top = \sum_k e^{-\lambda_k t}\,v_k v_k^\top $$

となります(中央の式で各対角成分が独立に $e^{-\lambda_k t}$ になります)。初期残差を固有基底で $r_0=\sum_k c_k v_k$($c_k=v_k^\top r_0$)と分解して (A4) に入れると、$v_jv_j^\top v_k=\delta_{jk}v_j$(正規直交性)より

$$ \begin{equation} r_t = e^{-\Theta_0 t}\,r_0 = \sum_k e^{-\lambda_k t}\,v_kv_k^\top\Big(\sum_j c_j v_j\Big) = \sum_k c_k\,e^{-\lambda_k t}\,v_k \tag{A6} \end{equation} $$

を得ます。各固有モードが独立に時定数 $\tau_k=1/\lambda_k$ で減衰する、というのが本文の式の中身です。二乗残差はモードごとに分離して

$$ \|r_t\|^2 = \sum_k c_k^2\,e^{-2\lambda_k t} $$

となり(再び正規直交性)、大きい固有値 $\lambda_k$ の成分は速く($e^{-2\lambda_k t}$ が急減)、小さい固有値の成分はゆっくり消えることが明示的に読み取れます。これが「スペクトルバイアス」の正体です。NTKの固有関数は典型的になめらかな(低周波の)関数ほど大きな固有値を持つため、ネットワークは低周波成分を先に学び、高周波成分を後回しにします。本文の固有値スペクトル図と、規格化二乗残差が固有値順に減衰する図は、この式 (A6) をそのまま可視化したものです。

補遺4:無限幅でNTKが決定的に収束する(1隠れ層のスケーリング論証)

なぜ幅 $m\to\infty$ で $\Theta_0$ が初期化によらない決定的な値に収束し、しかも訓練中も動かないのか。完全な測度収束の証明は重いので、ここでは1隠れ層ネットワークで「期待値への集中」と「分散の $O(1/m)$ オーダー」を計算で示します。本文のコードと同じパラメータ化、

$$ f(x;\theta) = \frac{1}{\sqrt{m}}\sum_{r=1}^{m} v_r\,\sigma(w_r^\top x + b_r), \qquad w_r,b_r,v_r \overset{\text{iid}}{\sim}\mathcal{N}(0,1) $$

を考えます。$\sigma$ は活性化(コードでは $\tanh$)、$1/\sqrt{m}$ が決定的収束を生むスケーリングです。各ユニット $r$ のパラメータをまとめて $\theta_r=(w_r,b_r,v_r)$ と書きます。

NTKがユニットごとの寄与の和になること。 パラメータ勾配を成分ごとに見ます。$v_r$ についての勾配は $\partial f/\partial v_r = \frac{1}{\sqrt m}\sigma(w_r^\top x+b_r)$、$w_r$ についての勾配は $\partial f/\partial w_r = \frac{1}{\sqrt m}v_r\,\sigma'(w_r^\top x+b_r)\,x$、$b_r$ も同様です。NTKはこれらの内積の総和なので、

$$ \Theta_m(x,x’) = \sum_{r=1}^{m}\Big[\frac{\partial f(x)}{\partial v_r}\frac{\partial f(x’)}{\partial v_r} + \frac{\partial f(x)}{\partial w_r}\!\cdot\!\frac{\partial f(x’)}{\partial w_r} + \frac{\partial f(x)}{\partial b_r}\frac{\partial f(x’)}{\partial b_r}\Big] =: \frac{1}{m}\sum_{r=1}^{m} \xi_r(x,x’) $$

と書けます。ここで各 $\xi_r(x,x’)$ は、上の各勾配の $1/\sqrt m$ を外に括り出した後に残る、ユニット $r$ のパラメータ $\theta_r$ だけで決まる確率変数です。具体的には

$$ \xi_r(x,x’) = \sigma(w_r^\top x+b_r)\sigma(w_r^\top x’+b_r) + v_r^2\,\sigma'(w_r^\top x+b_r)\sigma'(w_r^\top x’+b_r)\,(x^\top x’ + 1) $$

です(第2項の $x^\top x’+1$ は $w$ 勾配の $x$ 同士の内積と $b$ 勾配の $1$ の和)。重要なのは、$\{\xi_r\}_{r=1}^m$ が 独立同分布(iid) であることです。各ユニットのパラメータが独立に引かれているからです。

大数の法則による期待値への集中。 $\Theta_m(x,x’)=\frac{1}{m}\sum_r \xi_r$ は iid 確率変数の標本平均です。よって大数の法則により、$m\to\infty$ で

$$ \Theta_m(x,x’) \;\xrightarrow{\;m\to\infty\;}\; \mathbb{E}[\xi_1(x,x’)] =: \Theta_\infty(x,x’) $$

という決定的な極限(初期化のランダムさが消えた、$x,x’$ だけの関数)に収束します。$\Theta_\infty$ は1ユニットあたりの期待値で、本文で「サイコロを1万個振って平均すれば約3.5」と例えたものの正体です。

分散が $O(1/m)$ で消えること。 集中の速さを分散で評価します。iid の標本平均の分散公式から、

$$ \mathrm{Var}\big[\Theta_m(x,x’)\big] = \mathrm{Var}\Big[\frac{1}{m}\sum_{r=1}^m \xi_r\Big] = \frac{1}{m^2}\sum_{r=1}^m \mathrm{Var}[\xi_r] = \frac{1}{m^2}\cdot m\,\mathrm{Var}[\xi_1] = \frac{\mathrm{Var}[\xi_1]}{m} $$

となります(2つ目の等号で独立性により共分散項が消え、3つ目で同分布より各分散が等しいことを使いました)。$\mathrm{Var}[\xi_1]$ は $m$ に依らない定数なので、

$$ \mathrm{Var}\big[\Theta_m(x,x’)\big] = O\!\left(\frac{1}{m}\right), \qquad \text{標準偏差} = O\!\left(\frac{1}{\sqrt{m}}\right) $$

です。これが「幅を増やすと分散が $1/m$ で消える」ことの計算による根拠です。標準偏差は $1/\sqrt m$ で減るので、幅を100倍にすると標準偏差は約10分の1になります。本文の実測(幅10で標準偏差約0.168、幅3000で約0.009)はこの $1/\sqrt m$ 則とよく整合します。実際、$0.168\times\sqrt{10/3000}=0.168/\sqrt{300}\approx0.0097$ となり、観測値0.009とほぼ一致します。$1/\sqrt m$ スケーリングが見事に効いているわけです。

訓練中の変化が $O(1/\sqrt m)$ で消えること(lazy training)。 最後に、なぜ訓練中も $\Theta_t\approx\Theta_0$ なのかをスケーリングで論じます。勾配流の解の直感として、出力を $O(1)$ だけ動かすのに必要なパラメータ変位を見積もります。各パラメータの出力への感度は $\partial f/\partial\theta_p = O(1/\sqrt m)$(上の勾配式の $1/\sqrt m$ 因子)です。パラメータは $P=O(m)$ 個あるので、全パラメータをそろえて動かしたときの出力変化は $\|\nabla_\theta f\|\cdot\|\Delta\theta\| \sim \sqrt{m}\cdot(1/\sqrt m)\cdot\|\Delta\theta\|_{\text{per param}}$ のオーダー、すなわち出力を $O(1)$ 動かすのに各パラメータは $\|\Delta\theta_p\|=O(1/\sqrt m)$ しか動かなくて済みます。

一方、$\Theta_t$ の $\theta$ についての変化率はヘッセ的な量($\nabla_\theta^2 f$)で、これも $1/\sqrt m$ のオーダーを持ちます。よってカーネルの訓練中の変化は

$$ \|\Theta_t – \Theta_0\| \sim \underbrace{\|\nabla_\theta \Theta\|}_{O(1/\sqrt m)}\cdot\underbrace{\|\theta_t-\theta_0\|}_{O(1/\sqrt m)\,\text{(出力 }O(1)\text{ あたり)}} = O\!\left(\frac{1}{\sqrt m}\right)\to 0 $$

となり、$m\to\infty$ で消えます。これが lazy training(パラメータも特徴写像もほとんど動かない)のスケーリング上の根拠です。本文の実測(幅10でカーネル相対変化約190%、幅4000で約0.3%)は、まさにこの「幅を広げるほど訓練中のカーネル変化が消える」傾向を示しています。なお、この $1/\sqrt m$ 論証は2次テイラー項(補遺1で落としたヘッセ項)の寄与が広い幅で潰れることの根拠でもあり、線形化 (A1) が訓練の最後まで正当化される理由になっています。

以上で、(i) なぜ線形化が許されるか、(ii) なぜ訓練が線形ODEになりカーネル回帰に至るか、(iii) なぜ収束速度が固有値で決まるか、(iv) なぜ無限幅でカーネルが決定的になり訓練中も凍るか、の4点を計算で追えました。完全な測度収束の証明は原論文に譲りますが、各現象が「なぜそうなるか」は上のスケーリングと初等的な確率計算で見通せます。

まとめ

本記事では、ニューラルタンジェントカーネル(NTK)について解説しました。

  • 線形化:幅が広いと訓練中もパラメータが初期値の近くにとどまり、ネットワークは初期点まわりの1次テイラー展開=線形モデルに化けます。
  • NTKの定義:その線形モデルの内積を決めるのが $\Theta(x,x’)=\nabla_\theta f(x)^\top \nabla_\theta f(x’)$ で、ヤコビアンの積 $\Theta=JJ^\top$ として計算できます。
  • 無限幅とlazy training:幅を広げるとNTKは初期化によらず決定的な値に収束し、訓練中もほぼ不変になります。
  • 訓練ダイナミクス:2乗損失なら出力は線形ODE $\dot f_t=-\Theta_0(f_t-y)$ に従い、収束速度はNTKの固有値で決まります(スペクトルバイアス)。
  • 予測:最終予測はカーネル回帰の閉形式 $f_\infty=\Theta(x,X)\Theta(X,X)^{-1}y$ になり、ガウス過程回帰と同じ形です。
  • 含意と限界:大域収束と汎化がカーネルの言葉で説明できる一方、NTKは特徴が固定のlazy regimeなので、実用の深層学習の核心である特徴学習(rich regime)は捉えられません。

NTKは「なぜ深層学習が効くのか」を理解するための、強力で誠実な出発点です。ここからさらに、特徴学習を扱う平均場理論や、有限幅補正の研究へと話は広がっていきます。

次のステップとして、以下の記事も参考にしてください。