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

1. Transformer模型基础回顾Transformer模型最早由Google团队在2017年发表的论文《Attention Is All You Need》中提出彻底改变了自然语言处理领域的格局。与传统RNN和LSTM不同Transformer完全基于自注意力机制(self-attention mechanism)能够并行处理整个序列显著提升了训练效率。核心组件包括多头注意力机制(Multi-Head Attention)允许模型同时关注不同位置的表示子空间位置编码(Positional Encoding)为输入序列注入位置信息前馈神经网络(Feed Forward Network)对每个位置进行独立变换残差连接(Residual Connection)和层归一化(Layer Normalization)促进深层网络训练2. 自注意力机制深度解析2.1 缩放点积注意力自注意力机制的核心计算过程可以用以下公式表示Attention(Q, K, V) softmax(QK^T/√d_k)V其中Q(Query)、K(Key)、V(Value)分别表示查询、键和值矩阵d_k是键向量的维度√d_k的缩放因子防止点积结果过大导致softmax梯度消失实际实现时我们通常会使用矩阵运算来批量处理整个序列def scaled_dot_product_attention(q, k, v, maskNone): matmul_qk tf.matmul(q, k, transpose_bTrue) # (..., seq_len_q, seq_len_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) # (..., seq_len_q, seq_len_k) output tf.matmul(attention_weights, v) # (..., seq_len_q, depth_v) return output, attention_weights2.2 多头注意力实现多头注意力将Q、K、V通过不同的线性变换投影到多个子空间使模型能够关注不同方面的信息class MultiHeadAttention(tf.keras.layers.Layer): def __init__(self, d_model, num_heads): super(MultiHeadAttention, self).__init__() self.num_heads num_heads self.d_model d_model assert d_model % self.num_heads 0 self.depth d_model // self.num_heads self.wq tf.keras.layers.Dense(d_model) self.wk tf.keras.layers.Dense(d_model) self.wv tf.keras.layers.Dense(d_model) self.dense tf.keras.layers.Dense(d_model) def split_heads(self, x, batch_size): x tf.reshape(x, (batch_size, -1, self.num_heads, self.depth)) return tf.transpose(x, perm[0, 2, 1, 3]) def call(self, v, k, q, maskNone): batch_size tf.shape(q)[0] q self.wq(q) # (batch_size, seq_len, d_model) k self.wk(k) v self.wv(v) q self.split_heads(q, batch_size) # (batch_size, num_heads, seq_len_q, depth) k self.split_heads(k, batch_size) v self.split_heads(v, batch_size) scaled_attention, attention_weights scaled_dot_product_attention( q, k, v, mask) scaled_attention tf.transpose(scaled_attention, perm[0, 2, 1, 3]) # (batch_size, seq_len_q, num_heads, depth) concat_attention tf.reshape(scaled_attention, (batch_size, -1, self.d_model)) # (batch_size, seq_len_q, d_model) output self.dense(concat_attention) return output, attention_weights提示在实际应用中多头注意力的头数(num_heads)通常选择4-16之间需要确保d_model能被num_heads整除。3. Transformer编码器实现3.1 位置编码由于Transformer不包含循环和卷积结构需要显式地注入位置信息def get_angles(pos, i, d_model): angle_rates 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model)) return pos * angle_rates def positional_encoding(position, d_model): angle_rads get_angles(np.arange(position)[:, np.newaxis], np.arange(d_model)[np.newaxis, :], d_model) # 对数组中的偶数索引应用sin函数 angle_rads[:, 0::2] np.sin(angle_rads[:, 0::2]) # 对数组中的奇数索引应用cos函数 angle_rads[:, 1::2] np.cos(angle_rads[:, 1::2]) pos_encoding angle_rads[np.newaxis, ...] return tf.cast(pos_encoding, dtypetf.float32)3.2 编码器层实现单个编码器层包含多头注意力机制和前馈神经网络class EncoderLayer(tf.keras.layers.Layer): def __init__(self, d_model, num_heads, dff, rate0.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(epsilon1e-6) self.layernorm2 tf.keras.layers.LayerNormalization(epsilon1e-6) self.dropout1 tf.keras.layers.Dropout(rate) self.dropout2 tf.keras.layers.Dropout(rate) def call(self, x, training, maskNone): attn_output, _ self.mha(x, x, x, mask) # (batch_size, input_seq_len, d_model) attn_output self.dropout1(attn_output, trainingtraining) out1 self.layernorm1(x attn_output) # (batch_size, input_seq_len, d_model) ffn_output self.ffn(out1) # (batch_size, input_seq_len, d_model) ffn_output self.dropout2(ffn_output, trainingtraining) out2 self.layernorm2(out1 ffn_output) # (batch_size, input_seq_len, d_model) return out24. Transformer解码器实现4.1 解码器层结构解码器层比编码器层更复杂包含两个多头注意力机制class DecoderLayer(tf.keras.layers.Layer): def __init__(self, d_model, num_heads, dff, rate0.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(epsilon1e-6) self.layernorm2 tf.keras.layers.LayerNormalization(epsilon1e-6) self.layernorm3 tf.keras.layers.LayerNormalization(epsilon1e-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_maskNone, padding_maskNone): # enc_output.shape (batch_size, input_seq_len, d_model) attn1, attn_weights_block1 self.mha1(x, x, x, look_ahead_mask) # (batch_size, target_seq_len, d_model) attn1 self.dropout1(attn1, trainingtraining) out1 self.layernorm1(attn1 x) attn2, attn_weights_block2 self.mha2( enc_output, enc_output, out1, padding_mask) # (batch_size, target_seq_len, d_model) attn2 self.dropout2(attn2, trainingtraining) out2 self.layernorm2(attn2 out1) # (batch_size, target_seq_len, d_model) ffn_output self.ffn(out2) # (batch_size, target_seq_len, d_model) ffn_output self.dropout3(ffn_output, trainingtraining) out3 self.layernorm3(ffn_output out2) # (batch_size, target_seq_len, d_model) return out3, attn_weights_block1, attn_weights_block24.2 解码器中的注意力掩码解码器需要两种类型的注意力掩码填充掩码(Padding Mask)忽略填充标记前瞻掩码(Look-ahead Mask)防止解码器看到未来信息def create_padding_mask(seq): seq tf.cast(tf.math.equal(seq, 0), tf.float32) return seq[:, tf.newaxis, tf.newaxis, :] # (batch_size, 1, 1, seq_len) def create_look_ahead_mask(size): mask 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0) return mask # (seq_len, seq_len)5. Transformer训练技巧5.1 学习率调度Transformer使用特殊的学习率调度策略在训练初期快速升温然后按步数平方根的倒数衰减class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule): def __init__(self, d_model, warmup_steps4000): 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)5.2 标签平滑为改善模型泛化能力可以使用标签平滑技术def loss_function(real, pred): mask tf.math.logical_not(tf.math.equal(real, 0)) loss_object tf.keras.losses.SparseCategoricalCrossentropy( from_logitsTrue, reductionnone) loss loss_object(real, pred) mask tf.cast(mask, dtypeloss.dtype) loss * mask return tf.reduce_sum(loss)/tf.reduce_sum(mask)6. Transformer变体与改进6.1 模型压缩技术知识蒸馏使用大模型(teacher)指导小模型(student)训练量化感知训练在训练中模拟量化过程剪枝移除不重要的注意力头或权重6.2 高效注意力机制稀疏注意力限制每个位置只能关注局部区域线性注意力将softmax注意力近似为线性变换内存压缩使用低秩近似减少内存消耗7. Transformer应用实践7.1 文本分类实现class TransformerClassifier(tf.keras.Model): def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, maximum_position_encoding, rate0.1, num_classes2): super(TransformerClassifier, self).__init__() self.embedding tf.keras.layers.Embedding(input_vocab_size, d_model) self.pos_encoding positional_encoding(maximum_position_encoding, d_model) self.enc_layers [EncoderLayer(d_model, num_heads, dff, rate) for _ in range(num_layers)] self.dropout tf.keras.layers.Dropout(rate) self.final_layer tf.keras.layers.Dense(num_classes, activationsoftmax) def call(self, x, training, maskNone): seq_len tf.shape(x)[1] x self.embedding(x) # (batch_size, input_seq_len, d_model) x * tf.math.sqrt(tf.cast(self.d_model, tf.float32)) x self.pos_encoding[:, :seq_len, :] x self.dropout(x, trainingtraining) for i in range(self.num_layers): x self.enc_layers[i](x, training, mask) # 使用[CLS]位置的表示进行分类 cls_output x[:, 0, :] return self.final_layer(cls_output)7.2 超参数选择经验根据实际项目经验以下配置在大多数NLP任务中表现良好参数小模型中模型大模型d_model1285121024num_layers4612num_heads4816dff51220484096batch_size3264128dropout0.10.10.2注意实际应用中需要根据具体任务和计算资源调整这些参数。较小的batch size配合梯度累积可以达到与大batch size相似的效果。