制從原理到代碼：QKV、多頭、掩碼與KV Cache實(shí)戰(zhàn)指南)
簡(jiǎn)介這是一份面向深度學(xué)習(xí)初學(xué)者和自然語(yǔ)言處理研究者的注意力機(jī)制專題課件系統(tǒng)講解序列到序列模型中的編碼器—解碼器結(jié)構(gòu)為何需要注意力以及固定長(zhǎng)度語(yǔ)義向量在長(zhǎng)輸入下的信息丟失問(wèn)題。課件從循環(huán)神經(jīng)網(wǎng)絡(luò)只能處理等長(zhǎng)序列的局限出發(fā)逐步引出注意力思想的來(lái)源并重點(diǎn)剖析Bahdanau提出的加性注意力、Luong提出的注意力變體以及硬注意力、局部注意力與自注意力等主流分支分別說(shuō)明其計(jì)算邏輯、優(yōu)缺點(diǎn)和適用任務(wù)。結(jié)合機(jī)器翻譯、文本摘要、語(yǔ)音識(shí)別、圖像描述等典型場(chǎng)景通過(guò)圖解和權(quán)重配色示意直觀呈現(xiàn)注意力分?jǐn)?shù)如何在不同輸入位置上分配幫助讀者理解每類注意力的工作機(jī)制。資源為1個(gè)PPT文件大小約2.7MB結(jié)構(gòu)緊湊、圖文并茂既適合教師備課參考也適合學(xué)生課前預(yù)習(xí)、課堂展示或考研復(fù)試快速回顧。目前已有1183人學(xué)習(xí)下載能夠顯著降低理解門(mén)檻也為后續(xù)學(xué)習(xí)Transformer架構(gòu)打下扎實(shí)基礎(chǔ)。1. 注意力機(jī)制不是玄學(xué)你在Transformer、Seq2Seq和圖像里早就見(jiàn)過(guò)它搜索“深度學(xué)習(xí)中attention注意力機(jī)制.ppt”的人多半是要準(zhǔn)備技術(shù)分享或者被面試題逼到了。注意力機(jī)制看起來(lái)像一套優(yōu)雅的數(shù)學(xué)公式但真正困擾開(kāi)發(fā)者的不是讀不懂論文而是“我怎么把它用起來(lái)”。這篇筆記會(huì)把你從“知道Attention”拉到“能寫(xiě)最小實(shí)現(xiàn)、能調(diào)參數(shù)、能排坑”的狀態(tài)。我不打算復(fù)制PPT里的定義而是把QKV、多頭、mask、緩存這些實(shí)操點(diǎn)拆開(kāi)每個(gè)環(huán)節(jié)都給出能直接跑的代碼和踩過(guò)的坑。你在Seq2Seq、Transformer、圖像分類里見(jiàn)到的各種注意力底層都是同一套加權(quán)求和邏輯。新手可以照步驟做熟手可以對(duì)照參數(shù)和邊界。2. 從“對(duì)齊”到“加權(quán)求和”Attention的核心數(shù)學(xué)直覺(jué)與QKV拆解2.1 早期Seq2Seq的Attention為什么要在譯每個(gè)詞時(shí)“回頭看”注意力機(jī)制最早出圈是在機(jī)器翻譯里。當(dāng)時(shí)常見(jiàn)的做法是用一個(gè)編碼器把整個(gè)輸入句子壓成一個(gè)固定長(zhǎng)度的向量解碼器每一步都從這個(gè)向量里取信息。句子一長(zhǎng)最后那個(gè)向量扛不住信息量前面的內(nèi)容基本被后面蓋掉了。后來(lái)某篇工作提出解碼時(shí)每一步不要把目光鎖死在那個(gè)固定向量上而是回到編碼器的每個(gè)隱藏狀態(tài)上“按需取用”。具體做法是計(jì)算當(dāng)前解碼狀態(tài)和每個(gè)編碼器隱藏狀態(tài)的相關(guān)程度得到一個(gè)權(quán)重分布然后用這個(gè)權(quán)重對(duì)編碼器隱藏狀態(tài)做加權(quán)求和得到這一步的上下文向量。這個(gè)“回頭看”的操作本質(zhì)就是給序列里的每個(gè)位置分配一個(gè)重要性分?jǐn)?shù)。分?jǐn)?shù)高的位置對(duì)當(dāng)前輸出的影響大分?jǐn)?shù)低的位置幾乎被忽略。這就是Attention最初的含義注意力權(quán)重告訴模型“該看哪個(gè)位置”加權(quán)和把“看到了什么”送到解碼器。它不需要額外的監(jiān)督信號(hào)權(quán)重完全由數(shù)據(jù)驅(qū)動(dòng)學(xué)出來(lái)所以經(jīng)常被形容為“網(wǎng)絡(luò)學(xué)會(huì)了自動(dòng)對(duì)齊”?，F(xiàn)在很多文章一上來(lái)就講Q、K、V但如果你沒(méi)看過(guò)這個(gè)背景就會(huì)覺(jué)得QKV只是花哨的改名。其實(shí)Q就是“當(dāng)前的查詢”K是“序列里每個(gè)位置的鑰匙”V是“序列里每個(gè)位置的信息”。相關(guān)程度就是拿查詢和每一把鑰匙做匹配匹配度越高對(duì)應(yīng)信息被取出的比例就越大。這個(gè)視角貫穿所有變體包括Transformer里的自注意力。2.2 縮放點(diǎn)積Attention的完整公式為什么除以根號(hào)d_k現(xiàn)在主流公式長(zhǎng)這樣Attention(Q,K,V)softmax(QK^T/√d_k)V其中Q、K、V分別是查詢、鍵、值矩陣d_k是鍵向量的維度。計(jì)算分四步先把Q和K做點(diǎn)積得到相似度分?jǐn)?shù)再除以√d_k做縮放然后softmax歸一化成權(quán)重最后用權(quán)重對(duì)V做加權(quán)求和。QK^T如果形狀分別是[B, T_q, d_k]和[B, d_k, T_k]得到的就是[B, T_q, T_k]的分?jǐn)?shù)矩陣表示每個(gè)查詢對(duì)所有鍵的匹配程度。為什么要除以√d_k這是很多初學(xué)者最容易忽略的點(diǎn)。假設(shè)Q和K的元素是獨(dú)立的隨機(jī)變量均值0、方差1那么點(diǎn)積后的期望是0方差卻會(huì)變成d_k。d_k越大分?jǐn)?shù)的方差越大softmax輸入值會(huì)落到梯度極小的區(qū)域訓(xùn)練時(shí)梯度很容易消失。除以√d_k后方差被拉回接近1softmax的輸入范圍更穩(wěn)定。這個(gè)縮放是固定操作不需要學(xué)習(xí)也不要自作主張改成其他值它已經(jīng)是被反復(fù)驗(yàn)證過(guò)的默認(rèn)項(xiàng)。還有一種常見(jiàn)的變體叫加性注意力用一個(gè)小網(wǎng)絡(luò)來(lái)算匹配度而不是直接用點(diǎn)積。它在早期Seq2Seq里常見(jiàn)但縮放點(diǎn)積因?yàn)橛?jì)算更高效、更容易向量化后來(lái)成了主流。所以你現(xiàn)在看到“Attention”默認(rèn)都是指縮放點(diǎn)積注意力除非特別說(shuō)明。2.3 用PyTorch 20行實(shí)現(xiàn)單頭Attention最小可用代碼import torch import torch.nn.functional as F def scaled_dot_product_attention(query, key, value, maskNone): d_k query.size(-1) # 點(diǎn)積分?jǐn)?shù)形狀 [B, T_q, T_k] scores torch.matmul(query, key.transpose(-2, -1)) / (d_k ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) weights F.softmax(scores, dim-1) output torch.matmul(weights, value) return output, weights這段代碼里query、key、value的最后一維分別是d_k最前面兩維可以是[B, seq_len]。如果query長(zhǎng)度是T_q、key長(zhǎng)度是T_k那么輸出形狀是[B, T_q, d_v]。注意value的維度d_v可以和d_k不同但實(shí)際工程里通常相等方便做殘差拼接。mask參數(shù)采用“0表示屏蔽”的約定。把對(duì)應(yīng)分?jǐn)?shù)填充成負(fù)無(wú)窮softmax之后權(quán)重就是0。為什么要用負(fù)無(wú)窮而不是直接置0因?yàn)槿绻逊謹(jǐn)?shù)直接置0softmax后仍然會(huì)有小概率分配到該位置負(fù)無(wú)窮可以保證數(shù)學(xué)上徹底忽略。這個(gè)習(xí)慣在所有Transformer實(shí)現(xiàn)里都成立后面排查問(wèn)題時(shí)你也會(huì)用到它。寫(xiě)這個(gè)函數(shù)時(shí)我一般會(huì)順手打印出中間張量形狀。你拿一個(gè)B2、T_q3、T_k4的輸入跑一次就能看到scores變成[2,3,4]。先確認(rèn)維度再談效果這是避免后面對(duì)著頭數(shù)調(diào)半天卻沒(méi)用的一步。3. Self-Attention與Multi-HeadTransformer的基石怎么搭起來(lái)3.1 Self-Attention如何讓每個(gè)token看到全序列自注意力Self-Attention是Transformer的核心。一句話概括序列里的每個(gè)token都要看著整個(gè)序列重新計(jì)算自己的表示。它跟上一章的Seq2Seq注意力最大的區(qū)別是Q、K、V都來(lái)自同一個(gè)輸入。每個(gè)token先經(jīng)過(guò)三個(gè)線性投影變成自己的Q、K、V然后執(zhí)行縮放點(diǎn)積注意力。這樣一來(lái)每個(gè)位置都能直接和任意遠(yuǎn)的位置建立依賴卷積和循環(huán)網(wǎng)絡(luò)那種“隔著好幾層才能接觸”的問(wèn)題被繞開(kāi)了。舉個(gè)例子“他喜歡吃蘋(píng)果”和“她喜歡吃蘋(píng)果”只差一個(gè)字但語(yǔ)義完全變了。要理解“他”網(wǎng)絡(luò)需要看“喜歡”和“吃”等遠(yuǎn)距離信息。用自注意力這句話里每個(gè)詞兩兩之間都算一個(gè)分?jǐn)?shù)模型可以自由決定哪個(gè)詞對(duì)當(dāng)前詞更重要。訓(xùn)練時(shí)不帶任何人工規(guī)則全靠任務(wù)loss反向傳播來(lái)塑造這些分?jǐn)?shù)。但自注意力有個(gè)代價(jià)分?jǐn)?shù)矩陣是O(T^2)的。序列長(zhǎng)度從512漲到1024顯存占用直接翻4倍這在后面第5章會(huì)提到。所以工程上經(jīng)常會(huì)限制序列長(zhǎng)度或者用稀疏注意力、線形注意力等近似方案來(lái)處理超長(zhǎng)文本。3.2 Multi-Head Attention的拼接與投影多子空間帶來(lái)的增益單頭注意力只能學(xué)一種“匹配模式”實(shí)際任務(wù)里需要同時(shí)關(guān)注語(yǔ)法關(guān)系、指代關(guān)系、局部連貫性等多種信息。多頭注意力就是把d_model維拆成h個(gè)d_k維子空間每個(gè)頭獨(dú)立做注意力最后拼接起來(lái)再投影回d_model。谷歌最早提出的配置是d_model512、h8每個(gè)子空間d_k64。h個(gè)頭并行計(jì)算彼此不共享參數(shù)所以不同的頭會(huì)傾向于學(xué)到不同的注意力模式有的頭關(guān)注前后詞有的頭關(guān)注主語(yǔ)謂語(yǔ)這就是“多子空間”的價(jià)值。實(shí)現(xiàn)時(shí)要注意每個(gè)頭的維度d_k d_model / h所以d_model必須能整除h。這不只是一個(gè)數(shù)學(xué)要求更是一個(gè)工程約束。如果你選了不規(guī)則的頭數(shù)比如h7、d_model512就麻煩了。所以選頭數(shù)前先看d_model能不能被它整除。為了兼容預(yù)訓(xùn)練模型一般不要試圖去改這個(gè)比例。多頭注意力在計(jì)算結(jié)束后并不是簡(jiǎn)單把h個(gè)結(jié)果相加而是把所有頭的輸出在特征維度上拼接起來(lái)生成一個(gè)長(zhǎng)度為d_model的向量再做一次線性投影W_o。這次投影的目的是讓不同頭的信息做一次交叉融合而不是各說(shuō)各話。如果你在調(diào)試時(shí)發(fā)現(xiàn)效果不好檢查一下是不是忘了最后這個(gè)投影或者維度拼接時(shí)contiguous搞丟了這類問(wèn)題很常見(jiàn)。3.3 從零拼一個(gè)多頭注意力模塊代碼與維度核對(duì)import torch import torch.nn as nn import torch.nn.functional as F class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads, dropout0.1): super().__init__() assert d_model % num_heads 0 self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads self.W_q nn.Linear(d_model, d_model) self.W_k nn.Linear(d_model, d_model) self.W_v nn.Linear(d_model, d_model) self.W_o nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): batch_size, seq_len, _ x.size() # 投影并拆出頭維度: [B, T, h, d_k] - [B, h, T, d_k] Q self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) K self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) V self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) # 縮放點(diǎn)積注意力scores形狀 [B, h, T, T] scores Q K.transpose(-2, -1) / (self.d_k ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn F.softmax(scores, dim-1) attn self.dropout(attn) # 注意dropout在softmax之后 context attn V # [B, h, T, d_k] # 把頭拼回去: [B, T, h, d_k] - [B, T, d_model] context context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) output self.W_o(context) return output這段代碼里最容易被搞混的是view和transpose的順序。先用view把[B, T, d_model]拆成[B, T, h, d_k]再transpose成[B, h, T, d_k]這樣Q和K做矩陣乘時(shí)每個(gè)頭是獨(dú)立的。如果你先transpose再view會(huì)得到完全錯(cuò)亂的排列算出來(lái)的score沒(méi)有任何意義。mask參數(shù)可以傳兩種形狀。一是[B, 1, 1, T]的padding mask代表每個(gè)位置是否是有效token二是[B, 1, T, T]的causal mask下三角為1上三角為0用于自回歸生成。因?yàn)榫S度廣播機(jī)制[B, 1, T, T]會(huì)自動(dòng)廣播到[B, h, T, T]所以不需要復(fù)制多份。attention dropout放在softmax之后這是很多實(shí)現(xiàn)里的標(biāo)準(zhǔn)位置它隨機(jī)讓一部分注意力權(quán)重歸零相當(dāng)于正則化。我調(diào)試這段代碼時(shí)習(xí)慣用一個(gè)小例子batch_size2、seq_len5、d_model8、num_heads2跑一次打印Q、K、V的shape確認(rèn)最后輸出shape又變回[2,5,8]。如果哪一步維度對(duì)不上問(wèn)題一定出現(xiàn)在view或transpose上不要瞎調(diào)dropout。4. 訓(xùn)練與推理中的Attention參數(shù)怎么設(shè)d_model、頭數(shù)、dropout與mask4.1 頭數(shù)、維度和投影一個(gè)表格看清常見(jiàn)配置模型規(guī)模d_modelnum_headsd_k典型用途微型64416臨時(shí)Demo、小數(shù)據(jù)集小號(hào)128432輕量分類、短文本基礎(chǔ)512864常見(jiàn)Transformer基準(zhǔn)大號(hào)10241664長(zhǎng)文本、預(yù)訓(xùn)練模型頭數(shù)不是越大越好。頭越多每個(gè)頭的維度越小單個(gè)頭能表達(dá)的信息容量就越少。頭數(shù)增大到一定程度后不同頭會(huì)開(kāi)始“冗余”。很多實(shí)驗(yàn)表明8頭兩頭效果接近但不是翻倍就有翻倍效果。常見(jiàn)做法是保持每個(gè)頭d_k在32到64之間。如果你自己搭模型建議先從d_model128、heads4開(kāi)始跑通性能達(dá)標(biāo)后再往大擴(kuò)。直接照搬5128組合在小數(shù)據(jù)集上很容易過(guò)擬合。還有一個(gè)容易被忽略的參數(shù)輸出投影W_o是否帶bias。默認(rèn)nn.Linear是帶bias的但在某些實(shí)現(xiàn)里為了追求算力優(yōu)化會(huì)把它設(shè)成False。這不會(huì)顯著影響效果但如果你要加載預(yù)訓(xùn)練權(quán)重必須和原配置保持一致。另一個(gè)相關(guān)參數(shù)是注意力dropout和FFN的dropout它們通常是兩個(gè)獨(dú)立值。注意別把a(bǔ)ttention dropout設(shè)成0.5這會(huì)嚴(yán)重干擾學(xué)習(xí)一般0.1是個(gè)穩(wěn)定起點(diǎn)。4.2 訓(xùn)練階段必調(diào)的三個(gè)參數(shù)dropout、縮放因子和mask訓(xùn)練時(shí)最需要關(guān)注的三個(gè)點(diǎn)是dropout、縮放因子、mask它們各自都可能成為性能瓶頸。第一dropout。模型里通常有兩個(gè)dropout注意力權(quán)重上的dropout和FeedForward后的dropout。注意力權(quán)重dropout作用于softmax之后的權(quán)重它會(huì)讓某些位置的整條信息傳遞被阻斷防止個(gè)別token過(guò)度主導(dǎo)。如果你發(fā)現(xiàn)訓(xùn)練loss下降很慢先看它是不是設(shè)太大。我的習(xí)慣是注意力dropout0.1FFN dropout0.1到0.2。若你的數(shù)據(jù)集非常小可以考慮降到0.05但不要完全去掉。第二縮放因子。公式里固定用√d_k這個(gè)不需要調(diào)。不過(guò)有一種常見(jiàn)誤用把縮放因子當(dāng)作超參數(shù)去搜或者直接去掉。去掉縮放后大d_k會(huì)導(dǎo)致softmax梯度消失這在長(zhǎng)序列上尤其明顯。所以縮放因子不是可選項(xiàng)而是標(biāo)準(zhǔn)操作。你看到的某些第三方代碼里寫(xiě)了“scale”參數(shù)其實(shí)只是把除以√d_k換成了乘以scale本質(zhì)一樣。第三mask。訓(xùn)練時(shí)最常見(jiàn)的是padding mask和causal mask。padding mask負(fù)責(zé)把不足固定長(zhǎng)度的批量樣本填0的位置屏蔽掉causal mask負(fù)責(zé)保證自回歸模型看不到未來(lái)。如果你兩個(gè)都加需要把兩個(gè)mask做位與操作一個(gè)位置要想?yún)⑴c注意力必須同時(shí)滿足不是padding且在因果范圍內(nèi)。我見(jiàn)過(guò)很多翻車現(xiàn)場(chǎng)就是只加了padding mask忘了causal mask導(dǎo)致訓(xùn)練時(shí)模型提前看到了答案loss異常低但推理效果一塌糊涂。檢查方法很簡(jiǎn)單打印scores矩陣?yán)锉籱ask后的分布確認(rèn)右側(cè)上三角全是負(fù)無(wú)窮。4.3 推理階段的速度優(yōu)化緩存KV與attention的代價(jià)到了推理階段尤其是自回歸生成Attention的代價(jià)會(huì)變得非常扎眼。假設(shè)當(dāng)前生成了第n個(gè)token普通實(shí)現(xiàn)會(huì)把這n個(gè)token重新過(guò)一遍編碼器重新計(jì)算所有Q、K、V再做一次完整n×n的注意力。這就導(dǎo)致生成第n步的計(jì)算開(kāi)銷是O(n^2)總開(kāi)銷O(n^3)。實(shí)際項(xiàng)目里n稍微上千延遲和顯存就撐不住了。常用的優(yōu)化手段是KV Cache。因?yàn)榻獯a第n個(gè)token時(shí)前面所有token的K、V其實(shí)已經(jīng)算過(guò)而且不會(huì)變。只有Q是當(dāng)前步新算的。所以可以把每一步的K、V緩存下來(lái)下一步只需要算新token的K、V然后拼接到緩存矩陣上注意力分?jǐn)?shù)矩陣雖然還是n×n但計(jì)算量從“重新算所有n個(gè)token”降為“只算一個(gè)新的token”。在很多推理框架里你看到的關(guān)鍵詞“past_key_values”就是這個(gè)東西。實(shí)現(xiàn)KV Cache時(shí)要注意緩存的形狀是[B, h, T_prev, d_k]每一步新token的K/V是[B, h, 1, d_k]拼接后變成[B, h, T_prev1, d_k]。拼接發(fā)生在時(shí)間維度不要搞錯(cuò)。還有緩存必須和輸入長(zhǎng)度對(duì)得上如果你在批量生成不同長(zhǎng)度的序列需要給每個(gè)樣本維護(hù)一個(gè)獨(dú)立的cache否則短的序列會(huì)讀到長(zhǎng)的序列的K/V輸出直接崩掉。顯存方面即便有KV Cacheattention矩陣本身仍是O(T^2)。長(zhǎng)序列生成時(shí)這一步就會(huì)成為顯存瓶頸。常見(jiàn)的處理是把a(bǔ)ttention計(jì)算切成塊或者用FlashAttention這類kernel融合把中間矩陣不落地到全局顯存。作為一個(gè)獨(dú)立開(kāi)發(fā)者你最快能做的優(yōu)化是推理時(shí)把batch_size調(diào)小或者限制max_new_tokens讓attention矩陣不要爆掉。要是還不滿足再考慮FlashAttention或改用線性復(fù)雜度注意力。5. Attention的5個(gè)常見(jiàn)“玄學(xué)”坑排查順序與修正方案5.1 現(xiàn)象訓(xùn)練loss不降——原因mask把pad位置也參與注意力計(jì)算新手第一次搭Transformer最常遇到的問(wèn)題是loss一直在一個(gè)高位震蕩怎么調(diào)學(xué)習(xí)率都不動(dòng)。我見(jiàn)過(guò)某開(kāi)發(fā)者調(diào)試了一整晚最后發(fā)現(xiàn)是padding位置沒(méi)屏蔽。原因是批量輸入先pad到統(tǒng)一長(zhǎng)度比如一批里有3個(gè)樣本長(zhǎng)度分別是3、5、4pad到5。沒(méi)加mask的話第2個(gè)樣本的attention會(huì)看到后面兩個(gè)pad位置pad位置通常被初始化成0于是模型學(xué)會(huì)去關(guān)注這些廢位置有效信息被稀釋loss自然降不下去。解決方法是構(gòu)造padding mask把pad位置置0然后傳入attention函數(shù)在softmax之前把那些位置的分?jǐn)?shù)改為負(fù)無(wú)窮。mask形狀可以是[B, 1, T]廣播到所有頭。我在代碼里習(xí)慣用mask 0作為填充條件這樣mask里1表示有效、0表示pad。如果你用的庫(kù)是0表示有效、1表示pad記得反轉(zhuǎn)一下否則會(huì)把有效位置全屏蔽了loss直接飛起來(lái)。排查時(shí)先打印scores矩陣看pad位置的分?jǐn)?shù)是不是負(fù)無(wú)窮。如果是-1e9或者某個(gè)有限值說(shuō)明mask根本沒(méi)生效。再檢查mask是否在Q、K、V投影之后才傳入或者在多頭模塊里用的是mask而不是~mask。這些細(xì)節(jié)不看代碼一眼根本發(fā)現(xiàn)不了。5.2 現(xiàn)象多卡訓(xùn)練時(shí)顯存爆炸——原因attention矩陣O(n^2)被復(fù)制到每卡有人用單卡訓(xùn)練512長(zhǎng)度的序列很正常換成多卡DDP后明明batch_size沒(méi)變顯存卻直線上漲。attention機(jī)制本身有一個(gè)O(T^2)的分?jǐn)?shù)矩陣如果序列長(zhǎng)度是1024就是100萬(wàn)個(gè)浮點(diǎn)數(shù)單卡還能扛。但多卡DDP不是專人負(fù)責(zé)某一塊序列而是每張卡保留完整的batch副本和梯度副本所以attention矩陣也會(huì)在每個(gè)進(jìn)程里各自存一份。顯存占用是單卡的總和不會(huì)因?yàn)槟慵恿丝ň头值?。還有一個(gè)隱性膨脹點(diǎn)如果用PyTorch原生DDP且沒(méi)有開(kāi)啟gradient checkpointing反向傳播會(huì)保存所有中間激活包括注意力分?jǐn)?shù)和softmax權(quán)重。序列越長(zhǎng)激活越大。我對(duì)付這個(gè)問(wèn)題的順序是先開(kāi)自動(dòng)混合精度AMP把中間張量從FP32減到FP16顯存能砍一半左右然后開(kāi)gradient checkpointing犧牲一點(diǎn)重算時(shí)間換顯存最后如果還爆就只能降低batch_size或序列長(zhǎng)度。不要一開(kāi)始就懷疑是多卡配置問(wèn)題多數(shù)時(shí)候是attention矩陣大小乘以卡數(shù)超出物理顯存。5.3 現(xiàn)象推理結(jié)果重復(fù)——原因softmax溫度與attention dropout不一致生成任務(wù)里模型經(jīng)常輸出同一個(gè)詞反復(fù)循環(huán)比如“我愛(ài)吃我愛(ài)吃我愛(ài)吃”。很多人以為是解碼策略的懲罰參數(shù)沒(méi)調(diào)好但背后還有一個(gè)更隱蔽的坑訓(xùn)練時(shí)attention dropout隨機(jī)把一部分權(quán)重置零而推理時(shí)沒(méi)有開(kāi)dropout這會(huì)導(dǎo)致推理時(shí)的注意力分布和訓(xùn)練時(shí)不一樣尤其對(duì)短序列影響更大。如果訓(xùn)練時(shí)dropout較高模型在推理時(shí)看到的信息過(guò)于確定會(huì)更容易陷入重復(fù)模式。解決思路有兩類。一類是推理時(shí)也開(kāi)啟dropout但這會(huì)讓生成結(jié)果隨機(jī)性變大通常不建議。另一類是把注意力softmax的溫度調(diào)到小于1讓分布更尖銳減少隨機(jī)重復(fù)。具體做法是在softmax之前把scores除以一個(gè)溫度系數(shù)溫度越低輸出越確定。但溫度太低又會(huì)坍縮到最高概率的token反而重復(fù)。我的經(jīng)驗(yàn)是先確認(rèn)訓(xùn)練和推理的dropout設(shè)置一致再調(diào)溫度為0.8到0.9結(jié)合重復(fù)懲罰因子比單獨(dú)調(diào)一個(gè)參數(shù)有效。如果你用的是預(yù)訓(xùn)練模型還得檢查加載權(quán)重時(shí)是否強(qiáng)制設(shè)置了model.eval()。很多庫(kù)在推理時(shí)會(huì)自動(dòng)關(guān)閉dropout但如果你手工切換模型狀態(tài)晚了整個(gè)推理結(jié)果都會(huì)帶隨機(jī)性看起來(lái)就像“模型變笨了”。這種問(wèn)題耗時(shí)很長(zhǎng)但排查起來(lái)很簡(jiǎn)單打印推理次數(shù)如果兩次完全相同輸入給不同結(jié)果說(shuō)明dropout沒(méi)有關(guān)干凈。5.4 現(xiàn)象位置信息丟失——原因沒(méi)有加位置編碼或者加了但沒(méi)參與attention自注意力是置換等變的你把序列順序打亂得到的輸出只是同樣位置被打亂不會(huì)攜帶任何順序信息。這一點(diǎn)在理解上很容易踩坑。某同學(xué)在搭建情感分類模型時(shí)直接拿Embedding后的向量進(jìn)Transformer沒(méi)加任何位置編碼。模型訓(xùn)練出來(lái)準(zhǔn)確率很低因?yàn)樗X(jué)得attention已經(jīng)“記住了”每個(gè)詞的位置。其實(shí)attention只關(guān)注詞與詞之間的匹配關(guān)系并不會(huì)告訴你“第一個(gè)詞”是哪個(gè)絕對(duì)位置。解決方案是加位置編碼。常見(jiàn)的有固定正弦編碼和可學(xué)習(xí)位置編碼。正弦編碼根據(jù)公式生成不同頻率的正弦余弦值疊加到embedding上可學(xué)習(xí)編碼則是構(gòu)造一個(gè)形狀為[max_seq_len, d_model]的變量隨模型一起訓(xùn)練。無(wú)論哪種都要確保位置編碼在attention計(jì)算之前被加到詞嵌入向量上而不是加在輸出之后。否則attention看到的仍是原始embedding位置信息進(jìn)不了Q、K、V。調(diào)試時(shí)可以做一個(gè)簡(jiǎn)單驗(yàn)證輸入兩個(gè)完全相同詞但順序不同的句子比如“貓追狗”和“狗追貓”檢查模型最后的輸出表示是否不同。如果兩者完全一樣問(wèn)題幾乎肯定出在位置編碼沒(méi)生效。還要注意位置編碼的角色在多頭注意力中會(huì)被每個(gè)頭以不同方式調(diào)制所以加位置編碼的幅度不能太大一般和embedding同尺度否則會(huì)淹沒(méi)語(yǔ)義信息。5.5 現(xiàn)象可視化attention像個(gè)噪聲圖——原因多頭平均后互相抵消或者沒(méi)對(duì)齊tokenizer做模型可視化是注意力機(jī)制最有成就感的一步但經(jīng)常會(huì)得到一團(tuán)亂麻。有人把多頭注意力權(quán)重直接平均后畫(huà)出來(lái)發(fā)現(xiàn)完全看不出對(duì)角線或者局部模式。原因很簡(jiǎn)單每個(gè)頭關(guān)注的信息可能完全不同一個(gè)頭關(guān)注語(yǔ)法關(guān)系另一個(gè)頭關(guān)注長(zhǎng)距離依賴強(qiáng)行平均會(huì)把這些模式全部抵消掉。正確做法是分頭可視化先確認(rèn)每個(gè)頭的權(quán)重分布再合并。另一個(gè)崩潰點(diǎn)是tokenizer對(duì)齊。如果原始句子是中文按字切分但訓(xùn)練時(shí)用的tokenizer是BPE可視化時(shí)你拿字粒度的句子去對(duì)應(yīng)attention權(quán)重位置就會(huì)錯(cuò)位。需要先用同一個(gè)tokenizer把輸入轉(zhuǎn)換再拿解碼后的tokenid去映射注意力矩陣的行列。不然你圖上顯示的“注意力”根本不在這個(gè)字上看起來(lái)當(dāng)然像噪聲。最后的常見(jiàn)問(wèn)題是對(duì)角線權(quán)重過(guò)高。如果attention權(quán)重集中在對(duì)角線說(shuō)明模型幾乎沒(méi)有跨位置交互這通常發(fā)生在訓(xùn)練不足或序列長(zhǎng)度不夠時(shí)。這時(shí)不要急于下結(jié)論說(shuō)代碼有問(wèn)題先檢查是不是位置編碼沒(méi)有參與注意力或者是不是層數(shù)太少。我的習(xí)慣是同時(shí)打印scores的標(biāo)準(zhǔn)差如果標(biāo)準(zhǔn)差接近0說(shuō)明softmax已經(jīng)接近均勻分布模型還在“偷懶”需要增加訓(xùn)練步數(shù)。6. 把這張“PPT”講成落地技能三張圖講清原理一個(gè)腳本做原型驗(yàn)證如果你是要做分享與其堆公式不如準(zhǔn)備三張圖。第一張畫(huà)Seq2Seq不帶attention的瓶頸就是那個(gè)被固定向量壓扁的過(guò)程第二張畫(huà)QKV的映射關(guān)系用三個(gè)不同顏色的方塊表示Q、K、V再把點(diǎn)積到softmax到加權(quán)求和的過(guò)程用箭頭連起來(lái)第三張畫(huà)多頭在d_model維度上切出4條豎條每條算一次attention再拼回去。這三張圖能讓聽(tīng)的人五分鐘內(nèi)建立起直覺(jué)剩下的時(shí)間留給代碼和現(xiàn)象。驗(yàn)證自己是否真的理解了注意力機(jī)制我建議寫(xiě)一個(gè)極小的腳本隨機(jī)生成一組整數(shù)序列給每個(gè)位置接一個(gè)可學(xué)習(xí)的embedding然后只套一個(gè)單頭attention不加位置編碼用“預(yù)測(cè)序列中是否出現(xiàn)某個(gè)數(shù)字”這種玩具任務(wù)訓(xùn)練。如果attention能學(xué)會(huì)把權(quán)重放到那個(gè)數(shù)字出現(xiàn)的位置說(shuō)明你的實(shí)現(xiàn)是通的。這比直接跑一個(gè)大模型更省時(shí)間也能快速暴露mask或維度問(wèn)題。參考腳本的輸入輸出形狀對(duì)照第2章的函數(shù)即可。我早期做這個(gè)方向時(shí)曾經(jīng)為了一條維度對(duì)不上的錯(cuò)誤把多頭實(shí)現(xiàn)重寫(xiě)了三遍。后來(lái)養(yǎng)成一個(gè)習(xí)慣每個(gè)新模塊先跑一個(gè)batch把中間張量的shape全部打印出來(lái)確認(rèn)無(wú)誤再往上搭。注意力機(jī)制本身并不復(fù)雜復(fù)雜的是它和周圍模塊的邊界條件。你真正需要記住的是QKV是三種投影scores是匹配度mask是邊界dropout是正則KV Cache是代價(jià)。這些點(diǎn)串聯(lián)起來(lái)你再去讀任何Transformer源碼都會(huì)順暢很多。希望這篇筆記能幫你在注意力機(jī)制這條路上少走幾段彎路。本文還有配套的精品資源點(diǎn)擊獲取 注意力機(jī)制從原理到代碼：QKV、多頭、掩碼與KV Cache實(shí)戰(zhàn)指南
簡(jiǎn)介這是一份面向深度學(xué)習(xí)初學(xué)者和自然語(yǔ)言處理研究者的注意力機(jī)制專題課件系統(tǒng)講解序列到序列模型中的編碼器—解碼器結(jié)構(gòu)為何需要注意力以及固定長(zhǎng)度語(yǔ)義向量在長(zhǎng)輸入下的信息丟失問(wèn)題。課件從循環(huán)神經(jīng)網(wǎng)絡(luò)只能處理等長(zhǎng)序列的局限出發(fā)逐步引出注意力思想的來(lái)源并重點(diǎn)剖析Bahdanau提出的加性注意力、Luong提出的注意力變體以及硬注意力、局部注意力與自注意力等主流分支分別說(shuō)明其計(jì)算邏輯、優(yōu)缺點(diǎn)和適用任務(wù)。結(jié)合機(jī)器翻譯、文本摘要、語(yǔ)音識(shí)別、圖像描述等典型場(chǎng)景通過(guò)圖解和權(quán)重配色示意直觀呈現(xiàn)注意力分?jǐn)?shù)如何在不同輸入位置上分配幫助讀者理解每類注意力的工作機(jī)制。資源為1個(gè)PPT文件大小約2.7MB結(jié)構(gòu)緊湊、圖文并茂既適合教師備課參考也適合學(xué)生課前預(yù)習(xí)、課堂展示或考研復(fù)試快速回顧。目前已有1183人學(xué)習(xí)下載能夠顯著降低理解門(mén)檻也為后續(xù)學(xué)習(xí)Transformer架構(gòu)打下扎實(shí)基礎(chǔ)。1. 注意力機(jī)制不是玄學(xué)你在Transformer、Seq2Seq和圖像里早就見(jiàn)過(guò)它搜索“深度學(xué)習(xí)中attention注意力機(jī)制.ppt”的人多半是要準(zhǔn)備技術(shù)分享或者被面試題逼到了。注意力機(jī)制看起來(lái)像一套優(yōu)雅的數(shù)學(xué)公式但真正困擾開(kāi)發(fā)者的不是讀不懂論文而是“我怎么把它用起來(lái)”。這篇筆記會(huì)把你從“知道Attention”拉到“能寫(xiě)最小實(shí)現(xiàn)、能調(diào)參數(shù)、能排坑”的狀態(tài)。我不打算復(fù)制PPT里的定義而是把QKV、多頭、mask、緩存這些實(shí)操點(diǎn)拆開(kāi)每個(gè)環(huán)節(jié)都給出能直接跑的代碼和踩過(guò)的坑。你在Seq2Seq、Transformer、圖像分類里見(jiàn)到的各種注意力底層都是同一套加權(quán)求和邏輯。新手可以照步驟做熟手可以對(duì)照參數(shù)和邊界。2. 從“對(duì)齊”到“加權(quán)求和”Attention的核心數(shù)學(xué)直覺(jué)與QKV拆解2.1 早期Seq2Seq的Attention為什么要在譯每個(gè)詞時(shí)“回頭看”注意力機(jī)制最早出圈是在機(jī)器翻譯里。當(dāng)時(shí)常見(jiàn)的做法是用一個(gè)編碼器把整個(gè)輸入句子壓成一個(gè)固定長(zhǎng)度的向量解碼器每一步都從這個(gè)向量里取信息。句子一長(zhǎng)最后那個(gè)向量扛不住信息量前面的內(nèi)容基本被后面蓋掉了。后來(lái)某篇工作提出解碼時(shí)每一步不要把目光鎖死在那個(gè)固定向量上而是回到編碼器的每個(gè)隱藏狀態(tài)上“按需取用”。具體做法是計(jì)算當(dāng)前解碼狀態(tài)和每個(gè)編碼器隱藏狀態(tài)的相關(guān)程度得到一個(gè)權(quán)重分布然后用這個(gè)權(quán)重對(duì)編碼器隱藏狀態(tài)做加權(quán)求和得到這一步的上下文向量。這個(gè)“回頭看”的操作本質(zhì)就是給序列里的每個(gè)位置分配一個(gè)重要性分?jǐn)?shù)。分?jǐn)?shù)高的位置對(duì)當(dāng)前輸出的影響大分?jǐn)?shù)低的位置幾乎被忽略。這就是Attention最初的含義注意力權(quán)重告訴模型“該看哪個(gè)位置”加權(quán)和把“看到了什么”送到解碼器。它不需要額外的監(jiān)督信號(hào)權(quán)重完全由數(shù)據(jù)驅(qū)動(dòng)學(xué)出來(lái)所以經(jīng)常被形容為“網(wǎng)絡(luò)學(xué)會(huì)了自動(dòng)對(duì)齊”?，F(xiàn)在很多文章一上來(lái)就講Q、K、V但如果你沒(méi)看過(guò)這個(gè)背景就會(huì)覺(jué)得QKV只是花哨的改名。其實(shí)Q就是“當(dāng)前的查詢”K是“序列里每個(gè)位置的鑰匙”V是“序列里每個(gè)位置的信息”。相關(guān)程度就是拿查詢和每一把鑰匙做匹配匹配度越高對(duì)應(yīng)信息被取出的比例就越大。這個(gè)視角貫穿所有變體包括Transformer里的自注意力。2.2 縮放點(diǎn)積Attention的完整公式為什么除以根號(hào)d_k現(xiàn)在主流公式長(zhǎng)這樣Attention(Q,K,V)softmax(QK^T/√d_k)V其中Q、K、V分別是查詢、鍵、值矩陣d_k是鍵向量的維度。計(jì)算分四步先把Q和K做點(diǎn)積得到相似度分?jǐn)?shù)再除以√d_k做縮放然后softmax歸一化成權(quán)重最后用權(quán)重對(duì)V做加權(quán)求和。QK^T如果形狀分別是[B, T_q, d_k]和[B, d_k, T_k]得到的就是[B, T_q, T_k]的分?jǐn)?shù)矩陣表示每個(gè)查詢對(duì)所有鍵的匹配程度。為什么要除以√d_k這是很多初學(xué)者最容易忽略的點(diǎn)。假設(shè)Q和K的元素是獨(dú)立的隨機(jī)變量均值0、方差1那么點(diǎn)積后的期望是0方差卻會(huì)變成d_k。d_k越大分?jǐn)?shù)的方差越大softmax輸入值會(huì)落到梯度極小的區(qū)域訓(xùn)練時(shí)梯度很容易消失。除以√d_k后方差被拉回接近1softmax的輸入范圍更穩(wěn)定。這個(gè)縮放是固定操作不需要學(xué)習(xí)也不要自作主張改成其他值它已經(jīng)是被反復(fù)驗(yàn)證過(guò)的默認(rèn)項(xiàng)。還有一種常見(jiàn)的變體叫加性注意力用一個(gè)小網(wǎng)絡(luò)來(lái)算匹配度而不是直接用點(diǎn)積。它在早期Seq2Seq里常見(jiàn)但縮放點(diǎn)積因?yàn)橛?jì)算更高效、更容易向量化后來(lái)成了主流。所以你現(xiàn)在看到“Attention”默認(rèn)都是指縮放點(diǎn)積注意力除非特別說(shuō)明。2.3 用PyTorch 20行實(shí)現(xiàn)單頭Attention最小可用代碼import torch import torch.nn.functional as F def scaled_dot_product_attention(query, key, value, maskNone): d_k query.size(-1) # 點(diǎn)積分?jǐn)?shù)形狀 [B, T_q, T_k] scores torch.matmul(query, key.transpose(-2, -1)) / (d_k ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) weights F.softmax(scores, dim-1) output torch.matmul(weights, value) return output, weights這段代碼里query、key、value的最后一維分別是d_k最前面兩維可以是[B, seq_len]。如果query長(zhǎng)度是T_q、key長(zhǎng)度是T_k那么輸出形狀是[B, T_q, d_v]。注意value的維度d_v可以和d_k不同但實(shí)際工程里通常相等方便做殘差拼接。mask參數(shù)采用“0表示屏蔽”的約定。把對(duì)應(yīng)分?jǐn)?shù)填充成負(fù)無(wú)窮softmax之后權(quán)重就是0。為什么要用負(fù)無(wú)窮而不是直接置0因?yàn)槿绻逊謹(jǐn)?shù)直接置0softmax后仍然會(huì)有小概率分配到該位置負(fù)無(wú)窮可以保證數(shù)學(xué)上徹底忽略。這個(gè)習(xí)慣在所有Transformer實(shí)現(xiàn)里都成立后面排查問(wèn)題時(shí)你也會(huì)用到它。寫(xiě)這個(gè)函數(shù)時(shí)我一般會(huì)順手打印出中間張量形狀。你拿一個(gè)B2、T_q3、T_k4的輸入跑一次就能看到scores變成[2,3,4]。先確認(rèn)維度再談效果這是避免后面對(duì)著頭數(shù)調(diào)半天卻沒(méi)用的一步。3. Self-Attention與Multi-HeadTransformer的基石怎么搭起來(lái)3.1 Self-Attention如何讓每個(gè)token看到全序列自注意力Self-Attention是Transformer的核心。一句話概括序列里的每個(gè)token都要看著整個(gè)序列重新計(jì)算自己的表示。它跟上一章的Seq2Seq注意力最大的區(qū)別是Q、K、V都來(lái)自同一個(gè)輸入。每個(gè)token先經(jīng)過(guò)三個(gè)線性投影變成自己的Q、K、V然后執(zhí)行縮放點(diǎn)積注意力。這樣一來(lái)每個(gè)位置都能直接和任意遠(yuǎn)的位置建立依賴卷積和循環(huán)網(wǎng)絡(luò)那種“隔著好幾層才能接觸”的問(wèn)題被繞開(kāi)了。舉個(gè)例子“他喜歡吃蘋(píng)果”和“她喜歡吃蘋(píng)果”只差一個(gè)字但語(yǔ)義完全變了。要理解“他”網(wǎng)絡(luò)需要看“喜歡”和“吃”等遠(yuǎn)距離信息。用自注意力這句話里每個(gè)詞兩兩之間都算一個(gè)分?jǐn)?shù)模型可以自由決定哪個(gè)詞對(duì)當(dāng)前詞更重要。訓(xùn)練時(shí)不帶任何人工規(guī)則全靠任務(wù)loss反向傳播來(lái)塑造這些分?jǐn)?shù)。但自注意力有個(gè)代價(jià)分?jǐn)?shù)矩陣是O(T^2)的。序列長(zhǎng)度從512漲到1024顯存占用直接翻4倍這在后面第5章會(huì)提到。所以工程上經(jīng)常會(huì)限制序列長(zhǎng)度或者用稀疏注意力、線形注意力等近似方案來(lái)處理超長(zhǎng)文本。3.2 Multi-Head Attention的拼接與投影多子空間帶來(lái)的增益單頭注意力只能學(xué)一種“匹配模式”實(shí)際任務(wù)里需要同時(shí)關(guān)注語(yǔ)法關(guān)系、指代關(guān)系、局部連貫性等多種信息。多頭注意力就是把d_model維拆成h個(gè)d_k維子空間每個(gè)頭獨(dú)立做注意力最后拼接起來(lái)再投影回d_model。谷歌最早提出的配置是d_model512、h8每個(gè)子空間d_k64。h個(gè)頭并行計(jì)算彼此不共享參數(shù)所以不同的頭會(huì)傾向于學(xué)到不同的注意力模式有的頭關(guān)注前后詞有的頭關(guān)注主語(yǔ)謂語(yǔ)這就是“多子空間”的價(jià)值。實(shí)現(xiàn)時(shí)要注意每個(gè)頭的維度d_k d_model / h所以d_model必須能整除h。這不只是一個(gè)數(shù)學(xué)要求更是一個(gè)工程約束。如果你選了不規(guī)則的頭數(shù)比如h7、d_model512就麻煩了。所以選頭數(shù)前先看d_model能不能被它整除。為了兼容預(yù)訓(xùn)練模型一般不要試圖去改這個(gè)比例。多頭注意力在計(jì)算結(jié)束后并不是簡(jiǎn)單把h個(gè)結(jié)果相加而是把所有頭的輸出在特征維度上拼接起來(lái)生成一個(gè)長(zhǎng)度為d_model的向量再做一次線性投影W_o。這次投影的目的是讓不同頭的信息做一次交叉融合而不是各說(shuō)各話。如果你在調(diào)試時(shí)發(fā)現(xiàn)效果不好檢查一下是不是忘了最后這個(gè)投影或者維度拼接時(shí)contiguous搞丟了這類問(wèn)題很常見(jiàn)。3.3 從零拼一個(gè)多頭注意力模塊代碼與維度核對(duì)import torch import torch.nn as nn import torch.nn.functional as F class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads, dropout0.1): super().__init__() assert d_model % num_heads 0 self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads self.W_q nn.Linear(d_model, d_model) self.W_k nn.Linear(d_model, d_model) self.W_v nn.Linear(d_model, d_model) self.W_o nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): batch_size, seq_len, _ x.size() # 投影并拆出頭維度: [B, T, h, d_k] - [B, h, T, d_k] Q self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) K self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) V self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) # 縮放點(diǎn)積注意力scores形狀 [B, h, T, T] scores Q K.transpose(-2, -1) / (self.d_k ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn F.softmax(scores, dim-1) attn self.dropout(attn) # 注意dropout在softmax之后 context attn V # [B, h, T, d_k] # 把頭拼回去: [B, T, h, d_k] - [B, T, d_model] context context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) output self.W_o(context) return output這段代碼里最容易被搞混的是view和transpose的順序。先用view把[B, T, d_model]拆成[B, T, h, d_k]再transpose成[B, h, T, d_k]這樣Q和K做矩陣乘時(shí)每個(gè)頭是獨(dú)立的。如果你先transpose再view會(huì)得到完全錯(cuò)亂的排列算出來(lái)的score沒(méi)有任何意義。mask參數(shù)可以傳兩種形狀。一是[B, 1, 1, T]的padding mask代表每個(gè)位置是否是有效token二是[B, 1, T, T]的causal mask下三角為1上三角為0用于自回歸生成。因?yàn)榫S度廣播機(jī)制[B, 1, T, T]會(huì)自動(dòng)廣播到[B, h, T, T]所以不需要復(fù)制多份。attention dropout放在softmax之后這是很多實(shí)現(xiàn)里的標(biāo)準(zhǔn)位置它隨機(jī)讓一部分注意力權(quán)重歸零相當(dāng)于正則化。我調(diào)試這段代碼時(shí)習(xí)慣用一個(gè)小例子batch_size2、seq_len5、d_model8、num_heads2跑一次打印Q、K、V的shape確認(rèn)最后輸出shape又變回[2,5,8]。如果哪一步維度對(duì)不上問(wèn)題一定出現(xiàn)在view或transpose上不要瞎調(diào)dropout。4. 訓(xùn)練與推理中的Attention參數(shù)怎么設(shè)d_model、頭數(shù)、dropout與mask4.1 頭數(shù)、維度和投影一個(gè)表格看清常見(jiàn)配置模型規(guī)模d_modelnum_headsd_k典型用途微型64416臨時(shí)Demo、小數(shù)據(jù)集小號(hào)128432輕量分類、短文本基礎(chǔ)512864常見(jiàn)Transformer基準(zhǔn)大號(hào)10241664長(zhǎng)文本、預(yù)訓(xùn)練模型頭數(shù)不是越大越好。頭越多每個(gè)頭的維度越小單個(gè)頭能表達(dá)的信息容量就越少。頭數(shù)增大到一定程度后不同頭會(huì)開(kāi)始“冗余”。很多實(shí)驗(yàn)表明8頭兩頭效果接近但不是翻倍就有翻倍效果。常見(jiàn)做法是保持每個(gè)頭d_k在32到64之間。如果你自己搭模型建議先從d_model128、heads4開(kāi)始跑通性能達(dá)標(biāo)后再往大擴(kuò)。直接照搬5128組合在小數(shù)據(jù)集上很容易過(guò)擬合。還有一個(gè)容易被忽略的參數(shù)輸出投影W_o是否帶bias。默認(rèn)nn.Linear是帶bias的但在某些實(shí)現(xiàn)里為了追求算力優(yōu)化會(huì)把它設(shè)成False。這不會(huì)顯著影響效果但如果你要加載預(yù)訓(xùn)練權(quán)重必須和原配置保持一致。另一個(gè)相關(guān)參數(shù)是注意力dropout和FFN的dropout它們通常是兩個(gè)獨(dú)立值。注意別把a(bǔ)ttention dropout設(shè)成0.5這會(huì)嚴(yán)重干擾學(xué)習(xí)一般0.1是個(gè)穩(wěn)定起點(diǎn)。4.2 訓(xùn)練階段必調(diào)的三個(gè)參數(shù)dropout、縮放因子和mask訓(xùn)練時(shí)最需要關(guān)注的三個(gè)點(diǎn)是dropout、縮放因子、mask它們各自都可能成為性能瓶頸。第一dropout。模型里通常有兩個(gè)dropout注意力權(quán)重上的dropout和FeedForward后的dropout。注意力權(quán)重dropout作用于softmax之后的權(quán)重它會(huì)讓某些位置的整條信息傳遞被阻斷防止個(gè)別token過(guò)度主導(dǎo)。如果你發(fā)現(xiàn)訓(xùn)練loss下降很慢先看它是不是設(shè)太大。我的習(xí)慣是注意力dropout0.1FFN dropout0.1到0.2。若你的數(shù)據(jù)集非常小可以考慮降到0.05但不要完全去掉。第二縮放因子。公式里固定用√d_k這個(gè)不需要調(diào)。不過(guò)有一種常見(jiàn)誤用把縮放因子當(dāng)作超參數(shù)去搜或者直接去掉。去掉縮放后大d_k會(huì)導(dǎo)致softmax梯度消失這在長(zhǎng)序列上尤其明顯。所以縮放因子不是可選項(xiàng)而是標(biāo)準(zhǔn)操作。你看到的某些第三方代碼里寫(xiě)了“scale”參數(shù)其實(shí)只是把除以√d_k換成了乘以scale本質(zhì)一樣。第三mask。訓(xùn)練時(shí)最常見(jiàn)的是padding mask和causal mask。padding mask負(fù)責(zé)把不足固定長(zhǎng)度的批量樣本填0的位置屏蔽掉causal mask負(fù)責(zé)保證自回歸模型看不到未來(lái)。如果你兩個(gè)都加需要把兩個(gè)mask做位與操作一個(gè)位置要想?yún)⑴c注意力必須同時(shí)滿足不是padding且在因果范圍內(nèi)。我見(jiàn)過(guò)很多翻車現(xiàn)場(chǎng)就是只加了padding mask忘了causal mask導(dǎo)致訓(xùn)練時(shí)模型提前看到了答案loss異常低但推理效果一塌糊涂。檢查方法很簡(jiǎn)單打印scores矩陣?yán)锉籱ask后的分布確認(rèn)右側(cè)上三角全是負(fù)無(wú)窮。4.3 推理階段的速度優(yōu)化緩存KV與attention的代價(jià)到了推理階段尤其是自回歸生成Attention的代價(jià)會(huì)變得非常扎眼。假設(shè)當(dāng)前生成了第n個(gè)token普通實(shí)現(xiàn)會(huì)把這n個(gè)token重新過(guò)一遍編碼器重新計(jì)算所有Q、K、V再做一次完整n×n的注意力。這就導(dǎo)致生成第n步的計(jì)算開(kāi)銷是O(n^2)總開(kāi)銷O(n^3)。實(shí)際項(xiàng)目里n稍微上千延遲和顯存就撐不住了。常用的優(yōu)化手段是KV Cache。因?yàn)榻獯a第n個(gè)token時(shí)前面所有token的K、V其實(shí)已經(jīng)算過(guò)而且不會(huì)變。只有Q是當(dāng)前步新算的。所以可以把每一步的K、V緩存下來(lái)下一步只需要算新token的K、V然后拼接到緩存矩陣上注意力分?jǐn)?shù)矩陣雖然還是n×n但計(jì)算量從“重新算所有n個(gè)token”降為“只算一個(gè)新的token”。在很多推理框架里你看到的關(guān)鍵詞“past_key_values”就是這個(gè)東西。實(shí)現(xiàn)KV Cache時(shí)要注意緩存的形狀是[B, h, T_prev, d_k]每一步新token的K/V是[B, h, 1, d_k]拼接后變成[B, h, T_prev1, d_k]。拼接發(fā)生在時(shí)間維度不要搞錯(cuò)。還有緩存必須和輸入長(zhǎng)度對(duì)得上如果你在批量生成不同長(zhǎng)度的序列需要給每個(gè)樣本維護(hù)一個(gè)獨(dú)立的cache否則短的序列會(huì)讀到長(zhǎng)的序列的K/V輸出直接崩掉。顯存方面即便有KV Cacheattention矩陣本身仍是O(T^2)。長(zhǎng)序列生成時(shí)這一步就會(huì)成為顯存瓶頸。常見(jiàn)的處理是把a(bǔ)ttention計(jì)算切成塊或者用FlashAttention這類kernel融合把中間矩陣不落地到全局顯存。作為一個(gè)獨(dú)立開(kāi)發(fā)者你最快能做的優(yōu)化是推理時(shí)把batch_size調(diào)小或者限制max_new_tokens讓attention矩陣不要爆掉。要是還不滿足再考慮FlashAttention或改用線性復(fù)雜度注意力。5. Attention的5個(gè)常見(jiàn)“玄學(xué)”坑排查順序與修正方案5.1 現(xiàn)象訓(xùn)練loss不降——原因mask把pad位置也參與注意力計(jì)算新手第一次搭Transformer最常遇到的問(wèn)題是loss一直在一個(gè)高位震蕩怎么調(diào)學(xué)習(xí)率都不動(dòng)。我見(jiàn)過(guò)某開(kāi)發(fā)者調(diào)試了一整晚最后發(fā)現(xiàn)是padding位置沒(méi)屏蔽。原因是批量輸入先pad到統(tǒng)一長(zhǎng)度比如一批里有3個(gè)樣本長(zhǎng)度分別是3、5、4pad到5。沒(méi)加mask的話第2個(gè)樣本的attention會(huì)看到后面兩個(gè)pad位置pad位置通常被初始化成0于是模型學(xué)會(huì)去關(guān)注這些廢位置有效信息被稀釋loss自然降不下去。解決方法是構(gòu)造padding mask把pad位置置0然后傳入attention函數(shù)在softmax之前把那些位置的分?jǐn)?shù)改為負(fù)無(wú)窮。mask形狀可以是[B, 1, T]廣播到所有頭。我在代碼里習(xí)慣用mask 0作為填充條件這樣mask里1表示有效、0表示pad。如果你用的庫(kù)是0表示有效、1表示pad記得反轉(zhuǎn)一下否則會(huì)把有效位置全屏蔽了loss直接飛起來(lái)。排查時(shí)先打印scores矩陣看pad位置的分?jǐn)?shù)是不是負(fù)無(wú)窮。如果是-1e9或者某個(gè)有限值說(shuō)明mask根本沒(méi)生效。再檢查mask是否在Q、K、V投影之后才傳入或者在多頭模塊里用的是mask而不是~mask。這些細(xì)節(jié)不看代碼一眼根本發(fā)現(xiàn)不了。5.2 現(xiàn)象多卡訓(xùn)練時(shí)顯存爆炸——原因attention矩陣O(n^2)被復(fù)制到每卡有人用單卡訓(xùn)練512長(zhǎng)度的序列很正常換成多卡DDP后明明batch_size沒(méi)變顯存卻直線上漲。attention機(jī)制本身有一個(gè)O(T^2)的分?jǐn)?shù)矩陣如果序列長(zhǎng)度是1024就是100萬(wàn)個(gè)浮點(diǎn)數(shù)單卡還能扛。但多卡DDP不是專人負(fù)責(zé)某一塊序列而是每張卡保留完整的batch副本和梯度副本所以attention矩陣也會(huì)在每個(gè)進(jìn)程里各自存一份。顯存占用是單卡的總和不會(huì)因?yàn)槟慵恿丝ň头值?。還有一個(gè)隱性膨脹點(diǎn)如果用PyTorch原生DDP且沒(méi)有開(kāi)啟gradient checkpointing反向傳播會(huì)保存所有中間激活包括注意力分?jǐn)?shù)和softmax權(quán)重。序列越長(zhǎng)激活越大。我對(duì)付這個(gè)問(wèn)題的順序是先開(kāi)自動(dòng)混合精度AMP把中間張量從FP32減到FP16顯存能砍一半左右然后開(kāi)gradient checkpointing犧牲一點(diǎn)重算時(shí)間換顯存最后如果還爆就只能降低batch_size或序列長(zhǎng)度。不要一開(kāi)始就懷疑是多卡配置問(wèn)題多數(shù)時(shí)候是attention矩陣大小乘以卡數(shù)超出物理顯存。5.3 現(xiàn)象推理結(jié)果重復(fù)——原因softmax溫度與attention dropout不一致生成任務(wù)里模型經(jīng)常輸出同一個(gè)詞反復(fù)循環(huán)比如“我愛(ài)吃我愛(ài)吃我愛(ài)吃”。很多人以為是解碼策略的懲罰參數(shù)沒(méi)調(diào)好但背后還有一個(gè)更隱蔽的坑訓(xùn)練時(shí)attention dropout隨機(jī)把一部分權(quán)重置零而推理時(shí)沒(méi)有開(kāi)dropout這會(huì)導(dǎo)致推理時(shí)的注意力分布和訓(xùn)練時(shí)不一樣尤其對(duì)短序列影響更大。如果訓(xùn)練時(shí)dropout較高模型在推理時(shí)看到的信息過(guò)于確定會(huì)更容易陷入重復(fù)模式。解決思路有兩類。一類是推理時(shí)也開(kāi)啟dropout但這會(huì)讓生成結(jié)果隨機(jī)性變大通常不建議。另一類是把注意力softmax的溫度調(diào)到小于1讓分布更尖銳減少隨機(jī)重復(fù)。具體做法是在softmax之前把scores除以一個(gè)溫度系數(shù)溫度越低輸出越確定。但溫度太低又會(huì)坍縮到最高概率的token反而重復(fù)。我的經(jīng)驗(yàn)是先確認(rèn)訓(xùn)練和推理的dropout設(shè)置一致再調(diào)溫度為0.8到0.9結(jié)合重復(fù)懲罰因子比單獨(dú)調(diào)一個(gè)參數(shù)有效。如果你用的是預(yù)訓(xùn)練模型還得檢查加載權(quán)重時(shí)是否強(qiáng)制設(shè)置了model.eval()。很多庫(kù)在推理時(shí)會(huì)自動(dòng)關(guān)閉dropout但如果你手工切換模型狀態(tài)晚了整個(gè)推理結(jié)果都會(huì)帶隨機(jī)性看起來(lái)就像“模型變笨了”。這種問(wèn)題耗時(shí)很長(zhǎng)但排查起來(lái)很簡(jiǎn)單打印推理次數(shù)如果兩次完全相同輸入給不同結(jié)果說(shuō)明dropout沒(méi)有關(guān)干凈。5.4 現(xiàn)象位置信息丟失——原因沒(méi)有加位置編碼或者加了但沒(méi)參與attention自注意力是置換等變的你把序列順序打亂得到的輸出只是同樣位置被打亂不會(huì)攜帶任何順序信息。這一點(diǎn)在理解上很容易踩坑。某同學(xué)在搭建情感分類模型時(shí)直接拿Embedding后的向量進(jìn)Transformer沒(méi)加任何位置編碼。模型訓(xùn)練出來(lái)準(zhǔn)確率很低因?yàn)樗X(jué)得attention已經(jīng)“記住了”每個(gè)詞的位置。其實(shí)attention只關(guān)注詞與詞之間的匹配關(guān)系并不會(huì)告訴你“第一個(gè)詞”是哪個(gè)絕對(duì)位置。解決方案是加位置編碼。常見(jiàn)的有固定正弦編碼和可學(xué)習(xí)位置編碼。正弦編碼根據(jù)公式生成不同頻率的正弦余弦值疊加到embedding上可學(xué)習(xí)編碼則是構(gòu)造一個(gè)形狀為[max_seq_len, d_model]的變量隨模型一起訓(xùn)練。無(wú)論哪種都要確保位置編碼在attention計(jì)算之前被加到詞嵌入向量上而不是加在輸出之后。否則attention看到的仍是原始embedding位置信息進(jìn)不了Q、K、V。調(diào)試時(shí)可以做一個(gè)簡(jiǎn)單驗(yàn)證輸入兩個(gè)完全相同詞但順序不同的句子比如“貓追狗”和“狗追貓”檢查模型最后的輸出表示是否不同。如果兩者完全一樣問(wèn)題幾乎肯定出在位置編碼沒(méi)生效。還要注意位置編碼的角色在多頭注意力中會(huì)被每個(gè)頭以不同方式調(diào)制所以加位置編碼的幅度不能太大一般和embedding同尺度否則會(huì)淹沒(méi)語(yǔ)義信息。5.5 現(xiàn)象可視化attention像個(gè)噪聲圖——原因多頭平均后互相抵消或者沒(méi)對(duì)齊tokenizer做模型可視化是注意力機(jī)制最有成就感的一步但經(jīng)常會(huì)得到一團(tuán)亂麻。有人把多頭注意力權(quán)重直接平均后畫(huà)出來(lái)發(fā)現(xiàn)完全看不出對(duì)角線或者局部模式。原因很簡(jiǎn)單每個(gè)頭關(guān)注的信息可能完全不同一個(gè)頭關(guān)注語(yǔ)法關(guān)系另一個(gè)頭關(guān)注長(zhǎng)距離依賴強(qiáng)行平均會(huì)把這些模式全部抵消掉。正確做法是分頭可視化先確認(rèn)每個(gè)頭的權(quán)重分布再合并。另一個(gè)崩潰點(diǎn)是tokenizer對(duì)齊。如果原始句子是中文按字切分但訓(xùn)練時(shí)用的tokenizer是BPE可視化時(shí)你拿字粒度的句子去對(duì)應(yīng)attention權(quán)重位置就會(huì)錯(cuò)位。需要先用同一個(gè)tokenizer把輸入轉(zhuǎn)換再拿解碼后的tokenid去映射注意力矩陣的行列。不然你圖上顯示的“注意力”根本不在這個(gè)字上看起來(lái)當(dāng)然像噪聲。最后的常見(jiàn)問(wèn)題是對(duì)角線權(quán)重過(guò)高。如果attention權(quán)重集中在對(duì)角線說(shuō)明模型幾乎沒(méi)有跨位置交互這通常發(fā)生在訓(xùn)練不足或序列長(zhǎng)度不夠時(shí)。這時(shí)不要急于下結(jié)論說(shuō)代碼有問(wèn)題先檢查是不是位置編碼沒(méi)有參與注意力或者是不是層數(shù)太少。我的習(xí)慣是同時(shí)打印scores的標(biāo)準(zhǔn)差如果標(biāo)準(zhǔn)差接近0說(shuō)明softmax已經(jīng)接近均勻分布模型還在“偷懶”需要增加訓(xùn)練步數(shù)。6. 把這張“PPT”講成落地技能三張圖講清原理一個(gè)腳本做原型驗(yàn)證如果你是要做分享與其堆公式不如準(zhǔn)備三張圖。第一張畫(huà)Seq2Seq不帶attention的瓶頸就是那個(gè)被固定向量壓扁的過(guò)程第二張畫(huà)QKV的映射關(guān)系用三個(gè)不同顏色的方塊表示Q、K、V再把點(diǎn)積到softmax到加權(quán)求和的過(guò)程用箭頭連起來(lái)第三張畫(huà)多頭在d_model維度上切出4條豎條每條算一次attention再拼回去。這三張圖能讓聽(tīng)的人五分鐘內(nèi)建立起直覺(jué)剩下的時(shí)間留給代碼和現(xiàn)象。驗(yàn)證自己是否真的理解了注意力機(jī)制我建議寫(xiě)一個(gè)極小的腳本隨機(jī)生成一組整數(shù)序列給每個(gè)位置接一個(gè)可學(xué)習(xí)的embedding然后只套一個(gè)單頭attention不加位置編碼用“預(yù)測(cè)序列中是否出現(xiàn)某個(gè)數(shù)字”這種玩具任務(wù)訓(xùn)練。如果attention能學(xué)會(huì)把權(quán)重放到那個(gè)數(shù)字出現(xiàn)的位置說(shuō)明你的實(shí)現(xiàn)是通的。這比直接跑一個(gè)大模型更省時(shí)間也能快速暴露mask或維度問(wèn)題。參考腳本的輸入輸出形狀對(duì)照第2章的函數(shù)即可。我早期做這個(gè)方向時(shí)曾經(jīng)為了一條維度對(duì)不上的錯(cuò)誤把多頭實(shí)現(xiàn)重寫(xiě)了三遍。后來(lái)養(yǎng)成一個(gè)習(xí)慣每個(gè)新模塊先跑一個(gè)batch把中間張量的shape全部打印出來(lái)確認(rèn)無(wú)誤再往上搭。注意力機(jī)制本身并不復(fù)雜復(fù)雜的是它和周圍模塊的邊界條件。你真正需要記住的是QKV是三種投影scores是匹配度mask是邊界dropout是正則KV Cache是代價(jià)。這些點(diǎn)串聯(lián)起來(lái)你再去讀任何Transformer源碼都會(huì)順暢很多。希望這篇筆記能幫你在注意力機(jī)制這條路上少走幾段彎路。本文還有配套的精品資源點(diǎn)擊獲取