自注意力机制------Transformer详细介绍-开发者社区

Transformer 是一种完全基于注意力机制（attention）构建的神经网络架构，由编码器和解码器两部分组成：编码器负责把输入序列的每个词元映射为融合全局上下文的表示，解码器则在此基础上逐步生成输出序列。它最核心的创新是用自注意力取代了 RNN 的循环结构，使序列中任意两个位置都能直接建立联系，从而既捕捉长距离依赖，又让整个计算可以高度并行，训练效率远超传统循环模型。凭借这一设计，Transformer 成为 GPT、BERT 等大模型的基础架构，也奠定了当今自然语言处理乃至多模态领域的主流范式。
看着很复杂，但是原理其实很简单，我们详细介绍。

一.自注意力机制和位置编码

1.1自注意力机制

自注意力机制（self-attention）是让序列中每个位置都能“看”到同一序列里所有其他位置、并根据相关性动态分配权重来聚合信息的一种计算方式。它把每个词元分别线性变换成查询（Query）、键（Key）、值（Value）三个向量，用当前词的 Q 与所有词的 K 做点积并除以 √d 缩放、经 softmax 归一化得到注意力权重，再用这些权重对所有词的 V 加权求和，作为该位置的新表示——这样一来，每个词的表示都融合了全局上下文，且所有位置的计算可完全并行，不像 RNN 那样必须逐步循环，同时通过 mask 还能灵活控制“哪些位置可见”（如解码器里遮住未来词元），是 Transformer 的核心组件。

1.2位置编码

在处理词元序列时，循环神经网络是逐个的重复地处理词元的，而自注意力则因为并行计算而放弃了顺序操作。为了使用序列的顺序信息，通过在输入表示中添加位置编码（positional encoding）来注入绝对的或相对的位置信息。位置编码可以通过学习得到也可以直接固定得到。

这里我们介绍基于正弦函数和余弦函数的固定位置编码：

#@save class PositionalEncoding(nn.Module): """位置编码""" def __init__(self, num_hiddens, dropout, max_len=1000): super(PositionalEncoding, self).__init__() self.dropout = nn.Dropout(dropout) # 创建一个足够长的P self.P = torch.zeros((1, max_len, num_hiddens)) X = torch.arange(max_len, dtype=torch.float32).reshape( -1, 1) / torch.pow(10000, torch.arange( 0, num_hiddens, 2, dtype=torch.float32) / num_hiddens) self.P[:, :, 0::2] = torch.sin(X) self.P[:, :, 1::2] = torch.cos(X) def forward(self, X): X = X + self.P[:, :X.shape[1], :].to(X.device) return self.dropout(X)

除了捕获绝对位置信息之外，上述的位置编码还允许模型学习得到输入序列中相对位置信息。这是因为对于任何确定的位置偏移a，位置处于i的位置编码可以线性投影位置处i+a的位置编码来表示。

利用正弦和余弦函数的数学特性，模型可以通过一个固定的线性变换（矩阵乘法），极其轻松地计算出两个词之间的“相对距离”。

二.Transformer架构

左右两边分别是编码器和解码器的俩个块，transformer由多个相同的块叠加而成。

2.1逐位前馈网络

#@save class PositionWiseFFN(nn.Module): """基于位置的前馈网络""" def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs, **kwargs): super(PositionWiseFFN, self).__init__(**kwargs) self.dense1 = nn.Linear(ffn_num_input, ffn_num_hiddens) self.relu = nn.ReLU() self.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs) def forward(self, X): return self.dense2(self.relu(self.dense1(X)))

2.2残差链接和规范化

注意层规范化和批量规范化的区别：

#@save class AddNorm(nn.Module): """残差连接后进行层规范化""" def __init__(self, normalized_shape, dropout, **kwargs): super(AddNorm, self).__init__(**kwargs) self.dropout = nn.Dropout(dropout) self.ln = nn.LayerNorm(normalized_shape) def forward(self, X, Y): return self.ln(self.dropout(Y) + X)

2.3编码器

编码器由多个块组成，每个块代码如下：

#@save class EncoderBlock(nn.Module): """Transformer编码器块""" def __init__(self, key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, num_heads, dropout, use_bias=False, **kwargs): super(EncoderBlock, self).__init__(**kwargs) self.attention = d2l.MultiHeadAttention( key_size, query_size, value_size, num_hiddens, num_heads, dropout, use_bias) self.addnorm1 = AddNorm(norm_shape, dropout) self.ffn = PositionWiseFFN( ffn_num_input, ffn_num_hiddens, num_hiddens) self.addnorm2 = AddNorm(norm_shape, dropout) def forward(self, X, valid_lens): Y = self.addnorm1(X, self.attention(X, X, X, valid_lens)) return self.addnorm2(Y, self.ffn(Y))

整体代码如下：

#@save class TransformerEncoder(d2l.Encoder): """Transformer编码器""" def __init__(self, vocab_size, key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, num_heads, num_layers, dropout, use_bias=False, **kwargs): super(TransformerEncoder, self).__init__(**kwargs) self.num_hiddens = num_hiddens self.embedding = nn.Embedding(vocab_size, num_hiddens) self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout) self.blks = nn.Sequential() for i in range(num_layers): self.blks.add_module("block"+str(i), EncoderBlock(key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, num_heads, dropout, use_bias)) def forward(self, X, valid_lens, *args):# 因为位置编码值在-1和1之间， # 因此嵌入值乘以嵌入维度的平方根进行缩放， # 然后再与位置编码相加。X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens)) self.attention_weights = [None] * len(self.blks) for i, blk in enumerate(self.blks): X = blk(X, valid_lens) self.attention_weights[ i] = blk.attention.attention.attention_weights return X

2.4解码器

块代码：

class DecoderBlock(nn.Module): """解码器中第i个块""" def __init__(self, key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, num_heads, dropout, i, **kwargs): super(DecoderBlock, self).__init__(**kwargs) self.i = i self.attention1 = d2l.MultiHeadAttention( key_size, query_size, value_size, num_hiddens, num_heads, dropout) self.addnorm1 = AddNorm(norm_shape, dropout) self.attention2 = d2l.MultiHeadAttention( key_size, query_size, value_size, num_hiddens, num_heads, dropout) self.addnorm2 = AddNorm(norm_shape, dropout) self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens, num_hiddens) self.addnorm3 = AddNorm(norm_shape, dropout) def forward(self, X, state): enc_outputs, enc_valid_lens = state[0], state[1] # 训练阶段，输出序列的所有词元都在同一时间处理， # 因此state[2][self.i]初始化为None。 # 预测阶段，输出序列是通过词元一个接着一个解码的， # 因此state[2][self.i]包含着直到当前时间步第i个块解码的输出表示 if state[2][self.i] is None: key_values = X else: key_values = torch.cat((state[2][self.i], X), axis=1) state[2][self.i] = key_values if self.training: batch_size, num_steps, _ = X.shape # dec_valid_lens的开头:(batch_size,num_steps), # 其中每一行是[1,2,...,num_steps] dec_valid_lens = torch.arange( 1, num_steps + 1, device=X.device).repeat(batch_size, 1) else: dec_valid_lens = None # 自注意力 X2 = self.attention1(X, key_values, key_values, dec_valid_lens) Y = self.addnorm1(X, X2) # 编码器－解码器注意力。 # enc_outputs的开头:(batch_size,num_steps,num_hiddens) Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_lens) Z = self.addnorm2(Y, Y2) return self.addnorm3(Z, self.ffn(Z)), state

这里介绍一下state:

要关注的是有俩个注意力层，底下的自注意力层，上面的是用的解码器输出状态的注意力层。

对于自注意力层，情况主要分为预测和训练，预测时自注意力层q=x,key,value=x和已经生成的序列。训练时q=x, key和 value=x和当前时间步能看到的序列。所以训练是有效长度要单独处理，预测时有效长度为none就可以。

用的解码器输出状态的注意力层，q=y(自注意力层输出后的残差连接和规范化)，key,value=编码器输出状态。

X (B, T, H) ──────────────────────────────┐
│ │(残差)
│ │
▼ │
┌──────────────────────────────────────────┐ │
│ attention1 Q=X (B,T,H) │ │
│ K=V=key_values (B,T,H) │ │
│ mask=dec_valid_lens (B,T) │ │
│ 因果掩码: [1,2,...,T] │ │
└──────────────────────────────────────────┘ │
│ X2 (B, T, H) │
▼ │
┌─────────┐ │
│AddNorm1 │ ◄─────────────────────────────┘
└─────────┘
│ Y (B, T, H) ───────────────────┐
│ │(残差)
│ │
▼ │
┌──────────────────────────────────────────┐ │
│ attention2 Q=Y (B,T,H) │ │
│ K=V=enc_outputs (B,S,H) │ │
│ mask=enc_valid_lens (B,) │ │
│ padding 掩码 │ │
└──────────────────────────────────────────┘ │
│ Y2 (B, T, H) │
▼ │
┌─────────┐ │
│AddNorm2 │ ◄──────────────────────────┘
└─────────┘
│ Z (B, T, H) ─────────────────┐
│ │(残差)
▼ │
┌─────────┐ │
│ FFN │ 逐位置全连接 │
└─────────┘ │
│ ffn(Z) (B, T, H) │
▼ │
┌─────────┐ │
│AddNorm3 │ ◄────────────────────────┘
└─────────┘
│
▼
输出 (B, T, H) + state
（state[2][i] 仍为 X，训练时无缓存增长）

整体代码：

class TransformerDecoder(d2l.AttentionDecoder): def __init__(self, vocab_size, key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, num_heads, num_layers, dropout, **kwargs): super(TransformerDecoder, self).__init__(**kwargs) self.num_hiddens = num_hiddens self.num_layers = num_layers self.embedding = nn.Embedding(vocab_size, num_hiddens) self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout) self.blks = nn.Sequential() for i in range(num_layers): self.blks.add_module("block"+str(i), DecoderBlock(key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, num_heads, dropout, i)) self.dense = nn.Linear(num_hiddens, vocab_size) def init_state(self, enc_outputs, enc_valid_lens, *args): return [enc_outputs, enc_valid_lens, [None] * self.num_layers] def forward(self, X, state): X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens)) self._attention_weights = [[None] * len(self.blks) for _ in range (2)] for i, blk in enumerate(self.blks): X, state = blk(X, state) # 解码器自注意力权重 self._attention_weights[0][ i] = blk.attention1.attention.attention_weights # “编码器－解码器”自注意力权重 self._attention_weights[1][ i] = blk.attention2.attention.attention_weights return self.dense(X), state @property def attention_weights(self): return self._attention_weights

这里初始化代码：