NTKの導出と性質

ニューラルタンジェントカーネル(Neural Tangent Kernel, NTK)は、ニューラルネットワークの学習ダイナミクスを関数空間で記述するために導入されるカーネルであり、その導出は「勾配降下法における関数値の時間発展」を厳密に追跡することから得られる。以下では、その導出過程と数学的性質を体系的に述べる。

基本設定と目的

パラメータ $\theta \in \mathbb{R}^P$ を持つニューラルネットワーク $f(x;\theta)$ を考え、訓練データ $(x_i, y_i)_{i=1}^n$ に対する二乗損失

\[\mathcal{L}(\theta) = \frac{1}{2} \sum_{i=1}^n (f(x_i;\theta) - y_i)^2\]

を最小化する問題を考える。目的は、勾配降下法によりパラメータが更新されるとき、関数値 $f(x;\theta)$ がどのように変化するかを直接記述することである。

勾配降下法と関数の時間発展

連続時間極限(gradient flow)では、パラメータは

\[\frac{d\theta_t}{dt} = - \nabla_\theta \mathcal{L}(\theta_t)\]

に従う。このとき任意の入力 $x$ に対する出力の時間変化は連鎖律により

\[\frac{d}{dt} f(x;\theta_t) = \nabla_\theta f(x;\theta_t)^\top \frac{d\theta_t}{dt}\]

と書ける。

損失の勾配は

\[\nabla_\theta \mathcal{L}(\theta) = \sum_{i=1}^n (f(x_i;\theta) - y_i)\, \nabla_\theta f(x_i;\theta)\]

であるため、代入すると

\[\frac{d}{dt} f(x;\theta_t) = - \sum_{i=1}^n (f(x_i;\theta_t) - y_i)\, \nabla_\theta f(x;\theta_t)^\top \nabla_\theta f(x_i;\theta_t)\]

を得る。

NTKの導出

ここで、カーネル関数

\[K_{\theta_t}(x, x') = \nabla_\theta f(x;\theta_t)^\top \nabla_\theta f(x';\theta_t)\]

を定義すると、関数の時間発展は

\[\frac{d}{dt} f(x;\theta_t) = - \sum_{i=1}^n K_{\theta_t}(x, x_i)\, (f(x_i;\theta_t) - y_i)\]

と書き直される。

特に訓練点上では、ベクトル $f_t = (f(x_1;\theta_t),\dots,f(x_n;\theta_t))$ に対して

\[\frac{d}{dt} f_t = - K_{\theta_t} (f_t - y)\]

という閉じた常微分方程式が得られる。ここで $K_{\theta_t}$ はカーネル行列である。

無限幅極限とカーネルの収束

各層の幅を無限大にした極限では、適切な初期化のもとで以下が成立する:

\[K_{\theta_t}(x,x') \xrightarrow{n \to \infty} K(x,x') \quad \text{(確率収束)}\]

さらに重要な点として、この極限では

\[K_{\theta_t}(x,x') \approx K_{\theta_0}(x,x')\]

が時間に依存せず成立する。すなわち、NTKは学習中ほぼ一定となる。

この結果により、ダイナミクスは線形化される:

\[\frac{d}{dt} f_t = - K (f_t - y)\]

解の構造と指数収束

この線形方程式の解は

\[f_t = y + e^{-Kt}(f_0 - y)\]

であり、固有値分解 $K = U \Lambda U^\top$ を用いると、各モードごとに指数的収束

\[e^{-\lambda_i t}\]

が起こる。したがって、固有値の大きい方向ほど速く学習される。

線形化の厳密化

NTK理論の核心は、以下の線形近似が全時間にわたって有効となる点にある:

\[f(x;\theta_t) \approx f(x;\theta_0) + \nabla_\theta f(x;\theta_0)^\top (\theta_t - \theta_0)\]

無限幅極限では、パラメータの変化が $O(n^{-1/2})$ に抑えられるため、高次項が消失し、この近似が厳密化される。

NTKの数学的性質

(1)正定値性
NTKは内積として定義されるため、任意の有限点集合に対してカーネル行列は半正定値である:

\[\sum_{i,j} c_i c_j K(x_i,x_j) \geq 0\]

(2)再生核ヒルベルト空間
NTKに対応するRKHS $\mathcal{H}_K$ が存在し、学習結果はこの空間内の関数として表現される。

(3)不変性(lazy training)
無限幅ではカーネルが固定されるため、特徴表現は学習によってほとんど変化しない。このレジームは「lazy training」と呼ばれる。

(4)スペクトル構造
カーネルの固有値分解により、学習ダイナミクスは各固有モードに分解される。これはスペクトルバイアス(低周波優先学習)と関係する。

(5)カーネル回帰との同値性
最終的な解は

\[f(x) = K(x,X) K(X,X)^{-1} y\]

となり、カーネルリッジ回帰(正則化なし極限)と一致する。

NNGPとの関係

NNGPは「初期関数分布」を与えるのに対し、NTKは「学習ダイナミクス」を与える。両者は同じ無限幅極限から導かれるが、

  • NNGP:確率的(ガウス過程)
  • NTK:決定論的(線形ダイナミクス)

という対照的な役割を持つ。

有限幅における拡張

有限幅ではNTKは時間とともに変化し、

\[K_{\theta_t} \neq K_{\theta_0}\]

となる。このとき線形化は破れ、特徴学習が生じる。この非線形性こそが実際の深層学習の性能向上に寄与すると考えられている。

まとめ

NTKは、勾配降下法における関数の時間発展を厳密に追跡することで導出されるカーネルであり、無限幅極限において学習ダイナミクスを線形な微分方程式として記述することを可能にする。このとき学習はカーネル法と等価となり、ニューラルネットワークは関数空間上の線形モデルとして振る舞う。この理論は、深層学習の一側面を厳密に理解するための基盤を与えると同時に、有限幅における非線形性との対比を通じて、その本質的な特徴を浮き彫りにするものである。

Mathematics is the language with which God has written the universe.





















数理統計学 機械学習