Transformer模型学习【附代码】

关注

原创 2025-10-07 12:39:28 1k 阅读 AI 写同款· GEO 优化 ›

深度学习 | Transformer模型及代码实现 Transformer 与 RNN 不同，可以比较好地并行训练。Transformer 本身是不能利用单词的顺序信息的，因此需要在输入中添加位置 Embedding，否则 Transformer 就是一个词袋模型了。Transformer 的重点是 Self-Attention 结构，其中用到的 Q, K, V矩阵通过输出进行线性变换得到。阅读详情

前言

Transformer是一种基于注意力机制的深度学习模型，在2017年的论文[1706.03762] Attention Is All You Need中提出。

架构

在这里插入图片描述

嵌入（Embedding）

Embedding 是一种将离散的符号（如单词、句子、类别等）转换为连续向量空间中的实数向量的技术,这些向量能够捕捉原始符号的语义、语法或结构信息。

简单来说，Embedding 就是把高维稀疏的表示（比如 one-hot 编码）映射到低维稠密的向量空间中，使得语义相似的对象在向量空间中距离更近。

自注意力机制

自注意力机制允许模型在处理序列时，动态地为每个位置分配权重，从而捕捉序列不同位置间的依赖关系。

注意力评分函数通过计算查询Q和键K得到注意力评分，再经过softmax运算得到注意力权重，注意力汇聚的输出就是注意力权重和对应的值的加权和。

具体过程：

相似性计算：首先计算查询和每个键之间的相似性，常见的计算方法有点积、缩放点积、加性注意力等。
权重计算：将相似性分数通过softmax进行归一化，得到每个键的权重。
加权求和：将每个键对应的值乘以权重，然后加权求和，得到最终输出。这个输出就是模型在当前查询下关注到的最重要的信息。

掩码机制

参考手撕Transformer（二）| Transformer掩码机制的两个功能，三个位置的解析及其代码_transformer 掩码-优快云博客

缩放点积注意力（Scaled Dot-Product Attention）

点积要求查询Q和键K有相同长度d。

$Attention(Q,K,V)=softmax(\dfrac{QK^T}{\sqrt{d}})V$ ，其中，Q表示查询矩阵，K表示键矩阵，V表示值矩阵， $d$ 是键向量维度。

除以 $\sqrt{d}$ 可以控制点积结果的方差，防止梯度消失或爆炸。假设查询和键的所有元素都是独立的随机变量，并且都满足零均值和单位方差，那么两个向量的点积的均值为0，方差为d。为确保无论向量长度如何，点积的方差在不考虑向量长度的情况下仍然是1，我们将点积除以 $\sqrt{d}$ 。
import torch.nn as nn
import torch.nn.functional as F
import torch
import math


class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super().__init__()

  def forward(self, queries, keys, values, mask=None):
        # queries:(batch_size, num_queries, d)
        # keys:(batch_size, num_kvs, d)
        # values:(batch_size, num_kvs, v)
        # mask:(batch_size, num_queries, num_kvs)
        scores = torch.bmm(queries, keys.transpose(1, 2)) / \
            math.sqrt(queries.shape[-1])
       if mask is not None:
           scores = scores.masked_fill(mask, -1e9)
       attention_weights = F.softmax(scores, dim=-1)
       # 注意是dim=-1，对每一批次每一行归一化
       output = torch.bmm(attention_weights, values)
       return output


queries = torch.normal(0, 1, (2, 1, 2))
keys = torch.normal(0, 1, (2, 5, 2))
values = torch.normal(0, 1, (2, 5, 3))
net = ScaledDotProductAttention()
net.eval()
ans = net(queries, keys, values)
print(ans.shape) # [2, 1, 3]
加性注意力

查询和键可以有不同长度。
注意力评分函数 $a(q,k)= W_v^Ttanh(W_qq+W_kk)$
$Attention(Q,K,V)=softmax(W_v^Ttanh(QW_q^T+KW_k^T))V$

多头注意力（Multi-Head Attention）

多头注意力是自注意力的扩展，它通过将输入序列分成多个不同的“头”，并分别对每个头进行自注意力操作，然后将这些头的输出进行拼接或加权求和，从而得到最终的输出。
多头注意力的优点是可以从不同的角度学习输入序列的特征，使得模型能够捕捉到更丰富的信息

class MultiHeadAttention(nn.Module):
    def __init__(self, query_size, key_size, value_size, num_heads, num_hiddens, bias=False):
        super().__init__()
        self.wq = nn.Linear(query_size, num_hiddens, bias=bias)
        self.wk = nn.Linear(key_size, num_hiddens, bias=bias)
        self.wv = nn.Linear(value_size, num_hiddens, bias=bias)
        self.wo = nn.Linear(num_hiddens, num_hiddens, bias=bias)
        self.num_heads = num_heads
        self.attention = ScaledDotProductAttention()
        self.layernorm = nn.LayerNorm(num_hiddens)

    def split_heads(self, x):
        # x:(batch_size, num_tokens, num_hiddens)
        batch_size, num_tokens, num_hiddens = x.shape
        x = x.view(batch_size, num_tokens, self.num_heads,
                   num_hiddens // self.num_heads)
        # (batch_size, num_tokens, num_heads, num_hiddens/num_heads)
        x = x.permute(0, 2, 1, 3)
        # (batch_size, num_heads, num_tokens, num_hiddens/num_heads)
        return x.reshape(-1, num_tokens, num_hiddens // self.num_heads)
        # (batch_size*num_heads, num_tokens, num_hiddens/num_heads)

    def combine_heads(self, x):
        # x:(batch_size*num_heads, num_tokens, num_hiddens/num_heads)
        batch_size_times_num_heads, num_tokens, num_hiddens_div_num_heads = x.shape
        x = x.view(-1, self.num_heads, num_tokens, num_hiddens_div_num_heads)
        # (batch_size, num_heads, num_tokens, num_hiddens/num_heads)
        x = x.permute(0, 2, 1, 3)
        # (batch_size, num_tokens, num_heads, num_hiddens/num_heads)
        batch_size = batch_size_times_num_heads // self.num_heads
        return x.reshape(batch_size, num_tokens, -1)
        # (batch_size, num_tokens, num_hiddens)

    def forward(self, queries, keys, values, mask=None):
        # queries:(batch_size, num_queries, query_size)
        # keys:(batch_size, num_kvs, key_size)
        # values:(batch_size, num_kvs, value_size)
        # mask:(batch_size, num_queries, num_kvs)
        # 线性变换
        queries = self.wq(queries)  # (batch_size, num_queries, num_hiddens)
        original_queries = queries
        keys = self.wk(keys)  # (batch_size, num_kvs, num_hiddens)
        values = self.wv(values)  # (batch_size, num_kvs, num_hiddens)
        # 多头拆分
        # (batch_size * num_heads, num_queries, num_hiddens/num_heads)
        queries = self.split_heads(queries)
        # (batch_size * num_heads, num_kvs, num_hiddens/num_heads)
        keys = self.split_heads(keys)
        # (batch_size * num_heads, num_kvs, num_hiddens/num_heads)
        values = self.split_heads(values)
        if mask is not None:
            # mask: (batch_size, num_queries, num_kvs)
            mask = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1)
            # (batch_size, num_heads, num_queries, num_kvs)
            mask = mask.view(-1, mask.shape[2], mask.shape[3])
            # (batch_size * num_heads, num_queries, num_kvs)
        # 计算注意力
        # (batch_size * num_heads, num_queries, num_hiddens/num_heads)
        scores = self.attention(queries, keys, values, mask)
        # 合并多头
        # (batch_size, num_queries, num_hiddens)
        scores = self.combine_heads(scores)
        # 线性变换
        outputs = self.wo(scores)  # (batch_size, num_queries, num_hiddens)
        # residual + layernorm
        outputs = outputs + original_queries
        outputs = self.layernorm(outputs)
        return outputs

位置编码（Positional Encoding）

$PE_{(pos,2i)}=sin(\dfrac{pos}{10000^{2i/d_{model}}})$

$PE_{(pos,2i+1)}=cos(\dfrac{pos}{10000^{2i/d_{model}}})$

其中， $d_{model}$ 表示输入向量的维度， $p os$ 是词的位置， $i$ 是维度索引。

使用这种位置编码的好处有：

随着序列长度的增加，PE并不会无限增加。
pos和pos+k的位置编码是线性关系。
class PositionEncoding(nn.Module):
    def __init__(self, num_hiddens, max_len=1000):
        super().__init__()
        self.pe = torch.zeros(1, max_len, num_hiddens)
        x = torch.arange(max_len, dtype=torch.float).reshape(-1, 1)
        x /= torch.pow(10000, torch.arange(0, num_hiddens,
                       2, dtype=torch.float) / num_hiddens)
        self.pe[:, :, 0::2] = torch.sin(x)
        self.pe[:, :, 1::2] = torch.cos(x)

    def forward(self, x):
        # x:(batch_size, num_tokens, num_hiddens)
        self.pe = self.pe.to(x.device)
        x = x + self.pe[:, :x.shape[1], :]
        return x

编码器（Encoder）

**编码器层：**包含一个自注意力机制和一个前馈网络，每个子层后接残差连接和层归一化。

Transformer的编码器是由多个相同的层叠加而成的，每个层都有两个子层。第一个子层是多头自注意力汇聚；第二个子层是基于位置的前馈网络（positionwise feed-forward network）。在计算编码器的自注意力时，查询、键和值都来自前一个编码器层的输出。每个子层都采用了残差连接。对于序列中任何位置的任何输入 $\in \mathbf{R^d}$ ，都要求满足 $\in \mathbf{R^d}$ ，以便残差连接满足 $\in \mathbf{R^d}$ 。在残差连接的加法计算之后，应用层归一化（layer normalization），输出一个d维向量。因此编码器中的任何层都不会改变其输入的形状。

编码器的作用：将输入序列（如文本、特征序列）逐层处理，抽取上下文相关的表示，最终输出包含全局语义信息的隐藏表示。

class EncoderBlock(nn.Module):
    def __init__(self, query_size, key_size, value_size, num_heads, num_hiddens, ffn_num_inputs, ffn_num_hiddens, ffn_outputs, dropout, norm_shape, bias=False):
        super().__init__()
        self.attention = MultiHeadAttention(
            query_size, key_size, value_size, num_heads, num_hiddens, bias)
        # (batch_size, num_queries, num_hiddens)
        self.addnorm1 = AddNorm(norm_shape, dropout)
        self.ffn = PositionWiseFFN(
            ffn_num_inputs, ffn_num_hiddens, ffn_outputs, bias)
        self.addnorm2 = AddNorm(norm_shape, dropout)

    def forward(self, x, mask=None):
        y = self.attention(x, x, x, mask)
        x = self.addnorm1(x, y)
        y = self.ffn(x)
        x = self.addnorm2(x, y)
        return x

解码器（Decoder）

Transformer解码器也是由多个相同的层叠加而成的。除了编码器中描述的两个子层之外，解码器还在这两个子层之间插入了第三个子层，称为编码器－解码器注意力（encoder-decoder attention）层。在解码器自注意力中，查询、键和值都来自上一个解码器层的输出。在编码器－解码器注意力中，查询来自当前解码器层的自注意力子层的输出，而键和值来自编码器最后一层的输出，注意，所有解码器共享同一份编码器输出。但是，解码器中的每个位置只能考虑自身和之前的所有位置。

hidden states 是指序列在某一层的向量化语义表示

阶段 Self-Attention Cross-Attention Fnn
训练 Q,K,V 都来自目标序列 hidden states（整段并行输入）；未来位置由 mask 屏蔽 Q 来自Self-Attention 输出的 hidden states；K,V 来自编码器输出输入为上一子层的输出
测试 Q 来自当前步 hidden state；K,V 是历史所有步的 hidden states（缓存累积）与训练时相同与训练时相同
作用建立目标序列内部的依赖关系，保证自回归让解码器在生成时参考源序列的全局语义非线性变换，增强表示能力

解码器的作用：在已有目标序列部分的条件下，结合编码器输出的上下文表示，逐步生成目标序列，保证生成的内容既符合输入语义，又符合序列生成的因果约束。

阶段	Self-Attention	Cross-Attention	Fnn
训练	Q,K,V 都来自目标序列 hidden states（整段并行输入）；未来位置由 mask 屏蔽	Q 来自Self-Attention 输出的 hidden states；K,V 来自编码器输出	输入为上一子层的输出
测试	Q 来自当前步 hidden state；K,V 是历史所有步的 hidden states（缓存累积）	与训练时相同	与训练时相同
作用	建立目标序列内部的依赖关系，保证自回归	让解码器在生成时参考源序列的全局语义	非线性变换，增强表示能力

class DecoderBlock(nn.Module):
    def __init__(self, i, query_size, key_size, value_size, num_heads, num_hiddens, ffn_num_inputs, ffn_num_hiddens, ffn_outputs, dropout, norm_shape, bias=False):
        super().__init__()
        self.i = i
        self.attention1 = MultiHeadAttention(
            query_size, key_size, value_size, num_heads, num_hiddens, bias)
        self.addnorm1 = AddNorm(norm_shape, dropout)
        self.attention2 = MultiHeadAttention(
            query_size, key_size, value_size, num_heads, num_hiddens, bias)
        self.addnorm2 = AddNorm(norm_shape, dropout)
        self.ffn = PositionWiseFFN(
            ffn_num_inputs, ffn_num_hiddens, ffn_outputs, bias)
        self.addnorm3 = AddNorm(norm_shape, dropout)

    def forward(self, x, state):
        # 训练阶段 x:(batch_size, num_steps, features) 
        # 每个样本有num_steps个词元，所有词元在同一时间步被处理
        # 预测阶段 x:(batch_size, 1, features)
        # 每个时间步只处理一个词元
        enc_outputs, enc_mask = state[0], state[1]
        if state[2][self.i] is None: # 训练阶段
            key_values = x #后续会被mask遮蔽，所以全部给出
        else: # 测试阶段
            key_values = torch.cat((state[2][self.i], x), dim=1)
        # state[2][i]
        state[2][self.i] = key_values
        if self.training:
            batch_size, num_steps, _ = x.shape
            # 生成掩码矩阵
            dec_mask = torch.triu(
                torch.ones((num_steps, num_steps), device=x.device), diagonal=1).bool()
            dec_mask = dec_mask.unsqueeze(0).expand(batch_size, -1, -1)
        else:
            dec_mask = None
        x1 = self.attention1(x, key_values, key_values, dec_mask)
        # x是目标序列或者上一个解码器的输出
        y1 = self.addnorm1(x, x1)
        x2 = self.attention2(y1, enc_outputs, enc_outputs, enc_mask)
        y2 = self.addnorm2(y1, x2)
        x3 = self.ffn(y2)
        y3 = self.addnorm3(y2, x3)
        return y3, state