圖靈獎(jiǎng)得主 Geoffrey Hinton 等研究者讓前向梯度學(xué)習(xí)變得實(shí)用了。

我們知道，在人工智能領(lǐng)域里，反向傳播是個(gè)最基本的概念。

反向傳播（Backpropagation，BP）是一種與最優(yōu)化方法（如梯度下降）結(jié)合使用的，用來(lái)訓(xùn)練人工神經(jīng)網(wǎng)絡(luò)的常見方法。該方法計(jì)算對(duì)網(wǎng)絡(luò)中所有權(quán)重計(jì)算損失函數(shù)的梯度。這個(gè)梯度會(huì)反饋給最優(yōu)化方法，用來(lái)更新權(quán)值以最小化損失函數(shù)。

簡(jiǎn)而言之，BP 的核心思路其實(shí)就是負(fù)反饋，我們?cè)噲D用這種方式實(shí)現(xiàn)神經(jīng)網(wǎng)絡(luò)系統(tǒng)面對(duì)給定目標(biāo)的自動(dòng)迭代、校準(zhǔn)。隨著算力、數(shù)據(jù)和更多技術(shù)改進(jìn)的提升，在 AI 領(lǐng)域人們使用反向傳播訓(xùn)練的多層神經(jīng)網(wǎng)絡(luò)在部分任務(wù)上已經(jīng)足以與人類競(jìng)爭(zhēng)。

很多人把這項(xiàng)技術(shù)的發(fā)現(xiàn)歸功于深度學(xué)習(xí)先驅(qū)、2019 年圖靈獎(jiǎng)得主 Geoffrey Hinton，但 Hinton 本人表示，自己的貢獻(xiàn)在于明確提出了反向傳播可以學(xué)習(xí)有趣的內(nèi)部表征，并讓這一想法推廣開來(lái)：「我通過(guò)讓神經(jīng)網(wǎng)絡(luò)學(xué)習(xí)詞向量表征，使之基于之前詞的向量表征預(yù)測(cè)序列中的下一個(gè)詞實(shí)現(xiàn)了這一點(diǎn)。」

其例證在于 Nature 1986 年發(fā)表的論文《Learning representations by back-propagating errors》上。

無(wú)論如何，反向傳播技術(shù)推動(dòng)了現(xiàn)代深度學(xué)習(xí)的發(fā)展。但曾被冠以「反向傳播之父」的 Geoffrey Hinton，近年來(lái)卻經(jīng)常表示自己在構(gòu)思下一代神經(jīng)網(wǎng)絡(luò)，他對(duì)于反向傳播「非常懷疑」，并提出「應(yīng)該拋棄它并重新開始」。

可以說(shuō)自 2017 年起，Hinton 就已開始尋找新的方向。機(jī)器之心先前曾介紹 Hinton 在前向 - 前向網(wǎng)絡(luò)方面的思考（近萬(wàn)人圍觀 Hinton 最新演講：前向 - 前向神經(jīng)網(wǎng)絡(luò)訓(xùn)練算法，論文已公開）。

最近，我們又看到了重要的進(jìn)展。近日，由 Mengye Ren、Simon Kornblith、Renjie Liao、Geoffrey Hinton 完成的論文被工智能頂會(huì) ICLR 2023 接收。

前向梯度學(xué)習(xí)通常用于計(jì)算含有噪聲的方向梯度，是一種符合生物學(xué)機(jī)制、可替代反向傳播的深度神經(jīng)網(wǎng)絡(luò)學(xué)習(xí)方法。然而，當(dāng)要學(xué)習(xí)的參數(shù)量很大時(shí)，標(biāo)準(zhǔn)的前向梯度算法會(huì)出現(xiàn)較大的方差。

基于此，圖靈獎(jiǎng)得主 Geoffrey Hinton 等研究者提出了一系列新架構(gòu)和算法改進(jìn)，使得前向梯度學(xué)習(xí)對(duì)標(biāo)準(zhǔn)深度學(xué)習(xí)基準(zhǔn)測(cè)試任務(wù)具有實(shí)用性。

論文鏈接：https://arxiv.org/abs/2210.03310

GitHub 鏈接：https://github.com/google-research/google-research/tree/master/local_forward_gradient

該研究表明，通過(guò)將擾動(dòng)（perturbation）應(yīng)用于激活而不是權(quán)重，可以顯著減少前向梯度估計(jì)器的方差。研究團(tuán)隊(duì)通過(guò)引入大量局部貪心損失函數(shù)（每個(gè)損失函數(shù)只涉及少量可學(xué)習(xí)參數(shù)）和更適合局部學(xué)習(xí)的新架構(gòu) LocalMixer（受 MLPMixer 啟發(fā)），進(jìn)一步提高了前向梯度的可擴(kuò)展性。該研究提出的方法在 MNIST 和 CIFAR-10 上與反向傳播性能相當(dāng)，并且明顯優(yōu)于之前 ImageNet 上的無(wú)反向傳播算法。

當(dāng)前，大多數(shù)深度神經(jīng)網(wǎng)絡(luò)都使用反向傳播算法（Werbos, 1974; LeCun, 1985; Rumelhart et al., 1986）進(jìn)行訓(xùn)練，該算法通過(guò)從損失函數(shù)向每一層反向傳播誤差信號(hào)來(lái)有效地計(jì)算權(quán)重參數(shù)的梯度。盡管人工神經(jīng)網(wǎng)絡(luò)最初受到生物神經(jīng)元的啟發(fā)，但反向傳播一直被認(rèn)為不符合生物學(xué)機(jī)理，因?yàn)榇竽X不會(huì)形成對(duì)稱的反向連接或執(zhí)行同步計(jì)算。從工程的角度講，反向傳播與大規(guī)模模型的并行性不兼容，并且限制了潛在的硬件設(shè)計(jì)。這些問(wèn)題表明我們需要一種截然不同的深度網(wǎng)絡(luò)學(xué)習(xí)算法。

Hinton 等研究者重新審視了權(quán)重?cái)_動(dòng)的替代方法 —— 活動(dòng)擾動(dòng)（activity perturbation，Le Cun et al., 1988; Widrow & Lehr, 1990; Fiete & Seung, 2006），探索了該方法對(duì)視覺任務(wù)訓(xùn)練的普遍適用性。

該研究表明：活動(dòng)擾動(dòng)能比權(quán)重?cái)_動(dòng)產(chǎn)生方差更低的梯度估計(jì)，并且能夠?yàn)樵撗芯刻岢龅乃惴ㄌ峁┗谶B續(xù)時(shí)間速率（continuous-time rate-based）的解釋。

研究團(tuán)隊(duì)通過(guò)設(shè)計(jì)具有大量局部貪心損失函數(shù)的架構(gòu)，解決了前向梯度學(xué)習(xí)的可擴(kuò)展性問(wèn)題，其中將網(wǎng)絡(luò)隔離為局部模塊，從而減少了每個(gè)損失函數(shù)的可學(xué)習(xí)參數(shù)量。與僅沿深度維度添加局部損失的先前工作不同，該研究發(fā)現(xiàn) patch-wise 和 channel-group-wise 損失函數(shù)也非常關(guān)鍵。最后，受 MLPMixer (Tolstikhin et al., 2021) 的啟發(fā)，該研究設(shè)計(jì)了一個(gè)名為 LocalMixer 的網(wǎng)絡(luò)。LocalMixer 具有線性 token 混合層和分組通道（channel），以更好地與局部學(xué)習(xí)兼容。

該研究在監(jiān)督和自監(jiān)督圖像分類問(wèn)題上評(píng)估了其局部貪婪前向梯度算法。在 MNIST 和 CIFAR-10 上，該研究提出的學(xué)習(xí)算法性與反向傳播性能相當(dāng)，而在 ImageNet 上，其性能明顯優(yōu)于其他使用不對(duì)稱前向和后向權(quán)重的方案。雖然該研究提出的算法在更大規(guī)模的問(wèn)題上還達(dá)不到反向傳播算法的性能，但局部損失函數(shù)設(shè)計(jì)可能是生物學(xué)上合理的學(xué)習(xí)算法，也將成為下一代模型并行計(jì)算的關(guān)鍵因素。

該研究分析了前向梯度估計(jì)器的期望和方差的特性，并將分析重點(diǎn)放在了權(quán)重矩陣的梯度上，具體的理論分析結(jié)果如下表 1 所示，批大小為 N 時(shí)，獨(dú)立擾動(dòng)（independent perturbation）可以將方差減少為 1/N，而共享擾動(dòng)（shared perturbation）具有由平方梯度范數(shù)支配的常數(shù)方差項(xiàng)。然而，當(dāng)執(zhí)行獨(dú)立的權(quán)重?cái)_動(dòng)時(shí)，矩陣乘法不能進(jìn)行批處理，因?yàn)槊總€(gè)樣本的激活向量都要與不同的權(quán)重矩陣相乘。相比之下，獨(dú)立的活動(dòng)擾動(dòng)算法允許批量矩陣乘法。

與權(quán)重?cái)_動(dòng)相比，活動(dòng)擾動(dòng)的方差更小，因?yàn)閿_動(dòng)元素的數(shù)量是輸出單元的數(shù)量，而不是整個(gè)權(quán)重矩陣的大小。活動(dòng)擾動(dòng)的唯一缺點(diǎn)是存儲(chǔ)中間激活需要一定量的內(nèi)存。

此外，該研究發(fā)現(xiàn)在具有 ReLU 激活的網(wǎng)絡(luò)中，可以利用 ReLU 稀疏性來(lái)進(jìn)一步減少方差，因?yàn)槲醇せ畹膯卧荻葹榱悖虼瞬粦?yīng)該擾動(dòng)這些單元。

用局部損失函數(shù)進(jìn)行擴(kuò)展

由于擾動(dòng)學(xué)習(xí)可能會(huì)遭受「維數(shù)災(zāi)難」：方差隨著擾動(dòng)維數(shù)的增加而增加，并且深度網(wǎng)絡(luò)中通常有數(shù)百萬(wàn)個(gè)參數(shù)同時(shí)發(fā)生變化。限制可學(xué)習(xí)維度數(shù)量的一種方法是將網(wǎng)絡(luò)劃分為子模塊，每個(gè)子模塊都有一個(gè)單獨(dú)的損失函數(shù)。因此，該研究通過(guò)增加局部損失函數(shù)的數(shù)量來(lái)抑制方差，具體包括：

1）Blockwise 損失。首先，該研究將網(wǎng)絡(luò)深度劃分為多個(gè)模塊。每個(gè)模塊由幾個(gè)層組成，在每個(gè)模塊的末端計(jì)算一個(gè)損失函數(shù)，該損失函數(shù)用于更新該模塊中的參數(shù)。這種方法相當(dāng)于在模塊之間添加了一個(gè)「停止梯度（stop gradient）」算子，Belilovsky et al. (2019) 和 L?we et al. (2019) 曾探索過(guò)這種局部貪心損失函數(shù)。

2）patchwise 損失。圖像等感官輸入信號(hào)具有空間維度。該研究沿著這些空間維度，逐塊應(yīng)用單獨(dú)的損失。在 Vision Transformer 架構(gòu)中（Vaswani et al., 2017; Dosovitskiy et al., 2021），每個(gè)空間標(biāo)記代表圖像中的一個(gè) patch。在現(xiàn)代深度網(wǎng)絡(luò)中，每個(gè)空間位置的參數(shù)通常是共享的，以提高數(shù)據(jù)效率并降低內(nèi)存帶寬利用率。雖然簡(jiǎn)單的權(quán)重共享在生物學(xué)上是不合理的，但該研究仍然在這項(xiàng)工作中考慮共享權(quán)重。通過(guò)在 patch 之間添加知識(shí)蒸餾 (Hinton et al., 2015) 損失來(lái)模擬權(quán)重共享的效果是可能的。

3) Groupwise 損失。最后，該研究轉(zhuǎn)向通道（channel）維度。為了創(chuàng)建多個(gè)損失，該研究將通道分成多個(gè)組，每個(gè)組都附加到一個(gè)損失函數(shù)（Patel et al., 2022）。為了防止組之間相互通信，通道僅連接同一組內(nèi)的其他通道。

具有局部損失的 LocalMixer 殘差塊。

特征聚合器

簡(jiǎn)單地將損失分別應(yīng)用于空間和通道維度，會(huì)帶來(lái)次優(yōu)性能，因?yàn)槊總€(gè)維度僅包含局部信息。對(duì)于分類等標(biāo)準(zhǔn)任務(wù)的損失，模型需要輸入的全局視圖來(lái)做出決策。標(biāo)準(zhǔn)架構(gòu)通過(guò)在最終分類層之前執(zhí)行全局平均池化層，來(lái)獲得此全局視圖。因此，該研究探索了在局部損失函數(shù)之前聚合來(lái)自其他組和空間塊的信息的策略。

特征聚合器設(shè)計(jì)。

實(shí)現(xiàn)

網(wǎng)絡(luò)架構(gòu)：該研究提出了更適合局部學(xué)習(xí)的 LocalMixer 架構(gòu)。它的靈感來(lái)自 MLPMixer (Tolstikhin et al., 2021)，它由全連接網(wǎng)絡(luò)和殘差塊組成。該研究利用全連接網(wǎng)絡(luò)，使每個(gè)空間塊在不干擾其他 patch 的情況下執(zhí)行計(jì)算，這更符合局部學(xué)習(xí)目標(biāo)。圖 1 顯示了高級(jí)架構(gòu)，圖 2 顯示了一個(gè)殘差塊的詳細(xì)圖。

歸一化。在跨不同張量維度的神經(jīng)網(wǎng)絡(luò)中執(zhí)行歸一化的方法有很多種（Krizhevsky et al., 2012; Ioffe & Szegedy, 2015; Ba et al., 2016; Ren et al., 2017; Wu & He, 2018）。該研究選擇了層歸一化的局部變體，它在每個(gè)局部空間特征塊內(nèi)進(jìn)行歸一化（Ren et al., 2017）。對(duì)于分組的線性層，每組單獨(dú)進(jìn)行歸一化（Wu & He, 2018）。

該研究通過(guò)實(shí)驗(yàn)發(fā)現(xiàn)這種局部歸一化在對(duì)比學(xué)習(xí)中表現(xiàn)更好，并且與監(jiān)督學(xué)習(xí)中的層歸一化大致相同。局部歸一化在生物學(xué)上也更合理，因?yàn)樗粓?zhí)行全局通信。

通常，歸一化層放置在線性層之后。在 MLPMixer（Tolstikhin et al., 2021）中，層歸一化被放置在每個(gè)殘差塊的開頭。該研究發(fā)現(xiàn)最好在每個(gè)線性層之前和之后放置歸一化，如圖 2 所示。實(shí)驗(yàn)結(jié)果表明，這種設(shè)計(jì)選擇對(duì)反向傳播沒(méi)有太大影響，但它允許前向梯度學(xué)習(xí)更快地學(xué)習(xí)并實(shí)現(xiàn)更低的訓(xùn)練錯(cuò)誤。

有效實(shí)施復(fù)制損失。由于特征聚合和復(fù)制損失的設(shè)計(jì)，組（groups）的簡(jiǎn)單實(shí)現(xiàn)在內(nèi)存消耗和計(jì)算方面可能非常低效。然而，每個(gè)空間組實(shí)際上計(jì)算相同的聚合特征和損失函數(shù)。這意味著在執(zhí)行反向傳播和正向梯度時(shí)，可以跨損失函數(shù)共享大部分計(jì)算。該研究實(shí)現(xiàn)了自定義 JAX JVP/VJP 函數(shù)（Bradbury et al., 2018）并觀察到顯著的內(nèi)存節(jié)省和復(fù)制損失的計(jì)算速度提升，否則這在現(xiàn)代硬件上運(yùn)行是不可行的，結(jié)果如下圖所示。

復(fù)制損失的簡(jiǎn)單和融合實(shí)現(xiàn)時(shí)，內(nèi)存和計(jì)算使用情況。

實(shí)驗(yàn)

該研究將提出的算法與包括反向傳播、反饋對(duì)齊和其他前向梯度全局變體在內(nèi)的替代方案進(jìn)行比較。反向傳播是生物學(xué)上難以置信的神諭，因?yàn)樗?jì)算的是真實(shí)梯度，而該方法計(jì)算的是噪聲梯度。反饋對(duì)齊通過(guò)使用一組隨機(jī)向后權(quán)重來(lái)計(jì)算近似梯度。

各項(xiàng)實(shí)驗(yàn)的結(jié)果如下：

圖像分類自監(jiān)督學(xué)習(xí)。

有線性 readout 的自監(jiān)督對(duì)比學(xué)習(xí)。

在不同位置添加局部損失對(duì)前向梯度性能的影響。

使用不同數(shù)量組時(shí)，在 CIFAR-10 上訓(xùn)練時(shí) M/8/* 的錯(cuò)誤率。

總結(jié)

人們通常認(rèn)為基于擾動(dòng)的學(xué)習(xí)無(wú)法擴(kuò)展到大型深度網(wǎng)絡(luò)。該研究表明，這在某種程度上是正確的，因?yàn)樘荻裙烙?jì)方差隨著擾動(dòng)隱藏維度的數(shù)量而增長(zhǎng)，并且對(duì)共享權(quán)重?cái)_動(dòng)而言甚至更遭。

但樂(lè)觀的是，該研究表明大量的局部貪心損失可以幫助更好地推進(jìn)梯度學(xué)習(xí)規(guī)模，探索了 blockwise、patchwise 以及 groupwise 以及這三者組合的局部損失，在一個(gè)較大的網(wǎng)絡(luò)中總共有四分之一的損失，表現(xiàn)最好。局部活動(dòng)擾動(dòng)前向梯度在更大的網(wǎng)絡(luò)上比以前的無(wú)反向傳播算法表現(xiàn)更好。局部損失的想法為不同的損失設(shè)計(jì)開辟了機(jī)會(huì)，并闡明了如何在大腦和替代計(jì)算設(shè)備中尋找生物學(xué)上合理的學(xué)習(xí)算法。

編輯：黃飛

?