Sakana AI 推出通用記憶模型:讓模型學會遺忘,支援零樣本跨模態
Sakana AI 推出神經網路注意力記憶模型 NAMMs,模仿人類大腦機制讓模型主動剔除冗餘 token。該系統通過演化演算法最佳化注意力矩陣,在壓縮視訊記憶體的同時提升了長文本推理表現,還能零樣本直接遷移到視覺和強化學習模型中。

人在記住重要事情的同時,大腦其實一直在拼命遺忘無聊的細節。但現在的 AI 並不是這樣:主流 Transformer 模型就像一個收納強迫症,不管讀了幾萬行程式碼還是幾個小時的影片,所有 token 都會原封不動塞在記憶體裡,上下文越長,推理成本越嚇人。
獲得日本經濟產業省 GENIAC 超算專案資助的 Sakana AI,決定教大模型學會「主動遺忘」。他們公佈了論文《An Evolved Universal Transformer Memory》,推出了一套名為 NAMMs(神經注意力記憶模型)的注意力修剪系統。
用進化演算法,教神經網路斷舍離
在此之前,學術界減少長文本視訊記憶體佔用的常見做法是 H2O 或 L2 這類手工規則,比如死板地保留最近幾個詞或注意力最高的前幾個詞。問題是,硬性裁剪往往會傷到模型原本的理解能力,屬於拿效能換速度。
Sakana AI 的思路不同,他們把「留還是扔」做成了一個小巧的神經網路分類器。但這裡有個很棘手的問題:每個 token 要麼留下來,要麼徹底丟掉,這種非黑即白的離散決策是不可導的,沒法用反向傳播來梯度訓練。
Sakana AI 拿出了他們最擅長的進化演算法。通過模擬自然界的突變和試錯篩選,系統不斷演化這套小模型的網路引數,在沒有梯度指導的情況下,找到了最會挑重點的淘汰策略。
具體執行時,NAMMs 會抓取注意力矩陣的數值,先把它轉換成音訊和地震學中常見的語譜圖,再用指數移動平均壓縮成緊湊特徵,最後打出一個保留分數,低於門檻的直接丟棄。
讀文本演化出來的經驗,直接能給影片挑刺
更有意思的是這套系統的遷移能力。因為所有 Transformer 每一層都有注意力矩陣,NAMMs 盯的是注意力分佈本身,而不是具體文字。
研究團隊只在 Llama 3 8B 語言模型上演化了 NAMMs,然後直接做零樣本遷移。把它掛到 70B 的大尺寸 Llama 上,或者給影片模型 Llava Next Video 剔除冗餘幀,甚至給強化學習模型 Decision Transformer 扔掉次優動作,全部無需微調就能直接生效。
團隊拆解發現,NAMMs 在不同層和不同任務上的行為完全不同。在模型淺層,它會保留任務前言和全域性關鍵詞;到了深層,這些資訊已經被消化,它就會果斷忘掉前言,專心處理區域性細節。
面對程式碼時,被它修剪掉的通常是整塊的空格、註釋和樣板程式碼;換成普通英文文本,它丟掉的反而是語法中冗餘的虛詞,把核心人名和概念留下來。為了測試長文本表現,他們還開源了日語長文本評測基準 ChouBun,涵蓋 36 項長文本與程式設計任務。
為什麼重要
長上下文現在的最大痛點不是模型讀不懂,而是視訊記憶體放不下、算力跑不起。對於開發者和部署端來說,NAMMs 這種外掛式的修剪機制提供了一個不用重新預訓練大模型就能給推理瘦身的新路子,而且官方開源了全部程式碼,值得長文本場景的工程師跟進驗證。
對整個行業而言,這打破了過去大家只能在全量 KV 快取硬抗、或者換成 Mamba 這種線性架構之間二選一的僵局。它證明了 Transformer 現有的注意力權重本身就藏著篩選特徵的通用語言,不同模態之間竟然能直接共享這套剪枝邏輯。
我們認為,Sakana AI 這一步走得相當狡黠。他們巧妙繞開了不可導的數學死衚衕,把進化計算的優勢發揮得淋漓盡致。不過看起來它的侷限也很實在:進化演算法雖然訓練出來的結果好用,但訓練過程本身極其消耗算力,需要超算叢集長時間試錯。普通團隊很難低成本復現演化過程,未來更現實的路徑,或許是直接白嫖他們演化成熟的通用分類權重。
來源 (1)
- sakana.ai(在新分頁開啟)原文sakana.ai