谷歌重磅開源Neural Tangents:5行代碼打造無限寬神經網絡模型

魚羊 假裝發自 凹非寺
量子位 報道 | 公眾號 QbitAI

只要網絡足夠寬,深度學習動態就能大大簡化,並且更易於理解。

最近的許多研究結果表明,無限寬度的DNN會收斂成一類更為簡單的模型,稱為高斯過程(Gaussian processes)。

於是,複雜的現象可以被歸結為簡單的線性代數方程,以瞭解AI到底是怎樣工作的。

谷歌重磅開源Neural Tangents:5行代碼打造無限寬神經網絡模型

所謂的無限寬度(infinite width),指的是完全連接層中的隱藏單元數,或卷積層中的通道數量有無窮多。

但是,問題來了:推導有限網絡的無限寬度限制需要大量的數學知識,並且必須針對不同研究的體系結構分別進行計算。對工程技術水平的要求也很高。

谷歌最新開源的 Neural Tangents,旨在解決這個問題,讓研究人員能夠輕鬆建立、訓練無限寬神經網絡。

甚至只需要5行代碼,就能夠打造一個無限寬神經網絡模型。

這一研究成果已經中了ICLR 2020。戳進文末Colab鏈接,即可在線試玩。

開箱即用,5行代碼打造無限寬神經網絡模型

Neural Tangents 是一個高級神經網絡 API,可用於指定複雜、分層的神經網絡,在 CPU/GPU/TPU 上開箱即用。

該庫用 JAX編寫,既可以構建有限寬度神經網絡,亦可輕鬆創建和訓練無限寬度神經網絡。

有什麼用呢?舉個例子,你需要訓練一個完全連接神經網絡。通常,神經網絡是隨機初始化的,然後採用梯度下降進行訓練。

研究人員通過對一組神經網絡中不同成員的預測取均值,來提升模型的性能。另外,每個成員預測中的方差可以用來估計不確定性。

如此一來,就需要大量的計算預算。

但當神經網絡變得無限寬時,網絡集合就可以用高斯過程來描述,其均值和方差可以在整個訓練過程中進行計算。

而使用 Neural Tangents ,僅需5行代碼,就能完成對無限寬網絡集合的構造和訓練。

<code>from neural_tangents import predict, stax

init_fn, apply_fn, kernel_fn = stax.serial(
    stax.Dense(2048, W_std=1.5, b_std=0.05), stax.Erf(),
    stax.Dense(2048, W_std=1.5, b_std=0.05), stax.Erf(),
    stax.Dense(1, W_std=1.5, b_std=0.05))

y_mean, y_var = predict.gp_inference(kernel_fn, x_train, y_train, x_test, ‘ntk’, diag_reg=1e-4, compute_cov=True)/<code>


谷歌重磅開源Neural Tangents:5行代碼打造無限寬神經網絡模型

上圖中,左圖為訓練過程中輸出(f)隨輸入數據(x)的變化;右圖為訓練過程中的不確定性訓練、測試損失。

將有限神經網絡的集合訓練和相同體系結構的無限寬度神經網絡集合進行比較,研究人員發現,使用無限寬模型的精確推理,與使用梯度下降訓練整體模型的結果之間,具有良好的一致性。

這說明了無限寬神經網絡捕捉訓練動態的能力。

不僅如此,常規神經網絡可以解決的問題,Neural Tangents 構建的網絡亦不在話下。

研究人員在 CIFAR-10 數據集的圖像識別任務上比較了 3 種不同架構的無限寬神經網絡。

谷歌重磅開源Neural Tangents:5行代碼打造無限寬神經網絡模型

可以看到,無限寬網絡模擬有限神經網絡,遵循相似的性能層次結構,其全連接網絡的性能比卷積網絡差,而卷積網絡的性能又比寬殘餘網絡差。

但是,與常規訓練不同,這些模型的學習動力在封閉形式下是易於控制的,也就是說,可以用前所未有的視角去觀察其行為。

對於深入理解機器學習機制來說,該研究也提供了一種新思路。谷歌表示,這將有助於“打開機器學習的黑匣子”。

傳送門

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

谷歌博客:https://ai.googleblog.com/2020/03/fast-and-easy-infinitely-wide-networks.html

GitHub地址:https://github.com/google/neural-tangents

Colab地址:https://colab.research.google.com/github/google/neural-tangents/blob/master/notebooks/neural_tangents_cookbook.ipynb

— 完 —

量子位 QbitAI · 頭條號簽約

關注我們,第一時間獲知前沿科技動態


分享到:


相關文章: