现代大模型架构：组注意力机制（GQA）和 RMSNorm

MHA 通过多个小头从不同角度捕捉语义信息，增强模型的表达能力和稳定性。

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.0):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, attn_mask=None):
        B, L, _ = x.size()
        # 1. 线性投影
        Q = self.w_q(x)
        K = self.w_k(x)
        V = self.w_v(x)
        # 2. reshape 为 [B, H, L, Dh]
        def reshape_heads(t):
            return t.view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
        Q = reshape_heads(Q)
        K = reshape_heads(K)
        V = reshape_heads(V)
        # 3. 缩放点积注意力
        scores = Q @ K.transpose(-2, -1) / (self.head_dim ** 0.5)
        if attn_mask is not None:
            scores = scores.masked_fill(attn_mask == 0, float('-inf'))
        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        out = attn @ V
        # 4. 合并头
        out = out.transpose(1, 2).contiguous().view(B, L, self.d_model)
        return self.w_o(out)

2.2 Multi-Query Attention (MQA)

在 Decoder-Only + 长上下文 + 自回归生成场景下，MHA 暴露出 KV Cache 内存开销过大的问题。对于标准 MHA，每个注意力头维护一份自己的 $K_h, V_h$，显存开销约为 $O(h \cdot L \cdot d_{\text{head}})$。

MQA 提出所有头共享同一份 $K, V$，即只保留一组 $W^K, W^V$，而 $W_i^Q$ 仍然为每个头独立： $$Q_i = X W_i^Q, \quad K = X W^K, \quad V = X W^V$$

这使得 KV Cache 成本大幅降低，同时 Q 的多头结构仍能捕捉多种语义关系。

class MultiQueryAttention(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.0):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, self.head_dim)
        self.w_v = nn.Linear(d_model, self.head_dim)
        self.w_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, attn_mask=None):
        B, L, _ = x.size()
        # 1. 多头 Q
        Q = self.w_q(x)
        Q = Q.view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
        # 2. 单头 K/V
        K = self.w_k(x)
        V = self.w_v(x)
        # 3. 广播 K/V
        K = K.unsqueeze(1).expand(B, self.num_heads, L, self.head_dim)
        V = V.unsqueeze(1).expand(B, self.num_heads, L, self.head_dim)
        # 4. 缩放点积注意力
        scores = Q @ K.transpose(-2, -1) / (self.head_dim ** 0.5)
        if attn_mask is not None:
            scores = scores.masked_fill(attn_mask == 0, float('-inf'))
        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        out = attn @ V
        out = out.transpose(1, 2).contiguous().view(B, L, self.d_model)
        return self.w_o(out)

2.3 Grouped Query Attention (GQA)

GQA 是 MHA 与 MQA 的折中方案。Q 仍然是很多头，但 K/V 的头数减少为更少的组（num_kv_heads），每组 KV 服务若干个 Q 头。

核心思想：在'省 KV'和'头之间有点差异'之间找平衡。

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, num_q_heads, num_kv_heads, dropout=0.0):
        super().__init__()
        assert d_model % num_q_heads == 0
        assert num_q_heads % num_kv_heads == 0
        self.d_model = d_model
        self.num_q_heads = num_q_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = d_model // num_q_heads
        self.group_size = num_q_heads // num_kv_heads
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, num_kv_heads * self.head_dim)
        self.w_v = nn.Linear(d_model, num_kv_heads * self.head_dim)
        self.w_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, attn_mask=None):
        B, L, _ = x.size()
        # 1. Q: 多头；K/V: 少量头
        Q = self.w_q(x)
        K = self.w_k(x)
        V = self.w_v(x)
        Q = Q.view(B, L, self.num_q_heads, self.head_dim).transpose(1, 2)
        K = K.view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2)
        V = V.view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2)
        # 2. 扩展 KV 头
        K = K.repeat_interleave(self.group_size, dim=1)
        V = V.repeat_interleave(self.group_size, dim=1)
        # 3. 缩放点积注意力
        scores = Q @ K.transpose(-2, -1) / (self.head_dim ** 0.5)
        if attn_mask is not None:
            scores = scores.masked_fill(attn_mask == 0, float("-inf"))
        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        out = attn @ V
        # 4. 合并头
        out = out.transpose(1, 2).contiguous().view(B, L, self.d_model)
        return self.w_o(out)

三、归一化：LayerNorm → RMSNorm + Pre-Norm

归一化主要解决深层网络训练不稳定及分布漂移问题。现代 LLM 倾向于使用 RMSNorm + Pre-Norm。

• Post-Norm：原始 Transformer 用法，残差后接归一化。
• Pre-Norm：现代 LLM 常用，残差前接归一化，配合 RMSNorm 训练更稳定。

3.1 LayerNorm

给定隐藏表示 $x \in \mathbb{R}^d$，LayerNorm 对特征维度进行归一化： $$\mu = \frac{1}{d} \sum_{i=1}^{d} x_i, \quad \sigma^2 = \frac{1}{d} \sum_{i=1}^{d} (x_i - \mu)^2$$ $$\text{LN}(x) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta$$

为什么不用 BatchNorm？

• BN 依赖 batch 统计量，推理时 batch 较小易导致分布漂移。
• LN 不依赖 batch，适合变长序列和 Transformer 架构。

3.2 RMSNorm

RMSNorm 基于'层归一化中主要起作用的是缩放因子'的发现，去除了均值减法，只控制尺度： $$\text{RMS}(x) = \sqrt{\frac{1}{d} \sum_{i=1}^{d} x_i^2 + \epsilon}$$ $$\text{RMSNorm}(x) = \frac{x}{\text{RMS}(x)} \cdot \gamma$$

在 Decoder-only 大模型里，RMSNorm + Pre-Norm 组合在超深层网络上表现更稳定。

class RMSNorm(nn.Module):
    def __init__(self, d_model, eps=1e-8):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x):
        rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).sqrt()
        x_norm = x / rms
        return self.weight * x_norm

四、总结

现代大模型架构在注意力机制与归一化层上经历了显著演进。注意力方面，从 MHA 转向 GQA，通过减少 KV 头数量降低显存开销并维持推理速度。归一化方面，LayerNorm 逐渐被 RMSNorm 取代，配合 Pre-Norm 结构以增强深层网络训练稳定性。