Gradient Dynamics of Attention: How Cross-Entropy SculptsBayesian Manifolds

注意力梯度動(dòng)力學(xué)：交叉熵如何塑造貝葉斯流形

https://arxiv.org/pdf/2512.22473

Transformer 在精心構(gòu)建的“貝葉斯風(fēng)洞”和大規(guī)模語(yǔ)言模型中，經(jīng)驗(yàn)上展現(xiàn)出精確的概率推理能力，但基于梯度的學(xué)習(xí)如何生成所需的內(nèi)部幾何結(jié)構(gòu)仍不透明。我們對(duì)交叉熵訓(xùn)練如何重塑 Transformer 注意力頭中的注意力得分和值向量進(jìn)行了系統(tǒng)性的一階分析。我們的核心結(jié)果是一個(gè)基于優(yōu)勢(shì)（advantage-based）的注意力得分路由法則。

其中， u i
是位置 i i 處的上游梯度， α i j
是注意力權(quán)重。這些方程引發(fā)了一個(gè)正反饋循環(huán)，在該循環(huán)中路由和內(nèi)容協(xié)同專業(yè)化：查詢更強(qiáng)烈地路由到那些對(duì)其誤差信號(hào)高于平均值的值，而這些值又被拉向使用它們的查詢。我們表明，這種耦合的專業(yè)化行為類似于一個(gè)雙時(shí)間尺度的 EM 過(guò)程：注意力權(quán)重實(shí)現(xiàn) E 步（軟責(zé)任分配），而值實(shí)現(xiàn) M 步（基于責(zé)任加權(quán)的原型更新），同時(shí)查詢和鍵調(diào)整假設(shè)框架。通過(guò)受控模擬（包括一個(gè)粘性馬爾可夫鏈任務(wù)，我們?cè)谄渲袑㈤]式 EM 風(fēng)格更新與標(biāo)準(zhǔn) SGD 進(jìn)行比較），我們證明了最小化交叉熵的相同梯度動(dòng)力學(xué)也會(huì)塑造我們?cè)谙嚓P(guān)工作中識(shí)別出的低維流形，從而實(shí)現(xiàn)貝葉斯推斷。此外，在 EM 與 SGD 的比較中，EM 收斂得更快，且兩者的預(yù)測(cè)熵都接近最小貝葉斯熵。這產(chǎn)生了一幅統(tǒng)一的畫面：優(yōu)化（梯度流）催生幾何結(jié)構(gòu)（貝葉斯流形），而幾何結(jié)構(gòu)反過(guò)來(lái)又支持功能（上下文內(nèi)概率推理）。

1 引言

Transformer 已成為序列建模的主導(dǎo)架構(gòu)，但我們?nèi)匀狈?duì)梯度下降如何塑造其內(nèi)部表征的機(jī)制性理解。最近的研究表明，在受控的“貝葉斯風(fēng)洞”中，小型 Transformer 可以精確再現(xiàn)解析后驗(yàn)分布，其中鍵形成正交的假設(shè)軸，查詢實(shí)現(xiàn)漸進(jìn)式信念更新，值則沿著由后驗(yàn)熵參數(shù)化的一維流形展開 [1]。我們配套的縮放論文將這一圖景擴(kuò)展到生產(chǎn)規(guī)模模型，展示了 Pythia、Phi-2 和 LLaMA 中類似的幾何特征。

這些發(fā)現(xiàn)提出了一個(gè)自然的問(wèn)題：為什么單純的交叉熵訓(xùn)練會(huì)產(chǎn)生貝葉斯推斷所需的幾何結(jié)構(gòu)？理解這一點(diǎn)需要從靜態(tài)幾何轉(zhuǎn)向梯度動(dòng)力學(xué)：注意力分?jǐn)?shù)、查詢、鍵和值在優(yōu)化過(guò)程中如何共同演化？

1.1 貢獻(xiàn)

我們的主要貢獻(xiàn)如下：

一個(gè)由責(zé)任加權(quán)的上游梯度平均值。這引發(fā)了一個(gè)正反饋循環(huán)：查詢路由到有助于它們的值；這些值則向其使用者移動(dòng)，從而強(qiáng)化路由并促成專業(yè)化。

(4) 雙時(shí)間尺度 EM 解釋。我們表明，這些動(dòng)力學(xué)實(shí)現(xiàn)了一種隱式的類 EM 算法：注意力權(quán)重充當(dāng)軟責(zé)任分配（E 步），值作為在這些責(zé)任下更新的原型（M 步），而查詢/鍵則作為潛在分配模型的參數(shù)。注意力通常較早穩(wěn)定，而值則持續(xù)細(xì)化——這種框架-精度分離與我們?cè)陲L(fēng)洞實(shí)驗(yàn)和大型模型中的經(jīng)驗(yàn)觀察相符。

(5) 玩具實(shí)驗(yàn)及 EM 與 SGD 的比較。在合成任務(wù)中（包括一個(gè)粘性馬爾可夫鏈序列），我們將由我們的梯度分析所誘導(dǎo)的閉式 EM 風(fēng)格更新與標(biāo)準(zhǔn) SGD 進(jìn)行比較。EM 能顯著更快地達(dá)到低損失、高準(zhǔn)確率和尖銳的預(yù)測(cè)熵；SGD 收斂到類似解，但速度更慢且路由更分散。對(duì)值軌跡的 PCA 可視化揭示了新興的低維流形。

結(jié)合 [1]，我們的結(jié)果提供了一個(gè)統(tǒng)一的敘述：

梯度下降 ? 貝葉斯流形 ? 上下文內(nèi)推理。

關(guān)于“貝葉斯推斷”的澄清。本文中，“貝葉斯推斷”指的是對(duì)潛在任務(wù)變量的貝葉斯后驗(yàn)預(yù)測(cè)——而非對(duì)網(wǎng)絡(luò)權(quán)重的后驗(yàn)。我們表明，交叉熵訓(xùn)練塑造了幾何結(jié)構(gòu)，該結(jié)構(gòu)實(shí)現(xiàn)了對(duì)上下文內(nèi)假設(shè)的貝葉斯濾波，而非貝葉斯權(quán)重不確定性。

2 設(shè)置與符號(hào)說(shuō)明
我們分析一個(gè)作用于長(zhǎng)度為 T 的序列的單頭注意力機(jī)制。除非另有說(shuō)明，下標(biāo) i , j , k 的取值范圍均為 1 到 T 。

2.1向前傳球

2.2 輔助量
為簡(jiǎn)潔起見，我們定義：

3 一階梯度推導(dǎo)
我們現(xiàn)在不跳過(guò)任何步驟地推導(dǎo)所有相關(guān)梯度，重點(diǎn)關(guān)注能揭示其幾何意義的形式。

3.1 輸出梯度
對(duì)于每個(gè) i i，關(guān)于 logits 的交叉熵梯度為：

3.2 關(guān)于值的梯度

3.3 關(guān)于注意力權(quán)重的梯度

定義一個(gè)優(yōu)勢(shì)量是很方便的，其符號(hào)的選擇與梯度下降方向一致：

4 耦合動(dòng)力學(xué)與專業(yè)化

我們現(xiàn)在展開第 3 節(jié)中梯度流的含義，重點(diǎn)關(guān)注路由（通過(guò)分?jǐn)?shù)和注意力）與內(nèi)容（通過(guò)值）之間的相互作用。

4.1 基于優(yōu)勢(shì)的注意力重分配

公式 (23) 表明，對(duì)于固定的查詢 i i，定義一個(gè)符號(hào)與梯度下降方向一致的優(yōu)勢(shì)量是方便的：

4.2 值更新作為責(zé)任加權(quán)的原型
定義第 j j 列的注意力加權(quán)上游信號(hào)：

4.4 反饋循環(huán)與專業(yè)化

4.5 幾何示意圖

5 類 EM 的雙時(shí)間尺度動(dòng)力學(xué)

上述推導(dǎo)出的耦合動(dòng)力學(xué)與期望最大化（EM）算法存在一種有用的類比，但并非對(duì)顯式潛在變量似然函數(shù)的直接優(yōu)化，而是梯度流與責(zé)任加權(quán)更新之間的機(jī)制性對(duì)應(yīng)。注意力權(quán)重表現(xiàn)得如同對(duì)潛在源的責(zé)任分配，而值向量則作為在這些責(zé)任下更新的原型。與經(jīng)典 EM 不同，此處的更新由上游梯度驅(qū)動(dòng)，而非觀測(cè)數(shù)據(jù)，并且沒有獨(dú)立優(yōu)化關(guān)于值的似然函數(shù)。

5.1 注意力作為責(zé)任

5.3 近似 EM 與 SGD 的比較

在經(jīng)典 EM 中，E 步和 M 步是分離的：首先固定參數(shù)重新計(jì)算責(zé)任，然后固定責(zé)任更新參數(shù)。而在使用 SGD 訓(xùn)練的 Transformer 中，這些步驟是交織且嘈雜的，但其一階圖景仍類似于 EM：

在第 7.2 節(jié)中，我們將此類 EM 風(fēng)格的調(diào)度方案與標(biāo)準(zhǔn) SGD 在粘性馬爾可夫鏈任務(wù)上進(jìn)行比較，發(fā)現(xiàn)兩者均收斂至相似解，但 EM 風(fēng)格的更新能更快地達(dá)到低損失和尖銳、聚焦的注意力。

5.5 貝葉斯視角與 EM 視角的對(duì)比

EM 是一種優(yōu)化過(guò)程：它產(chǎn)生一個(gè)最大化（后驗(yàn)）似然的點(diǎn)估計(jì) θ ? 。而完整的貝葉斯處理方式則應(yīng)對(duì) θ 進(jìn)行積分，但這對(duì)于 Transformer 來(lái)說(shuō)是不可行的。因此，我們的分析停留在 EM/SGD 層面。

然而，至關(guān)重要的是，我們的配套工作 [1] 表明，以這種方式學(xué)習(xí)到的點(diǎn)估計(jì)參數(shù)支持表征空間中的貝葉斯計(jì)算：值流形、鍵框架和查詢軌跡在上下文中實(shí)現(xiàn)了貝葉斯信念更新。本文解釋了為什么交叉熵和梯度下降會(huì)自然地形成這些結(jié)構(gòu)。

6 從梯度流到貝葉斯流形

我們現(xiàn)在將上述推導(dǎo)出的梯度動(dòng)力學(xué)與在貝葉斯風(fēng)洞實(shí)驗(yàn)和生產(chǎn)模型中觀察到的幾何結(jié)構(gòu)聯(lián)系起來(lái)。

6.1 值流形展開

在風(fēng)洞實(shí)驗(yàn) [1] 中，我們觀察到：

• 訓(xùn)練早期，注意力熵下降，注意力聚焦于相關(guān)假設(shè)。
• 訓(xùn)練后期，注意力模式看似穩(wěn)定，但值表征沿一條平滑曲線展開；前幾個(gè)主成分解釋了大部分方差，且主軸與后驗(yàn)熵強(qiáng)相關(guān)。
• 即使注意力圖在視覺上保持不變，校準(zhǔn)誤差仍持續(xù)下降。

逐漸地，會(huì)沿著殘差誤差景觀的主方向?qū)R。在反復(fù)更新下，值向量最終落在由下游泛函（例如后驗(yàn)熵）參數(shù)化的低維流形上。

6.2 假設(shè)框架與鍵的正交性

如果不同的查詢子集持續(xù)地發(fā)現(xiàn)不同的鍵更有幫助，相應(yīng)的梯度貢獻(xiàn)會(huì)將這些鍵在 k k-空間中彼此推開，從而促進(jìn)不同假設(shè)軸之間的近似正交性。我們的“風(fēng)洞”論文正好測(cè)量了這種正交性，并將其與競(jìng)爭(zhēng)性假設(shè)的清晰分離聯(lián)系起來(lái)。

6.3 框架-精度分離

經(jīng)驗(yàn)觀察到的“框架-精度分離”——即注意力穩(wěn)定地定義一個(gè)假設(shè)框架，而校準(zhǔn)仍在持續(xù)改進(jìn)——現(xiàn)在很容易解釋：

因此，一個(gè)訓(xùn)練后期的 Transformer 具有固定的貝葉斯框架（假設(shè)軸與路由），但其后驗(yàn)幾何結(jié)構(gòu)仍會(huì)繼續(xù)變得更加銳利。

7 實(shí)驗(yàn)

我們現(xiàn)在通過(guò)受控模擬來(lái)說(shuō)明該理論。所有實(shí)驗(yàn)均使用一個(gè)單頭、單層注意力塊，不包含殘差連接或 LayerNorm，以保持動(dòng)力學(xué)過(guò)程清晰透明。

7.1 玩具注意力模擬

觀察結(jié)果： 在約 100 步內(nèi)，我們觀察到：

(1) 注意力熱圖逐漸變銳：每個(gè)查詢的注意力質(zhì)量集中于少數(shù)幾個(gè)位置（見圖 Figure 2、Figure 3）。

(2) 值向量在低維子空間中協(xié)同移動(dòng)；它們?cè)?PCA 投影下的軌跡顯示出新興的流形結(jié)構(gòu)（見圖 Figure 5）。

(3) 交叉熵?fù)p失平滑衰減（見圖 Figure 4），大部分收益發(fā)生在專業(yè)化現(xiàn)象出現(xiàn)之時(shí)。

7.2 粘性馬爾可夫鏈模擬：EM 與 SGD 對(duì)比

接下來(lái)，我們研究一個(gè)更具結(jié)構(gòu)的任務(wù)，在該任務(wù)中注意力可以利用時(shí)間持久性：一個(gè)基于符號(hào)的粘性馬爾可夫鏈。

要點(diǎn)： EM 風(fēng)格和 SGD 訓(xùn)練最終都收斂到相似的定性解：專業(yè)化的值向量和聚焦的注意力。然而，EM 風(fēng)格的調(diào)度方案以更少的步驟達(dá)到此狀態(tài)，并實(shí)現(xiàn)更銳利的專業(yè)化。這符合雙時(shí)間尺度的故事：責(zé)任分配（注意力）可被視為近似收斂，而閉式值更新則可利用這種穩(wěn)定性加速流形的形成。

8 實(shí)用影響與診斷方法

梯度分析為訓(xùn)練和解釋 Transformer 注意力機(jī)制提供了有用的診斷工具和設(shè)計(jì)原則。

8.1 診斷方法

8.2 正則化與穩(wěn)定性

• 值上的 LayerNorm 可穩(wěn)定范數(shù)，同時(shí)保持方向動(dòng)力學(xué)不變。
• 注意力 Dropout 會(huì)擾亂反饋循環(huán)，限制過(guò)度專業(yè)化，并鼓勵(lì)值被更均勻地使用。
• 學(xué)習(xí)率選擇 調(diào)節(jié)路由與內(nèi)容之間的時(shí)間尺度分離；較小的學(xué)習(xí)率使一階圖像更精確。

8.3 架構(gòu)選擇

• 多頭注意力 允許多個(gè)專業(yè)化的路由流形共存，減少單個(gè)頭內(nèi)部的競(jìng)爭(zhēng)。

• 深度 自然支持我們?cè)陲L(fēng)洞實(shí)驗(yàn)和大型模型中觀察到的綁定–消除–精煉層級(jí)結(jié)構(gòu)。
• 殘差連接 有助于在各個(gè)頭強(qiáng)烈專業(yè)化的同時(shí)，維持有用的中間表征。

9 相關(guān)工作

9.1 Transformer 的貝葉斯解釋

一些研究認(rèn)為，Transformer 通過(guò)行為方式或探針?lè)绞綄?shí)現(xiàn)了近似的貝葉斯推斷 [例如 7, 8]。我們的配套論文 [1] 在小型風(fēng)洞中展示了精確的貝葉斯行為和幾何特征，而一篇縮放論文則在生產(chǎn)級(jí)大語(yǔ)言模型中展示了類似的模式。本論文解釋了梯度動(dòng)力學(xué)如何產(chǎn)生這些幾何結(jié)構(gòu)。此外，本文還表明理論最小貝葉斯預(yù)測(cè)熵接近經(jīng)驗(yàn)熵，這一點(diǎn)在配套論文中有更詳細(xì)的探討。

9.2 機(jī)制性可解釋性

機(jī)制性可解釋性研究旨在識(shí)別執(zhí)行復(fù)制、歸納及其他算法任務(wù)的具體頭和電路 [3, 5]。我們的框架通過(guò)解釋專業(yè)化如何源于路由與內(nèi)容的相互作用，補(bǔ)充了這一方向，而非將專業(yè)化頭視為原始構(gòu)件。

9.3 優(yōu)化與隱式偏差

線性及深層網(wǎng)絡(luò)中梯度下降的隱式偏差已被廣泛研究 [2, 6]。我們將這些思想擴(kuò)展至注意力機(jī)制：梯度下降隱式偏好那些路由與誤差幾何對(duì)齊、且值位于支持貝葉斯更新的低維流形上的表征。

此處推導(dǎo)的責(zé)任加權(quán)值更新讓人聯(lián)想到神經(jīng) EM 和槽注意力模型，其中軟分配驅(qū)動(dòng)原型更新。關(guān)鍵區(qū)別在于，在 Transformer 中，責(zé)任是通過(guò)內(nèi)容可尋址的注意力計(jì)算的，而原型更新由反向傳播的誤差信號(hào)驅(qū)動(dòng)，而非重建似然。我們的重點(diǎn)并非提出一種新的 EM 風(fēng)格架構(gòu)，而是展示標(biāo)準(zhǔn)交叉熵訓(xùn)練在注意力層中如何作為梯度流的結(jié)果，自然誘導(dǎo)出類 EM 的專業(yè)化動(dòng)力學(xué)。

10 局限性與未來(lái)方向

我們的分析是有意保持簡(jiǎn)潔和受控的。

一階近似。 我們?cè)凇耙浑A”框架下工作，假設(shè)學(xué)習(xí)率較小，并忽略高階及隨機(jī)效應(yīng)（例如動(dòng)量、Adam、小批量噪聲）。將分析擴(kuò)展至更現(xiàn)實(shí)的優(yōu)化器是一個(gè)重要的下一步。

單頭單層聚焦。 我們分析的是一個(gè)孤立的單頭，不包含殘差路徑或 LayerNorm。多頭、多層動(dòng)力學(xué)——包括頭間協(xié)調(diào)與層級(jí)專業(yè)化——仍是開放問(wèn)題。

有限寬度 vs. 無(wú)限寬度。 我們并未明確將我們的分析與神經(jīng)切向核或無(wú)限寬度極限聯(lián)系起來(lái)。彌合這些機(jī)制可能有助于澄清 Transformer 何時(shí)在特征學(xué)習(xí)模式與懶惰訓(xùn)練模式下運(yùn)行。

大規(guī)模經(jīng)驗(yàn)驗(yàn)證。 我們的玩具模擬有意設(shè)計(jì)得規(guī)模很小。將第 8 節(jié)中的診斷工具應(yīng)用于全規(guī)模大語(yǔ)言模型（LLM）的訓(xùn)練過(guò)程，跟蹤優(yōu)勢(shì)矩陣和流形形成隨時(shí)間的變化，是一個(gè)有前景的方向。

11 結(jié)論

本文聚焦于最小化設(shè)定下的“一階”機(jī)制；配套工作已證實(shí)，相同的幾何結(jié)構(gòu)在更大規(guī)模上依然存在，并在受控風(fēng)洞實(shí)驗(yàn)和大型模型中支持精確的貝葉斯推斷。

我們的關(guān)鍵發(fā)現(xiàn)如下：

結(jié)合我們的風(fēng)洞實(shí)驗(yàn)和縮放研究論文，這形成了一個(gè)連貫的三部曲：優(yōu)化動(dòng)力學(xué)構(gòu)建貝葉斯幾何，而該幾何使 Transformer 能夠作為上下文內(nèi)的貝葉斯推理器運(yùn)作。

原文：https://arxiv.org/pdf/2512.22473