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

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

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

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

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

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

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

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

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

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

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

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

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

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

在一個(gè)大型transformer最先進(jìn)網(wǎng)絡(luò)架構(gòu)中演示了本文中的方法，該網(wǎng)絡(luò)由24層組成。我們的方法有1個(gè)存儲器和12層結(jié)構(gòu)，結(jié)果性能與24層transformer架構(gòu)相當(dāng)，推理時(shí)間則是后者的兩倍。實(shí)驗(yàn)表明，為各種復(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層。在實(shí)踐中，我們僅替換N層FFN層，其中N∈{0,1,2}

實(shí)驗(yàn)過程

數(shù)據(jù)集

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

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

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

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

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

研究人員使用具有16個(gè)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í)率計(jì)劃。

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

下圖是對比有內(nèi)和沒有內(nèi)存時(shí)，模型的測試?yán)Щ螅灰约霸跍y試集上的速度和困惑之間進(jìn)行權(quán)衡。

結(jié)果

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

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

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

消融研究（Ablation study）

在為了研究不同組件對內(nèi)存層的影響，并測量它們?nèi)绾斡绊懩Ｐ托阅芎蛢?nèi)存使用情況。除非特別注明，這里考慮存儲器為5122=262k插槽，4個(gè)存儲器磁頭，k=32個(gè)選定鍵，研究人員將其嵌入到第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可以提高性能，但也會增加計(jì)算時(shí)間。實(shí)驗(yàn)結(jié)果表明，head=4，k=32時(shí)可以在速度和性能之間取得良好的平衡。

總結(jié)

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

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

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