使用 PyTorch 从零开始实现 Transformer

开始

在今天的博客中，我们将深入了解变换器架构。变换器通过引入一种新颖的机制来捕捉序列中的依赖关系，彻底改变了自然语言处理（NLP）领域。这种机制就是注意力机制。让我们逐步解析，并使用 PyTorch 从头开始实现它。

import torch
import torch.nn as nn
import math

• torch：主要的 PyTorch 库。
• torch.nn：提供神经网络组件。
• math：提供数学函数。

输入嵌入（Input Embedding）

它允许将原始句子转换为 X 维度的向量（原始 Transformer 模型在基础版本中使用 512 作为维度大小（d_model），而在较大版本中使用 d_model = 1024）。

__init__()  方法的主要目的包括：

1. 初始化对象的状态（即，为对象的属性设置初始值）。
2. 定义神经网络模块将使用的层和组件。
3. 确保在创建对象时执行任何必要的设置或初始化代码。

super()
super() 函数用于调用父类的方法。

super() 返回一个临时的父类对象，使您能够调用其方法。
在 super().init() 的情况下，它调用父类（nn.Module）的 init 方法。
在继承的上下文中，使用 super() 尤其重要，因为它确保基类的初始化代码运行，从而设置子类可能依赖的任何必要内部结构和属性。

class InputEmbeddings(nn.Module):

    def __init__(self, d_model: int, vocab_size: int) -> None:
        super().__init__()
        self.d_model = d_model
        self.vocab_size = vocab_size
        self.embedding = nn.Embedding(vocab_size, d_model)

    def forward(self, x):
        # (batch, seq_len) --> (batch, seq_len, d_model)
        # 根据论文的要求，将嵌入乘以 sqrt(d_model) 进行缩放
        return self.embedding(x) * math.sqrt(self.d_model)

代码解释

nn.Embedding(vocab_size, d_model): 这创建了一个嵌入层，将索引（通常代表单词）映射到一个 d_model 维度的向量。嵌入层是随机初始化的，这些向量在训练过程中被学习。因此，给定一个数字，它每次都会提供相同的向量。要了解更多，请参考此链接。

self.embedding(x): 这里，x 是一个标记索引的张量。嵌入层查找 x 中每个标记索引的向量。例如，如果 x 是 [0, 1, 2]，它会查找索引 0、1 和 2 的向量。

* math.sqrt(self.d_model): 这将嵌入按 d_model 的平方根进行缩放。这通常是为了在嵌入通过网络传递时保持方差，帮助提高训练的稳定性。要阅读更多，请参考论文的第 3.4 节。

位置编码（PositionalEncoding 类）

位置编码是变换器模型中的一个关键组成部分，它帮助模型理解句子中每个单词的位置。由于变换器并不像 RNN（递归神经网络）那样固有地以顺序方式处理标记，因此需要一种方法来结合标记的顺序。这是通过位置编码实现的，位置编码是添加到单词嵌入中的向量。

位置编码的重要性

变换器同时处理整个标记序列，这使得高效的并行计算成为可能，但也意味着它们缺乏关于标记顺序的信息。位置编码提供了一种引入这种顺序信息的方法。

位置编码的工作原理

嵌入大小

位置编码向量的大小与词嵌入相同，通常为512维（在原始变换器模型的情况下）。这确保了位置信息可以无缝地添加到嵌入中，而不会改变其维度。

创建位置编码

位置编码向量旨在表示序列中每个单词的位置。这些向量是使用不同频率的正弦和余弦函数的组合创建的。这使得模型能够以平滑和连续的方式学习和区分不同的位置。

添加位置编码

一旦生成了位置编码向量，就将它们逐元素地添加到相应的词嵌入中。这种组合表示同时包含了来自词嵌入的语义信息和来自位置编码的位置信息。

数学方程

对于给定位置 pos 和嵌入维度 i

在位置编码中，每个位置 pos 和维度 i 的计算方式如下：

• 正弦函数和余弦函数的使用：对于偶数维度（i 是偶数），使用正弦函数：

• 对于奇数维度（i 是奇数），使用余弦函数：

这里，d_model 是嵌入的维度，pos 是词在序列中的位置，i 是嵌入维度的索引。

位置编码的特征：

• 这些函数的选择是为了确保每个位置的编码是唯一的，并且具有平滑的变化。这种设计使得模型能够有效地捕捉到相对和绝对位置的信息。
• 随着 i 的增加，频率会逐渐减小，从而为模型提供了多种尺度的位置信息。

示例

假设我们有一个序列中的某个位置 pos = 2，嵌入维度 d_model = 4，那么我们可以计算：

通过这种方法，我们为每个位置生成独特的编码，以便在输入到变换器模型时保留位置信息。

class PositionalEncoding(nn.Module):

    def __init__(self, d_model: int, seq_len: int, dropout: float) -> None:
        super().__init__()
        self.d_model = d_model
        self.seq_len = seq_len
        self.dropout = nn.Dropout(dropout)
        # Create a matrix of shape (seq_len, d_model)
        pe = torch.zeros(seq_len, d_model)
        # Create a vector of shape (seq_len)
        position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1) # (seq_len, 1)
        # Create a vector of shape (d_model)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) # (d_model / 2)
        # Apply sine to even indices
        pe[:, 0::2] = torch.sin(position * div_term) # sin(position * (10000 ** (2i / d_model))
        # Apply cosine to odd indices
        pe[:, 1::2] = torch.cos(position * div_term) # cos(position * (10000 ** (2i / d_model))
        # Add a batch dimension to the positional encoding
        pe = pe.unsqueeze(0) # (1, seq_len, d_model)
        # Register the positional encoding as a buffer
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False) # (batch, seq_len, d_model)
        return self.dropout(x)

__init__ 方法：

• 初始化参数：
• d_model（嵌入维度）、seq_len（最大序列长度）和 dropout（丢弃率）进行初始化。
• 创建位置编码矩阵：
• 创建一个形状为 (seq_len, d_model) 的矩阵，然后计算 position（分子）和 div_term（分母），我们在对数空间中计算这些以提高数值稳定性。详细信息请参考此链接。
• 生成位置编码：
• 使用正弦和余弦函数创建一个位置编码矩阵 pe，以编码位置信息。在每个词的嵌入中，偶数位置使用正弦函数，而奇数位置使用余弦函数。
• 添加批量维度：
• 由于会有多个句子的批处理，因此需要将批量维度添加到张量中（pe = pe.unsqueeze(0)），形状将从 (seq_len, d_model) 变为 (1, seq_len, d_model)。
• 使用 register_buffer：
• 使用 register_buffer 存储 pe，而不将其设置为可学习参数。

forward 方法：

• 添加位置编码：
• 将位置编码添加到输入嵌入 x，并应用丢弃。我们还告诉模型不要学习这些位置编码，因为它们是固定的，并且将始终保持不变，使用 requires_grad_(False) 来实现。

LayerNormalization 类

层归一化是一种用于改善深度神经网络训练的技术，通过对每个训练样本的特征进行归一化，帮助稳定学习过程，提高收敛性，并减少对参数精确初始化的依赖。

理解输入 x

• 词嵌入：在变换器模型的上下文中，句子中的每个单词或标记都被转换为一个固定大小的密集向量（嵌入）（例如，512 维度）。这些嵌入捕捉了单词的语义信息。
• 句子表示：句子表示为这些词嵌入的序列。这些向量是多维表示，捕捉单词含义和句法角色的各个方面。
• 特征维度：词嵌入的维度（例如，512 维度）称为“特征”。每个特征对应于单词表示的特定方面。这些特征可以包括句法角色、语义意义、上下文用法等。

应用层归一化

在变换器模型中，应用层归一化时，输入 x 通常是一个形状为 (batch_size, seq_len, d_model) 的 3 维张量，其中：

• batch_size：批次中的句子或序列数量。
• seq_len：每个句子或序列中的标记数量。
• d_model：词嵌入的维度（例如，512）。

示例

import torch

# 创建形状为 (batch_size, seq_len, d_model) 的虚拟输入
batch_size = 32   # 批次大小
seq_len = 10      # 每个句子的标记数量
d_model = 512     # 每个标记的嵌入维度

# 使用随机数生成输入数据
input_data = torch.randn(batch_size, seq_len, d_model)

print(input_data.shape)  # 输出: torch.Size([32, 10, 512])

输入张量的解释

• 输入张量 xxx 的每个元素对应于句子中一个单词的嵌入。
• 层归一化将在最后一个维度（嵌入维度，或称为特征）上应用于每个单词嵌入。

特征上的归一化

对于每个标记嵌入（大小为 dmodeld_{model}dmodel​ 的向量），归一化过程如下：

• 计算均值和方差：对于每个标记的嵌入，计算其特征的均值（mean）和方差（variance）。公式：
•  
• 归一化嵌入：通过从嵌入中减去均值并除以标准差来归一化：公式： '
• 
• 其中，ϵ\epsilonϵ 是一个小常数，防止除以零。

独立处理：此过程对每个标记独立进行，确保每个标记的嵌入都是基于其自身特征值进行归一化的。

层归一化的优点

• 稳定性：通过标准化嵌入，层归一化有助于提高训练的稳定性。
• 加速收敛：它可以减少对精细参数初始化的依赖，从而加速收敛过程。
• 适应性：与批量归一化不同，层归一化不受批量大小的影响，因此更适合处理变长的序列。

在层归一化中，除了进行标准化处理之外，还会应用缩放（scale）和平移（shift）操作。这些操作使得模型能够在归一化之后重新调整特征的分布，以便更好地适应特定任务。以下是有关缩放和平移的详细解释。

缩放和平移的原理

• 缩放（Scale）：层归一化后，输出的每个特征都会乘以一个可学习的缩放参数 γ\gammaγ。公式：

• 平移（Shift）：在缩放之后，输出会加上一个可学习的平移参数 β\betaβ。公式：

缩放和平移的实现

在 PyTorch 的 nn.LayerNorm 类中，默认情况下，这两个参数 γ\gammaγ 和 β\betaβ 都是可学习的，且在创建层归一化实例时会自动初始化。

归一化的完整过程

整合归一化、缩放和平移，完整的层归一化过程可以表示为：

注意事项——Gamma 和 Beta 的作用

• Gamma（γ\gammaγ）是乘法的，Beta（β\betaβ）是加法的：在层归一化中，γ\gammaγ 和 β\betaβ 的引入使得模型能够灵活调整输出特征的分布。

为什么使用 Gamma 和 Beta？

1. 灵活性和可调性：通过引入 γ\gammaγ 和 β\betaβ，模型可以在训练过程中调整每个特征的权重和偏移。这种灵活性使得模型可以根据需要放大或缩小特征值，以更好地适应不同的任务。
2. 放大和增强能力：模型学习到的 γ\gammaγ 参数允许其在必要时放大某些特征值。比如，当某个特征对分类或回归任务尤为重要时，模型能够自动学习到放大该特征的方式。例如，当输入数据的某个特征在特定情况下与目标输出有较强的相关性时，γ\gammaγ 的值会被调整得较大，从而增强该特征在输出中的影响力。
3. 改善学习能力：通过结合归一化、缩放和平移，模型能够更好地捕捉到输入数据的变化，从而提高其对复杂模式的学习能力。这样的设计减少了对手动特征选择的依赖，使得模型更具自适应性。

小记

总之，γ\gammaγ 和 β\betaβ 的引入为模型提供了更大的灵活性，使得模型能够根据不同的输入特征自适应地调整其输出。这种机制对于处理复杂数据和任务（尤其是在深度学习领域）至关重要，有助于提高模型的表现和稳定性。

示例

参数说明：

features：输入特征的维度，也就是每个嵌入的大小（如 512）。

eps：一个小的常数，用于防止除以零的错误。默认值为 10−610^{-6}10−6。

功能：调用父类的构造函数 super().__init__()。初始化 self.eps，self.alpha 和 self.bias：self.alpha 是一个可学习的参数，初始化为全 1 的张量，其形状为 (features,)。self.bias 是另一个可学习的参数，初始化为全 0 的张量，其形状也为 (features,)。

class LayerNormalization(nn.Module):

    def __init__(self, features: int, eps: float = 10**-6) -> None:
        super().__init__()
        self.eps = eps
        self.alpha = nn.Parameter(torch.ones(features))  # alpha 是可学习的参数
        self.bias = nn.Parameter(torch.zeros(features))   # bias 是可学习的参数

    def forward(self, x):
        # x: (batch, seq_len, hidden_size)
        # 保持维度以便传播
        mean = x.mean(dim=-1, keepdim=True)  # (batch, seq_len, 1)
        # 保持维度以便广播
        std = x.std(dim=-1, keepdim=True)    # (batch, seq_len, 1)
        # eps 用于防止除以零或标准差非常小的情况
        return self.alpha * (x - mean) / (std + self.eps) + self.bias

FeedForwardBlock 类

FeedForwardBlock 类是一个实现前馈神经网络的模块，通常在 Transformer 模型的编码器和解码器中使用。前馈神经网络由两个线性变换和一个 ReLU 激活函数组成。这种结构增加了模型的非线性能力，使其能够学习更复杂的模式。

线性层

• self.linear_1 和 self.linear_2 是线性变换，它们将输入投影到更高的维度空间（d_ff），然后再返回到原始维度（d_model）。

数学公式

给定来自前一层的输入 xxx，前馈网络执行以下操作：

import torch
import torch.nn as nn

class FeedForwardBlock(nn.Module):
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1) -> None:
        super().__init__()
        self.linear_1 = nn.Linear(d_model, d_ff)  # 第一层线性变换
        self.linear_2 = nn.Linear(d_ff, d_model)  # 第二层线性变换
        self.dropout = nn.Dropout(dropout)         # Dropout 层
        self.activation = nn.ReLU()                 # ReLU 激活函数

    def forward(self, x):
        # 应用第一层线性变换
        x = self.linear_1(x)                      # (batch_size, seq_len, d_ff)
        x = self.activation(x)                     # 应用 ReLU 激活
        x = self.dropout(x)                        # 应用 Dropout
        x = self.linear_2(x)                      # 应用第二层线性变换
        return x    

小记

FeedForwardBlock 类在 Transformer 模型中起到关键作用，通过两层线性变换和中间的 ReLU 激活，引入了非线性，使模型能够捕捉到更复杂的特征。Dropout 层则有助于防止过拟合。

MultiHeadAttentionBlock 类

多头注意力（Multi-head Attention）是 Transformer 架构的核心组件，使得模型能够同时关注输入序列的不同部分。接下来我们将分解多头注意力的工作原理以及其重要性。

1. 自注意力机制（Self-Attention Mechanism）

在理解多头注意力之前，理解自注意力机制至关重要。自注意力机制允许序列中的每个位置关注所有其他位置，并提供这些位置的加权和。这帮助模型捕捉序列中无论距离多远的依赖关系。

自注意力过程

• 输入向量：假设我们有一个输入向量序列 X=[x1,x2,…,xn]X = [x_1, x_2, \ldots, x_n]X=[x1​,x2​,…,xn​]，其中每个 xix_ixi​ 是一个 ddd 维向量。
• 查询、键和值矩阵：对于每个输入向量 xix_ixi​，我们计算三个向量：查询 qiq_iqi​、键 kik_iki​ 和值 viv_ivi​。这些向量通过学习的线性变换获得：

其中 WQW_QWQ​、WKW_KWK​ 和 WVW_VWV​ 是学习的权重矩阵。

• 缩放点积注意力（Scaled Dot-Product Attention）：注意力分数是通过查询和键向量计算的。分数决定每个位置在多大程度上关注其他位置：

这里 dkd_kdk​ 是键向量的维度，除以 dk\sqrt{d_k}dk​​ 是一个缩放因子，用以防止点积变得过大。

2. 多头注意力（Multi-Head Attention）

多头注意力通过允许模型在不同位置同时关注不同的表示子空间来扩展自注意力机制。模型并不是执行一个单一的注意力函数，而是并行执行 hhh 个注意力函数（头）。

多头注意力过程

• 多个线性投影：输入向量被线性投影 hhh 次，以创建多个查询、键和值的集合：
• 每个头的缩放点积注意力：每组查询、键和值用于计算注意力分数和输出。
• 头的拼接： hhh 个注意力头的输出被拼接：

其中 WOW_OWO​ 是学习的权重矩阵，用于将拼接的输出投影回所需的维度。

多头注意力的优势

• 增强的表示能力：通过同时关注序列的不同部分，模型能够捕捉到更复杂的关系。
• 信息损失的缓解：多个头确保即使某些头未能捕捉某些依赖关系，其他头可能会成功，从而提供更稳健的表示。

代码实现

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1):
        super().__init__()
        assert d_model % num_heads == 0  # 确保 d_model 可以被 num_heads 整除
        self.num_heads = num_heads
        self.d_head = d_model // num_heads  # 每个头的维度

        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.W_O = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        batch_size, seq_len, d_model = x.size()

        # 线性变换得到 Q, K, V
        Q = self.W_Q(x)  # (batch_size, seq_len, d_model)
        K = self.W_K(x)  # (batch_size, seq_len, d_model)
        V = self.W_V(x)  # (batch_size, seq_len, d_model)

        # 将 Q, K, V 切分成多个头
         # (batch_size, num_heads, seq_len, d_head)
        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2)  
        # (batch_size, num_heads, seq_len, d_head)
        K = K.view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2)
        # (batch_size, num_heads, seq_len, d_head)
        V = V.view(batch_size, seq_len, self.num_heads, self.d_head).transpose(1, 2) 

        # 缩放点积注意力
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_head) 
        attention_weights = torch.softmax(scores, dim=-1)  # 计算注意力权重

        # 加权值
        context = torch.matmul(attention_weights, V)  # (batch_size, num_heads, seq_len, d_head)
        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)  

        # 线性变换投影回原始维度
        output = self.W_O(context)
        return self.dropout(output)  # 返回输出

分析 MultiHeadAttentionBlock 类代码

在这一部分，我们将逐步分析 MultiHeadAttentionBlock 类的代码，了解每个步骤的目的和重要性。

__init__ 方法

def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1):
    super().__init__()
    assert d_model % num_heads == 0  # 确保 d_model 可以被 num_heads 整除
    self.num_heads = num_heads
    self.d_head = d_model // num_heads  # 每个头的维度

    self.W_Q = nn.Linear(d_model, d_model)
    self.W_K = nn.Linear(d_model, d_model)
    self.W_V = nn.Linear(d_model, d_model)
    self.W_O = nn.Linear(d_model, d_model)
    self.dropout = nn.Dropout(dropout)

参数初始化：初始化时设置模型维度 d_model、注意力头的数量 num_heads 和 dropout 比率 dropout。确保 d_model 可以被 num_heads 整除，以便将其均匀分配到各个头中。

线性层定义：定义线性层 W_Q、W_K、W_V 和 W_O，用于生成查询、键、值和最终输出的线性变换。

attention 静态方法

@staticmethod
def attention(query, key, value, mask=None, dropout=None):
    # 计算注意力分数和输出

计算缩放点积注意力：该静态方法用于计算注意力分数。如果提供了掩码（mask），它会应用于注意力分数，确保模型只关注特定位置。使用 softmax 计算注意力权重，最后如果设置了 dropout，还会在注意力权重上应用 dropout。

forward 方法

def forward(self, x):
    # Linear Projections
    query = self.W_Q(x)
    key = self.W_K(x)
    value = self.W_V(x)

• 线性投影：
• 通过学习的权重矩阵 W_Q、W_K 和 W_V 对输入张量 x 进行线性变换，以生成查询、键和值向量。这是计算注意力的第一步。

重塑和转置以进行多头注意力：

# Reshape and Transpose for Multi-Head Attention
query = query.view(query.shape[0], query.shape[1], self.num_heads, self.d_head).transpose(1, 2)
key = key.view(key.shape[0], key.shape[1], self.num_heads, self.d_head).transpose(1, 2)
value = value.view(value.shape[0], value.shape[1], self.num_heads, self.d_head).transpose(1, 2)

• view 方法将查询、键和值张量的形状重塑为 (Batch, Seq_len, h, d_k)，其中 h 是头的数量，d_k 是每个头的维度。
• transpose 方法交换序列长度和头维度，得到的张量形状为 (Batch, h, Seq_len, d_k)。这样，每个注意力头就可以独立处理其对应的 d_model 维度部分。
  

缩放点积注意力：

# Scaled Dot-Product Attention
x, self.attention_scores = MultiHeadAttentionBlock.attention(query, key, value, mask, self.dropout)

• 调用静态方法 attention 来计算注意力分数和加权值的总和。该方法使用查询、键和值张量，以及可能的掩码。
  

重塑和拼接头：

# Reshape and Concatenate Heads
x = x.transpose(1, 2).contiguous().view(x.shape[0], -1, self.h * self.d_k)

• x.transpose(1, 2) 将张量的形状从 (Batch, h, Seq_len, d_k) 变为 (Batch, Seq_len, h, d_k)。
• contiguous() 确保张量的内存布局是连续的，有助于后续操作的效率。
• view 方法将张量重塑为 (Batch, Seq_len, d_model)，其中 d_model = h * d_k，将所有头的输出拼接在一起。
  

最终线性层：

# Final Linear Layer
return self.W_O(x)

• 最后，使用权重矩阵 W_O 进行线性变换，输出张量的形状为 (Batch, Seq_len, d_model)，与输入形状相同。

小记

每个步骤在多头注意力机制中都起着至关重要的作用。通过线性投影，模型能够生成查询、键和值向量；通过重塑和转置，模型能够独立处理多个头的输出；而最终的线性

层则确保输出形状与输入形状一致，为后续处理做好准备。这种设计允许模型在不同的子空间中并行关注输入序列的不同部分，从而增强了模型的表达能力。

什么是 Mask？

在 MultiHeadAttentionBlock 类中，mask 的类型可以是填充 mask 或者前瞻 mask（因果 mask），具体取决于在调用 forward 方法时如何应用。mask 被作为参数传递到 forward 方法中，然后在注意力静态方法中使用。

填充 Mask（Padding Mask）

填充 mask 用于确保输入序列中的填充标记不会影响注意力机制。填充标记是在序列中添加的，以确保它们在一个批次中具有相同的长度，但这些标记并不包含有意义的信息，模型在训练和推理时应该忽略它们。

• 作用：通过将填充标记的注意力分数设置为一个非常大的负值，来确保填充标记不会影响注意力机制，实际上是忽略它们。
• 使用场景：

示例：填充 Mask

考虑一个具有不同长度的序列批次，这些序列已经被填充：

import torch

# 示例输入序列 (batch_size=2, seq_len=5)
input_sequences = [
    [1, 2, 3, 0, 0],  # 序列 1，填充了 0
    [4, 5, 6, 7, 8]   # 序列 2，没有填充
]

# 转换为张量
input_tensor = torch.tensor(input_sequences)

# 创建填充 mask，其中 1 表示有效位置，0 表示填充
padding_mask = (input_tensor != 0).unsqueeze(1).unsqueeze(2)

print("输入张量:")
print(input_tensor)
print("填充 Mask:")
print(padding_mask)

输出：

输入张量:
tensor([[1, 2, 3, 0, 0],
        [4, 5, 6, 7, 8]])

填充 Mask:
tensor([[[[ True,  True,  True, False, False]]],
        [[[ True,  True,  True,  True,  True]]]])

在这个填充 mask 中：

• True 表示有效的位置。
• False 表示填充的位置。

前瞻 Mask（Look-Ahead Mask / Causal Mask）

前瞻 mask 的目的是确保在训练和推理过程中，输出序列中的每个位置只能关注之前的位置和当前的位置，而不能关注任何未来的位置。这在自回归任务（如语言建模和文本生成）中至关重要，因为模型需要基于之前的标记来预测序列中的下一个标记。

• 使用场景：解码器块。

示例：解码器中的前瞻 Mask

import torch

# 创建前瞻 mask
seq_len = 5
look_ahead_mask = torch.triu(torch.ones((seq_len, seq_len)), diagonal=1).bool()

print("前瞻 Mask:")
print(look_ahead_mask)

输出：

前瞻 Mask:
tensor([[False,  True,  True,  True,  True],
        [False, False,  True,  True,  True],
        [False, False, False,  True,  True],
        [False, False, False, False,  True],
        [False, False, False, False, False]])

在这个前瞻 mask 中：

• True 值表示不应该被关注的位置（未来位置）。
• False 值表示可以被关注的位置（当前和过去的位置）。

残差连接（Residual Connection）类

残差连接的目的：残差连接或跳跃连接用于帮助深度神经网络的训练，使梯度能够更容易地通过网络流动。它们本质上“跳过”一个或多个层，将一个层的输入直接连接到更深层的输出。这有助于缓解消失梯度的问题，并允许非常深的网络更有效地训练。

Transformer 中的典型结构

在 Transformer 模型中，残差连接通常与层归一化（Layer Normalization）结合使用。在多头注意力层或前馈层之后，原始输入会添加回输出（经过该层的处理后），然后再应用层归一化。

class ResidualConnection(nn.Module):
    def __init__(self, dropout: float) -> None:
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.norm = LayerNormalization()

    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))

__init__ 方法：

• 初始化时设定 dropout（丢弃率）。
• 定义了 dropout 和层归一化。

forward 方法：

• 应用层归一化、dropout，并添加输入 x（实现残差连接）。

编码器块（EncoderBlock）类

现在我们将创建编码器块，其中包含一个多头注意力层、两个 Add & Norm 层以及一个前馈层。

class EncoderBlock(nn.Module):
    def __init__(self, self_attention_block: MultiHeadAttentionBlock,
             feed_forward_block: FeedForwardBlock, dropout: float) -> None:
        super().__init__()
        self.self_attention_block = self_attention_block
        self.feed_forward_block = feed_forward_block
        self.residual_connections = nn.ModuleList([ResidualConnection(dropout) for _ in range(2)])

    def forward(self, x, src_mask):
        x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, src_mask))
        x = self.residual_connections[1](x, self.feed_forward_block)
        return x

__init__ 方法：

• 初始化时接收 self_attention_block（一个 MultiHeadAttentionBlock 的实例）、feed_forward_block（一个 FeedForwardBlock 的实例）以及 dropout。
• 创建两个用于自注意力和前馈块的残差连接。

forward 方法：

• 应用自注意力块和残差连接：这里自注意力块的输入 x 作为查询（q）、键（k）和值（v），这就是为什么称之为“自注意力”。实际上，输入 x 用于关注自身。即句子中的每个单词都与同一句子中的其他单词交互。
• 在解码器中，注意力机制的工作方式有所不同，因为存在交叉注意力。在交叉注意力中，查询来自解码器，而键和值来自编码器。这使得解码器可以关注编码器生成的输入序列的相关部分，而不仅仅是关注自己的输出。
• 因此，自注意力能够实现句内交互，而交叉注意力则促进编码器输出与解码器输入之间的交互。
• 应用前馈块和残差连接。

编码器（Encoder）类

class Encoder(nn.Module):
    def __init__(self, layers: nn.ModuleList) -> None:
        super().__init__()
        self.layers = layers
        self.norm = LayerNormalization()

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

__init__ 方法：

• 初始化时接收层列表（layers），并定义层归一化。

forward 方法：

• 对输入 x 应用每一层（layer），并使用 mask 作为输入。
• 最后返回归一化后的输出。

解码器块（DecoderBlock）类

DecoderBlock 类表示 Transformer 解码器的一个单元块。每个解码器块包含自注意力机制、交叉注意力机制（关注编码器的输出）和前馈网络，这些都被残差连接和层归一化包围。

class DecoderBlock(nn.Module):
    def __init__(self, self_attention_block: MultiHeadAttentionBlock, 
            cross_attention_block: MultiHeadAttentionBlock, 
            feed_forward_block: FeedForwardBlock, dropout: float) -> None:
        super().__init__()
        self.self_attention_block = self_attention_block
        self.cross_attention_block = cross_attention_block
        self.feed_forward_block = feed_forward_block
        self.residual_connections = nn.ModuleList([ResidualConnection(dropout) for _ in range(3)])

    def forward(self, x, encoder_output, src_mask, tgt_mask):
        x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, tgt_mask))
        x = self.residual_connections[1](x, lambda x: self.cross_attention_block(x, encoder_output, encoder_output, src_mask))
        x = self.residual_connections[2](x, self.feed_forward_block)
        return x

__init__ 方法：

• 参数：
• self_attention_block: 自注意力的 MultiHeadAttentionBlock 实例。
• cross_attention_block: 交叉注意力的 MultiHeadAttentionBlock 实例。
• feed_forward_block: 前馈块的实例。
• dropout: 用于正则化的 dropout 率。
• 属性：
• self.residual_connections: 包含三个残差连接实例的列表，分别用于解码器块中的每个子层。

forward 方法：

• 参数：
• x: 输入张量（解码器输入）。
• encoder_output: 来自编码器的输出。
• src_mask: 源掩码，防止模型关注源输入中的填充标记。
• tgt_mask: 目标掩码，防止模型关注目标序列中的未来标记（看向前掩码）。

解码器（Decoder）类

Decoder 类由一系列解码器块组成。它将这些块依次应用于输入，最后进行层归一化。

class Decoder(nn.Module):
    def __init__(self, features: int, layers: nn.ModuleList) -> None:
        super().__init__()
        self.layers = layers
        self.norm = LayerNormalization(features)

    def forward(self, x, encoder_output, src_mask, tgt_mask):
        for layer in self.layers:
            x = layer(x, encoder_output, src_mask, tgt_mask)
        return self.norm(x)

__init__ 方法：

• 参数：
• features: 输入和输出的特征数量（维度）。
• layers: 构成解码器的 DecoderBlock 实例列表。
• 属性：
• self.layers: 存储解码器块的列表。
• self.norm: 对解码器堆栈的输出应用层归一化。

forward 方法：

• 参数：
• x: 输入张量（解码器输入）。
• encoder_output: 来自编码器的输出。
• src_mask: 源掩码，防止模型关注源输入中的填充标记。
• tgt_mask: 目标掩码，防止模型关注目标序列中的未来标记（看向前掩码）。

投影层（ProjectionLayer）类

ProjectionLayer 类用于将高维向量（解码器的输出）转换为词汇表上的 logits。这一投影通常是 Transformer 模型解码器中的最后一层。

class ProjectionLayer(nn.Module):
    def __init__(self, d_model, vocab_size) -> None:
        super().__init__()
        self.proj = nn.Linear(d_model, vocab_size)

    def forward(self, x) -> None:
        # (batch, seq_len, d_model) --> (batch, seq_len, vocab_size)
        return self.proj(x)

__init__ 方法：

• 参数：
• d_model: 模型内部表示的维度（即隐藏层大小）。
• vocab_size: 词汇表的大小，表示可能的输出标记数量。
• 属性：、
• self.proj: 一个 nn.Linear 实例，用于将 d_model 维度映射到 vocab_size 维度。

forward 方法：

• 参数：
• x: 形状为 (batch_size, seq_len, d_model) 的张量，其中 batch_size 是批次中的序列数量，seq_len 是每个序列的长度，d_model 是模型隐藏状态的维度。
• 返回：
• 通过线性变换将 d_model 维度的向量投影到 vocab_size 维度的 logits，用于生成序列中每个位置的概率分布。

小记

解码器块、解码器以及投影层共同构成了 Transformer 解码器的核心部分。解码器块利用自注意力和交叉注意力机制来增强模型的表示能力，而投影层则将内部高维表示转换为可以用于生成输出的 logits。这种结构允许 Transformer 在处理序列生成任务（如翻译或文本生成）时，灵活地关注输入序列的不同部分。

Transformer 类

Transformer 类封装了整个 Transformer 模型，集成了编码器和解码器组件，以及嵌入层和位置编码。下面是类中各部分的详细说明：

class Transformer(nn.Module):
    def __init__(self, encoder: Encoder, decoder: Decoder, src_embed: InputEmbeddings, 
                    tgt_embed: InputEmbeddings, src_pos: PositionalEncoding, 
                    tgt_pos: PositionalEncoding, projection_layer: ProjectionLayer) -> None:
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.src_pos = src_pos
        self.tgt_pos = tgt_pos
        self.projection_layer = projection_layer

    def encode(self, src, src_mask):
        # (batch, seq_len, d_model)
        src = self.src_embed(src)
        src = self.src_pos(src)
        return self.encoder(src, src_mask)
    
    def decode(self, encoder_output: torch.Tensor, src_mask: torch.Tensor, 
                    tgt: torch.Tensor, tgt_mask: torch.Tensor):
        # (batch, seq_len, d_model)
        tgt = self.tgt_embed(tgt)
        tgt = self.tgt_pos(tgt)
        return self.decoder(tgt, encoder_output, src_mask, tgt_mask)
    
    def project(self, x):
        # (batch, seq_len, vocab_size)
        return self.projection_layer(x)

参数：

• encoder: 一个 Encoder 类的实例，负责编码源序列。
• decoder: 一个 Decoder 类的实例，负责解码编码后的表示并生成输出序列。
• src_embed: 用于嵌入源序列标记的 InputEmbeddings 实例。
• tgt_embed: 用于嵌入目标序列标记的 InputEmbeddings 实例。
• src_pos: 用于将位置编码添加到源嵌入的 PositionalEncoding 实例。
• tgt_pos: 用于将位置编码添加到目标嵌入的 PositionalEncoding 实例。
• projection_layer: 一个 ProjectionLayer 实例，用于将解码器的输出投影到词汇表大小。

encode 方法

• 参数：
• src: 形状为 (batch_size, seq_len) 的源序列张量。
• src_mask: 用于处理源序列中填充标记的源掩码张量。
• 执行过程：
• 将源张量 src 应用源嵌入层。
• 将位置编码添加到嵌入后的源张量。
• 将结果张量通过编码器和源掩码传递。
• 输出：
• 返回形状为 (batch_size, seq_len, d_model) 的源序列的编码表示。

decode 方法

• 参数：
• encoder_output: 编码器的源序列编码表示。
• src_mask: 在编码期间使用的源掩码。
• tgt: 形状为 (batch_size, seq_len) 的目标序列张量。
• tgt_mask: 用于处理目标序列中填充和未来标记的目标掩码张量。
• 执行过程：
• 将目标张量 tgt 应用目标嵌入层。
• 将位置编码添加到嵌入后的目标张量。
• 将结果张量通过解码器，传入编码器输出和相应的掩码。
• 输出：
• 返回形状为 (batch_size, seq_len, d_model) 的目标序列的解码表示。

project 方法

• 参数：
• x: 解码器输出的张量，形状为 (batch_size, seq_len, d_model)。
• 执行过程：
• 应用投影层，将 d_model 维度的输出映射到 vocab_size 维度的 logits。
• 输出：
• 返回形状为 (batch_size, seq_len, vocab_size) 的 logits 张量。

构建 Transformer 方法

build_transformer 构建完整的 Transformer 模型，将各种组件组合在一起，例如嵌入层、位置编码、编码器和解码器块，以及最终的投影层。

def build_transformer(src_vocab_size: int, tgt_vocab_size: int, src_seq_len: int, 
                        tgt_seq_len: int, d_model: int = 512, N: int = 6, 
                        h: int = 8, dropout: float = 0.1, d_ff: int = 2048) -> Transformer:
    # 创建嵌入层
    src_embed = InputEmbeddings(d_model, src_vocab_size)
    tgt_embed = InputEmbeddings(d_model, tgt_vocab_size)

    # 创建位置编码层
    src_pos = PositionalEncoding(d_model, src_seq_len, dropout)
    tgt_pos = PositionalEncoding(d_model, tgt_seq_len, dropout)
    
    # 创建编码器块
    encoder_blocks = []
    for _ in range(N):
        encoder_self_attention_block = MultiHeadAttentionBlock(d_model, h, dropout)
        feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)
        encoder_block = EncoderBlock(d_model, encoder_self_attention_block, feed_forward_block, dropout)
        encoder_blocks.append(encoder_block)

    # 创建解码器块
    decoder_blocks = []
    for _ in range(N):
        decoder_self_attention_block = MultiHeadAttentionBlock(d_model, h, dropout)
        decoder_cross_attention_block = MultiHeadAttentionBlock(d_model, h, dropout)
        feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)
        decoder_block = DecoderBlock(d_model, decoder_self_attention_block, 
                                            decoder_cross_attention_block, feed_forward_block, dropout)
        decoder_blocks.append(decoder_block)
    
    # 创建编码器和解码器
    encoder = Encoder(d_model, nn.ModuleList(encoder_blocks))
    decoder = Decoder(d_model, nn.ModuleList(decoder_blocks))
    
    # 创建投影层
    projection_layer = ProjectionLayer(d_model, tgt_vocab_size)
    
    # 创建 Transformer
    transformer = Transformer(encoder, decoder, src_embed, tgt_embed, src_pos, tgt_pos, projection_layer)
    
    # 初始化参数
    for p in transformer.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    
    return transformer

在此感谢您的阅读~

参考内容

Chefer, H., Gur, S., & Wolf, L. (2021). Generic attention-model explainability for interpreting bi-modal and encoder-decoder Transformers. In arXiv [cs.CV]. http://arxiv.org/abs/2103.15679

Jain, S., & Wallace, B. C. (2019). Attention is not Explanation. In arXiv [cs.CL]. http://arxiv.org/abs/1902.10186

Lynn-Evans, S. (2018, September 27). How to code The Transformer in Pytorch. Towards Data Science. https://towardsdatascience.com/how-to-code-the-transformer-in-pytorch-24db27c8f9ec

Jamil, U. [@umarjamilai]. (2023, May 25). Coding a Transformer from scratch on PyTorch, with full explanation, training and inference. Youtube. https://www.youtube.com/watch?v=ISNdQcPhsts

Vaswani, A., Shazeer, N., Parmar, N., & Uszkoreit, J. (n.d.). Attention is all you need. Arxiv.org. Retrieved June 15, 2024, from http://arxiv.org/abs/1706.03762

(N.d.). Datacamp.com. Retrieved June 15, 2024, from https://www.datacamp.com/tutorial/building-a-transformer-with-py-torch

原文链接

点击这里