新智元報導
來源：Twietter
編輯：大明、鵬飛
【新智元導讀】Facebook AI又出新作，LeCun力薦！在BERT等先進架構中插入一個“記憶體”層，能極大提升網絡容量和性能，同時保持計算成本基本不變。實驗表明，採用新模型的12層網絡性能即與24層BERT-large模型的性能相當，運行時間減少一半。更多AI話題歡迎來新智元和群查看。
Facebook AI又出新作，在神經網絡中加入一個結構化記憶體層，極大擴充了網絡容量和參數數量，性能提升明顯，但計算成本基本保持不變，尤其適合NLP任務。
圖靈獎得主，AI大神Yann LeCun發Twitter推薦。LeCun認為，本文用product key memory層替代了BERT模型中的結構層，實現了與後者相當的性能，計算量降低了一半。
本文介紹了一種可以輕鬆集成到神經網絡中的結構化記憶體。該記憶體在設計上非常大，架構的容量顯著增加，參數數量可達十億個，而增加的計算成本基本上可忽略不計。記憶體的設計和訪問模式基於產品密鑰，可實現快速準確的最近鄰搜索。
論文地址：
https://arxiv.org/pdf/1907.05242.pdf
這一新方法在保持計算成本不增加的同時，大幅增加了參數數量，使得整個系統在訓練和測試時，能夠在預測準確度和計算效率之間進行更優化的權衡。這個記憶體層能夠處理超大規模的語言建模任務。
在實驗中，我們使用一個包含高達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∈
實驗過程
數據集
最大的公開語言建模數據集是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模型的性能，運行時間縮短了一半。
論文鏈接：
https://arxiv.org/pdf/1907.05242.pdf
香港九龍灣馬來街興發大廈15樓D室