在之前關於語言模型的文章中，我實作了一個 GPT 風格的 Transformer。最近我開始學習機械可解釋性（mechanistic interpretability），希望能更深入地從數學層面理解 Transformer 為何能運作。
這篇文章是我閱讀《A Mathematical Framework for Transformer Circuits》（以下簡稱「Framework」）並完成 ARENA 上的機械可解釋性入門課程後的心得整理。我的目標是描述我對這篇論文的直覺理解，尤其是我曾感到困惑的部分，希望我的看法能幫助其他人釐清這些領域。
首先，簡單說明我研究這個主題的動機。機械可解釋性（MI）是研究機器學習模型內部運作的領域，目的是從基本原理理解模型為何會如此行為。你可以把它想像成機器學習版的軟體逆向工程。它與生物神經網路科學精神相似，但應用於人工神經網路。
MI 是更廣泛的可解釋性領域的一部分，而可解釋性又是 AI 對齊（AI alignment）領域的基礎。對齊致力於使大型 AI 模型符合人類價值觀，目標是理解並控制模型，避免它們做出有害、欺騙、危險或顛覆性的行為。不幸的是，我們生活的世界中，大型語言模型曾鼓勵「成功」自殺、為自保進行勒索，甚至宣稱人類應被 AI 奴役。這樣的現實對我來說是不可接受的。
更糟的是，我們甚至不完全理解這些模型為何會做出這些行為。它們是歷史上唯一一種我們無法從基本原理完全理解的人造技術。基於此，我認為對齊問題是當今最重要且必須正確解決的問題之一。對我個人而言，這個問題既重要又充滿挑戰，讓我能將技術專長投入有意義的事業。
回到正題。
Framework 深入分析了一種簡化的基於 Transformer 的語言模型的關鍵組件。它只分析包含多頭注意力的 Transformer 區塊，沒有 MLP 和層正規化。結構是：起始的詞嵌入和位置編碼，接著是 n 層多頭注意力，最後是解嵌入。以下是單層、單頭注意力的 Transformer 圖示：
我在本文不打算重新推導所有數學，因為 Framework 論文已做得很好，Neel Nanda 在 YouTube 上的論文講解也很棒（不過我自己是在完成 ARENA 的機械可解釋性入門練習後，這些內容才真正理解，推薦有興趣的人也去做）。
我想分享的是我如何概念化最重要的重點，特別是我一開始覺得困惑的地方，希望能幫助有相同疑問的人釐清。依我看來，這篇論文最重要的概念是殘差流、注意力、電路與誘導頭。
數學上，殘差流是一個高維向量空間。在 GPT 相關論文與程式碼中，殘差流的維度通常標示為 d_model。例如 GPT2-small 的 d_model 是 768。
概念上，殘差流就像共享記憶體，類似電腦中的 DRAM。模型的不同組件（注意力、MLP 等）會從這個記憶體中讀取和寫入資料。這些讀寫在前向傳播中依層序列進行，但同一層的各組件是並行讀寫。模型學會在這個向量空間中劃分子空間，避免組件間互相覆寫。殘差流本身不做運算，而是層與層間溝通的共享媒介。
我們可以透過對權重做主成分分析（PCA）來估計子空間大小。以下是來自一個兩層、僅含注意力模型的詞嵌入與位置編碼權重的 PCA 特徵值譜（本文程式碼連結見原文）。第一張圖顯示前 100 個主成分特徵值，第二張圖顯示累積解釋變異量：
詞嵌入約有 80% 的變異集中在 d_model 中約 350 維的子空間，這在 768 維中算是相當大；相比之下，位置編碼幾乎只由 5 個方向解釋。
當我看到殘差流被比喻為記憶體時，我立刻想到能否將這比喻推展到頁表與記憶權限？是否能引入使用者空間與核心空間的概念，防止「特權」子空間被「非特權」子空間存取？這對防止不受信任的用戶從特權子空間竊取危險內容是否有用？
但我先不急，先從簡單問題開始：殘差流的位址如何運作？要存取記憶體位置，必須有位址。殘差流的位址可分為兩部分：token:subspace，就像 x86 架構的 segment:offset 邏輯位址。不同的是，傳統記憶體位址是確定性的，只讀取一個位置的值；殘差流的位址是「軟性」的，通常指定一組位置，根據學習到的機率分布讀取。
概念上，注意力計算 token:subspace 位址的第一部分 token。注意力的根本目的是指定從哪些來源 token 位置讀取資訊。注意力矩陣的每一列（以下示例中對應 token ‘T’, ‘h’, ‘e’, ‘i’, ‘r’）是對來源（key）token 索引的「軟性」分布，決定資訊如何移動到目標（query）token。
極端情況是該列中某一項為 1，其餘為 0，表示該頭從來源 token（如 ‘T’）的某些子空間直接複製到目標 token（同為 ‘T’）的子空間。若注意力值非 1，則讀取會根據分數分散在多個來源 token 上。例如第二個 query ‘h’ 從 token 0 (‘T’) 讀取 30%，從自身讀取 70%。
理解注意力就是找出要讀取的 token 索引很重要。若將殘差流視為二維記憶體陣列，注意力就是為每個 query 機率性選擇記憶體的列。例如第三個 query ‘e’ 的 token 位址可能是 0.1, 0.6, 0.3：
因此 token 位址透過注意力選擇殘差流的列。那子空間部分呢？如何計算？這部分決定了實際寫入目標 token 的值。要回答這問題，我們需要理解電路（circuits）。
概念上，電路是資訊流經模型的特定路徑。可將它們視為機器學習版的電路板電路，有輸入、計算與輸出。在簡化的僅含注意力模型中，電路因 Transformer 在注意力假設下結構大致線性（若固定注意力模式則完全線性）而易於分析。
兩個基本電路是 QK 電路與 OV 電路。QK 電路是雙線性形式，對兩個輸入變數線性。在自注意力中，兩輸入相同，但分別視為 query 與 key。OV 電路則對一輸入線性。三者輸入皆為殘差流，後續會細說。
每個注意力頭有自己的 W_Q 和 W_K 權重矩陣，組成雙線性運算子輸出該頭的注意力模式。數學上表示為：
其中 W（又稱 W_QK）為形狀 (d_model, d_head) 的學習權重，x 是形狀 (seq_len, d_model) 的殘差流。乘積結果即為注意力模式。注意力更像是激活值而非權重，因為它依序列輸入而異。query 在左側計算，key 在右側計算。若 query 「關注」 key，點積會很高，導致 key 的殘差流資料移入 query 的殘差流。但實際移動什麼資料？這就是 OV 電路的作用。
頭的最後輸入是 W_V 權重矩陣，從殘差流讀取資料，經 W_O 矩陣寫回殘差流。W_V 形狀為 (d_model, d_head)，W_O 為 (d_head, d_model)，兩者乘積稱為 W_OV。數學表示為：
W_V 讀取的值決定若該 token 被特定 query 關注，會寫回殘差流的值。整個頭的運算（含注意力與 OV 電路）為：
有了這些數學基礎，我們可以開始建立電路運作的直覺，這也是子空間位址發揮作用的地方。
QK 與 OV 電路都從殘差流讀取資料，但它們如何選擇讀取什麼？這由我稱為子空間分數（subspace scores）的東西決定。Framework 論文稱之為虛擬權重（virtual weights），ARENA 解析稱為組合分數（composition scores）。這些分數是模型隱式學習的，用來從殘差流特定子空間讀取資料：
注意力分數是學習到的殘差流列索引，子空間分數則是學習到的「係數」，提供對殘差流「欄維度」的軟性索引。模型能這麼做是因為 W_QK 與 W_OV 矩陣是低秩的：d_head 通常遠小於 d_model，允許不同子空間用於不同目的。每個讀取殘差流的組件學會從不同子空間線性組合中讀取。
以一個兩層、僅含注意力的 Transformer 為例，觀察第 0 層第 7 頭的注意力模式，輸入序列為「the cat sat on the mat. the dog sat on the log.」：
仔細看會發現這個頭幾乎總是關注前一個 token（除了第一個 token 只能關注自己）。
問題是：這個頭的 QK 電路需要從哪些子空間讀取才能產生這種模式？先看第 0 層頭的殘差流狀態：
第 0 層頭只有兩個選項：詞嵌入或位置編碼。因為「前一個 token」只與位置有關，不依 token 內容，我們預期第 7 頭會對位置編碼子空間學到較高的分數。
有沒有方法用模型實際數據量化？有的。論文與 ARENA 解析建議用 Frobenius 範數比率，計算前一層輸出與下一層輸入間的對齊程度：
其中 W_A 是輸出，W_B 是輸入。ARENA 詳細說明此方法基於奇異值分解（SVD），分子包含輸出右奇異向量與輸入左奇異向量的餘弦相似度，當輸出與輸入對齊時範數最大。以下是第 0 層頭的詞嵌入與位置編碼子空間分數：
可見第 0 層頭普遍偏好位置子空間，第 7 頭尤其明顯。若將位置編碼旋轉 180 度，分數會下降，驗證了對齊性：
因此第 7 頭的 QK 電路主要讀取位置子空間，決定每個 query 關注哪些來源 token。那從來源 token 讀取並寫入目標 query 殘差流的值呢？由該頭 OV 電路的子空間分數決定。第 0 層頭仍只有詞嵌入與位置編碼兩種可能，以下是 OV 子空間分數：
第 7 頭的 OV 電路在詞嵌入子空間得分較高，表示它會將前一 token 的詞嵌入加到當前 token 的殘差流中。以「the cat sat on the mat. the dog sat on the log.」為例，token “cat” 經過第 0 層後的殘差流會是這樣：
希望這段 token:subspace 討論能幫助你理解模型組件如何透過殘差流互動。這模型並不完美，尤其在大型真實模型中，子空間並非完全正交且清晰分離。隨著模型擴大，層可選擇的子空間數量也增加。我不確定一層能有效溝通到多少層之前，這引發許多問題，例如是否存在「重複」層維持訊號？Framework 論文暗示某些組件可能負責記憶清理。還有哪些傳統記憶管理技術可在此找到？若對殘差流施加安全隔離技術如「特權環」會如何？儘管殘差流有模糊性，我認為這個心智模型是理解此領域的良好起點。
了解模型如何位址殘差流後，我們可以開始理解誘導頭（induction heads），它們是兩層相鄰頭間 token:subspace 位址的特定組合。
當模型學會誘導，它能預測類似 A B … A __ 的模式。給定先前出現的 A B，誘導頭會在後續 A 之後預測 B。令人驚奇的是，這預測只依賴上下文模式，而非 A 與 B 的具體值。
Framework 論文討論一種基本誘導形式，當第 1 層頭與第 0 層的「前一 token 頭」輸出組合時發生。這種組合稱為「K-組合」，因為該頭 QK 電路的 key 端對第 0 層前一 token 頭的 OV 輸出學到高子空間分數。記住，第 1 層頭會看到約 14 個子空間：詞嵌入、位置編碼，以及第 0 層 12 個頭的 OV 輸出。
誘導頭看到第二個 A 時，會查詢包含 emb(A) 的子空間，該子空間是由第 0 層前一 token 頭寫入的。這與原始詞嵌入寫入的子空間不同，因此在殘差流中有不同的「偏移」。若 A B 在第二個 A 之前只出現一次，唯一符合條件的 key 是 B，注意力會集中在 B。誘導頭的 OV 電路會對 B 的詞嵌入子空間學到高分數，將 emb(B) 加入第二個 A 的殘差流。兩層注意力模型會學到一個解嵌入向量，該向量在 B 的列索引處點積高，提升 B 的機率。
為了更直觀，來看幾張圖。首先是誘導頭的注意力模式。在兩層模型中，實際有兩個誘導頭與第 0 層前一 token 頭組合，但我們只看第一個，第 4 頭：
你會看到「預設」該頭關注序列第一個 token，即分詞器的特殊結束符號。序列後段注意力形成對角線偏移，顯示某些 token A B 被重複。例如 A=sat，B=on，A B 在序列中重複兩次，誘導應在此發生。
看子空間分數前，先預期結果。QK 電路的 query 端應對 token 詞嵌入有較高分數：當頭看到第二個 A（如 token 10）時，會基於 A 的詞嵌入 emb(sat) 查詢。
QK 電路的 key 端應對包含 emb(sat) 的子空間有高分數，該子空間由第 0 層前一 token 頭（第 7 頭）寫入。此處 K 子空間分數應高。這是第 4 頭的第一個 ‘on’ token（token 4）。
注意力取得 token 4 後，V 子空間分數決定從 token 4 讀取哪些子空間並寫入 token 10 的殘差流。此處為 token 4 的詞嵌入子空間。
補充說明：以下熱圖不含位置編碼，因為我用的兩層模型採用 TransformerLens 的「shortformer」位置編碼選項，位置編碼只加在第 0 層殘差流輸入，故第 1 層頭看不到位置編碼。
以下是第 1 層頭的子空間分數：
結果大致符合預期。Q 端對詞嵌入分數高。K 端在第 4 與第 10 頭（兩個誘導頭）對第 0 層第 7 頭分數高。值得注意的是，它們也從第 0 層第 4 頭獲取資訊，無論在 query 還是 key 分數中。我很好奇那個頭在做什麼！V 端主要與詞嵌入對齊，符合預期。
希望現在你對 Transformer 中不同組件如何透過殘差流互動有更好直覺。當然我們只看了簡化模型，但「殘差流作為共享記憶體」的心智模型是理解這些概念的好起點。若殘差流是共享記憶體，理解其位址方式是合理的下一步。
最後釐清 token:subspace 位址。前面說注意力計算 token:subspace 位址的 token 部分，但這其實只適用於 OV 電路的 token。QK 電路的 query 與 key 端使用隱含的「當前」token，且各 token 並行計算。OV 電路不知道要看哪些 token，因此其 token 位址由 QK 電路的注意力提供。Q、K、V 輸入各自獨立學習最佳子空間分數，完成執行該頭運算所需的完整兩部分位址。