聲明:本文來自于微信公眾號(hào) 機(jī)器之心（ID:almosthuman2014），作者:機(jī)器之心，授權(quán)微新創(chuàng)想轉(zhuǎn)載發(fā)布。

如果能以3D 方式展示矩陣乘法的執(zhí)行過程，當(dāng)年學(xué)習(xí)矩陣乘法時(shí)也就不會(huì)那么吃力了。

現(xiàn)如今，矩陣乘法已經(jīng)成為機(jī)器學(xué)習(xí)模型的構(gòu)建模塊，是各種強(qiáng)大 AI 技術(shù)的基礎(chǔ)，了解其執(zhí)行方式必然有助于我們更深入地理解這個(gè) AI 以及這個(gè)日趨智能化的世界。

這篇來自 PyTorch 博客的文章將介紹一種用于矩陣乘法和矩陣乘法組合的可視化工具 mm。

因?yàn)?mm 使用了所有三個(gè)空間維度，所以相比于通常的二維圖表，mm 有助于直覺化的展示和激發(fā)思路，所使用的認(rèn)知開銷也更小，尤其是（但不限于）對(duì)于擅長視覺和空間思考的人。

而且如果有三個(gè)維度來組合矩陣乘法，再加上加載已訓(xùn)練權(quán)重的能力，mm 能可視化大型復(fù)合表達(dá)式（如注意力頭）并觀察它們的實(shí)際行為模式。

mm 是完全交互式的，運(yùn)行在瀏覽器或筆記本 iframe 中，并且其會(huì)將完整狀態(tài)保存在 URL 中，因此鏈接就是可共享的會(huì)話（本文中的截圖和視頻都有一個(gè)鏈接，可在該工具中打開對(duì)應(yīng)的可視化，具體請(qǐng)參看原博客）。本參考指南會(huì)描述所有可用的功能。

• 工具地址:https://bhosmer.github.io/mm/ref.html

• 博客原文:https://pytorch.org/blog/inside-the-matrix

本文首先會(huì)介紹可視化方法，通過可視化一些簡單矩陣乘法和表達(dá)式來建立直覺，然后再深入一些擴(kuò)展示例:

• 介紹:為什么這種可視化方式更好?

• 熱身:動(dòng)畫 —— 查看規(guī)范的矩陣乘法分解的工作過程

• 熱身:表達(dá)式 —— 速覽一些基本的表達(dá)式構(gòu)建模塊

• 深入注意力頭:通過 NanoGPT 深度觀察 GPT-2的一對(duì)注意力頭的結(jié)構(gòu)、值和計(jì)算行為

• 并行化注意力:使用來自近期的 Blockwise Parallel Transformer 論文中的示例可視化注意力頭的并行化。

• 注意力層的大小:當(dāng)我們將整個(gè)注意力層可視化為單個(gè)結(jié)構(gòu)，則注意力層的 MHA 半邊和 FFA 半邊合起來是什么樣子?在自回歸解碼過程中，其圖像會(huì)發(fā)生怎樣的變化?

• LoRA:對(duì)這種注意力頭架構(gòu)的詳細(xì)闡釋的可視化解釋

1介紹

mm 的可視化方法基于這一前提:矩陣乘法本質(zhì)上是一種三維運(yùn)算。

換句話說:

其實(shí)可以描繪成這種形式:

當(dāng)我們以這種方式將矩陣乘法包裹在一個(gè)立方體中時(shí)，參數(shù)形狀、結(jié)果形狀和共享維度之間的正確關(guān)系就全部就位了。

現(xiàn)在矩陣乘法計(jì)算就有了幾何意義:結(jié)果矩陣中的每個(gè)位置 i，j 都錨定了一個(gè)沿立方體內(nèi)部的深度（depth）維度 k 運(yùn)行的向量，其中從 L 的第 i 行延伸出來的水平面與從 R 的第 j 列延伸出來的垂直面相交。沿著這一向量，來自左邊參數(shù)和右邊參數(shù)的成對(duì)的 (i， k) (k， j) 元素會(huì)相遇并相乘，再沿 k 對(duì)所得積求和，所得結(jié)果放入結(jié)果的 i， j 位置。

這就是矩陣乘法的直觀含義:

1. 將兩個(gè)正交矩陣投影到一個(gè)立方體的內(nèi)部;

2. 將每個(gè)交叉點(diǎn)的一對(duì)值相乘，得到一個(gè)乘積網(wǎng)格;

3. 沿第三個(gè)正交維度進(jìn)行求和，以生成結(jié)果矩陣。

對(duì)于方向，該工具會(huì)在立方體內(nèi)部顯示一個(gè)指向結(jié)果矩陣的箭頭，其中藍(lán)色箭羽來自左側(cè)參數(shù)，紅色箭羽來自右側(cè)參數(shù)。該工具還會(huì)顯示白色指示線來指示每個(gè)矩陣的行軸線，盡管這些線在此截圖中很模糊。

其布局約束條件簡單又直接:

• 左側(cè)參數(shù)和結(jié)果必須沿它們共享的高度 （i） 維度鄰接

• 右側(cè)參數(shù)和結(jié)果必須沿它們共享的寬度 （j） 維度鄰接

• 左側(cè)參數(shù)和右側(cè)參數(shù)必須沿它們共享的（左寬度 / 右高度）維度鄰接，該維度成為矩陣乘法的深度 (k) 維度

這種幾何表示方法能為可視化所有標(biāo)準(zhǔn)的矩陣乘法分解提供堅(jiān)實(shí)的基礎(chǔ)，并能為探索非平凡的復(fù)雜矩陣乘法組合提供直觀的基礎(chǔ)，接下來我們就能看到這一點(diǎn)。

2熱身:動(dòng)畫

在深入介紹更復(fù)雜的示例之前，我們先來看看這種可視化風(fēng)格看起來是什么樣的，從而建立起對(duì)該工具的直覺認(rèn)知。

2a 點(diǎn)積

首先來看一個(gè)經(jīng)典算法 —— 通過計(jì)算對(duì)應(yīng)左側(cè)行和右側(cè)列的點(diǎn)積來計(jì)算每個(gè)結(jié)果元素。從這里的動(dòng)畫可以看到，相乘的值向量掃過立方體內(nèi)部，每一次都在相應(yīng)位置提交一個(gè)求和后的結(jié)果。

這里，L 具有填充有1（藍(lán)色）或 -1(紅色)的行塊;R 具有類似填充的列塊。這里 k 是24，所以結(jié)果矩陣 (L @ R) 的藍(lán)色值為24，紅色值為 -24。

2b 矩陣 – 向量積

分解為矩陣 – 向量積的矩陣乘法看起來像一個(gè)垂直平面（左側(cè)參數(shù)與右側(cè)參數(shù)每一列的積），當(dāng)它水平掃過立方體內(nèi)部時(shí)，將列繪制到結(jié)果上:

觀察一個(gè)分解的中間值可能很有意思，即使示例很簡單。

舉個(gè)例子，請(qǐng)注意當(dāng)我們使用隨機(jī)初始化的參數(shù)時(shí)，中間的矩陣 – 向量積突出的垂直模式 —— 這反映了一個(gè)事實(shí):每個(gè)中間值都是左側(cè)參數(shù)的列縮放的副本:

2c 向量 – 矩陣積

分解為向量 – 矩陣積的矩陣乘法看起來像一個(gè)水平平面，其在向下穿過立方體內(nèi)部時(shí)將行繪制到結(jié)果上:

切換成隨機(jī)初始化的參數(shù)，可以看到類似矩陣 – 向量積的模式 —— 只不過這次是水平模式，對(duì)應(yīng)的事實(shí)是每個(gè)中間向量 – 矩陣積都是右側(cè)參數(shù)的行縮放的副本。

在思考矩陣乘法如何表示其參數(shù)的秩和結(jié)構(gòu)時(shí)，一種有用的做法是設(shè)想這兩種模式在計(jì)算中同時(shí)發(fā)生:

這里還有另一個(gè)使用向量 – 矩陣積來構(gòu)建直覺的示例，其中展示了單位矩陣的作用就像是一面呈45度角擺放的鏡子，反射著其對(duì)應(yīng)參數(shù)和結(jié)果:

2d 對(duì)外積求和

第三次平面分解是沿著 k 軸，通過對(duì)向量外積逐點(diǎn)求和來計(jì)算矩陣乘法結(jié)果。這里我們可以看到外積平面「從后到前」掃過立方體，累積到結(jié)果中:

使用隨機(jī)初始化的矩陣進(jìn)行此分解，我們不僅可以看到值，還可以看到結(jié)果中的秩累積，因?yàn)槊總€(gè)秩為1的外積都被添加到其中。

這也從直覺上說明了為什么「低秩因式分解」（即通過構(gòu)造參數(shù)在深度維度上較小的矩陣乘法來近似矩陣）在被近似的矩陣為低秩矩陣時(shí)的效果最好。這是后面會(huì)提到的 LoRA:

3熱身:表達(dá)式

我們可以怎樣的方式將這種可視化方法擴(kuò)展用于矩陣乘法的分解?之前的示例可視化的是矩陣 L 和 R 的單次矩陣乘法 L @ R，但要是 L 和 / 或 R 本身也是矩陣乘法呢?

事實(shí)證明這種方法可以很好地?cái)U(kuò)展用于復(fù)合表達(dá)式。關(guān)鍵規(guī)則很簡單:子表達(dá)式（子）矩陣乘法是另一個(gè)立方體，其受到與父矩陣乘法一樣的布局約束;子矩陣乘法的結(jié)果面同時(shí)也是父矩陣乘法對(duì)應(yīng)的參數(shù)面，就像是共價(jià)共享的電子。

在這些約束限制中，我們可以按自己的需求排布子矩陣乘法的各個(gè)面。這里使用該工具的默認(rèn)方案，這會(huì)生成交替的凸面和凹面立方體 —— 這種布局的實(shí)踐效果很好，可以最大化地利用空間，同時(shí)盡可能減少遮擋。（但是布局是完全可定制的，詳情訪問 mm 工具頁面。）

這一節(jié)將可視化機(jī)器學(xué)習(xí)模型的一些關(guān)鍵構(gòu)建模塊，以便讀者熟悉這種視覺表示并從中獲得新的直覺認(rèn)識(shí)。

3a 左結(jié)合表達(dá)式

下面將會(huì)介紹兩個(gè) （A @ B） @ C 形式的表達(dá)式，每一個(gè)都有自己的獨(dú)特形狀和特征。(注意:mm 遵循矩陣乘法是左結(jié)合的約定，所以 (A @ B) @ C 可簡單寫為 A @ B @ C。)

首先為 A @ B @ C 賦予很有特點(diǎn)的 FFN 形狀，其中「隱藏維度」比「輸入」或「輸出」維度寬。（具體來說，就此示例而言，這意味著 B 的寬度大于 A 或 C 的寬度。）

和單次矩陣乘法示例一樣，浮動(dòng)的箭頭指向結(jié)果矩陣，其中藍(lán)色箭羽來自左側(cè)參數(shù)，紅色箭羽來自右側(cè)參數(shù)。

而當(dāng) B 的寬度小于 A 或 C 的寬度時(shí)，對(duì) A @ B @ C 的可視化則會(huì)有一個(gè)瓶頸，類似自動(dòng)編碼器的形狀。

交替的凹凸模塊的模式還可以擴(kuò)展成任意長度的鏈:比如這個(gè)多層瓶頸:

3b 右結(jié)合表達(dá)式

接下來可視化右結(jié)合表達(dá)式 A @ （B @ C）。

與左結(jié)合表達(dá)式的水平擴(kuò)展類似 —— 可以說是從根表達(dá)式的左側(cè)參數(shù)發(fā)端，右結(jié)合表達(dá)式鏈?zhǔn)且源怪狈绞綌U(kuò)展，從根表達(dá)式的右側(cè)參數(shù)發(fā)端。

人們有時(shí)候可以看到一個(gè)以右結(jié)合形式形成的 MLP，即右側(cè)是柱狀輸入，權(quán)重層從右到左運(yùn)行。使用上面描繪的二層 FFN 示例的矩陣（適當(dāng)轉(zhuǎn)置后），看起來會(huì)是這樣，C 現(xiàn)在是輸入，B 是第一層，A 是第二層:

另外，除了箭羽的顏色（左側(cè)為藍(lán)色，右側(cè)為紅色），區(qū)分左右參數(shù)的第二個(gè)視覺提示是它們的方向:左側(cè)參數(shù)的行與結(jié)果的行共面 —— 它們沿同一根軸 (i) 堆疊。比如上面的 (B @ C)，這兩個(gè)提示都能告訴我們 B 是左側(cè)參數(shù)。

3c 二元表達(dá)式

對(duì)于可視化工具，要有用，就不能只用于簡單的教學(xué)式示例，也要能方便地用于更復(fù)雜的表達(dá)式。在真實(shí)世界用例中，一個(gè)關(guān)鍵性結(jié)構(gòu)組件是二元表達(dá)式 —— 左側(cè)和右側(cè)都有子表達(dá)式的矩陣乘法。

這里可視化了此類表達(dá)式中形狀最簡單的一個(gè) （A @ B） @ (C @ D):

3d 一點(diǎn)注解:分區(qū)和并行性

完整闡述該主題超出了本文的范圍，但后面我們會(huì)在注意力頭部分看到它的實(shí)際效用。但這里熱個(gè)身，看兩個(gè)簡單示例，了解下這種可視化風(fēng)格可以如何讓對(duì)并行化復(fù)合表達(dá)式的推理非常直觀 —— 只需通過簡單的幾何分區(qū)。

第一個(gè)示例是將典型的「數(shù)據(jù)并行」分區(qū)應(yīng)用于上面的左結(jié)合多層瓶頸示例。我們沿 i 分區(qū)，對(duì)初始左側(cè)參數(shù)（「批」）和所有中間結(jié)果(「激活」)進(jìn)行分段，但不對(duì)后續(xù)參數(shù)(「權(quán)重」)分段 —— 這種幾何結(jié)構(gòu)使得表達(dá)式中的哪些參與者被分段以及哪些保持完整變得顯而易見:

第二個(gè)示例如果沒有清晰的幾何支持，就很難直覺地理解:它展示了如何通過沿 j 軸對(duì)左側(cè)子表達(dá)式分區(qū)、沿 i 軸對(duì)右側(cè)子表達(dá)式分區(qū)以及沿 k 軸對(duì)父表達(dá)式進(jìn)行分區(qū)來并行化一個(gè)二元表達(dá)式:

4深入注意力頭

現(xiàn)在來看看 GPT-2的注意力頭 —— 具體來說是來自 NanoGPT 的5層第4頭的 「gpt2」（small） 配置(層數(shù) =12，頭數(shù) =12，嵌入數(shù) =768)，通過 HuggingFace 使用了來自 OpenAI 的權(quán)重。輸入激活取自在含256個(gè) token 的 OpenWebText 訓(xùn)練樣本上一次前向通過。

這個(gè)特定的頭并無任何特殊之處，選擇它主要是因?yàn)槠溆?jì)算的是一個(gè)非常常見的注意力模式，并且它位于模型中部，其中激活已經(jīng)變得結(jié)構(gòu)化并顯示出一些有趣的紋理。

4a 結(jié)構(gòu)

這個(gè)完整注意力頭被可視化成了單個(gè)復(fù)合表達(dá)式，其始于輸入，終于投影的輸出。（注:為了保證自足性，這里按照 Megatron-LM 的描述對(duì)每個(gè)頭執(zhí)行輸出投影。）

這一計(jì)算包含六次矩陣乘法:

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

簡單描述一下這里在做什么:

• 風(fēng)車的葉片為矩陣乘法1、2、3和6:前一組是輸入到 Q、K 和 V 的內(nèi)投影;后者是從 attn @ V 回到嵌入維度的外投影。

• 中心有兩個(gè)矩陣乘法;第一個(gè)計(jì)算的是注意力分?jǐn)?shù)（后面的凸立方體），然后使用它們基于值向量得到輸出 token(前面的凹立方體)。因果關(guān)系意味著注意力分?jǐn)?shù)形成一個(gè)下三角形。

但讀者最好能親自詳細(xì)探索這個(gè)工具，而不是只看截圖或下面的視頻，以便更詳細(xì)地理解 —— 不管是其結(jié)構(gòu)還是流過計(jì)算過程的實(shí)際值。

4b 計(jì)算和值

這里是注意力頭的計(jì)算過程動(dòng)畫。具體來說，我們是看

sdpa（input@wQ@K_t）@V@wO

（即上面的矩陣乘法1、4、5和6，其中 K_t 和 V 已經(jīng)預(yù)先計(jì)算）是作為向量 – 矩陣積的融合鏈來計(jì)算:序列中的每一項(xiàng)都在一步之內(nèi)從輸入穿過注意力到輸出。后面關(guān)于并行化的部分會(huì)提到更多有關(guān)這個(gè)動(dòng)畫的選擇，但我們先看看計(jì)算的值能告訴我們什么。

我們可以看到很多有趣的東西:

• 在討論注意力計(jì)算之前，可以看到低秩 Q 和 K_t 的形態(tài)是多么驚人。放大 Q @ K_t 向量 – 矩陣積動(dòng)畫，看起來會(huì)更加生動(dòng):Q 和 K 中大量通道（嵌入位置）在序列中看起來或多或少是恒定的，這意味著有用的注意力信號(hào)可能僅由一小部分嵌入驅(qū)動(dòng)。理解和利用這種現(xiàn)象是 SysML ATOM transformer 效率項(xiàng)目的一部分。

• 也許人們最熟悉的是注意力矩陣中出現(xiàn)的強(qiáng)大但不完美的對(duì)角線。這是一種常見模式，出現(xiàn)在該模型（以及許多 Transformer）的許多注意力頭中。它能產(chǎn)生局部注意力:緊鄰輸出 token 位置之前的小鄰域中的值 token 很大程度上決定了輸出 token 的內(nèi)容模式。

• 然而，這個(gè)鄰域的大小和其中各個(gè) token 的影響變化很大 —— 這可以在注意力網(wǎng)格中的非對(duì)角 frost 中看到，也能在注意力矩陣沿序列下降時(shí) attn [i] @ V 向量 – 矩陣積平面的波動(dòng)模式中看到。

• 但請(qǐng)注意，局部鄰域并不是唯一值得注意的東西:注意力網(wǎng)格的最左列（對(duì)應(yīng)于序列的第一個(gè) token ）完全填充了非零(但波動(dòng))的值，這意味著每個(gè)輸出 token 都會(huì)受到第一個(gè)值 token 一定程度的影響。

• 此外，當(dāng)前 token 鄰域和初始 token 之間的注意力分?jǐn)?shù)主導(dǎo)性存在不精確但可辨別的振蕩。該振蕩的周期各有不同，但一般來說，一開始很短，然后沿序列向下移動(dòng)而變長（類似地，在給定因果關(guān)系的情況下，與每一行的候選注意力 token 的數(shù)量相關(guān)）。

• 為了了解 （attn @ V） 的形成方式，不單獨(dú)關(guān)注注意力是很重要的 ——V 也同等重要。每個(gè)輸出項(xiàng)都是整個(gè) V 向量的加權(quán)平均值:在注意力是完美對(duì)角線的極端情況下，attn @ V 只是 V 的精確副本。這里我們看到更有紋理的東西:可見的帶狀結(jié)構(gòu)，其中特定 token 在注意力行的連續(xù)子序列上的得分很高，疊加在與 V 明顯相似的矩陣上，但由于對(duì)角線較粗而有一些垂直遮擋。(旁注:根據(jù) mm 參考指南，長按或按住 Control 鍵單擊將顯示可視化元素的實(shí)際數(shù)值。)

• 請(qǐng)記住，由于我們位于中間層（5層），因此該注意力頭的輸入是一個(gè)中間表示，而不是原始 token 化文本。因此，在輸入中看到的模式本身就發(fā)人深省 —— 特別是，強(qiáng)大的垂直線條是特定的嵌入位置，其值在序列的長段上統(tǒng)一具有高的幅度 —— 有時(shí)幾乎是占滿了。

• 但有趣的是，輸入序列中的第一個(gè)向量是獨(dú)特的，不僅打破了這些高幅度列的模式，而且?guī)缀踉诿總€(gè)位置都攜帶著非典型值（旁注:這里沒有可視化，但這種模式反復(fù)出現(xiàn)在多個(gè)樣本輸入上）。

• 注意:關(guān)于最后兩個(gè)要點(diǎn)，值得重申的是，這里可視化的是對(duì)單個(gè)樣本輸入的計(jì)算。在實(shí)踐中，可以發(fā)現(xiàn)每個(gè)頭都有一個(gè)特征模式，能在相當(dāng)多的樣本集合上一致（盡管不等同）地表達(dá)，但當(dāng)查看任意包含激活的可視化時(shí)，需要記住:輸入的完整分布可能會(huì)以微妙的方式影響它激發(fā)的想法和直覺。

最后，再次建議直接探索動(dòng)畫!

4c 注意力頭有很多有趣的不同之處

繼續(xù)之前，這里再通過一個(gè)演示展現(xiàn)簡單地研究模型以了解其詳細(xì)工作方式的有用性。

這是 GPT-2的另一個(gè)注意力頭。其行為模式與5層第4頭大有不同 —— 這符合人們預(yù)期，畢竟它位于模型一個(gè)非常不同的部分。這個(gè)頭位于第一層:0層的第2頭:

值得注意的點(diǎn):

• 這個(gè)頭的注意力分布很均勻。這會(huì)產(chǎn)生一個(gè)效果:將 V 的相對(duì)未加權(quán)的平均值（或者說 V 的合適的因果前綴）交到 attn @ V 的每一行;如動(dòng)畫所示:當(dāng)我們向下移動(dòng)注意力分?jǐn)?shù)三角時(shí)，attn [i] @ V 向量 – 矩陣積有很小的波動(dòng)，而不是簡單的 V 的縮小比例的、逐漸揭示的副本。

• attn @ V 具有驚人的垂直均勻性 —— 在嵌入的大柱狀區(qū)域中，相同的值模式在整個(gè)序列中持續(xù)存在。人們可以將這些看作是每個(gè) token 共享的屬性。

• 旁注:一方面，考慮到注意力分布非常均勻的效果，人們可能會(huì)期望 attn @ V 具有一定的一致性。但是每一行都是由 V 的因果子序列而不是整個(gè)序列構(gòu)成 —— 為什么這不會(huì)導(dǎo)致更多的變化，就像沿著序列向下移動(dòng)時(shí)的漸進(jìn)變形一樣?通過視覺檢查可知，V 沿其長度并不均勻，因此答案必定在于其值分布的一些更微妙的屬性。

• 最后，在外投影后，這個(gè)頭的輸出在垂直方向上還要更加均勻。

• 我們能得到一個(gè)強(qiáng)烈的印象:該注意力頭傳遞的大部分信息由序列中每個(gè) token 共享的屬性組成。其輸出投影權(quán)重的構(gòu)成能強(qiáng)化這種直覺。

總的來說，我們不由得會(huì)想:這個(gè)注意力頭產(chǎn)生的極其規(guī)則、高度結(jié)構(gòu)化的信息可能是通過稍微…… 不那么奢華的計(jì)算手段獲得的。當(dāng)然，這不是一個(gè)未經(jīng)探索的領(lǐng)域，但可視化計(jì)算信號(hào)的明確性和豐富性對(duì)于產(chǎn)生新想法和推理現(xiàn)有想法都非常有用。

4d 重返介紹:免費(fèi)的不變性

回頭看，需要重申:我們之所以能夠?qū)⒆⒁饬︻^等非平凡的復(fù)合操作可視化并讓它們保持直觀，是因?yàn)橹匾拇鷶?shù)性質(zhì)（例如參數(shù)形狀的限制方式或者哪些并行軸與哪些操作相交），這些性質(zhì)不需要額外的思考:它們直接來自可視化對(duì)象的幾何屬性，而不是需要記住的額外規(guī)則。

舉個(gè)例子，在這些注意力頭可視化中，可以明顯看出:

• Q 和 attn @ V 的長度一樣，K 和 V 的長度一樣，這些配對(duì)的長度都彼此獨(dú)立;

• Q 和 K 的寬度一樣，V 和 attn @ V 的寬度一樣，這些配對(duì)的寬度都彼此獨(dú)立。

這些結(jié)構(gòu)在構(gòu)造上就是真實(shí)的，就是結(jié)構(gòu)組分位于復(fù)合結(jié)構(gòu)的哪個(gè)部分以及它們的方向如何的簡單結(jié)果。

這種「免費(fèi)性質(zhì)」的優(yōu)勢在探索典型結(jié)構(gòu)的變體時(shí)特別有用 —— 一個(gè)明顯的例子是一次解碼一個(gè)自回歸 token 中的單行高的注意力矩陣:

5并行化注意力

上面5層第4頭的動(dòng)畫可視化了注意力頭中6個(gè)矩陣乘法中的4個(gè)。

它們被可視化為了一條向量 – 矩陣積的融合鏈，從而證實(shí)了一個(gè)幾何直覺:從輸入到輸出的整個(gè)左結(jié)合鏈沿共享 i 軸呈層狀，且可并行化。

5a 示例:沿 i 分區(qū)

為了在實(shí)踐中并行計(jì)算，我們可將輸入沿 i 軸劃分為塊。我們可以在該工具中可視化這種分區(qū)，通過指定將給定軸劃分為特定數(shù)量的塊 —— 在這些示例中將使用8，但該數(shù)字并無特別之處。

除此之外，這種可視化清楚地表明，每次并行計(jì)算都需要完整的 wQ（用于內(nèi)投影）、K_t 和 V(用于注意力)和 wO(用于外投影)，因?yàn)樗鼈冄剡@些矩陣的未分區(qū)維度與已分區(qū)矩陣相鄰接:

5b 示例:雙重分區(qū)

這里也給出沿多個(gè)軸進(jìn)行分區(qū)的示例。為此，這里選擇可視化該領(lǐng)域一個(gè)近期的創(chuàng)新成果，即 Block Parallel Transformer（BPT），其基于 Flash Attention 等一些研究成果，參閱論文:https://arxiv.org/pdf/2305.19370.pdf

首先，BPT 如上所述沿 i 進(jìn)行分區(qū) —— 并且實(shí)際上也將序列的這種水平分區(qū)一直延伸到注意力層的另一半邊（FFN）。(對(duì)此的可視化將在后面展現(xiàn)。)

為了完全解決這個(gè)上下文長度問題，向 MHA 添加第二個(gè)分區(qū) —— 注意力計(jì)算本身的分區(qū)（即沿 Q @ K_t 的 j 軸的分區(qū)）。這兩個(gè)分區(qū)一起可將注意力分成塊構(gòu)成的網(wǎng)格:

從這個(gè)可視化可以清楚看到:

• 這種雙重分區(qū)能有效解決上下文長度問題，因?yàn)槲覀儸F(xiàn)在能以視覺劃分注意力計(jì)算中每次出現(xiàn)的序列長度。

• 第二次分區(qū)的「范圍」:根據(jù)幾何結(jié)構(gòu)可以明顯看出，K 和 V 的內(nèi)投影計(jì)算可與核心的雙矩陣乘法一起分區(qū)。

• 注意一個(gè)微妙細(xì)節(jié):這里的視覺暗示是我們還可以沿 k 并行化后續(xù)的矩陣乘法 attn @ V 并以 split-k 風(fēng)格對(duì)部分結(jié)果求和，從而并行化整個(gè)雙重矩陣乘法。但 sdpa （） 中的逐行 softmax 增加了要求:在計(jì)算 attn @ V 的相應(yīng)行之前，每一行都要將其所有分段歸一化，這會(huì)在注意力計(jì)算和最終矩陣乘法之間添加一個(gè)額外的逐行步驟。

6注意力層的大小

眾所周知，注意力層的前一半（MHA）由于其二次復(fù)雜度而有很高的計(jì)算需求，但后一半(FFN)也有自己的需求，這要?dú)w因于其隱藏維度的寬度，其通常是模型嵌入維度的寬度的4倍。可視化完整注意力層的生物量有助于建立關(guān)于該層兩半部分如何相互比較的直覺認(rèn)識(shí)。

6a 可視化完整的注意力層

下面是一個(gè)完整的注意力層，前一半（MHA）位于后面，后一半(FFN)位于前面。同樣，箭頭指向計(jì)算的方向。

注:

該可視化描繪的不是單個(gè)注意力頭，而是顯示了未切片的 Q/K/V 權(quán)重和圍繞中心雙重矩陣乘法的投影。當(dāng)然這沒有將完整的 MHA 運(yùn)算忠實(shí)地可視化出來 —— 但這里的目標(biāo)是更清楚地了解該層的兩半中的相對(duì)矩陣大小，而不是每半執(zhí)行的相對(duì)計(jì)算量。（此外，這里的權(quán)重使用了隨機(jī)值而非真實(shí)權(quán)重。）

這里使用的維度有所收縮以保證瀏覽器（相對(duì)）能帶得動(dòng)，但比例保持一樣(來自 NanoGPT 的 small 配置):模型嵌入維度 =192(原本是768)、FFN 嵌入維度 =768(原本是3072)、序列長度 =256(原本是1024)，盡管序列長度對(duì)模型沒有根本性影響。(從視覺上看，序列長度的變化將表現(xiàn)為輸入葉片寬度的變化，從而導(dǎo)致注意力中心大小和下游垂直平面高度的變化。)

6b 可視化 BPT 分區(qū)層

簡單地回顧一下 Blockwise Parallel Transformer，這里是在整個(gè)注意力層的語境中可視化 BPT 的并行化方案（和上面一樣省略了各個(gè)頭）。特別要注意，沿 i(序列塊)的分區(qū)以怎樣的方式擴(kuò)展通過 MHA 和 FFN 兩半邊:

6c 對(duì) FFN 進(jìn)行分區(qū)

這種可視化方法建議進(jìn)行額外的分區(qū)，該分區(qū)與上面描述的分區(qū)正交 —— 在注意力層的 FFN 半邊，將雙重矩陣乘法 （attn_out @ FFN_1） @ FFN_2分開，首先沿 j 進(jìn)行 attn_out @ FFN_1，然后沿 k 與 FFN_2執(zhí)行后續(xù)的矩陣乘法。這種分區(qū)會(huì)對(duì)兩個(gè) FFN 權(quán)重層進(jìn)行切片，以減少計(jì)算中每個(gè)參與組分的容量要求，但代價(jià)是部分結(jié)果的最終求和。

下面是將這種分區(qū)方法應(yīng)用于未分區(qū)的注意力層的樣子:

下面則是應(yīng)用于以 BPT 方式分區(qū)的層的情況:

6d 可視化一次一個(gè) token 解碼的過程

在自回歸式的一次一個(gè) token 的解碼過程中，查詢向量由單個(gè) token 構(gòu)成。你可以在頭腦中想象一下這種情況下的注意力層會(huì)是什么樣子，這很有啟發(fā)性 —— 單個(gè)嵌入行穿過一個(gè)巨大的平鋪的權(quán)重平面。

除了強(qiáng)調(diào)與激活相比權(quán)重的巨大性之外，這種觀點(diǎn)還能讓人想起這樣一個(gè)概念:K_t 和 V 的功能類似于一個(gè)6層 MLP 中動(dòng)態(tài)生成的層，盡管 MHA 本身的 mux/demux 計(jì)算會(huì)使這種對(duì)應(yīng)關(guān)系不精確:

7LoRA

近期的 LoRA 論文《LoRA: Low-Rank Adaptation of Large Language Models》描述了一種高效的微調(diào)技術(shù)，該技術(shù)基于這一思想:微調(diào)期間引入的權(quán)重 δ 是低秩的。根據(jù)這篇論文，這「允許我們通過在適應(yīng)過程中優(yōu)化密集層變化的秩分解矩陣來間接地訓(xùn)練神經(jīng)網(wǎng)絡(luò)中的一些密集層…… 同時(shí)保持預(yù)訓(xùn)練權(quán)重處于凍結(jié)狀態(tài)。」

7a 基本思想

簡而言之，關(guān)鍵一步是訓(xùn)練權(quán)重矩陣的因子而不是矩陣本身:用一個(gè) I x K 張量和 K x J 張量的矩陣乘法來替代 I x J 權(quán)重張量，其中要保證 K 為一個(gè)較小值。

如果 K 足夠小，則尺寸方面可能會(huì)有很大贏面，但也有權(quán)衡之處:降低 K 也會(huì)降低積可以表達(dá)的秩。這里通過一個(gè)示例說明尺寸上的節(jié)省與對(duì)結(jié)果的結(jié)構(gòu)化影響，這里是隨機(jī)的128×4左側(cè)參數(shù)和4×128右側(cè)參數(shù)的矩陣乘法 —— 即一個(gè)128×128矩陣的秩為4的分解。注意 L @ R 中的垂直和水平模式:

7b 將 LoRA 應(yīng)用于注意力頭

LoRA 將這種分解方法應(yīng)用于微調(diào)過程的方式是:

• 為每個(gè)權(quán)重張量創(chuàng)建一個(gè)要進(jìn)行微調(diào)的低秩分解，并訓(xùn)練其因子，同時(shí)保持原始權(quán)重凍結(jié);

• 微調(diào)之后，將每對(duì)低秩因子相乘，得到一個(gè)原始權(quán)重張量形狀的矩陣，并將其添加到原始的預(yù)訓(xùn)練權(quán)重張量中。

下面的可視化顯示了一個(gè)注意力頭，其權(quán)重張量 wQ、wK_t、wV、wO 被低秩分解 wQ_A @ wQ_B 等替換。從視覺上看，因子矩陣呈現(xiàn)為沿風(fēng)車葉片邊緣的低柵欄