打破「反向傳播」壟斷，「正向自動(dòng)微分」也能計(jì)算梯度，且訓(xùn)練時(shí)間減少一半

本文作者：我在思考中

2022-03-17 10:17

導(dǎo)語(yǔ)：實(shí)驗(yàn)證明，在一些問(wèn)題中，正向梯度的計(jì)算時(shí)間是反向傳播的二分之一。?

打破「反向傳播」壟斷，「正向自動(dòng)微分」也能計(jì)算梯度，且訓(xùn)練時(shí)間減少一半

用反向傳播（backpropagation）來(lái)計(jì)算優(yōu)化目標(biāo)函數(shù)的梯度，是當(dāng)前機(jī)器學(xué)習(xí)領(lǐng)域的主流方法。近日，牛津與微軟等機(jī)構(gòu)的多位學(xué)者聯(lián)合提出一種名為「正向梯度」（forward gradient）的自動(dòng)微分模式，可以完全拋棄反向傳播進(jìn)行梯度計(jì)算。實(shí)驗(yàn)證明，在一些問(wèn)題中，正向梯度的計(jì)算時(shí)間是反向傳播的二分之一。

編譯 | 張倩

編輯 | 陳彩嫻

反向傳播和基于梯度的優(yōu)化是近年來(lái)機(jī)器學(xué)習(xí)（ML）取得重大突破的核心技術(shù)。

人們普遍認(rèn)為，機(jī)器學(xué)習(xí)之所以能夠快速發(fā)展，是因?yàn)檠芯空邆兪褂昧说谌娇蚣埽ㄈ鏟yTorch、TensorFlow）來(lái)解析ML代碼。這些框架不僅具有自動(dòng)微分（AD）功能，還為本地代碼提供了基礎(chǔ)的計(jì)算功能。而ML所依賴的這些軟件框架都是圍繞 AD 的反向模式所構(gòu)建的。這主要是因?yàn)樵贛L中，當(dāng)輸入的梯度為海量時(shí)，可以通過(guò)反向模式的單次評(píng)估進(jìn)行精確有效的評(píng)估。

自動(dòng)微分算法分為正向模式和反向模式。但正向模式的特點(diǎn)是只需要對(duì)一個(gè)函數(shù)進(jìn)行一次正向評(píng)估（即沒(méi)有用到任何反向傳播），計(jì)算成本明顯降低。為此，來(lái)自劍橋與微軟等機(jī)構(gòu)的研究者們探索這種模式，展示了僅使用正向自動(dòng)微分也能在一系列機(jī)器學(xué)習(xí)框架上實(shí)現(xiàn)穩(wěn)定的梯度下降。

論文地址：https://arxiv.org/pdf/2202.08587v1.pdf

他們認(rèn)為，正向梯度有利于改變經(jīng)典機(jī)器學(xué)習(xí)訓(xùn)練管道的計(jì)算復(fù)雜性，減少訓(xùn)練的時(shí)間和精力成本，影響機(jī)器學(xué)習(xí)的硬件設(shè)計(jì)，甚至對(duì)大腦中反向傳播的生物學(xué)合理性產(chǎn)生影響。

自動(dòng)微分的兩種模式

首先，我們來(lái)簡(jiǎn)要回顧一下自動(dòng)微分的兩種基本模式。

正向模式

給定一個(gè)函數(shù) f: θ∈R n，v∈R n，正向模式的AD會(huì)計(jì)算 f(θ) 和雅可比向量乘積Jf (θ) v，其中Jf (θ) ∈R m×n是f在θ處評(píng)估的所有偏導(dǎo)數(shù)的雅可比矩陣，v是擾動(dòng)向量。對(duì)于 f : R n → R 的情況，在雅可比向量乘積對(duì)應(yīng)的方向?qū)?shù)用 ?f(θ)- v表示，即在θ處的梯度?f對(duì)方向向量v的映射，代表沿著該方向的變化率。

值得注意的是，正向模式在一次正向運(yùn)行中同時(shí)評(píng)估了函數(shù) f 及其雅可比向量乘積 Jf v。此外，獲得 Jf v 不需要計(jì)算雅可比向量Jf，這一特點(diǎn)被稱為無(wú)矩陣計(jì)算。

反向模式

給定一個(gè)函數(shù) f : R n → R m，數(shù)值 θ∈R n，v∈R m，AD反向模式會(huì)計(jì)算f(θ)和雅可比向量乘積v |Jf (θ)，其中Jf∈R m×n是f在θ處求值的所有偏導(dǎo)數(shù)的雅可比矩陣，v∈R m是一個(gè)鄰接的矢量。對(duì)于f : R n → R和v = 1的情況，反向模式計(jì)算梯度，即f對(duì)所有n個(gè)輸入的偏導(dǎo)數(shù)?f(θ)=h ?f ?θ1,. . . , ?f ?θn i| 。

請(qǐng)注意，v |Jf 是在一次前向-后向評(píng)估中進(jìn)行計(jì)算的，而不需要計(jì)算雅可比Jf 。

運(yùn)行時(shí)間成本

兩種AD模式的運(yùn)行時(shí)間以運(yùn)行正在微分的函數(shù) f 所需時(shí)間的恒定倍數(shù)為界。

反向模式的成本比正向模式高，因?yàn)樗婕暗綌?shù)據(jù)流的反轉(zhuǎn)，而且需要保留正向過(guò)程中所有操作結(jié)果的記錄，因?yàn)樵诮酉聛?lái)的反向過(guò)程中需要這些記錄來(lái)評(píng)估導(dǎo)數(shù)。內(nèi)存和計(jì)算成本特征最終取決于AD系統(tǒng)實(shí)現(xiàn)的功能，如利用稀疏性。

成本可以通過(guò)假設(shè)基本操作的計(jì)算復(fù)雜性來(lái)分析，如存儲(chǔ)、加法、乘法和非線性操作。將評(píng)估原始函數(shù) f 所需的時(shí)間表示設(shè)為 runtime(f)，我們可以將正向和反向模式所需的時(shí)間分別表示為 Rf×runtime(f) 和 Rb×runtime(f)。在實(shí)踐中，Rf 通常在1到3之間，Rb通常在5到10之間，不過(guò)這些結(jié)果都與程序高度相關(guān)。

方法

正向梯度

定義1

給定一個(gè)函數(shù) f : R n → R，他們將「正向梯度」 g : R n → R n 定義為：

其中，θ∈R n 是評(píng)估梯度的關(guān)鍵點(diǎn)，v∈R n 是一個(gè)擾動(dòng)向量，被視為一個(gè)多元隨機(jī)變量v～p(v)，這樣 v 的標(biāo)量分量 vi 是獨(dú)立的，對(duì)所有 i 都有零均值和單位方差，?f(θ)-v∈R 是 f 在在 v 方向上 θ 點(diǎn)的方向?qū)?shù)。

簡(jiǎn)要地談一下這個(gè)定義的由來(lái)。

如前所述，正向模式直接給我們提供了方向?qū)?shù)?f(θ) - v = P i ?f ?θi vi，無(wú)需計(jì)算?f。將 f 正向評(píng)估 n 次，方向向量取為標(biāo)準(zhǔn)基（獨(dú)熱碼）向量ei∈R n，i=1 ... n，其中ei表示在第i個(gè)坐標(biāo)上為1、其他地方為0的向量，這時(shí)，只用正向模式就可以計(jì)算?f。這樣就可以分別評(píng)估f對(duì)每個(gè)輸入?f ?θi的敏感性，把所有結(jié)果合并后就可以得到梯度?f。

為了獲得比反向傳播更優(yōu)的運(yùn)行時(shí)間優(yōu)勢(shì)，我們需要在每個(gè)優(yōu)化迭代中運(yùn)行一次正向模式。在一次正向運(yùn)行中，我們可以將方向v理解為敏感度加權(quán)和中的權(quán)重向量，即P i ?f ?θi vi，盡管這沒(méi)辦法區(qū)分每個(gè)θi在最終總數(shù)中的貢獻(xiàn)。因此，我們使用權(quán)重向量v將總體敏感度歸因于每個(gè)單獨(dú)的參數(shù)θi，與每個(gè)參數(shù)θi的權(quán)重vi成正比（例如，權(quán)重小的參數(shù)在總敏感度中的貢獻(xiàn)小，權(quán)重大的參數(shù)貢獻(xiàn)大）。

總之，每次評(píng)估正向梯度時(shí)，我們只需做以下工作：

對(duì)一個(gè)隨機(jī)擾動(dòng)向量v～p(v)進(jìn)行采樣，其大小與f的第一個(gè)參數(shù)相同。
通過(guò)AD正向模式運(yùn)行f函數(shù)，在一次正向運(yùn)行中同時(shí)評(píng)估f(θ)和?f(θ)-v，在此過(guò)程中無(wú)需計(jì)算?f。得到的方向?qū)?shù)（?f(θ)-v）是一個(gè)標(biāo)量，并且由AD精確計(jì)算（不是近似值）。
將標(biāo)量方向?qū)?shù)?f(θ)-v與矢量v相乘，得到g(θ)，即正向梯度。

圖 1 顯示了 Beale函數(shù)的幾個(gè)正向梯度的評(píng)估結(jié)果。我們可以看到擾動(dòng)vk（橙色）如何在k∈[1，5]的情況下轉(zhuǎn)化為正向梯度（?f-vk）vk（藍(lán)色），在受到指向限制時(shí)偶爾也會(huì)指向正確的梯度（紅色）。綠色箭頭表示通過(guò)平均正向梯度來(lái)評(píng)估蒙特卡洛梯度，即1 K PK k=1(?f - vk)vk≈E[(?f - v)v]。

正向梯度下降

他們構(gòu)建了一個(gè)正向梯度下降（FGD）算法，用正向梯度g代替標(biāo)準(zhǔn)梯度下降中的梯度?f（算法1）。

在實(shí)踐中，他們使用小型隨機(jī)版本，其中 ft 在每次迭代中都會(huì)發(fā)生變化，因?yàn)樗鼤?huì)被訓(xùn)練中使用的每一小批數(shù)據(jù)影響。研究者注意到，算法 1 中的方向?qū)?shù)dt可以為正負(fù)數(shù)。如果為負(fù)數(shù)，正向梯度gt的方向會(huì)發(fā)生逆轉(zhuǎn)，指向預(yù)料中的真實(shí)梯度。圖1顯示的兩個(gè)vk樣本，證明了這種行為。

在本文中，他們將范圍限制在FGD上，單純研究了這一基礎(chǔ)算法，并將其與標(biāo)準(zhǔn)反向傳播進(jìn)行比較，不考慮動(dòng)量或自適應(yīng)學(xué)習(xí)率等其他各種干擾因素。筆者認(rèn)為，正向梯度算法是可以應(yīng)用到其他基于梯度算法的優(yōu)化算法系列中的。

實(shí)驗(yàn)

研究者在PyTorch中執(zhí)行正向AD來(lái)進(jìn)行實(shí)驗(yàn)。他們發(fā)現(xiàn)，正向梯度與反向傳播這兩種方法在內(nèi)存上沒(méi)有實(shí)際差異（每個(gè)實(shí)驗(yàn)的差異都小于0.1%）。

邏輯回歸

圖 3 給出了多叉邏輯回歸在MNIST數(shù)字分類上的幾次運(yùn)行結(jié)果。我們觀察到，相比基本運(yùn)行時(shí)間，正向梯度和反向傳播的運(yùn)行時(shí)間成本分別為 Rf=2.435 和 Rb=4.389，這與人們對(duì)典型AD系統(tǒng)的預(yù)期相符。

Rf/Rb=0.555和Tf/Tb=0.553的比率表明，在運(yùn)行時(shí)間和損失性能方面，正向梯度大約比反向傳播快兩倍。

在簡(jiǎn)單的模型中，這些比率是一致的，因?yàn)檫@兩種技術(shù)在空間行為的迭代損失上幾乎相同，這意味著運(yùn)行時(shí)收益幾乎直接反映在每個(gè)時(shí)間空間的損失上。

多層神經(jīng)網(wǎng)絡(luò)

圖4顯示了用多層神經(jīng)網(wǎng)絡(luò)在不同學(xué)習(xí)率下進(jìn)行MNIST分類的兩個(gè)實(shí)驗(yàn)。他們使用了三個(gè)架構(gòu)大小分別為1024、1024、10的全連接層。在這個(gè)模型架構(gòu)中，他們觀察到正向梯度和反向傳播相對(duì)于基礎(chǔ)運(yùn)行時(shí)間的運(yùn)行成本為Rf=2.468和Rb=4.165，相對(duì)測(cè)量 Rf/Rb 平均為0.592，與邏輯回歸的情況大致相同。

有趣的是，在第二個(gè)實(shí)驗(yàn)中（學(xué)習(xí)率為2×10-4），我們可以看到正向梯度在每個(gè)迭代損失圖中都實(shí)現(xiàn)了快速的下降。作者認(rèn)為，這種行為是由于常規(guī)SGD（反向傳播）和正向SGD算法的隨機(jī)性不同所導(dǎo)致的，因此他們推測(cè)：正向梯度引入的干擾可能有利于探索損失平面。

我們可以從時(shí)間曲線圖看到，正向模式減少了運(yùn)行時(shí)間。我們看到，損失性能指標(biāo)Tf/Tb值為0.211，這表明在驗(yàn)證實(shí)驗(yàn)損失的過(guò)程中，正向梯度的速度是反向傳播的四倍以上。

卷積神經(jīng)網(wǎng)絡(luò)

圖 5 展示了一個(gè)卷積神經(jīng)網(wǎng)絡(luò)對(duì)同一MNIST分類任務(wù)的正向梯度和反向傳播的比較。

在這個(gè)架構(gòu)中，他們觀察到，相對(duì)于基本運(yùn)行時(shí)間，正向AD的性能最好，其中正向模式的Rf=1.434，代表了在基本運(yùn)行時(shí)間之上的開銷只有 43%。Rb=2.211 的反向傳播非常接近反向 AD 系統(tǒng)中所期待的理想情況。Rf/Rb=0.649 代表了正向AD運(yùn)行時(shí)間相對(duì)于反向傳播的一個(gè)顯著優(yōu)勢(shì)。在損失空間，他們得到一個(gè)比率 Tf /Tb=0.514，這表明在驗(yàn)證損失的實(shí)驗(yàn)中，正向梯度的速度比反向傳播的速度要快兩倍。

可擴(kuò)展性

前面的幾個(gè)結(jié)果表明：

不用反向傳播也可以在一個(gè)典型的ML訓(xùn)練管道中進(jìn)行訓(xùn)練，并且以一種競(jìng)爭(zhēng)計(jì)算的方式來(lái)實(shí)現(xiàn)；
在相同參數(shù)（學(xué)習(xí)率和學(xué)習(xí)率衰減）的情況下，正向AD比反向傳播所消耗的時(shí)間要少很多。

相對(duì)于基礎(chǔ)運(yùn)行時(shí)的成本，我們看到，對(duì)于大部分實(shí)驗(yàn)，反向傳播在Rb∈[4,5]內(nèi)，正向梯度在Rf∈[3,4]內(nèi)。我們還觀察到，正向梯度算法在整個(gè)范圍內(nèi)對(duì)運(yùn)行都是有利的。Rf/Rb比率在10層以內(nèi)保持在0.6以下，在100層時(shí)略高于0.8。重要的是，這兩種方法在內(nèi)存消耗上幾乎沒(méi)有差別。

結(jié)論

總的來(lái)說(shuō)，這篇工作的幾點(diǎn)貢獻(xiàn)主要如下：

他們將「正向梯度」（forward gradient）定義為：一個(gè)無(wú)偏差的、基于正向自動(dòng)微分且毫不涉及到反向傳播的梯度估算器。
他們?cè)赑yTorch中從零開始，實(shí)現(xiàn)了正向模式的自動(dòng)微分系統(tǒng)，且完全不依賴PyTorch中已有的反向傳播。
他們把正向梯度模式應(yīng)用在各類隨機(jī)梯度下降（SGD）優(yōu)化中，最后的結(jié)果充分證明了：一個(gè)典型的現(xiàn)代機(jī)器學(xué)習(xí)訓(xùn)練管道可以只使用自動(dòng)微分正向傳播來(lái)構(gòu)建。
他們比較了正向梯度和反向傳播的運(yùn)行時(shí)間和損失消耗等等，證明了在一些情況下，正向梯度算法的速度比反向傳播快兩倍。