一、定义

Transformer是一种深度学习模型架构,专门为处理序列数据而设计,尤其是在自然语言处理(NLP)领域取得了革命性的成功。它在2017年由Google Brain团队在论文《Attention Is All You Need》中首次提出。

它的核心思想是:完全摒弃了传统的循环神经网络(RNN)和卷积神经网络(CNN),仅依赖一种名为“自注意力(Self-Attention)”的机制来捕捉序列中各个元素之间的全局依赖关系。

二、内部结构

Transformer是一种基于注意力机制的编码器-解码器结构,具有很好的并行性能,同时利用注意力机制很好地解决了长序列的长程依赖问题。

以翻译的任务来解释整个过程

输入:I love you

输出:我爱你

1.输入编码:计算机无法直接理解单词,所以需要首先将每个单词转换为一个数字向量

2.位置编码:由于transformer没有RNN循环结构,它本身无法感知单词的顺序,因此要向词向量中添加位置信息(使用正弦函数与余弦函数生成一组代表每个位置的独特向量)

3.编码器(Encoder)

编码器由N个完全相同的层堆叠而成,每一层有两个核心子层

(1)多头注意力机制:多头注意力层接收输入词嵌入与位置编码之和,并进行多头注意力的计算,目的是从多个不同角度提取交互信息,从而更好的理解上下文

a)创建Q,K,V:自注意力的Q、K、V均来源于同一个输入X

b)  计算注意力分数

c)  缩放和Softmax

d)  加权求和

(2)前馈神经网络:一个简单的全连接网络,对自注意力层的输出进行进一步的非线性变换。

在编码器的每个子层都紧跟着惨超链接与层归一化

  • 残差连接 (Residual Connection):将子层的输入直接加到其输出上。这有助于缓解深层网络中的梯度消失问题。

  • 层归一化 (Layer Normalization):对相加后的结果进行归一化,使训练更加稳定。

4.解码器(Decoder)

每一层包含三个子层:

  1. 掩码多头自注意力机制:与编码器的自注意力类似,但有一个关键区别——masking。在训练时,为了防止模型“作弊”(看到未来的答案),它会将当前生成位置之后的所有位置都掩盖掉(设置为负无穷,经Softmax后变为0)。例如,在生成“爱”这个词时,它只能关注到 <start> 和 “我”,而不能看到未来的“你”。

  2. 编码器-解码器注意力机制 (Encoder-Decoder Attention):这是连接编码器和解码器的桥梁。它的Query来自解码器上一层的输出,而Key和Value则来自编码器的最终输出。这让解码器在生成每一个目标词时,都能够有选择地“聚焦”在输入序列的不同部分上

  3. 前馈神经网络:与编码器中的相同。

解码器最终会输出一个向量序列,每个向量经过一个Softmax层后,被转换成一个概率分布,覆盖目标词汇表中的所有单词。我们选择概率最高的那个词作为当前时间步的输出。

三、代码详解

先确定好transformer的一些全局参数

d_k = 64  # Q 和 K 向量的维度
d_v = 64  # V 向量的维度
d_embedding = 128  # 词嵌入的维度
n_heads = 8  # 多头注意力机制中头的数量
n_layers = 6  # 编码器和解码器的层数

缩放点积注意力:对应这个公式

class ScaledDotProductAttention(nn.Module):
    """
    缩放点积注意力模块
    根据 Q、K、V 计算注意力分数和上下文向量
    """

    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()

    def forward(self, Q, K, V, attn_mask):
        # Q, K, V 的维度: [batch_size, n_heads, len_q/k/v, dim_q=k/v]
        # attn_mask 的维度: [batch_size, n_heads, len_q, len_k]

        # 1. 计算注意力分数
        # 将 Q 与 K 的转置相乘,并除以 sqrt(d_k) 进行缩放
        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)

        # 2. 应用注意力掩码
        # 将 attn_mask 中为 True 的位置的 scores 替换为一个极小值 (-1e9)
        # 这样在 softmax 之后,这些位置的权重将接近于0
        scores.masked_fill_(attn_mask, -1e9)

        # 3. 对分数进行 softmax 归一化
        # 沿着最后一个维度 (len_k) 进行 softmax
        weights = nn.Softmax(dim=-1)(scores)

        # 4. 计算上下文向量
        # 将归一化后的权重与 V 相乘
        context = torch.matmul(weights, V)

        # 返回上下文向量和注意力权重
        return context, weights

多头注意力机制模块:对应公式

class MultiHeadAttention(nn.Module):
    """
    多头注意力机制模块
    将输入投影到多个子空间,并并行计算注意力,最后将结果拼接
    """

    def __init__(self):
        super(MultiHeadAttention, self).__init__()
        # 线性投影层,用于生成 Q, K, V [128,512]
        self.W_Q = nn.Linear(d_embedding, d_k * n_heads)
        self.W_K = nn.Linear(d_embedding, d_k * n_heads)
        self.W_V = nn.Linear(d_embedding, d_v * n_heads)
        # 最后的线性层,将拼接后的多头输出投影回原始维度 [512,128]
        self.linear = nn.Linear(n_heads * d_v, d_embedding)
        # 层归一化
        self.layer_norm = nn.LayerNorm(d_embedding)

    def forward(self, Q, K, V, attn_mask):
        # Q, K, V 的维度: [batch_size, len_q/k/v, d_embedding]

        residual, batch_size = Q, Q.size(0)

        # 1. 线性投影并重塑
        # 将输入 Q, K, V 投影到多头子空间,并调整维度以便并行计算
        # .view() 和 .transpose() 操作将维度变为: [batch_size, n_heads, len_q/k/v, d_q=k/v]
        q_s = self.W_Q(Q).view(batch_size, -1, n_heads, d_k).transpose(1, 2)
        k_s = self.W_K(K).view(batch_size, -1, n_heads, d_k).transpose(1, 2)
        v_s = self.W_V(V).view(batch_size, -1, n_heads, d_v).transpose(1, 2)

        # 2. 复制注意力掩码
        # 将 attn_mask 复制 n_heads 次,以适应多头注意力
        # attn_mask 的维度: [batch_size, n_heads, len_q, len_k]
        attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1)

        # 3. 计算缩放点积注意力
        context, weights = ScaledDotProductAttention()(q_s, k_s, v_s, attn_mask)
        # context 的维度: [batch_size, n_heads, len_q, dim_v]
        # weights 的维度: [batch_size, n_heads, len_q, len_k]

        # 4. 拼接多头结果
        # 调整维度并使用 .contiguous() 确保内存连续,然后用 .view() 将多头结果拼接
        # context 的维度: [batch_size, len_q, n_heads * dim_v]
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, n_heads * d_v)

        # 5. 最终线性投影和层归一化
        # 将拼接后的结果通过线性层投影回 d_embedding 维度
        output = self.linear(context)
        # 将输出与残差连接相加,并进行层归一化
        # output 的维度: [batch_size, len_q, d_embedding]
        output = self.layer_norm(output + residual)

        # 返回最终输出和注意力权重
        return output, weights

前馈神经网络:

class PoswiseFeedForwardNet(nn.Module):
    """
    逐位置前馈网络模块
    对序列中的每个位置独立地应用一个全连接前馈网络
    """

    def __init__(self, d_ff=2048):
        super(PoswiseFeedForwardNet, self).__init__()
        # 两个一维卷积层,相当于两个全连接层
        self.conv1 = nn.Conv1d(in_channels=d_embedding, out_channels=d_ff, kernel_size=1)
        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_embedding, kernel_size=1)
        # 层归一化
        self.layer_norm = nn.LayerNorm(d_embedding)

    def forward(self, inputs):
        # inputs 的维度: [batch_size, len_q, d_embedding]

        residual = inputs

        # 1. 维度转换并应用第一个卷积层
        # inputs.transpose(1, 2) 将维度变为 [batch_size, d_embedding, len_q]
        # 卷积操作后,output 的维度: [batch_size, d_ff, len_q]
        output = nn.ReLU()(self.conv1(inputs.transpose(1, 2)))

        # 2. 应用第二个卷积层并转换维度
        # 卷积操作后,output 的维度: [batch_size, d_embedding, len_q]
        # .transpose(1, 2) 将维度恢复为 [batch_size, len_q, d_embedding]
        output = self.conv2(output).transpose(1, 2)

        # 3. 残差连接和层归一化
        output = self.layer_norm(output + residual)

        return output

位置编码:

def get_sin_enc_table(n_position, embedding_dim):
    """
    生成正弦位置编码表
    用于在序列中引入词语的绝对位置信息
    """
    # 初始化正弦编码表
    sinusoid_table = np.zeros((n_position, embedding_dim))
    # 计算不同位置和维度的角度
    for pos_i in range(n_position):
        for hid_j in range(embedding_dim):
            angle = pos_i / np.power(10000, 2 * (hid_j // 2) / embedding_dim)
            sinusoid_table[pos_i, hid_j] = angle

    # 将偶数维应用 sin 函数,奇数维应用 cos 函数
    sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])
    sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])

    # 转换为 PyTorch 张量
    return torch.FloatTensor(sinusoid_table)

掩码机制(填充掩码和后续掩码)

def get_attn_pad_mask(seq_q, seq_k):
    """
    生成填充注意力掩码
    用于在注意力计算中忽略填充 <pad> 词语
    """
    # seq_q 的维度: [batch_size, len_q]
    # seq_k 的维度: [batch_size, len_k]

    batch_size, len_q = seq_q.size()
    batch_size, len_k = seq_k.size()

    # 找到 seq_k 中所有值为 0 (<pad>) 的位置
    # pad_attn_mask 的维度: [batch_size, 1, len_k]
    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1)

    # 将掩码扩展到与注意力分数相同的形状
    # pad_attn_mask 的维度: [batch_size, len_q, len_k]
    pad_attn_mask = pad_attn_mask.expand(batch_size, len_q, len_k)

    return pad_attn_mask


def get_attn_subsequent_mask(seq):
    """
    生成后续注意力掩码 (仅用于解码器)
    用于在注意力计算中忽略当前位置之后的信息,防止信息泄露
    """
    # seq 的维度: [batch_size, seq_len]
    attn_shape = [seq.size(0), seq.size(1), seq.size(1)]
    # 创建一个上三角矩阵,k=1 表示主对角线上方的元素为 1
    # subsequent_mask 的维度: [batch_size, seq_len, seq_len]
    subsequent_mask = np.triu(np.ones(attn_shape), k=1)
    # 转换为 PyTorch 字节张量 (布尔类型)
    subsequent_mask = torch.from_numpy(subsequent_mask).byte()

    return subsequent_mask

编码器模块:

class EncoderLayer(nn.Module):
    """
    编码器的一层
    包含一个多头自注意力层和一个位置前馈网络
    """

    def __init__(self):
        super(EncoderLayer, self).__init__()
        self.enc_self_attn = MultiHeadAttention()
        self.pos_ffn = PoswiseFeedForwardNet()

    def forward(self, enc_inputs, enc_self_attn_mask):
        # enc_inputs 的维度: [batch_size, seq_len, d_embedding]
        # enc_self_attn_mask 的维度: [batch_size, seq_len, seq_len]

        # 将相同的 Q, K, V 输入多头自注意力层
        enc_outputs, attn_weights = self.enc_self_attn(enc_inputs, enc_inputs,
                                                       enc_inputs, enc_self_attn_mask)

        # 将自注意力输出输入位置前馈网络
        enc_outputs = self.pos_ffn(enc_outputs)

        # 返回最终输出和注意力权重
        return enc_outputs, attn_weights


class Encoder(nn.Module):
    """
    Transformer 编码器
    由词嵌入层、位置嵌入层和多个编码器层组成
    """

    def __init__(self, corpus):
        super(Encoder, self).__init__()
        self.src_emb = nn.Embedding(len(corpus.src_vocab), d_embedding)
        # 从预计算的位置编码表初始化位置嵌入层,并冻结参数
        self.pos_emb = nn.Embedding.from_pretrained(
            get_sin_enc_table(corpus.src_len + 1, d_embedding), freeze=True)
        # 堆叠 n_layers 个编码器层
        self.layers = nn.ModuleList(EncoderLayer() for _ in range(n_layers))

    def forward(self, enc_inputs):
        # enc_inputs 的维度: [batch_size, source_len]

        # 生成位置索引序列
        pos_indices = torch.arange(1, enc_inputs.size(1) + 1).unsqueeze(0).to(enc_inputs)

        # 词嵌入和位置嵌入相加
        enc_outputs = self.src_emb(enc_inputs) + self.pos_emb(pos_indices)

        # 生成填充注意力掩码
        enc_self_attn_mask = get_attn_pad_mask(enc_inputs, enc_inputs)

        enc_self_attn_weights = []
        # 逐层通过编码器层
        for layer in self.layers:
            enc_outputs, enc_self_attn_weight = layer(enc_outputs, enc_self_attn_mask)
            enc_self_attn_weights.append(enc_self_attn_weight)

        # 返回最终输出和所有层的注意力权重
        return enc_outputs, enc_self_attn_weights

解码器模块:

class DecoderLayer(nn.Module):
    """
    解码器的一层
    包含一个多头自注意力层、一个编码器-解码器注意力层和一个位置前馈网络
    """

    def __init__(self):
        super(DecoderLayer, self).__init__()
        self.dec_self_attn = MultiHeadAttention()
        self.dec_enc_attn = MultiHeadAttention()
        self.pos_ffn = PoswiseFeedForwardNet()

    def forward(self, dec_inputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask):
        # dec_inputs 的维度: [batch_size, target_len, d_embedding]
        # enc_outputs 的维度: [batch_size, source_len, d_embedding]

        # 1. 第一个注意力子层: 多头自注意力
        # Q, K, V 都来自解码器输入
        dec_outputs, dec_self_attn = self.dec_self_attn(dec_inputs, dec_inputs,
                                                        dec_inputs, dec_self_attn_mask)

        # 2. 第二个注意力子层: 编码器-解码器注意力
        # Q 来自解码器输出,K, V 来自编码器输出
        dec_outputs, dec_enc_attn = self.dec_enc_attn(dec_outputs, enc_outputs,
                                                      enc_outputs, dec_enc_attn_mask)

        # 3. 位置前馈网络
        dec_outputs = self.pos_ffn(dec_outputs)

        # 返回最终输出和两个注意力层的权重
        return dec_outputs, dec_self_attn, dec_enc_attn


class Decoder(nn.Module):
    """
    Transformer 解码器
    由词嵌入层、位置嵌入层和多个解码器层组成
    """

    def __init__(self, corpus):
        super(Decoder, self).__init__()
        self.tgt_emb = nn.Embedding(len(corpus.tgt_vocab), d_embedding)
        # 从预计算的位置编码表初始化位置嵌入层,并冻结参数
        self.pos_emb = nn.Embedding.from_pretrained(
            get_sin_enc_table(corpus.tgt_len + 1, d_embedding), freeze=True)
        # 堆叠 n_layers 个解码器层
        self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)])

    def forward(self, dec_inputs, enc_inputs, enc_outputs):
        # dec_inputs 的维度: [batch_size, target_len]
        # enc_inputs 的维度: [batch_size, source_len]
        # enc_outputs 的维度: [batch_size, source_len, d_embedding]

        # 生成位置索引序列
        pos_indices = torch.arange(1, dec_inputs.size(1) + 1).unsqueeze(0).to(dec_inputs)

        # 词嵌入和位置嵌入相加
        dec_outputs = self.tgt_emb(dec_inputs) + self.pos_emb(pos_indices)

        # 1. 生成解码器自注意力掩码
        dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs)
        dec_self_attn_subsequent_mask = get_attn_subsequent_mask(dec_inputs)
        # 将填充掩码和后续掩码相加,结果大于 0 的位置为 True
        dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask
                                       + dec_self_attn_subsequent_mask), 0)

        # 2. 生成编码器-解码器注意力掩码 (仅考虑填充)
        dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs)

        dec_self_attns, dec_enc_attns = [], []
        # 逐层通过解码器层
        for layer in self.layers:
            dec_outputs, dec_self_attn, dec_enc_attn = layer(dec_outputs, enc_outputs,
                                                             dec_self_attn_mask, dec_enc_attn_mask)
            dec_self_attns.append(dec_self_attn)
            dec_enc_attns.append(dec_enc_attn)

        # 返回最终输出和所有层的注意力权重
        return dec_outputs, dec_self_attns, dec_enc_attns

Transformer的模型结构

class Transformer(nn.Module):
    """
    Transformer 模型的总框架
    由编码器、解码器和最后的线性投影层组成
    """

    def __init__(self, corpus):
        super(Transformer, self).__init__()
        self.encoder = Encoder(corpus)
        self.decoder = Decoder(corpus)
        # 最后的线性层,将解码器输出映射到目标词汇表大小
        self.projection = nn.Linear(d_embedding, len(corpus.tgt_vocab), bias=False)

    def forward(self, enc_inputs, dec_inputs):
        # enc_inputs 的维度: [batch_size, source_seq_len]
        # dec_inputs 的维度: [batch_size, target_seq_len]

        # 1. 编码器前向传播
        enc_outputs, enc_self_attns = self.encoder(enc_inputs)
        # enc_outputs 的维度: [batch_size, source_len, d_embedding]

        # 2. 解码器前向传播
        dec_outputs, dec_self_attns, dec_enc_attns = self.decoder(dec_inputs, enc_inputs, enc_outputs)
        # dec_outputs 的维度: [batch_size, target_len, d_embedding]

        # 3. 线性投影
        dec_logits = self.projection(dec_outputs)
        # dec_logits 的维度: [batch_size, target_len, tgt_vocab_size]

        # 返回逻辑值和所有注意力权重
        return dec_logits, enc_self_attns, dec_self_attns, dec_enc_attns

Logo

火山引擎开发者社区是火山引擎打造的AI技术生态平台,聚焦Agent与大模型开发,提供豆包系列模型(图像/视频/视觉)、智能分析与会话工具,并配套评测集、动手实验室及行业案例库。社区通过技术沙龙、挑战赛等活动促进开发者成长,新用户可领50万Tokens权益,助力构建智能应用。

更多推荐