制：自注意力、交叉注意力與因果自注意力)
1. 從“注意力”到“Transformer”一場(chǎng)認(rèn)知范式的革命如果你在2017年之后才開(kāi)始接觸深度學(xué)習(xí)尤其是自然語(yǔ)言處理或計(jì)算機(jī)視覺(jué)那么“注意力機(jī)制”和“Transformer”這兩個(gè)詞對(duì)你來(lái)說(shuō)可能就像空氣一樣自然存在。但回到那個(gè)時(shí)間點(diǎn)之前整個(gè)領(lǐng)域完全是另一番景象。循環(huán)神經(jīng)網(wǎng)絡(luò)RNN及其變體LSTM、GRU是處理序列數(shù)據(jù)的絕對(duì)王者我們花費(fèi)大量精力去設(shè)計(jì)門(mén)控結(jié)構(gòu)、解決梯度消失只為讓模型能“記住”更久遠(yuǎn)的信息。然而一個(gè)根本性的瓶頸始終存在序列計(jì)算的順序性。你必須一個(gè)接一個(gè)地處理輸入這既限制了訓(xùn)練時(shí)的并行效率也使得模型難以直接建立長(zhǎng)距離依賴(lài)關(guān)系尤其是在處理兩個(gè)相距很遠(yuǎn)的單詞或像素時(shí)信息需要經(jīng)過(guò)漫長(zhǎng)的“跋涉”才能關(guān)聯(lián)起來(lái)。Transformer架構(gòu)的橫空出世徹底打破了這一僵局。它最核心、最革命性的貢獻(xiàn)就是完全摒棄了循環(huán)結(jié)構(gòu)轉(zhuǎn)而依賴(lài)一種名為“自注意力”的機(jī)制來(lái)建立序列內(nèi)部所有元素之間的全局關(guān)聯(lián)。你可以把它想象成在一個(gè)會(huì)議室里開(kāi)會(huì)。RNN就像是一個(gè)只能聽(tīng)上一個(gè)人發(fā)言然后自己發(fā)言的串行模式而Transformer的自注意力機(jī)制則像是讓會(huì)議室里的每個(gè)人同時(shí)發(fā)言并且每個(gè)人都能瞬間聽(tīng)到并理解房間里所有人的發(fā)言?xún)?nèi)容然后綜合這些信息形成自己的觀點(diǎn)。這種“全局視野”的能力是之前任何模型都難以企及的。今天我們不再滿足于僅僅知道Transformer“很強(qiáng)大”。作為從業(yè)者我們需要深入其核心拆解它的引擎——注意力機(jī)制。很多人初學(xué)時(shí)會(huì)混淆“注意力機(jī)制”、“自注意力”和“多頭注意力”這些概念或者對(duì)“位置編碼”為何如此關(guān)鍵一知半解。這篇文章的目的就是幫你徹底厘清Transformer架構(gòu)中這三種核心的注意力機(jī)制自注意力、編碼器-解碼器注意力交叉注意力和因果自注意力掩碼自注意力。我們會(huì)從最直觀的比喻出發(fā)深入到數(shù)學(xué)計(jì)算和代碼實(shí)現(xiàn)最后探討它們?nèi)绾螀f(xié)同工作構(gòu)成了當(dāng)今大模型時(shí)代的基石。無(wú)論你是想徹底理解BERT、GPT、ViT這些明星模型的工作原理還是計(jì)劃在自己的項(xiàng)目中引入注意力模塊搞懂這三種機(jī)制都是必經(jīng)之路。2. Transformer架構(gòu)總覽與注意力機(jī)制的核心角色在深入三種具體的注意力機(jī)制之前我們有必要先快速回顧一下Transformer的整體架構(gòu)理解注意力機(jī)制在其中扮演的“發(fā)動(dòng)機(jī)”角色。原始的Transformer論文《Attention Is All You Need》提出的是一個(gè)用于序列到序列Seq2Seq任務(wù)的模型比如機(jī)器翻譯。它由編碼器Encoder和解碼器Decoder兩部分堆疊而成。編碼器由N個(gè)原論文中N6完全相同的層堆疊而成。每一層都包含兩個(gè)核心子層多頭自注意力機(jī)制Multi-Head Self-Attention前饋神經(jīng)網(wǎng)絡(luò)Position-wise Feed-Forward Network每個(gè)子層周?chē)疾捎昧藲埐钸B接Residual Connection和層歸一化Layer Normalization。編碼器的任務(wù)是接收輸入序列例如一句英文并通過(guò)自注意力機(jī)制提取其內(nèi)部豐富的上下文信息將每個(gè)輸入詞元Token編碼成一個(gè)蘊(yùn)含了全局上下文信息的向量表示。解碼器同樣由N個(gè)相同的層堆疊。每一層包含三個(gè)核心子層掩碼多頭自注意力機(jī)制Masked Multi-Head Self-Attention這是“因果自注意力”的具體實(shí)現(xiàn)。多頭編碼器-解碼器注意力機(jī)制Multi-Head Encoder-Decoder Attention這就是“交叉注意力”。前饋神經(jīng)網(wǎng)絡(luò)解碼器的任務(wù)是在已知編碼器輸出和已生成部分目標(biāo)序列的前提下自回歸地預(yù)測(cè)下一個(gè)詞元。注意雖然原始Transformer是編解碼結(jié)構(gòu)但后續(xù)很多著名模型只用了其中一部分。例如BERT只用了編碼器GPT系列只用了解碼器并去掉了其中的交叉注意力子層Vision Transformer (ViT) 則將圖像塊序列視為輸入使用編碼器進(jìn)行處理。因此理解這三種注意力機(jī)制是理解所有Transformer變體的基礎(chǔ)。那么什么是“注意力”的通用計(jì)算模式呢拋開(kāi)具體的變體其核心思想可以概括為給定一個(gè)“查詢(xún)”Query集合和一個(gè)“鍵-值”Key-Value對(duì)集合注意力機(jī)制通過(guò)計(jì)算Query與所有Key的相似度相關(guān)性來(lái)得到每個(gè)Key對(duì)應(yīng)的Value的加權(quán)和從而輸出一個(gè)聚焦于最重要信息的表示。用公式表示就是 scaled dot-product attention縮放點(diǎn)積注意力Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V這里Q、K、V分別是查詢(xún)、鍵、值矩陣。sqrt(d_k)是一個(gè)縮放因子用于防止點(diǎn)積結(jié)果過(guò)大導(dǎo)致softmax梯度消失。這個(gè)通用公式是三種注意力機(jī)制的共同數(shù)學(xué)基礎(chǔ)區(qū)別僅在于Q、K、V的來(lái)源以及是否應(yīng)用掩碼Mask。3. 核心機(jī)制一自注意力Self-Attention—— 建立序列內(nèi)部的全局關(guān)聯(lián)自注意力是Transformer的基石也是其得名“注意力就是你所需的一切”的底氣所在。它的核心思想非常簡(jiǎn)單讓序列中的每個(gè)元素都與該序列中的所有元素包括它自己進(jìn)行注意力計(jì)算。3.1 自注意力如何工作從“詞袋”到“關(guān)系網(wǎng)”想象一下我們要理解一句話“蘋(píng)果公司發(fā)布了新款手機(jī)它的設(shè)計(jì)很驚艷。” 對(duì)于傳統(tǒng)的詞嵌入模型每個(gè)詞“蘋(píng)果”、“公司”、“發(fā)布”等都被映射為一個(gè)獨(dú)立的向量它們之間缺乏明確的關(guān)聯(lián)。自注意力要做的事情就是為“蘋(píng)果”這個(gè)詞計(jì)算一個(gè)新的表示這個(gè)表示不僅包含“蘋(píng)果”自己的信息還包含了它與“公司”、“發(fā)布”、“手機(jī)”、“設(shè)計(jì)”、“驚艷”等所有詞的關(guān)系強(qiáng)度。具體計(jì)算步驟如下線性變換對(duì)于輸入序列中每個(gè)詞元的嵌入向量我們通過(guò)三個(gè)不同的權(quán)重矩陣W_QW_KW_V分別將其投影生成對(duì)應(yīng)的查詢(xún)向量q、鍵向量k和值向量v。這一步的目的是將原始嵌入轉(zhuǎn)換到適合計(jì)算注意力的空間。計(jì)算注意力分?jǐn)?shù)對(duì)于目標(biāo)詞元i例如“蘋(píng)果”我們?nèi)∑洳樵?xún)向量q_i與序列中所有詞元包括自身的鍵向量k_j進(jìn)行點(diǎn)積得到一系列分?jǐn)?shù)score_{ij} q_i · k_j。這個(gè)分?jǐn)?shù)直觀地反映了詞元i對(duì)詞元j的“關(guān)注程度”?？s放與歸一化將分?jǐn)?shù)除以鍵向量維度d_k的平方根進(jìn)行縮放然后通過(guò)softmax函數(shù)進(jìn)行歸一化得到注意力權(quán)重α_{ij}?？s放是為了在d_k較大時(shí)防止點(diǎn)積結(jié)果過(guò)大導(dǎo)致softmax進(jìn)入梯度極小的區(qū)域。加權(quán)求和最后將歸一化的注意力權(quán)重α_{ij}與對(duì)應(yīng)詞元的值向量v_j相乘并求和得到詞元i新的表示z_i Σ α_{ij} v_j。經(jīng)過(guò)這個(gè)過(guò)程“蘋(píng)果”的新向量z_蘋(píng)果就不再是一個(gè)孤立的“水果”或“品牌”概念而是一個(gè)融合了“蘋(píng)果公司發(fā)布了蘋(píng)果手機(jī)它的設(shè)計(jì)很驚艷”的豐富上下文信息的表示。它知道在這個(gè)句子里自己與“公司”和“手機(jī)”強(qiáng)相關(guān)。3.2 多頭注意力Multi-Head Attention并行化的多視角理解如果自注意力是讓模型學(xué)會(huì)關(guān)注全局那么多頭注意力就是讓模型同時(shí)從多個(gè)不同的角度進(jìn)行關(guān)注。這是Transformer性能強(qiáng)大的一個(gè)關(guān)鍵設(shè)計(jì)。單一的自注意力機(jī)制可以學(xué)習(xí)到一種特定的依賴(lài)模式。但在復(fù)雜的語(yǔ)言或視覺(jué)模式中詞語(yǔ)或圖像塊之間的關(guān)系是多元的。例如“蘋(píng)果”和“公司”之間可能是“所屬”關(guān)系而“蘋(píng)果”和“設(shè)計(jì)”之間可能是“評(píng)價(jià)”關(guān)系。單頭注意力可能難以同時(shí)捕獲所有這些不同類(lèi)型的關(guān)系。多頭注意力的做法是將查詢(xún)Q、鍵K、值V通過(guò)不同的線性投影矩陣并行地投影到h個(gè)例如8個(gè)更低維的子空間。在每個(gè)子空間稱(chēng)為一個(gè)“頭”上獨(dú)立執(zhí)行縮放點(diǎn)積注意力計(jì)算。將h個(gè)頭計(jì)算出的結(jié)果拼接Concat起來(lái)。最后通過(guò)一個(gè)線性投影矩陣W_O將拼接后的結(jié)果映射回預(yù)期的維度。公式表示為MultiHead(Q, K, V) Concat(head_1 ... head_h) W_Owhere head_i Attention(QW_i^Q KW_i^K VW_i^V)這相當(dāng)于讓模型擁有了多組“感官”或“專(zhuān)家”每組專(zhuān)注于捕捉不同子空間中的依賴(lài)關(guān)系——有的頭可能專(zhuān)注于語(yǔ)法結(jié)構(gòu)有的頭可能專(zhuān)注于指代關(guān)系有的頭可能專(zhuān)注于情感色彩。最后再將所有視角的信息融合得到更全面、更魯棒的表示。實(shí)操心得在實(shí)現(xiàn)多頭注意力時(shí)一個(gè)高效的技巧是使用矩陣運(yùn)算一次性完成所有頭的計(jì)算。我們并不需要真的創(chuàng)建h個(gè)獨(dú)立的矩陣然后循環(huán)計(jì)算。而是將QKV的維度從(batch_size seq_len d_model)通過(guò)線性變換和重塑reshape為(batch_size num_heads seq_len d_head)其中d_head d_model / num_heads。然后利用矩陣乘法的廣播機(jī)制一次性計(jì)算出所有頭的注意力輸出再重塑回來(lái)。這能極大利用GPU的并行計(jì)算能力。4. 核心機(jī)制二編碼器-解碼器注意力交叉注意力—— 連接兩個(gè)世界的橋梁在序列到序列的任務(wù)中如機(jī)器翻譯、文本摘要我們需要將源語(yǔ)言序列編碼器輸出的信息有效地傳遞并指導(dǎo)目標(biāo)語(yǔ)言序列解碼器的生成。這就是編碼器-解碼器注意力也稱(chēng)為交叉注意力的用武之地。4.1 為何需要交叉注意力解碼器在生成目標(biāo)序列的每一個(gè)詞時(shí)它需要知道源序列的哪些部分是當(dāng)前最相關(guān)的。例如將英文“I love machine learning”翻譯成中文“我熱愛(ài)機(jī)器學(xué)習(xí)”。當(dāng)解碼器生成“學(xué)習(xí)”這個(gè)詞時(shí)它應(yīng)該高度關(guān)注源句中的“l(fā)earning”而不是“I”或“l(fā)ove”。交叉注意力機(jī)制提供了這種動(dòng)態(tài)的、軟對(duì)齊的能力它比傳統(tǒng)的基于硬對(duì)齊如統(tǒng)計(jì)機(jī)器翻譯中的對(duì)齊表的方法更靈活、更強(qiáng)大。4.2 交叉注意力的計(jì)算邏輯交叉注意力發(fā)生在解碼器的每一層中在掩碼自注意力層之后。它的計(jì)算模式與自注意力完全相同都遵循Attention(Q K V) softmax(QK^T / sqrt(d_k)) V的公式。關(guān)鍵區(qū)別在于QKV的來(lái)源不同。查詢(xún)Q來(lái)源于解碼器上一層的輸出。可以理解為解碼器當(dāng)前正在努力構(gòu)建的目標(biāo)序列表示它提出“問(wèn)題”我現(xiàn)在該關(guān)注源序列的什么信息鍵K和值V都來(lái)源于編碼器最后一層的輸出。這是源序列經(jīng)過(guò)深度理解后的上下文表示它提供了可供查詢(xún)的“知識(shí)庫(kù)”。計(jì)算過(guò)程如下解碼器當(dāng)前的狀態(tài)作為Q去“查詢(xún)”編碼器的記憶K。計(jì)算Q和K中所有元素的相似度得到注意力權(quán)重。這個(gè)權(quán)重矩陣的每一行代表了在生成目標(biāo)序列某個(gè)位置時(shí)對(duì)源序列所有位置的關(guān)注程度。用這個(gè)權(quán)重對(duì)編碼器的值V通常V與K相同進(jìn)行加權(quán)求和得到一個(gè)“上下文向量”。這個(gè)向量融合了源序列中與當(dāng)前生成步驟最相關(guān)的信息。將這個(gè)上下文向量與解碼器原有的表示結(jié)合輸入到后續(xù)的前饋網(wǎng)絡(luò)最終預(yù)測(cè)出下一個(gè)詞元。通過(guò)這種方式解碼器在生成每一個(gè)詞時(shí)都能動(dòng)態(tài)地、有選擇地從源序列中提取信息實(shí)現(xiàn)了真正意義上的“內(nèi)容感知”生成。注意事項(xiàng)在訓(xùn)練時(shí)交叉注意力層的梯度會(huì)通過(guò)注意力權(quán)重反向傳播到編碼器。這意味著編碼器會(huì)被訓(xùn)練成能夠產(chǎn)生對(duì)解碼器有用的鍵值表示。這是一種端到端的聯(lián)合優(yōu)化編碼器和解碼器在訓(xùn)練中相互適應(yīng)。5. 核心機(jī)制三因果自注意力掩碼自注意力—— 確保未來(lái)的不可見(jiàn)性因果自注意力也稱(chēng)為掩碼自注意力是Transformer解碼器獨(dú)有的設(shè)計(jì)用于實(shí)現(xiàn)自回歸生成。自回歸生成是像GPT這樣的語(yǔ)言模型的核心工作方式根據(jù)已經(jīng)生成的詞來(lái)預(yù)測(cè)下一個(gè)詞。5.1 自回歸生成的核心約束在訓(xùn)練或生成時(shí)對(duì)于一個(gè)目標(biāo)序列“我 熱愛(ài) 機(jī)器 學(xué)習(xí)”當(dāng)模型預(yù)測(cè)第三個(gè)詞“機(jī)器”時(shí)它只能看到“我”和“熱愛(ài)”而不能看到未來(lái)的詞“學(xué)習(xí)”。如果它看到了“學(xué)習(xí)”那么預(yù)測(cè)“機(jī)器”就變成了一個(gè)平凡的任務(wù)因?yàn)橹懒舜鸢改Ｐ蜔o(wú)法學(xué)會(huì)真正的語(yǔ)言建模能力。因此必須施加一個(gè)約束每個(gè)位置只能關(guān)注該位置之前包括當(dāng)前位置的位置而不能關(guān)注之后的位置。5.2 掩碼的實(shí)現(xiàn)機(jī)制這個(gè)約束是通過(guò)一個(gè)“注意力掩碼”來(lái)實(shí)現(xiàn)的。具體來(lái)說(shuō)在計(jì)算注意力分?jǐn)?shù)矩陣S QK^T之后在送入softmax之前我們對(duì)這個(gè)矩陣進(jìn)行處理。我們創(chuàng)建一個(gè)下三角矩陣Lower Triangular Matrix其對(duì)角線及左下角元素為0或一個(gè)很小的負(fù)數(shù)如-1e9右上角未來(lái)位置元素為一個(gè)極大的負(fù)數(shù)如 -inf。有效位置為0 未來(lái)位置為 -inf [[0 -inf -inf -inf] [0 0 -inf -inf] [0 0 0 -inf] [0 0 0 0]]將這個(gè)掩碼矩陣加到注意力分?jǐn)?shù)矩陣S上。再進(jìn)行softmax計(jì)算。由于softmax對(duì)-inf的輸入會(huì)輸出0因此未來(lái)位置的注意力權(quán)重被強(qiáng)制設(shè)為0。這樣對(duì)于序列中的第i個(gè)位置其輸出z_i就只依賴(lài)于位置1到i的輸入完美滿足了自回歸的要求。實(shí)操心得與常見(jiàn)陷阱訓(xùn)練與推理的一致性在訓(xùn)練解碼器時(shí)我們通常使用“教師強(qiáng)制”策略即即便在預(yù)測(cè)第i個(gè)詞時(shí)我們也知道第i個(gè)詞的真實(shí)標(biāo)簽作為輸入。但即便如此掩碼也必須存在以確保模型在訓(xùn)練時(shí)學(xué)到的分布與推理時(shí)只能看到自己生成的歷史的分布是一致的。這是成功訓(xùn)練自回歸模型的關(guān)鍵。并行化訓(xùn)練掩碼的存在使得解碼器在訓(xùn)練時(shí)仍然是并行的我們一次性輸入整個(gè)目標(biāo)序列右移一位后的即起始符真實(shí)序列但通過(guò)掩碼每個(gè)位置的計(jì)算在數(shù)學(xué)上等價(jià)于只能看到歷史信息。這比RNN的串行訓(xùn)練快得多。掩碼的添加時(shí)機(jī)務(wù)必在縮放點(diǎn)積之后、softmax之前添加掩碼。如果在softmax之后添加由于softmax已經(jīng)將-inf轉(zhuǎn)換為了0再乘0沒(méi)有意義但更關(guān)鍵的是在softmax之前添加-inf能確保未來(lái)位置的概率嚴(yán)格為0梯度也為0。6. 位置編碼Positional Encoding為無(wú)位置感知的注意力注入順序信息這是一個(gè)必須與注意力機(jī)制放在一起討論的核心組件。細(xì)心的你可能已經(jīng)發(fā)現(xiàn)自注意力機(jī)制對(duì)輸入序列的處理是完全置換不變的。也就是說(shuō)打亂輸入序列的順序計(jì)算出的注意力權(quán)重和輸出在數(shù)學(xué)上只會(huì)隨之置換但模型本身無(wú)法感知“第一個(gè)詞”和“第二個(gè)詞”在原始序列中的先后關(guān)系。對(duì)于語(yǔ)言、時(shí)間序列等依賴(lài)順序的信息來(lái)說(shuō)這是致命的缺陷。位置編碼就是為了解決這個(gè)問(wèn)題而生的。它的作用是為序列中每個(gè)位置的詞元嵌入向量添加一個(gè)代表其位置信息的獨(dú)特向量。6.1 正弦余弦位置編碼的原理原始Transformer論文提出了一種非常巧妙且固定的位置編碼方式使用不同頻率的正弦和余弦函數(shù)PE(pos 2i) sin(pos / 10000^(2i/d_model))PE(pos 2i1) cos(pos / 10000^(2i/d_model))其中pos是位置i是維度索引。這種編碼方式有幾個(gè)優(yōu)良特性唯一性每個(gè)位置都有唯一的編碼。相對(duì)位置關(guān)系可學(xué)習(xí)對(duì)于固定的偏移量kPE(posk)可以表示為PE(pos)的線性函數(shù)這意味著模型可以很容易地學(xué)會(huì)關(guān)注相對(duì)位置信息。能處理比訓(xùn)練時(shí)更長(zhǎng)的序列由于是三角函數(shù)它可以外推到訓(xùn)練時(shí)未見(jiàn)過(guò)的更長(zhǎng)序列位置盡管效果可能下降。6.2 位置編碼與注意力的結(jié)合在實(shí)際操作中位置編碼PE與詞嵌入WE直接相加作為編碼器和解碼器第一層的輸入Input WE PE。通過(guò)這種方式模型在計(jì)算注意力時(shí)QKV中已經(jīng)包含了位置信息因此它可以學(xué)習(xí)到像“相鄰詞匯通常關(guān)系更緊密”這樣的模式。擴(kuò)展與變體可學(xué)習(xí)的位置編碼直接將位置編碼作為可訓(xùn)練的參數(shù)與模型一起學(xué)習(xí)。這在BERT等模型中很常見(jiàn)。它更靈活但可能缺乏三角函數(shù)編碼的外推能力。相對(duì)位置編碼上述的絕對(duì)位置編碼有一個(gè)潛在問(wèn)題它假設(shè)位置信息是絕對(duì)重要的。但很多時(shí)候兩個(gè)元素之間的相對(duì)距離如相距3個(gè)詞比它們的絕對(duì)位置第5個(gè)詞和第8個(gè)詞更重要。因此像Transformer-XL、T5等模型采用了相對(duì)位置編碼將位置信息注入到注意力分?jǐn)?shù)的計(jì)算中直接建模元素間的相對(duì)距離關(guān)系效果通常更好尤其是對(duì)于長(zhǎng)文本。在視覺(jué)Transformer中的應(yīng)用在ViT中除了對(duì)圖像塊序列添加一維位置編碼有時(shí)還會(huì)使用二維位置編碼將行、列位置信息分開(kāi)編碼再合并來(lái)更好地保留圖像的空間結(jié)構(gòu)信息。7. 三種注意力機(jī)制的協(xié)同與變體模型解析理解了這三種核心機(jī)制我們就能像搭積木一樣理解各種著名的Transformer變體模型。BERT (Bidirectional Encoder Representations from Transformers)它只使用了Transformer編碼器。因此它的核心是多頭自注意力并且是“雙向”的即每個(gè)詞能同時(shí)看到左右上下文因?yàn)闆](méi)有解碼器的掩碼。BERT通過(guò)MLM掩碼語(yǔ)言模型任務(wù)進(jìn)行預(yù)訓(xùn)練學(xué)習(xí)強(qiáng)大的上下文表征。GPT (Generative Pre-trained Transformer)它只使用了Transformer解碼器并且去掉了其中的編碼器-解碼器注意力層。因此GPT的核心是因果自注意力掩碼多頭自注意力和前饋網(wǎng)絡(luò)。它通過(guò)標(biāo)準(zhǔn)的自回歸語(yǔ)言模型任務(wù)進(jìn)行預(yù)訓(xùn)練學(xué)習(xí)生成文本的能力。原始Transformer / T5 / BART這些是完整的編碼器-解碼器架構(gòu)。編碼器使用自注意力理解源序列解碼器則依次使用因果自注意力關(guān)注已生成目標(biāo)歷史和交叉注意力關(guān)注編碼器源信息來(lái)生成目標(biāo)序列。適用于翻譯、摘要、問(wèn)答等需要從源到目標(biāo)轉(zhuǎn)換的任務(wù)。Vision Transformer (ViT)它將圖像分割成塊視為一個(gè)序列然后直接輸入到Transformer編碼器自注意力中進(jìn)行處理。它需要強(qiáng)大的位置編碼通常是可學(xué)習(xí)的一維或二維編碼來(lái)彌補(bǔ)自注意力對(duì)空間位置不敏感的缺陷。常見(jiàn)問(wèn)題與排查技巧實(shí)錄訓(xùn)練時(shí)Loss震蕩或不收斂檢查點(diǎn)首先檢查縮放因子sqrt(d_k)是否被正確應(yīng)用。如果忘記縮放在d_k較大時(shí)如512、1024點(diǎn)積QK^T的值會(huì)非常大導(dǎo)致softmax梯度接近于0訓(xùn)練困難。檢查點(diǎn)檢查位置編碼是否正確添加。如果忘記添加模型將無(wú)法學(xué)習(xí)序列順序在語(yǔ)言任務(wù)上性能會(huì)極差。檢查點(diǎn)檢查注意力掩碼對(duì)于解碼器或BERT的MLM是否正確實(shí)現(xiàn)。錯(cuò)誤的掩碼會(huì)導(dǎo)致信息泄露如未來(lái)信息被看到或該被關(guān)注的詞未被關(guān)注。推理時(shí)生成結(jié)果重復(fù)或無(wú)意義檢查點(diǎn)對(duì)于自回歸模型GPT類(lèi)確保在推理時(shí)每一步都正確應(yīng)用了因果掩碼。一個(gè)常見(jiàn)錯(cuò)誤是在推理循環(huán)中重復(fù)使用了訓(xùn)練時(shí)整個(gè)序列的掩碼而不是動(dòng)態(tài)地為當(dāng)前生成的序列長(zhǎng)度生成掩碼。檢查點(diǎn)檢查交叉注意力的輸入是否正確。在解碼器推理時(shí)每一步傳遞給交叉注意力層的編碼器輸出KV應(yīng)該是固定不變的整個(gè)源序列的編碼結(jié)果而查詢(xún)Q是當(dāng)前解碼器的隱狀態(tài)。多頭注意力效果不如單頭檢查點(diǎn)檢查d_model是否能被num_heads整除。d_head d_model / num_heads必須是整數(shù)。如果不是需要調(diào)整維度。檢查點(diǎn)檢查每個(gè)頭的輸出在拼接后經(jīng)過(guò)最終線性投影W_O的維度是否正確。輸出維度應(yīng)等于d_model。檢查點(diǎn)在資源有限時(shí)頭數(shù)并非越多越好。過(guò)多的頭數(shù)可能導(dǎo)致每個(gè)頭的表征能力d_head過(guò)小反而影響性能。通常d_head在64左右是一個(gè)經(jīng)驗(yàn)值。處理長(zhǎng)序列時(shí)內(nèi)存溢出OOM根源注意力矩陣的大小是(batch_size num_heads seq_len seq_len)。當(dāng)seq_len很大時(shí)如超過(guò)1024或2048這個(gè)矩陣會(huì)消耗巨大的內(nèi)存。解決方案優(yōu)化技巧使用梯度檢查點(diǎn)以時(shí)間換空間在反向傳播時(shí)重新計(jì)算部分前向結(jié)果。算法改進(jìn)研究并應(yīng)用高效注意力機(jī)制如Linformer低秩近似、Longformer局部全局注意力、BigBird稀疏注意力、FlashAttentionIO感知的精確注意力算法等。這些方法能顯著降低內(nèi)存和計(jì)算復(fù)雜度。工程手段減少批次大小batch size或使用模型并行。8. 從原理到實(shí)踐一個(gè)簡(jiǎn)化的注意力模塊代碼剖析理論說(shuō)了這么多最后我們通過(guò)一個(gè)高度簡(jiǎn)化但核心邏輯完整的PyTorch實(shí)現(xiàn)將自注意力的過(guò)程串聯(lián)起來(lái)。這能幫助你鞏固理解并作為自己實(shí)現(xiàn)的起點(diǎn)。import torch import torch.nn as nn import torch.nn.functional as F import math class SimpleSelfAttention(nn.Module): 一個(gè)簡(jiǎn)化的單頭自注意力模塊包含位置編碼。 def __init__(self d_model d_k d_v): super().__init__() self.d_k d_k # 定義線性變換層 self.W_q nn.Linear(d_model d_k) # 查詢(xún)變換 self.W_k nn.Linear(d_model d_k) # 鍵變換 self.W_v nn.Linear(d_model d_v) # 值變換 self.W_o nn.Linear(d_v d_model) # 輸出變換 def forward(self x maskNone): Args: x: 輸入張量形狀為 (batch_size seq_len d_model) mask: 可選掩碼形狀為 (batch_size seq_len seq_len) 或 (seq_len seq_len) Returns: 輸出張量形狀同輸入x batch_size seq_len _ x.shape # 1. 線性投影得到Q K V Q self.W_q(x) # (batch seq d_k) K self.W_k(x) # (batch seq d_k) V self.W_v(x) # (batch seq d_v) # 2. 計(jì)算縮放點(diǎn)積注意力分?jǐn)?shù) scores torch.matmul(Q K.transpose(-2 -1)) / math.sqrt(self.d_k) # (batch seq seq) # 3. 應(yīng)用注意力掩碼如果提供 if mask is not None: # 通常mask中1表示需要被掩蓋的位置我們將其替換為一個(gè)很大的負(fù)數(shù) scores scores.masked_fill(mask 1 -1e9) # 4. 應(yīng)用softmax得到注意力權(quán)重 attn_weights F.softmax(scores dim-1) # (batch seq seq) # 5. 加權(quán)求和得到上下文向量 context torch.matmul(attn_weights V) # (batch seq d_v) # 6. 輸出投影 output self.W_o(context) # (batch seq d_model) return output attn_weights # 返回輸出和注意力權(quán)重用于可視化或分析 # 示例使用正弦余弦位置編碼 class PositionalEncoding(nn.Module): def __init__(self d_model max_len5000): super().__init__() pe torch.zeros(max_len d_model) position torch.arange(0 max_len).unsqueeze(1) div_term torch.exp(torch.arange(0 d_model 2) * -(math.log(10000.0) / d_model)) pe[: 0::2] torch.sin(position * div_term) pe[: 1::2] torch.cos(position * div_term) pe pe.unsqueeze(0) # (1 max_len d_model) self.register_buffer(pe pe) # 不是模型參數(shù)但會(huì)隨模型保存/加載 def forward(self x): # x: (batch seq_len d_model) return x self.pe[: :x.size(1) :] # 組合使用示例 d_model 512 d_k d_v 64 seq_len 10 batch_size 4 # 模擬輸入詞嵌入 input_embeddings torch.randn(batch_size seq_len d_model) # 添加位置編碼 pos_encoder PositionalEncoding(d_model) x_with_pos pos_encoder(input_embeddings) # 創(chuàng)建自注意力層 attn_layer SimpleSelfAttention(d_model d_k d_v) # 前向傳播無(wú)掩碼即編碼器自注意力 output attn_weights attn_layer(x_with_pos) print(f輸入形狀 {input_embeddings.shape}) print(f輸出形狀 {output.shape}) print(f注意力權(quán)重形狀 {attn_weights.shape}) # 模擬一個(gè)因果掩碼下三角矩陣為0上三角為1 causal_mask torch.triu(torch.ones(seq_len seq_len) diagonal1).bool() print(因果掩碼True表示需要掩蓋) print(causal_mask) # 使用掩碼模擬解碼器 output_masked _ attn_layer(x_with_pos maskcausal_mask.unsqueeze(0)) # 增加batch維度這段代碼清晰地展示了從輸入到輸出的完整流程。在實(shí)際的Transformer實(shí)現(xiàn)中你會(huì)看到更復(fù)雜的多頭注意力類(lèi)它封裝了多個(gè)這樣的“頭”并處理了所有的線性變換、重塑和拼接操作。但核心的scaled_dot_product_attention計(jì)算邏輯是完全一致的。理解這三種注意力機(jī)制就如同掌握了Transformer這座大廈的承重結(jié)構(gòu)。自注意力賦予了模型理解上下文的能力交叉注意力搭建了信息傳遞的橋梁而因果自注意力則保證了生成過(guò)程的合理性。位置編碼則是讓這一切成為可能的“膠水”。當(dāng)你下次閱讀BERT、GPT或ViT的論文或代碼時(shí)嘗試去識(shí)別其中哪些部分是自注意力哪些部分應(yīng)用了何種掩碼你會(huì)有一種豁然開(kāi)朗的感覺(jué)。這不僅僅是理解一個(gè)模型更是掌握了一套構(gòu)建下一代智能系統(tǒng)的核心范式。 深入解析Transformer三大注意力機(jī)制：自注意力、交叉注意力與因果自注意力
1. 從“注意力”到“Transformer”一場(chǎng)認(rèn)知范式的革命如果你在2017年之后才開(kāi)始接觸深度學(xué)習(xí)尤其是自然語(yǔ)言處理或計(jì)算機(jī)視覺(jué)那么“注意力機(jī)制”和“Transformer”這兩個(gè)詞對(duì)你來(lái)說(shuō)可能就像空氣一樣自然存在。但回到那個(gè)時(shí)間點(diǎn)之前整個(gè)領(lǐng)域完全是另一番景象。循環(huán)神經(jīng)網(wǎng)絡(luò)RNN及其變體LSTM、GRU是處理序列數(shù)據(jù)的絕對(duì)王者我們花費(fèi)大量精力去設(shè)計(jì)門(mén)控結(jié)構(gòu)、解決梯度消失只為讓模型能“記住”更久遠(yuǎn)的信息。然而一個(gè)根本性的瓶頸始終存在序列計(jì)算的順序性。你必須一個(gè)接一個(gè)地處理輸入這既限制了訓(xùn)練時(shí)的并行效率也使得模型難以直接建立長(zhǎng)距離依賴(lài)關(guān)系尤其是在處理兩個(gè)相距很遠(yuǎn)的單詞或像素時(shí)信息需要經(jīng)過(guò)漫長(zhǎng)的“跋涉”才能關(guān)聯(lián)起來(lái)。Transformer架構(gòu)的橫空出世徹底打破了這一僵局。它最核心、最革命性的貢獻(xiàn)就是完全摒棄了循環(huán)結(jié)構(gòu)轉(zhuǎn)而依賴(lài)一種名為“自注意力”的機(jī)制來(lái)建立序列內(nèi)部所有元素之間的全局關(guān)聯(lián)。你可以把它想象成在一個(gè)會(huì)議室里開(kāi)會(huì)。RNN就像是一個(gè)只能聽(tīng)上一個(gè)人發(fā)言然后自己發(fā)言的串行模式而Transformer的自注意力機(jī)制則像是讓會(huì)議室里的每個(gè)人同時(shí)發(fā)言并且每個(gè)人都能瞬間聽(tīng)到并理解房間里所有人的發(fā)言?xún)?nèi)容然后綜合這些信息形成自己的觀點(diǎn)。這種“全局視野”的能力是之前任何模型都難以企及的。今天我們不再滿足于僅僅知道Transformer“很強(qiáng)大”。作為從業(yè)者我們需要深入其核心拆解它的引擎——注意力機(jī)制。很多人初學(xué)時(shí)會(huì)混淆“注意力機(jī)制”、“自注意力”和“多頭注意力”這些概念或者對(duì)“位置編碼”為何如此關(guān)鍵一知半解。這篇文章的目的就是幫你徹底厘清Transformer架構(gòu)中這三種核心的注意力機(jī)制自注意力、編碼器-解碼器注意力交叉注意力和因果自注意力掩碼自注意力。我們會(huì)從最直觀的比喻出發(fā)深入到數(shù)學(xué)計(jì)算和代碼實(shí)現(xiàn)最后探討它們?nèi)绾螀f(xié)同工作構(gòu)成了當(dāng)今大模型時(shí)代的基石。無(wú)論你是想徹底理解BERT、GPT、ViT這些明星模型的工作原理還是計(jì)劃在自己的項(xiàng)目中引入注意力模塊搞懂這三種機(jī)制都是必經(jīng)之路。2. Transformer架構(gòu)總覽與注意力機(jī)制的核心角色在深入三種具體的注意力機(jī)制之前我們有必要先快速回顧一下Transformer的整體架構(gòu)理解注意力機(jī)制在其中扮演的“發(fā)動(dòng)機(jī)”角色。原始的Transformer論文《Attention Is All You Need》提出的是一個(gè)用于序列到序列Seq2Seq任務(wù)的模型比如機(jī)器翻譯。它由編碼器Encoder和解碼器Decoder兩部分堆疊而成。編碼器由N個(gè)原論文中N6完全相同的層堆疊而成。每一層都包含兩個(gè)核心子層多頭自注意力機(jī)制Multi-Head Self-Attention前饋神經(jīng)網(wǎng)絡(luò)Position-wise Feed-Forward Network每個(gè)子層周?chē)疾捎昧藲埐钸B接Residual Connection和層歸一化Layer Normalization。編碼器的任務(wù)是接收輸入序列例如一句英文并通過(guò)自注意力機(jī)制提取其內(nèi)部豐富的上下文信息將每個(gè)輸入詞元Token編碼成一個(gè)蘊(yùn)含了全局上下文信息的向量表示。解碼器同樣由N個(gè)相同的層堆疊。每一層包含三個(gè)核心子層掩碼多頭自注意力機(jī)制Masked Multi-Head Self-Attention這是“因果自注意力”的具體實(shí)現(xiàn)。多頭編碼器-解碼器注意力機(jī)制Multi-Head Encoder-Decoder Attention這就是“交叉注意力”。前饋神經(jīng)網(wǎng)絡(luò)解碼器的任務(wù)是在已知編碼器輸出和已生成部分目標(biāo)序列的前提下自回歸地預(yù)測(cè)下一個(gè)詞元。注意雖然原始Transformer是編解碼結(jié)構(gòu)但后續(xù)很多著名模型只用了其中一部分。例如BERT只用了編碼器GPT系列只用了解碼器并去掉了其中的交叉注意力子層Vision Transformer (ViT) 則將圖像塊序列視為輸入使用編碼器進(jìn)行處理。因此理解這三種注意力機(jī)制是理解所有Transformer變體的基礎(chǔ)。那么什么是“注意力”的通用計(jì)算模式呢拋開(kāi)具體的變體其核心思想可以概括為給定一個(gè)“查詢(xún)”Query集合和一個(gè)“鍵-值”Key-Value對(duì)集合注意力機(jī)制通過(guò)計(jì)算Query與所有Key的相似度相關(guān)性來(lái)得到每個(gè)Key對(duì)應(yīng)的Value的加權(quán)和從而輸出一個(gè)聚焦于最重要信息的表示。用公式表示就是 scaled dot-product attention縮放點(diǎn)積注意力Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V這里Q、K、V分別是查詢(xún)、鍵、值矩陣。sqrt(d_k)是一個(gè)縮放因子用于防止點(diǎn)積結(jié)果過(guò)大導(dǎo)致softmax梯度消失。這個(gè)通用公式是三種注意力機(jī)制的共同數(shù)學(xué)基礎(chǔ)區(qū)別僅在于Q、K、V的來(lái)源以及是否應(yīng)用掩碼Mask。3. 核心機(jī)制一自注意力Self-Attention—— 建立序列內(nèi)部的全局關(guān)聯(lián)自注意力是Transformer的基石也是其得名“注意力就是你所需的一切”的底氣所在。它的核心思想非常簡(jiǎn)單讓序列中的每個(gè)元素都與該序列中的所有元素包括它自己進(jìn)行注意力計(jì)算。3.1 自注意力如何工作從“詞袋”到“關(guān)系網(wǎng)”想象一下我們要理解一句話“蘋(píng)果公司發(fā)布了新款手機(jī)它的設(shè)計(jì)很驚艷。” 對(duì)于傳統(tǒng)的詞嵌入模型每個(gè)詞“蘋(píng)果”、“公司”、“發(fā)布”等都被映射為一個(gè)獨(dú)立的向量它們之間缺乏明確的關(guān)聯(lián)。自注意力要做的事情就是為“蘋(píng)果”這個(gè)詞計(jì)算一個(gè)新的表示這個(gè)表示不僅包含“蘋(píng)果”自己的信息還包含了它與“公司”、“發(fā)布”、“手機(jī)”、“設(shè)計(jì)”、“驚艷”等所有詞的關(guān)系強(qiáng)度。具體計(jì)算步驟如下線性變換對(duì)于輸入序列中每個(gè)詞元的嵌入向量我們通過(guò)三個(gè)不同的權(quán)重矩陣W_QW_KW_V分別將其投影生成對(duì)應(yīng)的查詢(xún)向量q、鍵向量k和值向量v。這一步的目的是將原始嵌入轉(zhuǎn)換到適合計(jì)算注意力的空間。計(jì)算注意力分?jǐn)?shù)對(duì)于目標(biāo)詞元i例如“蘋(píng)果”我們?nèi)∑洳樵?xún)向量q_i與序列中所有詞元包括自身的鍵向量k_j進(jìn)行點(diǎn)積得到一系列分?jǐn)?shù)score_{ij} q_i · k_j。這個(gè)分?jǐn)?shù)直觀地反映了詞元i對(duì)詞元j的“關(guān)注程度”?？s放與歸一化將分?jǐn)?shù)除以鍵向量維度d_k的平方根進(jìn)行縮放然后通過(guò)softmax函數(shù)進(jìn)行歸一化得到注意力權(quán)重α_{ij}?？s放是為了在d_k較大時(shí)防止點(diǎn)積結(jié)果過(guò)大導(dǎo)致softmax進(jìn)入梯度極小的區(qū)域。加權(quán)求和最后將歸一化的注意力權(quán)重α_{ij}與對(duì)應(yīng)詞元的值向量v_j相乘并求和得到詞元i新的表示z_i Σ α_{ij} v_j。經(jīng)過(guò)這個(gè)過(guò)程“蘋(píng)果”的新向量z_蘋(píng)果就不再是一個(gè)孤立的“水果”或“品牌”概念而是一個(gè)融合了“蘋(píng)果公司發(fā)布了蘋(píng)果手機(jī)它的設(shè)計(jì)很驚艷”的豐富上下文信息的表示。它知道在這個(gè)句子里自己與“公司”和“手機(jī)”強(qiáng)相關(guān)。3.2 多頭注意力Multi-Head Attention并行化的多視角理解如果自注意力是讓模型學(xué)會(huì)關(guān)注全局那么多頭注意力就是讓模型同時(shí)從多個(gè)不同的角度進(jìn)行關(guān)注。這是Transformer性能強(qiáng)大的一個(gè)關(guān)鍵設(shè)計(jì)。單一的自注意力機(jī)制可以學(xué)習(xí)到一種特定的依賴(lài)模式。但在復(fù)雜的語(yǔ)言或視覺(jué)模式中詞語(yǔ)或圖像塊之間的關(guān)系是多元的。例如“蘋(píng)果”和“公司”之間可能是“所屬”關(guān)系而“蘋(píng)果”和“設(shè)計(jì)”之間可能是“評(píng)價(jià)”關(guān)系。單頭注意力可能難以同時(shí)捕獲所有這些不同類(lèi)型的關(guān)系。多頭注意力的做法是將查詢(xún)Q、鍵K、值V通過(guò)不同的線性投影矩陣并行地投影到h個(gè)例如8個(gè)更低維的子空間。在每個(gè)子空間稱(chēng)為一個(gè)“頭”上獨(dú)立執(zhí)行縮放點(diǎn)積注意力計(jì)算。將h個(gè)頭計(jì)算出的結(jié)果拼接Concat起來(lái)。最后通過(guò)一個(gè)線性投影矩陣W_O將拼接后的結(jié)果映射回預(yù)期的維度。公式表示為MultiHead(Q, K, V) Concat(head_1 ... head_h) W_Owhere head_i Attention(QW_i^Q KW_i^K VW_i^V)這相當(dāng)于讓模型擁有了多組“感官”或“專(zhuān)家”每組專(zhuān)注于捕捉不同子空間中的依賴(lài)關(guān)系——有的頭可能專(zhuān)注于語(yǔ)法結(jié)構(gòu)有的頭可能專(zhuān)注于指代關(guān)系有的頭可能專(zhuān)注于情感色彩。最后再將所有視角的信息融合得到更全面、更魯棒的表示。實(shí)操心得在實(shí)現(xiàn)多頭注意力時(shí)一個(gè)高效的技巧是使用矩陣運(yùn)算一次性完成所有頭的計(jì)算。我們并不需要真的創(chuàng)建h個(gè)獨(dú)立的矩陣然后循環(huán)計(jì)算。而是將QKV的維度從(batch_size seq_len d_model)通過(guò)線性變換和重塑reshape為(batch_size num_heads seq_len d_head)其中d_head d_model / num_heads。然后利用矩陣乘法的廣播機(jī)制一次性計(jì)算出所有頭的注意力輸出再重塑回來(lái)。這能極大利用GPU的并行計(jì)算能力。4. 核心機(jī)制二編碼器-解碼器注意力交叉注意力—— 連接兩個(gè)世界的橋梁在序列到序列的任務(wù)中如機(jī)器翻譯、文本摘要我們需要將源語(yǔ)言序列編碼器輸出的信息有效地傳遞并指導(dǎo)目標(biāo)語(yǔ)言序列解碼器的生成。這就是編碼器-解碼器注意力也稱(chēng)為交叉注意力的用武之地。4.1 為何需要交叉注意力解碼器在生成目標(biāo)序列的每一個(gè)詞時(shí)它需要知道源序列的哪些部分是當(dāng)前最相關(guān)的。例如將英文“I love machine learning”翻譯成中文“我熱愛(ài)機(jī)器學(xué)習(xí)”。當(dāng)解碼器生成“學(xué)習(xí)”這個(gè)詞時(shí)它應(yīng)該高度關(guān)注源句中的“l(fā)earning”而不是“I”或“l(fā)ove”。交叉注意力機(jī)制提供了這種動(dòng)態(tài)的、軟對(duì)齊的能力它比傳統(tǒng)的基于硬對(duì)齊如統(tǒng)計(jì)機(jī)器翻譯中的對(duì)齊表的方法更靈活、更強(qiáng)大。4.2 交叉注意力的計(jì)算邏輯交叉注意力發(fā)生在解碼器的每一層中在掩碼自注意力層之后。它的計(jì)算模式與自注意力完全相同都遵循Attention(Q K V) softmax(QK^T / sqrt(d_k)) V的公式。關(guān)鍵區(qū)別在于QKV的來(lái)源不同。查詢(xún)Q來(lái)源于解碼器上一層的輸出。可以理解為解碼器當(dāng)前正在努力構(gòu)建的目標(biāo)序列表示它提出“問(wèn)題”我現(xiàn)在該關(guān)注源序列的什么信息鍵K和值V都來(lái)源于編碼器最后一層的輸出。這是源序列經(jīng)過(guò)深度理解后的上下文表示它提供了可供查詢(xún)的“知識(shí)庫(kù)”。計(jì)算過(guò)程如下解碼器當(dāng)前的狀態(tài)作為Q去“查詢(xún)”編碼器的記憶K。計(jì)算Q和K中所有元素的相似度得到注意力權(quán)重。這個(gè)權(quán)重矩陣的每一行代表了在生成目標(biāo)序列某個(gè)位置時(shí)對(duì)源序列所有位置的關(guān)注程度。用這個(gè)權(quán)重對(duì)編碼器的值V通常V與K相同進(jìn)行加權(quán)求和得到一個(gè)“上下文向量”。這個(gè)向量融合了源序列中與當(dāng)前生成步驟最相關(guān)的信息。將這個(gè)上下文向量與解碼器原有的表示結(jié)合輸入到后續(xù)的前饋網(wǎng)絡(luò)最終預(yù)測(cè)出下一個(gè)詞元。通過(guò)這種方式解碼器在生成每一個(gè)詞時(shí)都能動(dòng)態(tài)地、有選擇地從源序列中提取信息實(shí)現(xiàn)了真正意義上的“內(nèi)容感知”生成。注意事項(xiàng)在訓(xùn)練時(shí)交叉注意力層的梯度會(huì)通過(guò)注意力權(quán)重反向傳播到編碼器。這意味著編碼器會(huì)被訓(xùn)練成能夠產(chǎn)生對(duì)解碼器有用的鍵值表示。這是一種端到端的聯(lián)合優(yōu)化編碼器和解碼器在訓(xùn)練中相互適應(yīng)。5. 核心機(jī)制三因果自注意力掩碼自注意力—— 確保未來(lái)的不可見(jiàn)性因果自注意力也稱(chēng)為掩碼自注意力是Transformer解碼器獨(dú)有的設(shè)計(jì)用于實(shí)現(xiàn)自回歸生成。自回歸生成是像GPT這樣的語(yǔ)言模型的核心工作方式根據(jù)已經(jīng)生成的詞來(lái)預(yù)測(cè)下一個(gè)詞。5.1 自回歸生成的核心約束在訓(xùn)練或生成時(shí)對(duì)于一個(gè)目標(biāo)序列“我 熱愛(ài) 機(jī)器 學(xué)習(xí)”當(dāng)模型預(yù)測(cè)第三個(gè)詞“機(jī)器”時(shí)它只能看到“我”和“熱愛(ài)”而不能看到未來(lái)的詞“學(xué)習(xí)”。如果它看到了“學(xué)習(xí)”那么預(yù)測(cè)“機(jī)器”就變成了一個(gè)平凡的任務(wù)因?yàn)橹懒舜鸢改Ｐ蜔o(wú)法學(xué)會(huì)真正的語(yǔ)言建模能力。因此必須施加一個(gè)約束每個(gè)位置只能關(guān)注該位置之前包括當(dāng)前位置的位置而不能關(guān)注之后的位置。5.2 掩碼的實(shí)現(xiàn)機(jī)制這個(gè)約束是通過(guò)一個(gè)“注意力掩碼”來(lái)實(shí)現(xiàn)的。具體來(lái)說(shuō)在計(jì)算注意力分?jǐn)?shù)矩陣S QK^T之后在送入softmax之前我們對(duì)這個(gè)矩陣進(jìn)行處理。我們創(chuàng)建一個(gè)下三角矩陣Lower Triangular Matrix其對(duì)角線及左下角元素為0或一個(gè)很小的負(fù)數(shù)如-1e9右上角未來(lái)位置元素為一個(gè)極大的負(fù)數(shù)如 -inf。有效位置為0 未來(lái)位置為 -inf [[0 -inf -inf -inf] [0 0 -inf -inf] [0 0 0 -inf] [0 0 0 0]]將這個(gè)掩碼矩陣加到注意力分?jǐn)?shù)矩陣S上。再進(jìn)行softmax計(jì)算。由于softmax對(duì)-inf的輸入會(huì)輸出0因此未來(lái)位置的注意力權(quán)重被強(qiáng)制設(shè)為0。這樣對(duì)于序列中的第i個(gè)位置其輸出z_i就只依賴(lài)于位置1到i的輸入完美滿足了自回歸的要求。實(shí)操心得與常見(jiàn)陷阱訓(xùn)練與推理的一致性在訓(xùn)練解碼器時(shí)我們通常使用“教師強(qiáng)制”策略即即便在預(yù)測(cè)第i個(gè)詞時(shí)我們也知道第i個(gè)詞的真實(shí)標(biāo)簽作為輸入。但即便如此掩碼也必須存在以確保模型在訓(xùn)練時(shí)學(xué)到的分布與推理時(shí)只能看到自己生成的歷史的分布是一致的。這是成功訓(xùn)練自回歸模型的關(guān)鍵。并行化訓(xùn)練掩碼的存在使得解碼器在訓(xùn)練時(shí)仍然是并行的我們一次性輸入整個(gè)目標(biāo)序列右移一位后的即起始符真實(shí)序列但通過(guò)掩碼每個(gè)位置的計(jì)算在數(shù)學(xué)上等價(jià)于只能看到歷史信息。這比RNN的串行訓(xùn)練快得多。掩碼的添加時(shí)機(jī)務(wù)必在縮放點(diǎn)積之后、softmax之前添加掩碼。如果在softmax之后添加由于softmax已經(jīng)將-inf轉(zhuǎn)換為了0再乘0沒(méi)有意義但更關(guān)鍵的是在softmax之前添加-inf能確保未來(lái)位置的概率嚴(yán)格為0梯度也為0。6. 位置編碼Positional Encoding為無(wú)位置感知的注意力注入順序信息這是一個(gè)必須與注意力機(jī)制放在一起討論的核心組件。細(xì)心的你可能已經(jīng)發(fā)現(xiàn)自注意力機(jī)制對(duì)輸入序列的處理是完全置換不變的。也就是說(shuō)打亂輸入序列的順序計(jì)算出的注意力權(quán)重和輸出在數(shù)學(xué)上只會(huì)隨之置換但模型本身無(wú)法感知“第一個(gè)詞”和“第二個(gè)詞”在原始序列中的先后關(guān)系。對(duì)于語(yǔ)言、時(shí)間序列等依賴(lài)順序的信息來(lái)說(shuō)這是致命的缺陷。位置編碼就是為了解決這個(gè)問(wèn)題而生的。它的作用是為序列中每個(gè)位置的詞元嵌入向量添加一個(gè)代表其位置信息的獨(dú)特向量。6.1 正弦余弦位置編碼的原理原始Transformer論文提出了一種非常巧妙且固定的位置編碼方式使用不同頻率的正弦和余弦函數(shù)PE(pos 2i) sin(pos / 10000^(2i/d_model))PE(pos 2i1) cos(pos / 10000^(2i/d_model))其中pos是位置i是維度索引。這種編碼方式有幾個(gè)優(yōu)良特性唯一性每個(gè)位置都有唯一的編碼。相對(duì)位置關(guān)系可學(xué)習(xí)對(duì)于固定的偏移量kPE(posk)可以表示為PE(pos)的線性函數(shù)這意味著模型可以很容易地學(xué)會(huì)關(guān)注相對(duì)位置信息。能處理比訓(xùn)練時(shí)更長(zhǎng)的序列由于是三角函數(shù)它可以外推到訓(xùn)練時(shí)未見(jiàn)過(guò)的更長(zhǎng)序列位置盡管效果可能下降。6.2 位置編碼與注意力的結(jié)合在實(shí)際操作中位置編碼PE與詞嵌入WE直接相加作為編碼器和解碼器第一層的輸入Input WE PE。通過(guò)這種方式模型在計(jì)算注意力時(shí)QKV中已經(jīng)包含了位置信息因此它可以學(xué)習(xí)到像“相鄰詞匯通常關(guān)系更緊密”這樣的模式。擴(kuò)展與變體可學(xué)習(xí)的位置編碼直接將位置編碼作為可訓(xùn)練的參數(shù)與模型一起學(xué)習(xí)。這在BERT等模型中很常見(jiàn)。它更靈活但可能缺乏三角函數(shù)編碼的外推能力。相對(duì)位置編碼上述的絕對(duì)位置編碼有一個(gè)潛在問(wèn)題它假設(shè)位置信息是絕對(duì)重要的。但很多時(shí)候兩個(gè)元素之間的相對(duì)距離如相距3個(gè)詞比它們的絕對(duì)位置第5個(gè)詞和第8個(gè)詞更重要。因此像Transformer-XL、T5等模型采用了相對(duì)位置編碼將位置信息注入到注意力分?jǐn)?shù)的計(jì)算中直接建模元素間的相對(duì)距離關(guān)系效果通常更好尤其是對(duì)于長(zhǎng)文本。在視覺(jué)Transformer中的應(yīng)用在ViT中除了對(duì)圖像塊序列添加一維位置編碼有時(shí)還會(huì)使用二維位置編碼將行、列位置信息分開(kāi)編碼再合并來(lái)更好地保留圖像的空間結(jié)構(gòu)信息。7. 三種注意力機(jī)制的協(xié)同與變體模型解析理解了這三種核心機(jī)制我們就能像搭積木一樣理解各種著名的Transformer變體模型。BERT (Bidirectional Encoder Representations from Transformers)它只使用了Transformer編碼器。因此它的核心是多頭自注意力并且是“雙向”的即每個(gè)詞能同時(shí)看到左右上下文因?yàn)闆](méi)有解碼器的掩碼。BERT通過(guò)MLM掩碼語(yǔ)言模型任務(wù)進(jìn)行預(yù)訓(xùn)練學(xué)習(xí)強(qiáng)大的上下文表征。GPT (Generative Pre-trained Transformer)它只使用了Transformer解碼器并且去掉了其中的編碼器-解碼器注意力層。因此GPT的核心是因果自注意力掩碼多頭自注意力和前饋網(wǎng)絡(luò)。它通過(guò)標(biāo)準(zhǔn)的自回歸語(yǔ)言模型任務(wù)進(jìn)行預(yù)訓(xùn)練學(xué)習(xí)生成文本的能力。原始Transformer / T5 / BART這些是完整的編碼器-解碼器架構(gòu)。編碼器使用自注意力理解源序列解碼器則依次使用因果自注意力關(guān)注已生成目標(biāo)歷史和交叉注意力關(guān)注編碼器源信息來(lái)生成目標(biāo)序列。適用于翻譯、摘要、問(wèn)答等需要從源到目標(biāo)轉(zhuǎn)換的任務(wù)。Vision Transformer (ViT)它將圖像分割成塊視為一個(gè)序列然后直接輸入到Transformer編碼器自注意力中進(jìn)行處理。它需要強(qiáng)大的位置編碼通常是可學(xué)習(xí)的一維或二維編碼來(lái)彌補(bǔ)自注意力對(duì)空間位置不敏感的缺陷。常見(jiàn)問(wèn)題與排查技巧實(shí)錄訓(xùn)練時(shí)Loss震蕩或不收斂檢查點(diǎn)首先檢查縮放因子sqrt(d_k)是否被正確應(yīng)用。如果忘記縮放在d_k較大時(shí)如512、1024點(diǎn)積QK^T的值會(huì)非常大導(dǎo)致softmax梯度接近于0訓(xùn)練困難。檢查點(diǎn)檢查位置編碼是否正確添加。如果忘記添加模型將無(wú)法學(xué)習(xí)序列順序在語(yǔ)言任務(wù)上性能會(huì)極差。檢查點(diǎn)檢查注意力掩碼對(duì)于解碼器或BERT的MLM是否正確實(shí)現(xiàn)。錯(cuò)誤的掩碼會(huì)導(dǎo)致信息泄露如未來(lái)信息被看到或該被關(guān)注的詞未被關(guān)注。推理時(shí)生成結(jié)果重復(fù)或無(wú)意義檢查點(diǎn)對(duì)于自回歸模型GPT類(lèi)確保在推理時(shí)每一步都正確應(yīng)用了因果掩碼。一個(gè)常見(jiàn)錯(cuò)誤是在推理循環(huán)中重復(fù)使用了訓(xùn)練時(shí)整個(gè)序列的掩碼而不是動(dòng)態(tài)地為當(dāng)前生成的序列長(zhǎng)度生成掩碼。檢查點(diǎn)檢查交叉注意力的輸入是否正確。在解碼器推理時(shí)每一步傳遞給交叉注意力層的編碼器輸出KV應(yīng)該是固定不變的整個(gè)源序列的編碼結(jié)果而查詢(xún)Q是當(dāng)前解碼器的隱狀態(tài)。多頭注意力效果不如單頭檢查點(diǎn)檢查d_model是否能被num_heads整除。d_head d_model / num_heads必須是整數(shù)。如果不是需要調(diào)整維度。檢查點(diǎn)檢查每個(gè)頭的輸出在拼接后經(jīng)過(guò)最終線性投影W_O的維度是否正確。輸出維度應(yīng)等于d_model。檢查點(diǎn)在資源有限時(shí)頭數(shù)并非越多越好。過(guò)多的頭數(shù)可能導(dǎo)致每個(gè)頭的表征能力d_head過(guò)小反而影響性能。通常d_head在64左右是一個(gè)經(jīng)驗(yàn)值。處理長(zhǎng)序列時(shí)內(nèi)存溢出OOM根源注意力矩陣的大小是(batch_size num_heads seq_len seq_len)。當(dāng)seq_len很大時(shí)如超過(guò)1024或2048這個(gè)矩陣會(huì)消耗巨大的內(nèi)存。解決方案優(yōu)化技巧使用梯度檢查點(diǎn)以時(shí)間換空間在反向傳播時(shí)重新計(jì)算部分前向結(jié)果。算法改進(jìn)研究并應(yīng)用高效注意力機(jī)制如Linformer低秩近似、Longformer局部全局注意力、BigBird稀疏注意力、FlashAttentionIO感知的精確注意力算法等。這些方法能顯著降低內(nèi)存和計(jì)算復(fù)雜度。工程手段減少批次大小batch size或使用模型并行。8. 從原理到實(shí)踐一個(gè)簡(jiǎn)化的注意力模塊代碼剖析理論說(shuō)了這么多最后我們通過(guò)一個(gè)高度簡(jiǎn)化但核心邏輯完整的PyTorch實(shí)現(xiàn)將自注意力的過(guò)程串聯(lián)起來(lái)。這能幫助你鞏固理解并作為自己實(shí)現(xiàn)的起點(diǎn)。import torch import torch.nn as nn import torch.nn.functional as F import math class SimpleSelfAttention(nn.Module): 一個(gè)簡(jiǎn)化的單頭自注意力模塊包含位置編碼。 def __init__(self d_model d_k d_v): super().__init__() self.d_k d_k # 定義線性變換層 self.W_q nn.Linear(d_model d_k) # 查詢(xún)變換 self.W_k nn.Linear(d_model d_k) # 鍵變換 self.W_v nn.Linear(d_model d_v) # 值變換 self.W_o nn.Linear(d_v d_model) # 輸出變換 def forward(self x maskNone): Args: x: 輸入張量形狀為 (batch_size seq_len d_model) mask: 可選掩碼形狀為 (batch_size seq_len seq_len) 或 (seq_len seq_len) Returns: 輸出張量形狀同輸入x batch_size seq_len _ x.shape # 1. 線性投影得到Q K V Q self.W_q(x) # (batch seq d_k) K self.W_k(x) # (batch seq d_k) V self.W_v(x) # (batch seq d_v) # 2. 計(jì)算縮放點(diǎn)積注意力分?jǐn)?shù) scores torch.matmul(Q K.transpose(-2 -1)) / math.sqrt(self.d_k) # (batch seq seq) # 3. 應(yīng)用注意力掩碼如果提供 if mask is not None: # 通常mask中1表示需要被掩蓋的位置我們將其替換為一個(gè)很大的負(fù)數(shù) scores scores.masked_fill(mask 1 -1e9) # 4. 應(yīng)用softmax得到注意力權(quán)重 attn_weights F.softmax(scores dim-1) # (batch seq seq) # 5. 加權(quán)求和得到上下文向量 context torch.matmul(attn_weights V) # (batch seq d_v) # 6. 輸出投影 output self.W_o(context) # (batch seq d_model) return output attn_weights # 返回輸出和注意力權(quán)重用于可視化或分析 # 示例使用正弦余弦位置編碼 class PositionalEncoding(nn.Module): def __init__(self d_model max_len5000): super().__init__() pe torch.zeros(max_len d_model) position torch.arange(0 max_len).unsqueeze(1) div_term torch.exp(torch.arange(0 d_model 2) * -(math.log(10000.0) / d_model)) pe[: 0::2] torch.sin(position * div_term) pe[: 1::2] torch.cos(position * div_term) pe pe.unsqueeze(0) # (1 max_len d_model) self.register_buffer(pe pe) # 不是模型參數(shù)但會(huì)隨模型保存/加載 def forward(self x): # x: (batch seq_len d_model) return x self.pe[: :x.size(1) :] # 組合使用示例 d_model 512 d_k d_v 64 seq_len 10 batch_size 4 # 模擬輸入詞嵌入 input_embeddings torch.randn(batch_size seq_len d_model) # 添加位置編碼 pos_encoder PositionalEncoding(d_model) x_with_pos pos_encoder(input_embeddings) # 創(chuàng)建自注意力層 attn_layer SimpleSelfAttention(d_model d_k d_v) # 前向傳播無(wú)掩碼即編碼器自注意力 output attn_weights attn_layer(x_with_pos) print(f輸入形狀 {input_embeddings.shape}) print(f輸出形狀 {output.shape}) print(f注意力權(quán)重形狀 {attn_weights.shape}) # 模擬一個(gè)因果掩碼下三角矩陣為0上三角為1 causal_mask torch.triu(torch.ones(seq_len seq_len) diagonal1).bool() print(因果掩碼True表示需要掩蓋) print(causal_mask) # 使用掩碼模擬解碼器 output_masked _ attn_layer(x_with_pos maskcausal_mask.unsqueeze(0)) # 增加batch維度這段代碼清晰地展示了從輸入到輸出的完整流程。在實(shí)際的Transformer實(shí)現(xiàn)中你會(huì)看到更復(fù)雜的多頭注意力類(lèi)它封裝了多個(gè)這樣的“頭”并處理了所有的線性變換、重塑和拼接操作。但核心的scaled_dot_product_attention計(jì)算邏輯是完全一致的。理解這三種注意力機(jī)制就如同掌握了Transformer這座大廈的承重結(jié)構(gòu)。自注意力賦予了模型理解上下文的能力交叉注意力搭建了信息傳遞的橋梁而因果自注意力則保證了生成過(guò)程的合理性。位置編碼則是讓這一切成為可能的“膠水”。當(dāng)你下次閱讀BERT、GPT或ViT的論文或代碼時(shí)嘗試去識(shí)別其中哪些部分是自注意力哪些部分應(yīng)用了何種掩碼你會(huì)有一種豁然開(kāi)朗的感覺(jué)。這不僅僅是理解一個(gè)模型更是掌握了一套構(gòu)建下一代智能系統(tǒng)的核心范式。