記憶是認(rèn)知的關(guān)鍵組成部分，讓人類(lèi)能夠從充斥我們生活的無(wú)盡噪音中選擇性存儲(chǔ)和提取重要信息。相比之下，大語(yǔ)言模型（Large Language Model，LLM）則缺乏這種能力，只能無(wú)差別地存儲(chǔ)和處理所有過(guò)去的輸入，這在長(zhǎng)時(shí)間任務(wù)中對(duì)它們的性能和成本產(chǎn)生了嚴(yán)重影響。

就像人類(lèi)大腦會(huì)選擇性地保留重要信息并逐漸淡忘不重要的細(xì)節(jié)一樣，人工智能系統(tǒng)也需要這樣的智能記憶管理機(jī)制。否則，隨著模型規(guī)模的不斷擴(kuò)大，其訓(xùn)練和部署對(duì)計(jì)算資源和內(nèi)存的需求也會(huì)無(wú)限上升。

長(zhǎng)期以來(lái)，研究人員一直在探索如何讓 AI 系統(tǒng)具備更接近人類(lèi)的記憶能力。傳統(tǒng)的解決方案主要依賴(lài)預(yù)設(shè)的規(guī)則來(lái)管理模型的記憶，比如基于時(shí)間順序或注意力分?jǐn)?shù)（Attention Score）來(lái)選擇性地保留或丟棄信息。

然而，這些方法往往過(guò)于機(jī)械，無(wú)法像人類(lèi)記憶那樣智能地區(qū)分信息的重要性，導(dǎo)致在提高效率的同時(shí)往往會(huì)損害模型的性能。

在此背景下，日本初創(chuàng)公司 Sakana AI 的研究團(tuán)隊(duì)提出了一種新的解決方案——神經(jīng)注意力記憶模型（Neural Attention Memory Models，NAMMs）。

這一方案借鑒了自然進(jìn)化在塑造人類(lèi)記憶系統(tǒng)中發(fā)揮的關(guān)鍵作用，通過(guò)進(jìn)化算法訓(xùn)練一個(gè)專(zhuān)門(mén)的神經(jīng)網(wǎng)絡(luò)來(lái)，這種方法能夠像人類(lèi)大腦一樣主動(dòng)選擇和保留重要信息，從而在提高效率的同時(shí)提升模型的性能。

就像人類(lèi)大腦會(huì)根據(jù)信息的長(zhǎng)期使用價(jià)值來(lái)決定是否保留一樣，NAMMs 通過(guò)對(duì)注意力模式的分析來(lái)評(píng)估信息的重要性。其核心機(jī)制包括三個(gè)關(guān)鍵組成部分：特征提取系統(tǒng)、記憶管理網(wǎng)絡(luò)和進(jìn)化優(yōu)化策略。

打開(kāi)網(wǎng)易新聞 查看精彩圖片

首先是特征提取機(jī)制。NAMMs 采用短時(shí)傅里葉變換（STFT）來(lái)處理注意力矩陣的列向量。具體來(lái)說(shuō)，它使用大小為 32 的 Hann 窗口進(jìn)行處理，生成 17 個(gè)復(fù)值頻率的頻譜圖表示。

這種表示方式非常巧妙，因?yàn)樗缺Ａ袅俗⒁饬χ惦S時(shí)間變化的頻率特征，又大大壓縮了數(shù)據(jù)量。研究團(tuán)隊(duì)通過(guò)實(shí)驗(yàn)發(fā)現(xiàn)，這種頻譜表示比直接使用原始注意力值或手工設(shè)計(jì)的特征更有效。

其次是向后注意力記憶（BAM）架構(gòu)的設(shè)計(jì)。這是 NAMMs 的核心創(chuàng)新之一，它引入了一種特殊的注意力機(jī)制，允許 token 只能關(guān)注其在 KV 緩存中的“未來(lái)”相關(guān)內(nèi)容。

這種設(shè)計(jì)的妙處在于，它能夠建立 token 之間的競(jìng)爭(zhēng)關(guān)系，使得模型能夠?qū)W會(huì)保留最有信息量的 token。例如，當(dāng)出現(xiàn)重復(fù)的句子或詞語(yǔ)時(shí)，模型會(huì)傾向于保留最新的出現(xiàn)，因?yàn)樗烁暾纳舷挛男畔ⅰ?/p>

打開(kāi)網(wǎng)易新聞 查看精彩圖片

在優(yōu)化策略上，研究團(tuán)隊(duì)采用了 CMA-ES（協(xié)方差矩陣自適應(yīng)進(jìn)化策略）算法。傳統(tǒng)的梯度下降法難以處理記憶管理這種具有離散決策的問(wèn)題，而 CMA-ES 通過(guò)模擬自然進(jìn)化過(guò)程，能夠直接優(yōu)化非可微的目標(biāo)函數(shù)。

具體來(lái)說(shuō)，團(tuán)隊(duì)采用了增量進(jìn)化的方式，從單個(gè)任務(wù)開(kāi)始，逐步增加訓(xùn)練任務(wù)的數(shù)量，這種方式能夠提供更好的正則化效果，提高模型的泛化能力。

研究團(tuán)隊(duì)選擇 Llama 3-8b 作為基礎(chǔ)模型訓(xùn)練了 NAMMs，并在 LongBench、InfiniteBench 以及 ChouBun 上進(jìn)行了全面評(píng)估。結(jié)果顯示，NAMM 為 Llama 3-8b Transformer 帶來(lái)了明顯的性能提升，在總體表現(xiàn)上超越了已有的 H2O 和 L2 這兩種手工設(shè)計(jì)的內(nèi)存管理方法。

例如，在 LongBench 基準(zhǔn)測(cè)試中，NAMMs 不僅將 KV 緩存大小減少到原來(lái)的 25%，還實(shí)現(xiàn)了 11% 的性能提升。在 InfiniteBench 測(cè)試中模型性能從基線(xiàn)的 1.05% 提升到了 11%，同時(shí)將緩存大小減少到原來(lái)的 40%。

打開(kāi)網(wǎng)易新聞 查看精彩圖片

NAMMs 的另一個(gè)重要特性是其出色的零樣本遷移能力。研究團(tuán)隊(duì)發(fā)現(xiàn)，僅在語(yǔ)言任務(wù)上訓(xùn)練的 NAMMs 可以直接應(yīng)用到其他架構(gòu)和模態(tài)上。

比如，當(dāng)應(yīng)用到 Llava Next Video-7B 模型時(shí)，NAMMs 在 LongVideoBench 和 MLVU 基準(zhǔn)測(cè)試中都取得了不錯(cuò)的表現(xiàn)，視覺(jué)任務(wù)的性能提升了 1%，同時(shí)將視頻幀的緩存大小減少到原來(lái)的 72%。

在強(qiáng)化學(xué)習(xí)方面，使用決策轉(zhuǎn)換器（Decision Transformer）時(shí)，NAMMs 在 D4RL 基準(zhǔn)測(cè)試中實(shí)現(xiàn)了 9% 的性能提升，同時(shí)將緩存大小減少到原來(lái)的 81%。

深入分析 NAMMs 的工作機(jī)制，研究團(tuán)隊(duì)發(fā)現(xiàn)它學(xué)會(huì)了一種智能的記憶管理策略。通過(guò)觀察不同層的記憶保留模式，發(fā)現(xiàn)模型在早期和中期層傾向于保留更多和更老的 token，這可能是因?yàn)檫@些層負(fù)責(zé)處理和聚合長(zhǎng)距離信息。而在信息密度較高的代碼任務(wù)中，模型則學(xué)會(huì)了保留相對(duì)更多的 token。

實(shí)際上，NAMMs 延續(xù)了 Sakana AI 此前的研究方法，即從自然界獲取靈感，通過(guò)模擬自然進(jìn)化的過(guò)程來(lái)優(yōu)化 AI 系統(tǒng)。這一研究思路與該公司在模型合并和進(jìn)化優(yōu)化方面的技術(shù)積累具有內(nèi)在的一致性。

同 Sakana AI 此前開(kāi)發(fā)的自動(dòng)化“進(jìn)化”算法能夠自主識(shí)別和合并優(yōu)秀模型一樣，NAMMs 也采用了進(jìn)化算法來(lái)優(yōu)化記憶管理系統(tǒng)，無(wú)需人工干預(yù)即可實(shí)現(xiàn)性能的持續(xù)提升。

其特殊的研發(fā)思路，已經(jīng)為這家僅成立了一年的初創(chuàng)公司贏得了 2.1 億美元的 A 輪融資，在這融資中，其估值已達(dá) 15 億美元。

未來(lái)，研究團(tuán)隊(duì)可能會(huì)探索更復(fù)雜的記憶模型設(shè)計(jì)，比如考慮更細(xì)粒度的特征提取方法，或者研究如何將 NAMMs 與其他優(yōu)化技術(shù)結(jié)合使用。

他們表示：“這項(xiàng)工作才剛剛開(kāi)始挖掘我們新類(lèi)記憶模型的潛力，我們預(yù)計(jì)這可能會(huì)為未來(lái)幾代 Transformer 的發(fā)展提供許多新的機(jī)會(huì)。”

參考資料：

1.https://sakana.ai/namm/

2.https://arxiv.org/abs/2410.13166

運(yùn)營(yíng)/排版：何晨龍