当前位置：首页 > news >正文

ChatGLM2-6B的GLMBlock里到底发生了什么？一次注意力与MLP的深度游

news 2026/6/12 2:48:17

ChatGLM2-6B的GLMBlock解剖：从注意力机制到SwiGLU的微观世界

当我们将ChatGLM2-6B这样的语言模型置于手术台上时，最迷人的部分往往藏在那些重复堆叠的基础模块中。GLMBlock作为这个6B参数巨人的基本组成单元，其内部精妙的设计决定了模型最终的推理能力和效率。本文将带您深入这个微观世界，用放大镜观察每一个计算步骤的数学本质和工程实现。

1. GLMBlock的整体架构与设计哲学

GLMBlock作为ChatGLM2-6B的核心计算单元，其整体结构可以看作是对标准Transformer block的创造性改造。与原始Transformer相比，它保留了注意力机制和前馈网络的基本框架，但在细节上进行了多处关键改进：

[输入(4096维)] │ ▼ RMSNorm # 前置归一化 │ ▼ Attention # 带RoPE的多头注意力 │ ▼ Dropout → Add → RMSNorm # 残差连接与归一化 │ ▼ MLP(SwiGLU) # 扩展维度至27392的前馈网络 │ ▼ Dropout → Add # 最终残差连接 │ ▼ [输出(4096维)]

这种设计最显著的特点是采用了前置归一化(Pre-LN)而非原始Transformer的后置归一化(Post-LN)。前置归一化将层归一化移到残差分支之前，这种结构在训练稳定性和收敛速度上表现更好。在实际测试中，Pre-LN结构可以使模型在相同训练步数下获得更低的损失值。

另一个关键设计是残差连接的位置选择。GLMBlock中存在两处残差连接：

注意力模块后的残差连接：将注意力输出与模块原始输入相加
MLP模块后的残差连接：将MLP输出与注意力后的归一化结果相加

这种分阶段的残差设计使得梯度能够更有效地在深层网络中传播。从参数文件分析，28个GLMBlock的参数差异主要体现在以下几个方面：

参数类型	变化规律	影响范围
注意力头参数	各层独立，无明显规律	QKV投影矩阵
MLP权重矩阵	随深度增加呈现平滑变化趋势	扩展/收缩层权重
归一化参数	几乎保持不变	scale参数

2. 注意力模块的深度解析

GLMBlock的注意力模块是其处理上下文关系的核心引擎。与标准Transformer相比，它在以下方面进行了优化：

2.1 改进的RoPE位置编码

ChatGLM2-6B采用了Rotary Position Embedding(RoPE)而非绝对位置编码。RoPE的独特之处在于将位置信息通过旋转矩阵注入到注意力计算中：

def apply_rotary_pos_emb(q, k, pos_emb): # q,k shape: [seq_len, num_heads, head_dim] # pos_emb shape: [seq_len, head_dim] cos, sin = pos_emb q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed

这种编码方式具有以下优势：

相对位置感知：模型能够更好地捕捉token之间的相对距离
长度外推性：对超过训练长度的序列有更好的处理能力
计算效率：不需要额外的位置编码参数

2.2 多头注意力的具体实现

GLMBlock中的多头注意力将输入投影到32个注意力头（每个头128维），计算过程可分为四个阶段：

QKV投影：将输入分别映射到查询(Query)、键(Key)和值(Value)空间

q = linear(x, W_q) # [seq_len, 32, 128] k = linear(x, W_k) # [seq_len, 32, 128] v = linear(x, W_v) # [seq_len, 32, 128]

注意力分数计算：采用缩放点积注意力机制
```
scores = torch.matmul(q, k.transpose(-2,-1)) / sqrt(d_k)
```
注意力权重应用：通过softmax和dropout
```
attn = dropout(softmax(scores, dim=-1))
```
输出投影：将多头输出合并并投影回原空间
```
output = linear(concat(heads), W_o)
```

注意：在实际推理时，KV缓存机制会保存历史Key和Value，避免重复计算。这是ChatGLM2-6B能够高效生成文本的关键优化。

2.3 注意力掩码策略

GLMBlock采用了两种注意力掩码的组合：

因果掩码：防止当前位置关注未来信息
前缀注意力掩码：允许特定位置关注整个前缀上下文

这种混合策略使得模型在对话场景中既能保持生成的连贯性，又能充分利用预设的对话历史信息。

3. MLP模块的维度魔术

GLMBlock中的MLP模块看似简单，却隐藏着模型强大表示能力的关键。其结构特点如下：

3.1 SwiGLU激活函数

传统Transformer使用ReLU作为MLP的激活函数，而GLMBlock采用了更先进的SwiGLU(Swished Gated Linear Unit)：

def SwiGLU(x, W1, W2, W3): return swish(x @ W1) * (x @ W3) @ W2

其中swish函数定义为：

def swish(x): return x * torch.sigmoid(x)

SwiGLU相比ReLU有以下优势：

更平滑的梯度流动
更强的非线性表示能力
门控机制可以动态控制信息流

3.2 维度扩展策略

MLP模块最引人注目的是其惊人的维度扩展：

输入: 4096 ↓ 扩展层: 4096 → 27392 (×6.68) ↓ 收缩层: 27392 → 4096

这种"扩展-收缩"模式为模型提供了以下能力：

高阶特征交互：在高维空间中可以形成更复杂的特征组合
信息解耦：不同特征可以在扩展维度上获得独立的表示
容量提升：大幅增加了可学习参数的数量

从参数文件分析，各层GLMBlock的MLP权重呈现出有趣的模式：

浅层(1-7层)：权重分布较广，捕捉基础语言特征
中层(8-21层)：权重出现明显的结构化模式
深层(22-28层)：权重高度特化，处理高级语义

4. 归一化与残差连接的协同设计

GLMBlock中的归一化和残差连接系统是保证深层模型稳定训练的关键组件。其设计有几个精妙之处：

4.1 RMSNorm的轻量替代

不同于传统的LayerNorm，GLMBlock采用了RMSNorm(Root Mean Square Layer Normalization)：

class RMSNorm(nn.Module): def __init__(self, dim): super().__init__() self.scale = dim ** -0.5 self.gamma = nn.Parameter(torch.ones(dim)) def forward(self, x): norm = x.norm(2, dim=-1, keepdim=True) * self.scale return x / norm * self.gamma

RMSNorm省去了计算均值的步骤，具有以下特点：

计算量减少约20%
在深层网络中表现更稳定
保留了缩放和平移的可学习参数

4.2 残差路径的精心设计

GLMBlock中的两处残差连接各有特点：

注意力后残差：直接将模块输出与原始输入相加
```
x = x + dropout(attention_output)
```
MLP后残差：将输出与归一化后的中间结果相加
```
x = mid_norm_output + dropout(mlp_output)
```

这种设计形成了两条并行的梯度传播路径：

短路径：保留原始信息
长路径：承载变换后的特征

实验表明，这种双残差结构可以使模型在保持训练稳定性的同时，达到更深的网络深度。

5. 28层GLMBlock的协同工作

ChatGLM2-6B的28个GLMBlock并非简单重复，而是形成了精密的处理流水线。通过分析各层注意力模式和MLP激活情况，我们可以观察到明显的层次分化：

层级范围	主要功能	注意力特点	MLP激活模式
1-7层	基础语法处理	局部注意力主导	激活值分布广泛
8-14层	短语级语义组合	开始出现跨片段注意力	特定神经元开始特化
15-21层	句子级逻辑推理	全局注意力模式明显	高度特化的神经元集群
22-28层	篇章级连贯与高级推理	稀疏的专家注意力模式	极稀疏的激活模式