Bootstrap

Transformer细节(六)——详解Transformer各层结构和组成及代码实现

Transformer 模型的架构是由多个编码器(Encoder)和解码器(Decoder)层堆叠而成的。

一、编码器(Encoder)

        编码器由多个相同的编码器层(Encoder Layer)堆叠而成。每个编码器层包含两个主要子层:自注意力(Self-Attention)子层和前馈神经网络(Feed Forward Neural Network,FFN)子层。

1. 自注意力子层(Self-Attention Layer)

        输入与输出维度

         输入张量的形状为 \((N, L, d_{\text{model}})\)

        计算过程

        注意力机制:通过查询(Query)、键(Key)和值(Value)来计算注意力得分,公式为:

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

        多头注意力(Multi-Head Attention)

        将输入分成多头,每个头独立计算注意力,再将结果拼接并通过线性变换。多头注意力公式为:\[ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)W^O \]
        其中,每个头的计算为:\[ \text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) \]

        层归一化(Layer Normalization)

        在自注意力操作之后进行归一化,并添加残差连接(Residual Connection)。

2. 前馈神经网络子层(Feed Forward Neural Network, FFN)

        结构

        由两个线性变换和一个激活函数(通常是ReLU)组成,公式为:\[ \text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2 \]

        层归一化(Layer Normalization)

        在前馈神经网络之后进行归一化,并添加残差连接。

3. 输出维度

        每个编码器层的输出维度仍然是 \((N, L, d_{\text{model}})\)。

二、解码器(Decoder)

        解码器也由多个相同的解码器层(Decoder Layer)堆叠而成。解码器层在编码器层的基础上增加了一个编码器-解码器注意力(Encoder-Decoder Attention)子层。

1. 掩码自注意力子层(Masked Self-Attention Layer)

        类似于编码器的自注意力子层,但在计算注意力得分时使用了掩码(mask),以防止当前位置看到未来的信息。

2. 编码器-解码器注意力子层(Encoder-Decoder Attention Layer)

        该子层的输入来自编码器层的输出和解码器层的前一层输出。计算过程与自注意力相同,但查询(Query)来自解码器,键(Key)和值(Value)来自编码器。

3. 前馈神经网络子层(Feed Forward Neural Network, FFN)

        与编码器层中的前馈神经网络相同。

4. 输出维度

        每个解码器层的输出维度也是 \((N, L, d_{\text{model}})\)。

解码器更多细节,参考我之前的博客Transformer细节(五)——详解Transformer解码器的自注意力层和编码器-解码器注意力层数据处理机制-CSDN博客

三、Transformer 模型的整体结构

1. 输入嵌入(Input Embedding)

        输入序列首先通过嵌入层(Embedding Layer)映射为 \((N, L, d_{\text{model}})\) 的张量。
        添加位置编码(Positional Encoding),以便模型捕获序列中的位置信息。

2. 编码器堆叠(Stack of Encoders)

        输入经过多个编码器层的处理,得到编码后的序列表示。

3. 解码器堆叠(Stack of Decoders)

        解码器接收编码器的输出和目标序列的嵌入,经过多个解码器层的处理,生成最终的输出序列。

4. 输出层(Output Layer)

        解码器的输出通过线性变换和 softmax 层,生成每个位置的预测分布。

四、代码实现

import torch
import torch.nn as nn
import torch.nn.functional as F

class Transformer(nn.Module):
    def __init__(self, num_layers, d_model, num_heads, d_ff):
        super(Transformer, self).__init__()
        self.encoder_layers = nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff) for _ in range(num_layers)])
        self.decoder_layers = nn.ModuleList([DecoderLayer(d_model, num_heads, d_ff) for _ in range(num_layers)])
        
    def forward(self, src, tgt, src_mask, tgt_mask):
        enc_output = src
        
        # 编码器层的前向传播
        for layer in self.encoder_layers:
            enc_output = layer(enc_output, src_mask)
        
        dec_output = tgt
        
        # 解码器层的前向传播
        for layer in self.decoder_layers:
            dec_output = layer(dec_output, enc_output, src_mask, tgt_mask)
        
        return dec_output

# 简化的编码器层和解码器层
class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff):
        super(EncoderLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(d_model, num_heads)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model)
        )
        self.layer_norm1 = nn.LayerNorm(d_model)
        self.layer_norm2 = nn.LayerNorm(d_model)

    def forward(self, x, src_mask):
        # 自注意力
        self_attn_output, _ = self.self_attn(x, x, x, attn_mask=src_mask)
        x = x + self_attn_output
        x = self.layer_norm1(x)
        
        # 前馈神经网络
        ff_output = self.feed_forward(x)
        x = x + ff_output
        x = self.layer_norm2(x)
        
        return x

class DecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff):
        super(DecoderLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(d_model, num_heads)
        self.cross_attn = nn.MultiheadAttention(d_model, num_heads)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model)
        )
        self.layer_norm1 = nn.LayerNorm(d_model)
        self.layer_norm2 = nn.LayerNorm(d_model)
        self.layer_norm3 = nn.LayerNorm(d_model)

    def forward(self, x, enc_output, src_mask, tgt_mask):
        # 自注意力
        self_attn_output, _ = self.self_attn(x, x, x, attn_mask=tgt_mask)
        x = x + self_attn_output
        x = self.layer_norm1(x)
        
        # 编码器-解码器注意力
        cross_attn_output, _ = self.cross_attn(x, enc_output, enc_output, attn_mask=src_mask)
        x = x + cross_attn_output
        x = self.layer_norm2(x)
        
        # 前馈神经网络
        ff_output = self.feed_forward(x)
        x = x + ff_output
        x = self.layer_norm3(x)
        
        return x

# 示例参数
num_layers = 6
d_model = 512
num_heads = 8
d_ff = 2048
sequence_length = 10
batch_size = 32

# 初始化Transformer模型
transformer = Transformer(num_layers, d_model, num_heads, d_ff)

# 生成示例输入
src = torch.rand(sequence_length, batch_size, d_model)  # 编码器输入
tgt = torch.rand(sequence_length, batch_size, d_model)  # 解码器输入
src_mask = torch.ones(sequence_length, sequence_length).bool()  # 编码器掩码
tgt_mask = torch.ones(sequence_length, sequence_length).bool()  # 解码器掩码

# 前向传播
output = transformer(src, tgt, src_mask, tgt_mask)

print(output.shape)  # 应输出: torch.Size([10, 32, 512])

五、总结

        这种架构设计使得 Transformer 模型可以高效并行地处理序列数据,并且能够捕获序列中的长距离依赖关系。具体如何并行处理数据,参考我之前的博客Transformer细节(三)——Transformer是并行or顺序处理数据?-CSDN博客

 

;