Meta開源“記憶層”,重塑Transformer架構大模型
朱道义
11月3日凌晨三點,全球社交巨頭Meta分享了一個創新研究——Memory layers(記憶層)。

目前,Transformer架構的預訓練大模型在存儲、查詢數據時,隨着參數的變大對算力的需求呈指數級增長。“記憶層”則提出了新的高效查詢機制替代了傳統的查詢方法,通過比較查詢鍵與兩個較小集合中的鍵,可以快速找到最相關的鍵,而無需遍歷模型的整個記憶層。
換言之,可以在不增加算力的情況下顯著增加大模型的參數。在傳統的鍵值查找中,每個查詢都需要與記憶層中的每個鍵進行比較,以找到最匹配的值。該方法在鍵的數量較少時是可行的,但隨着記憶層規模的增長,這種暴力搜索的方式變得非常低效,需要消耗巨大算力和時間。
Product - Key Lookup是“記憶層”的核心算法之一,使用了一種分而治之的策略,將傳統的單一鍵集合分解為兩個較小的鍵集合,通過兩個階段的查找來減少必要的比較次數,從而提高查找效率。
除了計算效率之外,Product-Key Lookup模塊還優化了內存和帶寬的使用。由於每個GPU只需要處理一半的鍵,因此內存的使用量減少了一半。由於每個GPU只需要返回與自己處理的鍵相關的值,所以內存帶寬的需求也得到了優化。
Product-Key Lookup算法不僅提高了記憶層的查詢效率,還為記憶層的應用開闢了新的可能性,使得記憶層可以被應用於更大規模的數據集和更復雜的任務中,包括大規模知識圖譜的查詢、長文本的語義檢索等。
並行記憶層則主要是用於對硬件GPU的優化。在傳統的Transformer架構模型中,隨着模型規模的增加,計算和內存需求也隨之增長。特別是在處理大規模數據集時,單一的計算單元很難滿足這種需求。並行記憶層通過在多個GPU之間分配任務,有效解決這一難題。在並行記憶層的設計中,每個計算單元只負責處理一部分數據,這樣可以減少單個計算單元的負擔,同時提高整體的處理速度。
共享記憶參數則是另外一個重要優化方法,允許不同層的記憶層共享同一個參數集合。這種設計的優勢在於,它減少了模型的總參數數量,同時提高了參數的利用率。
為了應對訓練期間可能出現的變化,研究人員開發了一套動態調整策略。每當有新的鍵加入或舊有的鍵被更新時,系統會自動調整相應的子集,而無需對整個記憶池進行全面改造。這樣的設計既簡化了維護流程,又提高了系統的靈活性和適應性。
本文系觀察者網獨家稿件,未經授權,不得轉載。