NNGPとNeural Tangent Kernel (NTK) — 無限幅ニューラルネットワークとガウス過程の等価性

ニューラルネットワークが「なぜ」「どう」学習するのかは、長らくブラックボックスのままでした。損失曲面は非凸で、パラメータは数百万から数十億。勾配降下がどこに向かうのか、理論的にはほとんど何も言えない——これが2010年代半ばまでの常識でした。

ところが、ネットワークの幅を無限大に飛ばすという大胆な極限を取ると、突然すべてが見通せるようになります。無限幅では、初期化時のネットワークは厳密にガウス過程になり(NNGP)、勾配降下による学習はカーネル回帰に帰着します(NTK)。非線形で手に負えなかった深層学習が、線形代数で完全に解ける問題に化けるのです。これは深層学習理論における最も美しい結果の一つです。

この理論は実用面でも効いてきます。たとえば、ニューラルネットを訓練せずにそのカーネルだけで予測を出す「無限幅ネットワーク回帰」は小データのベンチマークで強力な性能を示します。また、学習率や初期化スケールの選び方、なぜ過剰パラメータのモデルが過学習しないのか(Double Descent)といった実務的な謎にも、NTK理論が定量的な答えを与えます。

本記事の内容

  • なぜ幅を無限大にすると物事が単純になるのか、その直感
  • 中心極限定理によるNNGP対応(初期化時のネットワーク = ガウス過程)
  • 層ごとのカーネル再帰式 $K^{(l+1)} = f(K^{(l)})$ とReLU/erf活性のarccosカーネルの導出
  • 学習動力学を支配するNeural Tangent Kernel (NTK) の定義と勾配流の線形化
  • lazy training(訓練中もパラメータがほとんど動かない)と関数の線形化
  • 有限幅での補正と、Double Descentとの関係
  • Pythonでの数値実験(有限幅アンサンブルのNNGP収束、NTKカーネル回帰)

前提知識

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

無限幅で何が起きるのか — 直感

まず、なぜ「幅を無限大にする」と話が単純になるのか、その感覚をつかみましょう。

1隠れ層のネットワークの出力を考えます。隠れユニットが $H$ 個あって、それぞれが入力 $x$ に対して何らかの値 $\phi_j(x)$ を出し、出力ニューロンがそれらを重み $v_j$ で足し合わせるとします。

$$ f(x) = \sum_{j=1}^{H} v_j\, \phi_j(x) $$

ここで重み $v_j$ をランダムに初期化したとします。すると $f(x)$ は、$H$ 個の独立な確率変数 $v_j \phi_j(x)$ のです。独立な確率変数をたくさん足し合わせると何が起きるか——そう、中心極限定理(CLT) です。適切にスケールすれば、$H \to \infty$ で $f(x)$ の分布はガウス分布に収束します。

これは1点 $x$ での話ですが、複数の入力点 $x_1, \dots, x_n$ をまとめて考えても、出力ベクトル $(f(x_1), \dots, f(x_n))$ は多変量ガウス分布になります。任意の有限個の点で多変量ガウスになる確率過程——これはまさにガウス過程(GP) の定義そのものです。つまり、無限幅のランダム初期化ネットワークは、ガウス過程と等価になります。これが NNGP(Neural Network Gaussian Process) です。

歴史的には、Radford Neal が1996年の博士論文(Bayesian Learning for Neural Networks)で、1隠れ層の無限幅ネットワークがガウス過程になることを最初に示しました。それから20年以上を経て、Jaehoon Lee らが2018年の論文「Deep Neural Networks as Gaussian Processes」で、これを任意の深さの深層ネットワークに拡張しました。

一方で、これは「初期化時」の話に過ぎません。学習させたらどうなるのか? 驚くべきことに、無限幅では学習動力学さえも線形化され、解析的に追えるようになります。これを記述するのが Arthur Jacot らが2018年に提唱した Neural Tangent Kernel (NTK) です。

無限幅ネットワークの出力分布が中心極限定理でガウスに収束する様子

上の図は、1隠れ層ネットワーク(tanh活性)の固定入力点 $x_0$ における出力 $f(x_0)$ の分布を、幅 $H$ を変えてヒストグラムで示したものです。$H$ が小さいうちは分布がいびつですが、$H$ を大きくすると赤線のガウス近似にぴったり重なっていきます。中心極限定理が、まさに目に見える形で働いていることがわかります。

ここまでで「無限幅 → ガウス過程」という大枠の直感をつかみました。次は、これを数式できちんと正当化していきましょう。

NNGP — 初期化時のガウス過程対応

1隠れ層での中心極限定理

具体的な設定で、NNGP対応を導出します。入力 $x \in \mathbb{R}^{d}$、1隠れ層、隠れ幅 $H$ のネットワークを考えます。

$$ \begin{aligned} z_j(x) &= \sum_{k=1}^{d} W^{(1)}_{jk}\, x_k + b^{(1)}_j \quad (j = 1, \dots, H) \\ f(x) &= b^{(2)} + \sum_{j=1}^{H} W^{(2)}_j\, \phi\big(z_j(x)\big) \end{aligned} $$

ここで $\phi$ は活性化関数(ReLU や tanh など)、$z_j$ は $j$ 番目の隠れユニットの「前活性(pre-activation)」です。重みとバイアスは独立に、平均 $0$ のガウス分布から初期化します。

$$ W^{(1)}_{jk} \sim \mathcal{N}\!\left(0, \frac{\sigma_w^2}{d}\right), \quad b^{(1)}_j \sim \mathcal{N}(0, \sigma_b^2), \quad W^{(2)}_j \sim \mathcal{N}\!\left(0, \frac{\sigma_v^2}{H}\right) $$

ポイントは出力重み $W^{(2)}_j$ の分散を $\sigma_v^2/H$ と、幅 $H$ で割ってスケールしていることです。これがNNGPスケーリング(別名 “standard parametrization”)です。

出力 $f(x)$ を、$H$ 個の項の和として見ます。各項 $g_j(x) = W^{(2)}_j\, \phi(z_j(x))$ は独立同分布(i.i.d.)です(重みが独立なので)。期待値は $\mathbb{E}[g_j] = \mathbb{E}[W^{(2)}_j]\, \mathbb{E}[\phi(z_j)] = 0$($W^{(2)}_j$ の平均が $0$)。分散は有限です。

ここで中心極限定理を適用します。$H$ 個の i.i.d. な確率変数の和が、$H \to \infty$ で正規分布に収束します。$1/H$ スケーリングにより分散が発散せず収束するため、極限で $f(x)$ は平均 $0$ のガウス分布になります。

$$ f(x) \xrightarrow{H \to \infty} \mathcal{N}\big(0,\ K(x, x)\big) $$

多入力での共分散 = カーネル

複数の入力点でも同じ議論が成り立ち、$(f(x), f(x’))$ は2次元ガウスになります。その共分散を計算しましょう。

$$ \begin{aligned} \mathrm{Cov}\big(f(x), f(x’)\big) &= \mathbb{E}\!\left[\Big(\sum_j W^{(2)}_j \phi(z_j(x))\Big)\Big(\sum_{j’} W^{(2)}_{j’} \phi(z_{j’}(x’))\Big)\right] \end{aligned} $$

$W^{(2)}_j$ は独立で平均 $0$ なので、$j \neq j’$ の交差項は期待値が $0$ になり、対角項 $j = j’$ だけが残ります。$\mathbb{E}[(W^{(2)}_j)^2] = \sigma_v^2/H$ を使うと、

$$ \mathrm{Cov}\big(f(x), f(x’)\big) = \sum_{j=1}^{H} \frac{\sigma_v^2}{H}\, \mathbb{E}\big[\phi(z_j(x))\, \phi(z_j(x’))\big] $$

となります。和の中身は $j$ によらず同じ値(i.i.d.)なので、$H$ 個足して $1/H$ を掛けると相殺し、次のカーネルが得られます。

$$ \boxed{\,K(x, x’) = \sigma_b^2 + \sigma_v^2\, \mathbb{E}_{z \sim \mathcal{N}(0,\Sigma)}\big[\phi(z(x))\, \phi(z(x’))\big]\,} $$

ここで $z(x), z(x’)$ は前活性で、それらの2変量ガウス分布の共分散行列 $\Sigma$ は入力から決まります($\Sigma_{xx’} = \sigma_b^2 + \frac{\sigma_w^2}{d} x^\top x’$)。バイアス項 $\sigma_b^2$ は出力バイアス $b^{(2)}$ の寄与です。

この $K(x, x’)$ こそが NNGPカーネル です。任意の有限個の入力点で多変量ガウスになり、その共分散がこのカーネルで与えられる——だからこのネットワークはガウス過程 $\mathcal{GP}(0, K)$ そのものなのです。

ここで自然な疑問が湧きます。層を深くしたら、このカーネルはどう変わるのでしょうか? 実は、深さ方向にカーネルの再帰式が現れます。それを次に導きます。

層ごとのカーネル再帰式

再帰の構造

深さ $L$ のネットワークでは、各層の前活性 $z^{(l)}(x)$ が、前の層の出力から計算されます。無限幅では各層の前活性もガウス過程になり、その共分散カーネル $K^{(l)}(x, x’)$ が層ごとに定義されます。

ここで決定的に重要な事実があります。無限幅では、層を1つ進めるカーネルの更新が、決定的な写像 $f$ で書けるのです。

$$ K^{(l+1)}(x, x’) = \sigma_b^2 + \sigma_w^2\, \mathbb{E}_{(u, v) \sim \mathcal{N}(0, \Lambda^{(l)})}\big[\phi(u)\, \phi(v)\big] $$

ここで $\Lambda^{(l)}$ は、入力 $x, x’$ の第 $l$ 層カーネルから決まる $2 \times 2$ の共分散行列です。

$$ \Lambda^{(l)} = \begin{pmatrix} K^{(l)}(x,x) & K^{(l)}(x,x’) \\ K^{(l)}(x’,x) & K^{(l)}(x’,x’) \end{pmatrix} $$

つまり、第 $l$ 層のカーネル値だけから第 $l+1$ 層のカーネル値が決まる閉じた再帰になっています。これを略記して $K^{(l+1)} = f(K^{(l)})$ と書きます。$f$ は活性化関数 $\phi$ で決まる「カーネル写像」です。

初期条件は入力カーネル $K^{(0)}(x, x’) = \frac{\sigma_w^2}{d} x^\top x’ + \sigma_b^2$ で、これを $L$ 回反復すれば最終層のNNGPカーネル $K^{(L)}$ が得られます。深層NNGPは、この単純な反復で完全に記述されるのです。

ReLU活性のarccosカーネル

問題は、期待値 $\mathbb{E}[\phi(u)\phi(v)]$ を具体的に計算できるかどうかです。幸い、ReLU と erf(誤差関数)では閉じた形が知られています。これは Cho と Saul の2009年の論文「Kernel Methods for Deep Learning」が導いた arccosカーネル です。

ReLU $\phi(u) = \max(u, 0)$ の場合を導きましょう。$(u, v)$ が平均 $0$ の2変量ガウスで、$u$ の分散 $q_1$、$v$ の分散 $q_2$、相関係数 $c = \frac{K^{(l)}(x,x’)}{\sqrt{q_1 q_2}}$ とします。ゴールは $\mathbb{E}[\max(u,0)\max(v,0)]$ を $q_1, q_2, c$ で表すことです。

角度 $\theta = \arccos(c)$ を導入すると、ReLUの2次のarccosカーネルは次の形になります。

$$ \mathbb{E}[\max(u,0)\max(v,0)] = \frac{\sqrt{q_1 q_2}}{2\pi}\Big(\sin\theta + (\pi – \theta)\cos\theta\Big) $$

導出の要点だけ追います。極座標的な積分 $\int\int \max(u,0)\max(v,0)\, p(u,v)\, du\, dv$ を、相関 $c = \cos\theta$ をパラメータとして実行します。被積分関数が両方正($u>0, v>0$)の象限でのみ非ゼロになることを使い、ガウス積分を角度で評価すると、上の閉じた形が出てきます(詳細はCho-Saulの原論文を参照)。

これを再帰式に代入すると、ReLUネットワークのNNGPカーネル再帰が得られます。$q_1 = K^{(l)}(x,x)$、$q_2 = K^{(l)}(x’,x’)$ として、

$$ K^{(l+1)}(x, x’) = \sigma_b^2 + \frac{\sigma_w^2}{2\pi}\sqrt{q_1 q_2}\,\Big(\sin\theta^{(l)} + (\pi – \theta^{(l)})\cos\theta^{(l)}\Big) $$

$$ \theta^{(l)} = \arccos\!\left(\frac{K^{(l)}(x, x’)}{\sqrt{q_1 q_2}}\right) $$

これがReLU深層NNGPの心臓部です。erf活性の場合も同様に、$\arcsin$ を含む閉じた形になります(後の図で比較します)。

NNGPカーネルの層ごとの再帰による相関の平坦化

上の図は、単位円上に並べた入力点について、各層でのカーネル相関を「入力角度差 $\Delta\theta$ の関数」として描いたものです。第 $l=0$ 層(入力)では相関が角度差とともに大きく変化しますが、層を深くするにつれて相関の曲線がだんだん平坦になり、すべての入力ペアの相関が $1$ に近づいていきます。これは深いReLUネットワークの「相関が飽和する」性質を示しており、深さを増やすと入力の区別がつきにくくなる(情報が失われる)危険があることを意味します。

カーネル写像の不動点

カーネル写像 $f$ の振る舞いを、相関係数 $c$ の反復として見るとさらに見通しが良くなります。正規化したカーネル写像を $c_{l+1} = f_{\text{norm}}(c_l)$ と書くと、ReLUでは $c^* = 1$ が安定な不動点になります。

arccosカーネル写像とその不動点構造のコブウェブ図

左図は、入力相関 $c = \cos\theta$ に対する1層分のカーネル写像 $f(c)$ を、ReLU・erf・線形の3つで比較しています。ReLU(赤)は $c$ に対して凸な単調増加曲線で、対角線 $c_{l+1}=c_l$ より上にあります。右図のコブウェブ図は、初期相関 $c_0 = -0.8$ から写像を反復すると、緑の階段が次第に右上の不動点 $c^* = 1$ に吸い寄せられる様子を示します。つまり層を重ねると、どんな入力ペアの相関も最終的に $1$ に飽和する——これが先ほど見た「相関の平坦化」の数学的な正体です。

ここまでで初期化時のNNGPカーネルが、層再帰とarccosカーネルで完全に計算できることがわかりました。このカーネルを使えば、ネットワークを訓練せずにガウス過程回帰の予測が出せます。実際に予測を見てみましょう。

NNGPによるベイズ推論

NNGPカーネル $K$ が手に入れば、あとは普通のガウス過程回帰です。訓練データ $\{(x_i, y_i)\}_{i=1}^n$ に対し、テスト点 $x_*$ での予測分布は、ガウス過程回帰の公式そのままで与えられます。

$$ \begin{aligned} \mu_*(x_*) &= K(x_*, X)\,\big(K(X,X) + \sigma_n^2 I\big)^{-1}\, y \\ \sigma_*^2(x_*) &= K(x_*, x_*) – K(x_*, X)\,\big(K(X,X) + \sigma_n^2 I\big)^{-1}\, K(X, x_*) \end{aligned} $$

ここで $K(X,X)$ は訓練点どうしのカーネル行列、$\sigma_n^2$ は観測ノイズ分散です。重要なのは、ここにニューラルネットの訓練が一切ないことです。無限幅ネットワークのベイズ事後分布が、カーネル回帰の閉形式解にぴったり一致するのです。

深さ別のNNGP事後予測の比較

上の図は、深さ $L=1$ と $L=4$ のReLU-NNGPカーネルでガウス過程回帰を行った結果です。観測点(橙)の近くでは平均予測(青)が真の関数(緑破線)をよくなぞり、信頼区間(薄青帯)が狭くなります。一方、観測点から離れると信頼区間が大きく広がり、不確実性が正直に表現されます。深さ $L=4$ のほうが観測点付近でより柔軟にフィットしているのが見て取れ、深層NNGPのカーネルが浅いものより表現力に富むことがわかります。

ここまでは「初期化時」「ベイズ事後」の話でした。しかし実際の深層学習は、勾配降下でパラメータを動かして学習します。この学習動力学を無限幅で記述するのが、いよいよ本題のNeural Tangent Kernelです。

Neural Tangent Kernel (NTK) — 学習動力学

関数空間での勾配流

ネットワークの出力を $f(x; \theta)$ と書きます($\theta$ は全パラメータ)。二乗損失 $\mathcal{L} = \frac{1}{2}\sum_i (f(x_i;\theta) – y_i)^2$ を勾配降下で最小化するとき、パラメータの更新は連続時間の極限(勾配流)で次のように書けます。

$$ \frac{d\theta}{dt} = -\nabla_\theta \mathcal{L} = -\sum_i \big(f(x_i;\theta) – y_i\big)\, \nabla_\theta f(x_i; \theta) $$

ここで知りたいのは、パラメータ $\theta$ そのものではなく、関数 $f$ がどう動くかです。連鎖律で $f(x;\theta)$ の時間変化を計算します。

$$ \frac{d f(x; \theta)}{dt} = \nabla_\theta f(x;\theta)^\top \frac{d\theta}{dt} = -\sum_i \big(f(x_i;\theta) – y_i\big)\, \nabla_\theta f(x;\theta)^\top \nabla_\theta f(x_i;\theta) $$

ここに現れた内積こそが Neural Tangent Kernel です。

$$ \boxed{\,\Theta(x, x’) = \nabla_\theta f(x; \theta)^\top \nabla_\theta f(x’; \theta) = \sum_p \frac{\partial f(x;\theta)}{\partial \theta_p}\frac{\partial f(x’;\theta)}{\partial \theta_p}\,} $$

NTKは「2つの入力に対する出力勾配の内積」です。これを使うと、関数の学習動力学が次のシンプルな形になります。

$$ \frac{d f(x;\theta)}{dt} = -\sum_i \Theta(x, x_i)\big(f(x_i;\theta) – y_i\big) $$

つまり、関数 $f$ はNTKをカーネルとするカーネル勾配流に従って進化します。これはまだ厳密な式で、$\Theta$ は $\theta$ に依存して時々刻々と変化します。一般のネットワークでは $\Theta$ が動くので解けません。

NTKの概念図 — 勾配降下による関数の進化を支配するカーネル

上の概念図は、NTKがパラメータ空間と関数空間をつなぐ役割を整理したものです。パラメータ $\theta$ が勾配降下で更新されると出力関数 $f$ が動きますが、その動き方を支配するのが勾配の内積 = NTK です。そして右側の枠が示すように、幅を無限大にするとNTKは一定になり、学習が線形なカーネル勾配流に帰着します。

無限幅でのNTKの一定性

Jacot らの中心的な定理は次の2点です。

  1. 初期化時: 幅 $\to \infty$ で、ランダム初期化のNTK $\Theta_0(x,x’)$ は決定的なカーネル(ランダムでない固定値)に収束する。
  2. 訓練中: 幅 $\to \infty$ で、NTKは訓練を通じて一定に保たれる($\Theta_t = \Theta_0$ for all $t$)。

この2つが成り立つと、学習動力学の式で $\Theta$ を定数 $\Theta_0$ に置き換えられます。残差 $u_t(X) = f_t(X) – y$ について、

$$ \frac{d u_t(X)}{dt} = -\Theta_0(X, X)\, u_t(X) $$

という線形の常微分方程式になります。これは行列指数で解けて、

$$ u_t(X) = e^{-\Theta_0(X,X)\, t}\, u_0(X) $$

$t \to \infty$ で $u_\infty(X) = 0$、つまり訓練データを完全に補間します。テスト点 $x_*$ での予測は、NTKによるカーネル回帰の解に収束します。

$$ f_\infty(x_*) = f_0(x_*) + \Theta_0(x_*, X)\,\Theta_0(X,X)^{-1}\big(y – f_0(X)\big) $$

これは驚くべき結果です。無限幅ネットワークの勾配降下学習は、NTKをカーネルとするカーネルリッジ回帰(リッジなし極限)と数式的に等価なのです。非線形な深層学習が、線形代数だけで解ける問題に化けました。

なぜNTKが一定に保たれるのか——その鍵が「lazy training」という現象です。次に見ていきます。

lazy training と関数の線形化

なぜパラメータは動かないのか

NTKが訓練中も一定である理由は、無限幅では各パラメータがほとんど動かないからです。これを lazy training(怠惰な学習)と呼びます。

直感的に説明しましょう。出力 $f = \frac{1}{\sqrt{H}}\sum_j W^{(2)}_j \phi(z_j)$ で、勾配降下による各重み $W^{(2)}_j$ の更新量はおよそ $O(1/\sqrt{H})$ です。一方、出力 $f$ への寄与は $H$ 個の重みの和なので $O(\sqrt{H} \times 1/\sqrt{H}) = O(1)$ になり、関数はちゃんと $O(1)$ だけ動きます。つまり、個々のパラメータの相対的な移動量は幅とともにゼロに近づくのに、それらが大量に集まることで関数全体は有限に変化するのです。

各パラメータがほとんど動かないなら、出力を初期パラメータ $\theta_0$ のまわりで1次のテイラー展開(線形化)しても、訓練の全期間で良い近似になります。

$$ f(x; \theta) \approx f(x; \theta_0) + \nabla_\theta f(x;\theta_0)^\top (\theta – \theta_0) $$

この線形化モデルでは、$f$ はパラメータ $\theta$ について線形です。線形モデルの勾配降下はカーネル回帰そのものですから、NTKが一定($\Theta_t = \Theta_0$)になり、前節の解析が厳密に成立します。lazy training は、無限幅での学習を「特徴 $\nabla_\theta f(x;\theta_0)$ を固定した線形回帰」に変えてしまうのです。

lazy training — 幅が広いほどパラメータの相対移動量が小さい

上の図は、幅 $H$ を変えて同じ回帰問題を勾配降下で学習させた結果です。左図(a)では、どの幅でも訓練損失はきちんと下がっています(関数は学習している)。一方、右図(b)の「相対パラメータ変化 $\|\Delta\theta\|/\|\theta_0\|$」を見ると、幅 $H=2000$ ではパラメータがほとんど動いていません(ほぼ平らな緑線)。狭い $H=10$ ではパラメータが大きく動いています。これが lazy training の核心です。幅が広いほど、関数は学習するのに各パラメータはほとんど動かない——だからこそ線形化が成り立つのです。

NTKの安定性を確かめる

NTKの一定性も数値で確認できます。

幅別のNTKの安定性 — 経験NTKの収束と訓練中の不変性

左図(a)は、有限幅ネットワークの「経験的NTK」$\Theta(x_0, x)$ を幅別にプロットしたものです。幅 $H=50$ と $H=2000$ の曲線はほぼ重なっており、幅を増やすと経験NTKが決定的な形に収束していることがわかります(狭い $H=5$ だけ少しずれます)。右図(b)は、訓練に伴うNTKの相対変化量を示します。幅 $H=2000$ ではNTKが訓練中ほとんど変化しない(緑線が低いまま)のに対し、狭い $H=5$ では大きく変動します。理論が予言する「無限幅でNTKは決定的かつ一定」が、有限幅でも幅とともに成り立っていく様子が確認できます。

ここまでで無限幅の理論を完成させました。しかし実際のネットワークは有限幅です。理論からのズレはどう扱えばよいのでしょうか?

有限幅補正と実用的な含意

有限幅でのズレ

無限幅の理論は美しいですが、実際のネットワークは有限幅です。有限幅では、

  • NNGPカーネルやNTKがランダムになる(初期化ごとにばらつく。幅とともにそのばらつきは $O(1/\sqrt{H})$ で減少)
  • NTKが訓練中に変化する(lazy training が完全には成り立たず、特徴学習が起こる)

という補正が入ります。これらの有限幅効果は $1/H$ のオーダーで展開でき、近年は系統的な摂動論(Roberts・Yaida らの「The Principles of Deep Learning Theory」など)で扱われています。

重要なのは、有限幅の「特徴学習」こそが実用ネットワークの性能の源泉だという点です。無限幅NTK領域では特徴($\nabla_\theta f$)が初期値に固定されるため、データに合わせて表現を学習しません。これがNTK領域の限界で、実際の深層学習がNTK予測を上回ることが多い理由でもあります。

スケーリングと初期化の指針

それでもNTK理論は実用的な指針を与えます。たとえば前活性の分散が層を通じて発散も消滅もしない「臨界初期化」($\sigma_w^2$ の適切な選択)は、カーネル再帰の不動点解析から導けます。これは勾配消失・爆発を避ける初期化の理論的根拠になります。学習率の上限($\eta < 2/\lambda_{\max}(\Theta)$)も、NTKの最大固有値から決まります。

アンサンブルでのNNGP収束

NNGP対応も実験で直接見えます。有限幅ネットワークを多数(アンサンブル)用意してその統計を取ると、幅を増やすにつれてNNGP(ガウス過程)の予測に収束します。

有限幅NNのアンサンブルがNNGPに収束する様子

左図(a)は、固定点 $x_0$ での出力分散をアンサンブルで推定し、幅 $H$ に対してプロットしたものです。幅を増やすと分散がNNGPの値(赤破線)に収束していきます。右図(b)は関数の事前分布の「帯」(平均 $\pm 2\sigma$)で、狭い幅($H=3$)と広い幅($H=3000$)でほぼ同じガウス過程事前に一致することが見て取れます。有限幅NNのアンサンブルは、確かに無限幅のガウス過程に近づいていくのです。

そしてNTK理論が説明する最も有名な現象が、過剰パラメータ化のパラドックス——Double Descent です。次に見ましょう。

Double Descentとの関係

古典的な統計学では、モデルの複雑さを上げるとテスト誤差は「U字型」を描くと教わります(バイアス・バリアンスのトレードオフ)。ところが深層学習では、パラメータ数を訓練データ数よりはるかに増やしても過学習せず、むしろ性能が上がります。この一見矛盾した現象が Double Descent(二重降下) です(Belkin らの「Reconciling modern machine learning and the bias-variance trade-off」)。

直感的にはこうです。モデル容量がデータ数 $n$ にちょうど等しい「補間しきい値」では、モデルが訓練データを無理やり通すために激しく振動し、テスト誤差がピークになります。ところが容量をさらに増やすと、無数にある補間解の中から「最小ノルム解」(最も滑らかな解)が選ばれるようになり、テスト誤差が再び下がるのです。

NTK理論はこれを定量的に説明します。無限幅NTKは最小ノルム補間解に対応しており、過剰パラメータ領域で滑らかな解を実現します。これが、巨大なネットワークが過学習しない理由の理論的根拠の一つです。

Double Descent — 補間しきい値でテスト誤差がピークになり過剰パラメータ域で再低下

上の図は、ランダム特徴回帰でDouble Descentを再現したものです。横軸はモデル容量 $P$(ランダム特徴数)、縦軸は誤差です。訓練誤差(青)は補間しきい値 $P = n = 40$ でゼロに落ち込みます。一方、テスト誤差(赤)はちょうど $P = n$ でピークを作り、そこを超えて過剰パラメータ領域に入ると再び下がっていきます。古典的なU字曲線(左半分)と、現代的な「過剰パラメータでむしろ良くなる」領域(右半分)が、1枚の図に同居しているのがわかります。NTK領域での最小ノルム解の振る舞いが、この二度目の降下を生み出しています。

理論の予言が数値実験と一致するか、最後にNTKカーネル回帰と実際のNN学習を直接比較しましょう。

Pythonでの実験 — NTKカーネル回帰とNN学習の一致

NTK理論の核心的な主張は「無限幅NNの勾配降下学習 = NTKカーネル回帰」でした。これを直接確かめます。同じ初期パラメータを共有する広幅($H=4000$)ネットワークについて、(1) NTK線形化のカーネル回帰公式による予測と、(2) 実際に勾配降下で訓練した予測を比較します。

まず経験的NTK行列を計算する関数を用意します。1隠れ層ReLU網(NTKパラメトリゼーション $f = \frac{1}{\sqrt{H}} W^{(2)\top}\mathrm{ReLU}(\frac{\sigma_w}{\sqrt{d}} W^{(1)} x + b)$)の出力勾配の内積を解析的に組み立てます。

import numpy as np

def empirical_ntk_1d(X, H, rng, sw=np.sqrt(2.0), params=None):
    """1隠れ層ReLU網の経験的NTK行列 Θ(X,X) を解析式で計算する。
    X: (n,d) 入力, H: 隠れ幅。返り値: (Θ, (W1,b1,W2))."""
    d = X.shape[1]
    if params is None:
        W1 = rng.normal(0, 1, size=(d, H))
        b1 = np.zeros(H)
        W2 = rng.normal(0, 1, size=(H, 1))
    else:
        W1, b1, W2 = params
    pre = X @ W1 * sw / np.sqrt(d) + b1        # 前活性 (n,H)
    h = np.maximum(pre, 0)                      # ReLU出力 (n,H)
    act = (pre > 0).astype(float)               # ReLUの微分 (n,H)

    # 各パラメータでの出力勾配の内積を足し合わせる
    JW2 = h / np.sqrt(H)                         # ∂f/∂W2_j
    coef = (W2.ravel()[None, :] * act) / np.sqrt(H)  # 隠れ層側の係数
    Theta = JW2 @ JW2.T                          # W2 部の寄与
    Theta += (coef @ coef.T) * (sw**2 / d) * (X @ X.T)  # W1 部の寄与
    Theta += coef @ coef.T                        # b1 部の寄与
    return Theta, (W1, b1, W2)

この関数は、NTKの定義 $\Theta(x,x’) = \sum_p \partial_{\theta_p} f(x)\, \partial_{\theta_p} f(x’)$ を、各パラメータ群($W^{(2)}, W^{(1)}, b^{(1)}$)ごとに勾配を求めて内積を取り、足し合わせています。ReLUの微分が0/1のステップ関数になることを act で表現しているのがポイントです。

次に、このNTKを使ったカーネル回帰の予測と、実際の勾配降下学習を比較します。

import numpy as np

rng = np.random.default_rng(1)
Xtr = np.sort(rng.uniform(-2, 2, 12)).reshape(-1, 1)
ytr = np.sin(2.0 * Xtr)
Xte = np.linspace(-2.6, 2.6, 200).reshape(-1, 1)
H = 4000

# 訓練+テストを合同でNTKを作り、同じ初期パラメータを共有
Xall = np.vstack([Xtr, Xte]); ntr = len(Xtr)
Th_all, params = empirical_ntk_1d(Xall, H, np.random.default_rng(7))
Th_trtr = Th_all[:ntr, :ntr]
Th_tetr = Th_all[ntr:, :ntr]
W1, b1, W2 = params
d = 1; sw = np.sqrt(2.0)

def forward(Xb):
    pre = Xb @ W1 * sw / np.sqrt(d) + b1
    h = np.maximum(pre, 0)
    return h @ W2 / np.sqrt(H), h, pre

# 初期出力 f0(線形化の基点)
f0_tr = forward(Xtr)[0]
f0_te = forward(Xte)[0].ravel()
# NTK線形化の予測: f(x*) = f0(x*) + Θ(x*,X) Θ(X,X)^{-1} (y - f0(X))
reg = 1e-7 * np.trace(Th_trtr) / ntr
alpha = np.linalg.solve(Th_trtr + reg * np.eye(ntr), ytr - f0_tr)
f_ntk = f0_te + (Th_tetr @ alpha).ravel()

このコードは、NTK理論の予測公式 $f_\infty(x_*) = f_0(x_*) + \Theta(x_*, X)\Theta(X,X)^{-1}(y – f_0(X))$ をそのまま実装しています。初期出力 $f_0$ を基点に、残差 $y – f_0$ をNTKカーネル回帰で補間しています。続いて、同じ初期パラメータから実際に勾配降下で訓練します。

# 実際のNN(同じ初期パラメータ)を勾配降下で学習
lr = 0.2
for t in range(4000):
    out, h, pre = forward(Xtr)
    r = out - ytr
    g_out = r / len(Xtr)
    gW2 = (h.T @ g_out) / np.sqrt(H)
    gh = (g_out @ W2.T) / np.sqrt(H)
    gpre = gh * (pre > 0)
    gW1 = (Xtr.T @ gpre) * sw / np.sqrt(d)
    gb1 = gpre.sum(axis=0)
    W1 -= lr * gW1; b1 -= lr * gb1; W2 -= lr * gW2
f_nn = forward(Xte)[0].ravel()

print("NTK回帰とNN学習の最大差:", np.max(np.abs(f_ntk - f_nn)))

この実験を可視化したのが次の図です。

NTKカーネル回帰の予測と実際の広幅NN学習の予測の一致

このグラフから、NTKカーネル回帰(紫の理論曲線)と実際のNN学習(赤の点線)が同じ傾向で観測点を通っていることがわかります。両者とも、わずか12点の観測(橙)を補間する形で予測を出しています。理論で導いた「無限幅NNの勾配降下 = NTKカーネル回帰」という主張が、$H=4000$ という有限だが大きな幅でおおむね成り立っていることが確認できます。ここでの紫線が区分線形なのは、有限幅のNTKがReLUの折れ線特徴で構成されているためで、これは理論的に正しい振る舞いです。完全な一致には $H \to \infty$ が必要ですが、有限幅でも近似的に成立する様子が見て取れます。

まとめ

本記事では、無限幅ニューラルネットワークとガウス過程・カーネル回帰の等価性を、NNGPとNeural Tangent Kernelの観点から理論重視で解説しました。

  • NNGP(初期化時のGP対応): 中心極限定理により、無限幅のランダム初期化ネットワークはガウス過程 $\mathcal{GP}(0, K)$ になる。共分散カーネルは $K(x,x’) = \sigma_b^2 + \sigma_v^2\,\mathbb{E}[\phi(z(x))\phi(z(x’))]$ で与えられる(Neal 1996, Lee et al. 2018)。
  • カーネル再帰式: 深層では $K^{(l+1)} = f(K^{(l)})$ という決定的な層再帰が成り立つ。ReLU/erf活性では期待値が閉形式(arccosカーネル)で計算でき、深さとともに相関が不動点 $c^*=1$ に飽和する。
  • NTK(学習動力学): 勾配降下による関数の進化は $\frac{df}{dt} = -\Theta(x,X)(f(X)-y)$ という、NTK $\Theta(x,x’) = \nabla_\theta f^\top \nabla_\theta f’$ をカーネルとする勾配流で記述される(Jacot et al. 2018)。
  • lazy training と線形化: 無限幅では各パラメータがほとんど動かず(lazy training)、出力が $\theta_0$ まわりで線形化される。NTKは初期化時に決定的で、訓練中も一定。結果として学習はNTKカーネル回帰 $f_\infty(x_*) = f_0(x_*) + \Theta(x_*,X)\Theta(X,X)^{-1}(y-f_0(X))$ に帰着する。
  • 有限幅補正とDouble Descent: 有限幅ではカーネルがランダムかつ訓練中に変化し(特徴学習)、$1/H$ のオーダーで補正が入る。最小ノルム補間解の性質から、補間しきい値 $P=n$ でテスト誤差がピークを作り過剰パラメータ域で再び下がる Double Descent が説明できる。

無限幅という極限は、深層学習を線形代数で完全に解ける問題に変換し、初期化・学習率・汎化といった実務的な謎に理論的な答えを与えます。同時に「無限幅NTK領域は特徴学習をしない」という限界も明らかにし、有限幅の特徴学習こそが現代の深層学習の性能の源泉だという理解にもつながります。NNGPとNTKは、深層学習理論を学ぶうえでの出発点となる枠組みです。

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