连接Transformer块中的注意力层和线性层

约 2428 字大约 8 分钟

设计思路与核心概念

1. Transformer块的背景与动机

Transformer块是现代大语言模型的核心组件，由Vaswani等人在2017年的"Attention Is All You Need"论文中提出，主要解决以下问题：

序列建模效率：相比RNN/LSTM，Transformer可以并行处理序列中的所有位置
长距离依赖：通过自注意力机制有效捕获长距离的依赖关系
信息流动：通过残差连接和层归一化确保信息的有效传播
表达能力：结合注意力机制和前馈网络提供强大的表达能力

2. 核心设计思想

Transformer块的核心思想是将自注意力机制与前馈网络相结合：

多头注意力：捕获不同子空间的注意力模式
前馈网络：提供非线性变换和特征提取能力
残差连接：确保梯度有效传播，避免梯度消失
层归一化：稳定训练过程，加速收敛

3. 架构组成

Transformer块的标准结构

输入 → 层归一化 → 多头注意力 → Dropout → 残差连接 →
     → 层归一化 → 前馈网络 → Dropout → 残差连接 → 输出

数学表示

对于输入序列 $X$ ，Transformer块的计算过程为：

注意力子层：

X_1 = X + \text{Dropout}(\text{MultiHeadAttention}(\text{LayerNorm}(X)))

前馈子层：

X_2 = X_1 + \text{Dropout}(\text{FeedForward}(\text{LayerNorm}(X_1)))

执行流程

1. 整体执行流程图

2. 详细计算流程图

计算步骤详解

步骤	操作	输入形状	输出形状	说明
1	保存残差1	[batch, seq_len, emb_dim]	[batch, seq_len, emb_dim]	为第一个残差连接保存输入
2	层归一化1	[batch, seq_len, emb_dim]	[batch, seq_len, emb_dim]	标准化输入特征
3	多头注意力	[batch, seq_len, emb_dim]	[batch, seq_len, emb_dim]	计算自注意力
4	Dropout1	[batch, seq_len, emb_dim]	[batch, seq_len, emb_dim]	防止过拟合
5	残差连接1	[batch, seq_len, emb_dim]	[batch, seq_len, emb_dim]	添加残差连接
6	保存残差2	[batch, seq_len, emb_dim]	[batch, seq_len, emb_dim]	为第二个残差连接保存输入
7	层归一化2	[batch, seq_len, emb_dim]	[batch, seq_len, emb_dim]	标准化特征
8	前馈网络	[batch, seq_len, emb_dim]	[batch, seq_len, emb_dim]	非线性变换
9	Dropout2	[batch, seq_len, emb_dim]	[batch, seq_len, emb_dim]	防止过拟合
10	残差连接2	[batch, seq_len, emb_dim]	[batch, seq_len, emb_dim]	添加残差连接

完整代码实现

Transformer块实现.py
import torch
import torch.nn as nn
import math


class GELU(nn.Module):
    """GELU激活函数实现"""
    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(
            torch.sqrt(torch.tensor(2.0 / torch.pi)) *
            (x + 0.044715 * torch.pow(x, 3))
        ))


class LayerNorm(nn.Module):
    """层归一化实现"""
    def __init__(self, emb_dim):
        super().__init__()
        self.eps = 1e-5
        self.scale = nn.Parameter(torch.ones(emb_dim))
        self.shift = nn.Parameter(torch.zeros(emb_dim))

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        return self.scale * norm_x + self.shift


class MultiHeadAttention(nn.Module):
    """多头注意力机制实现"""
    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
        super().__init__()
        assert d_out % num_heads == 0, "d_out must be divisible by num_heads"

        self.d_out = d_out
        self.num_heads = num_heads
        self.head_dim = d_out // num_heads

        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.out_proj = nn.Linear(d_out, d_out)
        self.dropout = nn.Dropout(dropout)

        # 注册因果掩码
        self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))

    def forward(self, x):
        b, num_tokens, d_in = x.shape

        # 计算查询、键、值
        keys = self.W_key(x)
        queries = self.W_query(x)
        values = self.W_value(x)

        # 重塑为多头格式
        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)
        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)
        values = values.view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)

        # 计算注意力分数
        attn_scores = queries @ keys.transpose(2, 3)

        # 应用掩码
        mask_bool = self.mask.bool()[:num_tokens, :num_tokens]
        attn_scores.masked_fill_(mask_bool, -torch.inf)

        # 应用softmax和dropout
        attn_weights = torch.softmax(attn_scores / math.sqrt(self.head_dim), dim=-1)
        attn_weights = self.dropout(attn_weights)

        # 计算上下文向量
        context_vec = (attn_weights @ values).transpose(1, 2)
        context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)

        # 输出投影
        context_vec = self.out_proj(context_vec)

        return context_vec


class FeedForward(nn.Module):
    """前馈网络实现"""
    def __init__(self, cfg):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
            GELU(),
            nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
        )

    def forward(self, x):
        return self.layers(x)


class TransformerBlock(nn.Module):
    """Transformer块实现"""
    def __init__(self, cfg):
        super().__init__()
        self.att = MultiHeadAttention(
            d_in=cfg["emb_dim"],
            d_out=cfg["emb_dim"],
            context_length=cfg["context_length"],
            num_heads=cfg["n_heads"],
            dropout=cfg["drop_rate"],
            qkv_bias=cfg["qkv_bias"]
        )
        self.ff = FeedForward(cfg)
        self.norm1 = LayerNorm(cfg["emb_dim"])
        self.norm2 = LayerNorm(cfg["emb_dim"])
        self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

    def forward(self, x):
        # 第一个子层：多头注意力
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)
        x = self.drop_shortcut(x)
        x = x + shortcut  # 残差连接

        # 第二个子层：前馈网络
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut  # 残差连接

        return x


def main():
    """主函数：测试Transformer块"""
    print("=" * 60)
    print("Transformer块测试")
    print("=" * 60)

    # GPT配置
    GPT_CONFIG_124M = {
        "vocab_size": 50257,
        "context_length": 1024,
        "emb_dim": 768,
        "n_heads": 12,
        "n_layers": 12,
        "drop_rate": 0.1,
        "qkv_bias": False
    }

    print(f"配置信息:")
    for key, value in GPT_CONFIG_124M.items():
        print(f"  {key}: {value}")
    print()

    # 创建测试数据
    torch.manual_seed(123)
    batch_size, seq_len, emb_dim = 2, 4, 768
    x = torch.rand(batch_size, seq_len, emb_dim)

    print(f"输入张量形状: {x.shape}")
    print(f"输入数据示例 (前3个特征): {x[0, 0, :3].tolist()}")
    print()

    # 创建Transformer块
    block = TransformerBlock(GPT_CONFIG_124M)

    # 前向传播
    output = block(x)

    print(f"输出张量形状: {output.shape}")
    print(f"输出数据示例 (前3个特征): {output[0, 0, :3].tolist()}")
    print()

    # 验证形状保持
    assert x.shape == output.shape, "输入输出形状应该相同"
    print("✓ 形状验证通过：输入输出形状保持一致")

    # 计算参数量
    total_params = sum(p.numel() for p in block.parameters())
    print(f"✓ Transformer块总参数量: {total_params:,}")

    print("=" * 60)
    print("测试完成！")
    print("=" * 60)


if __name__ == "__main__":
    main()

运行结果

执行上述完整代码后，得到以下输出结果：

============================================================
Transformer块测试
============================================================
配置信息:
  vocab_size: 50257
  context_length: 1024
  emb_dim: 768
  n_heads: 12
  n_layers: 12
  drop_rate: 0.1
  qkv_bias: False

输入张量形状: torch.Size([2, 4, 768])
输入数据示例 (前3个特征): [0.29611194133758545, 0.516562283039093, 0.2516707181930542]

输出张量形状: torch.Size([2, 4, 768])
输出数据示例 (前3个特征): [-0.005541101098060608, 0.09722095727920532, -0.1122031882405281]

✓ 形状验证通过：输入输出形状保持一致
✓ Transformer块总参数量: 7,085,568
============================================================
测试完成！
============================================================

结果分析

通过测试结果，我们可以验证Transformer块的正确性：

关键验证点

验证项	期望结果	实际结果	状态
输入形状	[2, 4, 768]	[2, 4, 768]	✅ 正确
输出形状	[2, 4, 768]	[2, 4, 768]	✅ 正确
形状保持	输入输出相同	输入输出相同	✅ 正确
参数量	约700万	7,085,568	✅ 合理

关键发现

形状保持性：Transformer块成功保持了输入输出的形状一致性，这是Transformer架构的重要特性
数值变换：输入数据经过注意力机制和前馈网络的处理后，数值发生了显著变化，说明网络进行了有效的特征变换
参数规模：单个Transformer块包含约708万个参数，主要分布在：
- 多头注意力层：约236万参数
- 前馈网络：约472万参数
- 层归一化：约3千参数
配置兼容性：代码成功使用了GPT-124M的标准配置，验证了实现的正确性

代码详细解析

1. 类设计说明

TransformerBlock类结构

class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        # 初始化注意力层、前馈网络、层归一化和dropout
    def forward(self, x):
        # 实现两个子层的前向传播逻辑

设计要点：

继承自nn.Module，符合PyTorch模块化设计
包含两个主要子层：多头注意力和前馈网络
每个子层都有对应的层归一化和残差连接
使用配置字典统一管理超参数

2. 核心组件分析

多头注意力层

self.att = MultiHeadAttention(
    d_in=cfg["emb_dim"],
    d_out=cfg["emb_dim"],
    context_length=cfg["context_length"],
    num_heads=cfg["n_heads"],
    dropout=cfg["drop_rate"],
    qkv_bias=cfg["qkv_bias"]
)

关键特性：

多头机制：将注意力分解为多个头，捕获不同类型的依赖关系
因果掩码：确保只能关注当前位置之前的信息
缩放点积：使用缩放因子稳定训练过程

前馈网络层

class FeedForward(nn.Module):
    def __init__(self, cfg):
        self.layers = nn.Sequential(
            nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
            GELU(),
            nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
        )

设计特点：

维度扩展：中间层维度是输入的4倍，提供更强的表达能力
GELU激活：使用GELU激活函数，提供平滑的非线性变换
维度恢复：输出维度与输入相同，保持形状一致性

3. 前向传播流程

第一个子层：注意力机制

# 第一个子层：多头注意力
shortcut = x
x = self.norm1(x)
x = self.att(x)
x = self.drop_shortcut(x)
x = x + shortcut  # 残差连接

第二个子层：前馈网络

# 第二个子层：前馈网络
shortcut = x
x = self.norm2(x)
x = self.ff(x)
x = self.drop_shortcut(x)
x = x + shortcut  # 残差连接

关键组件解析：

组件	作用	位置
`LayerNorm`	标准化输入，稳定训练	每个子层之前
`Dropout`	防止过拟合	每个子层之后
`残差连接`	保持梯度流动	每个子层的输入输出
`MultiHeadAttention`	捕获序列依赖关系	第一个子层
`FeedForward`	非线性特征变换	第二个子层

4. 参数量分析

各组件参数分布

根据GPT-124M配置（emb_dim=768, n_heads=12）：

组件	参数量计算	参数数量
多头注意力	4 × 768 × 768 = 2,359,296	235.9万
前馈网络	768 × 3072 + 3072 × 768 = 4,718,592	471.9万
层归一化1	2 × 768 = 1,536	1.5千
层归一化2	2 × 768 = 1,536	1.5千
总计		708.6万

内存和计算复杂度

内存复杂度：O(batch_size × seq_len × emb_dim)
计算复杂度：O(batch_size × seq_len² × emb_dim) （主要来自注意力机制）
参数复杂度：O(emb_dim²) （与序列长度无关）

更新日志

2025/8/18 00:31

查看所有更新日志

bd1d0-迁移目录于 2025/8/18
8d9ff-feat: 新增大模型架构系列文档 - 快捷连接、Transformer块、GPT模型实现和文本生成于 2025/8/17

版权所有

版权归属：NateHHX

许可证：署名 4.0 国际 (CC-BY-4.0)

1.八皇后问题

2.手写数字识别

3.MCP协议

4.RAG系统

5.从零构建大模型

数据准备与采样

自注意力机制

大模型架构