導(dǎo)語(yǔ)

對(duì)于機(jī)器學(xué)習(xí)領(lǐng)域的從業(yè)者而言，鍵值（Key-Value, KV）記憶是再熟悉不過(guò)的基礎(chǔ)組件 —— 如Transformer 的自注意力機(jī)制，KV 架構(gòu)早已成為現(xiàn)代 AI 實(shí)現(xiàn)高效記憶與檢索的核心設(shè)計(jì)。但很少有人會(huì)追問(wèn)：這套我們天天在用的計(jì)算框架，會(huì)不會(huì)和人類大腦的記憶機(jī)制，共享著同一套底層邏輯？

這篇2025 年發(fā)表于神經(jīng)科學(xué)頂刊Neuron的論文《Key-value memory in the brain》，就為這個(gè)問(wèn)題給出了一套完整、自洽的跨學(xué)科解答。論文首次將機(jī)器學(xué)習(xí)領(lǐng)域的鍵值記憶框架，與認(rèn)知心理學(xué)、神經(jīng)科學(xué)中關(guān)于記憶編碼、存儲(chǔ)、遺忘的經(jīng)典理論與實(shí)驗(yàn)證據(jù)系統(tǒng)性整合，提出了一套貫通人工與自然智能的統(tǒng)一記憶框架。

本篇解讀面向具備基礎(chǔ)機(jī)器學(xué)習(xí)知識(shí)的讀者，拆解這套框架的核心數(shù)學(xué)邏輯、與大腦記憶系統(tǒng)的對(duì)應(yīng)關(guān)系，以及相關(guān)的計(jì)算驗(yàn)證實(shí)驗(yàn)，帶大家看懂這篇跨領(lǐng)域研究的核心洞見。

關(guān)鍵詞：鍵值記憶、腦科學(xué)、機(jī)器學(xué)習(xí)、持續(xù)學(xué)習(xí)、遺忘機(jī)制

賈梓杏丨作者

趙思怡丨審校

論文標(biāo)題：Key-value memory in the brain 論文鏈接：https://www.cell.com/neuron/abstract/S0896-6273(25)00172-2 發(fā)表時(shí)間：2025年6月4日 發(fā)表期刊：Neuron

引言

鍵值（Key-Value, KV）記憶是現(xiàn)代機(jī)器學(xué)習(xí)系統(tǒng)實(shí)現(xiàn)性能突破的核心基礎(chǔ)組件之一。從 Transformer 架構(gòu)的自注意力機(jī)制，到檢索增強(qiáng)生成（RAG）系統(tǒng)，再到經(jīng)典的聯(lián)想記憶網(wǎng)絡(luò)，KV 架構(gòu)通過(guò)分離存儲(chǔ)與檢索的表征需求，解決了海量信息的高保真存儲(chǔ)與精準(zhǔn)化提取的核心矛盾，已成為當(dāng)前人工智能領(lǐng)域的通用基礎(chǔ)設(shè)計(jì)。

長(zhǎng)期以來(lái)，認(rèn)知心理學(xué)與神經(jīng)科學(xué)領(lǐng)域圍繞人類記憶機(jī)制，積累了大量關(guān)于記憶編碼、存儲(chǔ)、提取與遺忘的實(shí)驗(yàn)證據(jù)和理論框架，包括海馬體在情景記憶中的核心作用、互補(bǔ)學(xué)習(xí)系統(tǒng)理論、記憶索引假說(shuō)等。但這些分散的研究結(jié)論，始終缺乏一個(gè)統(tǒng)一的、可量化的計(jì)算框架，與機(jī)器學(xué)習(xí)領(lǐng)域的 KV 記憶研究長(zhǎng)期處于平行發(fā)展、鮮有交叉的狀態(tài)。

2025 年 6 月，哈佛大學(xué) Samuel J. Gershman、麻省理工學(xué)院 Ila Fiete 等學(xué)者在神經(jīng)科學(xué)頂刊《Neuron》發(fā)表了題為Key-value memory in the brain[1]的觀點(diǎn)論文，首次系統(tǒng)完成了機(jī)器學(xué)習(xí)鍵值記憶框架與腦科學(xué)、心理學(xué)記憶理論的深度整合，提出了一個(gè)貫通人工與自然智能的核心假說(shuō)：

(1)人類大腦的記憶系統(tǒng)，本質(zhì)是一套實(shí)現(xiàn)了鍵值分離優(yōu)化的記憶架構(gòu) —— 海馬體負(fù)責(zé)編碼用于記憶尋址的鍵（Key），新皮層負(fù)責(zé)存儲(chǔ)記憶內(nèi)容的值（Value）。

(2)許多ML方法的成功（或者局限），都是源于其作為Key-Value System的本質(zhì)（令筆者吃驚的是，甚至包括了最最基本的多層感知機(jī)）。論文通過(guò)數(shù)學(xué)推導(dǎo)、神經(jīng)生物學(xué)證據(jù)梳理與計(jì)算模擬實(shí)驗(yàn)，完整論證了這一框架的合理性，同時(shí)揭示了自然智能與人工智能在記憶機(jī)制上的收斂性。

可能有讀者擔(dān)心自己缺乏神經(jīng)科學(xué)背景而錯(cuò)過(guò)了本文最有價(jià)值的部分，但其實(shí)完全不必有這等擔(dān)憂，只需要有一定的機(jī)器學(xué)習(xí)算法基礎(chǔ)即可。本文將從鍵值記憶的計(jì)算基礎(chǔ)出發(fā)，逐步拆解該框架與大腦記憶系統(tǒng)的神經(jīng)對(duì)應(yīng)關(guān)系、核心實(shí)驗(yàn)證據(jù)，論證 “遺忘的本質(zhì)是檢索失效而非記憶丟失” 的核心結(jié)論，同時(shí)梳理論文中關(guān)于 “全連接神經(jīng)網(wǎng)絡(luò)本質(zhì)是天然鍵值記憶系統(tǒng)” 的關(guān)鍵推導(dǎo)，為讀者呈現(xiàn)這一跨領(lǐng)域研究的完整邏輯。

1. 計(jì)算基礎(chǔ)

1.1 經(jīng)典聯(lián)想記憶的核心局限

經(jīng)典的心理學(xué)與神經(jīng)科學(xué)記憶模型，大多基于相似度驅(qū)動(dòng)的模式檢索，其核心是自聯(lián)想記憶架構(gòu) —— 最典型的代表是 Hopfield 網(wǎng)絡(luò)。這類模型的核心特征是：記憶的存儲(chǔ)與檢索復(fù)用同一套表征，輸入的檢索線索與存儲(chǔ)的記憶模式直接做相似度匹配，完成記憶提取。

這種架構(gòu)存在一個(gè)無(wú)法調(diào)和的核心矛盾：存儲(chǔ)保真度與檢索區(qū)分度的優(yōu)化目標(biāo)相互沖突。

• 對(duì)存儲(chǔ)而言，表征需要最大化對(duì)記憶內(nèi)容的還原度，保留盡可能多的細(xì)節(jié)信息；

• 對(duì)檢索而言，表征需要最大化不同記憶之間的區(qū)分度，讓相似但不相同的記憶能夠被精準(zhǔn)分離，避免檢索干擾。

同一套表征無(wú)法同時(shí)完成兩個(gè)相互沖突的優(yōu)化目標(biāo)，這也是經(jīng)典記憶模型在容量、抗干擾能力上存在瓶頸的核心原因。而鍵值記憶架構(gòu)的核心創(chuàng)新，正是徹底分離了存儲(chǔ)與檢索的表征，讓兩套表征可以獨(dú)立優(yōu)化，同時(shí)滿足兩個(gè)目標(biāo)。論文原文用書籍的索引與正文做了精準(zhǔn)類比：書籍的索引（鍵）僅負(fù)責(zé)定位內(nèi)容，優(yōu)化目標(biāo)是快速、精準(zhǔn)尋址，不需要包含正文的語(yǔ)義信息；正文（值）僅負(fù)責(zé)存儲(chǔ)內(nèi)容，優(yōu)化目標(biāo)是信息保真，不需要承擔(dān)檢索尋址的功能。

1.2 鍵值記憶的核心數(shù)學(xué)形式

鍵值記憶的最早形式化，是 Kohonen 于 1972 年提出的相關(guān)矩陣記憶模型[2]，論文基于此模型，用統(tǒng)一的數(shù)學(xué)語(yǔ)言定義了鍵值記憶的完整流程，我們將在下文逐一介紹。

我們定義：對(duì)于任意一個(gè)記憶條目，其由兩個(gè)獨(dú)立的向量構(gòu)成 ——鍵向量kn（對(duì)應(yīng)記憶的地址索引）、值向量vn（對(duì)應(yīng)記憶的內(nèi)容），n為記憶條目的編號(hào)。

（1）記憶寫入：赫布學(xué)習(xí)構(gòu)建關(guān)聯(lián)矩陣

鍵值記憶的寫入過(guò)程，是通過(guò)外積運(yùn)算，將鍵與值的關(guān)聯(lián)關(guān)系編碼到關(guān)聯(lián)矩陣 M中：

關(guān)聯(lián)矩陣 M 初始化為全零矩陣，每寫入一個(gè)記憶條目，就將鍵向量的轉(zhuǎn)置與值向量的外積，累加到 M 中。這一過(guò)程本質(zhì)是神經(jīng)科學(xué)中經(jīng)典的赫布學(xué)習(xí)規(guī)則：當(dāng)編碼鍵與值的神經(jīng)元同步激活時(shí)，二者之間的突觸連接強(qiáng)度會(huì)增強(qiáng)，矩陣 M 中的元素，對(duì)應(yīng)的正是鍵 - 值神經(jīng)元之間的突觸權(quán)重。

（2）記憶讀?。翰樵兤ヅ鋵?shí)現(xiàn)內(nèi)容提取

記憶的讀取過(guò)程，是通過(guò)查詢向量q（與鍵向量維度一致）與關(guān)聯(lián)矩陣 M 的內(nèi)積運(yùn)算，完成記憶內(nèi)容的提取:

讀取的核心邏輯是：查詢向量與所有存儲(chǔ)的鍵向量做相似度匹配，匹配度越高的鍵，其對(duì)應(yīng)的值在最終輸出中的權(quán)重越大。從更生物意義的神經(jīng)網(wǎng)絡(luò)來(lái)談，這等同于激活關(guān)鍵單元上的模式q，從而在下游價(jià)值單元中產(chǎn)生基于學(xué)習(xí)突觸強(qiáng)度M的活動(dòng)模式。為了更清晰地呈現(xiàn)這一邏輯，我們可以將公式改寫為對(duì)偶形式：

其中，αn為注意力權(quán)重，由查詢向量與鍵向量的相似度計(jì)算得到：α=σ(S(K,q))。這里的S(?,?)為相似度核函數(shù)，用于計(jì)算查詢與鍵的匹配程度；σ(?) 為分離算子，用于放大不同記憶的權(quán)重差異，提升檢索的區(qū)分度。

對(duì)偶形式的核心價(jià)值，是揭示了所有鍵值記憶系統(tǒng)的通用本質(zhì)：記憶檢索的結(jié)果，永遠(yuǎn)是所有存儲(chǔ)值的加權(quán)和，權(quán)重由查詢與對(duì)應(yīng)鍵的匹配度決定。這一形式可以將幾乎所有主流的記憶模型納入統(tǒng)一框架，比如：

• 當(dāng)分離算子σ(?)為恒等函數(shù)時(shí)，得到線性化注意力，對(duì)應(yīng)fast weight programmers；

• 當(dāng)分離算子σ(?)為 softmax 函數(shù)時(shí)，得到 Transformer 的標(biāo)準(zhǔn)自注意力機(jī)制(我們將在下一小節(jié)具體解釋)；

• 當(dāng)分離算子σ(?)為閾值函數(shù)時(shí)，得到經(jīng)典的稀疏分布記憶模型。

論文中明確指出，Transformer 的自注意力機(jī)制，是鍵值記憶框架最典型、最成功的實(shí)現(xiàn)之一，其核心計(jì)算完全可以納入上述對(duì)偶形式的框架中。

自注意力的核心計(jì)算分為兩步：

1. 相似度計(jì)算（縮放點(diǎn)積）：，其中D為鍵 / 查詢向量的維度，縮放的目的是避免向量維度升高后內(nèi)積值過(guò)大，導(dǎo)致 softmax 函數(shù)梯度飽和。

2. 權(quán)重歸一化（分離算子）：這里使用 softmax 函數(shù)作為分離算子，將相似度轉(zhuǎn)換為 0-1 之間的歸一化權(quán)重，且所有權(quán)重和為 1，實(shí)現(xiàn)對(duì)最匹配記憶的聚焦。

最終的注意力輸出，依然是值向量的加權(quán)和，與鍵值記憶的通用讀取邏輯完全一致。這也意味著，Transformer 的核心能力，本質(zhì)是建立在鍵值記憶的基礎(chǔ)之上 —— 這也是論文能夠打通人工與自然智能的核心前提：既然 Transformer 的成功源于鍵值記憶架構(gòu)，那么大腦的記憶能力，很可能也基于同一套通用計(jì)算邏輯。

1.4 兩種Key-Value的計(jì)算實(shí)現(xiàn)形式

Fig.1 Two architectures for key-value memory (圖源[1])

文中指出了實(shí)現(xiàn)Key-Value計(jì)算的架構(gòu)無(wú)非兩種。左圖的結(jié)構(gòu)，對(duì)于學(xué)習(xí)過(guò)Transformer的讀者而言，其實(shí)已經(jīng)相當(dāng)熟悉：輸入的x通過(guò)Query、Vaule和Key三個(gè)權(quán)重矩陣，分別映射為q,k,v，而最終的權(quán)重矩陣M則是由k,v通過(guò)Hebbian Learning來(lái)更新，其實(shí)也就是我們的公式(1)。

右圖的結(jié)構(gòu)，或許讀者們會(huì)覺(jué)得略顯陌生。我們進(jìn)行一個(gè)簡(jiǎn)單的拆解：

輸入層：接收原始輸入x，對(duì)應(yīng)檢索線索 / 記憶條目；

隱藏層：輸入經(jīng)線性映射得到隱藏層激活α，輸入層到隱藏層的權(quán)重矩陣，就是鍵值系統(tǒng)中存儲(chǔ)的鍵矩陣K；隱藏層的激活值α，就是查詢與所有鍵匹配后得到的注意力權(quán)重，對(duì)應(yīng)對(duì)偶形式中的相似度加權(quán)系數(shù)；

輸出層：隱藏層激活經(jīng)線性映射得到最終輸出，隱藏層到輸出層的權(quán)重矩陣，就是鍵值系統(tǒng)中存儲(chǔ)的值矩陣V；最終輸出就是注意力權(quán)重對(duì)值向量的加權(quán)求和，與對(duì)偶形式的核心公式(3)完全一致。

2. 作為Key-Value的海馬體與新皮層

在展開鍵值對(duì)應(yīng)之前，我們先明確兩個(gè)核心腦區(qū)的經(jīng)典功能定位，以及領(lǐng)域內(nèi)公認(rèn)的互補(bǔ)學(xué)習(xí)系統(tǒng)（Complementary Learning Systems, CLS）框架[3]—— 這也是論文鍵值分工理論的核心神經(jīng)科學(xué)基礎(chǔ)。

• 海馬體：位于大腦內(nèi)側(cè)顳葉，是情景記憶（特定時(shí)空背景下的事件記憶）的核心腦區(qū)，大量研究已證實(shí)，海馬體損傷會(huì)導(dǎo)致嚴(yán)重的順行性失憶，即無(wú)法形成新的情景記憶；

• 新皮層：大腦外層的折疊灰質(zhì)結(jié)構(gòu)，是語(yǔ)義記憶（關(guān)于世界的通用知識(shí)、事實(shí)、感官特征）的核心存儲(chǔ)區(qū)，負(fù)責(zé)編碼記憶的具體內(nèi)容與語(yǔ)義規(guī)律。

經(jīng)典 CLS 框架認(rèn)為，海馬體與新皮層存在明確的分工：海馬體負(fù)責(zé)快速編碼單次經(jīng)歷的情景記憶，新皮層負(fù)責(zé)慢學(xué)習(xí)、提取跨事件的通用語(yǔ)義規(guī)律。而論文的核心創(chuàng)新，是在 CLS 框架的基礎(chǔ)上，用鍵值記憶的計(jì)算邏輯，重新定義了二者的分工本質(zhì)：海馬體的核心功能不是存儲(chǔ)記憶內(nèi)容，而是編碼用于尋址的鍵（Key）；記憶的具體內(nèi)容（Value），全部存儲(chǔ)在新皮層中。

2.1 鍵值分工的實(shí)證證據(jù)

論文中引用了大量的實(shí)驗(yàn)證據(jù)來(lái)證明觀點(diǎn)，我們不妨介紹幾個(gè)最為核心的實(shí)驗(yàn)證據(jù)：

證據(jù) 1：海馬體損傷會(huì)導(dǎo)致記憶的過(guò)度泛化

鍵值框架中，鍵的核心作用是精準(zhǔn)區(qū)分不同記憶，避免檢索時(shí)的混淆。如果海馬體的鍵編碼功能受損，不同記憶的尋址邊界會(huì)消失，必然出現(xiàn)記憶的過(guò)度泛化。Winocur 等人[4]的經(jīng)典實(shí)驗(yàn)完美驗(yàn)證了這一點(diǎn)：

1. 訓(xùn)練大鼠在 A 場(chǎng)景中接受電擊，正常大鼠會(huì)在 A 場(chǎng)景中表現(xiàn)出特異性的凍結(jié)反應(yīng)，在無(wú)關(guān)的 B 場(chǎng)景中無(wú)反應(yīng)；

2. 一周后測(cè)試時(shí)，大鼠表現(xiàn)出泛化效應(yīng)（語(yǔ)境特異性喪失），在兩種情境中均會(huì)凍結(jié)。這種泛化可能是由于在這段時(shí)間內(nèi)獲得的記憶干擾所致。

3. 給正常大鼠短暫放回 A 場(chǎng)景（提供精準(zhǔn)的鍵線索），可以強(qiáng)化記憶的場(chǎng)景特異性。然而，這一 “提醒效應(yīng)” 在海馬體損傷的大鼠中完全消失[5]。

這一實(shí)驗(yàn)直接證明：海馬體的核心作用是為記憶生成特異性的鍵，實(shí)現(xiàn)精準(zhǔn)尋址；沒(méi)有海馬體的鍵，新皮層的恐懼記憶值會(huì)被任意線索激活，出現(xiàn)過(guò)度泛化，與鍵值框架的預(yù)測(cè)完全一致。

證據(jù) 2：海馬體表征會(huì)主動(dòng)優(yōu)化區(qū)分度，完全匹配鍵的優(yōu)化目標(biāo)

鍵值框架中，鍵的優(yōu)化目標(biāo)是最大化不同記憶間的區(qū)分度。論文指出，海馬體的神經(jīng)表征完全遵循這一優(yōu)化邏輯，最直接的證據(jù)是海馬體的表征排斥效應(yīng)：Chanales[6] 等人的研究發(fā)現(xiàn)，當(dāng)大鼠需要區(qū)分空間上高度重疊的兩條路線時(shí)，重疊區(qū)域?qū)?yīng)的海馬體位置細(xì)胞表征，會(huì)主動(dòng)向相反方向分離，最終完全逆轉(zhuǎn)原本的空間相似關(guān)系；且排斥效應(yīng)的強(qiáng)度，與大鼠的路線區(qū)分準(zhǔn)確率直接正相關(guān)。這表明鍵的表征會(huì)被優(yōu)化到空間中相互分離的位置，最大化檢索區(qū)分度，而這一優(yōu)化過(guò)程，僅發(fā)生在負(fù)責(zé)編碼鍵的海馬體中，新皮層的內(nèi)容表征不會(huì)出現(xiàn)此類排斥效應(yīng)。

2.2 記憶丟失？不，是索引失效

經(jīng)典的記憶衰退理論認(rèn)為，遺忘是記憶痕跡隨時(shí)間逐漸衰退、最終被永久擦除的過(guò)程。而論文基于鍵值記憶框架，提出了完全相反的核心假說(shuō)：大腦的記憶一旦被編碼存儲(chǔ)，就會(huì)永久保存在新皮層中，幾乎不會(huì)被永久擦除；我們?nèi)粘ｓw驗(yàn)到的遺忘，本質(zhì)是海馬體的鍵索引功能失效，導(dǎo)致系統(tǒng)無(wú)法通過(guò)查詢線索匹配到對(duì)應(yīng)的鍵，進(jìn)而無(wú)法提取新皮層中完整存儲(chǔ)的值。

證據(jù) 1：記憶的精度不會(huì)隨時(shí)間衰減，僅可訪問(wèn)性下降

如果遺忘是記憶痕跡的衰退，那么隨著時(shí)間推移，記憶的內(nèi)容會(huì)逐漸模糊、精度持續(xù)下降；但如果遺忘是檢索失效，那么記憶的內(nèi)容精度會(huì)保持不變，只是被成功提取的概率（可訪問(wèn)性）會(huì)下降。Berens 等人的實(shí)驗(yàn)精準(zhǔn)驗(yàn)證了這一點(diǎn)：

• 實(shí)驗(yàn)者讓受試者記憶 “單詞 - 空間位置” 配對(duì)，在不同的保留間隔后測(cè)試記憶表現(xiàn)，將記憶表現(xiàn)拆分為兩個(gè)維度：可訪問(wèn)性（能否成功回憶出位置）、精度（回憶出的位置與真實(shí)位置的誤差）；

• 結(jié)果顯示：隨著保留間隔的延長(zhǎng)，記憶的可訪問(wèn)性顯著下降，但回憶成功的記憶，其位置精度完全沒(méi)有衰減，與剛學(xué)習(xí)時(shí)的精度一致。

這一結(jié)果完全符合鍵值框架的預(yù)測(cè)：值的內(nèi)容完整保存在新皮層中，精度不會(huì)隨時(shí)間變化；只是隨著時(shí)間推移，海馬體的鍵與日常線索的匹配效率下降，導(dǎo)致可訪問(wèn)性降低，也就是我們體驗(yàn)到的 “遺忘”。

證據(jù) 2：記憶干擾是遺忘的核心誘因，而非時(shí)間本身

如果遺忘是記憶痕跡的衰退，那么時(shí)間應(yīng)該是遺忘的核心決定因素；但論文指出，記憶間的檢索干擾，才是遺忘的核心誘因，這也完全符合鍵值框架的邏輯：新記憶的鍵會(huì)對(duì)舊記憶的鍵產(chǎn)生干擾，降低舊鍵與查詢的匹配優(yōu)先級(jí)，最終導(dǎo)致檢索失敗。最經(jīng)典的證據(jù)來(lái)自 Shiffrin[7] 的序列列表記憶實(shí)驗(yàn)：

• 實(shí)驗(yàn)者給受試者依次呈現(xiàn)多個(gè)單詞列表，要求回憶倒數(shù)第二個(gè)列表的內(nèi)容；

• 結(jié)果顯示：回憶的準(zhǔn)確率，僅取決于被回憶列表的長(zhǎng)度，與最后一個(gè)列表的長(zhǎng)度完全無(wú)關(guān)。

這一結(jié)果對(duì) “新記憶覆蓋舊記憶” 的衰退假說(shuō)提出了重要挑戰(zhàn)：如果遺忘的本質(zhì)是新記憶替換了舊記憶，那么最后一個(gè)列表的長(zhǎng)度越長(zhǎng)，舊列表的遺忘程度應(yīng)該越嚴(yán)重；而實(shí)驗(yàn)結(jié)果顯示，遺忘的核心誘因是同一列表內(nèi)不同記憶的鍵相互干擾，導(dǎo)致檢索失敗，而非新記憶擦除了舊記憶的內(nèi)容。

在這一視角下，機(jī)器學(xué)習(xí)中的災(zāi)難性遺忘現(xiàn)象也可以得到一種統(tǒng)一的解釋：連續(xù)學(xué)習(xí)場(chǎng)景下，模型在學(xué)習(xí)新任務(wù)后舊任務(wù)的性能暴跌，并非必然是舊任務(wù)的記憶被從權(quán)重中擦除，更可能是新任務(wù)的鍵值對(duì)干擾了舊任務(wù)的檢索通路，這也正是論文 Figure.3 （對(duì)應(yīng)本文中的Figure.4）模擬實(shí)驗(yàn)的核心設(shè)計(jì)邏輯。

3. 作為Key-Value的MLP

在前文的內(nèi)容中，我們系統(tǒng)拆解了鍵值記憶的通用計(jì)算框架、與大腦海馬體 - 新皮層系統(tǒng)的神經(jīng)對(duì)應(yīng)關(guān)系，以及「遺忘本質(zhì)是檢索失效而非記憶丟失」的核心假說(shuō)。但對(duì)于機(jī)器學(xué)習(xí)從業(yè)者而言，最核心的疑問(wèn)依然存在：這套鍵值記憶框架，是否僅存在于 Transformer、聯(lián)想記憶這類專門設(shè)計(jì)的模型中？

本小節(jié)將基于論文中的數(shù)學(xué)推導(dǎo)，展示一個(gè)頗具啟發(fā)性的結(jié)論：任何經(jīng)標(biāo)準(zhǔn)梯度下降訓(xùn)練的線性層 / MLP，都可以被等價(jià)地重寫為一種鍵值記憶（Key-Value）形式；并通過(guò)可復(fù)現(xiàn)的模擬實(shí)驗(yàn)，說(shuō)明這一表述如何為理解神經(jīng)網(wǎng)絡(luò)的學(xué)習(xí)與遺忘行為提供新的視角，包括對(duì) “遺忘是否源于檢索失效” 這一假說(shuō)的計(jì)算驗(yàn)證與支持。

3.1 核心數(shù)學(xué)證明：線性層與鍵值記憶的嚴(yán)格等價(jià)性

論文引用 Irie 等人 2022 年的核心工作[8]，通過(guò)無(wú)任何近似的線性代數(shù)推導(dǎo)，證明了標(biāo)準(zhǔn) SGD 訓(xùn)練的線性層，與鍵值記憶系統(tǒng)存在 100% 的數(shù)學(xué)等價(jià)性。這一推導(dǎo)是整篇論文打通人工神經(jīng)網(wǎng)絡(luò)與大腦記憶系統(tǒng)的核心樞紐，也是模擬實(shí)驗(yàn)的理論基礎(chǔ)。

我們從讀者最熟悉的無(wú)偏置線性層開始：對(duì)于一個(gè)線性層，輸入為行向量$，輸出為行向量，可學(xué)習(xí)權(quán)重矩陣為 ，則前向傳播的核心公式為：y=xW

在模型訓(xùn)練階段，我們通過(guò)隨機(jī)梯度下降更新權(quán)重W。定義損失函數(shù)為L，根據(jù)鏈?zhǔn)椒▌t，損失對(duì)輸出y的梯度為誤差信號(hào)(η為學(xué)習(xí)率)，則損失對(duì)權(quán)重矩陣W的梯度為：。

當(dāng)模型完成N步訓(xùn)練后，我們可以將最終的權(quán)重矩陣W，從初始權(quán)重W0開始完整展開：

而對(duì)于我們最后的輸出，我們可以得到：

首先進(jìn)行引入一個(gè)等價(jià)的符號(hào)符號(hào)約定：

x=q還是比較容易理解的，而xn= kn，vn= en可能稍微有些費(fèi)解。關(guān)于前者，其實(shí)一旦我們將輸入的x理解為q，再看到，便能理解到其含義：訓(xùn)練樣本xn在訓(xùn)練的線性層中，天然承擔(dān)了鍵值記憶里「鍵」的全部作用 —— 它是對(duì)應(yīng)誤差記憶的地址索引，用來(lái)和查詢輸入算相似度、決定記憶的貢獻(xiàn)權(quán)重，二者的數(shù)學(xué)角色與功能完全等價(jià)。

而vn= en也是同樣的思路：訓(xùn)練過(guò)程中的誤差信號(hào)en，在訓(xùn)練的線性層中，天然承擔(dān)了鍵值記憶里「值」的全部作用 —— 它是和鍵綁定存儲(chǔ)的核心記憶內(nèi)容，被鍵匹配的相似度權(quán)重加權(quán)后，直接構(gòu)成模型的最終輸出，二者的數(shù)學(xué)角色與功能完全等價(jià)。

于是，最后經(jīng)過(guò)稍稍地變形，便可以寫作：

此時(shí)，我們已經(jīng)得到了類似公式(2)(3)的Key-Value的形式了。

上述證明相對(duì)簡(jiǎn)略，更完整的過(guò)程在 Irie 等人的工作中[8]。盡管簡(jiǎn)略，但我們依然得到了一個(gè)嚴(yán)謹(jǐn)結(jié)論：線性層的前向推理過(guò)程，本質(zhì)就是一次完整的鍵值記憶檢索。我們?nèi)粘Ｊ褂玫?MLP，是由多個(gè)線性層 + 激活函數(shù)堆疊而成的復(fù)合架構(gòu)，每一個(gè)線性層都是一個(gè)獨(dú)立的鍵值記憶系統(tǒng)，因此 MLP 本身就是一套多層級(jí)的鍵值記憶架構(gòu)。

3.2 MNIST模擬實(shí)驗(yàn)

論文設(shè)計(jì)了一套完整的連續(xù)學(xué)習(xí)模擬實(shí)驗(yàn)，核心設(shè)計(jì)目標(biāo)有兩點(diǎn)：一是通過(guò)可復(fù)現(xiàn)的計(jì)算結(jié)果，驗(yàn)證線性層 / MLP 可被等價(jià)重寫為鍵值記憶形式的合理性；二是在該鍵值記憶框架下考察災(zāi)難性遺忘的成因，展示其更可能源于檢索干擾而非記憶內(nèi)容的消失，同時(shí)在計(jì)算層面類比復(fù)現(xiàn)了神經(jīng)科學(xué)中 “激活沉默記憶可恢復(fù)行為表現(xiàn)” 的經(jīng)典現(xiàn)象。

3.2.1 實(shí)驗(yàn)設(shè)計(jì)

Fig.2 MNIST手寫數(shù)字?jǐn)?shù)據(jù)集 (圖源Tensorflow官網(wǎng))

實(shí)驗(yàn)完全遵循連續(xù)學(xué)習(xí)的標(biāo)準(zhǔn)范式，全程無(wú)任何特殊的抗遺忘設(shè)計(jì)（如記憶回放、正則化等），僅使用最基礎(chǔ)的單隱藏層 MLP：

1. 模型架構(gòu)：?jiǎn)坞[藏層前饋神經(jīng)網(wǎng)絡(luò)，無(wú)偏置項(xiàng)；輸入層 784 維（28×28 灰度圖像扁平化），隱藏層 64 維，ReLU 激活函數(shù)，輸出層 4 維（對(duì)應(yīng)兩個(gè)二分類任務(wù)的 4 個(gè)類別）。

Fig.3 實(shí)驗(yàn)?zāi)Ｐ徒Y(jié)構(gòu)示意圖：將MNIST 28*28的圖像展平為1d的向量之后，與Input Layer 大小對(duì)齊，然后經(jīng)過(guò)64d的隱藏層，最后通過(guò)輸出層的四個(gè)神經(jīng)元分別識(shí)別任務(wù)1和任務(wù)2的數(shù)字“0，1”或者“T-shirt，褲子”（圖源：筆者根據(jù)論文實(shí)現(xiàn)的可視化）

1. 任務(wù)設(shè)置：雙任務(wù)連續(xù)學(xué)習(xí)，模擬大腦的持續(xù)學(xué)習(xí)與遺忘：

• 任務(wù) 1：MNIST 手寫數(shù)字二分類，僅保留數(shù)字 0、1 的樣本，訓(xùn)練集 12665 張，測(cè)試集 2115 張；

• 任務(wù) 2：FashionMNIST 服飾二分類，僅保留 T 恤、褲子的樣本，訓(xùn)練集 12000 張，測(cè)試集 2000 張。

訓(xùn)練流程：

• 階段 1：僅用任務(wù) 1 數(shù)據(jù)訓(xùn)練模型 5 個(gè) epoch，直至任務(wù) 1 測(cè)試精度達(dá)到 99%，此時(shí)模型已完整寫入任務(wù) 1 的所有鍵值記憶；

• 階段 2：凍結(jié)任務(wù) 1 對(duì)應(yīng)的權(quán)重路徑，僅用任務(wù) 2 數(shù)據(jù)訓(xùn)練模型 5 個(gè) epoch，直至任務(wù) 2 測(cè)試精度達(dá)到 95%，全程不接觸任務(wù) 1 的任何訓(xùn)練數(shù)據(jù)。

核心干預(yù)操作：訓(xùn)練完全終止后，引入光遺傳學(xué)強(qiáng)度系數(shù) β≥1，僅在推理階段放大任務(wù) 1 對(duì)應(yīng)的鍵值記憶分量，全程無(wú)反向傳播、無(wú)重新訓(xùn)練、無(wú)任務(wù) 1 數(shù)據(jù)輸入，僅通過(guò)調(diào)整 β 測(cè)試任務(wù) 1 的精度變化。

這里的「光遺傳學(xué)強(qiáng)度」是嚴(yán)格的類比：神經(jīng)科學(xué)中，光遺傳學(xué)技術(shù)通過(guò)激光精準(zhǔn)激活海馬體中編碼記憶的沉默印跡細(xì)胞，即可喚醒失憶動(dòng)物的記憶；而這里的 β 系數(shù)，正是模擬激光的激活強(qiáng)度 ——β 越大，對(duì)任務(wù) 1 記憶的檢索增益越強(qiáng)。

3.2.2的具體實(shí)現(xiàn)

在一個(gè)分類任務(wù)上，突然出現(xiàn)了“光遺傳學(xué)強(qiáng)度”這樣非常生物的表述，的確是令人費(fèi)解的。那么論文到底是如何實(shí)現(xiàn)“僅在推理階段放大任務(wù) 1 對(duì)應(yīng)的鍵值記憶分量”這一操作的？事實(shí)上并不困難，我們知道，通過(guò)梯度下降更新的權(quán)重實(shí)際上是以加法的形式進(jìn)行作用的，即：

那么，我們只需要干兩件事： 1. 分別定義兩個(gè)Layer去學(xué)習(xí)Task1和Task2，分別得到兩個(gè)任務(wù)的權(quán)重變化量。

1. 在進(jìn)行學(xué)習(xí)的前向傳播過(guò)程中，自然還是使用Wfinal，只不過(guò)為Wtask1乘上一個(gè)系數(shù)β。

代碼實(shí)現(xiàn)也相當(dāng)簡(jiǎn)單：

        return out

Fig.4 Forgetting and reactivation of memory events (圖源[1])

從左圖中我們可以看出，模型完成任務(wù) 2 的訓(xùn)練后，任務(wù) 1 的測(cè)試精度從初始的 99% 暴跌至 9%，接近隨機(jī)猜測(cè)的水平，完美復(fù)現(xiàn)了連續(xù)學(xué)習(xí)中的經(jīng)典災(zāi)難性遺忘現(xiàn)象。

但根據(jù)前文的推導(dǎo)，任務(wù) 1 的所有鍵值對(duì)，已經(jīng)完整、無(wú)損地累加到了權(quán)重矩陣中，訓(xùn)練任務(wù) 2 的過(guò)程僅新增了任務(wù) 2 的鍵值對(duì)，未對(duì)任務(wù) 1 的記憶做任何修改、刪除。此時(shí)任務(wù) 1 的精度暴跌，本質(zhì)是任務(wù) 2 的鍵值對(duì)在檢索中占據(jù)了主導(dǎo)，任務(wù) 1 的記憶被干擾淹沒(méi)，成為了「沉默記憶」，而非記憶本身丟失。

于是，當(dāng)我們觀察右圖：通過(guò) β 系數(shù)放大任務(wù) 1 的鍵值記憶分量后，任務(wù) 1 的測(cè)試精度隨 β 的增大呈現(xiàn)顯著的單調(diào)回升，無(wú)需任何重新訓(xùn)練。當(dāng)β=1.0（無(wú)干預(yù)，原始權(quán)重）時(shí)，任務(wù) 1 精度維持 9% 的遺忘狀態(tài)；而當(dāng)β=1.8時(shí)，精度回升至接近初始的 99%，完全恢復(fù)任務(wù) 1 的分類性能。

如果任務(wù) 1 的記憶真的在訓(xùn)練任務(wù) 2 時(shí)被覆蓋、擦除，那么無(wú)論 β 放大多少倍，都不可能恢復(fù)任務(wù) 1 的精度。只有當(dāng)任務(wù) 1 的所有鍵值對(duì)完整無(wú)損地保存在權(quán)重中，放大操作才會(huì)生效。而這也說(shuō)明了，災(zāi)難性遺忘的核心原因，不是舊記憶被新記憶覆蓋，而是舊記憶的檢索優(yōu)先級(jí)被新記憶干擾，導(dǎo)致檢索失效。記憶本身始終完整存在，只是變成了無(wú)法被自然線索激活的沉默。

4. 結(jié)語(yǔ)

到這里，我們拆解了這篇打通腦科學(xué)與機(jī)器學(xué)習(xí)的核心研究。考慮到讀者群體以及受限于筆者本身水平，這篇解讀并非完整的解讀，僅選擇了文中筆者認(rèn)為驚艷的部分，并強(qiáng)烈建議感興趣的讀者閱讀原文。

論文提出的這套統(tǒng)一鍵值記憶框架，為我們理解大腦記憶機(jī)制提供了一套自洽的跨學(xué)科視角：在該框架下，海馬體與新皮層的分工可被對(duì)應(yīng)為鍵值記憶的尋址編碼與內(nèi)容存儲(chǔ)，而日常的遺忘現(xiàn)象，也可被解釋為檢索通路的失效，而非記憶內(nèi)容的永久丟失。同時(shí)，這套框架也為理解深度學(xué)習(xí)模型的學(xué)習(xí)與遺忘行為提供了統(tǒng)一的解釋邏輯：從 Transformer 的自注意力機(jī)制，到最基礎(chǔ)的 MLP，都可被等價(jià)重寫為鍵值記憶的形式；而領(lǐng)域內(nèi)長(zhǎng)期關(guān)注的災(zāi)難性遺忘問(wèn)題，在該框架下也可被歸因于檢索干擾，而非記憶內(nèi)容的擦除。

人工與自然智能在記憶機(jī)制上的驚人收斂，為兩個(gè)領(lǐng)域的交叉發(fā)展打開了全新思路。

參考文獻(xiàn)

原論文代碼：https://github.com/kazuki-irie/kv-memory-brain

Tensorflow MNIST數(shù)據(jù)集官網(wǎng)：https://tensorflow.google.cn/datasets/catalog/mnist

[1]Gershman et al., Key-value memory in the brain, Neuron (2025), https://doi.org/10.1016/j.neuron.2025.02.029

[2]Kohonen, T. (1972). Correlation matrix memories. IEEE Trans. Comput. C–21, 353–359. https://doi.org/10.1109/TC.1972.5008975.

[3]McClelland, J.L., McNaughton, B.L., and O’Reilly, R.C. (1995). Why there are complementary learning systems in the hippocampus and neocortex: insights from the successes and failures of connectionist models of learning and memory. Psychol. Rev. 102, 419–457. https://doi.org/10. 1037/0033-295X.102.3.419.

[4]Winocur, G., Frankland, P.W., Sekeres, M., Fogel, S., and Moscovitch, M. (2009). Changes in context-specificity during memory reconsolidation: selective effects of hippocampal lesions. Learn. Mem. 16, 722–729. https://doi.org/10.1101/lm.1447209.

[5]Wiltgen, B.J., Zhou, M., Cai, Y., Balaji, J., Karlsson, M.G., Parivash, S.N., Li, W., and Silva, A.J. (2010). The hippocampus plays a selective role in the retrieval of detailed contextual memories. Curr. Biol. 20, 13361344. https://doi.org/10.1016/j.cub.2010.06.068.

[6]Chanales, A.J.H., Oza, A., Favila, S.E., and Kuhl, B.A. (2017). Overlap among spatial memories triggers repulsion of hippocampal representations. Curr. Biol. 27, 2307–2317.e5. https://doi.org/10.1016/j.cub.2017. 06.057.

[7]Shiffrin, R.M. (1970). Forgetting: Trace erosion or retrieval failure? Science 168, 1601–1603. https://doi.org/10.1126/science.168.3939.1601.

[8]Irie, K., Csorda ? s, R., and Schmidhuber, J. (2022). The dual form of neural networks revisited: Connecting test time predictions to training patterns via spotlights of attention. In International Conference on Machine Learning, pp. 9639–9659.

參考文獻(xiàn)可上下滑動(dòng)查看

計(jì)算神經(jīng)科學(xué)第三季讀書會(huì)

從單個(gè)神經(jīng)元的放電到全腦范圍的意識(shí)涌現(xiàn)，理解智能的本質(zhì)與演化始終是一個(gè)關(guān)于尺度的問(wèn)題。更值得深思的是，無(wú)論是微觀的突觸可塑性、介觀的皮層模塊自組織，還是宏觀的全局信息廣播，不同尺度的動(dòng)力學(xué)過(guò)程都在共同塑造著認(rèn)知與意識(shí)。這說(shuō)明，對(duì)心智的研究從最初就必須直面一個(gè)核心挑戰(zhàn)：局部的神經(jīng)活動(dòng)如何整合為統(tǒng)一的體驗(yàn)？局域的網(wǎng)絡(luò)連接又如何支撐靈活的智能行為？

繼「」與「」讀書會(huì)后，集智俱樂(lè)部聯(lián)合來(lái)自數(shù)學(xué)、物理學(xué)、生物學(xué)、神經(jīng)科學(xué)和計(jì)算機(jī)的一線研究者共同發(fā)起，跨越微觀、介觀與宏觀的視角，探索意識(shí)與智能的跨尺度計(jì)算、演化與涌現(xiàn)。重點(diǎn)探討物理規(guī)律與人工智能如何幫助我們認(rèn)識(shí)神經(jīng)動(dòng)力學(xué)，以及神經(jīng)活動(dòng)跨尺度的計(jì)算與演化如何構(gòu)建微觀與宏觀、結(jié)構(gòu)與功能之間的橋梁。

詳情請(qǐng)見：

1.

2.

3.

4.

5.

6.

7.