谷歌開源Neural Tangents:簡單快速訓練無限寬度神經網絡

選自Google AI Blog

機器之心編譯

參與:Panda

無限寬的神經網絡可以訓練了?神經正切核果然可以發揮作用了。

近日,谷歌一項研究成果在社區引發了關注,該研究宣稱可通過神經正切核(Neural Tangent Kernel)使用貝葉斯推理或梯度下降分析式地訓練無限寬度的神經網絡。使用谷歌開源的軟件庫 Neural Tangents,這個過程不僅簡單且快速,而且效果非常好,甚至只需 5 行代碼就能一步到位地構建並訓練這種無限寬度網絡的集成模型!該論文已被 ICLR 2020 接收為 Spotlight 論文,相關鏈接見文末。

在自然語言處理、會話智能體和連接組學等許多領域,深度學習都已取得了廣泛的成功,機器學習領域的研究圖景也已經發生了變革。不過,仍還有一些有趣而又重要的問題有待解答,比如:為什麼即使在過度參數化時,深度神經網絡(DNN)也能取得非常好的泛化能力?深度網絡的架構、訓練和性能之間有何關係?如何提取出深度學習模型中的顯著特徵?

我們何以在近些年裡取得了如此的進步?一大關鍵理論見解是:增加 DNN 的寬度能使 DNN 的行為更有規律可循,也就人更容易理解它們。近來的許多研究已經表明,寬度可無限擴增的 DNN 可以收斂成另一類更為簡單的名為「高斯過程(Gaussian processes)」的模型。因此,貝葉斯推理與卷積神經網絡的梯度下降情況等複雜現象便可約簡為簡單的線性代數方程。從這些無限寬度網絡所得到的見解通常也適用於它們對應的有限版本。因此,無限寬度網絡可用作研究深度學習的透鏡,而且它們本身也可用作有用的模型。

谷歌開源Neural Tangents:簡單快速訓練無限寬度神經網絡

左圖:深度神經網絡變得無限寬時推導輸入/輸出映射關係的示意圖。右圖:隨著神經網絡寬度的增長,可以看到在網絡的不同隨機初始化上得到的輸出分佈逐漸變成高斯分佈。

不幸的是,推導有限網絡的無限寬度極限需要大量數學專業知識,而且必須針對每種不同架構專門研究。一旦推導得到了無限寬度模型,如果要得到一種高效且可擴展的實現,還進一步需要大量工程開發專業技能。總的來說,將有限寬度模型擴展成對應的無限寬度網絡的這個過程可能耗費數月時間,而且這個研究本身就足以支撐起一篇研究論文。

為了解決這個問題,谷歌開源了一個用 Jax 寫的開源軟件庫 Neural Tangents。谷歌宣稱這個軟件庫可讓無限寬度神經網絡的構建與訓練變得與有限寬度神經網絡一樣簡單。究其核心,Neural Tangents 提供了一個易於使用的神經網絡軟件庫,其可同時為神經網絡構建有限和無限寬度的版本。

舉個實際使用 Neural Tangents 的例子:假設要在某些數據上訓練一個全連接神經網絡,通常的做法是隨機初始化一個神經網絡,然後使用梯度下降來訓練它。初始化並訓練許多這樣的神經網絡可以得到一個集成模型。這個集成模型中各個模型成員都會給出一個預測結果,而為了得到最終結果,研究者和實踐者通常的做法是求取這些結果的平均。這種做法通常能取得更好的表現。此外,集成模型中各個模型成員的預測結果還可用於估計不確定度。但訓練集成模型也有缺點,即計算成本的預算很高,因此人們通常會避免這樣操作。但是,如果神經網絡的寬度變得無限大,就可以用一個高斯過程來描述這個集成模型了,而且這個高斯過程的均值和方差都可在訓練過程中計算得到。

谷歌表示,使用 Neural Tangents,只需 5 行代碼就能一步到位地構建並訓練這種無限寬度網絡的集成模型!下面展示了所得到的訓練過程,而且谷歌還發布了一個可以嘗試這一實驗的交互式 Colab 筆記:

https://colab.sandbox.google.com/github/google/neural-tangents/blob/master/notebooks/neural_tangents_cookbook.ipynb

谷歌開源Neural Tangents:簡單快速訓練無限寬度神經網絡

儘管事實上這種無限寬度的集成模型受一個簡單的閉式表達式的控制,但它卻能表現出與有限寬度集成模型的一致性。而且因為這種無限寬度的集成模型是一個高斯過程,所以它自然就能提供閉式的不確定度估計(上圖中填充顏色的區域)。這些不確定度估計與在訓練有限網絡的許多不同副本(虛線)時所觀察到的預測變化範圍是緊密一致的。

上面的例子表明,無限寬度神經網絡有習得並獲取訓練動態的能力。不僅如此,使用 Neural Tangents 構建的網絡適用於任意可應用常規神經網絡解決的問題。谷歌又給出了一個例子。下面比較了使用 CIFAR-10 數據集訓練的三個不同的無限寬度神經網絡架構。Neural Tangents 表現出色:不管是用梯度下降還是用完全貝葉斯的推理,都能以閉式形式評估無限寬殘差網絡等高度複雜的模型所組成的集成模型。而如果用有限寬度模型,這個任務根本無法解決。

谷歌開源Neural Tangents:簡單快速訓練無限寬度神經網絡

可以看到,無限寬度網絡與有限寬度網絡相似,遵循著類似的層次結構:全連接網絡的表現不及卷積網絡,而卷積網絡又趕不上寬殘差網絡。但是,不同於常規的訓練方法,這些模型的學習動態完全可用閉式方法解決,這能讓人以前所未有的方式理解這些動態,進而有助於揭開機器學習背後的奧秘,打開已困擾 AI 領域良久的深度學習黑箱。

想必現在你也想親自探索一番了,以下鏈接中一定有你最感興趣的探索形式:

論文:https://arxiv.org/abs/1912.02803

Colab 筆記教程:https://colab.sandbox.google.com/github/google/neural-tangents/blob/master/notebooks/neural_tangents_cookbook.ipynb

GitHub 代碼庫:https://github.com/google/neural-tangents

參考鏈接:

http://ai.googleblog.com/2020/03/fast-and-easy-infinitely-wide-networks.html


分享到:


相關文章: