Meta 發佈新型記憶層技術：突破參數限制，大幅提升 AI 事實準確性

Meta 公司近日發佈了一項突破性的研究成果，他們開發出一種新型的記憶層技術，可以顯著提升大型語言模型（LLM）的事實準確性，並在參數規模上實現了前所未有的擴展。這項技術不僅挑戰了傳統神經網絡的擴展方式，還爲未來的 AI 架構設計提供了新的方向。

這項研究的核心在於利用可訓練的鍵值查找機制，爲模型增加額外的參數，而無需增加計算量（FLOPs）。這種方法的核心思想是，通過稀疏激活的記憶層來補充計算密集的前饋層，從而提供專門的存儲和檢索信息的能力。

與傳統的稠密網絡相比，記憶層在處理信息存儲方面更具效率。例如，語言模型需要學習人名生日、國家首都等簡單的關聯信息，記憶層可以通過簡單的鍵值查找機制實現，這種方式比使用前饋網絡更高效。

該研究的主要貢獻在於將記憶層的規模擴展到了前所未有的程度，達到了1280億個參數。實驗結果表明，在下游任務中，配備改進型記憶層的語言模型不僅優於計算量翻倍的稠密模型，在計算量和參數量匹配的情況下，也勝過混合專家模型。尤其在事實性任務上，性能提升更爲顯著。

Meta 的研究人員通過將 Transformer 網絡中的一個或多個前饋網絡（FFN）替換爲記憶層來實現這一目標。這種替換方式在不同基礎模型大小(從1.34億到80億參數)和記憶容量(高達1280億參數)上都表現出了一致的優勢。實驗結果顯示，記憶層可以將語言模型的事實準確性提高100%以上，同時在代碼編寫和一般知識方面也有顯著提高。在許多情況下，配備記憶層的模型甚至可以達到需要4倍計算量的稠密模型的性能。

研究人員還對記憶層進行了多項改進，以克服其在規模化應用中的挑戰:

採用乘積鍵查找機制:爲了解決大規模記憶層中查詢鍵檢索的瓶頸，該研究採用了可訓練的乘積量化鍵，從而避免了對每個查詢鍵對進行比較。

並行化記憶層:爲了在多 GPU 環境下實現記憶層的並行化，研究人員將嵌入查找和聚合操作分佈在多個 GPU 上。

共享記憶機制:爲了最大限度地共享參數，研究人員在所有記憶層之間使用共享的記憶參數池。

優化性能和穩定性:研究人員使用自定義的 CUDA 內核優化了 EmbeddingBag 操作，顯著提高了內存帶寬利用率。此外，還引入了帶有 silu 非線性的輸入相關門控機制，以提高訓練性能和穩定性。

實驗結果還揭示了以下關鍵發現:

記憶層的大小對性能有顯著影響:隨着記憶層大小的增加，事實性問答的性能持續提高。

多個記憶層優於單個記憶層:使用多個共享參數的記憶層可以提高性能，但過多的記憶層會降低性能。最佳的記憶層數量爲三個。

記憶層能更快地學習事實:在訓練初期，配備記憶層的模型性能提升更快，表明記憶層有助於模型更快地學習事實.

記憶層與稠密層互補:實驗表明，稀疏的記憶層和稠密的前饋層都是必不可少的。

爲了驗證記憶層技術的有效性，研究人員在多個基準測試上進行了評估，包括:

事實性問題回答（NaturalQuestions， TriviaQA）

多跳問題回答（HotpotQA）

科學和常識性知識（MMLU， HellaSwag， OBQA， PIQA）

代碼編寫（HumanEval， MBPP）

結果顯示，配備記憶層的模型在這些測試中均優於基線模型，尤其是在事實性問題回答上，性能提升最爲明顯。

Meta 的這項研究不僅爲 AI 模型的擴展提供了新的思路，也爲解決事實性問題和提高模型性能開闢了新的道路。研究人員認爲，記憶層技術具有很強的可擴展性，未來有望在各種 AI 應用中得到廣泛應用。他們還指出，雖然記憶層在硬件加速方面仍面臨挑戰，但相信通過持續的研究和優化，其性能可以與傳統前饋網絡相媲美甚至超越.

此外，Meta 的研究團隊還希望通過新的學習方法進一步提升記憶層的性能，減少模型的遺忘、幻覺，並實現持續學習.

這項研究的發佈無疑爲 AI 領域注入了新的活力，也讓我們對未來 AI 的發展充滿了期待。

論文：https://arxiv.org/pdf/2412.09764