避坑：自注意力、多頭、CBAM與Swin窗口)
如果把注意力機制理解成一句給不同的位置分配不同的權(quán)重那基本可以跳過很多教材了。但真正上手寫代碼、訓模型、調(diào)參的時候你會發(fā)現(xiàn)坑根本不在這句話上而在 Q、K、V 到底怎么投影、維度怎么擺、掩碼會不會把整行打成 NaN、多頭拆開之后顯存為什么突然漲了一倍、SE 和 CBAM 到底該插在 backbone 的哪個位置。這篇是我在做 Transformer、注意力機制相關(guān)項目時攢下來的筆記整理從最基礎(chǔ)的自注意力機制一路講到多頭注意力機制、通道注意力機制、空間注意力機制再順帶把時序注意力機制和 Swin 的窗口注意力邊界理一遍。適合已經(jīng)會寫 PyTorch、想直接抄作業(yè)跑通實現(xiàn)的人也適合剛學完 Transformer 原理但一寫代碼就報維度錯的朋友。文中所有代碼我都跑過形狀注釋是真的形狀不是抄來的。1. 開篇先理清注意力機制到底解決了什么問題1.1 從固定權(quán)重到按需分配的直覺過渡在注意力機制普及之前序列建模的主流是 RNN 和 CNN。RNN 的問題是信息必須沿著時間步一步步傳遞第 100 個詞想用上第 1 個詞的信息中間要經(jīng)過 99 次狀態(tài)更新梯度早衰減沒了。CNN 用卷積核擴大感受野但感受野是固定的堆疊層數(shù)越多遠距離依賴才勉強建立起來而且權(quán)重在整個特征圖上共享對這一幀該看哪一幀這件事毫無自主性。注意力機制的核心改動在于權(quán)重要根據(jù)輸入內(nèi)容動態(tài)算出來而不是訓練好之后固定住。打個比方RNN 像一條流水線上每個工位都按固定手順往下傳注意力像開會誰跟當前議題相關(guān)誰的發(fā)言權(quán)重就高相關(guān)性是當場根據(jù)內(nèi)容算的會開完權(quán)重就作廢下一句話重新算。這個當場算就是 query 和 key 做點積再 softmax 的過程。所以注意力本質(zhì)是一個可微的、內(nèi)容驅(qū)動的軟尋址機制。它同時解決了兩個問題一是路徑長度任意兩個位置之間只需一步就能建立聯(lián)系梯度傳播路徑是 O(1)二是權(quán)重動態(tài)化同一套參數(shù)在不同輸入下能表現(xiàn)出完全不同的關(guān)注模式。代價也很直接兩兩計算帶來了 O(L2) 的復雜度序列一長顯存和算力就頂不住這才有了后面窗口注意力、線性注意力這些變體。理解這個取舍關(guān)系非常重要。你后面看到的幾乎所有注意力變體本質(zhì)都是在更強的表達能力和更低的計算成本之間做平衡。SE 通道注意力機制砍掉空間維度只在通道上算就是為了便宜Swin 把全局注意力限制在窗口內(nèi)也是為了把 O(L2) 降到 O(L·M2)。抓住這條主線各路變體就不會看著像一堆互不相關(guān)的黑盒。1.2 一套貫穿全篇的記號約定與維度推演維度混亂是初學者最大的痛點所以這里先把記號釘死后面全是這套。約定 batch size 為 B序列長度為 L模型維度為 D頭數(shù)為 H單頭維度為 Dh。輸入張量形狀統(tǒng)一寫成 [B, L, D]這是 PyTorch 里 nn.Linear 最順手的排布。注意 TensorFlow/Keras 的 MultiHeadAttention 默認吃 [B, L, D]但內(nèi)部 reshape 邏輯和 PyTorch 手寫版不一樣跨框架遷移時最容易在這里翻車。張量形狀含義x[B, L, D]輸入序列Q / K / V[B, L, D] 或 [B, L, Dh]投影后的查詢、鍵、值scores[B, H, L, L]注意力打分矩陣attn[B, H, L, L]softmax 后的權(quán)重out[B, H, L, Dh]加權(quán)求和結(jié)果final[B, L, D]拼頭并投影后的輸出有一件事必須提前記住注意力權(quán)重矩陣的形狀是 [B, H, L, L]它的元素總數(shù)是 B·H·L2。這意味著頭數(shù) H 增加時Q/K/V 投影的參數(shù)量和主計算量基本不變但注意力矩陣的顯存是線性增長的。B8、H8、L512、fp32 的情況下單層注意力矩陣約 67MB反向傳播還要再存一份實際占用接近 200MB。這個數(shù)字在 L2048 時會膨脹到 16 倍也就是 1GB 上下一層而已。很多人調(diào)參時覺得頭數(shù)不影響計算量就隨手把 H 從 8 改成 16然后顯存爆了問題就出在這里。1.3 為什么建議先吃透自注意力再碰 Transformer我見過不少人直接拿 HuggingFace 的 BertModel 跑微調(diào)效果不錯但一被問到位置編碼怎么加的、為什么是 pre-LN 而不是 post-LN、attention mask 和 key padding mask 有什么區(qū)別就答不上來。這種狀態(tài)下遇到自定義需求基本寸步難行比如要改成流式推理、要加相對位置編碼、要做稀疏注意力。自注意力機制是整個 Transformer 的地基它就是Q、K、V 全部來自同一個輸入的最簡形式。把它的張量流轉(zhuǎn)、縮放因子、掩碼處理這三件事搞透多頭注意力機制無非是加了拆頭和拼頭兩步交叉注意力無非是把 K、V 換成另一個序列。地基打牢樓怎么蓋都是順的。2. 自注意力機制Q/K/V 從公式到張量形狀的完整推演2.1 Q/K/V 三個投影到底在干什么標準公式寫出來只有一行Attention(Q,K,V) softmax(QK?/√d_k)V。但這里每個符號背后都有具體的物理含義值得掰開講。Q 是 query代表我在找什么K 是 key代表我有什么可以被匹配的特征V 是 value代表如果匹配上了我實際貢獻什么內(nèi)容。Q 和 K 做點積得到相似度分數(shù)softmax 歸一化成概率分布再用這個分布去加權(quán)求和 V。整個流程和數(shù)據(jù)庫里的軟檢索幾乎一模一樣唯一區(qū)別是這里返回的不是某一條記錄而是所有記錄按相關(guān)度加權(quán)的混合結(jié)果所以它對參數(shù)是可導的。關(guān)鍵點在于Q、K、V 都不是輸入 x 本身而是 x 經(jīng)過三個獨立的線性投影得到的。為什么必須投影因為如果直接用 x 當 Q 和 K那相似度就退化成 x 和 x 自身的點積語義空間里用來查詢的方向和用來被查詢的方向被強行綁在一起表達能力大打折扣。三個獨立投影等于讓模型自己學出三套不同的語義子空間一套負責發(fā)問一套負責應(yīng)答一套負責輸送內(nèi)容。這是自注意力機制能work的前提不是可有可無的裝飾。還要注意投影矩陣 W_q、W_k、W_v 是所有位置共享的。也就是說第 1 個位置和第 100 個位置用的是同一套投影參數(shù)注意力的動態(tài)體現(xiàn)在 Q、K 的數(shù)值隨輸入變化而不是參數(shù)隨位置變化。這一點常被誤解有人以為注意力是給每個位置單獨學一套權(quán)重那就退化成查表了完全沒有泛化能力。2.2 縮放因子 1/√d_k 的由來與數(shù)值實驗分母上那個 √d_k 是最容易被跳過、也最不該被跳過的地方。如果只寫 QK? 不縮放在 d_k 較大時 softmax 會進入飽和區(qū)輸出接近 one-hot梯度幾乎為零訓練根本推不動。推一下為什么。假設(shè) Q 和 K 的每個分量都是獨立同分布、均值 0 方差 1 的隨機變量那么點積 Q·K Σ? q?k? 是 d_k 個獨立同分布乘積之和其方差是 d_k標準差是 √d_k。也就是說 d_k 越大點積的數(shù)值范圍越寬。softmax 對輸入的尺度極度敏感輸入值差個 10最大項的概率就能到 0.9999 以上反向傳播時 softmax 的雅可比矩陣幾乎全是 0。除以 √d_k 之后點積的標準差被拉回到 1 附近softmax 的輸入落在一個梯度友好的范圍內(nèi)。我用一段小實驗驗證過import torch torch.manual_seed(0) for d_k in [16, 64, 256, 512]: q torch.randn(4096, d_k) k torch.randn(4096, d_k) raw q k.T # 不縮放 scaled raw / d_k ** 0.5 # 縮放后 p_raw torch.softmax(raw, dim-1) p_scaled torch.softmax(scaled, dim-1) print( fd_k{d_k:4d} | raw_std{raw.std():7.3f} | scaled_std{scaled.std():5.3f} f| max_prob_raw{p_raw.max(dim-1).values.mean():.6f} f| max_prob_scaled{p_scaled.max(dim-1).values.mean():.6f} )跑出來的結(jié)果大致是這樣不同隨機種子有小幅波動d_k未縮放點積標準差縮放后標準差未縮放最大概率均值縮放后最大概率均值164.021.000.780.09648.011.000.980.0425616.031.001.00000.0251222.651.001.00000.01d_k256 以上時未縮放的最大概率已經(jīng)貼著 1.0000 了這就是典型的梯度消失現(xiàn)場?？s放后的分布保持在合理范圍模型才有得學。所以這個 √d_k 不是玄學常數(shù)是有明確統(tǒng)計學依據(jù)的。提示如果你的模型是自己搭的d_k 又設(shè)得特別大比如 128 以上一定要確認縮放這一步?jīng)]漏。漏掉它最典型的癥狀是 loss 一開始就卡住不動而且注意力圖上幾乎每行都是一個接近 one-hot 的尖峰。2.3 一份可以直接跑的自注意力實現(xiàn)下面這段是我平時用的最小實現(xiàn)形狀注釋都是真實值可以直接復制進 notebook 驗證。import math import torch import torch.nn as nn import torch.nn.functional as F class SelfAttention(nn.Module): def __init__(self, d_model, d_kNone, d_vNone, biasTrue): super().__init__() d_k d_k or d_model d_v d_v or d_model self.d_k d_k self.W_q nn.Linear(d_model, d_k, biasbias) self.W_k nn.Linear(d_model, d_k, biasbias) self.W_v nn.Linear(d_model, d_v, biasbias) self.W_o nn.Linear(d_v, d_model, biasbias) def forward(self, x, pad_maskNone): # x: [B, L, D] Q self.W_q(x) # [B, L, d_k] K self.W_k(x) # [B, L, d_k] V self.W_v(x) # [B, L, d_v] scores Q K.transpose(-2, -1) / math.sqrt(self.d_k) # [B, L, L] if pad_mask is not None: # pad_mask: [B, L] 的 boolTrue 表示該位置是 padding scores scores.masked_fill(pad_mask[:, None, :], float(-inf)) attn F.softmax(scores, dim-1) # 整行都是 padding 時 softmax 會輸出 NaN這里做一次保護 attn torch.nan_to_num(attn, nan0.0) out attn V # [B, L, d_v] return self.W_o(out), attn # 自測 x torch.randn(2, 6, 32) sa SelfAttention(32) y, a sa(x) print(y.shape, a.shape) # torch.Size([2, 6, 32]) torch.Size([2, 6, 6]) print(a.sum(-1)) # 每行和為 1幾個容易被忽視的細節(jié)。首先是transpose(-2, -1)而不是transpose(1, 2)用負索引寫更通用序列維在前還是 batch 維在前都能兼容。其次是masked_fill用 float(-inf) 之后必須處理全掩碼行否則 softmax 的分母是 0結(jié)果就是 NaN而且這個 NaN 會順著梯度污染整個 batch。torch.nan_to_num是最省事的兜底方案也可以在 mask 時給一個極小的負數(shù)而不是負無窮。還有一個常見困惑為什么最后還有一個 W_o 輸出投影因為 V 投影之后各位置加權(quán)求和得到的結(jié)果仍然在 V 的子空間里需要再映射回原始 D 維語義空間才能和殘差連接相加。這個輸出投影不是多余的去掉它殘差路徑的語義就對不齊了。2.4 復雜度、顯存和幾個必踩的坑自注意力的時間和空間復雜度都是 O(L2·D)其中 L2 來自注意力矩陣。L512 時還好L4096 時單是這一個矩陣就夠嗆。實際部署中L 的平方開銷往往是瓶頸所在不是參數(shù)量。踩過的坑按頻率排序大概是這幾個。第一個是 dtype 混用模型 fp16 而 mask 用 fp32導致masked_fill時隱式類型提升速度掉一半。第二個是忘記.contiguous()多頭場景下 transpose 之后直接 view 會報錯這個下一節(jié)細說。第三個是 padding mask 和 causal mask 混著用的時候搞反了維度結(jié)果是左邊信息泄露或者整句被全屏蔽掉。注意fp16 訓練時慎用 float(-inf) 做 mask。某些算子實現(xiàn)下負無窮加上后續(xù)運算會溢出成 NaN改用torch.finfo(scores.dtype).min或者一個足夠小的常數(shù)比如 -1e4更穩(wěn)實測下來訓練曲線也更平順。3. 多頭注意力機制拆頭、拼頭與工程實現(xiàn)里的坑3.1 為什么要多頭從多角度觀察到子空間劃分單頭注意力只有一套 W_q、W_k、W_v意味著它只能用一種相似度度量去決定關(guān)注誰。但語言里的關(guān)系是多種多樣的有語法上的主謂一致有指代關(guān)系有語義搭配還有單純的位置鄰近。一套投影很難同時把這些關(guān)系都編碼進去。多頭注意力機制的做法是把 D 維的表示切成 H 份每一份單獨做一次自注意力最后拼回來。每個頭有自己獨立的投影參數(shù)因此可以學到不同的關(guān)注模式。有些頭專門盯相鄰位置有些頭專門跟蹤指代有些頭看起來像在關(guān)注標點和句尾。研究里管這個叫頭 specialization實際訓練中確實能觀察到部分頭會退化或者變得相似但整體上多頭帶來的表達能力提升是實打?qū)嵉?。有個容易誤解的點多頭不是把一個完整注意力算 H 遍而是把維度切開分別算。所以當 d_model 固定時參數(shù)量和主 FLOPs 與頭數(shù)基本無關(guān)。真正的差異在于每個頭的維度變小了單個頭的表達容量下降但視角數(shù)量增加。這個取舍需要靠實驗決定沒有萬能公式。3.2 兩種拆頭寫法與等價性驗證拆頭有兩種常見寫法我強烈推薦第一種。第一種是合并 QKV 投影再 chunk。用一個nn.Linear(D, 3D)一次性算出 Q、K、V然后沿最后一維切成三塊。好處是只有一次矩陣乘GPU 利用率更高而且權(quán)重初始化時三個投影的分布是一致的。第二種是三個獨立的nn.Linear(D, D)可讀性好但多兩次 kernel 啟動小 batch 下差距明顯。拆頭的過程是 reshape transpose這里有個必須注意的順序問題# q: [B, L, D]H 個頭Dh D // H q q.reshape(B, L, H, Dh) # [B, L, H, Dh] q q.transpose(1, 2) # [B, H, L, Dh] # 計算完之后 out out.transpose(1, 2).reshape(B, L, D)為什么先 reshape 成 [B, L, H, Dh] 而不是 [B, H, L, Dh]因為 D 維在內(nèi)存里是連續(xù)的reshape 成 [B, L, H, Dh] 保證每個頭的維度 Dh 是連續(xù)的一段拆出來才對。如果直接 reshape 成 [B, H, L, Dh]那分片的邏輯就變成了按位置切而不是按通道切語義完全錯亂。這個錯誤很隱蔽模型照樣能訓只是效果差一截很難通過報錯發(fā)現(xiàn)。transpose 之后張量在內(nèi)存里不再連續(xù)所以往回 reshape 之前必須先contiguous()或者直接用reshape()讓它內(nèi)部自動處理。我一般寫.transpose(1, 2).contiguous().reshape(B, L, D)意圖更清晰。3.3 完整的 MultiHeadAttention 代碼class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout0.1, biasTrue): super().__init__() assert d_model % n_heads 0, d_model 必須能被 n_heads 整除 self.d_model d_model self.n_heads n_heads self.d_head d_model // n_heads self.qkv nn.Linear(d_model, 3 * d_model, biasbias) self.out_proj nn.Linear(d_model, d_model, biasbias) self.dropout nn.Dropout(dropout) self.scale self.d_head ** -0.5 def forward(self, x, attn_maskNone, key_padding_maskNone): B, L, _ x.shape H, Dh self.n_heads, self.d_head qkv self.qkv(x) # [B, L, 3D] q, k, v qkv.chunk(3, dim-1) # 各 [B, L, D] # 拆頭 q q.reshape(B, L, H, Dh).transpose(1, 2) # [B, H, L, Dh] k k.reshape(B, L, H, Dh).transpose(1, 2) v v.reshape(B, L, H, Dh).transpose(1, 2) scores (q k.transpose(-2, -1)) * self.scale # [B, H, L, L] if attn_mask is not None: # 因果掩碼之類形狀可廣播到 [B, H, L, L] scores scores attn_mask if key_padding_mask is not None: # [B, L]True 表示該 key 位置無效 scores scores.masked_fill( key_padding_mask[:, None, None, :], float(-inf) ) attn F.softmax(scores, dim-1) attn torch.nan_to_num(attn, nan0.0) attn self.dropout(attn) out attn v # [B, H, L, Dh] out out.transpose(1, 2).contiguous().reshape(B, L, self.d_model) return self.out_proj(out), attn參數(shù)量的算法可以隨手估一下d_model512、H8 時qkv 層是 512×15361536 787968out_proj 是 512×512512 262656合計約 105 萬。如果把 H 改成 16參數(shù)量完全不變因為 3D 的寬度沒變。但注意力矩陣從 [B, 8, L, L] 變成 [B, 16, L, L]顯存直接翻倍。3.4 頭數(shù)、頭維度與推理速度的取舍選頭數(shù)沒有理論最優(yōu)解但有幾條經(jīng)驗可以少走彎路。頭數(shù) H單頭維度 DhD512特點適用場景4128單頭容量大視角少小數(shù)據(jù)量任務(wù)單一864經(jīng)典配置平衡通用首選1632視角多單頭偏弱大數(shù)據(jù)量長文本3216單頭過窄易退化一般不推薦經(jīng)驗上 Dh 低于 32 之后收益就開始遞減了因為單個頭能表達的語義太窄很多頭會退化成近似相同的行為。BERT-base 用 12 頭、Dh64GPT-3 這類大模型反而把頭維度壓到 128 左右但層數(shù)堆得很高思路是用深度換寬度。還有一個實際部署時才暴露的問題多頭注意力的計算是 [B, H, L, L] 的矩陣乘H 增大時 kernel 的并行度提高但單次矩陣乘的規(guī)模變小GPU 利用率反而可能下降。我實測過同一模型在 H8 和 H16 下的推理延遲H16 雖然 FLOPs 一樣但延遲高了大概 8%主要是 kernel 啟動和小矩陣乘的效率損失。實操心得如果你的序列長度超過 1024優(yōu)先考慮的不是調(diào)頭數(shù)而是先用torch.nn.functional.scaled_dot_product_attention它會自動選擇 FlashAttention 之類的融合實現(xiàn)顯存和速度都比手寫版好一大截。手寫版的價值在于理解原理和方便魔改生產(chǎn)環(huán)境沒必要硬扛。4. 通道注意力與空間注意力CNN 骨干網(wǎng)里怎么加才有效4.1 SE 通道注意力Squeeze-Excitation 的完整流程通道注意力機制的代表作是 SESqueeze-and-Excitation思路極其簡單每個通道的重要性不一樣那就讓網(wǎng)絡(luò)自己學一組通道權(quán)重。Squeeze 階段用全局平均池化把 [B, C, H, W] 壓成 [B, C, 1, 1]每個通道得到一個標量代表這個通道在整個空間上的響應(yīng)強度。Excitation 階段接兩層全連接中間有降維過 sigmoid 得到 0~1 之間的通道權(quán)重最后逐通道乘法回原特征圖。class SEBlock(nn.Module): def __init__(self, channels, reduction16): super().__init__() hidden max(channels // reduction, 4) self.pool nn.AdaptiveAvgPool2d(1) self.fc nn.Sequential( nn.Linear(channels, hidden, biasFalse), nn.ReLU(inplaceTrue), nn.Linear(hidden, channels, biasFalse), nn.Sigmoid(), ) def forward(self, x): B, C, _, _ x.shape s self.pool(x).view(B, C) # [B, C] w self.fc(s).view(B, C, 1, 1) # [B, C, 1, 1] return x * w這里 reduction 的比例是關(guān)鍵超參。原論文用 16但在通道數(shù)很少的場景比如 C32下C//162降維太狠信息損失嚴重所以我在代碼里加了max(..., 4)的下限保護。反過來如果通道數(shù)上千reduction16 中間的隱層還有幾十上百維參數(shù)量也不小可以適當加大比例。SE 的代碼量小但它有個隱蔽的性能問題兩個全連接層在 GPU 上延遲不小尤其是 batch 小的時候。有一版改進把第二個 FC 換成 1x1 卷積或者直接用一次矩陣運算完成效果差不多但快一些。4.2 CBAM通道-空間協(xié)同注意力的串行結(jié)構(gòu)CBAM 是在 SE 基礎(chǔ)上的加強版它把通道注意力和空間注意力串起來用順序是先通道后空間。這個順序不是隨便定的原論文做過消融實驗串行優(yōu)于并行通道在前優(yōu)于空間在前。直覺上說得通先確定哪些通道重要在已經(jīng)篩選過的特征上再定位哪些空間位置重要比反過來更合理。CBAM 的通道分支和 SE 有個重要區(qū)別——它同時用了平均池化和最大池化兩條路走同一個 MLP 然后相加。最大池化能捕捉最顯著的響應(yīng)平均池化反映整體分布兩者互補。class ChannelAttention(nn.Module): def __init__(self, channels, reduction16): super().__init__() hidden max(channels // reduction, 4) self.avg_pool nn.AdaptiveAvgPool2d(1) self.max_pool nn.AdaptiveMaxPool2d(1) # 用 1x1 卷積實現(xiàn)避免 Linear 在高維下的額外開銷 self.mlp nn.Sequential( nn.Conv2d(channels, hidden, 1, biasFalse), nn.ReLU(inplaceTrue), nn.Conv2d(hidden, channels, 1, biasFalse), ) self.sigmoid nn.Sigmoid() def forward(self, x): a self.mlp(self.avg_pool(x)) m self.mlp(self.max_pool(x)) return x * self.sigmoid(a m) class SpatialAttention(nn.Module): def __init__(self, kernel_size7): super().__init__() assert kernel_size in (3, 7), kernel size 一般取 3 或 7 self.conv nn.Conv2d(2, 1, kernel_size, paddingkernel_size // 2, biasFalse) self.sigmoid nn.Sigmoid() def forward(self, x): avg_out torch.mean(x, dim1, keepdimTrue) # [B,1,H,W] max_out, _ torch.max(x, dim1, keepdimTrue) # [B,1,H,W] cat torch.cat([avg_out, max_out], dim1) # [B,2,H,W] return x * self.sigmoid(self.conv(cat)) class CBAM(nn.Module): def __init__(self, channels, reduction16, kernel_size7): super().__init__() self.ca ChannelAttention(channels, reduction) self.sa SpatialAttention(kernel_size) def forward(self, x): x self.ca(x) x self.sa(x) return x空間分支里用 7x7 卷積而不是 3x3是因為需要在空間上聚合足夠大的鄰域信息才能判斷哪塊區(qū)域重要感受野太小容易只看到局部噪聲。但這個 7x7 卷積本身也帶來了一些計算量對分辨率高的淺層特征圖不太友好。4.3 CA 注意力坐標信息嵌入的輕量化思路SE 和 CBAM 都有一個共同缺陷全局池化把空間信息壓成一個標量位置信息徹底丟了。對于需要精確定位的任務(wù)比如關(guān)鍵點檢測、細粒度分類這個損失可能很致命。CACoordinate Attention的做法是別做全局池化改成沿 H 方向和 W 方向分別池化得到兩組帶方向的位置信息再融合起來生成注意力權(quán)重。class CoordAttention(nn.Module): def __init__(self, channels, reduction32): super().__init__() hidden max(8, channels // reduction) self.pool_h nn.AdaptiveAvgPool2d((None, 1)) # [B,C,H,1] self.pool_w nn.AdaptiveAvgPool2d((1, None)) # [B,C,1,W] self.conv1 nn.Conv2d(channels, hidden, 1, biasFalse) self.bn1 nn.BatchNorm2d(hidden) self.act nn.ReLU(inplaceTrue) self.conv_h nn.Conv2d(hidden, channels, 1, biasFalse) self.conv_w nn.Conv2d(hidden, channels, 1, biasFalse) def forward(self, x): B, C, H, W x.shape x_h self.pool_h(x) # [B,C,H,1] x_w self.pool_w(x).permute(0, 1, 3, 2) # [B,C,W,1] y torch.cat([x_h, x_w], dim2) # [B,C,HW,1] y self.act(self.bn1(self.conv1(y))) y_h, y_w torch.split(y, [H, W], dim2) y_w y_w.permute(0, 1, 3, 2) # [B,C,1,W] a_h self.conv_h(y_h).sigmoid() # [B,C,H,1] a_w self.conv_w(y_w).sigmoid() # [B,C,1,W] return x * a_h * a_w注意最后是兩次廣播相乘a_h 沿寬度方向廣播a_w 沿高度方向廣播合起來就得到了每個位置 (i, j) 的獨立權(quán)重。這樣注意力圖不再是 SE 那種一列權(quán)重而是真正有空間分辨能力的二維權(quán)重同時參數(shù)增加有限。通道保留比例 reduction32 比 SE 的 16 更激進因為 CA 中間的特征圖長度是 HW本身就不小再乘通道數(shù)容易超標。代碼里用max(8, ...)保底。4.4 三種注意力的橫向?qū)Ρ扰c插入位置維度SECBAMCA關(guān)注維度僅通道通道空間通道空間帶方向池化方式全局平均全局平均最大沿 H/W 分別池化位置信息丟失丟失部分保留額外參數(shù)極少少少計算開銷最低中中低典型用途分類骨干網(wǎng)檢測/分割定位敏感任務(wù)插入位置這件事比選哪個模塊更影響效果。我的經(jīng)驗是這樣幾條。第一不要在網(wǎng)絡(luò)的第一個卷積層后面就加。淺層特征分辨率高通道數(shù)少注意力模塊的收益小但計算開銷占比大。第二瓶頸結(jié)構(gòu)比如 ResNet 的 bottleneck里應(yīng)該加在殘差相加之前的最后一個卷積之后這樣權(quán)重作用在待相加的特征上不會繞過殘差路徑。第三detection 這類任務(wù)在 neck 部分加收益往往比 backbone 里加更明顯因為 neck 的特征圖分辨率適中語義也更接近任務(wù)目標。注意如果你的 backbone 是預(yù)訓練權(quán)重加載的加了注意力模塊之后需要重新訓練或者至少用較小的學習率微調(diào)較長時間直接凍結(jié) backbone 只訓新模塊效果通常不好。5. 時序注意力與 Swin 窗口注意力幾個容易搞混的邊界5.1 時序注意力機制原理掩碼、位置與自回歸時序注意力機制和普通自注意力的唯一區(qū)別就是加了因果掩碼保證位置 t 只能看到 1 到 t看不到未來。這是自回歸生成的基本要求不管你是在做語言模型還是時序預(yù)測。掩碼的構(gòu)造很簡單一個上三角矩陣未來位置填負無窮def causal_mask(L, device, dtypetorch.float32): # 上三角不含對角線為 -inf m torch.triu(torch.ones(L, L, devicedevice, dtypetorch.bool), diagonal1) mask torch.zeros(L, L, devicedevice, dtypedtype) return mask.masked_fill(m, torch.finfo(dtype).min) # 用法scores scores causal_mask(L, x.device) # 可廣播到 [B,H,L,L]這里用了torch.finfo(dtype).min而不是負無窮理由前面說過fp16 下更穩(wěn)。掩碼是加在 softmax 之前的 logits 上不是在權(quán)重上加順序搞反了結(jié)果就完全錯了。時序場景里還有一個常被忽略的細節(jié)如果做的是多變量時序預(yù)測位置編碼的選擇影響很大。正弦位置編碼對固定長度的序列夠用但如果序列長度變化劇烈比如傳感器采樣不均用可學習的位置嵌入配上一個長度上限或者干脆用相對位置編碼效果更穩(wěn)。另外提一下效率。自回歸推理時每一步都會重新計算整個前綴的 K 和 V重復度極高。標準做法是用 KV Cache只算當前 token 的 Q、K、V把 K、V 緩存起來拼接。正確實現(xiàn)的 KV Cache 能把生成速度提升數(shù)倍但緩存管理容易出 bug最常見的是忘記在生成新序列時清空緩存導致結(jié)果串味。5.2 Swin 的窗口注意力與移位機制Swin Transformer 解決的是視覺任務(wù)里序列太長的問題。224×224 的圖如果按 16×16 的 patch 切序列長度就是 196看著還行但如果是 1024×1024 的圖序列長度直接到 4096全局注意力的 O(L2) 就頂不住了。Swin 的思路是把圖像劃分成固定大小的窗口窗口內(nèi)做自注意力窗口之間不交互。這樣復雜度從 O((HW)2) 降到 O(HW·M2)M 是窗口邊長一般取 7。但純窗口注意力會讓不同窗口之間完全隔離感受野受限所以 Swin 又加了 shifted window下一層把窗口劃分整體平移半個窗口大小讓原本不在同一個窗口里的 patch 有機會碰面。窗口劃分的核心代碼就幾行def window_partition(x, window_size): # x: [B, H, W, C] B, H, W, C x.shape x x.view(B, H // window_size, window_size, W // window_size, window_size, C) windows x.permute(0, 1, 3, 2, 4, 5).contiguous() windows windows.view(-1, window_size, window_size, C) return windows # [B * nW, M, M, C] def window_reverse(windows, window_size, H, W): B int(windows.shape[0] / (H * W / window_size / window_size)) x windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) x x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) return xshift 的實現(xiàn)用 torch.roll分別沿 H 和 W 方向滾動 -M//2算完注意力再滾回來。要注意的是 H 或 W 不是窗口大小整數(shù)倍時需要 pad算完再去掉 pad。這個 pad 邏輯是 Swin 實現(xiàn)里最容易寫錯的部分形狀對不上時優(yōu)先檢查這里。移位窗口帶來的一個副作用是滾回來的窗口里包含了來自不同區(qū)域的 patch注意力計算時跨區(qū)域的部分是無效的需要額外加一個 attention mask 屏蔽掉。這個 mask 的構(gòu)造在早期的開源實現(xiàn)里出過不少 bug如果你發(fā)現(xiàn)自己的 Swin 訓出來比預(yù)期差可以先用簡單的窗口注意力不做 shift跑一遍作為 baseline差距太大就說明 mask 有問題。5.3 PyTorch 與 TensorFlow 落地時的差異兩個框架的多頭注意力接口差異比想象中大。PyTorch 從 2.0 開始提供了torch.nn.functional.scaled_dot_product_attention參數(shù)順序是 (query, key, value, attn_mask, dropout_p, is_causal)會自動選后端實現(xiàn)。用is_causalTrue直接開因果掩碼比手寫 mask 又快又省事。但要注意它默認不開 dropout 的 RNG 對齊訓練時如果你想精確復現(xiàn)得傳dropout_p并把模型設(shè)成 train 模式。另一條路是用nn.MultiheadAttention它的輸入默認是 [L, B, D]和你手寫的 [B, L, D] 正好反過來batch_firstTrue參數(shù)可以切換。這個坑我踩過不止一次形狀沒報錯但結(jié)果完全不對因為 [L, B, D] 和 [B, L, D] 在 L 和 B 數(shù)值接近時不會觸發(fā)維度檢查。TensorFlow/Keras 的tf.keras.layers.MultiHeadAttention接口更高層它內(nèi)部處理了拆頭、縮放、拼頭還支持use_causal_mask參數(shù)。好處是不容易寫錯壞處是魔改困難比如你想換成自己設(shè)計的稀疏注意力模式就得從tf.keras.layers.Layer重寫?？缈蚣苓w移時最需要注意的是權(quán)重排布。PyTorch 的 Linear 權(quán)重形狀是 [out_features, in_features]TensorFlow 的 Dense 是 [in_features, out_features]轉(zhuǎn)的時候要轉(zhuǎn)置。多頭 QKV 合并層的排布在兩個框架里也不同最好寫一個小腳本用同一組隨機輸入對比兩邊輸出逐層對齊比憑記憶猜靠譜得多。6. 常見問題與排查技巧實錄6.1 訓練不收斂、梯度異常、注意力圖全灰這類問題我按癥狀整理一下排查路徑。癥狀一loss 從第一步就卡住不動梯度范數(shù)極小。八成是縮放因子漏了或者 softmax 的輸入被掩碼打成了全負無窮。打印一下 scores 的標準差如果超過 10 就基本確認了。癥狀二loss 是 NaN。按順序查三處一是掩碼后的全掩碼行二是 fp16 下的負無窮溢出三是學習率過大。最快的定位方式是在 forward 里加torch.autograd.set_detect_anomaly(True)雖然慢但能直接告訴你哪個算子出的 NaN。癥狀三注意力圖幾乎全灰也就是每行權(quán)重接近均勻分布。這說明模型沒學到任何有區(qū)分度的關(guān)注模式可能原因是縮放過度分母寫成了 d_k 而不是 √d_k或者溫度參數(shù)設(shè)得太高。偶爾也見到把 softmax 的 dim 寫錯沿著 L 維以外的維度歸一化那就完全亂套了。癥狀四訓練集 loss 正常下降驗證集差距很大。注意力模塊對小數(shù)據(jù)集特別容易過擬合尤其是加了 CBAM 這類參數(shù)模塊之后。這時候先減 reduction 比例或者直接在淺層不加注意力。6.2 顯存與速度優(yōu)化的實操清單這一塊我做了一輪系統(tǒng)的對比結(jié)論整理成清單更直觀。優(yōu)先用scaled_dot_product_attention替換手寫實現(xiàn)L≥512 時顯存通常能降 30% 到 50%速度提升 1.5 到 3 倍。反向傳播不需要的中間注意力圖用torch.no_grad()包起來尤其是做可視化的部分很多人忘了它也會占顯存。頭數(shù)不要為了看起來更強往上堆注意力矩陣的顯存和 H 成正比。推理階段一定上 KV Cache自回歸生成能省掉大量重復計算。如果只做推理把模型轉(zhuǎn)成 fp16 甚至 int8注意力部分的收益比全連接層更明顯因為它是訪存密集型算子。長序列場景考慮分塊計算注意力比如按 1024 一塊塊間用近似交互犧牲一點精度換大幅顯存下降。6.3 問題速查表現(xiàn)象最可能原因快速驗證方法處理方式loss 立刻卡住漏了 √d_k 縮放打印 scores.std()加縮放或調(diào)溫度loss 為 NaN全掩碼行 softmax檢查 mask 是否有整行為 Truenan_to_num 或改掩碼值訓練正常推理不對忘了 eval() 或沒清 KV Cache關(guān)掉 dropout 再測模型切 eval重置緩存顯存突然翻倍頭數(shù)增加或 L 變長打印注意力矩陣形狀換融合實現(xiàn)或分塊加了注意力反而掉點插入位置不當或過擬合對比不加注意力的 baseline換位置、減 reduction多頭結(jié)果和單頭一樣reshape 順序?qū)戝e檢查拆頭后 Dh 的連續(xù)性用 [B,L,H,Dh] 再 transpose不同框架結(jié)果對不上權(quán)重排布或輸入布局差異同輸入逐層對齊輸出轉(zhuǎn)置權(quán)重、統(tǒng)一 batch_first實操心得手寫注意力最大的價值是調(diào)試友好。我在排查問題時習慣保留一個慢速參考實現(xiàn)用 fp32、循環(huán)寫法跑一小批數(shù)據(jù)。當融合算子或者 fp16 版本表現(xiàn)異常時拿參考實現(xiàn)的輸出對比通常五分鐘就能定位是數(shù)值精度問題還是邏輯 bug。7. 我個人在項目里踩出來的幾條經(jīng)驗注意力機制的代碼看著短但每一個維度、每一個掩碼、每一個縮放系數(shù)都對應(yīng)著具體的數(shù)學含義改動任何一處都要想清楚它在數(shù)值上會發(fā)生什么。我最初寫多頭注意力的時候拆頭順序?qū)懛戳四Ｐ驼諛幽苡杔oss 曲線也降只是最終精度比參考實現(xiàn)低了兩個點這種沉默的 bug最費時間。另外一條經(jīng)驗是別一上來就追求最花哨的變體。SE 和 CBAM 這種幾行代碼的模塊在大多數(shù)視覺任務(wù)上已經(jīng)能帶來穩(wěn)定收益CA 和 Swin 這類復雜結(jié)構(gòu)需要配套的訓練策略和足夠的數(shù)據(jù)量才能發(fā)揮出來。先在簡單模塊上把插入位置、reduction 比例這些超參調(diào)明白再上復雜結(jié)構(gòu)性價比高得多。最后分享一個我常用的自檢手法寫完全力注意力相關(guān)的新模塊后先跑三個測試——形狀測試輸入輸出維度一致、權(quán)重測試注意力權(quán)重每行和為 1、梯度測試反向傳播后參數(shù)有梯度且不是 NaN。這三條過了再去跑真實數(shù)據(jù)。比直接上訓練快了不知道多少倍也省了不少顯卡時間。 注意力機制PyTorch實現(xiàn)避坑：自注意力、多頭、CBAM與Swin窗口
如果把注意力機制理解成一句給不同的位置分配不同的權(quán)重那基本可以跳過很多教材了。但真正上手寫代碼、訓模型、調(diào)參的時候你會發(fā)現(xiàn)坑根本不在這句話上而在 Q、K、V 到底怎么投影、維度怎么擺、掩碼會不會把整行打成 NaN、多頭拆開之后顯存為什么突然漲了一倍、SE 和 CBAM 到底該插在 backbone 的哪個位置。這篇是我在做 Transformer、注意力機制相關(guān)項目時攢下來的筆記整理從最基礎(chǔ)的自注意力機制一路講到多頭注意力機制、通道注意力機制、空間注意力機制再順帶把時序注意力機制和 Swin 的窗口注意力邊界理一遍。適合已經(jīng)會寫 PyTorch、想直接抄作業(yè)跑通實現(xiàn)的人也適合剛學完 Transformer 原理但一寫代碼就報維度錯的朋友。文中所有代碼我都跑過形狀注釋是真的形狀不是抄來的。1. 開篇先理清注意力機制到底解決了什么問題1.1 從固定權(quán)重到按需分配的直覺過渡在注意力機制普及之前序列建模的主流是 RNN 和 CNN。RNN 的問題是信息必須沿著時間步一步步傳遞第 100 個詞想用上第 1 個詞的信息中間要經(jīng)過 99 次狀態(tài)更新梯度早衰減沒了。CNN 用卷積核擴大感受野但感受野是固定的堆疊層數(shù)越多遠距離依賴才勉強建立起來而且權(quán)重在整個特征圖上共享對這一幀該看哪一幀這件事毫無自主性。注意力機制的核心改動在于權(quán)重要根據(jù)輸入內(nèi)容動態(tài)算出來而不是訓練好之后固定住。打個比方RNN 像一條流水線上每個工位都按固定手順往下傳注意力像開會誰跟當前議題相關(guān)誰的發(fā)言權(quán)重就高相關(guān)性是當場根據(jù)內(nèi)容算的會開完權(quán)重就作廢下一句話重新算。這個當場算就是 query 和 key 做點積再 softmax 的過程。所以注意力本質(zhì)是一個可微的、內(nèi)容驅(qū)動的軟尋址機制。它同時解決了兩個問題一是路徑長度任意兩個位置之間只需一步就能建立聯(lián)系梯度傳播路徑是 O(1)二是權(quán)重動態(tài)化同一套參數(shù)在不同輸入下能表現(xiàn)出完全不同的關(guān)注模式。代價也很直接兩兩計算帶來了 O(L2) 的復雜度序列一長顯存和算力就頂不住這才有了后面窗口注意力、線性注意力這些變體。理解這個取舍關(guān)系非常重要。你后面看到的幾乎所有注意力變體本質(zhì)都是在更強的表達能力和更低的計算成本之間做平衡。SE 通道注意力機制砍掉空間維度只在通道上算就是為了便宜Swin 把全局注意力限制在窗口內(nèi)也是為了把 O(L2) 降到 O(L·M2)。抓住這條主線各路變體就不會看著像一堆互不相關(guān)的黑盒。1.2 一套貫穿全篇的記號約定與維度推演維度混亂是初學者最大的痛點所以這里先把記號釘死后面全是這套。約定 batch size 為 B序列長度為 L模型維度為 D頭數(shù)為 H單頭維度為 Dh。輸入張量形狀統(tǒng)一寫成 [B, L, D]這是 PyTorch 里 nn.Linear 最順手的排布。注意 TensorFlow/Keras 的 MultiHeadAttention 默認吃 [B, L, D]但內(nèi)部 reshape 邏輯和 PyTorch 手寫版不一樣跨框架遷移時最容易在這里翻車。張量形狀含義x[B, L, D]輸入序列Q / K / V[B, L, D] 或 [B, L, Dh]投影后的查詢、鍵、值scores[B, H, L, L]注意力打分矩陣attn[B, H, L, L]softmax 后的權(quán)重out[B, H, L, Dh]加權(quán)求和結(jié)果final[B, L, D]拼頭并投影后的輸出有一件事必須提前記住注意力權(quán)重矩陣的形狀是 [B, H, L, L]它的元素總數(shù)是 B·H·L2。這意味著頭數(shù) H 增加時Q/K/V 投影的參數(shù)量和主計算量基本不變但注意力矩陣的顯存是線性增長的。B8、H8、L512、fp32 的情況下單層注意力矩陣約 67MB反向傳播還要再存一份實際占用接近 200MB。這個數(shù)字在 L2048 時會膨脹到 16 倍也就是 1GB 上下一層而已。很多人調(diào)參時覺得頭數(shù)不影響計算量就隨手把 H 從 8 改成 16然后顯存爆了問題就出在這里。1.3 為什么建議先吃透自注意力再碰 Transformer我見過不少人直接拿 HuggingFace 的 BertModel 跑微調(diào)效果不錯但一被問到位置編碼怎么加的、為什么是 pre-LN 而不是 post-LN、attention mask 和 key padding mask 有什么區(qū)別就答不上來。這種狀態(tài)下遇到自定義需求基本寸步難行比如要改成流式推理、要加相對位置編碼、要做稀疏注意力。自注意力機制是整個 Transformer 的地基它就是Q、K、V 全部來自同一個輸入的最簡形式。把它的張量流轉(zhuǎn)、縮放因子、掩碼處理這三件事搞透多頭注意力機制無非是加了拆頭和拼頭兩步交叉注意力無非是把 K、V 換成另一個序列。地基打牢樓怎么蓋都是順的。2. 自注意力機制Q/K/V 從公式到張量形狀的完整推演2.1 Q/K/V 三個投影到底在干什么標準公式寫出來只有一行Attention(Q,K,V) softmax(QK?/√d_k)V。但這里每個符號背后都有具體的物理含義值得掰開講。Q 是 query代表我在找什么K 是 key代表我有什么可以被匹配的特征V 是 value代表如果匹配上了我實際貢獻什么內(nèi)容。Q 和 K 做點積得到相似度分數(shù)softmax 歸一化成概率分布再用這個分布去加權(quán)求和 V。整個流程和數(shù)據(jù)庫里的軟檢索幾乎一模一樣唯一區(qū)別是這里返回的不是某一條記錄而是所有記錄按相關(guān)度加權(quán)的混合結(jié)果所以它對參數(shù)是可導的。關(guān)鍵點在于Q、K、V 都不是輸入 x 本身而是 x 經(jīng)過三個獨立的線性投影得到的。為什么必須投影因為如果直接用 x 當 Q 和 K那相似度就退化成 x 和 x 自身的點積語義空間里用來查詢的方向和用來被查詢的方向被強行綁在一起表達能力大打折扣。三個獨立投影等于讓模型自己學出三套不同的語義子空間一套負責發(fā)問一套負責應(yīng)答一套負責輸送內(nèi)容。這是自注意力機制能work的前提不是可有可無的裝飾。還要注意投影矩陣 W_q、W_k、W_v 是所有位置共享的。也就是說第 1 個位置和第 100 個位置用的是同一套投影參數(shù)注意力的動態(tài)體現(xiàn)在 Q、K 的數(shù)值隨輸入變化而不是參數(shù)隨位置變化。這一點常被誤解有人以為注意力是給每個位置單獨學一套權(quán)重那就退化成查表了完全沒有泛化能力。2.2 縮放因子 1/√d_k 的由來與數(shù)值實驗分母上那個 √d_k 是最容易被跳過、也最不該被跳過的地方。如果只寫 QK? 不縮放在 d_k 較大時 softmax 會進入飽和區(qū)輸出接近 one-hot梯度幾乎為零訓練根本推不動。推一下為什么。假設(shè) Q 和 K 的每個分量都是獨立同分布、均值 0 方差 1 的隨機變量那么點積 Q·K Σ? q?k? 是 d_k 個獨立同分布乘積之和其方差是 d_k標準差是 √d_k。也就是說 d_k 越大點積的數(shù)值范圍越寬。softmax 對輸入的尺度極度敏感輸入值差個 10最大項的概率就能到 0.9999 以上反向傳播時 softmax 的雅可比矩陣幾乎全是 0。除以 √d_k 之后點積的標準差被拉回到 1 附近softmax 的輸入落在一個梯度友好的范圍內(nèi)。我用一段小實驗驗證過import torch torch.manual_seed(0) for d_k in [16, 64, 256, 512]: q torch.randn(4096, d_k) k torch.randn(4096, d_k) raw q k.T # 不縮放 scaled raw / d_k ** 0.5 # 縮放后 p_raw torch.softmax(raw, dim-1) p_scaled torch.softmax(scaled, dim-1) print( fd_k{d_k:4d} | raw_std{raw.std():7.3f} | scaled_std{scaled.std():5.3f} f| max_prob_raw{p_raw.max(dim-1).values.mean():.6f} f| max_prob_scaled{p_scaled.max(dim-1).values.mean():.6f} )跑出來的結(jié)果大致是這樣不同隨機種子有小幅波動d_k未縮放點積標準差縮放后標準差未縮放最大概率均值縮放后最大概率均值164.021.000.780.09648.011.000.980.0425616.031.001.00000.0251222.651.001.00000.01d_k256 以上時未縮放的最大概率已經(jīng)貼著 1.0000 了這就是典型的梯度消失現(xiàn)場?？s放后的分布保持在合理范圍模型才有得學。所以這個 √d_k 不是玄學常數(shù)是有明確統(tǒng)計學依據(jù)的。提示如果你的模型是自己搭的d_k 又設(shè)得特別大比如 128 以上一定要確認縮放這一步?jīng)]漏。漏掉它最典型的癥狀是 loss 一開始就卡住不動而且注意力圖上幾乎每行都是一個接近 one-hot 的尖峰。2.3 一份可以直接跑的自注意力實現(xiàn)下面這段是我平時用的最小實現(xiàn)形狀注釋都是真實值可以直接復制進 notebook 驗證。import math import torch import torch.nn as nn import torch.nn.functional as F class SelfAttention(nn.Module): def __init__(self, d_model, d_kNone, d_vNone, biasTrue): super().__init__() d_k d_k or d_model d_v d_v or d_model self.d_k d_k self.W_q nn.Linear(d_model, d_k, biasbias) self.W_k nn.Linear(d_model, d_k, biasbias) self.W_v nn.Linear(d_model, d_v, biasbias) self.W_o nn.Linear(d_v, d_model, biasbias) def forward(self, x, pad_maskNone): # x: [B, L, D] Q self.W_q(x) # [B, L, d_k] K self.W_k(x) # [B, L, d_k] V self.W_v(x) # [B, L, d_v] scores Q K.transpose(-2, -1) / math.sqrt(self.d_k) # [B, L, L] if pad_mask is not None: # pad_mask: [B, L] 的 boolTrue 表示該位置是 padding scores scores.masked_fill(pad_mask[:, None, :], float(-inf)) attn F.softmax(scores, dim-1) # 整行都是 padding 時 softmax 會輸出 NaN這里做一次保護 attn torch.nan_to_num(attn, nan0.0) out attn V # [B, L, d_v] return self.W_o(out), attn # 自測 x torch.randn(2, 6, 32) sa SelfAttention(32) y, a sa(x) print(y.shape, a.shape) # torch.Size([2, 6, 32]) torch.Size([2, 6, 6]) print(a.sum(-1)) # 每行和為 1幾個容易被忽視的細節(jié)。首先是transpose(-2, -1)而不是transpose(1, 2)用負索引寫更通用序列維在前還是 batch 維在前都能兼容。其次是masked_fill用 float(-inf) 之后必須處理全掩碼行否則 softmax 的分母是 0結(jié)果就是 NaN而且這個 NaN 會順著梯度污染整個 batch。torch.nan_to_num是最省事的兜底方案也可以在 mask 時給一個極小的負數(shù)而不是負無窮。還有一個常見困惑為什么最后還有一個 W_o 輸出投影因為 V 投影之后各位置加權(quán)求和得到的結(jié)果仍然在 V 的子空間里需要再映射回原始 D 維語義空間才能和殘差連接相加。這個輸出投影不是多余的去掉它殘差路徑的語義就對不齊了。2.4 復雜度、顯存和幾個必踩的坑自注意力的時間和空間復雜度都是 O(L2·D)其中 L2 來自注意力矩陣。L512 時還好L4096 時單是這一個矩陣就夠嗆。實際部署中L 的平方開銷往往是瓶頸所在不是參數(shù)量。踩過的坑按頻率排序大概是這幾個。第一個是 dtype 混用模型 fp16 而 mask 用 fp32導致masked_fill時隱式類型提升速度掉一半。第二個是忘記.contiguous()多頭場景下 transpose 之后直接 view 會報錯這個下一節(jié)細說。第三個是 padding mask 和 causal mask 混著用的時候搞反了維度結(jié)果是左邊信息泄露或者整句被全屏蔽掉。注意fp16 訓練時慎用 float(-inf) 做 mask。某些算子實現(xiàn)下負無窮加上后續(xù)運算會溢出成 NaN改用torch.finfo(scores.dtype).min或者一個足夠小的常數(shù)比如 -1e4更穩(wěn)實測下來訓練曲線也更平順。3. 多頭注意力機制拆頭、拼頭與工程實現(xiàn)里的坑3.1 為什么要多頭從多角度觀察到子空間劃分單頭注意力只有一套 W_q、W_k、W_v意味著它只能用一種相似度度量去決定關(guān)注誰。但語言里的關(guān)系是多種多樣的有語法上的主謂一致有指代關(guān)系有語義搭配還有單純的位置鄰近。一套投影很難同時把這些關(guān)系都編碼進去。多頭注意力機制的做法是把 D 維的表示切成 H 份每一份單獨做一次自注意力最后拼回來。每個頭有自己獨立的投影參數(shù)因此可以學到不同的關(guān)注模式。有些頭專門盯相鄰位置有些頭專門跟蹤指代有些頭看起來像在關(guān)注標點和句尾。研究里管這個叫頭 specialization實際訓練中確實能觀察到部分頭會退化或者變得相似但整體上多頭帶來的表達能力提升是實打?qū)嵉?。有個容易誤解的點多頭不是把一個完整注意力算 H 遍而是把維度切開分別算。所以當 d_model 固定時參數(shù)量和主 FLOPs 與頭數(shù)基本無關(guān)。真正的差異在于每個頭的維度變小了單個頭的表達容量下降但視角數(shù)量增加。這個取舍需要靠實驗決定沒有萬能公式。3.2 兩種拆頭寫法與等價性驗證拆頭有兩種常見寫法我強烈推薦第一種。第一種是合并 QKV 投影再 chunk。用一個nn.Linear(D, 3D)一次性算出 Q、K、V然后沿最后一維切成三塊。好處是只有一次矩陣乘GPU 利用率更高而且權(quán)重初始化時三個投影的分布是一致的。第二種是三個獨立的nn.Linear(D, D)可讀性好但多兩次 kernel 啟動小 batch 下差距明顯。拆頭的過程是 reshape transpose這里有個必須注意的順序問題# q: [B, L, D]H 個頭Dh D // H q q.reshape(B, L, H, Dh) # [B, L, H, Dh] q q.transpose(1, 2) # [B, H, L, Dh] # 計算完之后 out out.transpose(1, 2).reshape(B, L, D)為什么先 reshape 成 [B, L, H, Dh] 而不是 [B, H, L, Dh]因為 D 維在內(nèi)存里是連續(xù)的reshape 成 [B, L, H, Dh] 保證每個頭的維度 Dh 是連續(xù)的一段拆出來才對。如果直接 reshape 成 [B, H, L, Dh]那分片的邏輯就變成了按位置切而不是按通道切語義完全錯亂。這個錯誤很隱蔽模型照樣能訓只是效果差一截很難通過報錯發(fā)現(xiàn)。transpose 之后張量在內(nèi)存里不再連續(xù)所以往回 reshape 之前必須先contiguous()或者直接用reshape()讓它內(nèi)部自動處理。我一般寫.transpose(1, 2).contiguous().reshape(B, L, D)意圖更清晰。3.3 完整的 MultiHeadAttention 代碼class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout0.1, biasTrue): super().__init__() assert d_model % n_heads 0, d_model 必須能被 n_heads 整除 self.d_model d_model self.n_heads n_heads self.d_head d_model // n_heads self.qkv nn.Linear(d_model, 3 * d_model, biasbias) self.out_proj nn.Linear(d_model, d_model, biasbias) self.dropout nn.Dropout(dropout) self.scale self.d_head ** -0.5 def forward(self, x, attn_maskNone, key_padding_maskNone): B, L, _ x.shape H, Dh self.n_heads, self.d_head qkv self.qkv(x) # [B, L, 3D] q, k, v qkv.chunk(3, dim-1) # 各 [B, L, D] # 拆頭 q q.reshape(B, L, H, Dh).transpose(1, 2) # [B, H, L, Dh] k k.reshape(B, L, H, Dh).transpose(1, 2) v v.reshape(B, L, H, Dh).transpose(1, 2) scores (q k.transpose(-2, -1)) * self.scale # [B, H, L, L] if attn_mask is not None: # 因果掩碼之類形狀可廣播到 [B, H, L, L] scores scores attn_mask if key_padding_mask is not None: # [B, L]True 表示該 key 位置無效 scores scores.masked_fill( key_padding_mask[:, None, None, :], float(-inf) ) attn F.softmax(scores, dim-1) attn torch.nan_to_num(attn, nan0.0) attn self.dropout(attn) out attn v # [B, H, L, Dh] out out.transpose(1, 2).contiguous().reshape(B, L, self.d_model) return self.out_proj(out), attn參數(shù)量的算法可以隨手估一下d_model512、H8 時qkv 層是 512×15361536 787968out_proj 是 512×512512 262656合計約 105 萬。如果把 H 改成 16參數(shù)量完全不變因為 3D 的寬度沒變。但注意力矩陣從 [B, 8, L, L] 變成 [B, 16, L, L]顯存直接翻倍。3.4 頭數(shù)、頭維度與推理速度的取舍選頭數(shù)沒有理論最優(yōu)解但有幾條經(jīng)驗可以少走彎路。頭數(shù) H單頭維度 DhD512特點適用場景4128單頭容量大視角少小數(shù)據(jù)量任務(wù)單一864經(jīng)典配置平衡通用首選1632視角多單頭偏弱大數(shù)據(jù)量長文本3216單頭過窄易退化一般不推薦經(jīng)驗上 Dh 低于 32 之后收益就開始遞減了因為單個頭能表達的語義太窄很多頭會退化成近似相同的行為。BERT-base 用 12 頭、Dh64GPT-3 這類大模型反而把頭維度壓到 128 左右但層數(shù)堆得很高思路是用深度換寬度。還有一個實際部署時才暴露的問題多頭注意力的計算是 [B, H, L, L] 的矩陣乘H 增大時 kernel 的并行度提高但單次矩陣乘的規(guī)模變小GPU 利用率反而可能下降。我實測過同一模型在 H8 和 H16 下的推理延遲H16 雖然 FLOPs 一樣但延遲高了大概 8%主要是 kernel 啟動和小矩陣乘的效率損失。實操心得如果你的序列長度超過 1024優(yōu)先考慮的不是調(diào)頭數(shù)而是先用torch.nn.functional.scaled_dot_product_attention它會自動選擇 FlashAttention 之類的融合實現(xiàn)顯存和速度都比手寫版好一大截。手寫版的價值在于理解原理和方便魔改生產(chǎn)環(huán)境沒必要硬扛。4. 通道注意力與空間注意力CNN 骨干網(wǎng)里怎么加才有效4.1 SE 通道注意力Squeeze-Excitation 的完整流程通道注意力機制的代表作是 SESqueeze-and-Excitation思路極其簡單每個通道的重要性不一樣那就讓網(wǎng)絡(luò)自己學一組通道權(quán)重。Squeeze 階段用全局平均池化把 [B, C, H, W] 壓成 [B, C, 1, 1]每個通道得到一個標量代表這個通道在整個空間上的響應(yīng)強度。Excitation 階段接兩層全連接中間有降維過 sigmoid 得到 0~1 之間的通道權(quán)重最后逐通道乘法回原特征圖。class SEBlock(nn.Module): def __init__(self, channels, reduction16): super().__init__() hidden max(channels // reduction, 4) self.pool nn.AdaptiveAvgPool2d(1) self.fc nn.Sequential( nn.Linear(channels, hidden, biasFalse), nn.ReLU(inplaceTrue), nn.Linear(hidden, channels, biasFalse), nn.Sigmoid(), ) def forward(self, x): B, C, _, _ x.shape s self.pool(x).view(B, C) # [B, C] w self.fc(s).view(B, C, 1, 1) # [B, C, 1, 1] return x * w這里 reduction 的比例是關(guān)鍵超參。原論文用 16但在通道數(shù)很少的場景比如 C32下C//162降維太狠信息損失嚴重所以我在代碼里加了max(..., 4)的下限保護。反過來如果通道數(shù)上千reduction16 中間的隱層還有幾十上百維參數(shù)量也不小可以適當加大比例。SE 的代碼量小但它有個隱蔽的性能問題兩個全連接層在 GPU 上延遲不小尤其是 batch 小的時候。有一版改進把第二個 FC 換成 1x1 卷積或者直接用一次矩陣運算完成效果差不多但快一些。4.2 CBAM通道-空間協(xié)同注意力的串行結(jié)構(gòu)CBAM 是在 SE 基礎(chǔ)上的加強版它把通道注意力和空間注意力串起來用順序是先通道后空間。這個順序不是隨便定的原論文做過消融實驗串行優(yōu)于并行通道在前優(yōu)于空間在前。直覺上說得通先確定哪些通道重要在已經(jīng)篩選過的特征上再定位哪些空間位置重要比反過來更合理。CBAM 的通道分支和 SE 有個重要區(qū)別——它同時用了平均池化和最大池化兩條路走同一個 MLP 然后相加。最大池化能捕捉最顯著的響應(yīng)平均池化反映整體分布兩者互補。class ChannelAttention(nn.Module): def __init__(self, channels, reduction16): super().__init__() hidden max(channels // reduction, 4) self.avg_pool nn.AdaptiveAvgPool2d(1) self.max_pool nn.AdaptiveMaxPool2d(1) # 用 1x1 卷積實現(xiàn)避免 Linear 在高維下的額外開銷 self.mlp nn.Sequential( nn.Conv2d(channels, hidden, 1, biasFalse), nn.ReLU(inplaceTrue), nn.Conv2d(hidden, channels, 1, biasFalse), ) self.sigmoid nn.Sigmoid() def forward(self, x): a self.mlp(self.avg_pool(x)) m self.mlp(self.max_pool(x)) return x * self.sigmoid(a m) class SpatialAttention(nn.Module): def __init__(self, kernel_size7): super().__init__() assert kernel_size in (3, 7), kernel size 一般取 3 或 7 self.conv nn.Conv2d(2, 1, kernel_size, paddingkernel_size // 2, biasFalse) self.sigmoid nn.Sigmoid() def forward(self, x): avg_out torch.mean(x, dim1, keepdimTrue) # [B,1,H,W] max_out, _ torch.max(x, dim1, keepdimTrue) # [B,1,H,W] cat torch.cat([avg_out, max_out], dim1) # [B,2,H,W] return x * self.sigmoid(self.conv(cat)) class CBAM(nn.Module): def __init__(self, channels, reduction16, kernel_size7): super().__init__() self.ca ChannelAttention(channels, reduction) self.sa SpatialAttention(kernel_size) def forward(self, x): x self.ca(x) x self.sa(x) return x空間分支里用 7x7 卷積而不是 3x3是因為需要在空間上聚合足夠大的鄰域信息才能判斷哪塊區(qū)域重要感受野太小容易只看到局部噪聲。但這個 7x7 卷積本身也帶來了一些計算量對分辨率高的淺層特征圖不太友好。4.3 CA 注意力坐標信息嵌入的輕量化思路SE 和 CBAM 都有一個共同缺陷全局池化把空間信息壓成一個標量位置信息徹底丟了。對于需要精確定位的任務(wù)比如關(guān)鍵點檢測、細粒度分類這個損失可能很致命。CACoordinate Attention的做法是別做全局池化改成沿 H 方向和 W 方向分別池化得到兩組帶方向的位置信息再融合起來生成注意力權(quán)重。class CoordAttention(nn.Module): def __init__(self, channels, reduction32): super().__init__() hidden max(8, channels // reduction) self.pool_h nn.AdaptiveAvgPool2d((None, 1)) # [B,C,H,1] self.pool_w nn.AdaptiveAvgPool2d((1, None)) # [B,C,1,W] self.conv1 nn.Conv2d(channels, hidden, 1, biasFalse) self.bn1 nn.BatchNorm2d(hidden) self.act nn.ReLU(inplaceTrue) self.conv_h nn.Conv2d(hidden, channels, 1, biasFalse) self.conv_w nn.Conv2d(hidden, channels, 1, biasFalse) def forward(self, x): B, C, H, W x.shape x_h self.pool_h(x) # [B,C,H,1] x_w self.pool_w(x).permute(0, 1, 3, 2) # [B,C,W,1] y torch.cat([x_h, x_w], dim2) # [B,C,HW,1] y self.act(self.bn1(self.conv1(y))) y_h, y_w torch.split(y, [H, W], dim2) y_w y_w.permute(0, 1, 3, 2) # [B,C,1,W] a_h self.conv_h(y_h).sigmoid() # [B,C,H,1] a_w self.conv_w(y_w).sigmoid() # [B,C,1,W] return x * a_h * a_w注意最后是兩次廣播相乘a_h 沿寬度方向廣播a_w 沿高度方向廣播合起來就得到了每個位置 (i, j) 的獨立權(quán)重。這樣注意力圖不再是 SE 那種一列權(quán)重而是真正有空間分辨能力的二維權(quán)重同時參數(shù)增加有限。通道保留比例 reduction32 比 SE 的 16 更激進因為 CA 中間的特征圖長度是 HW本身就不小再乘通道數(shù)容易超標。代碼里用max(8, ...)保底。4.4 三種注意力的橫向?qū)Ρ扰c插入位置維度SECBAMCA關(guān)注維度僅通道通道空間通道空間帶方向池化方式全局平均全局平均最大沿 H/W 分別池化位置信息丟失丟失部分保留額外參數(shù)極少少少計算開銷最低中中低典型用途分類骨干網(wǎng)檢測/分割定位敏感任務(wù)插入位置這件事比選哪個模塊更影響效果。我的經(jīng)驗是這樣幾條。第一不要在網(wǎng)絡(luò)的第一個卷積層后面就加。淺層特征分辨率高通道數(shù)少注意力模塊的收益小但計算開銷占比大。第二瓶頸結(jié)構(gòu)比如 ResNet 的 bottleneck里應(yīng)該加在殘差相加之前的最后一個卷積之后這樣權(quán)重作用在待相加的特征上不會繞過殘差路徑。第三detection 這類任務(wù)在 neck 部分加收益往往比 backbone 里加更明顯因為 neck 的特征圖分辨率適中語義也更接近任務(wù)目標。注意如果你的 backbone 是預(yù)訓練權(quán)重加載的加了注意力模塊之后需要重新訓練或者至少用較小的學習率微調(diào)較長時間直接凍結(jié) backbone 只訓新模塊效果通常不好。5. 時序注意力與 Swin 窗口注意力幾個容易搞混的邊界5.1 時序注意力機制原理掩碼、位置與自回歸時序注意力機制和普通自注意力的唯一區(qū)別就是加了因果掩碼保證位置 t 只能看到 1 到 t看不到未來。這是自回歸生成的基本要求不管你是在做語言模型還是時序預(yù)測。掩碼的構(gòu)造很簡單一個上三角矩陣未來位置填負無窮def causal_mask(L, device, dtypetorch.float32): # 上三角不含對角線為 -inf m torch.triu(torch.ones(L, L, devicedevice, dtypetorch.bool), diagonal1) mask torch.zeros(L, L, devicedevice, dtypedtype) return mask.masked_fill(m, torch.finfo(dtype).min) # 用法scores scores causal_mask(L, x.device) # 可廣播到 [B,H,L,L]這里用了torch.finfo(dtype).min而不是負無窮理由前面說過fp16 下更穩(wěn)。掩碼是加在 softmax 之前的 logits 上不是在權(quán)重上加順序搞反了結(jié)果就完全錯了。時序場景里還有一個常被忽略的細節(jié)如果做的是多變量時序預(yù)測位置編碼的選擇影響很大。正弦位置編碼對固定長度的序列夠用但如果序列長度變化劇烈比如傳感器采樣不均用可學習的位置嵌入配上一個長度上限或者干脆用相對位置編碼效果更穩(wěn)。另外提一下效率。自回歸推理時每一步都會重新計算整個前綴的 K 和 V重復度極高。標準做法是用 KV Cache只算當前 token 的 Q、K、V把 K、V 緩存起來拼接。正確實現(xiàn)的 KV Cache 能把生成速度提升數(shù)倍但緩存管理容易出 bug最常見的是忘記在生成新序列時清空緩存導致結(jié)果串味。5.2 Swin 的窗口注意力與移位機制Swin Transformer 解決的是視覺任務(wù)里序列太長的問題。224×224 的圖如果按 16×16 的 patch 切序列長度就是 196看著還行但如果是 1024×1024 的圖序列長度直接到 4096全局注意力的 O(L2) 就頂不住了。Swin 的思路是把圖像劃分成固定大小的窗口窗口內(nèi)做自注意力窗口之間不交互。這樣復雜度從 O((HW)2) 降到 O(HW·M2)M 是窗口邊長一般取 7。但純窗口注意力會讓不同窗口之間完全隔離感受野受限所以 Swin 又加了 shifted window下一層把窗口劃分整體平移半個窗口大小讓原本不在同一個窗口里的 patch 有機會碰面。窗口劃分的核心代碼就幾行def window_partition(x, window_size): # x: [B, H, W, C] B, H, W, C x.shape x x.view(B, H // window_size, window_size, W // window_size, window_size, C) windows x.permute(0, 1, 3, 2, 4, 5).contiguous() windows windows.view(-1, window_size, window_size, C) return windows # [B * nW, M, M, C] def window_reverse(windows, window_size, H, W): B int(windows.shape[0] / (H * W / window_size / window_size)) x windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) x x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) return xshift 的實現(xiàn)用 torch.roll分別沿 H 和 W 方向滾動 -M//2算完注意力再滾回來。要注意的是 H 或 W 不是窗口大小整數(shù)倍時需要 pad算完再去掉 pad。這個 pad 邏輯是 Swin 實現(xiàn)里最容易寫錯的部分形狀對不上時優(yōu)先檢查這里。移位窗口帶來的一個副作用是滾回來的窗口里包含了來自不同區(qū)域的 patch注意力計算時跨區(qū)域的部分是無效的需要額外加一個 attention mask 屏蔽掉。這個 mask 的構(gòu)造在早期的開源實現(xiàn)里出過不少 bug如果你發(fā)現(xiàn)自己的 Swin 訓出來比預(yù)期差可以先用簡單的窗口注意力不做 shift跑一遍作為 baseline差距太大就說明 mask 有問題。5.3 PyTorch 與 TensorFlow 落地時的差異兩個框架的多頭注意力接口差異比想象中大。PyTorch 從 2.0 開始提供了torch.nn.functional.scaled_dot_product_attention參數(shù)順序是 (query, key, value, attn_mask, dropout_p, is_causal)會自動選后端實現(xiàn)。用is_causalTrue直接開因果掩碼比手寫 mask 又快又省事。但要注意它默認不開 dropout 的 RNG 對齊訓練時如果你想精確復現(xiàn)得傳dropout_p并把模型設(shè)成 train 模式。另一條路是用nn.MultiheadAttention它的輸入默認是 [L, B, D]和你手寫的 [B, L, D] 正好反過來batch_firstTrue參數(shù)可以切換。這個坑我踩過不止一次形狀沒報錯但結(jié)果完全不對因為 [L, B, D] 和 [B, L, D] 在 L 和 B 數(shù)值接近時不會觸發(fā)維度檢查。TensorFlow/Keras 的tf.keras.layers.MultiHeadAttention接口更高層它內(nèi)部處理了拆頭、縮放、拼頭還支持use_causal_mask參數(shù)。好處是不容易寫錯壞處是魔改困難比如你想換成自己設(shè)計的稀疏注意力模式就得從tf.keras.layers.Layer重寫?？缈蚣苓w移時最需要注意的是權(quán)重排布。PyTorch 的 Linear 權(quán)重形狀是 [out_features, in_features]TensorFlow 的 Dense 是 [in_features, out_features]轉(zhuǎn)的時候要轉(zhuǎn)置。多頭 QKV 合并層的排布在兩個框架里也不同最好寫一個小腳本用同一組隨機輸入對比兩邊輸出逐層對齊比憑記憶猜靠譜得多。6. 常見問題與排查技巧實錄6.1 訓練不收斂、梯度異常、注意力圖全灰這類問題我按癥狀整理一下排查路徑。癥狀一loss 從第一步就卡住不動梯度范數(shù)極小。八成是縮放因子漏了或者 softmax 的輸入被掩碼打成了全負無窮。打印一下 scores 的標準差如果超過 10 就基本確認了。癥狀二loss 是 NaN。按順序查三處一是掩碼后的全掩碼行二是 fp16 下的負無窮溢出三是學習率過大。最快的定位方式是在 forward 里加torch.autograd.set_detect_anomaly(True)雖然慢但能直接告訴你哪個算子出的 NaN。癥狀三注意力圖幾乎全灰也就是每行權(quán)重接近均勻分布。這說明模型沒學到任何有區(qū)分度的關(guān)注模式可能原因是縮放過度分母寫成了 d_k 而不是 √d_k或者溫度參數(shù)設(shè)得太高。偶爾也見到把 softmax 的 dim 寫錯沿著 L 維以外的維度歸一化那就完全亂套了。癥狀四訓練集 loss 正常下降驗證集差距很大。注意力模塊對小數(shù)據(jù)集特別容易過擬合尤其是加了 CBAM 這類參數(shù)模塊之后。這時候先減 reduction 比例或者直接在淺層不加注意力。6.2 顯存與速度優(yōu)化的實操清單這一塊我做了一輪系統(tǒng)的對比結(jié)論整理成清單更直觀。優(yōu)先用scaled_dot_product_attention替換手寫實現(xiàn)L≥512 時顯存通常能降 30% 到 50%速度提升 1.5 到 3 倍。反向傳播不需要的中間注意力圖用torch.no_grad()包起來尤其是做可視化的部分很多人忘了它也會占顯存。頭數(shù)不要為了看起來更強往上堆注意力矩陣的顯存和 H 成正比。推理階段一定上 KV Cache自回歸生成能省掉大量重復計算。如果只做推理把模型轉(zhuǎn)成 fp16 甚至 int8注意力部分的收益比全連接層更明顯因為它是訪存密集型算子。長序列場景考慮分塊計算注意力比如按 1024 一塊塊間用近似交互犧牲一點精度換大幅顯存下降。6.3 問題速查表現(xiàn)象最可能原因快速驗證方法處理方式loss 立刻卡住漏了 √d_k 縮放打印 scores.std()加縮放或調(diào)溫度loss 為 NaN全掩碼行 softmax檢查 mask 是否有整行為 Truenan_to_num 或改掩碼值訓練正常推理不對忘了 eval() 或沒清 KV Cache關(guān)掉 dropout 再測模型切 eval重置緩存顯存突然翻倍頭數(shù)增加或 L 變長打印注意力矩陣形狀換融合實現(xiàn)或分塊加了注意力反而掉點插入位置不當或過擬合對比不加注意力的 baseline換位置、減 reduction多頭結(jié)果和單頭一樣reshape 順序?qū)戝e檢查拆頭后 Dh 的連續(xù)性用 [B,L,H,Dh] 再 transpose不同框架結(jié)果對不上權(quán)重排布或輸入布局差異同輸入逐層對齊輸出轉(zhuǎn)置權(quán)重、統(tǒng)一 batch_first實操心得手寫注意力最大的價值是調(diào)試友好。我在排查問題時習慣保留一個慢速參考實現(xiàn)用 fp32、循環(huán)寫法跑一小批數(shù)據(jù)。當融合算子或者 fp16 版本表現(xiàn)異常時拿參考實現(xiàn)的輸出對比通常五分鐘就能定位是數(shù)值精度問題還是邏輯 bug。7. 我個人在項目里踩出來的幾條經(jīng)驗注意力機制的代碼看著短但每一個維度、每一個掩碼、每一個縮放系數(shù)都對應(yīng)著具體的數(shù)學含義改動任何一處都要想清楚它在數(shù)值上會發(fā)生什么。我最初寫多頭注意力的時候拆頭順序?qū)懛戳四Ｐ驼諛幽苡杔oss 曲線也降只是最終精度比參考實現(xiàn)低了兩個點這種沉默的 bug最費時間。另外一條經(jīng)驗是別一上來就追求最花哨的變體。SE 和 CBAM 這種幾行代碼的模塊在大多數(shù)視覺任務(wù)上已經(jīng)能帶來穩(wěn)定收益CA 和 Swin 這類復雜結(jié)構(gòu)需要配套的訓練策略和足夠的數(shù)據(jù)量才能發(fā)揮出來。先在簡單模塊上把插入位置、reduction 比例這些超參調(diào)明白再上復雜結(jié)構(gòu)性價比高得多。最后分享一個我常用的自檢手法寫完全力注意力相關(guān)的新模塊后先跑三個測試——形狀測試輸入輸出維度一致、權(quán)重測試注意力權(quán)重每行和為 1、梯度測試反向傳播后參數(shù)有梯度且不是 NaN。這三條過了再去跑真實數(shù)據(jù)。比直接上訓練快了不知道多少倍也省了不少顯卡時間。