「ディープラーニングにおける速度限界」の論文を読む

Speed Limits for Deep Learning
https://arxiv.org/abs/2307.14653

という論文でニューラルネットの学習過程における熱力学的速度限界の研究がされています。
先を越されてしまった感があるので紹介します。本文10ページくらいで解析式とテストデータでの興味深い結果が書かれています。

熱力学的速度限界とは非平衡な物理系において分布の変化の大きさL、エントロピー生成$σ\sigma$に対して変化に要した時間$τ\tau$が

$τ≥L22Λσ\tau \geq \frac{L^2}{2\Lambda \sigma}$

というような形で制限される法則です。

詳しくは
https://zenn.dev/xiangze/articles/68e6921f035003

ここでは機械学習の過程でパラメーター$θt\theta_t$が時間$tt$に対して

$dθt=η∇θV(θ;D)dt+1ηβ−1dBtd\theta_t=\eta\nabla_\theta V(\theta;D)dt+\sqrt{1\eta\beta^{-1}}dB_t$

(Dは学習データ$η\eta$は学習率)と変化するような状況を考え、

論文ではエントロピー生成(始点、終点分布間のKLダイバージェンスの中身)

$R=<log⁡p(θ0)p(θT)>+<log⁡p(θT∣θ0)q(θ0∣θT)>R=<\log\frac{p(\theta_0)}{p(\theta_T)}>+<\log\frac{p(\theta_T|\theta_0)}{q(\theta_0|\theta_T)}>$

をカノニカル分布$p(θt)p(\theta_t)$に対する自由エネルギー$ZtZ_t$

$p(θt)=e−ηV(θt)/(log⁡Zt)p(\theta_t)=e^{-\eta V(\theta_t)}/(\log Z_t)$

に対して

$β−1R=β−1log⁡Z∞−β−1log⁡Z0+<L(θ0)>\beta^{-1}R=\beta^{-1}\log Z_{\infty}-\beta^{-1}\log Z_0 +<L(\theta_0)>$

(LはポテンシャルV+正則化項)をとしています。これは有限時間Tの遷移の場合にさらに

$β−1RT=∫0T<∣∇θV∣2>−2β−1<ΔθV>+β−2<∣∇θlog⁡p∣2>dt\beta^{-1}R_T=\int_0^T <|\nabla_\theta V|^2>-2\beta^{-1}<\Delta_\theta V> +\beta^{-2}<|\nabla_\theta \log p|^2>dt$

と逆温度βに関して勾配とヘッシアンを係数として展開されるそうです。NTKの場合は$β−1=0\beta^{-1}=0$(温度無限大)として書かれるのでより簡単に

$β−1RT=<L(θ0)−L(θT)>\beta^{-1}R_T=<L(\theta_0)-L(\theta_T)>$

と書かれるそうです。

最適輸送との関係

「非平衡統計力学」などの本や論文ではコーシー・シュワルツの不等式などから熱力学的速度限界を導出していましたが、
最適輸送理論の考えを使ってエントロピー生成率と速度限界の関係を理解することもできます。

$L2L^2$-Wasserstein距離($W2(pA,pB):=min⁡π∈Π(pA,pB)∫RdxRd∣x−y∣2π(x,y)dxdyW_2(p^A,p^B):=\sqrt{\min_{\pi \in \Pi(p^A,p^B)}\int_{R^d x R^d }|x-y|^2\pi(x,y) dxdy}$)

を拡散係数Dを使ってBenamou-Brenier式

$W2(p0,pT)=min⁡∫0TDRtdtW_2(p_0,p_T)=\min \int_0^T\sqrt{DR_t} dt$

と表すことでエントロピー生成Rの下限が

$T≥TSL:=W2(p0,pT)β−1RT \ge T_{SL}:=\frac{W_2(p_0,p_T)}{\beta^{-1}R}$

という形で表せるそうです。
Benamou-Brenier式に関しては以下の論文にも書いてあります。
https://journals.aps.org/prx/abstract/10.1103/PhysRevX.13.011013

https://link.springer.com/article/10.1007/s41884-023-00102-3

解析解

線形モデルとネットワークの幅が無限大の極限でガウスカーネルを用いて書かれるNTK(Neural Tangent Kernel)の場合に対して解析解を導いています。各導出の詳細はAppendixにあります。線形モデル

$y^(x)=θTxL(θ;D={X,y})=12∣y−XTθ∣2\hat{y}(x)=\theta^T x \\ L(\theta;D=\{X,y\})=\frac{1}{2}|y-X^T\theta|^2$

の場合遷移時間$TSLT_{SL}$がパラメーター$α,λ\alpha,\lambda$とMarchenko-Pastor分布$ρ\rho$を用いて

$TSL(β−>∞)→21+αλ∫sdρ(s)T_{SL}(\beta->\infty)\rightarrow2\frac{1+\alpha\lambda}{\int sd\rho(s)}$

と書かれるそうです。興味深いのはサンプル数$n→∞n\rightarrow\infty$の場合は

$TSL→2λαT_{SL}\rightarrow2\lambda\alpha$

と一定の値になるのに入力次元$d→∞,d>>nd\rightarrow\infty, d>>n$の場合は

$TSL→0T_{SL}\rightarrow0$

となるということです。これはoverparametrazeの状態に相当し、最適解が初期値のすぐ近くにあるという状況を示していると主張しています。

NTKの場合は固有値$λ\lambda$の分布によって書かれまた、学習前後のパラメーター間の距離$lgeo=W2l_{geo}=\sqrt{W_2}$が

$lgeo∝T(α−1+1−δ/α)/2l_{geo}\propto T^{(\alpha^{-1} +1-\delta/\alpha)/2}$

というスケーリング則を満たすことを導出しています。

解釈

学習率と挙動の関係について自明なこと以上が言えるのか論文の内容も含めて個人的にはまだわかっていません。

数値実験

CIFAR-10を5層、128channelのCNN(Myrtle-5)で学習する過程に対して非効率性の比率
$T/TSL(T)T/T_{SL}(T)$と幾何学的非効率性比率$lNTK(T)/lgeo(T)l_{NTK}(T)/l_{geo}(T)$ をepoch,date
sizeに対して取っているのですが、それが速度限界とどう関係するのかいまいち見て取れませんでした。
また重みパラメーターの各成分は学習過程で直線的に変動すると主張しています。
NTKの大きな固有値が最初にすばやく学習されるということも言えるらしいです。

展望

• 数値実験のところをもっとわかりやすく取りたい。
• 熱力学的不確定性も言えるのだろうか
• 言うほどDeepではない。CNNならではの特徴、ネットワークの深さに対する依存はあるのだろうか。
• タスクによる違い、現実的なタスク、特に異なるプロンプトに対応できる画像生成や基盤モデル
• 有限要素の場合の解析的補正は難しそう(kernel-adaptation methodsというのがあるらしい、機械学習か？)

参考

NTKと有限要素の場合の補正について
「勾配ランジュバン動力学　平均場ランジュバン動力学」最後の方で拡散モデルの説明と誤差解析の研究紹介があります
http://ibis.t.u-tokyo.ac.jp/suzuki/lecture/2023/TohokuUniv/ランジュバン.pdf

A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem
最適輸送の話です。
https://www.iap.fr/actualites/laune/2022/TransportOptimal/ark _67375_VQC-XB4DR0Z3-2.pdf

matlabのコード
https://www.numerical-tours.com/matlab/optimaltransp_2_benamou_brenier/