一種可以減少RNN訓(xùn)練時(shí)內(nèi)存需求的新方法

多倫多大學(xué)的研究人員提出Reversible RNN，一種可以減少RNN訓(xùn)練時(shí)內(nèi)存需求的新方法，在保留模型性能的同時(shí)，將激活內(nèi)存成本降低了10-15倍。

循環(huán)神經(jīng)網(wǎng)絡(luò)（RNN）在處理序列數(shù)據(jù)方面能有很好的性能，但在訓(xùn)練時(shí)需要大量內(nèi)存，限制了可訓(xùn)練的RNN模型的靈活性。

近日，多倫多大學(xué)Vector Institute的研究人員提出Reversible RNN，描述了一種可以減少RNN訓(xùn)練時(shí)內(nèi)存需求的新方法。論文題為Reversible Recurrent Neural Networks，已被NIPS 2018接收。

https://arxiv.org/pdf/1810.10999.pdf

可逆RNN（Reversible RNN）是指網(wǎng)絡(luò)中hidden-to-hidden的轉(zhuǎn)換可以逆向進(jìn)行的RNN，這就提供了一個(gè)減少訓(xùn)練的內(nèi)存需求的路徑，因?yàn)殡[藏狀態(tài)不需要存儲(chǔ)，而是可以在反向傳播期間重新計(jì)算。

這篇論文先證明了完全可逆的RNN（perfectly reversible RNNs），即不需要存儲(chǔ)隱藏的激活，在根本是受到限制的，因?yàn)樗鼈儾荒芡涬[藏狀態(tài)的信息。

然后，論文提出一種存儲(chǔ)少量bits的方案，以允許在遺忘時(shí)實(shí)現(xiàn)完美的逆轉(zhuǎn)。

這一方法實(shí)現(xiàn)了與傳統(tǒng)模型相當(dāng)?shù)男阅埽瑫r(shí)將激活內(nèi)存成本降低了10-15倍。

研究人員將這一方法擴(kuò)展到基于注意力的sequence-to-sequence模型，實(shí)驗(yàn)證明能它能保持性能，同時(shí)在encoder中將激活內(nèi)存成本降低了5-10倍，在decoder中降低了10-15倍。

可逆循環(huán)結(jié)構(gòu)

用于構(gòu)建RevNets的技術(shù)可以與傳統(tǒng)的RNN模型結(jié)合，產(chǎn)生reversible RNN。在本節(jié)中，我們提出了GRU和LSTM的可逆版本。

Reversible GRU

讓我們首先回顧一下用于計(jì)算下一個(gè)隱藏狀態(tài)的GRU方程，給定當(dāng)前隱藏狀態(tài)和當(dāng)前輸入（省略偏差）：

這里，⊙表示elementwise乘法。為了使這個(gè)更新可逆，我們將隱藏狀態(tài)h分成兩組，。使用以下規(guī)則更新這些組：

注意，而不是被用于計(jì)算的更新。我們將此模型稱為可逆門控循環(huán)單元（Reversible Gated Recurrent Unit），簡稱RevGRU。

對(duì)于i = 1,2，，因?yàn)樗莝igmoid函數(shù)的輸出，映射到開放區(qū)間（0,1）。這意味著RevGRU更新在精確算術(shù)中是可逆的：給定，我們可以使用和。然后我們可以使用以下公式找到：

Reversible LSTM

接下來我們構(gòu)建一個(gè)reversible LSTM。LSTM將隱藏狀態(tài)分離為輸出狀態(tài)h和單元狀態(tài)c。更新方程是：

我們不能直接應(yīng)用可逆技術(shù)，因?yàn)?img src="http://file.elecfans.com/web1/M00/69/4A/o4YBAFvXryOABJ4qAAABP750eE0740.png" />的非零線性變換。但可以使用以下公式實(shí)現(xiàn)可逆性：

使用和RevLSTM。

No Forgetting的限制

我們已經(jīng)證明，通過確保不丟棄任何信息，可以構(gòu)建具有有限精度的reversible RNN。

但是，對(duì)于語言建模等任務(wù)，我們還是無法找到能獲得可接受性能的架構(gòu)。

我們認(rèn)為這是由于無遺忘可逆模型（no-forgetting reversible models）的基本限制導(dǎo)致的：如果任何隱藏狀態(tài)都不能被遺忘，那么任何給定時(shí)間步長的隱藏狀態(tài)必須包含足夠的信息來重建所有先前的隱藏狀態(tài)。因此，在一個(gè)時(shí)間步長上，存儲(chǔ)在隱藏狀態(tài)中的任何信息都必須保留在所有未來的時(shí)間步長，以確保精確的重構(gòu)，從而超過了模型的存儲(chǔ)容量。

圖1：在重復(fù)任務(wù)上展開完全可逆模型的反向計(jì)算，得到sequence-to-sequence計(jì)算。左：重復(fù)任務(wù)本身，其中模型重復(fù)每個(gè)輸入標(biāo)記。 右：展開逆轉(zhuǎn)。模型有效地使用最終隱藏狀態(tài)來重建所有輸入tokens，這意味著整個(gè)輸入序列必須存儲(chǔ)在最終隱藏狀態(tài)中。

我們通過考慮一個(gè)基本的序列學(xué)習(xí)任務(wù)，即重復(fù)任務(wù)，來說明這個(gè)問題。在這個(gè)任務(wù)中，RNN被輸入一個(gè)離散token的序列，并且必須在隨后的時(shí)間步長中簡單地重復(fù)每個(gè)token。

普通的RNN模型只需要少量的隱藏單元就可以輕松解決這個(gè)任務(wù)，因?yàn)樗恍枰ｉL距離依賴關(guān)系。但請(qǐng)考慮一個(gè)完全可逆的模型如何執(zhí)行重復(fù)任務(wù)。

展開反向計(jì)算，如圖1所示，顯示了sequence-to-sequence的計(jì)算，其中編碼器和解碼器權(quán)重相關(guān)聯(lián)。編碼器接收token并產(chǎn)生最終隱藏狀態(tài)。解碼器使用該最終隱藏狀態(tài)以反向順序產(chǎn)生輸入序列。

我們通過實(shí)驗(yàn)證實(shí)，容量有限的NF-RevGRU和NF-RevLSM網(wǎng)絡(luò)無法解決重復(fù)任務(wù)。

有限遺忘實(shí)現(xiàn)可逆性

由于No Forgetting不可能，我們需要探索實(shí)現(xiàn)可逆性的第二種可能：在正向運(yùn)算期間存儲(chǔ)隱藏狀態(tài)丟失的信息，然后在反向計(jì)算終恢復(fù)它。

我們研究了fractional forgetting，即允許遺忘一小部分bits。

算法1描述了可逆乘法的完整過程。

實(shí)驗(yàn)和結(jié)果

我們?cè)趦蓚€(gè)標(biāo)準(zhǔn)RNN任務(wù)上評(píng)估了可逆模型的性能：語言建模和機(jī)器翻譯。我們希望確定使用我們開發(fā)的技術(shù)可以節(jié)省多少內(nèi)存，這些節(jié)省跟使用理想緩沖區(qū)可能節(jié)省的內(nèi)存有可比性嗎，以及這些內(nèi)存節(jié)省是否以降低性能為代價(jià)。

表1:Penn TreeBank詞級(jí)語言建模上的驗(yàn)證困惑度(內(nèi)存節(jié)省)。當(dāng)遺忘被限制為每個(gè)timestep 每個(gè)隱藏單元2、3和5bits，以及沒有限制的情況下的結(jié)果。

表2：WikiText-2單詞級(jí)語言建模的驗(yàn)證困惑度。當(dāng)遺忘被限制為每個(gè)timestep 每個(gè)隱藏單元2、3和5bits，以及沒有限制的情況下的結(jié)果。

表3：具有不同遺忘限制時(shí)Multi30K數(shù)據(jù)集上的性能。P表示測試BLEU scores; M表示訓(xùn)練期間編碼器的平均內(nèi)存節(jié)省。

總的來說，雖然Emb attention實(shí)現(xiàn)了最佳的內(nèi)存節(jié)省，但Emb + 20H在性能和內(nèi)存節(jié)省之間實(shí)現(xiàn)了最佳平衡。

具有Emb + 20H attention且遺忘最多2bits的RevGRU實(shí)現(xiàn)了34.41的test BLEU score，優(yōu)于標(biāo)準(zhǔn)GRU，同時(shí)分別在編碼器和解碼器中將激活內(nèi)存要求降低了7.1倍和14.8倍。

具有Emb + 20H attention且遺忘最多3bits的RevLSTM的test BLEU score為37.23，優(yōu)于標(biāo)準(zhǔn)LSTM，同時(shí)分別在編碼器和解碼器中將激活內(nèi)存要求降低了8.9倍和11.1倍。

baseline GRU和LSTM模型的測試BLEU分?jǐn)?shù)分別是16.07和22.35。RevGRU的測試BLEU得分為20.70，優(yōu)于GRU，同時(shí)分別在編碼器和解碼器中節(jié)省內(nèi)存7.15倍和12.92倍。RevLSTM得分為22.34，與LSTM相比，分別在編碼器和解碼器中節(jié)省了8.32倍和6.57倍的內(nèi)存。兩種可逆模型都被限制為最多遺忘5 bits。