Facebook AI又出新作:性能翻倍,計算成本不增加

Facebook AI又出新作，LeCun力薦！在BERT等先進(jìn)架構(gòu)中插入一個“存儲器”層，能極大提升網(wǎng)絡(luò)容量和性能，同時保持計算成本基本不變。實驗表明，采用新模型的12層網(wǎng)絡(luò)性能即與24層BERT-large模型的性能相當(dāng)，運行時間減少一半。

圖靈獎得主，AI大神Yann LeCun發(fā)Twitter推薦。LeCun認(rèn)為，本文用product key memory層替代了BERT模型中的結(jié)構(gòu)層，實現(xiàn)了與后者相當(dāng)?shù)男阅?，計算量降低了一半?/p>

本文介紹了一種可以輕松集成到神經(jīng)網(wǎng)絡(luò)中的結(jié)構(gòu)化存儲器。該存儲器在設(shè)計上非常大，架構(gòu)的容量顯著增加，參數(shù)數(shù)量可達(dá)十億個，而增加的計算成本基本上可忽略不計。存儲器的設(shè)計和訪問模式基于產(chǎn)品密鑰，可實現(xiàn)快速準(zhǔn)確的最近鄰搜索。

這一新方法在保持計算成本不增加的同時，大幅增加了參數(shù)數(shù)量，使得整個系統(tǒng)在訓(xùn)練和測試時，能夠在預(yù)測準(zhǔn)確度和計算效率之間進(jìn)行更優(yōu)化的權(quán)衡。這個存儲器層能夠處理超大規(guī)模的語言建模任務(wù)。

在實驗中，我們使用一個包含高達(dá)300億字的數(shù)據(jù)集，并將存儲器層插入最先進(jìn)的transformer的架構(gòu)網(wǎng)絡(luò)中。結(jié)果表明，只有12層的存儲器增強(qiáng)模型的性能，優(yōu)于24層的基線transformer模型，而在推理時間方面，前者比后者快兩倍。相關(guān)代碼已經(jīng)發(fā)布，以用于重現(xiàn)實驗。

圖1：鍵值存儲器層概述：輸入x通過查詢網(wǎng)絡(luò)處理，該查詢網(wǎng)絡(luò)生成查詢向量q，將查詢向量q與所有鍵進(jìn)行比較。輸出是與所選鍵相關(guān)的存儲器的稀疏加權(quán)和。對于大量密鑰| K |，密鑰選擇過程在實踐中成本過高。我們的product key方法是精確的，整個搜索過程非?？?。

神奇的“存儲器層”：性能翻倍，計算成本不增加

本文提出了一個鍵值存儲器（key memory）層，可以擴(kuò)展到非常大的規(guī)模，同時保持對關(guān)鍵空間的搜索精度。該層顯著增加了整個系統(tǒng)的容量，而增加的計算成本可以忽略不計。與基于鍵值存儲器的現(xiàn)有模型（圖1）不同，本文將“鍵”定義為兩個子鍵的串聯(lián)。更多細(xì)節(jié)如圖2所示，該結(jié)構(gòu)隱含地定義了一組非常大的鍵，每個鍵與值存儲器槽相關(guān)。值向量集中引入了大量參數(shù)，因為參數(shù)數(shù)量與子鍵的數(shù)量成平方關(guān)系。

圖2：product key示意圖。我們定義了兩個離散的密鑰子集（子密鑰集1和子密鑰集2）。它們會產(chǎn)生更大的密鑰集，這些密鑰永遠(yuǎn)不會明文表示。對于給定的查詢，我們將其分為兩個子查詢（q1和q2）。在每個子集中選擇k個最接近的密鑰（圖中的k = 2），從而隱含地選擇k×k個密鑰。保證使用查詢最大化內(nèi)積的k個key屬于該子集，在該子集上可以更高效地進(jìn)行搜索。

盡管存儲器slot數(shù)量很大，但找到輸入的最精確鍵是非常有效的，通常需要O（p | K |）次向量比較，其中|K |是內(nèi)存插槽的總數(shù)。所有存儲器參數(shù)都是可訓(xùn)練的，但在訓(xùn)練時每個輸入只更新少量內(nèi)存slot。密鑰選擇和參數(shù)更新的稀疏性使訓(xùn)練和推理非常有效。

本文中加入的存儲器層，可以解決現(xiàn)有架構(gòu)在給定大量可用數(shù)據(jù)的情況下遇到的問題，也可以提升運行速度。我們以語言建模任務(wù)為例，將存儲器層整合到流行的transformer架構(gòu)中。這樣做的原因是，BERT 和GPT-2 取得了巨大成功，證明了增加大型模型的容量，能夠直接轉(zhuǎn)化為對語言建模性能的大幅改進(jìn)，反過來又能促進(jìn)雙語言理解任務(wù)和文本生成任務(wù)的性能提升。

總的來說，本文的主要貢獻(xiàn)如下：

引入了一個新的網(wǎng)絡(luò)層，大幅擴(kuò)充了神經(jīng)網(wǎng)絡(luò)的容量，在訓(xùn)練和測試時只需要很小的計算成本，幾乎可以忽略不計。

提出了新的快速索引策略，通過構(gòu)造提供精確的最近鄰域搜索，并避免了依賴在訓(xùn)練期間重新學(xué)習(xí)的索引結(jié)構(gòu)產(chǎn)生的缺陷。

在一個大型transformer最先進(jìn)網(wǎng)絡(luò)架構(gòu)中演示了本文中的方法，該網(wǎng)絡(luò)由24層組成。我們的方法有1個存儲器和12層結(jié)構(gòu)，結(jié)果性能與24層transformer架構(gòu)相當(dāng)，推理時間則是后者的兩倍。實驗表明，為各種復(fù)雜性的transformer網(wǎng)絡(luò)架構(gòu)添加更多存儲器層，可以為目標(biāo)任務(wù)提供系統(tǒng)而顯著的性能提升。

圖3：左：典型的transformer模塊由自注意力層和FFN層（雙層網(wǎng)絡(luò)）組成。右圖：在我們的系統(tǒng)用product存儲器層替換了FFN層，這類似于具有非常大的隱藏狀態(tài)的稀疏FFN層。在實踐中，我們僅替換N層FFN層，其中N∈{0,1,2}

實驗過程

數(shù)據(jù)集

最大的公開語言建模數(shù)據(jù)集是One Billion Word語料庫。在該數(shù)據(jù)集上獲得良好的性能需要繁瑣的正則化，因為它現(xiàn)在對于標(biāo)準(zhǔn)體系結(jié)構(gòu)來說太小了，本實驗亦然，且觀察到即使是小模型也足以過度擬合。

在此數(shù)據(jù)集上，對于維度為1024的16層模型，當(dāng)驗證困惑（perplexity）開始增加時，獲得25.3的測試?yán)Щ蠖取?/p>

研究人員從公共通用爬網(wǎng)中提取了30倍大的語料庫，訓(xùn)練集由280億字組成（140 GB的數(shù)據(jù)）。

與One Billion Word語料庫不同，研究人員的語料庫沒有改變句子，允許模型學(xué)習(xí)遠(yuǎn)程依賴。在這個數(shù)據(jù)集上，沒有觀察到任何過度擬合，并且系統(tǒng)地增加模型容量導(dǎo)致驗證集上的更好性能。

操作細(xì)節(jié)

研究人員使用具有16個attention heads和位置嵌入的transformer架構(gòu)。

研究人員考慮具有12、16或24層的模型，具有1024或1600維度，并使用Adam優(yōu)化器訓(xùn)練模型，學(xué)習(xí)率為2.5×10^(-4)，其中β1=0.9，β2=0.98，遵循Vaswani等人的學(xué)習(xí)率計劃。

由于使用稀疏更新來學(xué)習(xí)內(nèi)存值，研究人員發(fā)現(xiàn)以更高的Adam學(xué)習(xí)率10^(-3)來學(xué)習(xí)它們是極好的。于是用PyTorch實現(xiàn)模型、在32個Volta GPU上訓(xùn)練，并使用float16操作來加速訓(xùn)練、減少模型的GPU內(nèi)存使用。

下圖是對比有內(nèi)和沒有內(nèi)存時，模型的測試?yán)Щ?；以及在測試集上的速度和困惑之間進(jìn)行權(quán)衡。

結(jié)果

上圖顯示了CC-News語料庫測試集上，不同模型的困惑度。研究人員觀察到，增加維度或?qū)訑?shù)會導(dǎo)致在所有模型中顯著改善困惑。

但是，為模型添加內(nèi)存比增加層數(shù)更有利;例如，當(dāng)隱藏單元的數(shù)量為1024和1600時，具有單個存儲器和12層的模型優(yōu)于具有相同隱藏尺寸和24層的無記憶模型。

添加2或3個存儲層進(jìn)一步提高了性能。特別是，當(dāng)內(nèi)部隱藏狀態(tài)具有1600維時，具有12層和存儲器的模型比具有24層（與BERT large的配置相同）的模型獲得更好的困惑，速度幾乎快了兩倍。將內(nèi)存添加到內(nèi)部維度等于1600的large型模型時，推理時間幾乎不會增加。

消融研究（Ablation study）

在為了研究不同組件對內(nèi)存層的影響，并測量它們?nèi)绾斡绊懩Ｐ托阅芎蛢?nèi)存使用情況。除非特別注明，這里考慮存儲器為5122=262k插槽，4個存儲器磁頭，k=32個選定鍵，研究人員將其嵌入到第5層。

上圖顯示增加head數(shù)或k-NN數(shù)?？梢愿纳颇Ｐ偷膹?fù)雜性和內(nèi)存使用。

研究人員還注意到。具有相同h×k（h是head數(shù)，k是最近鄰數(shù)）的模型具有相似的內(nèi)存使用。比如，(h, k) ∈ {(1, 64),(2, 32),(4, 16),(8, 8)}所有內(nèi)存使用率約為70％，困惑度約為20.5。

總體上看，添加更多head可以提高性能，但也會增加計算時間。實驗結(jié)果表明，head=4，k=32時可以在速度和性能之間取得良好的平衡。

總結(jié)

論文介紹了一個內(nèi)存層，它允許以幾乎可以忽略不計的計算開銷大幅提高神經(jīng)網(wǎng)絡(luò)的容量。

該層的效率依賴于兩個關(guān)鍵因素：將key分解為產(chǎn)品集，以及對內(nèi)存值的稀疏讀/寫訪問。圖層被集成到現(xiàn)有的神經(jīng)網(wǎng)絡(luò)架構(gòu)中。

研究人員通過實驗證明它在大規(guī)模語言建模方面取得了重要進(jìn)展，12層的性能達(dá)到了24層BERT-large模型的性能，運行時間縮短了一半。