注意力机制与 Transformer 模型实战详解

学习目标与重点

掌握注意力机制的核心原理、经典算法，以及 Transformer 模型的架构设计与实战应用。重点理解自注意力与多头注意力的计算逻辑，学会使用 TensorFlow 搭建 Transformer 模型，完成机器翻译任务。

注意力机制的核心思想

为什么需要注意力机制

传统的 RNN 和 LSTM 在处理长序列时，存在长距离依赖捕捉能力不足和并行计算效率低的问题。注意力机制的出现，解决了这两个核心痛点。

注意力机制的本质是让模型学会'聚焦'——在处理序列数据时，自动分配不同的权重给输入序列中的各个元素，重点关注与当前任务相关的信息，弱化无关信息的干扰。比如在机器翻译任务中，翻译'我爱中国'时，模型会给'我''爱''中国'分配不同的注意力权重，从而更精准地生成对应的英文翻译。

注意力机制的基本框架

注意力机制的计算通常包含查询（Query）、键（Key）、值（Value）三个核心要素，简称 QKV 框架。其计算流程可以总结为三步：

计算 Query 和所有 Key 的相似度，得到注意力分数
对注意力分数进行归一化处理（常用 Softmax 函数），得到注意力权重
用归一化后的权重对 Value 进行加权求和，得到最终的注意力输出

基础注意力计算公式如下：

$$ \text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$

其中 $d_k$ 是 Key 的维度，除以 $\sqrt{d_k}$ 是为了防止内积结果过大，导致 Softmax 函数饱和。

import tensorflow as tf
import numpy as np

# 实现基础注意力计算
def scaled_dot_product_attention(q, k, v, mask=None):
    # 计算 Q 和 K 的点积
    matmul_qk = tf.matmul(q, k, transpose_b=True)
    # 获取 k 的维度
    dk = tf.cast(tf.shape(k)[-1], tf.float32)
    # 缩放点积
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)
    # 应用掩码（可选）
    if mask is not None:
        scaled_attention_logits += (mask * -1e9)
    # 计算注意力权重
    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)
    # 计算最终输出
    output = tf.matmul(attention_weights, v)
    return output, attention_weights

# 模拟输入：批次大小=2，序列长度=3，特征维度=4
q = tf.random.normal((, , ))
k = tf.random.normal((, , ))
v = tf.random.normal((, , ))
output, attn_weights = scaled_dot_product_attention(q, k, v)
(, output.shape)
(, attn_weights.shape)

# 定义前馈神经网络 def point_wise_feed_forward_network(d_model, dff): return tf.keras.Sequential([ tf.keras.layers.Dense(dff, activation='relu'), tf.keras.layers.Dense(d_model) ]) # 定义编码层 class EncoderLayer(tf.keras.layers.Layer): def __init__(self, d_model, num_heads, dff, rate=0.1): super(EncoderLayer, self).__init__() self.mha = MultiHeadAttention(d_model, num_heads) self.ffn = point_wise_feed_forward_network(d_model, dff) self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6) self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6) self.dropout1 = tf.keras.layers.Dropout(rate) self.dropout2 = tf.keras.layers.Dropout(rate) def call(self, x, training, mask): attn_output, _ = self.mha(x, x, x, mask) attn_output = self.dropout1(attn_output, training=training) out1 = self.layernorm1(x + attn_output) ffn_output = self.ffn(out1) ffn_output = self.dropout2(ffn_output, training=training) out2 = self.layernorm2(out1 + ffn_output) return out2 # 定义解码器层 class DecoderLayer(tf.keras.layers.Layer): def __init__(self, d_model, num_heads, dff, rate=0.1): super(DecoderLayer, self).__init__() self.mha1 = MultiHeadAttention(d_model, num_heads) self.mha2 = MultiHeadAttention(d_model, num_heads) self.ffn = point_wise_feed_forward_network(d_model, dff) self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6) self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6) self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6) self.dropout1 = tf.keras.layers.Dropout(rate) self.dropout2 = tf.keras.layers.Dropout(rate) self.dropout3 = tf.keras.layers.Dropout(rate) def call(self, x, enc_output, training, look_ahead_mask, padding_mask): attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask) attn1 = self.dropout1(attn1, training=training) out1 = self.layernorm1(attn1 + x) attn2, attn_weights_block2 = self.mha2(enc_output, enc_output, out1, padding_mask) attn2 = self.dropout2(attn2, training=training) out2 = self.layernorm2(attn2 + out1) ffn_output = self.ffn(out2) ffn_output = self.dropout3(ffn_output, training=training) out3 = self.layernorm3(ffn_output + out2) return out3, attn_weights_block1, attn_weights_block2 # 定义完整编码器 class Encoder(tf.keras.layers.Layer): def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, maximum_position_encoding, rate=0.1): super(Encoder, self).__init__() self.d_model = d_model self.num_layers = num_layers self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model) self.pos_encoding = positional_encoding(maximum_position_encoding, self.d_model) self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate) for _ in range(num_layers)] self.dropout = tf.keras.layers.Dropout(rate) def call(self, x, training, mask): seq_len = tf.shape(x)[1] x = self.embedding(x) x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32)) x += self.pos_encoding[:,:seq_len,:] x = self.dropout(x, training=training) for i in range(self.num_layers): x = self.enc_layers[i](x, training, mask) return x # 定义完整解码器 class Decoder(tf.keras.layers.Layer): def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size, maximum_position_encoding, rate=0.1): super(Decoder, self).__init__() self.d_model = d_model self.num_layers = num_layers self.embedding = tf.keras.layers.Embedding(target_vocab_size, d_model) self.pos_encoding = positional_encoding(maximum_position_encoding, d_model) self.dec_layers = [DecoderLayer(d_model, num_heads, dff, rate) for _ in range(num_layers)] self.dropout = tf.keras.layers.Dropout(rate) def call(self, x, enc_output, training, look_ahead_mask, padding_mask): seq_len = tf.shape(x)[1] attention_weights = {} x = self.embedding(x) x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32)) x += self.pos_encoding[:,:seq_len,:] x = self.dropout(x, training=training) for i in range(self.num_layers): x, block1, block2 = self.dec_layers[i](x, enc_output, training, look_ahead_mask, padding_mask) attention_weights['decoder_layer{}_block1'.format(i+1)] = block1 attention_weights['decoder_layer{}_block2'.format(i+1)] = block2 return x, attention_weights # 定义完整 Transformer 模型 class Transformer(tf.keras.Model): def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, target_vocab_size, pe_input, pe_target, rate=0.1): super(Transformer, self).__init__() self.encoder = Encoder(num_layers, d_model, num_heads, dff, input_vocab_size, pe_input, rate) self.decoder = Decoder(num_layers, d_model, num_heads, dff, target_vocab_size, pe_target, rate) self.final_layer = tf.keras.layers.Dense(target_vocab_size) def call(self, inp, tar, training, enc_padding_mask, look_ahead_mask, dec_padding_mask): enc_output = self.encoder(inp, training, enc_padding_mask) dec_output, attention_weights = self.decoder(tar, enc_output, training, look_ahead_mask, dec_padding_mask) final_output = self.final_layer(dec_output) return final_output, attention_weights # 设置模型参数 num_layers = 4 d_model = 128 dff = 512 num_heads = 8 input_vocab_size = tokenizer_en.vocab_size + 2 target_vocab_size = tokenizer_fr.vocab_size + 2 dropout_rate = 0.1 # 初始化模型 transformer = Transformer(num_layers, d_model, num_heads, dff, input_vocab_size, target_vocab_size, pe_input=1000, pe_target=1000, rate=dropout_rate)

# 定义学习率调度器 class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule): def __init__(self, d_model, warmup_steps=4000): super(CustomSchedule, self).__init__() self.d_model = d_model self.d_model = tf.cast(self.d_model, tf.float32) self.warmup_steps = warmup_steps def __call__(self, step): arg1 = tf.math.rsqrt(step) arg2 = step * (self.warmup_steps ** -1.5) return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2) # 初始化学习率和优化器 learning_rate = CustomSchedule(d_model) optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9) # 定义损失函数和评估指标 loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none') def loss_function(real, pred): mask = tf.math.logical_not(tf.math.equal(real, 0)) loss_ = loss_object(real, pred) mask = tf.cast(mask, dtype=loss_.dtype) loss_ *= mask return tf.reduce_sum(loss_) / tf.reduce_sum(mask) train_loss = tf.keras.metrics.Mean(name='train_loss') train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy') # 定义训练步骤 @tf.function def train_step(inp, tar): tar_inp = tar[:,:-1] tar_real = tar[:,1:] enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp, tar_inp) with tf.GradientTape() as tape: predictions, _ = transformer(inp, tar_inp, True, enc_padding_mask, combined_mask, dec_padding_mask) loss = loss_function(tar_real, predictions) gradients = tape.gradient(loss, transformer.trainable_variables) optimizer.apply_gradients(zip(gradients, transformer.trainable_variables)) train_loss(loss) train_accuracy(tar_real, predictions) # 定义掩码生成函数 def create_masks(inp, tar): enc_padding_mask = create_padding_mask(inp) dec_padding_mask = create_padding_mask(inp) look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1]) dec_target_padding_mask = create_padding_mask(tar) combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask) return enc_padding_mask, combined_mask, dec_padding_mask def create_padding_mask(seq): seq = tf.cast(tf.math.equal(seq, 0), tf.float32) return seq[:, tf.newaxis, tf.newaxis, :] def create_look_ahead_mask(size): mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0) return mask # 开始训练 EPOCHS = 20 for epoch in range(EPOCHS): train_loss.reset_states() train_accuracy.reset_states() for(batch,(inp, tar)) in enumerate(train_dataset): train_step(inp, tar) if batch % 50 == 0: print(f'Epoch {epoch+1} Batch {batch} Loss {train_loss.result():.4f} Accuracy {train_accuracy.result():.4f}') print(f'Epoch {epoch+1} Loss {train_loss.result():.4f} Accuracy {train_accuracy.result():.4f}')

注意力机制与 Transformer 模型实战详解