MemoryLLM：可插拔的可解釋前饋記憶Transformer

理解Transformer元件在大語言模型（LLM）中的運作方式至關重要，因為這是近年來人工智慧技術進步的核心。蘋果機器學習研究團隊在2026年7月發表的一篇論文中，重新審視了前饋模組（FFN）可解釋性面臨的挑戰，並提出了MemoryLLM。MemoryLLM旨在將FFN與自注意力機制解耦，使得解耦後的FFN能夠作為上下文無關的詞元級神經檢索記憶進行獨立研究。

具體而言，研究人員深入探究了輸入詞元如何訪問FFN引數中的記憶位置，以及FFN記憶在不同下游任務中的重要性。傳統的FFN透過輸入詞元嵌入與自注意力機制共同訓練，因此其行為依賴於上下文。MemoryLLM透過直接使用詞元嵌入將FFN與自注意力隔離訓練，成功實現了上下文無關的FFN。這種設計的關鍵優勢在於，FFN可以預先計算為詞元級查詢表（ToLs），從而支援按需在視訊記憶體與儲存之間遷移。這一特性顯著提升了推理效率，特別是在處理大規模模型時，能夠減少視訊記憶體佔用並加速計算。

為了彌補因使用上下文無關詞元級嵌入訓練FFN而導致的效能差距，研究團隊還引入了Flex-MemoryLLM架構，它定位在傳統Transformer設計與完全解耦的MemoryLLM之間。Flex-MemoryLLM透過部分保留上下文資訊，在保持可解釋性的同時，接近傳統Transformer的效能水平。

此外，MemoryLLM支援即插即用的部署方式，這意味著預計算的ToLs可以方便地整合到現有模型中，而無需重新訓練整個模型。這種靈活性使得MemoryLLM特別適用於需要快速迭代和部署的場景，如邊緣計算或即時應用。同時，由於FFN以查詢表形式存在，模型的可解釋性也得到了顯著提升——研究人員可以直觀地檢視每個詞元對應的FFN輸出，從而理解模型如何利用記憶進行推理。

在實驗中，MemoryLLM在多個自然語言處理基準任務上展示了與標準Transformer相當的效能，同時推理速度提升了數倍。Flex-MemoryLLM進一步縮小了與全上下文模型的差距，在某些任務上甚至超越了傳統方法。這些結果表明，將FFN作為獨立記憶模組的設想是可行的，並且有可能開闢新的研究方向。

MemoryLLM的設計靈感來源於對FFN引數的深入分析。研究發現，FFN引數中儲存了大量的“記憶”資訊，這些資訊以詞元級方式組織。透過將FFN與自注意力解耦，MemoryLLM使得這些記憶可以被顯式地檢索和使用。這種機制類似於外部儲存，但完全整合在Transformer內部，無需額外的硬體或複雜性。此外，ToLs的預計算特性意味著模型在推理時只需要一次表查詢操作，而不是計算完整的FFN前向傳播，從而大幅減少了計算量。

該論文由Ajay Jaiswal、Lauren Hannah、Han-Byul Kim、Duc Hoang、Arnav Kundu、Mehrdad Farajtabar和Minsik Cho共同撰寫。MemoryLLM的提出為理解FFN的角色提供了新的視角，並可能對未來的模型設計和推理效率最佳化產生重要影響。它展示瞭如何透過架構創新，在不犧牲效能的前提下，提高模型的可解釋性和計算效率。