如果能以 3D 方式顯示矩陣乘法的執行過程，當年學習矩陣乘法時也就不會那麼吃力了。
現在，矩陣乘法已經成為機器學習模型的建構模組，是各種強大AI 技術的基礎，了解其執行方式必然有助於我們更深入地理解這個AI 以及這個日趨智能化的世界。
這篇來自 PyTorch 部落格的文章將介紹一種用於矩陣乘法和矩陣乘法組合的視覺化工具 mm。
因為mm 使用了所有三個空間維度，所以相比於通常的二維圖表，mm 有助於直覺化的展示和激發思路，所使用的認知開銷也更小，尤其是（但不限於）對於擅長視覺和空間思考的人。
而且如果有三個維度來組合矩陣乘法，再加上載入已訓練權重的能力，mm 能可視化大型複合表達式（如注意力頭）並觀察它們的實際行為模式。
mm 是完全互動式的，運行在瀏覽器或筆記本iframe 中，並且其會將完整狀態保存在URL 中，因此連結就是可共享的會話（本文中的截圖和影片都有一個鏈接，可在該工具中打開對應的可視化，具體請參看原博客）。本參考指南會描述所有可用的功能。
工具網址：https://bhosmer.github.io/mm/ref.html
部落格原文：https://pytorch .org/blog/inside-the-matrix
本文首先會介紹視覺化方法，透過視覺化一些簡單矩陣乘法和表達式來建立直覺，然後再深入一些擴展範例：
介紹：為什麼這種視覺化方式比較好？
暖身：動畫－ 檢視規範的矩陣乘法分解的工作過程
暖身：表達式－ 速覽一些基本的表達式建構模組
深入註意力頭：透過NanoGPT 深度觀察GPT-2 的一對注意力頭的結構、值和計算行為
並行化注意力：使用來自近期的Blockwise Parallel Transformer 論文中的範例來視覺化注意力頭的並行化。
注意力層的大小：當我們將整個注意力層視覺化為單一結構，則注意力層的 MHA 半邊和 FFA 半邊合起來是什麼樣子？在自回歸解碼過程中，其影像會發生怎樣的變化？
LoRA：對此註意力頭架構的詳細闡釋的可視化解釋
1 介紹
mm 的視覺化方法基於這個前提：矩陣乘法本質上是一種三維運算。
換句話說：
其實可以描繪成這種形式：
當我們以這種方式將矩陣乘法包裹在一個立方體中時，參數形狀、結果形狀和共享維度之間的正確關係就全部就位了。
現在矩陣乘法計算就有了幾何意義：結果矩陣中的每個位置i,j 都錨定了一個沿著立方體內部的深度（depth）維度k 運行的向量，其中從L 的第i 行延伸出來的水平面與從R 的第j 列延伸出來的垂直面相交。沿著這一向量，來自左邊參數和右邊參數的成對的(i, k) (k, j) 元素會相遇並相乘，再沿k 對所得積求和，所得結果放入結果的i, j 位置。
這就是矩陣乘法的直覺意義：
1. 將兩個正交矩陣投影到一個立方體的內部；
2. 將每個交叉點的一對值相乘，得到一個乘積網格；
3. 沿第三個正交維度進行求和，以產生結果矩陣。
對於方向，工具會在立方體內部顯示一個指向結果矩陣的箭頭，其中藍色箭羽來自左側參數，紅色箭羽來自右側參數。該工具還會顯示白色指示線來指示每個矩陣的行軸線，儘管這些線在此截圖中很模糊。
其佈局約束條件簡單又直接：
左側參數和結果必須沿著它們共享的高度(i) 維度鄰接
右側參數和結果必須沿著它們共享的寬度(j) 維度鄰接
左側參數和右側參數必須沿它們共享的（左寬度/ 右高度）維度鄰接，此維度成為矩陣乘法的深度(k) 維度
這種幾何表示方法能為可視化所有標準的矩陣乘法分解提供堅實的基礎，並能為探索非平凡的複雜矩陣乘法組合提供直覺的基礎，接下來我們就能看到這一點。
2 暖身：動畫
在深入介紹更複雜的範例之前，我們先來看看這個視覺化風格看起來是什麼樣的，從而建立起對該工具的直覺認知。
2a 點積
首先來看一個經典演算法 —— 透過計算對應左側行和右側列的點積來計算每個結果元素。從這裡的動畫可以看到，相乘的值向量掃過立方體內部，每次都在對應位置提交一個求和後的結果。
這裡，L 具有填充 1（藍色）或 -1（紅色）的行塊；R 具有類似填充的列塊。這裡 k 是 24，所以結果矩陣 (L @ R) 的藍色值為 24，紅色值為 -24。 
2b 矩陣- 向量積
#分解為矩陣- 向量積的矩陣乘法看起來像垂直平面（左側參數與右側參數每一列的積），當它水平掃過立方體內部時，將列繪製到結果上：
觀察一個分解的中間值可能很有意思，即使範例很簡單。
舉個例子，請注意當我們使用隨機初始化的參數時，中間的矩陣- 向量積突出的垂直模式—— 這反映了一個事實：每個中間值都是左側參數的列縮放的副本：
2c 向量- 矩陣積
分解為向量- 矩陣積的矩陣乘法看起來像水平平面，其在向下穿過立方體內部時將行繪製到結果上：
#切換成隨機初始化的參數，可以看到類似矩陣- 向量積的模式—— 只不過這次是水平模式，對應的事實是每個中間向量- 矩陣積都是右側參數的行縮放的副本。
在思考矩陣乘法如何表示其參數的秩和結構時，一個有用的做法是設想這兩種模式在計算中同時發生：
這裡還有另一個使用向量- 矩陣積來建構直覺的範例，其中展示了單位矩陣的作用就像是一面呈45 度角擺放的鏡子，反射著其對應參數和結果： 
2d 對外積求和
第三次平面分解是沿著k 軸，透過向量外積逐點求和來計算矩陣乘法結果。這裡我們可以看到外積平面「從後到前」掃過立方體，累積到結果中：
使用隨機初始化的矩陣進行此分解，我們不僅可以看到數值，也可以看到結果中的秩累積，因為每個秩為1 的外積都被加到其中。
這也從直覺上說明了為什麼「低秩因式分解」（即透過構造參數在深度維度上較小的矩陣乘法來近似矩陣）在被近似的矩陣為低秩矩陣時的效果最好。這是後面會提到的LoRA：
3 暖身：表達式
我們可以怎樣的方式將這種視覺化方法擴展用於矩陣乘法的分解？先前的範例視覺化的是矩陣 L 和 R 的單次矩陣乘法 L @ R，但要是 L 和 / 或 R 本身也是矩陣乘法呢？
事實證明這種方法可以很好地擴展用於複合表達式。關鍵規則很簡單：子表達式（子）矩陣乘法是另一個立方體，其受到與父矩陣乘法一樣的佈局約束；子矩陣乘法的結果面同時也是父矩陣乘法對應的參數面，就像是共價共享的電子。
在這些限制限制中，我們可以依照自己的需求排布子矩陣乘法的各個面。這裡使用該工具的預設方案，這會產生交替的凸面和凹面立方體 —— 這種佈局的實踐效果很好，可以最大化地利用空間，同時盡可能減少遮擋。 （但是佈局是完全可自訂的，詳情訪問mm 工具頁面。）
這一節將可視化機器學習模型的一些關鍵構建模組，以便讀者熟悉這種視覺表示並從中獲得新的直覺認識。
3a 左結合表達式
下面將會介紹兩個(A @ B) @ C 形式的表達式，每一個都有自己的獨特形狀和特徵。 （注意：mm 遵循矩陣乘法是左結合的約定，所以(A @ B) @ C 可簡單寫為A @ B @ C。）
首先為A @ B @ C 賦予很有特點的FFN 形狀，其中「隱藏維度」比「輸入」或「輸出」維度寬。 （具體來說，就此範例而言，這意味著B 的寬度大於A 或C 的寬度。）
和單次矩陣乘法範例一樣，浮動的箭頭指向結果矩陣，其中藍色箭羽來自左側參數，紅色箭羽來自右側參數。
而當 B 的寬度小於 A 或 C 的寬度時，對 A @ B @ C 的視覺化則會有一個瓶頸，類似自動編碼器的形狀。
交替的凹凸模組的模式還可以擴展成任意長度的鏈：例如這個多層瓶頸：
##3b 右邊結合表達式
接下來視覺化右邊結合表達式A @ (B @ C)。 與左結合表達式的水平擴展類似—— 可以說是從根表達式的左側參數發端，右結合表達式鍊是以垂直方式擴展，從根表達式的右側參數發端。 人們有時可以看到一個以右側結合形式形成的 MLP，即右側是柱狀輸入，權重層從右到左運行。使用上面描繪的二層FFN 範例的矩陣（適當轉置後），看起來會是這樣，C 現在是輸入，B 是第一層，A 是第二層：
3c 二元表達式
對於視覺化工具，要有用，就不能只用於簡單的教學式範例，也要能方便地用於更複雜的表達式。在真實世界用例中，一個關鍵性結構元件是二元表達式 —— 左側和右側都有子表達式的矩陣乘法。 這裡視覺化了此類表達式中形狀最簡單的一個(A @ B) @ (C @ D)：
3d一點註解：分區和平行性
完整闡述主題超出了本文的範圍，但後面我們會在註意力頭部分看到它的實際效用。但在這裡熱身，看兩個簡單範例，了解下這種視覺化風格可以如何讓對平行化複合表達式的推理非常直觀 —— 只需通過簡單的幾何分區。 第一個範例是將典型的「資料並行」分割區套用到上面的左結合多層瓶頸範例。我們沿著i 分區，對初始左側參數（「批」）和所有中間結果（「激活」）進行分段，但不對後續參數（「權重」）分段— 這種幾何結構使得表達式中的哪些參與者被分段以及哪些保持完整變得顯而易見：
4 深入註意力頭
現在來看看GPT-2 的注意力頭－ 具體來說是來自NanoGPT 的5 層第4 頭的「gpt2」(small) 配置（層數= 12，頭數= 12，嵌入數= 768），透過HuggingFace 使用了來自OpenAI 的權重。輸入啟動取自在含 256 個 token 的 OpenWebText 訓練樣本上一次前向通過。
這個特定的頭並無任何特殊之處，選擇它主要是因為其計算的是一個非常常見的注意力模式，並且它位於模型中部，其中激活已經變得結構化並顯示出一些有趣的紋理。
4a 結構
這個完整注意力頭被視覺化成了單一複合表達式，其始於輸入，終於投影的輸出。 （註：為了確保自足性，這裡按照Megatron-LM 的描述對每個頭執行輸出投影。）
這項計算包含六次矩陣乘法：
Q = input @ wQ// 1K_t = wK_t @ input_t// 2V = input @ wV// 3attn = sdpa(Q @ K_t)// 4head_out = attn @ V // 5out = head_out @ wO // 6
簡單描述這裡在做什麼：
風車的葉片為矩陣乘法1、2、3 和6：前一組是輸入到Q、K 和V 的內投影；後者是從attn @ V 回到嵌入維度的外投影。
中心有兩個矩陣乘法；第一個計算的是注意力分數（後面的凸立方體），然後使用它們基於值向量得到輸出token（前面的凹立方體） 。因果關係意味著注意力分數形成一個下三角形。
但讀者最好能親自詳細探索這個工具，而不是只看截圖或下面的視頻，以便更詳細地理解—— 不管是其結構還是流過計算過程的實際值。
4b 計算和值
這裡是注意力頭的計算過程動畫。具體來說，我們是看
sdpa (input @ wQ @ K_t) @ V @ wO
（即上面的矩陣乘法1、4、5 和6，其中K_t 和V 已經預先計算）是作為向量- 矩陣積的融合鏈來計算：序列中的每一項都在一步之內從輸入穿過注意力到輸出。後面關於並行化的部分會提到更多關於這個動畫的選擇，但我們先看看計算的值能告訴我們什麼。 
我們可以看到很多有趣的東西：
在討論注意力計算之前，可以看到低秩Q 和K_t 的形態是多麼驚人。放大Q @ K_t 向量- 矩陣積動畫，看起來會更加生動：Q 和K 中大量通道（嵌入位置）在序列中看起來或多或少是恆定的，這意味著有用的注意力信號可能僅由一小部分嵌入驅動。理解和利用這種現像是 SysML ATOM transformer 效率專案的一部分。
也許人們最熟悉的是注意力矩陣中出現的強大但不完美的對角線。這是一種常見模式，出現在該模型（以及許多 Transformer）的許多注意力頭中。它能產生局部注意力：緊鄰輸出 token 位置之前的小鄰域中的值 token 很大程度上決定了輸出 token 的內容模式。
然而，這個鄰域的大小和其中各個token 的影響變化很大—— 這可以在註意力網格中的非對角frost 中看到，也能在注意力矩陣沿著序列下降時attn [i] @ V 向量- 矩陣積平面的波動模式中看到。
但請注意，局部鄰域並不是唯一值得注意的東西：注意力網格的最左列（對應於序列的第一個token ）完全填充了非零（但波動）的值，這表示每個輸出token 都會受到第一個值token 一定程度的影響。
此外，目前 token 鄰域和初始 token 之間的注意力分數主導性存在不精確但可辨別的振盪。此振蕩的週期各有不同，但一般來說，一開始很短，然後沿著序列向下移動而變長（類似地，在給定因果關係的情況下，與每一行的候選注意力token 的數量相關）。
為了了解 (attn @ V) 的形成方式，不單獨關註註意力是很重要的 ——V 也同等重要。每個輸出項都是整個 V 向量的加權平均值：在註意力是完美對角線的極端情況下，attn @ V 只是 V 的精確副本。這裡我們看到更有紋理的東西：可見的帶狀結構，其中特定token 在註意力行的連續子序列上的得分很高，疊加在與V 明顯相似的矩陣上，但由於對角線較粗而有一些垂直遮擋。 （附註：根據 mm 參考指南，長按或按住 Control 鍵單擊將顯示可視化元素的實際數值。）
請記住，由於我們位於中間層（5 層），因此該注意力頭的輸入是一個中間表示，而不是原始 token 化文字。因此，在輸入中看到的模式本身就發人深省 —— 特別是，強大的垂直線條是特定的嵌入位置，其值在序列的長段上統一具有高的幅度 —— 有時幾乎是佔滿了。
但有趣的是，輸入序列中的第一個向量是獨特的，不僅打破了這些高幅度列的模式，而且幾乎在每個位置都攜帶著非典型值（附註：這裡沒有視覺化，但這種模式會反覆出現在多個樣本輸入上）。
注意：關於最後兩個要點，值得重申的是，這裡視覺化的是單一樣本輸入的計算。在實踐中，可以發現每個頭都有一個特徵模式，能在相當多的樣本集合上一致（儘管不等同）地表達，但當查看任意包含激活的可視化時，需要記住：輸入的完整分佈可能會以微妙的方式影響它激發的想法和直覺。
最後，再次建議直接探索動畫！
4c 注意力頭有很多有趣的不同之處
繼續之前，這裡再透過一個演示展現簡單地研究模型以了解其詳細工作方式的有用性。
這是 GPT-2 的另一個注意力頭。其行為模式與 5 層第 4 頭大有不同 —— 這符合人們預期，畢竟它位於模型一個非常不同的部分。這個頭位於第一層：0 層的第2 頭：
值得注意的點：
這個頭的注意力分佈很均勻。這會產生一個效果：將V 的相對未加權的平均值（或者說V 的合適的因果前綴）交到attn @ V 的每一行；如動畫所示：當我們向下移動注意力分數三角時， attn [i] @ V 向量- 矩陣積有很小的波動，而不是簡單的V 的縮小比例的、逐漸揭示的副本。
attn @ V 具有驚人的垂直均勻性 —— 在嵌入的大柱狀區域中，相同的值模式在整個序列中持續存在。人們可以將這些視為每個 token 共享的屬性。
附註：一方面，考慮到注意力分佈非常均勻的效果，人們可能會期望 attn @ V 具有一定的一致性。但每一行都是由 V 的因果子序列而不是整個序列構成 —— 為什麼這不會導致更多的變化，就像沿著序列向下移動時的漸進變形一樣？透過視覺檢查可知，V 沿其長度並不均勻，因此答案必定在於其值分佈的一些更微妙的屬性。
最後，在外投影後，這個頭的輸出在垂直方向上還要更均勻。
我們可以得到一個強烈的印象：這個注意力頭傳遞的大部分資訊由序列中每個 token 共享的屬性組成。其輸出投影權重的構成能強化這種直覺。
總的來說，我們不由得會想：這個注意力頭產生的極其規則、高度結構化的資訊可能是透過稍微… 不那麼奢華的計算手段獲得的。當然，這不是一個未經探索的領域，但視覺化計算訊號的明確性和豐富性對於產生新想法和推理現有想法都非常有用。
4d 重返介紹：免費的不變性
回頭看，需要重申：我們之所以能夠將注意力頭等非平凡的複合操作視覺化並讓它們保持直觀，是因為重要的代數性質（例如參數形狀的限制方式或哪些並行軸與哪些操作相交），這些性質不需要額外的思考：它們直接來自可視化對象的幾何屬性，而不是需要記住的額外規則。
舉個例子，在這些注意力頭視覺化中，可以明顯看出：
Q 和attn @ V 的長度一樣，K 和V 的長度一樣，這些配對的長度都彼此獨立；
Q 和K 的寬度一樣，V 和attn @ V 的寬度一樣，這些配對的寬度都彼此獨立。
這些結構在構造上就是真實的，就是結構組分位於複合結構的哪個部分以及它們的方向如何的簡單結果。
這種「免費性質」的優勢在探索典型結構的變體時特別有用—— 一個明顯的例子是一次解碼一個自回歸token 中的單行高的注意力矩陣：
5 並行化注意力
上面5 層第4 頭的動畫視覺化了注意力頭中6 個矩陣乘法中的4 個。
它們被視覺化為了一條向量 - 矩陣積的融合鏈，從而證實了一個幾何直覺：從輸入到輸出的整個左結合鏈沿共享 i 軸呈層狀，且可並行化。
5a 範例：沿著 i 分區
為了在實務上並行計算，我們可將輸入沿著 i 軸劃分為區塊。我們可以在該工具中可視化這種分區，透過指定將給定軸劃分為特定數量的區塊 —— 在這些範例中將使用 8，但該數字並無特別之處。
除此之外，這種視覺化清楚地表明，每次並行計算都需要完整的wQ（用於內投影）、K_t 和V（用於注意力）和wO（用於外投影），因為它們沿著這些矩陣的未分區維度與已分區矩陣相鄰接：
5b 範例：雙重分區
#這裡也給出沿著多個軸進行分區的範例。為此，這裡選擇可視化該領域一個近期的創新成果，即Block Parallel Transformer（BPT），其基於Flash Attention 等一些研究成果，請參閱論文：https://arxiv.org/pdf/2305.19370.pdf
首先，BPT 如上所述沿著i 進行分區－ 並且實際上也將序列的這種水平分區一直延伸到注意力層的另一半邊（FFN）。 （對此的視覺化將在後面展現。）
為了完全解決這個上下文長度問題，向MHA 添加第二個分區——注意力計算本身的分區（即沿Q @ K_t 的j 軸的分區）。這兩個分區一起可將注意力分成區塊構成的網格：
從這個視覺化可以清楚看到：
這種雙重分區能有效解決上下文長度問題，因為我們現在能以視覺劃分注意力計算中每次出現的序列長度。
第二次分區的「範圍」：根據幾何結構可以明顯看出，K 和 V 的內投影計算可與核心的雙矩陣乘法一起分區。
注意一個微妙細節：這裡的視覺暗示是我們還可以沿著k 並行化後續的矩陣乘法attn @ V 並以split-k 風格對部分結果求和，從而並行化整個雙矩陣乘法。但sdpa () 中的逐行softmax 增加了要求：在計算attn @ V 的相應行之前，每一行都要將其所有分段歸一化，這會在註意力計算和最終矩陣乘法之間添加一個額外的逐行步驟。
6 注意力層的大小
眾所周知，注意力層的前一半（MHA）由於其二次複雜度而有很高的運算需求，但後一半（FFN）也有自己的需求，這要歸因於其隱藏維度的寬度，其通常是模型嵌入維度的寬度的4 倍。可視化完整注意力層的生物量有助於建立關於該層兩半部如何相互比較的直覺認識。
6a 視覺化完整的注意力層
下面是一個完整的注意力層，前一半（MHA）位於後面，後一半（FFN）位於前面。同樣，箭頭指向計算的方向。
註：
此視覺化描繪的不是單一注意力頭，而是顯示了未切片的 Q/K/V 權重和圍繞中心雙重矩陣乘法的投影。當然這並沒有將完整的 MHA 運算忠實地視覺化出來 —— 但這裡的目標是更清楚地了解該層的兩半中的相對矩陣大小，而不是每半執行的相對計算量。 （此外，這裡的權重使用了隨機值而非真實權重。）
這裡使用的維度有所收縮以保證瀏覽器（相對）能帶得動，但比例保持一樣（來自NanoGPT 的small配置）：模型嵌入維度= 192（原本是768）、FFN 嵌入維度= 768（原本是3072）、序列長度= 256（原本是1024），儘管序列長度對模型沒有根本影響。 （從視覺上看，序列長度的變化將表現為輸入葉片寬度的變化，從而導致注意力中心大小和下游垂直平面高度的變化。）
6b 視覺化BPT 分區層
簡單地回顧一下Blockwise Parallel Transformer，這裡是在整個注意力層的語境中可視化BPT 的平行化方案（和上面一樣省略了各個頭） 。特別要注意，沿著i（序列區塊）的分區以怎樣的方式擴展通過MHA 和FFN 兩半邊：
6c 對FFN 進行分區
這個視覺化方法建議進行額外的分區，該分區與上面描述的分區正交－ 在註意力層的FFN 半邊，將雙重矩陣乘法(attn_out @ FFN_1) @ FFN_2 分開，首先沿j 進行attn_out @ FFN_1，然後沿著k 與FFN_2 執行後續的矩陣乘法。這種分區會對兩個 FFN 權重層進行切片，以減少計算中每個參與組分的容量要求，但代價是部分結果的最終求和。
下面是將此分割區方法套用到未分割區的注意力層的樣子：
#下面則是套用於以BPT 方式分割區的層的情況：
6d 視覺化一次一個token 解碼的過程
在自迴歸式的一次一個token 的解碼過程中，查詢向量由單一token 構成。你可以在頭腦中想像一下這種情況下的注意力層會是什麼樣子，這很有啟發性 —— 單一嵌入行穿過一個巨大的平舖的權重平面。
除了強調與激活相比權重的巨大性之外，這種觀點還能讓人想起這樣一個概念：K_t 和V 的功能類似於一個6 層MLP 中動態生成的層，儘管MHA本身的mux/demux 計算會使這種對應關係不精確：
#7 LoRA
##最近的LoRA 論文《LoRA : Low-Rank Adaptation of Large Language Models》描述了一種高效的微調技術，該技術基於這一思想：微調期間引入的權重δ 是低秩的。根據這篇論文，這「允許我們透過在適應過程中優化密集層變化的秩分解矩陣來間接地訓練神經網路中的一些密集層… 同時保持預訓練權重處於凍結狀態。」7a 基本概念
簡而言之，關鍵一步是訓練權重矩陣的因子而不是矩陣本身：用一個I x K 張量和K x J 張量的矩陣乘法來取代I x J 權重張量，其中要保證K 為一個較小值。 如果 K 夠小，則尺寸方面可能會有很大贏面，但也有權衡之處：降低 K 也會降低積可以表達的秩。這裡透過一個範例說明尺寸上的節省與對結果的結構化影響，這裡是隨機的128 x 4 左側參數和4 x 128 右側參數的矩陣乘法—— 即一個128 x 128 矩陣的秩為4 的分解。注意L @ R 中的垂直和水平模式：
7b 將LoRA 應用於注意力頭
LoRA 將此分解方法應用於微調過程的方式是：
#
以上是以3D視角洞悉矩陣乘法，這就是AI思考的樣子的詳細內容。更多資訊請關注PHP中文網其他相關文章！
