注意力的梯度動力學：交叉熵如何塑造貝葉斯流形——貝葉斯注意力三部曲之Paper II

Gradient Dynamics of Attention: How Cross-Entropy SculptsBayesian Manifolds

注意力梯度動力學：交叉熵如何塑造貝葉斯流形

https://arxiv.org/pdf/2512.22473

Transformer 在精心構建的“貝葉斯風洞”和大規模語言模型中，經驗上展現出精確的概率推理能力，但基于梯度的學習如何生成所需的內部幾何結構仍不透明。我們對交叉熵訓練如何重塑 Transformer 注意力頭中的注意力得分和值向量進行了系統性的一階分析。我們的核心結果是一個基于優勢（advantage-based）的注意力得分路由法則。

其中， u i
是位置 i i 處的上游梯度， α i j
是注意力權重。這些方程引發了一個正反饋循環，在該循環中路由和內容協同專業化：查詢更強烈地路由到那些對其誤差信號高于平均值的值，而這些值又被拉向使用它們的查詢。我們表明，這種耦合的專業化行為類似于一個雙時間尺度的 EM 過程：注意力權重實現 E 步（軟責任分配），而值實現 M 步（基于責任加權的原型更新），同時查詢和鍵調整假設框架。通過受控模擬（包括一個粘性馬爾可夫鏈任務，我們在其中將閉式 EM 風格更新與標準 SGD 進行比較），我們證明了最小化交叉熵的相同梯度動力學也會塑造我們在相關工作中識別出的低維流形，從而實現貝葉斯推斷。此外，在 EM 與 SGD 的比較中，EM 收斂得更快，且兩者的預測熵都接近最小貝葉斯熵。這產生了一幅統一的畫面：優化（梯度流）催生幾何結構（貝葉斯流形），而幾何結構反過來又支持功能（上下文內概率推理）。

1 引言

Transformer 已成為序列建模的主導架構，但我們仍缺乏對梯度下降如何塑造其內部表征的機制性理解。最近的研究表明，在受控的“貝葉斯風洞”中，小型 Transformer 可以精確再現解析后驗分布，其中鍵形成正交的假設軸，查詢實現漸進式信念更新，值則沿著由后驗熵參數化的一維流形展開 [1]。我們配套的縮放論文將這一圖景擴展到生產規模模型，展示了 Pythia、Phi-2 和 LLaMA 中類似的幾何特征。

這些發現提出了一個自然的問題：為什么單純的交叉熵訓練會產生貝葉斯推斷所需的幾何結構？理解這一點需要從靜態幾何轉向梯度動力學：注意力分數、查詢、鍵和值在優化過程中如何共同演化？

1.1 貢獻

我們的主要貢獻如下：

一個由責任加權的上游梯度平均值。這引發了一個正反饋循環：查詢路由到有助于它們的值；這些值則向其使用者移動，從而強化路由并促成專業化。

(4) 雙時間尺度 EM 解釋。我們表明，這些動力學實現了一種隱式的類 EM 算法：注意力權重充當軟責任分配（E 步），值作為在這些責任下更新的原型（M 步），而查詢/鍵則作為潛在分配模型的參數。注意力通常較早穩定，而值則持續細化——這種框架-精度分離與我們在風洞實驗和大型模型中的經驗觀察相符。

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

結合 [1]，我們的結果提供了一個統一的敘述：

梯度下降 ? 貝葉斯流形 ? 上下文內推理。

關于“貝葉斯推斷”的澄清。本文中，“貝葉斯推斷”指的是對潛在任務變量的貝葉斯后驗預測——而非對網絡權重的后驗。我們表明，交叉熵訓練塑造了幾何結構，該結構實現了對上下文內假設的貝葉斯濾波，而非貝葉斯權重不確定性。

2 設置與符號說明
我們分析一個作用于長度為 T 的序列的單頭注意力機制。除非另有說明，下標 i , j , k 的取值范圍均為 1 到 T 。

2.1向前傳球

2.2 輔助量
為簡潔起見，我們定義：

3 一階梯度推導
我們現在不跳過任何步驟地推導所有相關梯度，重點關注能揭示其幾何意義的形式。

3.1 輸出梯度
對于每個 i i，關于 logits 的交叉熵梯度為：

3.2 關于值的梯度

3.3 關于注意力權重的梯度

定義一個優勢量是很方便的，其符號的選擇與梯度下降方向一致：

4 耦合動力學與專業化

我們現在展開第 3 節中梯度流的含義，重點關注路由（通過分數和注意力）與內容（通過值）之間的相互作用。

4.1 基于優勢的注意力重分配

公式 (23) 表明，對于固定的查詢 i i，定義一個符號與梯度下降方向一致的優勢量是方便的：

4.2 值更新作為責任加權的原型
定義第 j j 列的注意力加權上游信號：

4.4 反饋循環與專業化

4.5 幾何示意圖

5 類 EM 的雙時間尺度動力學

上述推導出的耦合動力學與期望最大化（EM）算法存在一種有用的類比，但并非對顯式潛在變量似然函數的直接優化，而是梯度流與責任加權更新之間的機制性對應。注意力權重表現得如同對潛在源的責任分配，而值向量則作為在這些責任下更新的原型。與經典 EM 不同，此處的更新由上游梯度驅動，而非觀測數據，并且沒有獨立優化關于值的似然函數。

5.1 注意力作為責任

5.3 近似 EM 與 SGD 的比較

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

在第 7.2 節中，我們將此類 EM 風格的調度方案與標準 SGD 在粘性馬爾可夫鏈任務上進行比較，發現兩者均收斂至相似解，但 EM 風格的更新能更快地達到低損失和尖銳、聚焦的注意力。

5.5 貝葉斯視角與 EM 視角的對比

EM 是一種優化過程：它產生一個最大化（后驗）似然的點估計 θ ? 。而完整的貝葉斯處理方式則應對 θ 進行積分，但這對于 Transformer 來說是不可行的。因此，我們的分析停留在 EM/SGD 層面。

然而，至關重要的是，我們的配套工作 [1] 表明，以這種方式學習到的點估計參數支持表征空間中的貝葉斯計算：值流形、鍵框架和查詢軌跡在上下文中實現了貝葉斯信念更新。本文解釋了為什么交叉熵和梯度下降會自然地形成這些結構。

6 從梯度流到貝葉斯流形

我們現在將上述推導出的梯度動力學與在貝葉斯風洞實驗和生產模型中觀察到的幾何結構聯系起來。

6.1 值流形展開

在風洞實驗 [1] 中，我們觀察到：

• 訓練早期，注意力熵下降，注意力聚焦于相關假設。
• 訓練后期，注意力模式看似穩定，但值表征沿一條平滑曲線展開；前幾個主成分解釋了大部分方差，且主軸與后驗熵強相關。
• 即使注意力圖在視覺上保持不變，校準誤差仍持續下降。

逐漸地，會沿著殘差誤差景觀的主方向對齊。在反復更新下，值向量最終落在由下游泛函（例如后驗熵）參數化的低維流形上。

6.2 假設框架與鍵的正交性

如果不同的查詢子集持續地發現不同的鍵更有幫助，相應的梯度貢獻會將這些鍵在 k k-空間中彼此推開，從而促進不同假設軸之間的近似正交性。我們的“風洞”論文正好測量了這種正交性，并將其與競爭性假設的清晰分離聯系起來。

6.3 框架-精度分離

經驗觀察到的“框架-精度分離”——即注意力穩定地定義一個假設框架，而校準仍在持續改進——現在很容易解釋：

因此，一個訓練后期的 Transformer 具有固定的貝葉斯框架（假設軸與路由），但其后驗幾何結構仍會繼續變得更加銳利。

7 實驗

我們現在通過受控模擬來說明該理論。所有實驗均使用一個單頭、單層注意力塊，不包含殘差連接或 LayerNorm，以保持動力學過程清晰透明。

7.1 玩具注意力模擬

觀察結果： 在約 100 步內，我們觀察到：

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

(2) 值向量在低維子空間中協同移動；它們在 PCA 投影下的軌跡顯示出新興的流形結構（見圖 Figure 5）。

(3) 交叉熵損失平滑衰減（見圖 Figure 4），大部分收益發生在專業化現象出現之時。

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

接下來，我們研究一個更具結構的任務，在該任務中注意力可以利用時間持久性：一個基于符號的粘性馬爾可夫鏈。

要點： EM 風格和 SGD 訓練最終都收斂到相似的定性解：專業化的值向量和聚焦的注意力。然而，EM 風格的調度方案以更少的步驟達到此狀態，并實現更銳利的專業化。這符合雙時間尺度的故事：責任分配（注意力）可被視為近似收斂，而閉式值更新則可利用這種穩定性加速流形的形成。

8 實用影響與診斷方法

梯度分析為訓練和解釋 Transformer 注意力機制提供了有用的診斷工具和設計原則。

8.1 診斷方法

8.2 正則化與穩定性

• 值上的 LayerNorm 可穩定范數，同時保持方向動力學不變。
• 注意力 Dropout 會擾亂反饋循環，限制過度專業化，并鼓勵值被更均勻地使用。
• 學習率選擇 調節路由與內容之間的時間尺度分離；較小的學習率使一階圖像更精確。

8.3 架構選擇

• 多頭注意力 允許多個專業化的路由流形共存，減少單個頭內部的競爭。

• 深度 自然支持我們在風洞實驗和大型模型中觀察到的綁定–消除–精煉層級結構。
• 殘差連接 有助于在各個頭強烈專業化的同時，維持有用的中間表征。

9 相關工作

9.1 Transformer 的貝葉斯解釋

一些研究認為，Transformer 通過行為方式或探針方式實現了近似的貝葉斯推斷 [例如 7, 8]。我們的配套論文 [1] 在小型風洞中展示了精確的貝葉斯行為和幾何特征，而一篇縮放論文則在生產級大語言模型中展示了類似的模式。本論文解釋了梯度動力學如何產生這些幾何結構。此外，本文還表明理論最小貝葉斯預測熵接近經驗熵，這一點在配套論文中有更詳細的探討。

9.2 機制性可解釋性

機制性可解釋性研究旨在識別執行復制、歸納及其他算法任務的具體頭和電路 [3, 5]。我們的框架通過解釋專業化如何源于路由與內容的相互作用，補充了這一方向，而非將專業化頭視為原始構件。

9.3 優化與隱式偏差

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

此處推導的責任加權值更新讓人聯想到神經 EM 和槽注意力模型，其中軟分配驅動原型更新。關鍵區別在于，在 Transformer 中，責任是通過內容可尋址的注意力計算的，而原型更新由反向傳播的誤差信號驅動，而非重建似然。我們的重點并非提出一種新的 EM 風格架構，而是展示標準交叉熵訓練在注意力層中如何作為梯度流的結果，自然誘導出類 EM 的專業化動力學。

10 局限性與未來方向

我們的分析是有意保持簡潔和受控的。

一階近似。 我們在“一階”框架下工作，假設學習率較小，并忽略高階及隨機效應（例如動量、Adam、小批量噪聲）。將分析擴展至更現實的優化器是一個重要的下一步。

單頭單層聚焦。 我們分析的是一個孤立的單頭，不包含殘差路徑或 LayerNorm。多頭、多層動力學——包括頭間協調與層級專業化——仍是開放問題。

有限寬度 vs. 無限寬度。 我們并未明確將我們的分析與神經切向核或無限寬度極限聯系起來。彌合這些機制可能有助于澄清 Transformer 何時在特征學習模式與懶惰訓練模式下運行。

大規模經驗驗證。 我們的玩具模擬有意設計得規模很小。將第 8 節中的診斷工具應用于全規模大語言模型（LLM）的訓練過程，跟蹤優勢矩陣和流形形成隨時間的變化，是一個有前景的方向。

11 結論

本文聚焦于最小化設定下的“一階”機制；配套工作已證實，相同的幾何結構在更大規模上依然存在，并在受控風洞實驗和大型模型中支持精確的貝葉斯推斷。

我們的關鍵發現如下：

結合我們的風洞實驗和縮放研究論文，這形成了一個連貫的三部曲：優化動力學構建貝葉斯幾何，而該幾何使 Transformer 能夠作為上下文內的貝葉斯推理器運作。

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