Neural Tangents:
Fast and Easy Infinite Neural Networks in Python
Ryuichi Kanoh (Mobility Technologies*DeNAより出向)
Twitter
ICLR2020オンライン読み会
無限の幅を持つニューラルネットワークを実装する話

Introduction
2

Over-parameterization
近年のニューラルネットワークは、(パラメータ数)>>(学習データ数)
• CIFAR10: 5万枚の訓練画像
• ResNet152: 約6000万個のパラメータ
• 正則化やData Augmentationなしでも、ある程度うまくいく
疑問
(パラメータ数) >> (学習データ数)の学習が、なぜうまくいくのか？
3

経験的に得られている事実
局所解にはまらずに学習は大域収束する
パラメータ数が増え続けても汎化する
論文リンク (ICLR2020)
4

(興味を持ってもらうための)
関連実験紹介 [1]
学習データを増やすほど性能が劣化するレンジが存在
• 二重降下の山の位置がデータ数に依存することに起因
• あなたもこの罠にはまっているかも？
論文リンク (ICLR2020)
5

(興味を持ってもらうための)
関連実験紹介 [2]
アンサンブルに適切なパラメータ数は、二重効果の山を少し超えたところ
• シンプルな定式化ではこの効果は現れず、よくわかっていない 論文リンク (ICML2020)
論文リンク
増やしすぎると劣化
6

課題提起:
無限への興味
無限にパラメータを増やすと、何が起こるか気になりませんか？
• 有限の議論が無限に拡張できるか(逆も然り)は理論研究にとって重要
• 単純に、機械学習モデルとして強力なものになりうる
?
7

無限幅ネットワークに対するアプローチ
Neural Tangent Kernel [定式化]
パラメータ(w)空間での学習の定式化
＊勾配法 (w: パラメータ, η: 学習率, L: 誤差関数)
式変形
η→0
＊誤差関数は二乗誤差とする
関数(y)空間での学習の定式化
Neural Tangent Kernel (NTK)
8

無限幅ネットワークに対するアプローチ
Neural Tangent Kernel [不変性]
隠れ層の幅が広いと、ひとつひとつの重みの変化が微小でも推論結果は大きく変化する
• 実験的/理論的に、幅が広いと重みと は変化しなくなることが示されてきた
論文リンク (NeurIPS 2019)
(証明が気になる方はこちら)
幅と共にNTKの変化が小さく
9

無限幅ネットワークに対するアプローチ
Neural Tangent Kernel [解の獲得]
が時間依存しないとすると、上式は単純な微分方程式となり、解を解析的に得られる
NTK
幅が無限の場合は時間依存しない
10

NTKの活用にあたっての障壁
の計算方法は、モデルの構成ごとに変わる
• 例えば、L層の(全結合層+活性化)のモデルは、以下のように定式化できる
• これらをパラメータで微分し、積をとることでNTKは計算できる
f: pre-activation
g: post-activation
W: NNの重み
σ: 活性化関数
cσ: 規格化定数
dh: h層目の幅
11

NTKの活用にあたっての障壁
の計算方法は、モデルの構成ごとに変わる
• 例えば、L層の(全結合層+活性化)のモデルは、以下のように定式化できる
• これらをパラメータで微分し、積をとることでNTKは計算できる
f: pre-activation
g: post-activation
W: NNの重み
σ: 活性化関数
cσ: 規格化定数
dh: h層目の幅
• 畳み込み層が入ったら？
• Pooling層が入ったら？
• グラフニューラルネットワークは？
• 活性化関数が変わったら？
• Skip Connectionを導入したら？
• …
• 初見での計算は厳しいのでキャッチアップが必要
• 具体的な計算が気になる方はこちら
• Dense (NeurIPS 2018)
• CNN, Pooling (NeurIPS 2019)
• GNN (NeurIPS 2019)
• などなど
12

NTKの活用にあたっての障壁
の計算方法は、モデルの構成ごとに変わる
• 例えば、L層の(全結合層+活性化)のモデルは、以下のように定式化できる
• これらをパラメータで微分し、積をとることでNTKは計算できる
f: pre-activation
g: post-activation
W: NNの重み
σ: 活性化関数
cσ: 規格化定数
dh: h層目の幅
• 畳み込み層が入ったら？
• Pooling層が入ったら？
• グラフニューラルネットワークは？
• 活性化関数が変わったら？
• Skip Connectionを導入したら？
• …
• 初見での計算は厳しいのでキャッチアップが必要
• 具体的な計算が気になる方はこちら
• Dense (NeurIPS 2018)
• CNN, Pooling (NeurIPS 2019)
• GNN (NeurIPS 2019)
• などなど
13
参入障壁が高い！
計算ミスや実装ミスが怖い…
使われる数学も難しい…

ここまでのまとめ
無限の幅を持つモデルを調べることは面白そう
• 有限の議論が無限まで拡張できるかどうか(逆も然り)は、理論的に事柄を示す際に重要な観点
• 単純に、機械学習モデルとして強力なものになりうる
NTKなど、無限幅のモデルを扱う理論も整備されてきている
• それらを用いれば、理論解析と数値実験を交えながら議論が展開できそう
• 中心極限定理や大数の法則が使えるので、かえって有限幅より議論しやすいこともある
モデルを変えるたびに複雑な定式化を行ってNTKを実装するのは大変
• 先行研究がいくつも存在しているが、実装が各論文でバラバラで読み解くのや横展開が大変
• CUDA, numpy, pytorch…
• NTKの計算は重いので、hardware-friendlyな形での実装を考える必要もある
14

How to implement
infinite neural network
15

neural-tangents
neural-tangents (ICLR2020)というライブラリを使用すると簡単。 (pipで入る)
• JAXというGoogleのライブラリをラップしている
• JAXは、numpyに自動微分とJITがくっついたようなもの。GPUやTPUでも動く
• numpy-likeな記法で、低レベルな部分も自分で実装しやすいのが特徴らしい
16

実装例 (NTKの獲得)
JAXの中のstaxというモジュールをラップする形で使用されることが多い
• モデルさえ定義すれば、モデル構造ごとに変わる面倒な設計は、全部ライブラリが内包してくれる
17
無限の幅を持つネットワークのカーネル関数
(Denseで定義している512というパラメータはこの関数に影響しない)
対応するNTKが格納された配列

実装例 (NTKを使用した推論)
カーネル関数とデータを渡せば、推論結果を得ることができる
18

ドキュメント類
初見でも、きちんと情報やハンズオンが整備されているので、使いやすい
ドキュメント
• https://neural-tangents.readthedocs.io/en/latest/#
Gitリポジトリ
• https://github.com/google/neural-tangents
Colab Notebook
• https://colab.research.google.com/github/google/neural-
tangents/blob/master/notebooks/neural_tangents_cookbook.ipynb
19

Experiments
20

有限の幅を持つモデルとの比較
Wide-ResNetなど幅の広いモデルの挙動は、かなり近い振る舞いが模擬できている
• tが大きくなるとズレ始めることにも注意
21
論文リンク (NeurIPS 2019)

機械学習モデルとしての性能
(小スケール)
UCIなどの小さなデータセット(N~5,000)では、既存のモデルを上回る性能が出ている
• 小さなデータを扱うKaggleコンペティションなどで活躍する未来もあり得るかも？
22
論文リンク (ICLR2020)

機械学習モデルとしての性能
(中スケール)
CIFAR10を分類させてみると、深いモデルでは有限幅のモデルの方が性能が良い
• この性能差が意味するところは何なのか？NTKの変化が重要? (論文リンク)
• 深さに対しての依存性も異なりそう
• ギャップを埋めていければ、理論と実践の壁は薄れていき、進化の方向性のヒントを得られるかもしれない
23
論文リンク (NeurIPS 2019)

処理時間
24
ICLR2020プレゼン資料より引用
ImageNet規模の大スケールの対象は厳しい
• 計算の高速化を考えることにも、大いに価値がある

並列処理性能
ほぼ並列処理数に対して線形に処理時間が減っていっていく
25

Summary
26

まとめ
Over-parameterizationについて考えるのは面白い
• 汎化性能、最適化の性質についてなど、不思議なことがたくさん
極限である無限幅のモデルを扱う手法が幾つも出てきている
• これらの活用は理論、実践共に今まさに盛り上がっているところ
• 一方、必要な手法は多岐にわたり、理解や実装が大変
neural-tangentsを使用すると、簡単に無限幅のモデルの実験ができる
• NTKは有限幅モデルでも扱うことはできるので、有限幅モデルを研究するためのツールにもなる
• NTKを利用して二重降下を説明する研究例
27