Meta探索大模型記憶層，擴充套件至1280億個引數，優於MoE

机器之心發表於2025-01-04

預訓練語言模型通常在其引數中編碼大量資訊，並且隨著規模的增加，它們可以更準確地回憶和使用這些資訊。對於主要將資訊編碼為線性矩陣變換權重的密集深度神經網路來說，引數大小的擴充套件直接與計算和能量需求的增加相關。語言模型需要學習的一個重要資訊子集是簡單關聯。雖然前饋網路原則上（給定足夠的規模）可以學習任何函式，但使用聯想記憶（associative memory）會更高效。

記憶層（memory layers）使用可訓練的鍵值查詢機制向模型新增額外的引數，而不會增加 FLOP。從概念上講，稀疏啟用的記憶層補充了計算量大的密集前饋層，提供了廉價地儲存和檢索資訊的專用容量。

最近，Meta 的一項新研究使記憶層超越了概念驗證，證明了它們在大型語言模型（LLM）擴充套件中的實用性。

論文標題：Memory Layers at Scale
論文地址：https://arxiv.org/pdf/2412.09764
專案地址：https://github.com/facebookresearch/memory

在下游任務中，透過改進的記憶層增強的語言模型的效能優於計算預算兩倍以上的密集模型，以及在計算和引數相當的專家混合（MoE）模型。

這項工作表明，當記憶層得到充分改進和擴充套件時，可以用於增強密集神經網路，從而帶來巨大的效能提升。透過用記憶層替換一個或多個 transformer 層的前饋網路（FFN）來實現這一點（保持其他層不變）。這些優勢在各種基本模型大小（從 1.34 億到 80 億引數）和記憶體容量（最多 1280 億引數）中都是一致的。這意味著儲存容量實現了兩個數量級的飛躍。

記憶增強架構

可訓練的記憶層類似於注意力機制。給定一個查詢

，一組鍵

，以及值

。輸出是值的軟組合，根據 q 和相應鍵之間的相似性進行加權。

在使用時，記憶層與注意力層之間存在兩個區別。

首先，記憶層中的鍵和值是可訓練引數，而不是啟用引數；
其次，記憶層在鍵和值的數量方面通常具有更大的規模，因此稀疏查詢和更新是必需的。

該研究將鍵-值對的數量擴充套件到數百萬。在這種情況下，只有 top-k 最相似的鍵和相應的值被輸出。一個簡單的記憶層可以用下面的等式來描述：

其中，I 是一組指標，

，輸出

。

擴充套件記憶層

擴充套件記憶層時面臨的一個瓶頸是「查詢 - 鍵」檢索機制。簡單的最近鄰搜尋需要比較每一對查詢 - 鍵，這對於大型記憶來說很快就變得不可行。雖然可以使用近似向量相似性技術，但當鍵正在不斷訓練並需要重新索引時，將它們整合起來是一個挑戰。相反，本文采用了可訓練的「product-quantized」鍵。

並行記憶。記憶層是記憶密集型的，主要是由於可訓練引數和相關最佳化器狀態的數量龐大導致的。該研究在多個 GPU 上並行化嵌入查詢和聚合，記憶值在嵌入維度上進行分片。在每個步驟中，索引都從程序組中收集，每個 worker 進行查詢，然後將嵌入的部分聚合到分片中。此後，每個 worker 收集與其自身索引部分相對應的部分嵌入。該過程如圖 2 所示。