训练具有部分双向性的混合块扩散语言模型

# 训练具有部分双向性的混合块扩散语言模型

**来源：** https://arxiv.org/html/2607.02805  
**作者：** Pranshu Chaturvedi（斯坦福大学，[email protected]）、Parth Shroff（斯坦福大学，[email protected]）、Tarun Suresh（斯坦福大学，[email protected]）、Hangoo Kang（斯坦福大学，[email protected]）、Kaiyue Wen（斯坦福大学，[email protected]）

###### 摘要

高吞吐量的长上下文生成是大型语言模型面临的核心挑战之一。生成过程通常受限于内存带宽而非计算能力：每个解码步骤都必须从内存中读取累积的键/值（KV）缓存，因此带宽需求随上下文长度增长，而每次只生成一个token。为此，出现了两种并行的方法：一是通过高效注意力变体和线性时间混合器（如Mamba）来减少内存访问，二是通过一次生成多个token的块来增加并行计算。然而，将这两种思路结合起来会面临技术挑战。早期的混合扩散模型（如DiffuMamba）使用了双向Mamba混合，其中包括相对于因果生成的逆向扫描。这种逆向扫描需要扫描整个序列，因此其状态并非仅基于前缀，即使在逐块扩散的情况下也无法精确地作为缓存复用。我们提出了一种BDLM Mamba-注意力混合模型，通过将逆向Mamba扫描限制在活跃的去噪块内来解决这一挑战，从而实现了跨块的精确缓存。在87M参数的DCLM扫描中，BDLM Mamba-H在C4-en验证困惑度上达到了最佳，优于BDLM注意力和全序列基线。在350M参数规模下，它与BDLM注意力相比仍具有竞争力。对于长上下文推理，BDLM Mamba-H在65K tokens时达到全序列DiffuMamba-H吞吐量的19.7倍，在262K tokens时达到BDLM注意力吞吐量的3.7倍，这表明Mamba混合架构是一种有潜力的长上下文扩散架构。

## 1 引言

在长上下文中服务语言模型主要是一个内存带宽问题。在自回归解码中，每个生成的token都会读取累积的键/值（KV）缓存，因此随着上下文长度增长，移动的字节数增加，而每次只生成一个token。两条研究路线从不同角度应对这一压力。混合自回归模型通过将许多注意力层替换为线性时间循环或状态空间混合器（如Mamba）来减少缓存流量（Gu和Dao，2024 (https://arxiv.org/html/2607.02805#bib.bib4)；Dao和Gu，2024 (https://arxiv.org/html/2607.02805#bib.bib5)；Qin等人，2024 (https://arxiv.org/html/2607.02805#bib.bib31)；Kimi Team等人，2025 (https://arxiv.org/html/2607.02805#bib.bib41)；Li等人，2025 (https://arxiv.org/html/2607.02805#bib.bib33)；Ling Team等人，2025 (https://arxiv.org/html/2607.02805#bib.bib42)）。掩码扩散语言模型则通过迭代去噪并行生成许多token（Austin等人，2021a (https://arxiv.org/html/2607.02805#bib.bib36)；Hoogeboom等人，2021 (https://arxiv.org/html/2607.02805#bib.bib37)；Sahoo等人，2024 (https://arxiv.org/html/2607.02805#bib.bib38)；Nie等人，2025 (https://arxiv.org/html/2607.02805#bib.bib39)）。块扩散语言模型（BDLM）通过从左到右生成块使扩散过程可缓存：完成的先前块成为固定的条件，当前块则并行去噪（Arriola等人，2025 (https://arxiv.org/html/2607.02805#bib.bib40)）。

DiffuMamba将扩散与Mamba结合，并展示了全序列去噪器在长上下文推理中的高吞吐量（Singh等人，2026 (https://arxiv.org/html/2607.02805#bib.bib13)）。然而，其完全双向的Mamba扫描包含对整个去噪窗口的逆向传递。这些逆向状态依赖于活跃的token，因此不是可复用的前缀状态。我们研究了一种原生BDLM Mamba-注意力混合架构BDLM Mamba-H，它将逆向Mamba扫描限制在活跃的去噪块内。已完成的先前块被存储为注意力层的注意力KV状态和Mamba层的正向Mamba边界状态。全块目标在训练期间构建这些相同的前缀缓存，并让下游块的损失直接训练它们，而活跃块则保持局部双向混合。

在这项工作中，我们探索可缓存的BDLM Mamba混合架构是否能在保持竞争力验证质量的同时，提高长上下文推理吞吐量。我们还描述了仅适用于BDLM Mamba-H当前活跃块的时间步条件分解，因此时间步调制可以在不使干净前缀缓存依赖于逆向扩散步骤的情况下应用（附录D (https://arxiv.org/html/2607.02805#A4)）。我们的贡献如下：

```mermaid
flowchart TD
    subgraph 干净前缀 c^(k)
        c
    end
    subgraph 当前块
        b1⋆
        b2⋆
        b3⋆
    end
    subgraph 未来块（掩码掉）
        b̃_t^(k)
        b_{t+1}⋆
        b_{t+2}⋆
    end
    Mamba前缀扫描 --> Cache(c)
    Cache(c) --> 活跃块去噪器
    活跃块去噪器 --> t
    t --> CE损失
    CE损失 -.->|梯度流经 Mamba前缀状态| Mamba前缀扫描
```

**图1：** BDLM Mamba-H 训练。全块目标在所有块位置上并行应用此活跃块计算。实线箭头表示前向计算，绿色虚线箭头表示梯度流，黑色虚线箭头表示可选的时间步输入。去噪器接收干净前缀缓存、带噪的当前块以及可选的时间步条件，而未来块保持掩码状态。对于Mamba混合层，缓存的前缀是仅正向的，逆向扫描仅局部于活跃块，因此活跃块的损失训练与采样器使用的相同前缀缓存对象。

- **缓存对齐的训练。** 我们构建了一个BDLM全块目标，其缓存对象与推理时的块缓存匹配：注意力层的注意力KV状态和循环层的Mamba前缀状态。对于BDLM Mamba-H架构，下游块的损失直接训练用于构建推理时复用循环状态的前缀扫描。
- **受控验证与扩展。** 我们在DCLM上使用调整后的超参数，以87M参数训练了全序列注意力、全序列DiffuMamba-H、BDLM注意力和BDLM Mamba-H。BDLM Mamba-H在87M规模下达到了最佳的C4-en验证困惑度61.6，而BDLM注意力为76.5，全序列基线分别为83.7和87.6。我们使用Complete(d)P风格的学习率迁移规则（Mlodozeniec等人，2025 (https://arxiv.org/html/2607.02805#bib.bib14)）将BDLM注意力和BDLM Mamba-H扩展到350M参数；在此规模下，BDLM Mamba-H在困惑度上仍与BDLM注意力相当。
- **长上下文吞吐量提升。** 我们在8x A100-80GB上评估了8192 token上下文下的350M训练吞吐量，并在1x A100-80GB上以batch size 1评估了700M随机初始化推理吞吐量。随着生成长度增长，BDLM Mamba-H在吞吐量上超越了BDLM注意力（以及全序列DiffuMamba-H），在65K时达到全序列DiffuMamba-H吞吐量的19.7倍，在262K时达到BDLM注意力吞吐量的3.7倍。

## 2 相关工作

#### 扩散语言模型。

离散扩散语言模型通过掩码或离散噪声破坏文本，并训练去噪器重建干净token（Austin等人，2021a (https://arxiv.org/html/2607.02805#bib.bib36)；Hoogeboom等人，2021 (https://arxiv.org/html/2607.02805#bib.bib37)；Sahoo等人，2024 (https://arxiv.org/html/2607.02805#bib.bib38)；Nie等人，2025 (https://arxiv.org/html/2607.02805#bib.bib39)）。BDLM增加了自回归块分解，使得每个完成的块成为后续块的干净前缀上下文（Arriola等人，2025 (https://arxiv.org/html/2607.02805#bib.bib40)）。这为块扩散提供了原生的缓存边界，这与全序列MDLM不同，后者的隐藏状态与当前带噪序列绑定。免训练缓存方法通过近似KV复用或自适应缓存，将复用技术改造应用于预训练的扩散语言模型（Wu等人，2026 (https://arxiv.org/html/2607.02805#bib.bib8)；Ma等人，2025 (https://arxiv.org/html/2607.02805#bib.bib9)；Nguyen-Tri等人，2025 (https://arxiv.org/html/2607.02805#bib.bib11)；Liu等人，2025 (https://arxiv.org/html/2607.02805#bib.bib10)；Wang等人，2025 (https://arxiv.org/html/2607.02805#bib.bib12)）。我们的重点则是训练一种块扩散架构，其缓存既是目标的一部分，也是采样器的一部分。

#### 用于长上下文生成的Mamba混合架构。

Mamba用选择性状态空间扫描替代了token对之间的注意力，其转移、输入和输出映射是当前token的函数（Gu和Dao，2024 (https://arxiv.org/html/2607.02805#bib.bib4)）。在下面使用的符号中，Mamba层在位置或块i之后、层ℓ中携带循环边界状态s_i^ℓ；该状态包括卷积状态和选择性SSM状态，这些状态是无需重新处理前缀即可继续正向扫描所必需的。Mamba-2通过结构化状态空间对偶性重新诠释了这一系列，产生了更高效的扫描内核以及与注意力更紧密的关系（Dao和Gu，2024 (https://arxiv.org/html/2607.02805#bib.bib5)）。最近的自回归系统越来越多地使用混合堆栈，保留一些全注意力层，同时用循环、线性注意力或状态空间模块替换许多层：例如Jamba的Transformer–Mamba混合专家设计、MiniMax-01的闪电注意力和softmax注意力混合、Qwen3-Next的门控DeltaNet/全注意力混合，以及Kimi Linear的Kimi Delta注意力/MLA混合（Lieber等人，2024 (https://arxiv.org/html/2607.02805#bib.bib43)；Li等人，2025 (https://arxiv.org/html/2607.02805#bib.bib33)；Qwen Team，2025 (https://arxiv.org/html/2607.02805#bib.bib44)；Kimi Team等人，2025 (https://arxiv.org/html/2607.02805#bib.bib41)）。这些自回归模型将混合作为实际的长上下文扩展方向；我们的工作研究块扩散的类似问题，其中缓存复用必须与去噪目标对齐，而不是与从左到右的下一token似然对齐。DiffuMamba-H将注意力和双向Mamba层交错用于掩码扩散去噪，并证明基于Mamba的去噪器提高了长上下文推理吞吐量（Singh等人，2026 (https://arxiv.org/html/2607.02805#bib.bib13)）。

## 3 背景

Mamba是一种选择性状态空间序列混合器，它用循环扫描替代了token间的注意力，并且该循环状态可以在序列片段间延续（Gu和Dao，2024 (https://arxiv.org/html/2607.02805#bib.bib4)）。Mamba-2通过结构化状态空间对偶性和更高效的内核改进了这一系列（Dao和Gu，2024 (https://arxiv.org/html/2607.02805#bib.bib5)）。在扩散语言模型中，DiffuMamba-H使用稀疏混合调度，将注意力和双向Mamba层交错，展示了全序列扩散去噪器可以从线性时间混合器中受益（Singh等人，2026 (https://arxiv.org/html/2607.02805#bib.bib13)）。当前工作采用了相同的广泛混合动机，但修改了扫描方向，使得BDLM生成可以在训练和推理时原生地复用前缀缓存。在层ℓ处，我们将简化的选择性状态空间更新写为

s_i^ℓ = A_i^ℓ s_{i-1}^ℓ + B_i^ℓ x_i^ℓ, y_i^ℓ = C_i^ℓ s_i^ℓ + D^ℓ x_i^ℓ, （1）

其中A_i^ℓ, B_i^ℓ, C_i^ℓ是依赖于token的投影，s_i^ℓ是token或块i后的循环边界状态。前缀缓存是在块边界处这些正向Mamba状态以及注意力层的注意力KV张量的集合。

掩码离散扩散通过用掩码替换token来破坏干净文本，并训练去噪器恢复原始token（Austin等人，2021b (https://arxiv.org/html/2607.02805#bib.bib1)；Sahoo等人，2024 (https://arxiv.org/html/2607.02805#bib.bib2)）。如果x̃_t是时间步t处被破坏的序列，全序列去噪器优化

L_full(θ) = E_{x,t,M_t} [ 1/|M_t| ∑_{i∈M_t} CE( f_θ(x̃_t, t)_i, x_i ) ]。 （2）

该目标为每个掩码token提供了覆盖去噪窗口的双向上下文。它也使整个窗口成为时间步条件计算的一部分。在生成过程中，隐藏状态与当前带噪序列绑定，因此稳定前缀没有可复用的缓存边界，必须在每个逆向步骤重新处理。

BDLM通过将序列分区为K个块b^{⋆(1)}, ..., b^{⋆(K)}来引入该边界。遵循BDLM公式（Arriola等人，2025 (https://arxiv.org/html/2607.02805#bib.bib40)），似然在块上是自回归的，

log p_θ(x) = ∑_{k=1}^K log p_θ( b^{⋆(k)} | b^{⋆(<k)} ), （3）

并且每个块的条件分布由仅针对当前块的离散扩散过程表示（Arriola等人，2025 (https://arxiv.org/html/2607.02805#bib.bib40)）。将扩散NELBO应用于每一项得到块求和目标

- log p_θ(x) ≤ L_BD(x;θ) := ∑_{k=1}^K L( b^{⋆(k)}, b^{⋆(<k)}; θ )。 （4）

对于BDLM风格训练中使用的掩码连续时间参数化，这可以写成加权当前块去噪损失的和，

L_BD(x;θ) = ∑_{k=1}^K E_{t,q} [ w(t) CE( f_θ( b_t^{(k)}, b^{⋆(<k)}, t ), b^{⋆(k)} ) ], w(t) = -α'_t / (1 - α_t), （5）

其中b_t^{(k)} ∼ q_t(· | b^{⋆(k)})是带噪的当前块，b^{⋆(<k)}是干净前缀，未来块被排除在外。全块训练目标在同一个训练窗口内的块边界上评估这个和，因此每个块都被训练为以其干净前缀为条件的当前去噪目标。

这种分解使得块扩散的缓存成为可能。对于注意力BDLM，块k的去噪器可以用与自回归Transformer相同的缓存接口写出：

z^{(k)}, K_k, V_k ← f_θ^{(k)}( b_t^{(k)}, K_{1:k-1}, V_{1:k-1}) ≡ f_θ^{(k)}( b_t^{(k)}, b^{⋆(<k)} ), （6）

其中z^{(k)}是当前块的干净token logits，K_{1:k-1}, V_{1:k-1}是从已完成的先前块计算出的键/值张量。在采样过程中，模型并行去噪当前块内的所有token，同时读取此前缀缓存；在块完成后，对新区块进行一次干净前向传递，将K_k, V_k附加以备后续使用。