手撕Transformer！！從每一模塊原理講解到代碼實現【超詳細！】

一、位置編碼

位置編碼（Positional Encoding）是Transformer模型中的一個重要組成部分，用于在序列數據中引入位置信息。由于Transformer模型本身不具備處理序列順序的能力（因為它完全基于自注意力機制，沒有遞歸或卷積結構），位置編碼的引入使得模型能夠利用序列的順序信息。

位置編碼的原理

位置編碼通過在輸入嵌入向量中添加一個與位置相關的向量來實現。具體來說，對于每個位置 ( pos ) 和每個維度 ( i )，位置編碼向量 ( PE(pos, 2i) ) 和 ( PE(pos, 2i+1) ) 分別由以下公式計算：

代碼解釋

以下是 PositionalEncoder 類的詳細解釋：

import torch
import torch.nn as nn
import math

class PositionalEncoder(nn.Module):

    def __init__(self, d_model, max_seq_len=80):
        super().__init__()
        self.d_model = d_model
        # 創建一個常量 PE 矩陣
        pe = torch.zeros(max_seq_len, d_model)
        for pos in range(max_seq_len):
            for i in range(0, d_model, 2):
                pe[pos, i] = math.sin(pos / (10000**((2 * i) / d_model)))
                pe[pos, i + 1] = math.cos(pos / (10000**((2 * (i + 1)) / d_model)))
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        # 使得單詞嵌入表示相對大一些
        x = x * math.sqrt(self.d_model)
        # 增加位置常量到單詞嵌入表示中
        seq_len = x.size(1)
        x = x + self.pe[:, :seq_len]
        return x

1. 初始化：

   • d_model：模型的維度。
   • max_seq_len：序列的最大長度。
   • pe：一個大小為 (max_seq_len, d_model) 的零矩陣，用于存儲位置編碼。

2. 計算位置編碼：

   • 對于每個位置 pos 和每個維度 i，計算 sin 和 cos 值，并將其存儲在 pe 矩陣中。
   • pe 矩陣通過 unsqueeze(0) 增加一個批次維度，使其形狀為 (1, max_seq_len, d_model)。

3. 注冊緩沖區：

   • self.register_buffer('pe', pe)：將 pe 注冊為一個緩沖區，這樣它會在模型保存和加載時被保存，但不會被優化器更新。

4. 前向傳播：

   • x = x * math.sqrt(self.d_model)：將輸入嵌入向量 x 放大，以確保嵌入向量的值不會被位置編碼淹沒。
   • x = x + self.pe[:, :seq_len]：將位置編碼添加到輸入嵌入向量中，其中 seq_len 是輸入序列的實際長度。

二、多頭注意力

多頭注意力機制（Multi-Head Attention）是Transformer模型中的一個關鍵組件，用于處理序列數據，特別是在自然語言處理任務中。它的主要思想是將輸入的查詢（Query）、鍵（Key）和值（Value）通過多個獨立的注意力頭（Attention Heads）進行處理，然后將這些頭的輸出拼接起來并通過一個線性層進行整合。這種機制可以捕捉序列中不同位置的多種復雜關系。

以下是對多頭注意力機制的詳細解釋：

1. 初始化：

   • d_model：輸入和輸出的維度。
   • heads：注意力頭的數量。
   • d_k：每個注意力頭的維度，計算方式為 d_model // heads。
   • 線性層：用于將輸入的查詢、鍵和值分別映射到 d_model 維度。
   • 丟棄層（Dropout）：用于防止過擬合。
   • 輸出線性層：用于將拼接后的多頭注意力輸出映射回 d_model 維度。

2. 注意力計算：

   • attention 方法計算注意力分數。首先，通過矩陣乘法計算查詢和鍵的點積，然后除以 sqrt(d_k) 進行縮放，以防止梯度消失或爆炸。
   • 如果提供了掩碼（mask），則將掩碼中為0的位置對應的分數設置為一個非常小的值（如 -1e9），以確保這些位置在softmax后為0。
   • 對分數進行softmax操作，使其成為一個概率分布。
   • 應用丟棄層（Dropout）。
   • 通過矩陣乘法將注意力分數與值相乘，得到加權的值。

3. 前向傳播：

   • 對輸入的查詢、鍵和值分別進行線性變換，然后重塑為多頭形式。
   • 將這些張量進行轉置，以便在注意力計算中正確對齊。
   • 調用 attention 方法計算多頭注意力。
   • 將多頭注意力的輸出進行轉置和拼接，然后通過輸出線性層進行整合。

以下是完整的代碼實現：

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class MultiHeadAttention(nn.Module):

    def __init__(self, heads, d_model, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.d_k = d_model // heads
        self.h = heads
        self.q_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)
        self.out = nn.Linear(d_model, d_model)

    def attention(self, q, k, v, d_k, mask=None, dropout=None):
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
        if mask is not None:
            mask = mask.unsqueeze(1)
            scores = scores.masked_fill(mask == 0, -1e9)
        scores = F.softmax(scores, dim=-1)
        if dropout is not None:
            scores = dropout(scores)
        output = torch.matmul(scores, v)
        return output

    def forward(self, q, k, v, mask=None):
        bs = q.size(0)
        k = self.k_linear(k).view(bs, -1, self.h, self.d_k).transpose(1, 2)
        q = self.q_linear(q).view(bs, -1, self.h, self.d_k).transpose(1, 2)
        v = self.v_linear(v).view(bs, -1, self.h, self.d_k).transpose(1, 2)
        scores = self.attention(q, k, v, self.d_k, mask, self.dropout)
        concat = scores.transpose(1, 2).contiguous().view(bs, -1, self.d_model)
        output = self.out(concat)
        return output

轉置操作 .transpose(1, 2) 是為了在多頭注意力計算中正確對齊每個頭的查詢、鍵和值，指的是，矩陣計算在sequence_length, d_k這兩個維度上進行。

三、前饋神經網絡（FeedForward）和層歸一化（NormLayer）

FeedForward 模塊

FeedForward 模塊是一個簡單的前饋神經網絡，通常緊跟在多頭注意力機制之后。它由兩個線性層和一個激活函數組成，中間包含一個丟棄層（Dropout）以防止過擬合。

代碼解析

class FeedForward(nn.Module):

    def __init__(self, d_model, d_ff=2048, dropout=0.1):
        super().__init__()
        # d_ff 默認設置為 2048
        self.linear_1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.linear_2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        x = self.dropout(F.relu(self.linear_1(x)))
        x = self.linear_2(x)
        return x

• 初始化：

  • d_model：輸入和輸出的維度。
  • d_ff：中間層的維度，默認設置為 2048。
  • dropout：丟棄層的丟棄率，默認設置為 0.1。
  • self.linear_1：第一個線性層，將輸入從 d_model 維度映射到 d_ff 維度。
  • self.dropout：丟棄層，用于防止過擬合。
  • self.linear_2：第二個線性層，將輸入從 d_ff 維度映射回 d_model 維度。

• 前向傳播：

  • self.linear_1(x)：將輸入 x 從 d_model 維度映射到 d_ff 維度。
  • F.relu(self.linear_1(x))：應用 ReLU 激活函數。
  • self.dropout(F.relu(self.linear_1(x)))：應用丟棄層。
  • self.linear_2(x)：將輸入從 d_ff 維度映射回 d_model 維度。

NormLayer 模塊

NormLayer 模塊是一個層歸一化層，用于對輸入進行歸一化處理。層歸一化通過對每個樣本的所有特征進行歸一化，使得每個樣本的特征具有相同的均值和標準差。

代碼解析

class NormLayer(nn.Module):

    def __init__(self, d_model, eps=1e-6):
        super().__init__()
        self.size = d_model
        # 層歸一化包含兩個可以學習的參數
        self.alpha = nn.Parameter(torch.ones(self.size))
        self.bias = nn.Parameter(torch.zeros(self.size))
        self.eps = eps

    def forward(self, x):
        norm = self.alpha * (x - x.mean(dim=-1, keepdim=True))         / (x.std(dim=-1, keepdim=True) + self.eps) + self.bias
        return norm

• 初始化：

  • d_model：輸入和輸出的維度。
  • eps：一個很小的數，用于防止除零錯誤，默認設置為 1e-6。
  • self.alpha：一個可學習的縮放參數，初始化為全1。
  • self.bias：一個可學習的偏移參數，初始化為全0。

• 前向傳播：

  • x.mean(dim=-1, keepdim=True)：計算輸入 x 在最后一個維度上的均值。
  • x.std(dim=-1, keepdim=True)：計算輸入 x 在最后一個維度上的標準差。
  • (x - x.mean(dim=-1, keepdim=True)) / (x.std(dim=-1, keepdim=True) + self.eps)：對輸入 x 進行歸一化處理。
  • self.alpha * ... + self.bias：應用可學習的縮放和偏移參數。

這兩個模塊在Transformer模型中通常一起使用，FeedForward 模塊用于增加模型的非線性能力，而 NormLayer 模塊用于穩定訓練過程和加速收斂。

四、Encoder

這里的 Encoder 類是 Transformer 模型中的編碼器部分。編碼器的主要作用是將輸入序列（例如一段文本）轉換成一系列高維特征向量，這些特征向量可以被解碼器用來生成輸出序列。下面是對 Encoder 類及其組成部分的詳細解釋：

Encoder 類

Encoder 類是整個編碼器的主要結構，它包含了以下幾個部分：

1. 嵌入層 (self.embed)：

   • 將輸入的詞匯索引序列（src）轉換為對應的詞嵌入向量。每個詞匯索引對應一個 d_model 維的向量。

2. 位置編碼器 (self.pe)：

   • 由于 Transformer 模型沒有遞歸和卷積結構，無法自然地利用序列的順序信息。位置編碼器通過在詞嵌入向量中添加位置信息來解決這個問題。位置編碼可以是固定的（如正弦和余弦函數），也可以是可學習的。

3. 編碼器層 (self.layers)：

   • 這是一個由 N 個 EncoderLayer 組成的列表。每個 EncoderLayer 包含一個多頭注意力機制和一個前饋神經網絡，以及相應的歸一化層和丟棄層。

4. 歸一化層 (self.norm)：

   • 在所有編碼器層之后，對輸出進行層歸一化，以穩定訓練過程。

class Encoder(nn.Module):

    def __init__(self, vocab_size, d_model, N, heads, dropout):
        super().__init__()
        self.N = N
        self.embed = nn.Embedding(vocab_size, d_model)
        self.pe = PositionalEncoder(d_model)
        self.layers = nn.ModuleList([EncoderLayer(d_model, heads, dropout) for _ in range(N)])
        self.norm = NormLayer(d_model)

    def forward(self, src, mask):
        x = self.embed(src)
        x = self.pe(x)
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

EncoderLayer 類

EncoderLayer 類是編碼器的基本組成單元，每個 EncoderLayer 包含以下幾個部分：

1. 歸一化層 (self.norm_1 和 self.norm_2)：

   • 在多頭注意力機制和前饋神經網絡之前，對輸入進行層歸一化。

2. 多頭注意力機制 (self.attn)：

   • 計算輸入序列的自注意力表示。自注意力機制允許模型在處理每個位置的輸入時，考慮到序列中所有其他位置的信息。

3. 前饋神經網絡 (self.ff)：

   • 一個簡單的兩層全連接神經網絡，用于對每個位置的輸入進行非線性變換。

4. 丟棄層 (self.dropout_1 和 self.dropout_2)：

   • 在多頭注意力機制和前饋神經網絡的輸出上應用丟棄操作，以防止過擬合。

class EncoderLayer(nn.Module):

    def __init__(self, d_model, heads, dropout=0.1):
        super().__init__()
        self.norm_1 = NormLayer(d_model)
        self.norm_2 = NormLayer(d_model)
        self.attn = MultiHeadAttention(heads, d_model, dropout=dropout)
        self.ff = FeedForward(d_model, dropout=dropout)
        self.dropout_1 = nn.Dropout(dropout)
        self.dropout_2 = nn.Dropout(dropout)

    def forward(self, x, mask):
        x2 = self.norm_1(x)
        x = x + self.dropout_1(self.attn(x2, x2, x2, mask))
        x2 = self.norm_2(x)
        x = x + self.dropout_2(self.ff(x2))
        return x

前向傳播過程

1. 嵌入和位置編碼：

   • 輸入序列 src 首先通過嵌入層轉換為詞嵌入向量，然后通過位置編碼器添加位置信息。

2. 編碼器層處理：

   • 將添加了位置信息的詞嵌入向量輸入到第一個編碼器層。每個編碼器層的輸出作為下一個編碼器層的輸入，依次經過所有 N 個編碼器層。

3. 歸一化：

   • 在所有編碼器層處理完畢后，對最終的輸出進行層歸一化。

五、Decoder

這個 Decoder 類是 Transformer 模型中的解碼器部分。解碼器的主要作用是生成輸出序列，例如在機器翻譯任務中，解碼器負責生成目標語言的句子。下面是對 Decoder 類及其組成部分的詳細解釋：

Decoder 類

Decoder 類是整個解碼器的主要結構，它包含了以下幾個部分：

1. 嵌入層 (self.embed)：

   • 將輸入的目標語言詞匯索引序列（trg）轉換為對應的詞嵌入向量。每個詞匯索引對應一個 d_model 維的向量。

2. 位置編碼器 (self.pe)：

   • 由于 Transformer 模型沒有遞歸和卷積結構，無法自然地利用序列的順序信息。位置編碼器通過在詞嵌入向量中添加位置信息來解決這個問題。位置編碼可以是固定的（如正弦和余弦函數），也可以是可學習的。

3. 解碼器層 (self.layers)：

   • 這是一個由 N 個 DecoderLayer 組成的列表。每個 DecoderLayer 包含兩個多頭注意力機制和一個前饋神經網絡，以及相應的歸一化層和丟棄層。

4. 歸一化層 (self.norm)：

   • 在所有解碼器層之后，對輸出進行層歸一化，以穩定訓練過程。

class Decoder(nn.Module):

    def __init__(self, vocab_size, d_model, N, heads, dropout):
        super().__init__()
        self.N = N
        self.embed = nn.Embedding(vocab_size, d_model)
        self.pe = PositionalEncoder(d_model)
        self.layers = nn.ModuleList([DecoderLayer(d_model, heads, dropout) for _ in range(N)])
        self.norm = NormLayer(d_model)

    def forward(self, trg, e_outputs, src_mask, trg_mask):
        x = self.embed(trg)
        x = self.pe(x)
        for layer in self.layers:
            x = layer(x, e_outputs, src_mask, trg_mask)
        return self.norm(x)

DecoderLayer 類

DecoderLayer 類是解碼器的基本組成單元，每個 DecoderLayer 包含以下幾個部分：

1. 歸一化層 (self.norm_1, self.norm_2, self.norm_3)：

   • 在多頭注意力機制和前饋神經網絡之前，對輸入進行層歸一化。

2. 丟棄層 (self.dropout_1, self.dropout_2, self.dropout_3)：

   • 在多頭注意力機制和前饋神經網絡的輸出上應用丟棄操作，以防止過擬合。

3. 多頭注意力機制 (self.attn_1, self.attn_2)：

   • self.attn_1 是自注意力機制，計算輸入序列的自注意力表示。自注意力機制允許模型在處理每個位置的輸入時，考慮到序列中所有其他位置的信息。
   • self.attn_2 是編碼器-解碼器注意力機制，允許解碼器在生成每個位置的輸出時，考慮到編碼器的輸出（即源語言的上下文信息）。

4. 前饋神經網絡 (self.ff)：

   • 一個簡單的兩層全連接神經網絡，用于對每個位置的輸入進行非線性變換。

class DecoderLayer(nn.Module):

    def __init__(self, d_model, heads, dropout=0.1):
        super().__init__()
        self.norm_1 = NormLayer(d_model)
        self.norm_2 = NormLayer(d_model)
        self.norm_3 = NormLayer(d_model)
        self.dropout_1 = nn.Dropout(dropout)
        self.dropout_2 = nn.Dropout(dropout)
        self.dropout_3 = nn.Dropout(dropout)
        self.attn_1 = MultiHeadAttention(heads, d_model, dropout=dropout)
        self.attn_2 = MultiHeadAttention(heads, d_model, dropout=dropout)
        self.ff = FeedForward(d_model, dropout=dropout)

    def forward(self, x, e_outputs, src_mask, trg_mask):
        x2 = self.norm_1(x)
        x = x + self.dropout_1(self.attn_1(x2, x2, x2, trg_mask))
        x2 = self.norm_2(x)
        x = x + self.dropout_2(self.attn_2(x2, e_outputs, e_outputs, src_mask))
        x2 = self.norm_3(x)
        x = x + self.dropout_3(self.ff(x2))
        return x

前向傳播過程

1. 嵌入和位置編碼：

   • 輸入序列 trg 首先通過嵌入層轉換為詞嵌入向量，然后通過位置編碼器添加位置信息。

2. 解碼器層處理：

   • 將添加了位置信息的詞嵌入向量輸入到第一個解碼器層。每個解碼器層的輸出作為下一個解碼器層的輸入，依次經過所有 N 個解碼器層。
   • 在每個解碼器層中，首先進行自注意力機制計算，然后進行編碼器-解碼器注意力機制計算，最后進行前饋神經網絡計算。

3. 歸一化：

   • 在所有解碼器層處理完畢后，對最終的輸出進行層歸一化。

六、Transformer整體框架

class Transformer(nn.Module):

    def __init__(self, src_vocab, trg_vocab, d_model, N, heads, dropout):
        super().__init__()
        self.encoder = Encoder(src_vocab, d_model, N, heads, dropout)
        self.decoder = Decoder(trg_vocab, d_model, N, heads, dropout)
        self.out = nn.Linear(d_model, trg_vocab)

    def forward(self, src, trg, src_mask, trg_mask):
        e_outputs = self.encoder(src, src_mask)
        d_output = self.decoder(trg, e_outputs, src_mask, trg_mask)
        output = self.out(d_output)
        return output

參考資料

https://zhuanlan.zhihu.com/p/657456977
《Attention is all you need》

版權聲明
本博客內容僅供學習交流，轉載請注明出處。