5. 线性注意力

概述

前面章节介绍的 MHA、GQA 等都属于标准注意力（Full Attention），其核心问题是计算复杂度为 O(n²)——每个 token 都要和所有其他 token 计算注意力权重。当序列长度为 8K、32K 甚至 128K 时，这成为性能瓶颈。

本章将介绍线性注意力的演进历程：

标准注意力 O(n²) → 线性注意力 O(n) → DeltaNet → Gated DeltaNet (Qwen3.5 采用)

5.1 标准注意力的问题

复杂度分析

标准 Scaled Dot-Product Attention：

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

• $Q,K\in {\mathbb{R}}^{n\times d}$：序列长度 $n$，维度 $d$
• $Q{K}^T\in {\mathbb{R}}^{n\times n}$：这是 O(n²) 的来源
• 每个 token 都要和所有其他 token 计算相关性

KV Cache 问题：

• 推理时，每生成一个新 token，都需要将其 K 和 V 追加到缓存中
• 缓存大小随序列长度线性增长：batch × heads × seq_len × head_dim
• 当 seq_len = 128K 时，单层 KV Cache 可达数百 MB

5.2 线性注意力的基本思想

核心观察

标准注意力可以写成：

$$O=\text{softmax}(Q{K}^T)V$$

如果我们把 softmax 拆开（假设用核函数 $\varphi$ 替代）：

$$O_i=\frac{\sum _j\varphi (q_i{)}^T\varphi (k_j)\cdot v_j}{\sum _j\varphi (q_i{)}^T\varphi (k_j)}$$

关键洞察：矩阵乘法的结合律！

$$\text{标准: }(\varphi (Q)\varphi (K{)}^T)V\to O(n^2)$$$$\text{线性: }\varphi (Q)(\varphi (K{)}^{T}V)\to O(n)$$

先计算 $\varphi (K{)}^{T}V\in {\mathbb{R}}^{d\times d}$（固定大小），再乘以 $\varphi (Q)$，复杂度从 O(n²) 降到 O(n)。

递归视角

线性注意力还有一个等价的递归视角。定义状态矩阵 $S_t$：

$$S_t=S_{t-1}+\varphi (k_t)v_t^T$$$$o_t=\varphi (q_t)S_t$$

每一步只更新一个 rank-1 矩阵，总复杂度 O(n)。

普通线性注意力的问题

上述公式中，$S_t=S_{t-1}+\varphi (k_t)v_t^T$ 只能累加，不能擦除。这意味着：

• 历史信息只能越积越多，无法遗忘
• 随着序列增长，$S_t$ 的范数越来越大
• 没有门控机制，无法选择性地遗忘旧信息

5.3 DeltaNet：修正错误的线性注意力

原理

DeltaNet（Yang et al., 2024）对普通线性注意力做了关键改进，引入 delta rule：

$$S_t=S_{t-1}+k_t(v_t-S_{t-1}^T{k}_t{)}^T{\beta }_t$$

其中：

• $v_t-S_{t-1}^T{k}_t$：delta，即"期望值 - 预测值"
• $S_{t-1}^T{k}_t$：用当前状态对 key 做预测，得到预测的 value
• $v_t-S_{t-1}^T{k}_t$：预测误差（"惊讶"程度）
• ${\beta }_t$：更新强度，控制写入多少新信息

直觉理解：

• 普通线性注意力：$S_t=S_{t-1}+k_t{v}_t^T$，总是累加新信息
• DeltaNet：$S_t=S_{t-1}+k_t(v_t-\text{预测值}{)}^T$，只更新预测错误的部分
• 如果模型已经能很好地预测 $v_t$，delta 接近 0，状态几乎不变
• 如果预测很差（"惊讶"），则大幅更新状态

类比：类似于人类学习——我们不会重复学习已经掌握的知识，而是专注于新知识和纠正错误。

与普通线性注意力的区别

	普通线性注意力	DeltaNet
更新公式	$S_t=S_{t-1}+k_t{v}_t^T$	$S_t=S_{t-1}+k_t(v_t-S_{t-1}^T{k}_t{)}^T{\beta }_t$
能否擦除旧信息	❌ 只能累加	✅ 可以修正/覆盖
更新机制	无差别累加	只更新"惊讶"的部分
门控	无	${\beta }_t$ 控制更新强度

5.4 Gated DeltaNet：加门控的 DeltaNet（Qwen3.5 采用）

原理

Gated DeltaNet 在 DeltaNet 基础上增加了遗忘门控 $g_t$：

$$S_t=S_{t-1}\cdot g_t+k_t(v_t-S_{t-1}^T{k}_t{)}^T{\beta }_t$$$$o_t=q_t{S}_t$$

其中：

• $g_t$：门控衰减（在 log-space，已取负），控制遗忘多少旧信息
• $g_t$ 越大 → 保留越多旧信息
• $g_t$ 越小 → 遗忘越多旧信息

三重控制：

1. delta $(v_t-S_{t-1}^T{k}_t)$：只更新预测错误的部分
2. beta ${\beta }_t$：控制写入多少新信息
3. gate $g_t$：控制遗忘多少旧信息

代码结构总览

Gated DeltaNet 的完整流程：

hidden_states

▼

in_proj_qkv

投影得到 Q, K, V

▼

causal_conv1d

因果卷积预处理（提取局部特征）

▼

in_proj_b / in_proj_a

投影得到 beta 和 alpha（门控参数）

▼

Delta Rule 核心

更新递归状态 S_t

▼

RMSNormGated

带门控的归一化

▼

out_proj

输出投影

▼

output

颜色含义：蓝色 = 线性投影层；绿色 = 特征处理与状态更新（Delta Rule 为核心，加粗边框突出）；橙色 = 归一化层。

因果卷积（Causal Conv1d）

为什么需要因果卷积？

在 Gated DeltaNet 中，QKV 投影之后、Delta Rule 之前，有一个因果卷积层。它的作用是：

1. 提取局部特征：DeltaNet 的递归状态 $S_t$ 只能看到当前 token 和之前的历史压缩信息，无法直接看到相邻 token 的原始特征。因果卷积让每个 token 能"看到"前面几个 token 的原始表示，提取局部模式（如词组、短语）。

2. 增强输入质量：经过卷积平滑后的 QKV 表示更稳定，Delta Rule 的更新也更有效。

3. 保持因果性：不能让未来的 token 信息泄露到当前 token，所以必须用"因果"卷积。

标准卷积 vs 因果卷积

标准卷积（如图像处理）：

输入: [x0, x1, x2, x3, x4]
kernel_size = 3

输出[0] = f(x0, x1, x2)      ← 看到了未来！
输出[1] = f(x1, x2, x3)      ← 看到了未来！
输出[2] = f(x2, x3, x4)      ← 看到了未来！

标准卷积的每个输出都依赖当前和未来的输入，这在语言模型中是不允许的。

因果卷积（Causal Conv1d）：

输入: [x0, x1, x2, x3, x4]
kernel_size = 3

输出[0] = f(0,   0,   x0)    ← 只看到过去
输出[1] = f(0,   x0,  x1)    ← 只看到过去
输出[2] = f(x0,  x1,  x2)    ← 只看到过去和当前
输出[3] = f(x1,  x2,  x3)    ← 只看到过去和当前
输出[4] = f(x2,  x3,  x4)    ← 只看到过去和当前

通过左侧填充 (padding) 来保证因果性。

实现方式

因果卷积的标准做法：

# kernel_size = 4 的因果卷积
conv1d = nn.Conv1d(
    in_channels=C,
    out_channels=C,
    kernel_size=4,
    groups=C,              # depthwise: 每个通道独立卷积
    padding=3,             # 因果卷积需要 kernel_size - 1 = 3 个 padding
)

# 前向传播
output = conv1d(x)[:, :, :seq_len]    # Conv1d 已经做了 padding，只需截断到原始长度
output = F.silu(output)                # 激活函数

逐步图解（kernel_size=4）：

原始序列:      [a, b, c, d, e, f]
左侧填充:   [0, 0, 0, a, b, c, d, e, f]

卷积核滑动:
  位置 0: [0, 0, 0, a] → output[0]   ← 只看到 a
  位置 1: [0, 0, a, b] → output[1]   ← 看到 a, b
  位置 2: [0, a, b, c] → output[2]   ← 看到 a, b, c
  位置 3: [a, b, c, d] → output[3]   ← 看到 a, b, c, d
  位置 4: [b, c, d, e] → output[4]   ← 看到 b, c, d, e
  位置 5: [c, d, e, f] → output[5]   ← 看到 c, d, e, f
  （截断到原始长度 6）

Depthwise Convolution

代码中的 groups=conv_dim 表示 depthwise convolution：

# 标准卷积: 所有通道共享同一个卷积核，通道间有交互
nn.Conv1d(in_channels=64, out_channels=64, kernel_size=4, groups=1)
# 参数量: 64 × 64 × 4 = 16,384

# Depthwise 卷积: 每个通道独立卷积，通道间无交互
nn.Conv1d(in_channels=64, out_channels=64, kernel_size=4, groups=64)
# 参数量: 64 × 1 × 4 = 256  ← 参数量大幅减少！

在 DeltaNet 中使用 depthwise 的原因：

• QKV 的每个维度（通道）独立做局部平滑，不需要跨维度交互
• 参数量从 C² × K 降到 C × K，节省显存和计算
• 跨维度的交互由后面的 Delta Rule 和线性层负责

代码对照

# 在 GatedDeltaNet.__init__ 中:
conv_dim = self.key_dim * 2 + self.value_dim  # Q+K+V 的总维度 = 6144
self.conv1d = nn.Conv1d(
    in_channels=conv_dim,
    out_channels=conv_dim,
    bias=False,
    kernel_size=self.conv_kernel_size,       # 4
    groups=conv_dim,                         # depthwise: 每个维度独立
    padding=self.conv_kernel_size - 1,       # 因果 padding
)

# 在 GatedDeltaNet.forward 中:
mixed_qkv = self.in_proj_qkv(hidden_states)  # (batch, seq, 6144)
mixed_qkv = mixed_qkv.transpose(1, 2)        # (batch, 6144, seq) ← Conv1d 需要

mixed_qkv = self.conv1d(mixed_qkv)[:, :, :seq_len]             # 卷积 + 截断
mixed_qkv = F.silu(mixed_qkv)                                  # 激活

mixed_qkv = mixed_qkv.transpose(1, 2)        # (batch, seq, 6144) ← 恢复

代码实现

代码文件

门控 RMSNorm

DeltaNet 的输出归一化用的是 RMSNormGated，它和第 2 章的标准 RMSNorm 不是同一套参数化，不能混淆：

	标准 RMSNorm（第 2 章）	RMSNormGated（DeltaNet）
weight 初始化	zeros（零初始化）	ones（一初始化）
公式	(1 + weight) * norm(x)	weight * norm(x)
用途	模型主体的 Pre-Norm	DeltaNet 输出 + 门控

官方权重下 linear_attn.norm.weight 聚类在 1 附近，正是一初始化的指纹；如果误用零初始化的 (1+weight) 公式加载，DeltaNet 输出会全部偏掉。

class RMSNormGated(nn.Module):
    """
    带门控的 RMSNorm（对应官方 Qwen3_5RMSNormGated）
    公式: output = weight * norm(x) * silu(z)
    注意：weight 初始化为 1（一初始化），与标准 RMSNorm 的零初始化不同。
    """

    def __init__(self, dim: int, eps: float = 1e-6, activation: str = "silu"):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))  # 一初始化！
        self.activation = activation

    def _norm(self, x: torch.Tensor) -> torch.Tensor:
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        # norm 用 float32 算，乘 weight 前先转回输入 dtype（对齐官方）
        input_dtype = x.dtype
        hidden_states = x.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
        hidden_states = self.weight * hidden_states.to(input_dtype)  # weight * norm(x)
        if self.activation == "silu":
            hidden_states = hidden_states * F.silu(z.to(torch.float32))
        else:
            hidden_states = hidden_states * z.to(torch.float32)
        return hidden_states.to(input_dtype)

Gated DeltaNet 主体

class GatedDeltaNet(nn.Module):
    """
    Gated DeltaNet 线性注意力层

    这是 Qwen3.5 的核心创新组件。
    """

    def __init__(self, config, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.num_v_heads = config.linear_num_value_heads  # 16
        self.num_k_heads = config.linear_num_key_heads  # 16
        self.head_k_dim = config.linear_key_head_dim  # 128
        self.head_v_dim = config.linear_value_head_dim  # 128
        self.key_dim = self.head_k_dim * self.num_k_heads  # 2048
        self.value_dim = self.head_v_dim * self.num_v_heads  # 2048
        self.conv_kernel_size = config.linear_conv_kernel_dim  # 4
        self.layer_idx = layer_idx

        # ---- QKV 投影 ----
        # 把 hidden_size 投影到 Q, K, V
        self.in_proj_qkv = nn.Linear(self.hidden_size, self.key_dim * 2 + self.value_dim, bias=False)
        # 门控投影 (z gate)
        self.in_proj_z = nn.Linear(self.hidden_size, self.value_dim, bias=False)
        # beta 投影（更新强度）
        self.in_proj_b = nn.Linear(self.hidden_size, self.num_v_heads, bias=False)
        # alpha 投影（遗忘门控）
        self.in_proj_a = nn.Linear(self.hidden_size, self.num_v_heads, bias=False)

        # ---- 因果卷积 ----
        # 对 QKV 做 depthwise 因果卷积，提取局部特征
        conv_dim = self.key_dim * 2 + self.value_dim
        self.conv1d = nn.Conv1d(
            in_channels=conv_dim,
            out_channels=conv_dim,
            bias=False,
            kernel_size=self.conv_kernel_size,
            groups=conv_dim,  # depthwise convolution
            padding=self.conv_kernel_size - 1,  # causal padding
        )

        # ---- 时间步参数 ----
        self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads))
        A = torch.empty(self.num_v_heads).uniform_(0, 16)
        self.A_log = nn.Parameter(torch.log(A))

        # ---- 带门控的 RMSNorm ----
        self.norm = RMSNormGated(self.head_v_dim, eps=config.rms_norm_eps)

        # ---- 输出投影 ----
        self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        Args:
            hidden_states: (batch, seq_len, hidden_size)
        Returns:
            output: (batch, seq_len, hidden_size)
        """
        batch_size, seq_len, _ = hidden_states.shape

        # ---- 投影 QKV ----
        mixed_qkv = self.in_proj_qkv(hidden_states)  # (batch, seq, 2*key_dim + value_dim)
        mixed_qkv = mixed_qkv.transpose(1, 2)  # (batch, channels, seq) for conv1d

        # ---- 因果卷积 ----
        mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, :seq_len])

        mixed_qkv = mixed_qkv.transpose(1, 2)  # (batch, seq, channels)

        # ---- 分离 Q, K, V ----
        query, key, value = mixed_qkv.split(
            [self.key_dim, self.key_dim, self.value_dim], dim=-1
        )
        query = query.reshape(batch_size, seq_len, self.num_k_heads, self.head_k_dim)
        key = key.reshape(batch_size, seq_len, self.num_k_heads, self.head_k_dim)
        value = value.reshape(batch_size, seq_len, self.num_v_heads, self.head_v_dim)

        # ---- 计算门控参数 ----
        z = self.in_proj_z(hidden_states)  # (batch, seq, value_dim)
        z = z.reshape(batch_size, seq_len, self.num_v_heads, self.head_v_dim)

        b = self.in_proj_b(hidden_states)  # (batch, seq, num_v_heads)
        a = self.in_proj_a(hidden_states)  # (batch, seq, num_v_heads)

        beta = b.sigmoid()
        # g = -exp(A_log) * softplus(a + dt_bias)，控制遗忘速率
        g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)

        # ---- KV head 扩展（如果 num_v_heads > num_k_heads）----
        # 这里 num_v_heads == num_k_heads == 16，所以不需要 repeat
        if self.num_v_heads // self.num_k_heads > 1:
            query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
            key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)

        # ---- Delta Rule 核心计算 ----
        core_attn_out, _ = torch_chunk_gated_delta_rule(
            query=query,
            key=key,
            value=value,
            g=g,
            beta=beta,
            chunk_size=64,
            use_qk_l2norm=True,
        )

        # ---- 门控 RMSNorm ----
        core_attn_out = core_attn_out.reshape(-1, self.head_v_dim)
        z = z.reshape(-1, self.head_v_dim)
        core_attn_out = self.norm(core_attn_out, z)
        core_attn_out = core_attn_out.reshape(batch_size, seq_len, -1)

        # ---- 输出投影 ----
        output = self.out_proj(core_attn_out)
        return output

Delta Rule 的 Chunk 实现

为什么需要 Chunk？

逐 token 的 Delta Rule 是完全串行的：

# 逐 token 递归（慢，无法并行）
for t in range(seq_len):
    pred = k_t @ S_{t-1}                # 用历史状态"预测"当前 value
    delta = (v_t - pred) * beta_t       # 误差 = 实际 - 预测
    S_t = S_{t-1} * exp(g_t) + outer(k_t, delta)  # 写入
    o_t = (q_t @ S_t) * scale           # 输出

每一步都依赖上一步的 $S_{t-1}$，无法并行。当 seq_len = 2048 时需要串行 2048 步。

Chunk 模式：把序列分成大小为 $C$ 的块（如 C=64），块内用矩阵运算并行化、块间仍然递归但只有 $n/C$ 步。

序列: [t0, t1, ..., t63, t64, ..., t127, ...]
       |---- Chunk 0 ----|---- Chunk 1 ----|----
       矩阵运算(并行)    矩阵运算(并行)    ...
              ↓                   ↓
         递归传 S ─────────→ 递归传 S ───→ ...

串行从 $n$ 步降到 $n/C$ 步，每步内部并行——这是 Gated DeltaNet 能实用部署的关键。

一个简单例子贯穿本节

为讲清 chunk 算法和它与朴素的等价性，本节用一个最小可手算的例子：

• 维度 $d=2$，chunk 大小 $C=2$，序列长度 $n=2$（1 个 chunk）
• $K$（键）$=V$（值）$=\left(\begin{array}{cc}1& 0\\ 0& 1\end{array}\right)$（$K=V$ = 单位基）
• $\beta =1$，$g$ 在本节暂设全 0（先看对齐），后面引入衰减

g 全 0 下，朴素和 chunk 数学上完全等价——这给我们一个干净的起点建立直觉。

朴素方法（baseline）

直接按教科书 delta rule 逐 token 跑：

S = zeros(2,2)
for t in [0, 1]:
    pred = K[t] @ S
    delta = (V[t] - pred) * beta[t]
    S = S * exp(g[t]) + outer(K[t], delta)
    o[t] = scale * (K[t] @ S)        # scale = 1/√d

$g=0$ 时步骤详细：

t	$S_{\text{prev}}$	$pred=K_t@S_{\text{prev}}$	$\delta =V_t-pred$	$S_{\text{after}}$	$o_t=K_t@S_{\text{after}}$
0	$\left(\begin{array}{cc}0& 0\\ 0& 0\end{array}\right)$	$\left(\begin{array}{c}0\\ 0\end{array}\right)$	$\left(\begin{array}{c}1\\ 0\end{array}\right)$	$\left(\begin{array}{cc}1& 0\\ 0& 0\end{array}\right)$	$\left(\begin{array}{c}1\\ 0\end{array}\right)$
1	$\left(\begin{array}{cc}1& 0\\ 0& 0\end{array}\right)$	$\left(\begin{array}{c}0\\ 0\end{array}\right)$	$\left(\begin{array}{c}0\\ 1\end{array}\right)$	$\left(\begin{array}{cc}1& 0\\ 0& 1\end{array}\right)$	$\left(\begin{array}{c}0\\ 1\end{array}\right)$

$o=[[1,0],[0,1]]$。这就是朴素 baseline。

Chunk 方法

目标：把同一个例子跑出完全相同的 $o$，但 chunk 内用矩阵运算（GPU 可并行）替代串行。

算法三步（先给主流程，原理细节放 <details>）：

def chunk_delta(K, V, g, beta, S_start=zeros, C=2):
    # === 阶段 1: chunk 内衰减矩阵（可并行）===
    gcum = cumsum(g)                          # (C,) 累积门控
    decay[i,j] = exp(gcum[i] - gcum[j])  if i>=j else 0   # (C,C) 下三角

    # === 阶段 2: chunk 内修正（可并行）===
    A = -( (K*beta) @ K.T ) * decay         # (C,C) 严格下三角
    A[上三角含对角] = 0                      # 因果 mask
    # 递推: attn = (I+A)^{-1}, 不显式求逆
    attn = A + eye(C)
    for i in 1..C-1:
        attn[i,:i] = attn[i,:i] + (attn[i,:i,None] * attn[:i,:i]).sum(0)
    Vp = attn @ (V*beta)                       # (C,d) 修正后的 value

    # === 阶段 3: chunk 间递归（仍串行，但只 C 步一次）===
    v_new = Vp - ((attn @ (K*beta * exp(gcum)[:,None])) @ S_start)
    S_end = S_start * exp(gcum[-1]) + (K * exp(gcum[-1]-gcum)[:,None]).T @ v_new
    return v_new, S_end

逐项原理解释（点击展开，不要求掌握所有细节——知道"chunk 内能并行、修正矩阵是什么"就够了）：

阶段 1 原理推导：衰减矩阵 $\text{decay}$ 的来历

朴素递归里，token $j$ 在 $t$ 时刻写入的状态 $S$ 会被后续门控逐级衰减。要推导 chunk 内的衰减矩阵，先看朴素 $S_t$ 中 token $j$ 的贡献怎么演化。

朴素递归（$S_t=S_{t-1}\cdot e^{g_t}+k_j{\delta }_j^{\mathrm{\top }}$，其中 ${\delta }_j$ 是 token $j$ 写入的 delta）：token $j$ 写入后，从 $j+1$ 到 $t$ 每步都乘 $e^{g_s}$，所以 token $j$ 的贡献在 $S_t$ 里被乘了 $e^{g_{j+1}}\cdot e^{g_{j+2}}\cdots e^{g_t}=e^{\sum _{s=j+1}^t{g}_s}$。

累积门控：记 $g_{\text{cum}}[i]=\sum _{s=0}^i{g}_s$（含 $g_0$ 到 $g_i$），则 $\sum _{s=j+1}^t{g}_s=g_{\text{cum}}[t]-g_{\text{cum}}[j]$。

关键时机：token $i$ 在读取 $S_{i-1}$（写入 ${\delta }_i$ 之前的状态）时，token $j$ 的贡献已经过 $g_{j+1},\dots ,g_{i-1}$ 衰减（不含 $g_i$，因为 $g_i$ 是写入 ${\delta }_i$ 时才应用的衰减）。所以 token $i$ 读到的 token $j$ 贡献的衰减因子是：

$$\text{decay}[i,j]=e^{g_{\text{cum}}[i-1]-g_{\text{cum}}[j]}(i>j)$$

对角线 $\text{decay}[i,i]=1$（自身不衰减，写入即读）。下三角结构（$i<j$ 为 0）来自因果性：未来 token 不影响过去。

为什么 chunk 代码里写的是 $\exp (g_{\text{cum}}[i]-g_{\text{cum}}[j])$ 而不是 $i-1$？ 因为 chunk 算法把"读取时的衰减"和"写入时的衰减"统一处理——在 chunk 末态整块更新时，token $j$ 的贡献要一直衰减到 chunk 末尾（含 $g_{\text{cum}}[C-1]$），所以用 $g_{\text{cum}}[i]$ 而非 $g_{\text{cum}}[i-1]$。这两种时机在朴素递归里是分开的（读取用 $i-1$，写入累积到末态用 $C-1$），chunk 把它们合并到一个矩阵里。

数值例子（$g=[0,\ln 0.5,0,\ln 0.5]$, $C=4$）：$g_{\text{cum}}=[0,\ln 0.5,\ln 0.5,\ln 0.25]\approx [0,-0.69,-0.69,-1.39]$。

$$\text{decay}=\left(\begin{array}{cccc}1& 0& 0& 0\\ e^{-0.69}& 1& 0& 0\\ e^{-0.69}& e^0& 1& 0\\ e^{-1.39}& e^{-0.69}& e^{-0.69}& 1\end{array}\right)=\left(\begin{array}{cccc}1& 0& 0& 0\\ 0.5& 1& 0& 0\\ 0.5& 1& 1& 0\\ 0.25& 0.5& 0.5& 1\end{array}\right)$$

读 $\text{decay}[3,1]=0.5$：按 chunk 公式 $e^{g_{\text{cum}}[3]-g_{\text{cum}}[1]}=e^{\ln 0.25-\ln 0.5}=e^{\ln 0.5}=0.5$。但按朴素"读取时机"$i-1$ 公式 $e^{g_{\text{cum}}[2]-g_{\text{cum}}[1]}=e^{\ln 0.5-\ln 0.5}=e^0=1$。chunk 用 0.5，朴素用 1——这正是下一节"进阶"里 $g$ 有衰减时两者分叉的数学根源：chunk 的衰减矩阵把"读取后到 chunk 末尾"的衰减也提前算进去了，朴素只算到"读取前"。$g$ 全 0 时两者都等于 1，所以对齐。

阶段 2 原理推导：attn 修正矩阵 $(I+A)^{-1}$ 为什么是这个形式

第 5.3 节讲了 delta rule 的核心是"减去预测值"：${\delta }_t=(v_t-k_t^{\mathrm{\top }}{S}_{t-1}){\beta }_t$。朴素递归里这是逐 token 做的——每个 token 用当时的 $S_{t-1}$ 预测、再写入误差。

chunk 的挑战：chunk 内 $C$ 个 token 要一次性并行算出各自的 $\delta$，但它们互相依赖（token $i$ 的 ${\delta }_i$ 依赖 $S_{i-1}$，而 $S_{i-1}$ 含 token $0..i-1$ 的 $\delta$）。这看似串行依赖，但有矩阵形式的解。

推导：记 $K_{\beta }=\beta \odot K$（带 beta 的 key），$V_{\beta }=\beta \odot V$。朴素递归里 token $i$ 的 ${\delta }_i=v_i{\beta }_i-k_i^{\mathrm{\top }}{S}_{i-1}{\beta }_i$，其中 $S_{i-1}$ 含 $\sum _{j<i}{k}_j{\delta }_j^{\mathrm{\top }}\cdot \text{decay}[i-1,j]$。代入得：

$${\delta }_i=v_i{\beta }_i-\sum _{j<i}(k_i^{\mathrm{\top }}{k}_j)\cdot {\delta }_j\cdot \text{decay}[i-1,j]\cdot {\beta }_i$$

写成矩阵形式（$\mathit{\delta }$ 是所有 token 的 delta 向量，$A_{ij}=-(k_i^{\mathrm{\top }}{k}_j){\beta }_i{\beta }_j\text{decay}[i,j]$ 严格下三角）：

$$\mathit{\delta }=V_{\beta }+A\mathit{\delta }⟹(I-A)\mathit{\delta }=V_{\beta }⟹\mathit{\delta }=(I-A{)}^{-1}{V}_{\beta }$$

注意符号：代码里的 attn = -(kb @ k^T) * decay 就是 $A$（带负号），所以 $(I-A{)}^{-1}=(I+A_{\text{code}}{)}^{-1}$。这就是为什么修正矩阵叫 $(I+A{)}^{-1}$。

级数展开：$(I+A{)}^{-1}=I-A+A^2-A^3+\cdots$（$A$ 严格下三角，级数有限，至多 $C-1$ 项就收敛，因为 $A^C=0$）。

• $I$（零阶）：token $i$ 的 $v_i{\beta }_i$ 自身
• $-A$（一阶）：减去 token $i$ 对前面 token 的直接预测——这就是朴素"减去 $k_i^{\mathrm{\top }}{S}_{i-1}$"的一阶部分
• $+A^2$（二阶）：经一个中间 token 的间接修正。比如 token 3 经 token 2 间接读 token 1 的贡献，符号反转（因为"减去的预测"本身又被前面的预测修正了一次）
• $\pm A^k$：经 $k-1$ 个中间 token 的间接修正，正负交替

为什么朴素只算一阶、chunk 要算所有阶？ 朴素逐 token 递归里，token $i$ 算 ${\delta }_i$ 时 $S_{i-1}$ 已经包含了 token $0..i-1$ 的完整 $\delta$（含它们之间的所有间接修正）——所以朴素"一次 $k_i^{\mathrm{\top }}{S}_{i-1}$"隐式吃掉了所有高阶项。chunk 并行算时没有这个"逐步累积"的便利，必须显式把 $A^2,A^3,\dots$ 都算出来。

递推代替求逆：直接算 $(I+A{)}^{-1}$ 要矩阵求逆（$O(C^3)$）。但 $A$ 严格下三角，逐行递推就能算：算第 $i$ 行时，前 $i-1$ 行已全部算好（含所有阶），用 attn[i,:i] = row + (row * sub).sum(-2) 一次性吸收所有经过中间 token 的间接路径，不必显式算 $A^2,A^3$ 每一项。这就是伪代码里那个 for i in range(1, C) 循环——$O(C^2)$ 比 $O(C^3)$ 求逆快，且是 FLA Triton kernel 的核心。

收敛性：$A$ 严格下三角 $\Rightarrow$ $A^C=0$（nilpotent），级数 $I-A+A^2-\cdots$ 至多 $C$ 项就截断，有限项精确成立，不是无穷级数近似。这是 chunk delta rule 数值上严格等价于朴素递归（在 $g$ 全 0 时）的数学保证。

阶段 3 原理推导：chunk 间整块衰减与 $v_{\text{new}}$

chunk 间递归把"chunk 内 $C$ 个 token 的 delta"压缩成"chunk 末态的一次状态更新"。要理解这个整块更新，先看朴素递归在 chunk 边界的状态演化。

朴素 chunk 边界：chunk 开始时状态 $S_{\text{start}}$，chunk 内 token $0..C-1$ 逐个写入 ${\delta }_j$ 并逐级衰减。到 chunk 末尾 $S_{\text{end}}$：

$$S_{\text{end}}=S_{\text{start}}\cdot e^{g_{\text{cum}}[C-1]}+\sum _{j=0}^{C-1}{k}_j{\delta }_j^{\mathrm{\top }}\cdot e^{g_{\text{cum}}[C-1]-g_{\text{cum}}[j]}$$

第一项：$S_{\text{start}}$ 经整个 chunk 的累积衰减 $e^{g_0+\cdots +g_{C-1}}=e^{g_{\text{cum}}[C-1]}$。

第二项：token $j$ 的 ${\delta }_j$ 写入后，从 $j+1$ 到 $C-1$ 继续衰减，到 chunk 末尾剩 $e^{g_{\text{cum}}[C-1]-g_{\text{cum}}[j]}$。定义 $k_{\text{weighted}}[j]=k_j\cdot e^{g_{\text{cum}}[C-1]-g_{\text{cum}}[j]}$，则：

$$S_{\text{end}}=S_{\text{start}}\cdot e^{g_{\text{cum}}[C-1]}+k_{\text{weighted}}^{\mathrm{\top }}\mathit{\delta }$$

这就是伪代码 S_end = S_start * exp(gcum[-1]) + k_weighted.T @ v_new，严格等价于朴素"逐 token 写入 + 后续衰减"的累积结果。

$v_{\text{new}}$ 的来历：chunk 内的 $\mathit{\delta }$ 不是 $V_{\beta }$，而是"修正后 value 减去从历史 $S_{\text{start}}$ 读出的预测"：

$$\mathit{\delta }=V_p-k_{\text{cumdecay}}\cdot S_{\text{start}}$$

其中 $k_{\text{cumdecay}}=\text{attn}@(K_{\beta }\odot e^{g_{\text{cum}}})$ 是"按 attn 修正权重加权、并含累积衰减"的 key——用它读 $S_{\text{start}}$ 得到"chunk 内各 token 对历史状态的预测"。减掉这个预测，剩下的就是 chunk 内的"新信息"（delta）。

输出 $o$：chunk 内 token $i$ 的输出 = 历史贡献 + chunk 内新信息：

$$o_i=\underset{\text{历史}}{\underbrace{(q_i\cdot e^{g_{\text{cum}}[i]})\cdot S_{\text{start}}}}+\underset{\text{chunk 内}}{\underbrace{\text{attn\_qk}[i]\cdot \mathit{\delta }}}$$

其中 $\text{attn\_qk}[i,j]=(q_i^{\mathrm{\top }}{k}_j)\cdot \text{decay}[i,j]$ 是 chunk 内 QK attention（未经修正矩阵处理，和阶段2的 attn_mod 不同）。

严格对齐条件：朴素和 chunk 的 $S_{\text{end}}$ 和 $o$ 严格相等，当且仅当 $g$ 全 0（此时 $\text{decay}$ 全 1，$A$ 的衰减因子消失，attn_mod 修正和朴素链式衰减归约一致）。$g$ 有衰减时，attn_mod 矩阵的"一次性内化间接修正"和朴素的"链式衰减累积"在中间过程上分叉，导致 chunk 末态 token 的输出和朴素有差——这是两种归约的固有差异，不是对错（见下一节"进阶"）。

同一例子跑 Chunk（g 全 0）

把 $K=V=I$, $\beta =1$, $C=2$, $g=[0,0]$ 代入 chunk 算法：

阶段	计算	结果
1	$\text{decay}=\left(\begin{array}{cc}1& 0\\ 1& 1\end{array}\right)$（g 全 0，exp=1）	$gcum=[0,0]$
2	$A=0$（$k_i\cdot k_j$ 内积对正交基单位阵全 0），$\text{attn}=I$	$\text{attn}=\left(\begin{array}{cc}1& 0\\ 0& 1\end{array}\right)$
2	$V_p=I@V=V=I$	$V_p=\left(\begin{array}{cc}1& 0\\ 0& 1\end{array}\right)$
3	$v_{\text{new}}=V_p-(I@K\odot 1)@0=V_p$	$v_{\text{new}}=\left(\begin{array}{cc}1& 0\\ 0& 1\end{array}\right)$
3	$S_{\text{end}}=0+(I\odot 1{)}^T@v_{\text{new}}=I$	$S_{\text{end}}=I$
输出	$o_t=\text{scale}\cdot k_t@S_{\text{end}}=(1/\sqrt{2})\cdot k_t$	$o=[[1/\sqrt{2},0],[0,1/\sqrt{2}]]$

对比朴素：

	t=0	t=1
朴素 $o$（未乘 scale）	$[1,0]$	$[0,1]$
chunk $o$（未乘 scale）	$[1,0]$	$[0,1]$
朴素 $o$（乘 scale=1/√2）	$[1/\sqrt{2},0]$	$[0,1/\sqrt{2}]$
chunk $o$（乘 scale=1/√2）	$[1/\sqrt{2},0]$	$[0,1/\sqrt{2}]$

逐元素一致 ✓。这说明 chunk 在 g 全 0 时数学等价于朴素——它把串行递归"重写"成"块内矩阵并行 + 块间递归"，不改变结果。

进阶：引入门控衰减时，两者怎么对齐？

问题来了：g 全 0 时 chunk = 朴素，但 Gated DeltaNet 的"gated"（门控衰减）就是 Gated DeltaNet 区别于普通 DeltaNet 的核心——如果 g 全 0 就退化成普通 DeltaNet 了。

跑同一组 K=V=单位基，但 g = $[0,\ln 0.5,0,\ln 0.5]$（即 $g_1=g_3=\ln 0.5$，$g_0=g_2=0$）：

t	朴素 ${\delta }_t$	朴素 $o_t$	chunk $v_{\text{new}}[t]$	chunk $o_t$（FLA 实跑）
0	$[1,0]$	$[0.707,0]$	$[1,0]$	$[0.707,0]$
1	$[0,1]$	$[0,0.707]$	$[0,1]$	$[0,0.707]$
2	$[1,0]$	$[1.061,0]$	$[1,0]$	$[1.061,0]$
3	$[0,1]$（朴素链式衰减已减 v[1] 贡献）	$[0,1.414]$	$[0,2]$（attn_mod 多减了一份）	$[0,2.121]$

结论：

1. g 全 0 时（t=0,1 那个 chunk），朴素和 chunk o 逐元素一致——chunk 是正确的并行实现
2. g 有衰减时（t=2,3 那个 chunk），朴素和 chunk o 在每个 chunk 的末态 token（t=1, t=3）开始分叉

为什么会分叉？ 朴素和 chunk 的状态累积方式不同：

• 朴素：S 在每个 t 都被 *exp(g[t]) 衰减，前 token 写入的 delta 在 S_prev 中已含"从写入到 t-1 的链式衰减"
• chunk：attn_mod 矩阵（(I+A)^{-1}）一次性内化了"经中间 token 间接读 v[j] 的修正"——这是一次性算完所有间接路径，不是链式累积

两者数学上不等价（在 g 有衰减时），但都是正确的 delta rule——只是"归约"方式不同。真实工程（FLA 官方 chunk delta rule kernel）就采用 chunk 这种形式，因为它块内可并行，是 GPU 友好的实现。

为什么这两种归约都对？

教科书 delta rule 的唯一定义是逐 token 形式。chunk 用 attn_mod (I+A)^{-1} 等价于逐 token 形式——前提是 g 全 0。g 有衰减时，$(I+A{)}^{-1}$ 和链式累积在数学上不等价，但两者都是"合理的 delta rule 近似"：

• 朴素 = "实时链式累积"（教学清晰，但串行）
• chunk = "矩阵一次性内化"（并行友好，FLA 工业选择）

选哪种取决于场景：

• 教学/单步推理/小序列 → 朴素
• 训练/批量推理/长序列 → chunk（FLA 的 chunk_gated_delta_rule / fused_recurrent_gated_delta_rule 都用 chunk 形式）

进阶读者注：FLA 实际还有第三种形式——fused_recurrent_gated_delta_rule（decode 用的纯递归 kernel，不分 chunk），它在 g 有衰减时和 chunk 形式数学上严格等价。教学时先掌握"朴素 vs chunk"这两层，再去看 FLA decode kernel 的简化递归即可。

完整代码（带详细注释）

def torch_chunk_gated_delta_rule(
    query, key, value, g, beta,
    chunk_size=64, initial_state=None, use_qk_l2norm=True,
):
    """
    Gated Delta Rule 的 Chunk 模式实现

    Args:
        query:  (batch, seq_len, num_heads, key_dim)     查询
        key:    (batch, seq_len, num_heads, key_dim)     键
        value:  (batch, seq_len, num_heads, value_dim)   值
        g:      (batch, seq_len, num_heads)              门控衰减 (log-space)
        beta:   (batch, seq_len, num_heads)              更新强度 (0~1)
        chunk_size: 块大小 (默认 64)

    Returns:
        output:      (batch, seq_len, num_heads, value_dim)
        last_state:  (batch, num_heads, key_dim, value_dim)  最终递归状态
    """
    # ---- 前处理 ----
    if use_qk_l2norm:
        query = l2norm(query, dim=-1)   # L2 归一化，稳定训练
        key = l2norm(key, dim=-1)

    # 转为 float32 计算（避免半精度溢出）
    query, key, value, beta, g = [x.to(torch.float32) for x in (query, key, value, beta, g)]

    # Padding 到 chunk_size 的整数倍
    pad_size = (chunk_size - seq_len % chunk_size) % chunk_size
    query = F.pad(query, (0, 0, 0, pad_size))
    # ... 对 key, value, beta, g 同样 padding

    # Q 缩放
    scale = 1 / sqrt(key_dim)
    query = query * scale

    # 预计算 beta * value 和 beta * key
    v_beta = value * beta[..., None]   # (batch, heads, seq, v_dim)
    k_beta = key * beta[..., None]     # (batch, heads, seq, k_dim)

    # ---- Reshape 为 chunks ----
    # (batch, heads, seq, dim) → (batch, heads, num_chunks, chunk_size, dim)
    query, key, value, k_beta, v_beta = [
        x.reshape(batch, heads, -1, chunk_size, dim)
        for x in (query, key, value, k_beta, v_beta)
    ]
    g = g.reshape(batch, heads, -1, chunk_size)
    num_chunks = total_seq_len // chunk_size

    # ============================================================
    # 阶段 1: Chunk 内预处理
    # ============================================================

    # 累积门控衰减: g[i] = sum(g[0:i+1])
    g = g.cumsum(dim=-1)

    # 衰减矩阵: decay[i,j] = exp(g[i] - g[j]) for i >= j
    # 表示 chunk 内 token j 对 token i 的衰减因子
    decay_mask = (g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp()

    # ============================================================
    # 阶段 2: Chunk 内 Delta Rule 修正
    # ============================================================

    # 初始 attn 矩阵: 记录 chunk 内 token 之间的关系
    # attn[i,j] = -beta[i] * <k[i], k[j]> * decay[i,j]
    attn = -(k_beta @ key.transpose(-1, -2)) * decay_mask
    attn = attn.masked_fill(上三角 mask, 0)  # 因果性

    # 递推修正: 让 attn 包含 delta rule 的修正效果
    # 这是算法的核心——用递推代替逐 token 递归
    for i in range(1, chunk_size):
        row = attn[..., i, :i].clone()      # token i 和前面 tokens 的直接关系
        sub = attn[..., :i, :i].clone()      # 前面 tokens 之间的修正矩阵
        # 修正: 加上 token i 通过中间 token 的间接影响
        attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)

    # 加上单位矩阵（自身到自身的映射）
    attn = attn + torch.eye(chunk_size)

    # 用修正后的 attn 矩阵计算 chunk 内的 value
    # value[i] = sum_j attn[i,j] * beta[j] * value[j]
    value = attn @ v_beta

    # 用于从历史状态读取的累积衰减
    k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))

    # ============================================================
    # 阶段 3: Chunk 间递归
    # ============================================================

    # 初始化递归状态 S: (batch, heads, key_dim, value_dim)
    # S 存储 "key → value" 的关联记忆
    S = initial_state or zeros(key_dim, value_dim)

    core_attn_out = torch.zeros_like(value)

    for i in range(num_chunks):
        q_i = query[:, :, i]   # 当前 chunk 的 Q
        k_i = key[:, :, i]     # 当前 chunk 的 K
        v_i = value[:, :, i]   # 当前 chunk 的 V（已修正）

        # ---- 读取历史记忆 ----
        # 用当前 chunk 的 key 去查历史状态 S，得到预测值
        v_prime = k_cumdecay[:, :, i] @ S     # (batch, heads, C, v_dim)

        # ---- 计算 Delta ----
        # 实际值(修正后) - 预测值 = 惊讶程度
        v_new = v_i - v_prime

        # ---- 生成输出 ----
        # 历史贡献: q 和历史状态的交互（带衰减）
        attn_inter = (q_i * exp(g[:, :, i, :, None])) @ S
        # chunk 内新信息: q 和 chunk 内修正后 value 的交互
        attn_intra = chunk_attn @ v_new
        # 最终输出 = 历史 + 新信息
        core_attn_out[:, :, i] = attn_inter + attn_intra

        # ---- 更新递归状态 ----
        # 旧状态衰减
        S = S * exp(g[:, :, i, -1, None, None])
        # 写入新信息: k^T @ v_new（带衰减权重）
        k_weighted = k_i * exp(g[:, :, i, -1, None] - g[:, :, i])[..., None]
        S = S + k_weighted.transpose(-1, -2) @ v_new

    # 去掉 padding，恢复原始长度
    core_attn_out = core_attn_out[:, :, :seq_len]
    return core_attn_out, S

Chunk 模式的复杂度分析

	标准注意力	逐 token DeltaNet	Chunk DeltaNet
时间复杂度	O(n²d)	O(nd²)	O(n·d² + n·C²)
空间复杂度	O(n²)	O(d²)	O(d² + C²)
并行性	完全并行	完全串行	块内并行

其中 $C$ 是 chunk_size（通常 64），$d$ 是维度。当 $d<n$ 时（大序列场景），DeltaNet 有显著优势。

5.5 Qwen3.5 的混合层布局

设计理念

Qwen3.5 不是所有层都用线性注意力，而是混合使用两种注意力：

# config.py
full_attention_interval: int = 4  # 每 4 层一个 full_attention

# 自动生成层布局
layer_types = []
for i in range(num_hidden_layers):
    if (i + 1) % full_attention_interval == 0:
        layer_types.append("full_attention")      # GQA
    else:
        layer_types.append("linear_attention")    # Gated DeltaNet

以 24 层为例：

Layer  0: GatedDeltaNet (线性注意力)
Layer  1: GatedDeltaNet
Layer  2: GatedDeltaNet
Layer  3: GQA           (全注意力) ← 每 4 层一个
Layer  4: GatedDeltaNet
Layer  5: GatedDeltaNet
Layer  6: GatedDeltaNet
Layer  7: GQA
Layer  8: GatedDeltaNet
Layer  9: GatedDeltaNet
Layer 10: GatedDeltaNet
Layer 11: GQA
...
Layer 20: GatedDeltaNet
Layer 21: GatedDeltaNet
Layer 22: GatedDeltaNet
Layer 23: GQA

24 层中有 18 层线性注意力 + 6 层全注意力。

为什么混合使用？

1. 效率：大部分层用 O(n) 的线性注意力，只有少数层用 O(n²) 的全注意力
2. 质量：全注意力层可以捕获长距离依赖，线性注意力层负责局部特征和高效推理
3. KV Cache：只有全注意力层需要 KV Cache，大幅减少显存占用

DecoderLayer 的统一接口

class DecoderLayer(nn.Module):
    def __init__(self, config, layer_idx):
        self.layer_type = config.layer_types[layer_idx]

        # 根据层类型选择 mixer
        if self.layer_type == "linear_attention":
            self.linear_attn = GatedDeltaNet(config, layer_idx)
        elif self.layer_type == "full_attention":
            self.self_attn = GQAAttention(config, layer_idx)

        # FFN 和归一化层是共享的
        self.mlp = SwiGLU(config.hidden_size, config.intermediate_size)
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(self, hidden_states, attention_mask=None, position_ids=None):
        # Token Mixer
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)

        if self.layer_type == "linear_attention":
            hidden_states = self.linear_attn(hidden_states)  # 不需要 mask 和 position_ids
        elif self.layer_type == "full_attention":
            hidden_states = self.self_attn(
                hidden_states, attention_mask=attention_mask, position_ids=position_ids
            )

        hidden_states = residual + hidden_states

        # FFN
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        return hidden_states

5.6 注意力复杂度对比总结

注意力类型	训练复杂度	推理复杂度 (per token)	KV Cache	代表模型
Full Attention (MHA)	O(n²d)	O(nd) per step	O(n·d)	GPT-3
Full Attention (GQA)	O(n²d)	O(nd) per step	O(n·d/g)	Qwen3.5 (部分层)
Linear Attention	O(nd²)	O(d²) per step	O(d²)	Linear Transformer
DeltaNet	O(nd²)	O(d²) per step	O(d²)	—
Gated DeltaNet	O(nd²)	O(d²) per step	O(d²)	Qwen3.5 (部分层)

5.7 总结

线性注意力的核心思想：

1. 利用矩阵结合律，将 O(n²) 降到 O(n)
2. 用递归状态矩阵 $S$ 压缩历史信息，大小固定为 O(d²)
3. 推理时每步只需 O(d²)，不随序列长度增长

Gated DeltaNet 的三重控制：

1. delta $(v_t-\text{预测值})$：只更新预测错误的部分
2. beta ${\beta }_t$：控制写入多少新信息
3. gate $g_t$：控制遗忘多少旧信息

Qwen3.5 的设计：

• 每 4 层中 3 层用 Gated DeltaNet，1 层用 GQA
• 全注意力层保证长距离依赖建模
• 线性注意力层提供高效推理
• 混合布局在效率和质量之间取得平衡

---

上一章：4. 注意力机制 | 下一章：6_miniqwen.md