CS336 · 第 4 课

注意力替代方案与混合专家

长上下文会放大标准 attention 的计算与数据移动成本。这一课分开研究两类稀疏性：减少 token 之间的 attention edges，以及让每个 token 只激活一部分专家参数；两者解决的不是同一个瓶颈。

本讲问题：标准 attention 在上下文、KV cache 和计算量上形成瓶颈，因此需要先明确想减少的是 KV heads、序列交互还是每 token 激活的参数量，再选择替代机制。

GQA: query heads 保留，KV heads 共享 → KV cache / KV projection 变小
local attention: 每个 query 可见的 key positions 变少 → attention edges 变少
MoE: total expert parameters 很大，每 token 只路由到 top-k → active FLOPs 变小

total parameters ≠ active parameters per token
router probability ≠ attention probability
attention sparsity ≠ MoE parameter sparsity
capacity ≠ guaranteed throughput

标准 Attention 的瓶颈

O(n²) 问题

标准 self-attention 对长度为 n 的序列：

scores = Q @ K^T    # (n, d) @ (d, n) → (n, n)  ← O(n²d) FLOPs
attn = softmax(scores)
output = attn @ V   # (n, n) @ (n, d) → (n, d)  ← O(n²d) FLOPs

朴素训练实现的中间量: 可能物化 (n, n) 的 attention matrix

当 n = 128K（长上下文），n² 超过 16G 元素，单个 bf16 矩阵约 32 GiB。FlashAttention 可以在不改变 full attention 数学结果的前提下避免物化该矩阵；若要降低 O(n²d) 计算量，则需要稀疏、局部、低秩或其他序列模型。

两条突围路径

注意力替代方案

1. Linear Attention

核心思想：把 softmax 去掉，用核函数近似：

# 标准 attention: O(n²)
attn = softmax(Q @ K^T / √d) @ V

# Linear attention: O(n)
# 利用 (Q @ K^T) @ V = Q @ (K^T @ V) 的结合律
# 先算 K^T @ V（d × d），再算 Q @ (K^T @ V)
output = φ(Q) @ (φ(K)^T @ V)  # φ 是核函数（如 ReLU, elu+1）

# 复杂度: O(nd²) —— 当 n >> d 时远优于 O(n²d)

代价：核化 linear attention 改变了 softmax attention 的相似度与归一化形式，表达能力和数值行为可能不同；token 仍由 Q/K 内容动态加权，并不是“所有 token 被平等对待”。

2. Sliding Window Attention

Mistral 7B 和 LongFormer 使用：

# 每个 token 只看前后 w 个 token
for i in range(n):
    scores[i] = Q[i] @ K[max(0,i-w):i+w]^T  # 只看窗口内的 key
    attn[i] = softmax(scores[i])
    output[i] = attn[i] @ V[max(0,i-w):i+w]

# 复杂度: O(nwd) —— 线性于 n

直觉：大多数语言建模任务中，附近的 token 比远处的更重要。

3. Sparse Attention（全局 + 局部混合）

4. State Space Models（Mamba）

完全抛弃 attention，用经典控制论的状态空间模型：

# SSM: 递推形式
h(t) = A × h(t-1) + B × x(t)
y(t) = C × h(t)

# 训练时: 用并行 scan / 专用 kernel 高效计算
# 推理时: 递推；状态内存相对序列长度为 O(1)

# Mamba 的改进: A, B, C 是输入相关的（选择性 SSM）
# 让模型学会"忘记"不相关的信息

方案对比

不要与 decode KV cache 混淆：上表这一列讨论单次 attention 计算的中间张量。标准 autoregressive decode 的历史 K/V storage 则随已缓存 context length 线性增长；二者是不同对象、不同时间尺度。

Grouped-Query Attention (GQA)

从 MHA 到 MQA 到 GQA

为什么 GQA 有效？

实际效果

Multi-head Latent Attention (MLA)

DeepSeek-V2 的进一步创新：

Mixture of Experts (MoE)

核心思想

MoE 层结构

class MoELayer(nn.Module):
    def __init__(self, d_model, d_ff, num_experts, top_k):
        self.gate = nn.Linear(d_model, num_experts)  # 路由器
        self.experts = nn.ModuleList([
            FFN(d_model, d_ff) for _ in range(num_experts)
        ])
        self.top_k = top_k  # 通常 2-4

    def forward(self, x):
        # 1. 路由: 每个 token 选择 top_k 个专家
        router_probs = self.gate(x).softmax(dim=-1)  # (batch*seq, num_experts)
        top_k_scores, top_k_indices = router_probs.topk(self.top_k, dim=-1)
        top_k_scores = top_k_scores / top_k_scores.sum(dim=-1, keepdim=True)

        # 2. 分发: 把 token 送到对应的专家
        output = torch.zeros_like(x)
        for k in range(self.top_k):
            expert_idx = top_k_indices[:, k]
            expert_mask = (
                expert_idx.unsqueeze(1)
                == torch.arange(len(self.experts), device=x.device)
            )
            # 把属于该专家的 token 送入计算
            for i, expert in enumerate(self.experts):
                token_mask = expert_mask[:, i]
                if token_mask.any():
                    weight = top_k_scores[token_mask, k].unsqueeze(-1)
                    output[token_mask] += expert(x[token_mask]) * weight

        return output

路由策略

DeepSeek-V3 使用的是 token-choice top-k 路由，并通过可学习的 expert bias 做近似 auxiliary-loss-free 的负载均衡；它不是 Expert Choice Routing。

负载均衡问题

MoE 最大的挑战：负载不均衡。热门专家过载，冷门专家饥饿。

辅助损失（Auxiliary Loss）

# Switch Transformer / DeepSeek 类辅助损失的简化示意
# 同时考虑离散 token 分配和连续路由概率

def load_balancing_loss(router_logits, top_k, num_experts):
    # router_logits: (num_tokens, num_experts)
    router_probs = router_logits.softmax(dim=-1)
    selected = router_probs.topk(top_k, dim=-1).indices

    # 实际离散路由后，每个专家接收的 token-assignment 比例
    assignments = F.one_hot(selected, num_classes=num_experts).float()
    token_fraction = assignments.mean(dim=(0, 1))

    # 路由器分配给每个专家的平均概率
    prob_fraction = router_probs.mean(dim=0)

    # Switch/DeepSeek 类辅助项的简化写法
    load_loss = num_experts * (token_fraction * prob_fraction).sum()
    return load_loss

# 总损失 = 主损失 + α × 辅助损失
# α 是需要按模型、路由方案与训练规模调节的超参数

MoE 的通信挑战

在多 GPU 上，不同专家分布在不同 GPU 上：

并行策略

现代 MoE 模型

MoE 的挑战与解决方案

挑战 1：训练不稳定

MoE 训练容易发散，特别是当某些专家被过度使用时。

解决方案：

• 辅助负载均衡损失
• 路由器计算保留 FP32，减少低精度带来的不稳定
• 路由器 z-loss（惩罚过大的 logits）
• 专家容量的硬限制（超载 token 被丢弃）

挑战 2：通信开销

all-to-all 通信在跨节点时成为瓶颈。

解决方案：

• 节点内 expert parallel，节点间 data parallel
• 通信与计算 overlap
• 局部性感知路由（让相关 token 去同一节点）

挑战 3：专家设计与初始化

专家粒度、共享专家是否有效以及如何从已有 dense 模型初始化，都需要通过模型与数据上的实验验证。

常见设计：

• 细粒度专家：把大 FFN 拆成更多、更小的 routed experts
• 共享专家：始终激活少量专家；DeepSeek 的消融显示有收益，但 OLMoE 的消融未观察到同样收益
• Upcycling：复制或拆分预训练 dense FFN，初始化新的稀疏 MoE

挑战 4：推理内存

虽然激活参数少，但推理时所有专家都要在内存中。

解决方案：

• 专家卸载（不活跃的专家放 CPU）
• 专家量化（低精度存储）
• 专家剪枝（移除不重要的专家）

混合架构趋势

最新的模型开始混合不同架构：

Jamba（AI21 Labs）

MiniMax-01

DeepSeek-V3

知识检查

如果只能记住 5 件事

延伸阅读