Transformer 是近幾年深度學(xué)習(xí)領(lǐng)域繞不開(kāi)的一個(gè)名字。很多同學(xué)從 CNN、RNN 一路學(xué)過(guò)來(lái)看到各類大模型、多模態(tài)模型都建立在 Transformer 之上難免會(huì)想這一套架構(gòu)到底是怎么設(shè)計(jì)的QKV 是什么位置編碼又是干什么的網(wǎng)上資料很多但要么偏數(shù)學(xué)推導(dǎo)要么只給調(diào)用代碼很難有一條線把“從 Token 到訓(xùn)練”完整串起來(lái)。本文就用 PyTorch 從零實(shí)現(xiàn)一個(gè)可運(yùn)行的 Transformer 編碼器把詞嵌入、位置編碼、自注意力、多頭注意力、殘差連接、層歸一化這些概念全部落實(shí)到代碼里最后再用一個(gè)簡(jiǎn)單的文本分類訓(xùn)練任務(wù)驗(yàn)證效果。適合剛?cè)腴T深度學(xué)習(xí)的同學(xué)也適合想要把理論真正“落地”的開(kāi)發(fā)者。1. 憑什么 Transformer 能成為主流架構(gòu)1.1 從序列建模的痛點(diǎn)說(shuō)起在 Transformer 出現(xiàn)之前處理文本、語(yǔ)音等序列數(shù)據(jù)最常用的模型是 RNN 和 LSTM。它們的基本思路是“逐個(gè)讀入”序列中的元素用一個(gè)隱藏狀態(tài)把前面的信息記憶下來(lái)再傳給下一個(gè)時(shí)刻。這個(gè)思路簡(jiǎn)單但也有兩個(gè)很難繞開(kāi)的毛病長(zhǎng)距離依賴問(wèn)題。信息要在時(shí)間方向上一步一步傳遞當(dāng)序列很長(zhǎng)時(shí)前面的內(nèi)容很容易在傳遞過(guò)程中被“稀釋”模型很難學(xué)到兩個(gè)相隔很遠(yuǎn)詞之間的關(guān)聯(lián)。串行計(jì)算效率低。每個(gè)時(shí)刻的輸出依賴前一個(gè)時(shí)刻的狀態(tài)GPU 很難把整個(gè)序列的計(jì)算并行起來(lái)訓(xùn)練速度上不占優(yōu)勢(shì)。CNN 在處理圖像時(shí)表現(xiàn)很好但卷積核的感受野是局部的想要覆蓋長(zhǎng)文本也需要堆很多層。于是一個(gè)既能直接建模全局關(guān)系、又能并行計(jì)算的架構(gòu)就成了剛需。1.2 Transformer 的核心思想注意力取代循環(huán)Transformer 在 2017 年由 Google 團(tuán)隊(duì)在論文Attention Is All You Need中提出。它徹底拋棄了循環(huán)結(jié)構(gòu)改用“自注意力機(jī)制”Self-Attention來(lái)刻畫序列中任意兩個(gè)位置之間的關(guān)系。一個(gè)詞在編碼時(shí)可以直接“看到”句子里的所有其他詞并根據(jù)相關(guān)性分配不同的注意力權(quán)重。你可以這樣通俗理解給模型一個(gè)句子“小明因?yàn)榭荚嚨昧藵M分所以非常開(kāi)心”當(dāng)模型編碼“開(kāi)心”這個(gè)詞時(shí)自注意力機(jī)制會(huì)幫它找到“滿分”“考試”這些關(guān)鍵證據(jù)并讓“開(kāi)心”的向量表示中融入這些相關(guān)信息。這種并行地捕捉全局依賴的能力正是 Transformer 在長(zhǎng)文本、大模型場(chǎng)景下表現(xiàn)優(yōu)異的重要原因?；?Transformer 延伸出的 BERT、GPT、ViT、Swin Transformer 等模型覆蓋了自然語(yǔ)言處理、圖像分類、目標(biāo)檢測(cè)、多模態(tài)理解等大量方向。因此把 Transformer 的核心組件親手實(shí)現(xiàn)一遍理解它的前向傳播過(guò)程對(duì)后續(xù)閱讀源碼、修改模型、甚至做模型部署都很有幫助。1.3 本文采用的實(shí)現(xiàn)路線Transformer 的完整結(jié)構(gòu)包含 Encoder 和 Decoder 兩大塊。Encoder 負(fù)責(zé)把輸入序列編碼成上下文相關(guān)的向量表示Decoder 負(fù)責(zé)根據(jù)這些表示逐步生成輸出序列。本文為了聚焦核心原理選擇實(shí)現(xiàn)一個(gè)Encoder-only 的簡(jiǎn)化版 Transformer并把任務(wù)設(shè)計(jì)為文本分類。這樣做的好處是不需要處理復(fù)雜的自回歸生成邏輯我們可以把主要精力放在 Embedding、位置編碼、多頭自注意力、前饋網(wǎng)絡(luò)、殘差連接與層歸一化這些通用組件上。掌握這些組件之后再去看 Encoder-Decoder 結(jié)構(gòu)或者 GPT 的因果自注意力都會(huì)輕松很多。2. 環(huán)境準(zhǔn)備與 PyTorch 安裝2.1 本地環(huán)境建議本文代碼基于 PyTorch 編寫。不同操作系統(tǒng)、不同 GPU 環(huán)境下安裝命令會(huì)有差異下面的步驟以最常見(jiàn)的 CPU/GPU 混合環(huán)境為例基本思路是通用的。# 使用 conda 創(chuàng)建獨(dú)立環(huán)境避免依賴沖突 conda create -n transformer-demo python3.10 conda activate transformer-demoPyTorch 的安裝建議直接參考 PyTorch 官網(wǎng)給出的命令。以 CUDA 11.8 為例可以使用類似下面的命令pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118如果你使用的是 CPU 環(huán)境可以把安裝命令換成pip install torch torchvision torchaudio如果你的顯卡是 AMD可以關(guān)注 PyTorch 官方針對(duì) ROCm 發(fā)布的安裝包如果只是做代碼學(xué)習(xí)CPU 也完全足夠跑通本文示例。需要注意的是PyTorch 版本迭代比較快不同版本的 API 有一些差異。本文代碼使用的是nn.Module、nn.Linear、nn.LayerNorm等穩(wěn)定 API在 PyTorch 1.13 到 2.x 的常見(jiàn)版本中都能正常運(yùn)行。如果你在安裝過(guò)程中遇到下載速度慢的問(wèn)題可以考慮把 pip 源切換為國(guó)內(nèi)鏡像源例如pip install torch -i https://pypi.tuna.tsinghua.edu.cn/simple但要注意如果使用 NVIDIA GPU 且需要 CUDA 加速建議還是在官網(wǎng)選擇正確的 CUDA 版本安裝命令避免額外踩坑。2.2 驗(yàn)證安裝結(jié)果安裝完成后在終端里輸入python進(jìn)入交互式環(huán)境執(zhí)行下面的代碼import torch print(torch.__version__) print(torch.cuda.is_available())如果torch.cuda.is_available()輸出True說(shuō)明 GPU 可用輸出False也不影響本次實(shí)驗(yàn)代碼會(huì)在 CPU 上運(yùn)行只是訓(xùn)練稍慢一些。3. Token 與 Embedding輸入是如何變成向量的3.1 Token 到底是什么在深度學(xué)習(xí)模型里我們沒(méi)法直接把“你好世界”這樣的字符串喂給網(wǎng)絡(luò)。第一步是把它切分成一個(gè)個(gè)最小單元這些單元就叫 Token。Token 可以是一個(gè)字、一個(gè)詞也可能是通過(guò) BPE 等算法切分出的子詞片段。例如“我喜歡深度學(xué)習(xí)”可以拆成[我, 喜歡, 深度, 學(xué)習(xí)]每個(gè) Token 在詞表中都有一個(gè)唯一編號(hào)。我們可以給每個(gè) Token 分配一個(gè)整數(shù) ID比如{我: 2, 喜歡: 7, 深度: 11, 學(xué)習(xí): 15}于是句子就變成了一個(gè)整數(shù)序列[2, 7, 11, 15]這個(gè)過(guò)程通常由分詞器Tokenizer完成。在本文的簡(jiǎn)化實(shí)現(xiàn)里我們會(huì)跳過(guò)分詞細(xì)節(jié)直接使用torch.randint生成模擬的 Token ID 序列來(lái)演示模型前向傳播和訓(xùn)練流程。實(shí)際項(xiàng)目中你可以使用 HuggingFace 的BertTokenizer、GPT2Tokenizer等工具完成真實(shí)分詞。3.2 詞嵌入層的作用有了 Token ID 之后還需要把它們映射成稠密向量。這個(gè)映射層就是 Embedding 層。在 PyTorch 中nn.Embedding(vocab_size, d_model)可以理解為一個(gè)查找表給定 Token ID它返回對(duì)應(yīng)的d_model維向量。import torch import torch.nn as nn vocab_size 1000 # 詞表大小 d_model 128 # 隱藏維度 embedding nn.Embedding(vocab_size, d_model) tokens torch.randint(0, vocab_size, (4, 10)) # 模擬 batch4, seq_len10 的 Token ID embedded embedding(tokens) print(embedded.shape) # torch.Size([4, 10, 128])這里輸出的張量形狀是(batch_size, seq_len, d_model)含義是每個(gè)批次里有 4 個(gè)句子每個(gè)句子的長(zhǎng)度是 10 個(gè) Token每個(gè) Token 被表示成一個(gè) 128 維向量。需要注意Embedding 層的參數(shù)是可學(xué)習(xí)的。在訓(xùn)練過(guò)程中網(wǎng)絡(luò)會(huì)不斷調(diào)整這些向量使語(yǔ)義相近的 Token 在向量空間中距離更近。這也是為什么經(jīng)過(guò)預(yù)訓(xùn)練的 Embedding 往往比隨機(jī)初始化的效果好很多。4. 位置編碼給模型補(bǔ)上“順序感”4.1 為什么必須加位置信息自注意力機(jī)制本身對(duì)輸入順序是不敏感的。它會(huì)把序列中所有 Token 當(dāng)作一個(gè)集合來(lái)計(jì)算相關(guān)性無(wú)論“我打你”還是“你打我”如果把 Token 順序打亂注意力計(jì)算的結(jié)果是完全一樣的。可語(yǔ)言是講究順序的“我打你”和“你打我”語(yǔ)義完全不同。因此必須在輸入中注入位置信息。解決方案是在 Embedding 向量上疊加一個(gè)“位置向量”。簡(jiǎn)單來(lái)說(shuō)如果第 i 個(gè) Token 的 Embedding 是e_i位置編碼是p_i那么真正輸入模型的就是x_i e_i p_i4.2 正余弦位置編碼公式原版 Transformer 使用正余弦函數(shù)生成位置編碼。對(duì)于位置pos的第i個(gè)維度編碼公式如下PE(pos, 2i) sin(pos / 10000^(2i / d_model)) PE(pos, 2i 1) cos(pos / 10000^(2i / d_model))這個(gè)公式最巧妙的性質(zhì)是不同位置產(chǎn)生的編碼向量是周期性的、唯一的而且不需要額外學(xué)習(xí)參數(shù)。高頻維度的變化較快負(fù)責(zé)區(qū)分相鄰位置低頻維度的變化較慢負(fù)責(zé)區(qū)分遠(yuǎn)距離位置。正余弦交替使用也讓不同位置向量之間具備一定的線性關(guān)系便于模型捕捉相對(duì)位置信息。4.3 PyTorch 實(shí)現(xiàn)位置編碼import math import torch import torch.nn as nn class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000, dropout0.1): super().__init__() self.dropout nn.Dropout(pdropout) # 預(yù)先計(jì)算好所有位置的編碼 pe torch.zeros(max_len, d_model) position torch.arange(0, max_len).unsqueeze(1).float() div_term torch.exp( torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model) ) # 偶數(shù)維度用 sin奇數(shù)維度用 cos pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) # 增加 batch 維度注冊(cè)為 buffer不參與梯度更新 pe pe.unsqueeze(0) self.register_buffer(pe, pe) def forward(self, x): # x 形狀: (batch, seq_len, d_model) seq_len x.size(1) x x self.pe[:, :seq_len, :] return self.dropout(x)這里有三個(gè)細(xì)節(jié)需要留意torch.arange(0, d_model, 2)取的是 0、2、4……這樣的偶數(shù)索引對(duì)應(yīng)公式中的2i。使用register_buffer注冊(cè)位置編碼它會(huì)被model.to(device)自動(dòng)搬運(yùn)到 GPU但不會(huì)參與梯度計(jì)算。forward中只取self.pe[:, :seq_len, :]是為了保證位置編碼長(zhǎng)度跟當(dāng)前輸入的序列長(zhǎng)度一致避免形狀不匹配。為什么選擇正余弦而不是直接學(xué)一個(gè)位置向量一個(gè)原因是它可以外推到比訓(xùn)練時(shí)更長(zhǎng)的序列長(zhǎng)度另一個(gè)原因是正余弦編碼包含大量的先驗(yàn)歸納偏置在序列長(zhǎng)度不大時(shí)效果也不比可學(xué)習(xí)位置編碼差。后續(xù)很多模型比如 BERT 使用的是可學(xué)習(xí)位置編碼兩種方案都可以沒(méi)有絕對(duì)優(yōu)劣。5. 自注意力機(jī)制讓 Token 之間互相“關(guān)注”5.1 從相似度到注意力權(quán)重自注意力的核心是計(jì)算序列中任意兩個(gè) Token 之間的相關(guān)度。假設(shè)輸入向量為x_1, x_2, ..., x_n我們希望為每個(gè)位置計(jì)算一個(gè)新的向量這個(gè)向量是其他所有位置向量的加權(quán)求和權(quán)重就由相關(guān)性決定。步驟可以拆解為計(jì)算查詢向量q和鍵向量k的內(nèi)積作為兩兩之間的相似度分?jǐn)?shù)。將分?jǐn)?shù)除以縮放因子sqrt(d_k)避免數(shù)值過(guò)大導(dǎo)致 Softmax 梯度消失。通過(guò) Softmax 歸一化成權(quán)重保證權(quán)重的和為 1。用權(quán)重對(duì)所有“值向量”v做加權(quán)求和。5.2 QKV 到底代表什么Q、K、V 是自注意力里最核心的三個(gè)概念很多人第一次接觸時(shí)會(huì)覺(jué)得抽象?？梢杂靡粋€(gè)檢索場(chǎng)景來(lái)比喻Query查詢表示你當(dāng)前正在關(guān)注的信息內(nèi)容。Key鍵表示序列中每個(gè) Token 的“標(biāo)簽”用來(lái)和 Query 做匹配。Value值表示每個(gè) Token 真正攜帶的信息內(nèi)容。注意力分?jǐn)?shù)Q * K^T就是 Query 和 Key 的匹配程度。匹配度越高對(duì)應(yīng)的 Value 在最終輸出中占的比重就越大。在實(shí)際代碼里Q、K、V 都是通過(guò)線性變換從輸入x得到的Q x W_Q K x W_K V x W_V這里的W_Q、W_K、W_V都是可學(xué)習(xí)的參數(shù)矩陣。它們給出的不是原始向量而是“經(jīng)過(guò)投影后的表示”這種設(shè)計(jì)增加了模型的表達(dá)能力。5.3 多頭注意力讓模型從多個(gè)角度觀察單頭注意力相當(dāng)于模型只用一組 QKV 投影觀察一種“關(guān)聯(lián)模式”。但序列中的關(guān)系往往是多層次的一個(gè)詞可能既和句法主謂有關(guān)又和語(yǔ)義情感有關(guān)。多頭注意力就是把d_model維空間切分成n_heads個(gè)子空間每個(gè)頭獨(dú)立做注意力計(jì)算最后拼接起來(lái)再做一次線性變換。例如d_model128、n_heads4每個(gè)頭的維度就是128 / 4 32。這樣既增加了模型的表達(dá)空間又不會(huì)顯著增加計(jì)算量還讓注意力權(quán)重更容易被理解和可視化。5.4 PyTorch 實(shí)現(xiàn)多頭自注意力下面這段代碼是本文的核心組件之一。我們先把代碼完整貼出來(lái)再逐段解釋。import math import torch import torch.nn as nn class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout0.1): super().__init__() assert d_model % n_heads 0, d_model 必須能被 n_heads 整除 self.d_model d_model self.n_heads n_heads self.head_dim d_model // n_heads self.scale math.sqrt(self.head_dim) 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(pdropout) def forward(self, x, maskNone): batch_size, seq_len, _ x.shape # 1. 生成 Q、K、V并拆分成多頭 Q self.w_q(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2) K self.w_k(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2) V self.w_v(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2) # 2. 計(jì)算縮放點(diǎn)積注意力分?jǐn)?shù) scores torch.matmul(Q, K.transpose(-2, -1)) / self.scale # 3. 處理 mask將需要屏蔽的位置設(shè)為負(fù)無(wú)窮 if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) # 4. Softmax 歸一化后加權(quán)求和 attn_weights torch.softmax(scores, dim-1) attn_weights self.dropout(attn_weights) context torch.matmul(attn_weights, V) # 5. 合并多頭結(jié)果還原形狀 context context.transpose(1, 2).contiguous().view( batch_size, seq_len, self.d_model ) output self.w_o(context) return output關(guān)鍵的形狀變化如下x的形狀是(batch, seq_len, d_model)。經(jīng)過(guò)w_q(x)后仍然是(batch, seq_len, d_model)。view把它變成(batch, seq_len, n_heads, head_dim)transpose(1, 2)后變成(batch, n_heads, seq_len, head_dim)。這樣每個(gè)頭都能獨(dú)立處理一整條序列。torch.matmul(Q, K.transpose(-2, -1))得到(batch, n_heads, seq_len, seq_len)其中每個(gè)元素表示對(duì)應(yīng)位置上 Query 和 Key 的相似度。scale math.sqrt(self.head_dim)的作用非常關(guān)鍵。如果不縮放當(dāng)head_dim很大時(shí)內(nèi)積結(jié)果可能非常大Softmax 之后概率分布會(huì)變得很“尖銳”造成梯度消失影響訓(xùn)練穩(wěn)定性。所以用head_dim的平方根做一次縮放。5.5 關(guān)于 mask 的說(shuō)明在 Encoder 的文本分類場(chǎng)景中一般不需要 mask但在 Decoder 或 GPT 這類自回歸模型中需要引入因果掩碼causal mask確保當(dāng)前位置只能看到自己及之前的位置看不到未來(lái)信息。這個(gè) mask 通常是一個(gè)上三角矩陣對(duì)角線以上的位置設(shè)為False或 0對(duì)應(yīng)位置在算注意力分?jǐn)?shù)時(shí)被替換成負(fù)無(wú)窮。在本文的分類任務(wù)里我們會(huì)先不使用 mask。但代碼中已經(jīng)預(yù)留了mask參數(shù)后續(xù)你嘗試實(shí)現(xiàn) GPT 時(shí)可以直接復(fù)用。6. 層歸一化、殘差連接與前饋網(wǎng)絡(luò)6.1 殘差連接為什么不能省網(wǎng)絡(luò)層數(shù)加深之后梯度在反向傳播過(guò)程中容易消失或爆炸。殘差連接的思路非常樸素把輸入直接加到輸出上。output x F(x)這樣一來(lái)即使F(x)這一路的梯度很小輸入那一路依然可以無(wú)損地把梯度回傳模型訓(xùn)練起來(lái)會(huì)更穩(wěn)定。Transformer 在自注意力和前饋網(wǎng)絡(luò)之后都加了殘差連接。6.2 LayerNorm 與 BatchNorm 的區(qū)別LayerNorm 是 Transformer 中使用的歸一化方式。它與 BatchNorm 的區(qū)別在于歸一化的維度不同BatchNorm 對(duì)每個(gè)特征維度跨 batch 做歸一化因此依賴 batch 大小并且需要維護(hù)全局統(tǒng)計(jì)量。LayerNorm 對(duì)每個(gè)樣本的單個(gè)層內(nèi)所有神經(jīng)元做歸一化不依賴 batch 大小更容易適配變長(zhǎng)序列。在(batch, seq_len, d_model)的張量上LayerNorm 通常對(duì)最后一維d_model做歸一化也就是對(duì)每個(gè) Token 的特征向量做標(biāo)準(zhǔn)化讓均值為 0、方差為 1再進(jìn)行可學(xué)習(xí)的縮放和平移。norm nn.LayerNorm(d_model)6.3 前饋網(wǎng)絡(luò)與 Encoder 層組裝Transformer 的 Encoder 層由兩個(gè)子層組成多頭自注意力子層配合殘差連接和 LayerNorm。前饋網(wǎng)絡(luò)子層配合殘差連接和 LayerNorm。前饋網(wǎng)絡(luò)就是兩個(gè)線性層中間夾一個(gè) ReLU 激活函數(shù)。第一個(gè)線性層會(huì)把維度從d_model放大到d_ff通常d_ff是d_model的四倍第二個(gè)線性層再把它壓縮回去。先看前饋網(wǎng)絡(luò)實(shí)現(xiàn)class FeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout0.1): super().__init__() self.linear1 nn.Linear(d_model, d_ff) self.linear2 nn.Linear(d_ff, d_model) self.dropout nn.Dropout(pdropout) def forward(self, x): return self.linear2(self.dropout(torch.relu(self.linear1(x))))再組裝完整 Encoder 層class EncoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, n_heads, dropout) self.feed_forward FeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout1 nn.Dropout(pdropout) self.dropout2 nn.Dropout(pdropout) def forward(self, x, maskNone): # 第一個(gè)子層多頭自注意力 殘差 LayerNorm attn_out self.self_attn(x, mask) x self.norm1(x self.dropout1(attn_out)) # 第二個(gè)子層前饋網(wǎng)絡(luò) 殘差 LayerNorm ff_out self.feed_forward(x) x self.norm2(x self.dropout2(ff_out)) return x很多實(shí)現(xiàn)會(huì)把dropout放在子層輸出之后也有的放在殘差之前效果差異不大關(guān)鍵是理解殘差“把原輸入繞過(guò)去”這一設(shè)計(jì)思想。實(shí)際中更常見(jiàn)的寫法是x norm(x dropout(sublayer(x)))。之所以把 LayerNorm 放在殘差之后而不是之前主要是為了保證每條殘差路徑上的梯度傳播更順暢。不同實(shí)現(xiàn)順序各有理由但在實(shí)踐中這兩種排列都能工作讀者不必過(guò)于糾結(jié)。7. 手撕一個(gè)可訓(xùn)練的簡(jiǎn)化版 Transformer 分類器7.1 總體模型結(jié)構(gòu)我們現(xiàn)在把前面的組件全部拼起來(lái)得到一個(gè)完整的 Encoder-only Transformer 分類器。它的前向流程是Token ID - Embedding - 位置編碼 - N 層 Encoder - LayerNorm - 取 [CLS] 向量 - 分類輸出這里用了一個(gè)常見(jiàn)技巧在序列最前面插入一個(gè)特殊 Token[CLS]最后一層 Encoder 輸出中對(duì)應(yīng)[CLS]位置的向量就可以作為整條序列的語(yǔ)義聚合表示送入分類頭。這個(gè)技巧在 BERT 中使用得非常多。在很多實(shí)際實(shí)現(xiàn)中[CLS]是一個(gè)專門學(xué)習(xí)的向量拼在輸入序列的最前面。這里我們簡(jiǎn)化一點(diǎn)因?yàn)榱钆菩蛄斜旧砭褪请S機(jī)數(shù)我們不需要單獨(dú)構(gòu)造特殊 Token直接約定“取第 0 個(gè)位置的向量做分類”即可。這樣能讓代碼保持簡(jiǎn)潔同時(shí)不改變整體思路。完整模型代碼如下import torch import torch.nn as nn class TransformerClassifier(nn.Module): def __init__(self, vocab_size, d_model, n_heads, d_ff, n_layers, max_len, num_classes, dropout0.1): super().__init__() self.embedding nn.Embedding(vocab_size, d_model) self.positional_encoding PositionalEncoding(d_model, max_len, dropout) self.layers nn.ModuleList([ EncoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_layers) ]) self.norm nn.LayerNorm(d_model) self.classifier nn.Linear(d_model, num_classes) def forward(self, tokens, maskNone): x self.embedding(tokens) x self.positional_encoding(x) for layer in self.layers: x layer(x, mask) x self.norm(x) # 取第 0 個(gè)位置的向量作為整條序列的表示 cls_repr x[:, 0, :] logits self.classifier(cls_repr) return logits7.2 構(gòu)造模擬訓(xùn)練數(shù)據(jù)因?yàn)楸疚闹攸c(diǎn)是模型實(shí)現(xiàn)這里用隨機(jī)生成的 Token 和標(biāo)簽?zāi)M一個(gè)分類數(shù)據(jù)集。batch_size表示每個(gè)批次的樣本數(shù)seq_len表示每條文本的長(zhǎng)度。def make_batch(batch_size, seq_len, vocab_size, num_classes): tokens torch.randint(0, vocab_size, (batch_size, seq_len)) labels torch.randint(0, num_classes, (batch_size,)) return tokens, labels你需要理解的是真實(shí)場(chǎng)景中 tokens 應(yīng)該來(lái)自 Tokenizer 對(duì)原始文本的編碼labels 來(lái)自人工標(biāo)注。上述代碼只是為了快速驗(yàn)證模型前向和訓(xùn)練流程是否正常。7.3 初始化模型并設(shè)置訓(xùn)練參數(shù)我們選擇一個(gè)偏小的模型配置這樣在 CPU 上也能快速完成訓(xùn)練演示model TransformerClassifier( vocab_size1000, d_model128, n_heads4, d_ff512, n_layers2, max_len64, num_classes2, dropout0.1 ) # 用一個(gè)很小的輸入測(cè)試前向 tokens, labels make_batch(batch_size4, seq_len32, vocab_size1000, num_classes2) logits model(tokens) print(logits.shape) # torch.Size([4, 2])如果輸出形狀為torch.Size([4, 2])說(shuō)明模型前向流程基本正確。接下來(lái)就可以進(jìn)入訓(xùn)練環(huán)節(jié)。7.4 訓(xùn)練循環(huán)訓(xùn)練過(guò)程采用標(biāo)準(zhǔn)的交叉熵?fù)p失和 Adam 優(yōu)化器。這里我們只迭代 20 輪每輪內(nèi)部隨機(jī)生成 100 個(gè)批次用來(lái)觀察 loss 是否在下降。import torch.optim as optim criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3) model.train() for epoch in range(20): total_loss 0.0 for _ in range(100): tokens, labels make_batch( batch_size32, seq_len32, vocab_size1000, num_classes2 ) optimizer.zero_grad() logits model(tokens) loss criterion(logits, labels) loss.backward() optimizer.step() total_loss loss.item() avg_loss total_loss / 100 print(fEpoch {epoch 1:02d}, Loss: {avg_loss:.4f})由于標(biāo)簽是隨機(jī)生成的模型不可能學(xué)到真正的規(guī)律loss 不會(huì)降到特別低但應(yīng)該會(huì)從大約 0.7 附近開(kāi)始下降。如果你把某一類樣本的比例固定下來(lái)或者構(gòu)造一個(gè)線性可分的規(guī)則數(shù)據(jù)集就能看到更明顯的訓(xùn)練效果。7.5 增加一個(gè)可學(xué)習(xí)的簡(jiǎn)單分類任務(wù)為了讓結(jié)果更具說(shuō)服力我們可以構(gòu)造一個(gè)非常簡(jiǎn)單的規(guī)則任務(wù)當(dāng)每條 Token 序列中出現(xiàn)某個(gè)特殊 ID 時(shí)標(biāo)簽為 1否則為 0。這樣模型確實(shí)能學(xué)到一些模式loss 會(huì)明顯下降。def make_rule_batch(batch_size, seq_len, vocab_size): special_id 99 tokens torch.randint(0, vocab_size, (batch_size, seq_len)) # 如果序列中出現(xiàn) special_id標(biāo)簽為 1 labels (tokens special_id).any(dim1).long() return tokens, labels替換訓(xùn)練循環(huán)里的make_batch為make_rule_batch再運(yùn)行同樣的訓(xùn)練代碼loss 會(huì)穩(wěn)定下降到較低水平。這說(shuō)明 Transformer 確實(shí)從數(shù)據(jù)中捕捉到了“某個(gè) Token 出現(xiàn)與否”這個(gè)特征并把它編碼到了[CLS]位置的表示中。7.6 完整代碼文件建議為了方便你保存和運(yùn)行建議把代碼按下面方式組織transformer-demo/ |-- model.py # 模型組件與 TransformerClassifier |-- train.py # 數(shù)據(jù)生成與訓(xùn)練循環(huán)model.py中存放PositionalEncoding、MultiHeadAttention、FeedForward、EncoderLayer、TransformerClassifier。train.py中存放make_batch、make_rule_batch和訓(xùn)練邏輯。8. 常見(jiàn)報(bào)錯(cuò)與排查思路手撕 Transformer 的過(guò)程中最常見(jiàn)的報(bào)錯(cuò)往往集中在形狀不匹配上。下面整理了幾個(gè)高頻問(wèn)題。問(wèn)題現(xiàn)象常見(jiàn)原因解決思路注意力機(jī)制報(bào)錯(cuò)說(shuō)維度無(wú)法相乘Q、K、V 的 shape 沒(méi)有對(duì)齊檢查view和transpose后的形狀重點(diǎn)確認(rèn)(batch, n_heads, seq_len, head_dim)是否正確訓(xùn)練時(shí) loss 不下降學(xué)習(xí)率過(guò)高或過(guò)低數(shù)據(jù)無(wú)規(guī)律先在小數(shù)據(jù)上跑通嘗試調(diào)整學(xué)習(xí)率觀察梯度是否正常位置編碼越界輸入序列長(zhǎng)度大于max_len把max_len設(shè)置得比訓(xùn)練/推理序列更長(zhǎng)或者在 forward 中限制長(zhǎng)度CUDA 內(nèi)存不足batch 太大或序列太長(zhǎng)減小 batch、減小d_model、使用混合精度訓(xùn)練mask 導(dǎo)致全為負(fù)無(wú)窮mask 形狀或數(shù)值邏輯錯(cuò)誤打印 mask 和 scores 的 shape確認(rèn)需要屏蔽的位置是否被設(shè)置為 0d_model無(wú)法被n_heads整除多頭注意力要求整除使用 assert 提前檢查或選擇合適配置排查這類問(wèn)題有一個(gè)通用技巧不要在報(bào)錯(cuò)后盲目改代碼先在關(guān)鍵張量之間插入print(shape)看每一層的輸出形狀是否符合預(yù)期。尤其是view、transpose、matmul這幾個(gè)操作是最容易出現(xiàn)形狀問(wèn)題的地方。另一個(gè)經(jīng)常被忽略的問(wèn)題是 device 不匹配。如果你使用 GPU 訓(xùn)練記得在模型和數(shù)據(jù)上都調(diào)用.to(device)device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) tokens tokens.to(device) labels labels.to(device)如果模型參數(shù)在 GPU 上、數(shù)據(jù)在 CPU 上PyTorch 會(huì)拋出類似Expected all tensors to be on the same device的報(bào)錯(cuò)。9. 工程實(shí)踐與后續(xù)學(xué)習(xí)建議9.1 讓代碼更接近真實(shí)項(xiàng)目的幾個(gè)改進(jìn)點(diǎn)本文為了把原理講清楚做了大量簡(jiǎn)化。實(shí)際工程中你還需要關(guān)注以下幾個(gè)方面詞表與分詞器。使用 BPE、WordPiece 或 SentencePiece 訓(xùn)練自己的分詞器而不是按字隨機(jī)切分。預(yù)訓(xùn)練好的 Tokenizer 通常更符合真實(shí)文本分布。Padding 與 Attention Mask。一個(gè) batch 里文本長(zhǎng)度不同需要把短序列填充到相同長(zhǎng)度并用 mask 讓模型忽略 padding 位置否則填充位會(huì)影響注意力計(jì)算。學(xué)習(xí)率調(diào)度。Transformer 訓(xùn)練通常配合 Warmup 和余弦退火常見(jiàn)做法是先用較小學(xué)習(xí)率熱身再逐步衰減。這里推薦使用 PyTorch 的torch.optim.lr_scheduler實(shí)現(xiàn)。梯度裁剪。對(duì)于序列模型clip_grad_norm_能有效防止梯度爆炸是訓(xùn)練穩(wěn)定性的重要保障。混合精度訓(xùn)練。用torch.cuda.amp可以顯著減少顯存占用并提升訓(xùn)練速度特別是模型較大時(shí)收益明顯。實(shí)驗(yàn)記錄。保存模型權(quán)重、記錄每次實(shí)驗(yàn)的超參數(shù)和 loss 曲線。可以使用tensorboard或簡(jiǎn)單的 JSON 文件記錄。9.2 從 Encoder 到 Encoder-Decoder 再到 GPT掌握本文的 Encoder 代碼后你的下一步可以按下面順序推進(jìn)實(shí)現(xiàn)完整的Transformer Encoder-Decoder把 Decoder 里的 Cross-Attention 加上用于機(jī)器翻譯任務(wù)。實(shí)現(xiàn)因果自注意力。只需要把 mask 換成上三角掩碼把[CLS]分類頭改成逐 Token 的預(yù)測(cè)頭就是一個(gè)簡(jiǎn)化版 GPT。使用 HuggingFacetransformers庫(kù)加載一個(gè)預(yù)訓(xùn)練模型對(duì)比自己的實(shí)現(xiàn)學(xué)習(xí)工業(yè)級(jí)代碼的組織方式。嘗試把 Transformer 用到圖像上。ViT 把圖片切成 Patch 后當(dāng)作 Token 序列幾乎是復(fù)用同一套代碼。9.3 幾個(gè)容易踩的坑不要一上來(lái)就追求大模型。先在小規(guī)模數(shù)據(jù)上跑通流程再逐步擴(kuò)大。不要隨意調(diào)整d_model和n_heads記得d_model必須能被n_heads整除。不要忽略數(shù)據(jù)順序。位置編碼是必須的不要因?yàn)闇y(cè)試時(shí)模型能跑就跳過(guò)去。不要把 dropout 設(shè)得過(guò)大。Transformer 對(duì) dropout 比較敏感小模型尤其如此一般 0.1 左右比較穩(wěn)妥。10. 總結(jié)本文從零開(kāi)始實(shí)現(xiàn)了 Transformer 的核心組件Token Embedding、正余弦位置編碼、多頭自注意力、前饋網(wǎng)絡(luò)、殘差連接與層歸一化并組裝成了一個(gè)可用于文本分類的 Transformer 編碼器。你可以跟著代碼在本地跑通一次完整的訓(xùn)練流程觀察 loss 的變化再逐步加入 mask、真實(shí)分詞器、學(xué)習(xí)率調(diào)度等工程細(xì)節(jié)。手撕一遍之后再去看 BERT、GPT 或者 ViT 的源碼你會(huì)發(fā)現(xiàn)自己對(duì)“自注意力如何并行建模全局依賴”已經(jīng)有了直觀認(rèn)識(shí)。剩下的就是在小數(shù)據(jù)集上多做實(shí)驗(yàn)親手調(diào)一調(diào)參數(shù)把每個(gè)張量的形狀變化印在腦子里。代碼寫熟之后Transformer 就不再是紙面上的公式而是一個(gè)隨時(shí)可以改動(dòng)、可以調(diào)試、可以擴(kuò)展的工具。