摘要

長序列建模一直面臨著一個根本性的權(quán)衡問題：類RNN模型中壓縮固定大小記憶的效率與基于注意力的Transformer中無損增長記憶的保真度之間的矛盾。ByteDance Seed團隊受認(rèn)知科學(xué)中多存儲模型（Multi-Store Model）的啟發(fā)，提出了一種創(chuàng)新的人工神經(jīng)網(wǎng)絡(luò)記憶框架——人工海馬體網(wǎng)絡(luò)（Artificial Hippocampus Networks, AHN）。該方法巧妙地將Transformer的KV緩存滑動窗口作為無損短期記憶，同時通過可學(xué)習(xí)的AHN模塊將窗口外信息遞歸壓縮為固定大小的緊湊長期記憶。

AHN框架示意圖

實驗結(jié)果表明，AHN增強模型在長上下文基準(zhǔn)測試中表現(xiàn)卓越，不僅持續(xù)優(yōu)于滑動窗口基線，甚至在某些任務(wù)上超越了全注意力模型，同時大幅降低了計算和內(nèi)存需求。特別值得注意的是，在Qwen2.5-3B-Instruct模型上應(yīng)用AHN后，僅增加0.4%的參數(shù)量，就實現(xiàn)了40.5%的FLOP減少和74.0%的內(nèi)存緩存降低，同時將LV-Eval（128k序列長度）的平均得分從4.41提升至5.88。

技術(shù)背景與動機

長序列建模的挑戰(zhàn)

在深度學(xué)習(xí)領(lǐng)域，處理長序列數(shù)據(jù)一直是一個核心挑戰(zhàn)。傳統(tǒng)的RNN類模型通過固定大小的隱藏狀態(tài)來維護記憶，這種設(shè)計確保了每步恒定的內(nèi)存和計算開銷，使其在處理長序列時具有良好的效率。然而，將所有歷史信息壓縮到固定大小的記憶中不可避免地會導(dǎo)致信息丟失，特別是在需要精確長距離信息回憶的任務(wù)中表現(xiàn)不佳。

為了解決RNN的局限性，注意力機制和Transformer架構(gòu)應(yīng)運而生。在因果注意力中，鍵值緩存充當(dāng)記憶：對每個輸入標(biāo)記生成新的鍵和值并附加到緩存中。與RNN不同，這種記憶本質(zhì)上是無損的，因為它保留了所有標(biāo)記級信息，從而提供了更高的記憶容量。然而，KV緩存的無損特性是一把雙刃劍：雖然它能夠?qū)崿F(xiàn)強大的記憶保留，但內(nèi)存大小隨序列長度線性增長，注意力更新的總計算成本呈二次方增長。

認(rèn)知科學(xué)的啟發(fā)

研究團隊從認(rèn)知科學(xué)和神經(jīng)科學(xué)中的多存儲記憶模型（MSM）理論中汲取靈感。該理論表明，盡管無損短期記憶（或稱工作記憶）的容量和持續(xù)時間有限，但海馬體會持續(xù)將它們鞏固為長期皮質(zhì)表征。人腦通過這種機制在成年早期和中期保持幾乎恒定的體積，同時仍然支持人類一生中對信息的高效處理。

這一生物學(xué)機制為解決人工神經(jīng)網(wǎng)絡(luò)中的長序列建模問題提供了重要啟示：可以設(shè)計一個系統(tǒng)，既保持近期信息的精確性（短期記憶），又能有效壓縮歷史信息（長期記憶），從而在效率和保真度之間找到最佳平衡點。

方法論詳解

人工海馬體網(wǎng)絡(luò)的核心設(shè)計

AHN的核心思想是模擬生物海馬體的功能，將無損短期記憶轉(zhuǎn)換為壓縮的長期記憶。具體而言，AHN與大小為W的滑動注意力窗口協(xié)同工作。對于步驟t > W的標(biāo)記，AHN通過處理剛剛退出滑動窗口的鍵值對(k_{t-W}, v_{t-W})來更新壓縮記憶。

這種遞歸記憶更新的數(shù)學(xué)表達(dá)式為：

ht?W=AHN((kt?W,vt?W),ht?W?1)

其中ht?Wht?W是更新后的壓縮記憶，總結(jié)了直到位置t-W（包括該位置）的上下文信息。由于方程的遞歸性質(zhì)，AHN可以用類RNN架構(gòu)實現(xiàn)，從而實現(xiàn)對長上下文歷史的可學(xué)習(xí)和高效壓縮。

AHN工作機制圖

與無損記憶的集成

在預(yù)定義的滑動窗口內(nèi)，系統(tǒng)應(yīng)用標(biāo)準(zhǔn)因果注意力來保持近期標(biāo)記的無損記憶。一旦輸入序列長度超過窗口大小，AHN就會被激活，將窗口外的KV對壓縮為固定大小的壓縮記憶。壓縮完成后，可以安全地丟棄窗口外的原始KV對，僅保留窗口內(nèi)的KV緩存。

最終，當(dāng)前查詢qtqt從壓縮記憶和無損記憶中訪問信息來產(chǎn)生輸出：

yt=f(ht?W,{(ki,vi)}i=t?W+1t,qt)

圖片

這種設(shè)計確保了模型既能利用精確的短期信息，又能從壓縮的歷史摘要中獲益，實現(xiàn)了效率與性能的雙重優(yōu)化。

具體實現(xiàn)方案

研究團隊使用三種現(xiàn)代線性遞歸模型來實例化AHN：Mamba2、DeltaNet（DN）和GatedDeltaNet（GDN），分別產(chǎn)生AHN-Mamba2、AHN-DN和AHN-GDN。

以AHN-GDN為例，其通過門控增量規(guī)則更新記憶：

ht?W=α(xt?W)(I?β(xt?W)kt?WTkt?W)ht?W?1+β(xt?W)kt?WTvt?W

與原始GatedDeltaNet不同，AHN-GDN只壓縮滑動窗口外的標(biāo)記。對于每個位置t，從xtxt派生的查詢qtqt用于訪問壓縮記憶ht?Wht?W。輸出進(jìn)一步由門函數(shù)γ(xt)γ(xt)調(diào)制，然后通過線性投影進(jìn)行轉(zhuǎn)換。

復(fù)雜度分析與效率優(yōu)勢

計算復(fù)雜度對比

復(fù)雜度分析圖

AHN的一個顯著優(yōu)勢在于其優(yōu)越的復(fù)雜度特性。傳統(tǒng)的全注意力機制具有關(guān)于序列長度的二次計算復(fù)雜度和線性內(nèi)存增長，而AHN將注意力的計算復(fù)雜度降低為關(guān)于序列長度的線性，同時保持恒定的內(nèi)存緩存大小。

具體而言，對于長度為L的序列，全注意力的計算復(fù)雜度為O(L2)，內(nèi)存復(fù)雜度為O(L)。而AHN增強的模型將計算復(fù)雜度降低至O(WL)（其中W為窗口大?。?，內(nèi)存復(fù)雜度降至O(W)。這種改進(jìn)在處理超長序列時尤為顯著。

實際性能表現(xiàn)

在PG19數(shù)據(jù)集的57K標(biāo)記段落上的實驗結(jié)果清晰地展現(xiàn)了AHN的優(yōu)勢。標(biāo)準(zhǔn)Qwen模型在超過32K標(biāo)記上下文窗口后困惑度急劇上升，而AHN-GDN增強模型保持了持續(xù)的低困惑度。更重要的是，在FlashAttention下，基礎(chǔ)模型的內(nèi)存使用量線性增長，而AHN-GDN保持CUDA內(nèi)存使用量幾乎恒定。

訓(xùn)練框架與優(yōu)化策略

自蒸餾訓(xùn)練方法

訓(xùn)練框架圖

研究團隊采用了一種計算高效的自蒸餾方法來訓(xùn)練AHN增強模型。該框架使用開源權(quán)重的LLM（如Qwen）作為教師模型，其輸出概率記為p'。學(xué)生模型是同一個LLM，但將其注意力機制修改為在每層上操作有限感受野的滑動窗口，然后用AHN增強這些窗口注意力層。

訓(xùn)練過程中，基礎(chǔ)模型的權(quán)重被凍結(jié)，僅優(yōu)化AHN參數(shù)。學(xué)生模型通過最小化Kullback-Leibler（KL）散度來模仿教師的輸出分布：

l=KL(p′∣∣p

這種方法的優(yōu)勢在于提供了密集的學(xué)習(xí)信號，迫使AHN學(xué)習(xí)更具泛化性的上下文表示，而不是依賴訓(xùn)練數(shù)據(jù)中的捷徑。

隨機化策略的重要性

為了鼓勵A(yù)HN學(xué)習(xí)適應(yīng)不同前瞻上下文的通用壓縮模塊，研究團隊采用了隨機化滑動窗口大小的訓(xùn)練策略。實驗表明，使用固定窗口訓(xùn)練的模型往往會過擬合特定配置，無法泛化到未見的上下文長度。相比之下，隨機化訓(xùn)練使AHN在各種窗口配置下都能保持強勁性能。

實驗結(jié)果與性能評估

超長上下文基準(zhǔn)測試

在LV-Eval和InfiniteBench的128k上下文子集上，AHN增強模型展現(xiàn)出了卓越的性能。實驗配置為32768標(biāo)記的無損記憶，包括128標(biāo)記的注意力匯聚點和32640標(biāo)記的滑動窗口。

結(jié)果顯示，AHN增強模型在幾乎所有任務(wù)上都持續(xù)優(yōu)于滑動窗口注意力基線，甚至匹配或超越全注意力模型的性能，同時顯著降低了計算和內(nèi)存緩存成本。這一結(jié)果強有力地證明了壓縮記憶機制的有效性。

長上下文場景評估

在LongBench的六個平均長度超過8192標(biāo)記的任務(wù)上，AHN增強模型再次展現(xiàn)出一致的優(yōu)越準(zhǔn)確性。這些結(jié)果強烈表明，遞歸隱藏狀態(tài)有效地捕獲和利用了歷史信息，在不同場景下都能帶來性能提升。

消融實驗分析

消融實驗揭示了兩個關(guān)鍵設(shè)計選擇的重要性：

訓(xùn)練目標(biāo)的影響：自蒸餾相比標(biāo)準(zhǔn)的下一標(biāo)記預(yù)測顯著提升了性能。這是因為交叉熵?fù)p失提供稀疏的學(xué)習(xí)信號，容易使小型AHN模塊走向訓(xùn)練數(shù)據(jù)中的捷徑，而自蒸餾提供了更密集的指導(dǎo)。

隨機化的作用：隨機化滑動窗口大小相比固定窗口配置帶來了明顯的性能提升，表明這種策略有助于學(xué)習(xí)更具泛化性的壓縮表示。

梯度可視化分析

梯度可視化圖

為了深入理解AHN如何壓縮和利用窗口外信息，研究團隊通過可視化自蒸餾損失的梯度來探測AHN增強模型的反向動態(tài)。梯度幅度較小的窗口外標(biāo)記表明其信息已被AHN的壓縮記憶很好地捕獲。

數(shù)學(xué)示例的可視化結(jié)果顯示，AHN傾向于保留數(shù)學(xué)符號和數(shù)字的信息，而忽略不太關(guān)鍵的代詞和特殊標(biāo)記，這證明了其作為有針對性壓縮模塊的有效性。

技術(shù)創(chuàng)新點與貢獻(xiàn)

理論創(chuàng)新

AHN的核心創(chuàng)新在于將認(rèn)知科學(xué)中的多存儲記憶模型成功轉(zhuǎn)化為實用的神經(jīng)網(wǎng)絡(luò)架構(gòu)。這種跨學(xué)科的方法不僅解決了長序列建模中的效率問題，還為未來的記憶機制設(shè)計提供了新的思路。

架構(gòu)設(shè)計

AHN框架的模塊化設(shè)計使其能夠與不同的RNN類架構(gòu)結(jié)合，展現(xiàn)了良好的通用性和擴展性。通過將壓縮模塊抽象為AHN概念，研究團隊提供了一個更通用的記憶框架。

訓(xùn)練方法

自蒸餾訓(xùn)練方案的引入使得AHN能夠高效地利用預(yù)訓(xùn)練模型的知識，避免了從頭訓(xùn)練的巨大計算開銷。這種方法為參數(shù)高效的模型增強提供了新的范式。

局限性與挑戰(zhàn)

精確召回任務(wù)的限制

盡管AHN在效率和性能之間取得了有效平衡，但其固定大小的壓縮記憶不可避免地會帶來一些信息損失。在RULER基準(zhǔn)測試的精確召回任務(wù)中，AHN的表現(xiàn)與滑動窗口注意力相當(dāng)，但明顯不如全注意力模型。這反映了有損壓縮的固有權(quán)衡：雖然AHN增強模型能夠?qū)崿F(xiàn)高效的長上下文推理，但在需要從壓縮記憶中精確召回的任務(wù)上仍有不足。

參數(shù)效率訓(xùn)練的性能上限

由于研究采用了參數(shù)高效的自蒸餾設(shè)置，性能仍受到底層基礎(chǔ)模型容量的限制。這表明全參數(shù)訓(xùn)練可能會進(jìn)一步釋放AHN的潛力。

未來發(fā)展方向與展望

增強記憶管理機制

未來的研究可以探索更智能的記憶管理策略，例如動態(tài)決定哪些信息應(yīng)保留在無損記憶中，哪些可以安全壓縮。這可能涉及開發(fā)注意力權(quán)重驅(qū)動的記憶分配算法，或者基于任務(wù)特定重要性的自適應(yīng)壓縮策略。

多模態(tài)長序列處理

AHN框架具有擴展到多模態(tài)場景的潛力。在處理長視頻序列、多模態(tài)對話或文檔理解任務(wù)時，可以設(shè)計專門的壓縮策略來處理不同模態(tài)的信息，實現(xiàn)更高效的跨模態(tài)長序列建模。

分層記憶架構(gòu)

借鑒人腦的分層記憶結(jié)構(gòu)，未來可以開發(fā)多層次的AHN架構(gòu)，其中不同層次的AHN負(fù)責(zé)壓縮不同時間尺度的信息。這種分層設(shè)計可能會進(jìn)一步提高長期依賴建模的能力。

硬件優(yōu)化與部署

AHN的固定內(nèi)存特性使其特別適合邊緣設(shè)備部署。未來的工作可以專注于針對特定硬件平臺的AHN優(yōu)化，開發(fā)專用的加速器設(shè)計，或者探索量化和剪枝技術(shù)在AHN中的應(yīng)用。

終身學(xué)習(xí)應(yīng)用

AHN的壓縮記憶機制天然適合終身學(xué)習(xí)場景。通過持續(xù)更新壓縮記憶，模型可以在不遺忘歷史知識的情況下學(xué)習(xí)新任務(wù)。這為開發(fā)真正的終身學(xué)習(xí)系統(tǒng)提供了新的可能性。

可解釋性研究

深入理解AHN的壓縮機制對于提高模型的可解釋性至關(guān)重要。未來的研究可以開發(fā)更精細(xì)的分析工具，揭示AHN如何選擇和壓縮信息，這不僅有助于改進(jìn)算法，還能為認(rèn)知科學(xué)研究提供新的洞察。

結(jié)論

人工海馬體網(wǎng)絡(luò)代表了長序列建模領(lǐng)域的一個重要突破。通過巧妙地結(jié)合認(rèn)知科學(xué)理論和現(xiàn)代深度學(xué)習(xí)技術(shù)，AHN成功解決了效率與性能之間的根本性權(quán)衡問題。其在多個長上下文基準(zhǔn)測試中的卓越表現(xiàn)，以及顯著的計算和內(nèi)存效率提升，證明了這一方法的實用價值。

更重要的是，AHN框架的提出為記憶機制的設(shè)計開辟了新的研究方向。它不僅為當(dāng)前的長序列處理任務(wù)提供了實用的解決方案，還為未來的人工智能系統(tǒng)設(shè)計提供了重要啟示。隨著對記憶機制理解的不斷深入和技術(shù)的持續(xù)改進(jìn)，我們有理由相信AHN及其衍生技術(shù)將在更廣泛的應(yīng)用場景中發(fā)揮重要作用，推動人工智能向更高效、更智能的方向發(fā)展。

這項工作不僅在技術(shù)層面取得了顯著進(jìn)展，更重要的是它展示了跨學(xué)科研究的巨大潛力。通過將認(rèn)知科學(xué)的洞察轉(zhuǎn)化為實用的工程解決方案，研究團隊為未來的AI研究提供了一個優(yōu)秀的范例，證明了基礎(chǔ)科學(xué)研究與應(yīng)用技術(shù)開發(fā)之間的深度融合能夠產(chǎn)生革命性的創(chuàng)新成果。

相關(guān)資源與代碼

研究團隊已在GitHub上開源了AHN的完整實現(xiàn)代碼，地址為：https://github.com/ByteDance-Seed/AHN

該代碼庫包含了所有三種AHN實例（AHN-Mamba2、AHN-DN、AHN-GDN）的實現(xiàn)，以及完整的訓(xùn)練和評估腳本。研究者可以基于此代碼進(jìn)行進(jìn)一步的研究和改進(jìn)。

論文：https://arxiv.org/abs/2510.07318