5行代碼打造無限寬神經(jīng)網(wǎng)絡(luò)模型

只要網(wǎng)絡(luò)足夠?qū)挘?a href="http://www.wenjunhu.com/v/tag/448/" target="_blank">深度學(xué)習(xí)動態(tài)就能大大簡化，并且更易于理解。

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

于是，復(fù)雜的現(xiàn)象可以被歸結(jié)為簡單的線性代數(shù)方程，以了解AI到底是怎樣工作的。

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

但是，問題來了：推導(dǎo)有限網(wǎng)絡(luò)的無限寬度限制需要大量的數(shù)學(xué)知識，并且必須針對不同研究的體系結(jié)構(gòu)分別進(jìn)行計算。對工程技術(shù)水平的要求也很高。

谷歌最新開源的Neural Tangents，旨在解決這個問題，讓研究人員能夠輕松建立、訓(xùn)練無限寬神經(jīng)網(wǎng)絡(luò)。

甚至只需要5行代碼，就能夠打造一個無限寬神經(jīng)網(wǎng)絡(luò)模型。

這一研究成果已經(jīng)中了ICLR 2020。戳進(jìn)文末Colab鏈接，即可在線試玩。

開箱即用，5行代碼打造無限寬神經(jīng)網(wǎng)絡(luò)模型

Neural Tangents 是一個高級神經(jīng)網(wǎng)絡(luò) API，可用于指定復(fù)雜、分層的神經(jīng)網(wǎng)絡(luò)，在 CPU/GPU/TPU 上開箱即用。

該庫用 JAX編寫，既可以構(gòu)建有限寬度神經(jīng)網(wǎng)絡(luò)，亦可輕松創(chuàng)建和訓(xùn)練無限寬度神經(jīng)網(wǎng)絡(luò)。

有什么用呢？舉個例子，你需要訓(xùn)練一個完全連接神經(jīng)網(wǎng)絡(luò)。通常，神經(jīng)網(wǎng)絡(luò)是隨機(jī)初始化的，然后采用梯度下降進(jìn)行訓(xùn)練。

研究人員通過對一組神經(jīng)網(wǎng)絡(luò)中不同成員的預(yù)測取均值，來提升模型的性能。另外，每個成員預(yù)測中的方差可以用來估計不確定性。

如此一來，就需要大量的計算預(yù)算。

但當(dāng)神經(jīng)網(wǎng)絡(luò)變得無限寬時，網(wǎng)絡(luò)集合就可以用高斯過程來描述，其均值和方差可以在整個訓(xùn)練過程中進(jìn)行計算。

而使用 Neural Tangents ，僅需5行代碼，就能完成對無限寬網(wǎng)絡(luò)集合的構(gòu)造和訓(xùn)練。

from neural_tangents import predict， staxinit_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）

上圖中，左圖為訓(xùn)練過程中輸出（f）隨輸入數(shù)據(jù)（x）的變化；右圖為訓(xùn)練過程中的不確定性訓(xùn)練、測試損失。

將有限神經(jīng)網(wǎng)絡(luò)的集合訓(xùn)練和相同體系結(jié)構(gòu)的無限寬度神經(jīng)網(wǎng)絡(luò)集合進(jìn)行比較，研究人員發(fā)現(xiàn)，使用無限寬模型的精確推理，與使用梯度下降訓(xùn)練整體模型的結(jié)果之間，具有良好的一致性。

這說明了無限寬神經(jīng)網(wǎng)絡(luò)捕捉訓(xùn)練動態(tài)的能力。

不僅如此，常規(guī)神經(jīng)網(wǎng)絡(luò)可以解決的問題，Neural Tangents 構(gòu)建的網(wǎng)絡(luò)亦不在話下。

研究人員在 CIFAR-10 數(shù)據(jù)集的圖像識別任務(wù)上比較了 3 種不同架構(gòu)的無限寬神經(jīng)網(wǎng)絡(luò)。

可以看到，無限寬網(wǎng)絡(luò)模擬有限神經(jīng)網(wǎng)絡(luò)，遵循相似的性能層次結(jié)構(gòu)，其全連接網(wǎng)絡(luò)的性能比卷積網(wǎng)絡(luò)差，而卷積網(wǎng)絡(luò)的性能又比寬殘余網(wǎng)絡(luò)差。

但是，與常規(guī)訓(xùn)練不同，這些模型的學(xué)習(xí)動力在封閉形式下是易于控制的，也就是說，可以用前所未有的視角去觀察其行為。

對于深入理解機(jī)器學(xué)習(xí)機(jī)制來說，該研究也提供了一種新思路。谷歌表示，這將有助于“打開機(jī)器學(xué)習(xí)的黑匣子”。