突破算力限制！Meta開源“記憶層”，重塑Transformer架構(gòu)大模型

　　今天凌晨3點，全球社交巨頭Meta分享了一個創(chuàng)新研究——Memory layers(記憶層)。

　　目前，Transformer架構(gòu)的預(yù)訓(xùn)練大模型在存儲、查詢數(shù)據(jù)時，隨著參數(shù)的變大對算力的需求呈指數(shù)級增長。“記憶層”提出了新的高效查詢機制替代了傳統(tǒng)的查詢方法，通過比較查詢鍵與兩個較小集合中的鍵，可以快速找到最相關(guān)的鍵，而無需遍歷模型的整個記憶層。

　　這也就是說，可以在不增加算力的情況下顯著增加大模型的參數(shù)。例如，研究人員在僅有1.3億參數(shù)的模型中添加了128億額外的記憶參數(shù)，其性能與Meta開源的Llama2-70相當(dāng)，而算力卻比它低了10倍左右。

　　Product - Key Lookup

　　在傳統(tǒng)的鍵值查找中，每個查詢都需要與記憶層中的每個鍵進行比較，以找到最匹配的值。該方法在鍵的數(shù)量較少時是可行的，但隨著記憶層規(guī)模的增長，這種暴力搜索的方式變得非常低效，需要消耗巨大算力和時間。

　　給大家舉一個簡單的例子，你想在一個巨大的圖書館里找一本書。這個圖書館有成千上萬本書，每本書都有一個唯一的編號(相當(dāng)于記憶層中的“鍵”)。如果你要找到一本特定的書(相當(dāng)于查詢)，傳統(tǒng)的方法是逐個檢查每一本書的編號來查找你要的那一本。

　　這種方法在圖書館只有幾百本本書時可能還行得通，當(dāng)圖書館藏書量達到數(shù)萬時，逐本查找方法就變得極其耗時和低效了。

　　Product - Key Lookup是“記憶層”的核心算法之一，使用了一種分而治之的策略，將傳統(tǒng)的單一鍵集合分解為兩個較小的鍵集合，通過兩個階段的查找來減少必要的比較次數(shù)，從而提高查找效率。

　　首先，查詢鍵被分割為兩個子查詢，每個子查詢分別與兩個半鍵集合進行比較。由于每個半鍵集合的大小只有原始鍵集合的平方根大小，因此這個階段的計算量大幅減少。在第一階段，每個半鍵集合中找到與子查詢最相似的k個鍵，這個過程稱為top-k查找。

　　在第二階段，兩個半鍵集合中找到的top-k鍵被合并，以確定最終的top-k鍵。這一步驟涉及到對兩個半鍵集合中找到的鍵進行綜合評分，以確定它們與原始查詢鍵的整體相似度。需要考慮到兩個半鍵集合中的鍵的組合，以找到最佳的匹配。

　　除了計算效率之外，Product-Key Lookup模塊還優(yōu)化了內(nèi)存和帶寬的使用。由于每個GPU只需要處理一半的鍵，因此內(nèi)存的使用量減少了一半。由于每個GPU只需要返回與自己處理的鍵相關(guān)的值，所以內(nèi)存帶寬的需求也得到了優(yōu)化。

　　Product-Key Lookup算法不僅提高了記憶層的查詢效率，還為記憶層的應(yīng)用開辟了新的可能性，使得記憶層可以被應(yīng)用于更大規(guī)模的數(shù)據(jù)集和更復(fù)雜的任務(wù)中，包括大規(guī)模知識圖譜的查詢、長文本的語義檢索等。

　　并行記憶層和共享記憶參數(shù)

　　并行記憶層主要是用于對硬件GPU的優(yōu)化。在傳統(tǒng)的Transformer架構(gòu)模型中，隨著模型規(guī)模的增加，計算和內(nèi)存需求也隨之增長。特別是在處理大規(guī)模數(shù)據(jù)集時，單一的計算單元很難滿足這種需求。并行記憶層通過在多個GPU之間分配任務(wù)，有效解決這一難題。

　　在并行記憶層的設(shè)計中，每個計算單元只負責(zé)處理一部分數(shù)據(jù)，這樣可以減少單個計算單元的負擔(dān)，同時提高整體的處理速度。這種設(shè)計允許模型在保持單個計算單元負載合理的同時，處理更大規(guī)模的記憶層。使得模型可以擴展到數(shù)十億甚至數(shù)百億的參數(shù)，而不會受到單個計算單元性能的限制。

　　共享記憶參數(shù)則是另外一個重要優(yōu)化方法，允許不同層的記憶層共享同一個參數(shù)集合。這種設(shè)計的優(yōu)勢在于，它減少了模型的總參數(shù)數(shù)量，同時提高了參數(shù)的利用率。

　　當(dāng)一個記憶層接收到輸入后，它會先從共享記憶池中查找最相似的記憶單元，然后根據(jù)查詢結(jié)果生成輸出。由于所有記憶層都指向同一個記憶池，因此它們可以在不影響彼此的情況下同時進行操作。

　　為了應(yīng)對訓(xùn)練期間可能出現(xiàn)的變化，研究人員開發(fā)了一套動態(tài)調(diào)整策略。每當(dāng)有新的鍵加入或舊有的鍵被更新時，系統(tǒng)會自動調(diào)整相應(yīng)的子集，而無需對整個記憶池進行全面改造。這樣的設(shè)計既簡化了維護流程，又提高了系統(tǒng)的靈活性和適應(yīng)性。