探秘Transformer系列之（10）--- 自注意力

從零開始解析Transformer，目標是：(1) 解析Transformer如何運作，以及為何如此運作，讓新同學可以入門；(2) 力爭融入一些比較新的或者有特色的論文或者理念，讓老鳥也可以有所收獲。

探秘Transformer系列之（10）--- 自注意力

探秘Transformer系列之（10）--- 自注意力

0x00 概述

Transformer的核心所在或者說與其他架構的關鍵區別之處是自注意力機制，其允許模型在處理一個句子時，考慮句子中每個單詞與其他所有單詞的依賴關系，并使用這些信息來捕捉句子的內部結構和表示，最終計算單詞之間的關聯度（權重）。我們可以把自注意力機制分為三個階段：

輸入：從前文我們可以了解到，注意力機制接受查詢（query）、鍵（key）和值（value）三個輸入，但是對于自注意力來說，只有一個輸入序列，Q、K、V都是來自于這個序列。這個輸入序列是一個向量列表，且向量之間有一定的關系。以機器翻譯為例，輸入序列就是源語句或者目標語句，語句中每個token對應一個向量。
計算：自注意力機制會計算序列中每個向量與序列中其他向量的關系（也就是每個單詞與句子中所有單詞的關系），使得序列中的每個token都能感知其他token。針對當前向量，自注意力機制會接受計算查詢（當前token）與所有鍵（感興趣的token）的點積，應用Softmax函數在點積上以獲取權重，并使用權重對所有與之關聯的值進行加權平均。這樣就可以把對其他單詞的“理解”融入到當前處理的單詞中。
輸出：一個序列，比如一個向量列表，但是列表之中所有向量都考慮了其上下文關系，是蘊含了序列內部關系的全局特征表示。

具體如下圖所示。

0x01 原理

1.1 設計思路

自注意力并非Transformer首創，但之前在其他模型上效果不甚理想。所以我們好奇為什么Transformer依然使用自注意力呢，論文是這樣解釋其設計思路：

Motivating our use of self-attention we consider three desiderata. One is the total computational complexity per layer. Another is the amount of computation that can be parallelized, as measured by the minimum number of sequential operations required. The third is the path length between long-range dependencies in the network.

三個考慮因素我們具體解析如下。

每層的總計算復雜度。Transformer的自注意力使用了縮放點積注意力評分函數，相比于加性注意力減少了計算量，效果也是相似的。
并行計算。自注意力是可以并行化計算的，并行化計算量可以用所需的序列操作的最小數目來衡量。
網絡中長距離依賴關系之間的路徑長度。RNN捕捉詞與詞之間關系需要把句子從頭看到尾，CNN需要層疊多個卷積層才能捕捉詞與詞之間關系，而自注意力是完全并行的，每個詞可以直接關聯。

雖然有若干優勢，但是知易行難，Transformer如何做到？我們接下來一步一步進行分析。

1.2 輸入

從宏觀角度來說，Transformer只有一個輸入序列，由這個序列派生出來Q、K和V。具體如下圖所示。

從微觀角度看，以編碼器為例，自注意力的Q、K、V的來源有兩種：

第一個編碼器層的QKV由輸向量x組成的矩陣X進行線性變化而來，線性變化就是用$W^Q,W^K,W^V$進行矩陣乘法。
后續編碼器層的QKV由上一個編碼器層的輸出經過線性變化而來。

我們接下來以第一個編碼器層為例，從源序列中的單個詞開始來跟蹤它們在 Transformer 中的路徑。為了解釋和可視化，我們暫時不用關心細節，只跟蹤每個詞對應的"行"。

假如我們進行英譯中，輸入中文：我愛你。假定模型維度為d，輸入序列長度為L。源序列首先通過嵌入和位置編碼層，該層為序列中的每個單詞生成嵌入向量，這些嵌入向量構成的矩陣就是X。

輸入序列接下來會通過三個矩陣$\mathbf{W}^K \in \mathbb{R}^{d \times d_k}$，$\mathbf{W}^Q \in \mathbb{R}^{d \times d_q}$，$\mathbf{W}^V \in \mathbb{R}^{d \times d_v}$進行轉換。具體來說，輸入序列的每一個元素$x_i \in \mathbb{R}^w0obha2h00$會分別乘以這三個矩陣，得到

\[q_i = x_i \mathbf{W}^Q \\ k_i = x_i \mathbf{W}^K \\ v_i = x_i \mathbf{W}^V \\ \]

把L個$q_i$對堆疊起來就得到矩陣$\mathbf{Q} \in \mathbb{R}^{L \times d_q}$，類似可以得到矩陣$\mathbf{K} \in \mathbb{R}^{L \times d_k}$，$\mathbf{V} \in \mathbb{R}^{L \times d_v}$?；蛘咧苯佑镁仃囆问奖磉_：

\[\mathbf{Q} = \mathbf{X}\mathbf{W}^Q \\ \mathbf{K} = \mathbf{X}\mathbf{W}^K \\ \mathbf{V} = \mathbf{X}\mathbf{W}^V \]

這個三個獨立的矩陣Q、K、V會被用來計算注意力得分。這些矩陣的每一 "行 "都是一個向量，對應于源序列中的一個詞。每個這樣的"行"都是通過一系列的諸如嵌入、位置編碼和線性變換等轉換，從其相應的源詞中產生。而所有這些的轉換都是可訓練的操作。這意味著在這些操作中使用的權重不是預先確定的，而是利用反向傳播機制進行學習得到的。

1.3 QKV解析

自注意力機制中第一步就是用Token來生成查詢向量、鍵向量和值向量，也就是用到了query，key，value（各種相關論文、網址之中也縮寫為q、k、v）這三個概念。Query向量代表當前正在處理的token或位置，它表示模型需要“查詢”的信息。Key向量代表序列中每個token的唯一標識，用于與Query進行比較。Value向量包含序列中每個token的實際內容或特征，它對生成當前token的輸出有貢獻。

要理解LLM的底層實現原理，就必須要了解Transformers Block里面的QKV矩陣，因為前沿的大模型研究工作很大一部分就是著QKV矩陣去做的，比如注意力、量化、低秩壓縮等等，目標是在保證效果不變壞的前提下，進行對性能和存儲的極致壓縮。其本質原因是因為QKV權重占比著大語言模型50%以上的權重比例，在推理過程中，QKV存儲量還會隨著上下文長度的增長而線性增長，計算量也平方增加?？梢哉f，query，key，value對于Transformers 和自注意力機制至關重要。

相信大家一直都有個疑問，為什么要取QKV這些名字？這一套思想到底怎么去理解？之前篇幅中介紹過QKV，但是始終沒有深入，本篇會進行詳細分析。因為深度學習其實是帶有實踐性質的科學，尚未找到確切的理論分析。所以接下來我們從不同理解的角度來闡釋，希望讀者能夠從其中對QKV有所理解。

心理學角度

有研究人員發現，注意力機制可以追溯到美國心理學之父威廉·詹姆斯在19世紀90年代提出的非自主性提示（nonvolitional cue）、自主性提示（volitional cue）和感官輸入（Sensory inputs）這幾個概念。而這三個概念就分別可以對應到Key張量，Query張量和與Key有對應關系的Value張量，然后由這三者構建了注意力機制。

我們用個通俗例子來分析下。本來你去買鹽（帶有目的性的關注度，即自主性提示），結果你到了商店，發現了變形金剛，你注意力都被變形金剛（下意識的關注度，即非自主性提示）吸引了。我們可以得到：

Key：一系列物品（鹽和變形金剛）。
Value：這一系列物品對人下意識的吸引力（在你下意識中，變形金剛的吸引力肯定比鹽要高）。
Query：你想要的物品（鹽）。

注意力作用就是讓鹽（目標物品）所對應的權重值變高。這樣使用該'權重向量’乘‘key’后，即使目標物品“下意識的吸引力（即key）”不夠高，但是因為目標物品對應的權重高，其他物品對應的權重小，故選擇到目標物品的可能性也會變大。

數據庫角度

query，key，value的名稱也暗示了整個注意力計算的思路，因此我們從搜索領域的業務來看Q、K、V可能更好理解。我們把注意力機制看作是一種模糊尋址，或者說是一個模糊的、可微分的、向量化的數據庫（或者字典）查找機制。Q、K、V這三者的關系就是：找到與現有數據（Query）相似或者相關的數據（Key）所對應的內容（Value）。其具體特點如下：

Key和Value是數據庫的組件。
數據庫中每個元素由地址Key和值Value組成（一個<Key,Value>數據對）?；蛘哒f，Key這個地址里面就存放了Value。
Key是地址，就是要查找的位置。該地址總結了地址中Value的特征，或者說Key可以體現Value上的語義信息。
Value是與地址Key相關聯的值，是表征語義的真實數據，是對外提供的使用者所需的內容。
Query是查詢信息，是任務相關的變量。假設當前有個Key=Query的查詢，該查詢會通過Query和存儲器內所有元素Key的地址進行相似性比較來尋址，其目的是取出數據庫中對應的Value值。自注意力機制中這種Q和K一問一答的形式，問的就是Q和K兩個詞之間的緊密程度。直觀地說，Key是Query（我們正在尋找什么）和Value（我們將實際獲得什么）之間的橋梁。
普通的字典查找是精確匹配，即依據匹配的鍵來返回其對應的值，而且只從存儲內容里面找出一條內容。而注意力機制是向量化+模糊匹配+合并的組合使用。其會根據Query和Key的相似性來計算每個Key的相似度或者匹配程度（即注意力權重）。然后從每個Key地址都會取出Value，并依據匹配程度對這些value做加權求和，這個相似度得分決定了相應Value在最終輸出中的權重。

用通俗例子來講解，假如我們在淘寶上進搜索”李寧鞋“，Query（Q）就是你在搜索欄輸入的查詢內容。Key（K）就是在頁面上返回的商品描述、標題，其實就是數據庫中與候選商品相關的關鍵字。Value（V）就是李寧商品本身。注意力機制就是這個查詢過程，即注意力是把你要查詢的Q與淘寶數據庫中的K進行比較，計算出這些K與Q的相似度，最終返回相似度最高的若干商品V。流程如下：

用Query與數據庫內所有Keys進行計算相似度（查詢的相關性，即你有多大概率是我要查的東西）。
得到相似度之后，對結果進行排序。
基于相似度排序結果，得到需要獲取的商品ID
依據商品ID來獲取對應的Values。

因此，自注意力機制中的QKV思想，本質是一個具有全局語義整合功能的數據庫。<Key,Value>數據對就是數據庫的元素，Q就是任務相關的查詢向量。

下圖從數據庫角度展示了自注意力的細節。

seq2seq角度

讓我們回到具體任務上來分析，可能會更加清晰一點。比如在機器翻譯任務中，query可以定義成解碼器中某一步的隱狀態，即對上一個詞的預測輸出。key是編碼器中每個時間步的隱狀態，我們用每一個query對所有key都做一個對齊，于是解碼器的每一步都會得到一個不一樣的對齊向量?；蛘哒f，編碼器的編碼序列（Encoded sequence）提供key和value 。Hidden state of Decoder 提供query 。這就好比解碼器要去編碼器的編碼序列那里查字典一樣。

重構詞向量角度

前面我們從數據庫角度來看到如何尋址獲取訊息，其最終目的是輸出一個新向量。在新的向量中，每一個維度的數值都是由幾個詞向量在這一維度的數值加權求和得來的。因此，自注意力機制的核心是重構詞向量（查詢+聚合），每個輸入單詞的編碼輸出都會通過注意力機制引入其余單詞的編碼信息。

相互操作

人類在讀一篇文章時，為了理解一句話的意思，你不僅會關注這句話本身，還會回看上下文中相關的其他句子或詞語。我們還是以之前兩個句子為例進行解釋。

Several distributor transformers had fallen from the poles, and secondary wires were down.
Transformer models have emerged as the most widely used architecture in applications such as natural language processing and image classification.

如何才能對“Transformer”這個多義詞進行語義區分？我們必須考慮單詞的上下文才能更好的把單詞的語義識別出來，即不僅僅要考慮到詞本身，還要考慮其他詞對這個詞的影響，也就是語境的影響。比如第一個句子的“pole”、”fallen”和“wires”這幾個鄰近單詞暗示了“Transformer”和物理環境相關。第二個句子的“model”和“natural language processing and image classification”則直接告訴我們此處的“Transformer”是深度學習相關概念。最終我們通過上下文語境（句子中的其他詞）可以推斷“Transformer”的準確含義。

原理我們知道，但是如何實踐？如何通過句子中的其它詞來推斷？人類可以知道哪些詞提供了上下文，但計算機卻毫無頭緒，因為計算機只處理數字。解決方案就是注意力機制在Transformer中所模擬的過程。

Transformer 通過點積這個提取特征的操作將輸入序列中的每個詞與其他詞關聯起來，也就是詞之間進行互相操作。然后通過加權求和把這些詞加起來，最終可以捕捉到某個特定的詞和句子中其他每個詞之間的一些互動。于是修改后的詞如下：

transformer 1 = 0.7 transformer + 0.1 pole + 0.1 fallen + 0.1 wires
transformer 2 = 0.6 transformer + 0.1 language + 0.1 image + 0.2 model

最終兩個transformer單詞就通過和句子中其它單詞的操作完成了對本身語義的重構。

我們接下來具體看看提取特征和加權求和這兩個操作。

提取特征

K提取的特征是如何獲得的呢？根據自注意力的思想和人腦的機制，我們需要先看過所有項才能準確地定義某一個項。因此，對于每一個查詢語句Q，注意力機制會：

用這個Query和每個Key通過內積的方式來計算出相似度或者相關性，以此來決定哪個元素會對目標元素造成多少影響（即Key和Query會得出對齊系數）。Key與Query越相似或者說越相關，Value的影響力就越大，越應該承擔更多的對輸入的預測。
然后注意力機制會對點積結果進行一個softmax操作，使得所有value的權重總和為1。這是為了保證所有源元素貢獻的特征總量保持一定。如果有多個key都與query高度相似，那么它們各自的通道都會只打開一部分（好像“注意力分散在這幾個源元素上”）。從這個角度來看，可以理解為輸出是在value之間根據key-query的相似度進行內插值的結果。很明顯，這個輸出表征攜帶了其它單詞的信息。
最后得到的矩陣 Y ，就是輸入矩陣 X 融合了上下文信息得到的在隱空間的語義矩陣，每一行代表一個token。

如果是訓練過程，則在拿到Y之后，模型會通過損失函數進行計算，最終經過反向傳播后，Q就能逐漸學習到V的特征。這個機制讓模型可以基于相同的注意力機制學習到不同的行為，并且能夠捕獲序列內各種范圍的依賴關系。

加權求和

讓我們通過一個真實例子來理解加權求和操作。

我們希望了解郁達夫（Query），但是因為人的精力有限，所以需要把有限的精力集中在重點信息上，這樣可以用更少的資源快速獲取最有用的信息，效果更好。圖書館（Source）里面有很多書（Value），我們看書就相當于獲取其書中的詳細信息（Value）。為了提升效率，我們給每本書做了編號和信息摘要（Key）。于是我們可以通過Key搜索出來很多書，比如《薄奠》，《沉淪》，《遲桂花》，《春風沉醉的晚上》，《歸航》等等，也能搜出來《大眾文藝》（郁達夫曾任主編）。

通過將Query與Key中攜帶的信息摘要相比較，我們可以知道它們的相關程度有多高。相關性越高的書，其權重越大。前面幾本書的權重就高，需要分配更多的注意力來重點看，《大眾文藝》的權重就顯然要略低，大致瀏覽即可。

假如我們一共要花費11小時在了解郁達夫上。我們會分別花費2小時在《薄奠》，《沉淪》，《遲桂花》，《春風沉醉的晚上》，《歸航》，花費1小時在《大眾文藝》上。我們把時間歸一化成和為1的概率值，得到[0.18, 0.18, 0.18, 0.18, 0.18, 0.09]，所以郁達夫 = 0.18《薄奠》+ 0.18《沉淪》+ 0.18《遲桂花》+ 0.18《春風沉醉的晚上》+0.18《歸航》+0.09《大眾文藝》。最終得到的信息是所有書籍內容按照權重綜合起來的結果。這樣當我們全部看完以上幾本書后，就對郁達夫有一個全面的了解，就是加權求和。

注意力機制本質就是使用 Q 和 K 來計算出“注意力權重“，然后利用注意力權重對Ｖ進行加權求和。從機制上看，注意力機制聚焦的過程體現在權重系數上，權重越大表示投射更多的注意力在對應的值上，即權重代表了信息的重要性。注意力機制可以被解釋為將多個局部信息源路由到一個局部表征的全局樹結構中。在這個例子中，我們計算相關性就相當于注意力機制中的 $QK^T$ ，歸一化就是softmax操作，然后通過加權求和取得最后的閱讀量/特征向量。

Elhage重構了注意力頭的表達形式（仍然等價于vanilla Transformer的設計），重構的表示可以表達為如下公式，可以看出來相互操作、提取特征和加權求和的特點。

1.4 小結

我們以”我吃了一個蘋果“為例來看看自注意力的流程。

首先確定哪個目標token來作自注意力機制，這個目標token是”一個“，即讓”一個“來判斷它和其他三個詞之間的關系
”一個“會對”我吃了一個蘋果“這句話中所有token都做一遍點積，然后做softmax操作做歸一化，生成權重。
用權重和V向量相乘，得到加權之后的向量，即”一個“可以用0.2 * 我、0.1 * 吃了、 0.5 * 一個和 0.3 * 蘋果來組合表示。所以”一個“的帶有詞關聯性的表示為：$一個 = 0.2 \times 我 + 0.1 \times 吃了 + 0.5 \times 一個 + 0.3 \times 蘋果$。在新的向量中，每一個維度的數值都是由幾個詞向量在這一維度的數值加權求和得來的，這個新的向量就是"一個"詞向量經過注意力機制加權求和之后的表示，該詞具備詞關聯性。
解碼輸出”an“。

可以看出，“一個”在全句中，除了自己之外，與“蘋果”關聯度最大，其次是“我”。所以“一個”這個詞也可以理解為“我-一個-蘋果”。這便把“一個”在這句話中的本質通過“變形”給體現出來了?！耙粋€”本身并沒有變，而是通過“變形”展示出了另外一種變體狀態“我-一個-蘋果”。外在沒變，靈魂變了。

可以看到，自注意力機制是一種動態的、數據驅動的變換，是對輸入向量空間的一種動態變換。這種變換不是固定的，而是依賴于輸入數據的內容來決定的。因此，注意力的本質思想可以改寫為如下公式：通過計算相似性得出權重最后加權求和。

\[Attention(Target,Source)=Attention(Query,Source)=\sum^{Length_{Source}}_{i=1}Similarity(Query,Key_i)*Value_i \]

0x02 實現

2.1 權重矩陣

現在的神經網絡很少有將 Word Embedding 直接參與一些網絡結構的計算，一般都會先做一個線性變換。實際上，Transformer是把每個 token 的 Embedding 向量x分別乘以三個不同的權重矩陣$W^T,W^Q,W^V$，作三次線性投影（或稱為線性變換），派生出Q、K、V三個矩陣（注意，這里提到的“權重”，是指神經網絡中的連接權重，與Attention中token之間的語義關聯權重不是一個意思）。而且，每個Transformer block都有自己的$W^T,W^Q,W^V$。

這三個權重矩陣是在模型訓練過程中通過反向傳播訓練出來的。在訓練階段，模型會對這三個權重矩陣進行隨機初始化。在模型的執行階段（預測階段），這三個矩陣是固定的，即 Transformer 神經網絡架構中固定的節點連接權重，是早就被預先訓練好的了（Pre-Trained）。$W^T,W^Q,W^V$這三個矩陣實際上是模型學會的分配Q，K，V的邏輯。

為什么要引入權重矩陣？或者說，為什么不直接使用 X 而要對其進行線性變換？為何要從一個序列中的每一個 token 的 Embedding 派生出三個向量Q、K、V（即查詢向量、鍵向量和值向量）呢？主要可以從如下方面進行思考。

首先看看直接使用embedding的缺點。

輸入的 Embedding 其實只做了一次線性變換，特征提取能力或者表示學習的能力及其有限。
一個點積操作中沒有什么可以學的參數。為了識別不一樣的模式，我們希望有不一樣的計算相似度的辦法以及更多的參數。
如果直接對原始的embedding做自注意力操作，則計算的相似度結果是個對稱矩陣，對角向上的值一定是最大的。因為每個字/詞必定最關心自己，這樣背離了自注意力操作的初衷。
- $Q*K^T$大概率會得到一個類似單位矩陣的attention矩陣，這樣self-attention就退化成一個point-wise線性映射，捕捉注意力的能力就會受限，i，j 位置之間的前后向注意力就會變得一樣。而我們一般期望兩個token在一句話中先后順序也能反映一定的不同信息。對于兩個詞語，A對于B的重要性，不一定等同于B對A的重要性。比如：”A愛B“和”B愛A“的程度不一定一樣。
- 這個對稱矩陣的對角線上的值大概率是本行最大的，這樣 softmax 后對角線上的注意力一定是本行最大，也就是不論任何位置任何搭配，每個token的注意力幾乎全在自己身上，這樣違背了Transformer用自注意力機制捕捉上下文信息的初衷。

我們再看看使用權重矩陣的優勢所在。

匹配。$\alpha_{ij}$要通過計算$?_i$和$s_j$之間的關系得到，一個最簡單的辦法就是把這兩個矩陣直接相乘。但是這樣可能會有問題：兩個矩陣可能形狀不匹配，沒法直接做矩陣乘法。而給這兩個矩陣分別左乘一個矩陣$W^K$和$W^Q$就可以解決上述兩個問題。
可學習。在注意力機制中，每一個單詞的query, key, value應該不僅僅只和該單詞本身有關，而應該是和對應任務相關。每個單詞的query, key, value不應該是人工指定的，而應該是可學習的。因此，我們可以用可學習的參數來描述從詞嵌入到query, key, value的變換過程。這樣經過大量訓練之后，每個元素都會找到完成各自任務所需的最合適的query、key和value。這些相關的訓練參數就在三個權重矩陣中。
增加擬合能力。三個權重矩陣都是可訓練的，這增加了模型可學習的參數量，擴展了特征空間，增加了模型的擬合能力。
信息交換。我們指定輸入矩陣X的第 i 行表示第 i 個時刻的輸入 $x_i$。對于此矩陣中的向量來說，權重矩陣$W^T,W^Q,W^V$在整體運行過程中是共享的。即，不同的$x_i$ 共享了同一個$W^T,W^Q,W^V$，通過這個操作，$x_1$和$x_2$已經發生了某種程度上的信息交換。也就是說，單詞和單詞通過共享權值已經相互發生了一定程度的信息交換。

另外，也許讀者會問，既然K和Q都是一樣維度，為什么不合用一個權重矩陣呢？或者說，為何要使用三個不同的權重矩陣？使用不同的權重矩陣生成的主要原因是：為了提供更靈活的模型表示能力和捕捉數據中的復雜依賴關系。其實，此處也回答了為何要區分Q、K和V。

增加表達能力。加入了不同的線性變換相當于對 x 做了不同的投影，將向量 x 投影到不同空間，這意味著Q和K可以在不同的語義空間中進行表達，有助于模型捕捉更豐富的語義信息和依賴關系。
區分不同角色或者說角色分離：在自注意力機制中，Q、K和V扮演著不同的角色。Q代表了我們要查詢的信息或者說當前位置希望獲得的信息，K代表了我們用來與Q匹配的鍵或者說序列中各位置能提供的信息，而V代表了一旦找到匹配，我們要提取的值或者說應該從各位置獲取的實際內容。使用不同的權重矩陣能夠更好地區分這些不同的角色，使用不同的權重矩陣為Q和K提供了能力去捕捉不同的依賴關系，增強了模型對輸入數據的理解。提高模型的效果。
增加靈活性。如果Q和K使用相同的權重矩陣，則其結果和使用X自身進行點乘的結果相同，那么它們之間的關系會被嚴格限制在一個固定的模式中，這限制了模型的靈活性。而直接從Q得到V會忽略了通過K來確定相關性的重要性，也減少了模型處理信息的靈活性。
并行處理。Transformer模型的設計允許在處理序列時進行高效的并行計算。Q、K、V的獨立使得模型可以同時計算整個序列中所有位置的注意力分數，這大大提高了計算效率。

總的來說，雖然在某些情況下使用相同的值進行自身的點乘（或者共享權重矩陣）可能也能工作得很好，但使用不同的權重矩陣為Q和K提供了更大的靈活性和表示能力，有助于提升模型性能和泛化能力。

我們再用三個權重矩陣來細化注意力公式如下。

2.2 計算過程

縮放點積注意力（Scaled Dot-Product Attention）模塊的公式如下：

此公式中，$d_k$是向量的維度，且$d_k=d_q=d_v$，如果只設置了一個頭，$d_k$那就是模型的維度$d_{model}$，如果設置了8個頭，則$d_k=d_{model}/8$，且如果模型的維度是512維，則 $\sqrt{d_k}$ 即等于8。Q和K的維度均是$(L,d_k)$，V的維度是$(L,d_v)$，其中L是輸入序列長度。$softmax(QK^T)$的維度是$(L,L)$，$Attention(Q,K,V)$的輸出維度是$(L,d_v)$。

我們梳理下計算過程如下（對應上圖中從下到上的順序）：

輸入。Q、K、V是把輸入映射成高維空間的點，它們之間的關系通過后續的變換來捕捉。
計算分數（score function）。Query和所有的Key進行相似度計算，得到注意力分數（查詢的相關性）。計算公式為$s_i = a(q, k_i)$。也就是Q 矩陣和 K矩陣的轉置之間做矩陣乘法（即點積）。這一步是計算在高維空間中度量向量之間的相似性。
縮放。對得分矩陣scores進行縮放，即將其除以向量維度的平方根$\sqrt{d_k}$。
掩碼。若存在掩碼矩陣，則將掩碼矩陣中值為True的位置對應的得分矩陣元素置為負無窮。這是由于在整個模型的運行過程中，可能需要根據實際情況來忽略掉一些輸入。我們將在下一篇進行詳細解釋。
歸一化（alignment function）。對點積結果進行歸一化，即使用softmax操作將權值進行歸一化，這樣可以更加突出重要的權重。計算公式為$a = softmax(s_i)$。這一步是將實數域的分數映射到概率分布上。
生成結果（context vector function）。使用a對Value進行加權平均的線性變換，可以理解為輸出y是在value之間根據key-query的相似度進行內插值。計算公式為$Attention\ Value=\sum_ia_iv_i$

從泛函分析的角度來看，Attention機制中的相關性計算和加權求和步驟可以看作是對輸入向量空間的一種動態變換。這種變換不是固定的，而是依賴于輸入數據的內容來決定的。

我們接下來對上面過程中的一些重點進行詳細梳理。

2.3 點積注意力函數

Transformer論文使用了乘法函數或者說點積注意力函數來計算相似度。從抽象代數的角度來看，注意力機制更像是一個關系運算：用兩個元素之間的關系（比如相似度）來決兩個元素是否屬于同一個類。分類之后才會基于輸入元素之間的相似性進行加權組合。

方案選擇

如下所示，常見的相似度計算有點積（相乘）和相加。

\[sim(q,k) = q^Tk\ [內積相似度] \\ sim(q,k)=w^T[q;k]\ [拼接相似度] \\ sim(q,k)=\frac{q^Tk}{||q||\ ||k||}\ [余弦相似度] \]

其中拼接相似度是將兩個向量拼接起來，然后利用一個可以學習的權重 ?? 求內積得到相似度，也稱為Additive Attention，意思是指的 $w^T[q;k]=w_1^Tq+w_2^Tk$。

對應到V則是：

\[attention(q,k,v) = W_v\ (tanh(W_k + W_q)) \ [相加] \\ attention(q,k,v) = W_v\frac{(W_q)(W_k^T)}{\sqrt d}\ [相乘] \]

從公式中可以看到，加法方案優點是：

可以處理不同維度的key與query（而點積操作要求query和key具有相同的長度）。
計算更簡單，但是外面套了tanh和v，相當于一個完整的隱層，因此整體復雜度其實和乘法方案接近。
雖然在大多數任務中，點乘注意力和加法注意力的性能差異不大，但是在某些長序列任務中，加法注意力可能會略優于點乘注意力。

而且，乘法方案還有一個劣勢是：隨著向量維度的增大，點乘結果的上限越來越高，點乘結果的差異越來越大，因此計算Attention權重需要加入scaled。

從表現效果來講，論文”Massive Exploration of Neural Machine Translation Architectures“對此做了對比實驗。

從結果上可以看得出，加法式注意機制略微但始終優于乘法式注意力機制。

那么為何Transformer為何選取點積注意力而非用加法注意力？網絡上也有一些討論，給出的一些思考點是：

點乘操作可以通過矩陣乘法高效地在硬件上并行化，從而實現快速計算。
點乘注意力能夠捕捉查詢和鍵之間的相似度，當查詢和鍵相似時給予更高的權重，這有助于模型捕捉輸入序列中的復雜依賴關系。
在將表示分割成不同個頭進行運算時，使用點乘會更加靈活方便計算。

論文中給出的原因是基于效率和建模能力的考慮，具體如下：

解讀

我們對公式作進一步的解讀。點積是兩個向量的夾角，表征一個向量在另一個向量上的投影。投影值越大，說明兩個向量的相關性越高，如果兩個向量的夾角為90度，則這兩個向量線性無關，完全沒有相關性。實際上，點積計算的是對齊后的長度的乘積。因為在機器翻譯中，這個向量是詞向量，是詞在高維空間的數值映射，而詞向量之間的高度相關性說明在一定程度上，在關注當前詞A的基礎上，也會給相似詞B更多的關注。

我們再來看看蘇劍林大神對公式的精彩解讀，借此可以對公式有更加深深刻的理解。

將$QK^T$進行拆解，得到兩個向量的乘積為：$q_i?k_j=∥q_i∥∥k_j∥cos(q_i,k_j)$，即將兩個向量的乘積分解為了兩個向量各自模長與夾角余弦的乘積。其中：

$∥q_i∥$只跟當前位置i有關，因此它不改變注意力的相對大小，而只改變稀疏程度。
$∥k_j∥$是其它位置的的張量模長，有能力改變條件概率$??(??|??)$的相對大小，但它不涉及到i,j的交互，只能用來表達一些絕對信號。
$cos(q_i,k_j)$就是用來表達??,??的交互，是自由度最大的一項。

為了提高某個位置j的相對重要性，模型有兩個選擇：

增大模長$∥k_j∥$。
增大$cos(q_i,k_j)$，即縮小$q_i$，$k_j$的夾角大小。

然而，由于“維度災難”的存在，在高維空間中顯著地改變夾角大小相對來說沒有那么容易，所以如果能靠增大模長$∥k_j∥$完成的，模型會優先選擇通過增大模長$∥k_j∥$來完成，這導致的直接后果是：$cos(q_i,k_j)$的訓練可能并不充分（指被訓練過的夾角只是一個有限的集合，而進行長度外推時，它要面對一個更大的集合，從而無法進行正確的預測），這可能是Attention無法長度外推的主要原因。

2.4 softmax

定義

Softmax操作的意義是歸一化。我們將沒有做softmax歸一化之前的結果稱為注意力分數，將注意力分數經過softmax歸一化后的結果稱為注意力權重。給定一個包含 $n$ 個實數的向量 $\mathbf{x} = [x_1, x_2, \ldots, x_n]$，Softmax函數將其轉換為一個概率分布 $\mathbf{p} = [p_1, p_2, \ldots, p_n]$，其中每個 $ p_i $ 的計算公式為：

\[p_i = softmax(x_i) = \frac{e^{x_i}}{\sum_{j=1}^n e^{x_j}} \]

Softmax關鍵性質如下：

非負性：對于任意 $i$， $p_i \geq 0$。
歸一化：所有輸出的和為1，即 $\sum_{i=1}^n p_i = 1$。
指數函數的使用：指數函數 $e^{x_i}$ 確保了輸出值為正，并且放大了較大的 $x_i$ 值的差異。

算法

算法流程需要兩個循環，首先需要迭代計算分母的和，然后再迭代計算向量中每一個值對應的softmax值，即縮放每一個元素。這個過程需要兩次從內存讀取和一次寫回內存操作。具體算法如下，其中$d_V$就是分母。

具體圖例如下。

必要性

softmax函數之所以在神經網絡中得到廣泛應用，是如下原因：

softmax在推理時可以將輸出層的原始輸出（logits）轉換為有效的概率分布，使得輸出具有概率的物理意義，便于解釋和計算損失函數，同時保持數值穩定性和對稱性。這使得softmax成為多類分類問題中的首選激活函數。
在訓練時指導模型學習。在神經網絡的訓練過程中，通常會使用交叉熵損失函數來配合softmax函數。交叉熵損失函數能夠量化模型輸出的概率分布與真實標簽之間的差異，而softmax函數的輸出提供了一個概率分布。這種機制使得模型能夠在訓練過程中更有效地調整權重，以提高對真實概率分布的估計準確性。

為何要在注意力機制中加入softmax？因為看起來softmax是有害無益，比如：

如果沒有Softmax，則計算復雜度會大幅度降低。因為去除softmax之后，我們得到三個矩陣連乘$QK^?V$，而矩陣乘法是滿足結合率的，所以我們可以先算$K^?V$，得到一個d×d的矩陣，然后再用Q左乘它，由于d?n，所以這樣算大致的復雜度只是O(n)。即，去掉Softmax的Attention的復雜度可以降到最理想的線性級別O(n)。
如果沒有Softmax，則內存占用會大幅度降低。這里我們提前看看FlashAttention所要解決的困境。如果沒有softmax的話，我們可以對矩陣采用分塊（Tiling）計算。比如，我們可以把Q，K，V沿著N（seqence length維度）切成塊，算完一塊Q和一塊$K^T$之后，立刻和一塊V進行矩陣矩陣乘法運算（GEMM）。一方面，避免在HBM和SRAM中移動P矩陣，另一方面，P矩陣也不需要被顯式分配出來，消除了HBM中$O(N^2) $級別的存儲開銷。

在注意力機制中加入softmax是因為其有如下優點或者功能。

Softmax操作實質上是在量化地衡量各個詞的信息貢獻度。因為在注意力機制中，我們在直覺上是希望關注語義上相關的單詞，并弱化不相關的單詞。這其實是一個多分類問題。在多分類問題中，我們希望神經網絡的輸出可以反映每個類別的概率，即每個輸出節點的值代表了相應類別的概率。這要求輸出值必須滿足兩個條件：首先，每個輸出值都應該在0到1之間；其次，所有輸出值的和應該等于1。這樣，輸出就可以被解釋為概率分布。直接的線性歸一化雖然可以滿足第一個條件，但往往不能滿足第二個條件，因為它沒有考慮分值間的相對差異，不能反映出原始分值中的相對強度或置信度。而softmax函數恰好能夠同時滿足這兩個條件。

Softmax函數通過對每個分值應用指數函數，然后對這些指數值進行歸一化處理來轉換為概率，這樣既保證了每個輸出值在0到1之間，又保證了所有輸出值之和為1。更重要的是，指數函數的使用放大了分值之間的差異，更好的反映了原始分值中的相對置信度。

在給定一組數組成的向量，Softmax先將這組數的差距拉大（由于exp函數），然后歸一化，它實質做的是一個soft版本的argmax操作，或者當作argmax的一種平滑近似。與argmax操作中暴力地選出一個最大值（產生一個one-hot向量）不同，softmax將這種輸出作了一定的平滑，即將one-hot輸出中最大值對應的1按輸入元素值的大小分配給其他位置。因此得到的向量接近一個one-hot向量（接近程度根據這組數的數量級有所不同）。

另外，KQ 兩個矩陣相乘，線性乘積下秩不會超過d，softmax后會有一定程度的增秩效果，如果不使用softmax，線性attention的秩更低，表達能力也更差。

缺點

Softmax也存在一些固有限制。比如論文" softmax is not enough (for sharp out-of-distribution)" 指出：softmax函數在輸入規模增大時，其輸出系數會趨于均勻分布（注意力分散）。即，即便這些token的注意力系數在分布內是尖銳的（even if they were sharp for in-distribution instances），輸入更多的token會導致注意力更加分散（或者說注意力的熵變大），從而導致訓練和預測的結果不一致。具體如下圖所示。

該論文提醒我們，即使是像softmax這樣廣泛使用的函數，也可能存在其適用范圍的局限性，尤其是在處理超出訓練分布的數據時。我們需要更加重視模型的泛化能力，并探索更魯棒的模型架構。

改進

人們對softmax也有很多改進。

Log-Softmax

標準的softmax公式涉及到了很多的求冪和除法，導致計算成本較高，我們可以通過對數值取log，使得計算成本降低。具體如下：

\[??????\ ??????????????(??_??)=log\frac{e^{x_i-c}}{\sum_{j-1}^d e^{x_i-c}} = x_i - c - log{\sum_{j-1}^d e^{x_i-c}} \]

這種方式叫作Log-Softmax。與Softmax 相比，使用 Log-Softmax 有許多優點，包括提高數值性能和梯度優化等。這些優勢對于實現非常重要，尤其是當訓練模型在計算上的成本很高時，其能帶來很客觀的收益。而且log 概率的使用，具有更好的信息理論可解釋性，當用于分類器時，Log-Softmax 在模型無法預測正確的類別時會懲罰模型。

Hierarchical Softmax(H-Softmax)

標準softmax的時間復雜度都是 ??(??) ，這在分類任務中可能影響不大，只需要在最后的一層后面進行一次 ??(??) 計算就可以了。但是在NLP的生成任務中，我們需要對詞表大小的向量進行Softmax，是為了把每個詞的Softmax值用作似然，從而挑選出似然最大的那個詞。一旦詞表特別大，計算量將是一個嚴峻的問題，每次預測一個token都需要O(|V|)的時間復雜度。所以需要對Softmax進行一定的改造來適應實際任務，而H-Softmax就是用來解決這個問題的。

H-Softmax就是Word2vector中的Hierarchical Softmax。H-Softmax的解決方案是將Huffman Tree融入進來，將原先用 softmax 做多分類分解成多個sigmoid，然后使用Logistic Regression判斷在哈夫曼樹中走左子樹還是右子樹，最后其輸出的值就是走某一條樹分支的概率。

分層softmax使用輸出層的二叉樹表示，假設 ?? 個詞分別作為葉子節點，每個節點都表示其子節點的相對概率。詞表中的每個詞都有一條從二叉樹根部到自身的路徑。用 ??(??,??) 來表示從根節點到 ?? 詞這條路徑上的第 ?? 個節點，用 ???(??) 來表示 ?? 的任意一個子節點，設如果 ?? 為真則 [??]=1 ，否則 [??]=?1，那么 Hierarchical Softmax 可以表示為下圖。

adaptive softmax

常見的對softmax改進方法可以大致區分為兩類：一是準確的模型產生近似概率分布，二是用近似模型產生準確的概率分布，這些都是從數理角度進行優化。但是從硬件的角度，應該如何配合GPU進行優化呢？

論文“Efficient softmax approximation for GPUs"提出的Adaptive Softmax就給出了一些思路。Adaptive Softmax借鑒的是Hierarchical Softmax和一些變型，與以往工作的不同在于，該方法結合了GPU的特點進行了計算加速，這樣可以提高softmax函數的運算效率，適用于一些具有非常大詞匯量的神經網絡。

想法很簡單，文章的大多數詞都由詞匯表里的少數詞構成，即長尾理論或28原則。而語言模型在預測詞的時候往往需要預測每個詞的概率（通常是softmax），詞匯表可能非常大，低頻詞非常多。那么就可以利用單詞分布不均衡的特點，在訓練時把詞語分成高頻詞和低頻次兩類，先預測詞屬于哪一類，然后再預測具體是哪個詞，這樣簡單的分類就使得softmax的計算量大大減少。原來是每個詞都要計算一次，現在是：V(高頻) + P(低頻) * V(低頻)，V(高頻) 會大幅變小，V(低頻)雖然大但是P(低頻) 很小。另外，Adaptive Softmax也通過結合現代架構和矩陣乘積操作的特點，通過使其更適合GPU單元的方式進一步加速計算。

Adaptive Softmax的層次模型被組織為（i）第一層，包括最頻繁的單詞和表示聚類的向量.（ii）第二層上是與罕見單詞相關的聚類，最大的聚類的與最不頻繁的單詞相關。具體如下圖所示，藍色表示高頻詞（總體分作一類），白色是表示低頻詞（這里有三個白色框，就代表三類低頻詞），每個time step中，先預測當前詞是哪一類（高頻詞還是低頻詞分類），然后再在所得的分類中進行Softmax計算，從而得到最終的結果。

那么如何進行分類呢？論文也給了一個計算的模型，假設B是batch size，d是hidden layer 的size，k 是feature的大小，具體見下圖上的公式。

2.5 縮放

縮放點積注意力（Scaled Dot-Product Attention）中的Scaled是縮放的意思，是在點乘之后除以一個分母。因此我們提出了一個問題，注意力的計算公式中 $\text{Attention}(Q,K,V) = \text{softmax}(\frac{QK^\top}{\sqrtw0obha2h00}) $為什么要除以 $\sqrtw0obha2h00 $？

結論

我們細化原始論文的解釋如下：當維度$d_k$較大時，$q$和$k$的的點積會容易出現較大數值（正比于維度），當這些大數值經過softmax函數時，點積結果的分布會趨近于陡峭（分布的方差大，分布集中在絕對值大的區域）。進而會把點積結果推向softmax函數的梯度平緩區，導致 softmax 函數的梯度將變得非常小，這意味著模型將難以收斂，加大學習難度。因此Transformer作者將乘法函數按照因子$1/\sqrt{d_k}$進行縮放，這樣$softmax(QK^T)$分布的陡峭程度就和$d_k$解耦了。正好抵消了維度增加造成的點積尺度放大效應，可以保證無論維度$d_k$取什么值，點乘的結果都處在一個合理的范圍內，從而有助于保持梯度的穩定性，加速模型的訓練過程。

對于較大數值這部分，可以參見論文腳注。

$\sqrt d$ 很像超參數 Temperature。Temperature 用于調整模型的 softmax 輸出層中預測詞的概率，控制生成文本的隨機性和創造性。

問題推導

我們把上述問題點梳理如下：

如果 $d_k$ 變大，$q \cdot k^\top $ 方差會變大。
方差變大會導致向量之間元素的差值變大。
元素的差值變大會導致 softmax 退化為 argmax，即最大值為 1，其他值則為0。
如果只有一個值為 1，其他都為 0 的話，反向傳播的梯度會變為 0, 也就是所謂的梯度消失。

我們接下來逐一分析。

方差變大

從根本上說，當我們將兩個矩陣相乘（例如 Q 和 $K^T$）時，我們對它們各自的列向量進行多次點乘。例如，我們將第一個向量的第一個值與第二個向量的第一列相乘，以此類推，得到輸出矩陣中的值。

當$q$和$k_l$中每個元素都是正態分布（其均值為零，方差為1）時，則$q$和$k_l$的的點積的方差為$d_k$。推導如下：

拿出Q矩陣中單獨一列$q_i$和K矩陣單獨一行$k_i$出來，假設$q_i$和$k_i$中每個元素均是均值為0，方差為1的獨立同分布隨機變量，則$q_ik_i^T$中每個元素也是均值為0，方差為1。也就是說，每個這樣的qk對都會產生一個期望為0，方差為$d_k$的新向量。當$d_k$很大時，這個向量的方差就很大。即，點積方差正比于維度數量，當我們對低維向量進行點積時，輸出的方差往往較小。當我們對高維向量進行點積運算時，輸出的數字往往具有較高的方差。

\[var[q_i \cdot k_i^T] = var[\sum^{d_k}_{i=1} q_i \times k_i] = \sum^{d_k}_{i=1} var[q_i \times k_i] = \sum^{d_k}_{i=1} var[q_i] \times var[k_i] = \sum^{d_k}_{i=1}1 = d_k \]

而在現實世界中，表示復雜語義概念需要很高的query和key維度，高維度導致點積$qk^T$值變大。

元素間差值變大

方差變大會導致向量之間元素的差值變大。這是一個顯而易見的結論，因為方差變大就是代表了數據之間的差異性變大。有些數字會非常大，而有些則會非常小。

softmax退化

softmax中每個分量$softmax(x_i)$如下：

\[softmax(x_i) = \frac{e^{x_i}}{\sum _{j=1}^V e^{x_j}} \]

將x的每個元素表達為最大元素$x_{max}$減去一個插值$\Delta_i$，即$x_i = x_{max} - \Delta_i$，則softmax重寫為：

\[softmax(x_i) = \frac{e^{x_i}}{\sum _{j=1}^V e^{x_j}} = \frac{e^{x_{max}-\Delta_i}}{\sum _{j=1}^V e^{x_{max}-\Delta_j}} \]

因為$e^{x_{max}}$是公因子，可以提出來約掉，則

\[softmax(x_i) = \frac{e^{x_i}}{\sum _{j=1}^V e^{x_j}} = \frac{e^{x_{max}-\Delta_i}}{\sum _{j=1}^V e^{x_{max}-\Delta_j}} = \frac{e^{x_{max}}e^{-\Delta_i}}{e^{x_{max}}\sum_{j=1}^V e^{-\Delta_j}} = \frac{e^{-\Delta_i}}{\sum_{j=1}^V e^{-\Delta_j}} \]

如果$x_i$遠小于$x_{max}$時，$e^{-\Delta_i}$會接近于0，因此除了$\Delta_{max} = 0$之外所有項都接近于0，即：

\[softmax(x_{max}) \approx \frac{1}{1} = 1 \\ softmax(x_i) \approx 0 \ \ \ \ \ i \neq max \]

所以方差變大時，softmax 函數會退化為 argmax 函數。當數字較大時，退化的SoftMax 傾向于分配較高的概率；當數字較小時，SoftMax 傾向于分配較低的概率。最終，退化的softmax 函數會把大部分概率分布分配給最大的元素，這趨近于把最大的元素賦值為1，其它元素賦值為0。

即softmax將one-hot輸出中最大值對應的1按輸入元素值的大小分配給其他位置。而當喂入的數組內部數量級相差較大時，“1分出去的部分”就會越來越少，當數量級相差到一定程度，softmax將幾乎全部的概率分布都分配給了最大值對應的標簽，其效果也就被削減了。

梯度消失

假設我們的輸入數量級很大，那么softmax將產生一個接近 one-hot 的向量。$Y≈[0，...,1,?,0]^?$，求導如下

\[\frac{\partial \mathbf{y}}{\partial \mathbf{x}} = \mathbf{diag}(\mathbf{y}) - \mathbf{yy}^\top \]

上面公式運算結果就接近于0，即幾乎所有梯度都接近于0。所以，絕對值很大的點積在訓練中會收到幾乎為0的梯度，這對于基于梯度下降法的優化非常不利。在反向傳播過程中，模型會不斷調整和改變較大的數字，而概率較小的數值幾乎不會得到更新。因此，由于并非所有參數都會更新，這會導致訓練進展緩慢。

如何降低方差？

因此，在應用 SoftMax 之前，我們需要找到一種解決方案，幫助減少這些數字的方差?，F在，我們的問題已經簡化為降低包含初始注意力分數的乘積矩陣的方差。

vanilla Transformer提出的方案就是縮放矩陣以獲得與之前相同的方差。因為除以維度的平方根意味著，即使增加維度（需要增加維度來捕捉更復雜的模式），我們也可以在增加維度的同時縮放矩陣以保持方差一致。

熵的作用

信息熵衡量的是不確定性。在注意力機制中，熵可以用來度量模型輸出的不確定性，或者說，某個查詢向量的注意力權重分布的熵可以用來衡量它對不同鍵向量的關注程度。

高熵表示模型的注意力分散，即這個查詢向量在多個輸入數據（鍵向量）上都有較高的注意力權重。模型在多個輸入之間分配了注意力，因此模型只能從多個地方提取信息而沒有明確的重點。因此，高熵分布可以捕捉更廣泛的上下文信息，對于需要全局信息的場景更加合適。
低熵表示模型的注意力集中，即這個查詢向量在少數幾個輸入數據（鍵向量）上有較高的注意力權重，模型對這些鍵向量的選擇更加明確和確定。因此，低熵分布傾向于選擇性地提取信息，對于需要聚焦于關鍵信息的場景更加合適。

為了使得模型結果能夠更好地泛化到未知長度，Attention機制的設計應該使得$a_{i,j}$盡量具備熵不變性。

熵不變性是指，熵值????應該對長度??不敏感。更具體一點，就是如果在已有的token基礎上，再補充幾個token，那么新算出來各個$??_{??,??}$自然也會有所改變，但我們希望????不要有太大改變。

換個角度想，我們可以將不確定性視為注意力的“聚焦程度”：如果熵為0，那么注意力將聚焦到某一個token上，如果熵為logn，那么注意力均勻分布到所有token上。我們希望熵不變，是希望引入新的token后，已有的token依舊能同樣地聚焦到原來的token上，而不希望新token的引入過多地“分攤”了原有的注意力，導致求和結果顯著發生變化。使用$log???$縮放注意力可以在一定程度上緩解這個問題，即將Attention從

\[Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt d})V \]

修改為

\[Attention(Q,K,V)=softmax(\frac{log_m n}{\sqrt d}QK^T)V \]

其中??是訓練長度，??是預測長度。經過這樣修改，注意力的熵隨著長度的變化更加平穩。

2.6 小結

最后，我們梳理自注意力機制的計算過程如下圖所示。

0x03 實現

3.1 哈佛代碼

attention()函數定義了標準注意力機制的操作過程，計算公式是：Attention(q,k,v) = softmax( ${q \times k^T}\over{\sqrt d_k}$ ) * v，這也是整個transformer的核心操作所在。

輸入&輸出

首先要說明的是，本小節下面行文和注釋中的一些術語解釋如下。

batch_size：樣本有多少個句子。
seq_length是句子長度。
d_model是模型的維度。
head_num是注意力頭數。
head_dim是單個頭的注意力維度，大小是d_model / head_num。
embedding_size：詞嵌入的大小。

其次，attention()函數的輸入參數如下：

query，key，value是輸入的向量組，就是論文公式提到的，經過$W^Q, W^K, W^V$計算后的Q, K, V，計算過程位于MultiHeadedAttention類的forward()函數之中。query，key，value的形狀有兩種可能：
- 單注意力頭，即自注意力調用到attention()函數，則形狀是(batch size, seq_length，d_model)。
- 多注意力頭，即多頭注意力機制調用到attention()函數，則形狀是(batch size, head_num，seq_length，head_dim)。Transformer中使用的是多頭注意力機制，所以query，key，value的Shape只會是第二種。
mask：用于遮掩某些位置，防止計算這些位置的注意力。mask有兩種，一種是src_mask，另一種是tgt_mask。
dropout：dropout率，用于添加隨機性，有助于防止過擬合。

attention()函數的輸出有兩個：

torch.matmul(p_attn, value)：value 的加權平均，權重來自p_attn。
p_attn：從 query 和 key 的計算結果得到的權重，后續沒有用到。

這里要對mask做一下說明，mask有兩種：

src_mask：形狀是(batch size, 1, 1, seq_length)。因為所有head的mask都一樣，所以第二維是1，masked_fill 時使用 broadcasting 就可以了。這里是 self-attention 的mask，所以每個時刻都可以 attend 到所有其它時刻的第三維也是 1，也使用broadcasting。
tgt_mask：形狀是(batch size, 1, seq_length, seq_length)。

圖例&代碼

我們可以通過下圖來和代碼互相印證。

對應代碼如下。

def attention(query, key, value, mask=None, dropout=None):
    """
    本函數計算縮放點積注意力（Scaled Dot Product Attention）
    query, key, value：是經過權重矩陣線性轉換過的Q,K,V矩陣，具體線性轉換是在后續會介紹的  MultiHeadedAttention類中。query, key, value有兩種可能的形狀：
    1. 若注意力為單頭自注意力，則形狀為(batch size, 詞數, d_model)。
    2. 若注意力為多頭自注意力，則形狀為(batch size, head數, 詞數，d_model/head數)
    在哈佛代碼中使用的是多頭自注意力，所以query, key, value的形狀只會是第二種。
    mask（掩碼矩陣）：有兩種參數，一種是src_mask，另一種是tgt_mask。   
    """
    
    """
    用query最后一個維度的大小來給d_model賦值。之所以這樣可以獲取，是因為query和輸入的shape相同，
    若注意力為單頭自注意力，則最后一維都是詞向量的維度，也就是d_model的值。
    若注意力為多頭自注意力，則最后一維是 d_model / h，h為head數    
    """
    d_k = query.size(-1) 
     
    """
    執行QK^T/√d_k，得到注意力分數。
    1. key.transpose(-2, -1)將將最后兩個維度進行轉置，得到key的轉置矩陣K^T。
    2. torch.matmul()函數做矩陣乘法，即將query矩陣的每一行與key的轉置矩陣的每一列進行點積（對應元素相乘并求和），得到新矩陣的每個元素。此處操作對于上圖的標號1。
    3. math.sqrt(d_k)操作會對矩陣相乘結果進行縮放。此處操作對于上圖的標號2。
    scores是一個方陣, 形狀為(batch_size, head數，詞數，詞數)  
    """
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
  
    # 判斷是否使用掩碼張量
    if mask is not None:
        # 如果使用掩碼，則屏蔽不想要的元素
        # masked_fill()將掩碼張量和scores張量每個位置一一比較, 如果掩碼張量處為0，則把注意力分數中對應元素設置為-1e9，因為后續還要進行softmax操作，softmax會讓負無窮變為0（是理論上還是用到了很少一點點未來的信息，因為畢竟還是有一個小小的數值）。此處對于上圖的標號3。       
        scores = scores.masked_fill(mask == 0, -1e9)

    """
	執行公式中的softmax，對scores的最后一維做歸一化操作，得到注意力權重，這些權重值加起來的和為1。此處對于上圖的標號4。 
    p_attn的形狀如下：
    1. 若p_attn是自注意力，則形狀為(batch size, seq_length, seq_length)
    2. 若p_attn是多頭注意力，則形狀為(batch size, head_num, seq_length，seq_length)  
    # 這樣獲得最終的注意力張量    
    """
    p_attn = scores.softmax(dim=-1) # 得到注意力權重
    
    # 判斷是否使用dropout進行隨機置0    
    if dropout is not None:
        # 如果提供了dropout，則對p_attn進行dropout操作。
        p_attn = dropout(p_attn)        
 
    """
    用注意力權重p_attn對value向量進行加權求和，得到最終的輸出softmax(QK^T/√d_k)V，我們命名為Z。
    1. 對于自注意力，Z的形狀為(batch size, seq_length, d_model)，即最終結果。
    2. 對于多頭注意力，Z的形狀為(batch size, head_num, seq_length，d_model/head_num),  而并非最終結果，后續在MultiHeadAttention的forward()函數中還要將多頭的結果進行合并，變為(batch size, seq_length, d_model)
    torch.matmul對應上圖的標號5
    """
    return torch.matmul(p_attn, value), p_attn # 返回Z和權重p_attn

再分析注意力

看過源碼之后，我們再來對編碼器和解碼器中的注意力機制進行回顧和溫習。

encoder使用自注意力的目的是：找到輸入序列自身的關系。

源序列之中，每個token都搜集到本字和源序列之中其他哪幾個字比較相關。得到相關權重矩陣score。
接下來對 score 求 softmax，把得分變成概率 p_attn.
利用p_attn，把源序列轉換為源隱狀態，具體操作是torch.matmul(p_attn, value)，其中每個token都是綜合了源序列之中所有token的相關信息。
最后返回的是源隱狀態torch.matmul(p_attn, value)和p_attn 。源隱狀態會作為參數Memory傳給解碼器。

decoder使用兩種注意力結構。

使用self-attention的目的是找到目標序列自身的關系。
- 讓目標序列之中，每個token都搜集到本字和目標序列之中其他哪幾個字比較相關。得到相關權重矩陣score。
- 接下來對 score 求 softmax，把得分變成概率 p_attn.
- 利用p_attn把目標序列轉換為目標隱狀態，具體操作是torch.matmul(p_attn, value)，其中每個token都是綜合了源序列之中所有token的相關信息。
- 最后返回的是目標隱狀態torch.matmul(p_attn, value)和p_attn 。
使用cross-attention的目的是讓源序列與目標序列對齊
- 用目標隱狀態作為 Q，源隱狀態（參數Memory)作為 K, V。
- 找到目標隱狀態之中每個token與源隱狀態之中哪幾個token比較相關。得到相關權重矩陣（就是score）。
- 接下來對 score 求 softmax，把得分變成概率 p_attn。
- 利用p_attn把目標隱狀態轉換成新的目標隱狀態，具體操作是torch.matmul(p_attn, value)，其中每個token都是綜合了源隱狀態之中所有token的相關信息。
- 最后返回的是目標隱狀態torch.matmul(p_attn, value)，p_attn。

3.2 llama3

我們再用工業界的代碼來進行學習。首先我們給出Transformer總體代碼如下。Transformer是整個模型的主體，它將詞嵌入層、TransformerBlock層、歸一化層和輸出層串聯起來。從下面代碼中可以看到，Transformer中包含了很多層TransformerBlock。每層都有自己的注意力機制。每層注意力機制中都有自己的$W^Q, W^K, W^V$參數。所以也有一種說法是attention是基于圖的一個信息傳遞機制，節點之間的邊是學習到的。

class Transformer(nn.Module):
    def __init__(self, params: ModelArgs):
        super().__init__()
        self.params = params
        self.vocab_size = params.vocab_size
        self.n_layers = params.n_layers

        # 詞嵌入層
        self.tok_embeddings = VocabParallelEmbedding(
            params.vocab_size, params.dim, init_method=lambda x: x
        )

        # 將32個TransformerBlock存于ModuleList
        self.layers = torch.nn.ModuleList()
        for layer_id in range(params.n_layers):
            self.layers.append(TransformerBlock(layer_id, params))

        # 歸一化層
        self.norm = RMSNorm(params.dim, eps=params.norm_eps)
        # 輸出層，輸入特征數為詞嵌入的維度，輸出特征數為詞表大小
        self.output = ColumnParallelLinear(
            params.dim, params.vocab_size, bias=False, init_method=lambda x: x
        )

        # 旋轉位置編碼中的旋轉矩陣
        self.freqs_cis = precompute_freqs_cis(
            params.dim // params.n_heads,
            params.max_seq_len * 2,
            params.rope_theta,
        )

    @torch.inference_mode()
    def forward(self, tokens: torch.Tensor, start_pos: int):
        # 批次大小和序列長度
        _bsz, seqlen = tokens.shape
        # 應用詞嵌入之后的輸入h，大小為(bsz, seqlen, dim)
        h = self.tok_embeddings(tokens)
        # 確保旋轉矩陣和輸入h位于同一設備(GPU)
        self.freqs_cis = self.freqs_cis.to(h.device)
        # 從旋轉矩陣中提取旋轉角度（頻率）
        freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]

        mask = None
        if seqlen > 1:
            # 創建大小為(seqlen, seqlen)的張量，并用“負無窮大”填充
            mask = torch.full((seqlen, seqlen), float("-inf"), device=tokens.device)
            mask = torch.triu(mask, diagonal=1) # 上三角化

            # When performing key-value caching, we compute the attention scores
            # only for the new sequence. Thus, the matrix of scores is of size
            # (seqlen, cache_len + seqlen), and the only masked entries are (i, j) for
            # j > cache_len + i, since row i corresponds to token cache_len + i.
            # 將大小為(seqlen, start_pos)的全0張量和mask進行水平拼接，形成最終的mask
            mask = torch.hstack(
                [torch.zeros((seqlen, start_pos), device=tokens.device), mask]
            ).type_as(h)

        # 輸入h被32個TransformerBlock逐個處理
        for layer in self.layers:
            h = layer(h, start_pos, freqs_cis, mask)
        # 將最后一個TransformerBlock的輸出進行歸一化
        h = self.norm(h)
        # 將歸一化之后的結果經過輸出層，進行線性計算
        output = self.output(h).float()
        return output # 大小為(bsz, seqlen, vocab_size)

接下來我們看看每個層的代碼。

class TransformerBlock(nn.Module):
    def __init__(self, layer_id: int, args: ModelArgs):
        super().__init__()
        self.n_heads = args.n_heads
        self.dim = args.dim
        self.head_dim = args.dim // args.n_heads
        self.attention = Attention(args)
        self.feed_forward = FeedForward(
            dim=args.dim,
            hidden_dim=4 * args.dim,
            multiple_of=args.multiple_of,
            ffn_dim_multiplier=args.ffn_dim_multiplier,
        )
        self.layer_id = layer_id
        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)

    def forward(
        self,
        x: torch.Tensor,
        start_pos: int,
        freqs_cis: torch.Tensor,
        mask: Optional[torch.Tensor],
    ):
        h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
        out = h + self.feed_forward(self.ffn_norm(h))
        return out

然后是Attention的代碼如下。

class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        # 注意力中K(Key)和V(Value)的頭數，n_kv_heads=n_heads時為多頭注意力，n_kv_heads=1時為多查詢注意力，否則為分組查詢注意力
        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        # 分布式訓練的進程數，也就是模型訓練的并行度或參與計算的設備(GPU)數量
        model_parallel_size = fs_init.get_model_parallel_world_size()
        # 將Q(Query)的頭數，根據并行度進行拆分
        self.n_local_heads = args.n_heads // model_parallel_size
        # 將K(Key)和V(Value)的頭數，根據并行度進行拆分
        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
        # 假如本地Q數為2，本地KV數為1，我們需要把本地KV復制為2份，以便矩陣相乘
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        # 每頭(Q)的輸出維度。每頭的輸出會concat，來生成注意力的最終輸出
        self.head_dim = args.dim // args.n_heads

        # ColumnParallelLinear和RowParallelLinear是用于并行訓練的兩種并行線性層
        self.wq = ColumnParallelLinear(
            args.dim, # 輸入的特征數
            args.n_heads * self.head_dim, # 輸出的特征
            bias=False, # 注意力里面的線性計算一般不使用偏置單元
            gather_output=False,
            init_method=lambda x: x,
        )
        # 大小為(dim, n_kv_heads * head_dim)
        self.wk = ColumnParallelLinear(
            args.dim,
            self.n_kv_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        # 大小為(dim, n_kv_heads * head_dim)
        self.wv = ColumnParallelLinear(
            args.dim,
            self.n_kv_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        # 大小為(n_heads * head_dim, dim)
        self.wo = RowParallelLinear(
            args.n_heads * self.head_dim,
            args.dim,
            bias=False,
            input_is_parallel=True,
            init_method=lambda x: x,
        )

        self.cache_k = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )
        ).cuda()
        self.cache_v = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )
        ).cuda()

    def forward(
        self,
        x: torch.Tensor,
        start_pos: int,
        freqs_cis: torch.Tensor,
        mask: Optional[torch.Tensor],
    ):
        bsz, seqlen, _ = x.shape
        # 輸入x分別和權重wq、wk和wv進行線性計算，得到xq、xk和xv
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

        # 改變xq、xk和xv的形狀
        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

        # 對xq和xk應用旋轉位置編碼
        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

        self.cache_k = self.cache_k.to(xq) # 確保cache_k和xq位于同一設備(GPU)
        self.cache_v = self.cache_v.to(xq) # 確保cache_v和xq位于同一設備(GPU)

        # 將xk和xv加載進緩存
        self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
        self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv

        # keys和values的取值從數組的最左邊開始
        keys = self.cache_k[:bsz, : start_pos + seqlen]
        values = self.cache_v[:bsz, : start_pos + seqlen]

        # 將keys和values根據n_rep復制出相應的份數，以便矩陣相乘
        # repeat k/v heads if n_kv_heads < n_heads
        keys = repeat_kv(
            keys, self.n_rep
        )  # (bs, cache_len + seqlen, n_local_heads, head_dim)
        values = repeat_kv(
            values, self.n_rep
        )  # (bs, cache_len + seqlen, n_local_heads, head_dim)

        xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
        keys = keys.transpose(1, 2)  # (bs, n_local_heads, cache_len + seqlen, head_dim)
        values = values.transpose(
            1, 2
        )  # (bs, n_local_heads, cache_len + seqlen, head_dim)
        # 計算注意力得分
        scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
        if mask is not None:
            scores = scores + mask  # (bs, n_local_heads, seqlen, cache_len + seqlen)
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
        output = torch.matmul(scores, values)  # (bs, n_local_heads, seqlen, head_dim)
        # 改變output的形狀為(bs, seqlen, n_local_heads * head_dim)，n_local_heads * head_dim等同于將每頭的head_dim進行concat
        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
        # 將output和權重wo進行線性計算，得到注意力的最終輸出
        return self.wo(output)

0x04 優化

4.1 優化策略

對于長序列的數據和高維度的模型，自注意力機制的計算量是非常大的，而且這種巨大的計算量也限制了模型擴展到訓練數據以外的更長的序列，即存在長度外推問題。因此需要對自注意力機制進行優化。有各種優化的分類，本文選取幾個典型的方面進行分析。

從序列角度優化

原始自注意力機制中，每個token都要和所有token作注意力計算。但是Transformer中的注意力分布其實是不均勻的。比如“How Do Language Models put Attention Weights over Long Context”就發現不同層的注意力分布有顯著差異。起始層的注意力分布大致均勻。然而中間層的注意力模式變得更加復雜，大部分概率質量集中在初始標記（注意力匯聚）和最近的/最后標記（近期偏見）上。即Transformer的中間層大部分都是“V形”注意力分布。這意味著中間層很多的token其實作用不大，也許可以通過壓縮token的方式減少context長度，從而達到加速推理或者支持更長context需求。

因此，人們提出了稀疏計算（或者說是限制注意力）的思想，即每個 token 只與一部分 token 做注意力計算（而非全部）。其中技巧在于當前token與哪一部分 token 做注意力計算。幾個典型做法如下。

Atrous self-attention(空洞自注意力）：每個 token 等間隔的和其他 token 做注意力計算。
Local self-attention（局部注意力）：類似于 n-gram，每個 token 與附近的幾個 token 做注意力計算。
Sparse attention（稀疏注意力）：結合了空洞注意力和局部注意力機制，根據規則（或者動態決定），每個 token 即有局部感受野，也有全局感受野。

幾個經典模型如下：

Longformer：是稀疏注意力的一種，參與注意力的包括局部注意力、空洞注意力和全局注意力。全局注意力就是對于一些特定位置的 token（例如 Bert 中的 CLS），做全部的注意力計算。
Reformer：本方案基于這樣的假設：經過 softmax 之后，一個 query 和其他所有的 token 計算的 attention score 主要是取決于高相似度的幾個 tokens，所以可以采用最相似的token來近似計算得到最終的 attention score。因此，問題就轉化成為了找最相似的 Top-N 問題。reformer 是通過 Locality sensitive hashing（LSH）方法將 attention score 相近的token分到一個 bucket 中。從而在低維空間中保留高維空間的相似性。這樣可以大幅減少計算復雜度，特別是在處理大規模數據時。
Linformer：證明了 Transformer 的 Attention 矩陣是低秩矩陣，這說明了矩陣的大部分的內容可以通過少部分最大的奇異值得到。層級越高，會有越多的信息集中在最大的奇異值上。Linformer 對 Self-Attention 的結構進行了一些修改（主要的區別在于增加了兩個線性映射層），使復雜度降到線性。
BigBird：使用了一種塊稀疏的注意力模式，結合了局部（窗口）、全局（固定點）和隨機的注意力。塊式稀疏注意力：將注意力機制應用在連續的局部塊上，而不是整個序列。全局注意力點：為了保持交互的全局性，添加了總是彼此進行注意力操作的標記，如 CLS 標記。隨機注意力：在序列中隨機引入額外的注意力連接，以確保任意兩個位置之間總有一個路徑（可能通過多個中間節點）使得彼此可以連接。
DCA（Dual Chunk Attention）：雙塊注意力（DCA）是一個無需訓練的框架，用于推斷LLMs的上下文窗口。DCA沒有使用線性縮放位置索引或增加RoPE的基頻。相反，它選擇重用預訓練模型中的原始位置索引及其嵌入，但重新設計了相對位置矩陣的構建，以盡可能準確地反映兩個標記之間的相對位置。DCA的創新點也在于它與Flash Attention的兼容性，僅需要對推理代碼進行修改，無需進行大量的重新訓練。

另外，長度外推性是一個訓練和預測的不一致問題，而解決這個不一致的主要思路是將注意力局部化，很多外推性好的改進某種意義上都是局部注意力的變體。局部化注意力就是通過限制注意力的感知范圍，來賦予整個模型“平移不變性”。

窗口截斷的方式也可以作為長度外推的一個不錯的Baseline，但這是通過強行截斷窗口外的注意力、犧牲遠程依賴換來的，因此還不是最終的解決方案。

DCA

因為Qwen2.5-Turbo采用了DCA來通過分塊處理長序列，減少了計算復雜度，提升了模型的推理速度，所以我們再重點介紹下DCA。

DCA由三個組件組成，這些各自的處理幫助模型有效捕捉序列中的長距離和短距離依賴。

塊內注意力，針對同一塊內的標記處理。
塊間注意力，用于處理不同塊之間的標記。
連續塊注意力，用于處理連續的、不同的塊中的標記。

下圖是塊內注意力。

下圖是塊間注意力。

下圖是連續塊注意力，這是塊間注意力的特殊case。連續塊注意力用來保持LLM的局部性，局部性意味著LLM傾向于主要依賴相鄰token來預測下一個token。

綜合三種注意力之后，得到如下。

從多頭角度優化

具體包括MHA、MQA、GQA，我們后續有專門篇幅介紹。

從軟硬件層面優化 MHA

硬件層面上，比如現在已在使用的 HBM（高速帶寬內存）提高讀取速度，或者更徹底些，拋棄馮諾依曼架構，改變計算單元從內存讀數據的方式，不再以計算單元為中心，而以存儲為中心，做成計算和存儲一體的“存內計算”。

軟件層面上的話，最近的很多優化，比如 Flash Attention，Paged Attention。我們后續有專門篇幅介紹。

從其它角度優化

比如下面要介紹的$Transformer^2$，其并不是對Transformer結構做了改變，而是在推理時分兩次推理。

4.2 案例

我們接下來介紹幾個特色案例。

注意力權重細化

傳統的自注意力模型可能不足以捕獲序列推薦場景中復雜的item依賴關系，這是因為缺乏對注意力權重的明確強調，而注意力權重在分配注意力和理解item間相關性方面起著關鍵作用。為更好發揮注意力權重的潛力并提高序列推薦在學習高階依賴關系方面的能力，論文"Pay Attention to Attention for Sequential Recommendation“提出了注意力權重細化（AWRSR，Attention Weight Reffnement for Sequential Recommendation）方法，通過額外關注注意力權重來增強自注意力機制的有效性，從而實現item間相關性更精細的注意力分布。

現有的注意力權重操作圍繞計算item embedding或分布之間的相似性或距離，如下圖右下角部分所示。這種方式并沒有考慮注意力權重內的潛在相關性，而這些相關性可能會進一步揭示注意力權重本身內的高階轉換。如下圖中右上角部分所示，注意力權重矩陣 A 源自序列中item的表示，它編碼了序列中每個元素相對于所有其他元素的關系。例如，$A_1$表示從第一個item到序列中所有其他item的注意力或重要性的分配，這意味著細化這些權重（本質上捕捉權重/注意力之間的相關性）具有建模高階轉換的潛力。鑒于此，比較$A_1$和$A_2$可以有助于對序列中兩個元素（位于第一和第二位置）之間依賴關系的理解。

論文共設計了4種細化機制來計算注意力權重之間的依賴關系，即：計算注意力權重之間的注意力權重（pay attention to attention）。

Simple refinement。該機制只是將一組新的可訓練矩陣應用于注意力權重矩陣 A，將其轉換為新的查詢和鍵，以計算更高級別的權重，形式如下圖標號1。其中W是可訓練的矩陣。接著對新的權重矩陣B使用softmax函數，softmax(B) 與item集合的乘積相當于對值進行求和，這意味著每個值都通過細化的高階注意力權重進行縮放。
Value-weighted refinement。通過attention進一步轉化Simple refinement，得到新權重B，希望在更復雜的高階空間中提煉并重新表達權重相關性，給出矩陣形式以簡化計算，形式如下圖標號2。
Additive refinement。該機制的目的是合并/平衡不同級別的注意力權重，形式如下圖標號3。
Stochastic refinement。該機制針對 STOSA（STOchastic Self-Attention）進行了量身定制，嘗試將 STOSA 原有的隨機權重 A 轉變為一種可能保持權重概率性質的新形式。對item k計算一個新的注意力分布。這里的$W_\mu$和$W_\sum$為兩個隨機矩陣，形式如下圖標號4。

線性注意力

Transformer 核心的自注意力機制是其計算成本的重要來源。為了優化，研究社區提出了稀疏注意力、低秩分解和基于核的線性注意力（KERNEL-BASED LINEAR ATTENTION）等許多技術。

vanilla Transformer使用Softmax注意力，需要為此構建一個N×N 的全連接矩陣，對于超長序列，這個矩陣會非常龐大。它會讓模型在處理長文本時復雜度成n的平方的增加，復雜度是$O (N^2)$。

為了緩解標準自注意力機制的效率瓶頸，在論文“Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention"中提出了基于核的線性注意力機制，該機制將相似度函數分解為特征映射的點積。按照 Linear Attention 工作里的符號，我們定義$SM(q,k)=exp(q^Tk)$為softmax 核函數。從數學上講，線性注意力的目標是使用 $?(q_i)?(k_j)^T $來近似 SM (?,?)。

通過“右積核技巧（right product kernel trick）”可以將傳統的二次計算復雜度轉化為線性復雜度，顯著降低了長序列處理的計算負擔。，每個頭的復雜度可以降低到 $O (Nd’^2)$，其中 d’是特征映射后的維度，與序列長度成線性關系。具體來說，線性注意力通過遞歸更新鍵值矩陣的乘積，避免了重復計算整個注意力矩陣，從而在推理過程中保持了恒定的計算復雜度。

下圖給出了 softmax注意力（左）和線性注意力（右）的計算說明。輸入長度為N 特征維數為d，d≥N。同一框中的張量與計算相關聯。線性化公式可以獲得O(N)的時間和空間復雜性。

論文“The Devil in Linear Transformer”中的 TransNormer 就是線性注意力的體現，具體參見下圖。為了保持整體架構線性復雜度，TransNormer 把靠底層的layer換成了local attention，然后后面的layer用去掉了分母的，normalize過后的線性attention。

PolaFormer

論文“PolaFormer: Polarity-aware Linear Attention for Vision Transformers”中，作者提出了一種極性感知線性注意力（PolaFormer）機制，旨在通過納入被忽略的負交互作用來解決先前線性注意力模型的局限性。與此同時，為了解決線性注意力中常見的注意力權重分布信息熵過高的問題，他們提供了數學理論基礎，表明如果一個逐元素計算的函數具有正的一階和二階導數，則可以重新縮放 q,k 響應以降低熵。這些增強功能共同提供了一個更穩健的解決方案，以縮小線性化和基于 Softmax 的注意力之間的差距。

研究背景

線性注意力，作為一種更具可行性的解決方案使用核化特征映射替換 q,k 點積中的 Softmax 操作，有效地將時間和空間復雜度從 O (N2d) 降低到 O (Nd2)。盡管線性注意力在計算效率上有所提升，但在表達能力方面仍不及基于 Softmax 的注意力，論文分析確定了造成這種不足的兩個主要原因，它們都源于 Softmax 近似過程中的信息丟失：

負值丟失。依賴非負特征映射（如 ReLU）的線性注意力模型無法保持與原始 q,k 點積的一致性。這些特征映射僅保留了正 - 正交互作用，而關鍵的正 - 負和負 - 負交互作用則完全丟失。這種選擇性表示限制了模型捕獲全面關系范圍的能力，導致注意力圖（attention maps）的表達能力減弱和判別力降低。
注意力分布的高信息熵。沒有 softmax 的指數縮放，線性注意力會導致權重分布更加均勻且熵更低。這種均勻性削弱了模型區分強弱 q,k 對的能力，損害了其對重要特征的關注，并在需要精細細節的任務中降低了性能。

思路

極性感知注意力背后的核心思想是為了解決現有線性注意力機制的局限性，這些機制通常會丟棄來自負成分的有價值信息。PolaFormer 在處理負成分時，極性感知注意力將 query 和 key 向量分解為它們的正部和負部。這種分解允許機制分別考慮正相似度和負相似度對注意力權重的影響。然后將這些分解代入 q 和 k 的內積中，得到如下圖所示。

之前的線性注意力方法，如基于 ReLU 的特征映射，通過將負成分映射到零來消除它們，這在近似 q,k 點積時會導致顯著的信息丟失。為了解決這個問題，極性感知注意力機制根據 q,k 的極性將它們分開，并獨立計算它們之間的相互作用。注意力權重的計算方式如下圖所示。

PolaFormer 根據極性明確地將 q,k 對分開，處理在內積計算過程中維度的同號和異號交互作用。這些交互作用在兩個流中處理，從而能夠更準確地重建原始的 softmax 注意力權重。為了避免不必要的復雜性，作者沿著通道維度拆分 v 向量，在不引入額外可學習參數的情況下處理這兩種類型的交互作用。然后，將輸出進行拼接，并通過$G^S$和$G^O$進行縮放，以確保準確重建 q,k 關系。$G^S$和$G^O$是兩個可學習的極性感知系數矩陣，應用了逐元素（element-wise）乘法，會學習相同符號值和相反符號值之間的互補關系。

具體也可以參見下圖。

為了解決線性注意力中常見的注意力權重分布信息熵過高的問題，作者提供了數學理論基礎，表明如果一個逐元素計算的函數具有正的一階和二階導數，則可以重新縮放 q,k 響應以降低熵。這一理論有助于闡明為什么先前的特征映射會提高信息熵，從而導致注意力分布過于平滑。為了簡化，作者采用通道級可學習的冪函數進行重新縮放，這保留了 Softmax 中固有的指數函數的尖銳性。這使得模型能夠捕獲尖銳的注意力峰值，提高了其區分強弱響應的能力。與此同時，為了區分開不同通道之間的主次關系，作者設計了可學習的冪次來捕捉每個維度的不同重要性。

MiniMax-01

MiniMax-01是第一個依賴線性注意力機制的大規模部署的模型。線性注意力則可以把復雜度控制在線性增加范圍，但它基本是一種“實驗”的狀態。MiniMax-01在注意力機制層面做了大膽的創新，在業內首次實現了新的線性注意力機制，第一次把它放到了生產環境里。它的80層注意力層里，每一層softmax attention層前放置了7層線性注意力（lightning attention）層。

模型架構

MiniMax-01是基于線性注意力機制，采用混合架構 (Hybrid-Lightning)，并集成了 MoE 架構。

Lightning Attention

MiniMax 的 Lightning Attention 便是一種線性注意力，是基于 TransNormer 實現的一個 I/O 感知型優化版本。Lightning Attention 的線性注意力機制解決了因果模型在計算單向注意力時，需要進行累加求和操作導致無法矩陣運算的情況，實現了單向注意力先計算右乘，成功將復雜度降為 $O(nd^2)$ 。具體而言，Lightning Attention是通過進行tiling（分塊計算），避免了因果語言建模中的累積求和操作，從而實現了理論上的線性復雜度。具體來說，對于超長序列，Lightning Attention將Q、K和V矩陣沿行維度劃分為多個塊，每個塊的大小固定。然后將注意力計算分為塊內計算（使用左積）和塊間計算（使用右積）兩部分，先獨立計算塊內部的詞之間的注意力分數（intra-block），接著再通過一種遞歸更新的方法，將塊與塊之間的信息逐步傳遞（inter-block），最終可以捕捉到全局語義關系。

這個過程類似于分組討論：先解決每組內部的問題，再匯總所有組的結果，最終得到全局的答案。這種優化大大減少了計算和內存需求，也從傳統 Softmax 注意力的平方復雜度降低為線性?；蛘咭部梢园涯Ｐ拖胂蟪稍诜喴槐揪薮蟮臅?，即使每次只能看幾頁，但它能記住之前的內容，最終把整本書的知識都處理一遍。

以下是 Lightning Attention 前向傳播的算法描述。

Hybrid-lightning

為了平衡效率與全局信息捕捉能力，MiniMax 還基于 Lightning Attention提出了一種 Hybrid-lightning，即在 Transformer 的每 8 層中，有 7 層使用 Lightning Attention，高效處理局部關系；而剩下 1 層保留傳統的 Softmax 注意力，從而既解決了 softmax 注意力的效率問題，確保能夠捕捉關鍵的全局上下文，也提升了 Lightning Attention 的 scaling 能力。和傳統的機制相比，一個是看書時候每個字都看，另一個是挑重點看，然后偶爾看一下目錄對照一下整體。效率自然不同。

此外，它還引入了Varlen Ring Attention，用來直接將整個文本拼接成一個連續的序列，從而讓變長序列的數據在模型中按需分配資源；在預訓練數據上使用數據打包（Data Packing），將不同長度的文本拼接成連續的長序列；在分布式計算時改進了 Linear Attention Sequence Parallelism (LASP+)，使模型能夠在多 GPU 之間高效協作，無需對文本進行窗口切分。

下表給出了根據層數 l、模型維度 d、批量大小 b 和序列長度 n 計算注意力架構參數量與 FLOPs 的公式。可以明顯看出，模型規模越大，Lightning Attention 與 Hybrid-lightning 相對于 softmax 注意力的優勢就越明顯。

Transformer2

論文"Transformer2 : SELF-ADAPTIVE LLMS" 中提出了一種可以根據不同任務動態調整模型權重的機器學習系統 ——Transformer2，其通過奇異值微調和權重自適應策略，實時選擇性地調整權重矩陣中的單一組件，提高了LLM的泛化和自適應能力，使LLM能夠適應未見過的任務。幾個有趣的思路如下:

雖然論文題目叫$Transformer^2$，但其實并不是對Transformer結構做了改變，而是在推理時分兩次推理。
使用SVD分解訓練，而非LoRA。
使用Reinforce方法，而非SFT。

研究背景

自適應性

在自然界，「適應」是一種非常普遍的現象。例如，章魚能夠迅速改變自身的膚色和紋理，以融入周圍環境，從而躲避天敵和捕捉獵物；人腦在受傷后能夠重新連接自身神經回路，使個體能夠恢復失去的功能并適應新的思維方式或行動方式。生物體展現出的適應能力使得生命能夠在不斷變化的環境中蓬勃發展。這些無不體現著那句經典的名言——「物競天擇，適者生存」。

在人工智能領域，適應的概念同樣具有巨大的吸引力。想象一個機器學習系統，它能夠動態地調整自身的權重以在陌生的環境中不斷學習、進化。與部署在環境中的靜態 AI 模型相比，這種有自適應能力的模型明顯學習效率更高，而且有望成為與現實世界動態本質始終保持一致的終生模型。

SVD

就像人類大腦通過互連的神經通路存儲知識和處理信息一樣，LLM 在其權重矩陣中存儲知識。這些矩陣是 LLM 的「大腦」，保存著它從訓練數據中學到的精髓。要理解這個「大腦」并確保它能夠有效地適應新任務，需要仔細研究其內部結構。而奇異值分解（SVD）提供了寶貴的洞察力。

SVD通過識別LLM權重矩陣中的主成分來實現這一目標。SVD 將存儲在 LLM 中龐大、復雜的知識分解成更小的、有意義的、獨立的部分（例如數學、語言理解等不同的組件）。在研究中發現，增強某些成分的信號，同時抑制其他部分的信號，可以提高LLM在下游任務中的表現。

可以將SVD看作是一名外科醫生，正在對LLM的大腦進行細致操作。這名外科醫生將LLM中存儲的龐大復雜的知識分解成更小、更有意義且獨立的部分（例如，針對數學、語言理解等的不同路徑或組件）。

或者說，預訓練已經使模型充分具備了解答問題的能力，而微調實際上是將該能力涌現出來，因此可以通過構建子矩陣并僅調整奇異值矩陣部分主成分的方法來達到輕量微調的效果。且由于僅微調奇異值尺度本身，因此可以使用少量數據訓練而不會導致大規模遺忘或崩潰。

研究動機

雖然組合性和可擴展性對于有效適應至關重要，但當前的 LLM 訓練方法難以同時實現這兩個特性。Sakana AI 的研究旨在提出一個開創性的解決方案來實現這一愿景并解決這些 gap。

傳統上，LLM 后訓練試圖通過一次全面的訓練來優化模型，使其具備廣泛的能力。對于LLM來說，想要加入哪怕只是一句話的新知識，都必須要再訓練一次。從簡化的角度，這種「one shot（一次性）」微調框架從簡單性的角度來看是理想的，但在實踐中很難實現。例如，后訓練仍然非常消耗資源，導致巨大的計算成本和超長的訓練時間。此外，在引入更多樣化的數據時，往往存在明顯的性能權衡，這使得同時克服過擬合和任務干擾變得具有挑戰性。另外，目前微調過程會產生災難性遺忘和泛化性低的問題。使用某個數據集對某個模型進行微調訓練，再進行評測，會發現大概率上，該模型其他任務的性能會全面下降，甚至本任務也有一定程度的下降（過擬合+遺忘）。這是因為微調的過程雖然簡單，但是卻破壞了模型原有的能力。

相比之下，自適應模型提供了更靈活和高效的方法。與其試圖一次性訓練 LLM 完成所有任務，不如開發專家模塊，然后根據需求將其離線開發并增強到基礎LLM中。專家模塊可以離線開發并按需增強到基礎 LLM 中。這使模型能夠根據當前任務動態修改其行為，而無需不斷重新調整。除了具有獨立組件的好處外，這種模塊化還支持持續學習，使模型能夠隨時間增加新技能而不會出現災難性遺忘。此外，自適應 LLM 反映了神經科學和計算生物學中一個公認的原理，即大腦根據當前任務激活特定區域，并動態重組其功能網絡以響應不斷變化的任務需求。

原則上，實現自適應 LLM 的第一步可以通過開發專門的專家模塊來實現，每個模塊都通過 LoRA 等技術進行微調。然后這些專家模塊可以根據任務需求在運行時動態組合，這個過程可以通過 MoE 類系統高效管理。然而，要使這種方法既可擴展又具有組合性，需要解決幾個挑戰。首先，微調 LLM 以創建多個專家模塊顯著增加了需要訓練的參數數量。實際上，即使使用 LoRA 等參數高效的方法，這些模塊的累積大小也會快速增加，導致存儲和計算需求增加。其次，這些專家模塊往往容易過擬合，這種現象在較小數據集或窄任務領域訓練時尤為普遍。第三，這些專家模塊的靈活組合也帶來了目前尚未解決的挑戰。

另外，多個LoRA也不具備可組合型。實際上LoRA在擴散模型算法的領域內是可以組合的，例如風格LoRA和人物LoRA可以同時在一個模型上應用起來分別生效，但是這個不是LoRA的特性，而是擴散算法本身的穩定性，而在大模型領域，多個LoRA聯合使用是不具備數學意義的。因此奇異值微調的數學意義是具備了充分的可解釋性（向量方向和尺度）和可組合型（例如插值），因此可以組合多個tuner共同使用。

思路

奇異值微調（SVF）

為了克服針對組合性和可擴展性的這些限制，Transformer2 作者首先提出了奇異值微調（SVF），這是一種使用強化學習（RL）來增強或抑制不同「大腦」組件的信號，也是一種新的參數高效微調（PEFT）方法，用于獲得自適應的有效構建塊。

SVF 使用SVD將LLM的「大腦」（即權重矩陣）分解為若干獨立的組件，組件可能在多個任務中共享。比如，某些組件在語言理解和推理任務之間是共享的。在訓練時，SVF利用RL訓練這些組件的組合以應對不同任務。在推理時，SVF首先識別任務類型，然后動態調整組件的組合。

具體而言，在訓練階段，SVF學習一組z向量，其中每個下游任務對應一個z向量。每個z向量可以視作該任務的專家，它是一個緊湊的表示，指定了權重矩陣中每個組件的期望強度，充當「放大器」或「衰減器」，調節不同組件對模型行為的影響。例如，假設SVD將權重矩陣分解為五個組件[A，B，C，D，E]。對于數學任務，學習到的z向量可能是[1，0.8，0，0.3，0.5]，這表明組件A對數學任務至關重要，而組件C幾乎不影響其表現。對于語言理解任務，z向量可能是[0.1，0.3，1，0.7，0.5]，表明盡管C組件對數學任務的貢獻較小，但它對語言理解任務至關重要。SVF利用RL在預定義的下游任務集上學習這些z向量。學習到的z向量使Transformer2能夠適應各種新的下游任務，同時僅引入最少量的附加參數（即z向量）。為何說這樣參數量非常低？是因為U和V矩陣不參與訓練，因此僅需要一個一維向量z，即可以表達某個具體dense矩陣的改變。

下圖給出了SVF的概況。在訓練時，Transformer2 使用SVF和RL來學習“專家”向量z，這些專家向量可以縮放權重矩陣的奇異值。在推理時，Transformer2 提出了三種不同的方法來自適應地選擇/組合學習到的專家向量。

另外，Transformer2使用Reinforce方式而非SFT。作者給出的解釋是RL對數據集的要求更低。由于RL是針對獎勵值進行優化而非next-token的邏輯關聯，因此可以忽視CoT帶來的影響，使用較為一般的數據集進行訓練。

自適應性

在推理階段，Transformer2采用兩階段適應策略，以有效地組合任務特定的 z 向量集。Transformer2這個名稱也反映了它的兩步過程。其核心在于能夠動態調整權重矩陣中的關鍵組件。

首先，模型分析傳入的任務以了解其要求，然后應用特定于任務的調整來生成最佳結果。即，第一次推理分析任務類型，第二次推理則會根據任務類型來選擇具體的tuner（或者類似MoE，進行專家間的組合推理）進行推理來解決具體問題。通過有選擇地調整模型權重的關鍵組成部分，該框架允許 LLM 實時動態地適應新任務。

在第一次階段，Transformer2 執行模型并觀察其測試時行為，收集相關信息以理解解決當前問題所需的技能。對于給定任務或單個輸入提示，Transformer2通過以下三種適配方法之一來識別任務的特征。
- 基于prompt的適配：專門設計的適配性提示，讓LLM對輸入prompt進行分類，并選擇合適的 z 向量。
- 基于分類器的適配：使用依靠 SVF 訓練的任務分類器，在推理過程中識別任務，并選擇合適的 z 向量。
- 少樣本適配：通過加權插值來組合多個預訓練的 z 向量。
這三種方法共同確保了Transformer2能夠實現強大且高效的任務適應，為其在多種場景下的出色表現奠定了基礎。
在第二階段，Transformer2框架通過組合 z 向量來相應地調制權重，從而生成最適合任務的最終響應。

下圖給出了Transformer2的總體架構。

Titans

論文“Titans: Learning to Memorize at Test Time）”提出了一種新型的神經長期記憶模塊，該模塊能夠在利用長遠歷史信息的同時，讓注意力機制專注于當前上下文。該神經記憶的優勢在于可以快速并行訓練，并保持快速推理。論文指出，由于上下文有限但依賴關系建模精確，注意力機制可以作為短期記憶；而神經記憶由于其記憶數據的能力，可以作為長期、更持久的記憶?；谶@兩個模塊，論文引入了一種新的架構——泰坦（Titans），并提出了三種變體來有效地將記憶融入架構中。

研究背景和動機

近年來，循環模型和注意力機制在深度學習領域得到了廣泛應用。循環模型旨在將數據壓縮成固定大小的記憶（隱藏狀態），而注意力機制則允許模型關注整個上下文窗口，捕捉所有標記的直接依賴關系。然而，這種更精確的依賴關系建模帶來了二次方的計算成本，限制了模型的上下文長度。

核心創新

這篇論文的核心創新在于提出了一個能夠在測試時學習記憶的神經長期記憶模塊（neural memory module）。它能夠學習記憶歷史上下文，并幫助注意力機制在利用過去已久信息的同時處理當前上下文。結果表明，這種神經記憶具有快速并行化訓練的優勢，同時還能保持快速推理。

這個模塊的工作方式如下：

記憶的獲取：該模塊將訓練過程視為在線學習問題，旨在將過去的信息壓縮到其參數中。受人類記憶啟發，該模塊將“違反預期”的事件（即令人驚訝的輸入）視為更值得記憶的。它通過計算神經網絡相對于輸入的梯度來衡量輸入的“驚訝程度”，并使用這一指標來更新記憶。
遺忘機制：為了解決有限記憶的問題，該模塊引入了一種自適應的遺忘機制，該機制考慮了記憶大小和數據驚訝程度，從而更好地管理記憶。
記憶的結構：論文探索了不同的記憶結構，發現深度記憶模塊（即使用多層感知機）比線性模型更有效。
記憶的檢索：該模塊通過簡單的正向傳遞（不更新權重）來檢索與查詢相對應的記憶。

Titans架構

基于長期神經記憶模塊，論文提出了泰坦架構，該架構包含三個分支：

核心分支（Core）：使用注意力機制進行數據處理，關注有限的上下文窗口。
長期記憶分支（Long-term Memory）：使用神經長期記憶模塊來存儲和回憶歷史信息。
持久記憶分支（Persistent Memory）：使用可學習但不依賴于數據的參數來編碼任務相關知識。

論文提出了三種不同的泰坦變體來有效地將記憶融合到該系統架構中。

上下文記憶（Memory as a Context，MAC）：將長期記憶視為當前信息的上下文，使用注意力機制融合這些信息。
門控記憶（Memory as a Gate，MAG）：使用門控機制將長期記憶與核心分支的信息融合。
層式記憶（Memory as a Layer，MAL）：將長期記憶模塊作為深度神經網絡的一層。

長期記憶

為了設計一個長期神經記憶模塊，我們需要模型能夠將過去歷史的抽象編碼到其參數中。因此，一個簡單的思路是訓練神經網絡并期望它能夠記住自己的訓練數據，然而記憶幾乎一直是神經網絡中令人頭疼的現象，它限制了模型的泛化能力，還引發隱私問題，因此導致測試時性能不佳。

基于此，谷歌認為需要一個在線元模型來學習如何在測試時記憶或忘記數據。在這種設置下，模型學習一個能夠記憶的函數，但不會過擬合訓練數據，從而在測試時實現更好的泛化性能。

學習過程和意外指標（Learning Process and Surprise Metric）。訓練長期記憶的關鍵思路是將訓練視為在線學習問題，將過去信息壓縮到長期神經記憶模塊中。人類往往能夠記住背離預期（令人驚訝）的事件，受此啟發，模型意外可以簡單定義為它相對于輸入的梯度。梯度越大，輸入數據與過去數據的偏差就越大。因此，使用這個意外分數，可以將記憶更新如下圖標號1。

意外指標可以導致在重大意外時刻之后出現重要信息缺失。從人類記憶的角度來看，即使一個事件令人難忘，但它可能不會在長時間內持續讓我們感到驚訝。為了改進這一現象，谷歌將意外指標分解稱為(1)過去意外，它衡量最近過去的意外程度；(2)瞬時意外，它衡量傳入數據的意外。具體如下圖標號2。
上述的意外指標是基于一個損失函數 $ l(.;.)$之上的，它就是記憶模塊在測試時學習的目標。也就是說，記憶模塊是一個元模型，它基于損失函數$l(.;.)$來學習一個函數。谷歌則專注于關聯（associative）記憶，目的是將過去的數據存儲為鍵（keys）和值（values）對。類似于 Transformer，在給定 x_t 的情況下，谷歌使用兩個線性層將 x_t 投影到鍵和值中。如下圖標號3。谷歌希望記憶模塊可以學習鍵和值之間的關聯，為此將損失定義如下圖標號4。
遺忘機制（Forgetting Mechanism）。在處理非常大的序列（比如百萬 tokens）時，管理哪些過去信息應該被遺忘非常重要，即使使用深度或者非常大的矩陣值記憶時也是如此。因此，谷歌使用了一種自適應遺忘機制，允許記憶忘記不再需要的信息，從而更好地管理有限的記憶容量。也就是說，給定下一個 token x_t，遺忘機制如下圖標號5：
檢索記憶（Retrieving a Memory）。在探討如何設計和訓練一個可以在測試時學習記憶的長期記憶模塊之后，剩下的關鍵問題便是如何從記憶中檢索信息？谷歌僅僅使用了沒有更新權重的前向傳遞（即推理）來檢索與查詢相對應的記憶。在形式上，給定一個輸入 x_t，谷歌使用線性層 W_Q 來投影輸入，即 q_t = x_tW_Q，并通過以下公式從記憶 y_t 中檢索相應（或有用）的信息，對應下圖標號6。

融合記憶

接下來需要解決的一個重要問題是：如何有效且高效地將神經記憶融合到深度學習架構中？

從記憶的角度來看，Transformer 中的 K 和 V 矩陣對可以解釋為聯想記憶塊。由于它們對依賴關系的精確建模以及有限的上下文窗口，它們可以被用作短期記憶模塊，以處理當前上下文窗口大小。另一方面，神經記憶能夠不斷從數據中學習并存儲在其權重中，因而可以發揮長期記憶的作用。谷歌通過三個不同的 Titans 變體來回答以上問題。

如果把token看作是Instruction， Test Time ScalingLaw可能就是對token序列進行修改，讓模型按照新的流程來執行。Titans可能就是這樣一個自我驗證糾錯機制的雛形。注意力機制可以類比為片上指令緩存，神經記憶模型可以類比為代碼段。

記憶作為上下文（Memory as a Context，MAC）

Titans 的第一個變體 MAC 的架構設計如上圖所示，將記憶作為當前信息的上下文。

該架構具有兩個關鍵優勢：一是注意力模塊同時具有歷史和當前上下文，能夠根據當前數據決定是否需要長期記憶信息，二是注意力模塊幫助長期記憶只存儲來自當前上下文的有用信息。這意味著，并非每個片段中的所有 token 都是有用的，記憶所有 token 可能會導致內存溢出。因此，注意力模塊幫助記憶了解哪些信息是有用的，從而更好地管理內存容量。

記憶作為門（Memory as a Gate，MAG）