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

圖靈獎得主，AI大神Yann LeCun發Twitter推薦。LeCun認為，本文用product key memory層替代了BERT模型中的結構層，實現了與后者相當的性能，計算量降低了一半。

本文介紹了一種可以輕松集成到神經網絡中的結構化存儲器。該存儲器在設計上非常大，架構的容量顯著增加，參數數量可達十億個，而增加的計算成本基本上可忽略不計。存儲器的設計和訪問模式基于產品密鑰，可實現快速準確的最近鄰搜索。

這一新方法在保持計算成本不增加的同時，大幅增加了參數數量，使得整個系統在訓練和測試時，能夠在預測準確度和計算效率之間進行更優化的權衡。這個存儲器層能夠處理超大規模的語言建模任務。

在實驗中，我們使用一個包含高達300億字的數據集，并將存儲器層插入最先進的transformer的架構網絡中。結果表明，只有12層的存儲器增強模型的性能，優于24層的基線transformer模型，而在推理時間方面，前者比后者快兩倍。相關代碼已經發布，以用于重現實驗。

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

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

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

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

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

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

總的來說，本文的主要貢獻如下：

引入了一個新的網絡層，大幅擴充了神經網絡的容量，在訓練和測試時只需要很小的計算成本，幾乎可以忽略不計。

提出了新的快速索引策略，通過構造提供精確的最近鄰域搜索，并避免了依賴在訓練期間重新學習的索引結構產生的缺陷。

在一個大型transformer最先進網絡架構中演示了本文中的方法，該網絡由24層組成。我們的方法有1個存儲器和12層結構，結果性能與24層transformer架構相當，推理時間則是后者的兩倍。實驗表明，為各種復雜性的transformer網絡架構添加更多存儲器層，可以為目標任務提供系統而顯著的性能提升。

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

實驗過程

數據集

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

在此數據集上，對于維度為1024的16層模型，當驗證困惑（perplexity）開始增加時，獲得25.3的測試困惑度。

研究人員從公共通用爬網中提取了30倍大的語料庫，訓練集由280億字組成（140 GB的數據）。

與One Billion Word語料庫不同，研究人員的語料庫沒有改變句子，允許模型學習遠程依賴。在這個數據集上，沒有觀察到任何過度擬合，并且系統地增加模型容量導致驗證集上的更好性能。

操作細節

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

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

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

下圖是對比有內和沒有內存時，模型的測試困惑；以及在測試集上的速度和困惑之間進行權衡。

結果

上圖顯示了CC-News語料庫測試集上，不同模型的困惑度。研究人員觀察到，增加維度或層數會導致在所有模型中顯著改善困惑。

但是，為模型添加內存比增加層數更有利;例如，當隱藏單元的數量為1024和1600時，具有單個存儲器和12層的模型優于具有相同隱藏尺寸和24層的無記憶模型。

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

消融研究（Ablation study）

在為了研究不同組件對內存層的影響，并測量它們如何影響模型性能和內存使用情況。除非特別注明，這里考慮存儲器為5122=262k插槽，4個存儲器磁頭，k=32個選定鍵，研究人員將其嵌入到第5層。

上圖顯示增加head數或k-NN數。可以改善模型的復雜性和內存使用。

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

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

總結

論文介紹了一個內存層，它允許以幾乎可以忽略不計的計算開銷大幅提高神經網絡的容量。

該層的效率依賴于兩個關鍵因素：將key分解為產品集，以及對內存值的稀疏讀/寫訪問。圖層被集成到現有的神經網絡架構中。

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