Blog Tags GitHub
Back to blog
image-generation diffusion-model llm-training

Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models

块级离散去噪扩散语言模型，在自回归和扩散之间插值，实现任意长度生成和 SOTA 似然

Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models

一、论文概述

项目内容
标题Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models
作者Marianne Arriola, Aaron Gokaslan, Justin T. Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Subham Sekhar Sahoo, Volodymyr Kuleshov
机构Cornell University
论文https://arxiv.org/abs/2503.09573
代码https://github.com/kuleshov-group/bd3lms
项目https://m-arriola.com/bd3lms
发布2025-03-12 (v1), 最新 v3
会议ICLR 2025 Oral
许可MIT

二、核心思想

问题定义

语言模型存在两种主要范式，各有优劣：

范式优势劣势
自回归 (AR)灵活长度生成、KV 缓存、高质量串行生成、无法并行
扩散 (Diffusion)并行生成、可控性固定长度、似然建模较弱

核心洞察

自回归和扩散并非互斥——可以通过块级分解将两者结合：

  • 将 token 序列分成固定大小的块 (blocks)
  • 块间使用自回归生成（支持灵活长度 + KV 缓存）
  • 块内使用离散去噪扩散（支持并行生成）

通过调整块大小 (block size)，可以在自回归和扩散之间连续插值：

  • 块大小 = 1 → 纯自回归
  • 块大小 = 序列长度 → 纯扩散

解决方案概述

BD3-LMs (Block Discrete Denoising Diffusion Language Models)：

  1. 块级自回归似然参数化：将序列分解为块，块间自回归，块内扩散
  2. 数据驱动噪声调度：最小化训练梯度方差
  3. 任意长度离散扩散采样器：突破固定长度限制
  4. 高效训练算法：梯度方差估计器

三、技术架构

块级分解

给定 token 序列 x=(x1,x2,...,xL)\mathbf{x} = (x_1, x_2, ..., x_L)，将其分为 BB 个块：

x=(x(1),x(2),...,x(B))\mathbf{x} = (\mathbf{x}^{(1)}, \mathbf{x}^{(2)}, ..., \mathbf{x}^{(B)})

每个块 x(b)\mathbf{x}^{(b)} 包含 KK 个 token（最后一个块可能更短）。

块级自回归似然

pθ(x)=∏b=1Bpθ(x(b)∣x(<b))p_\theta(\mathbf{x}) = \prod_{b=1}^{B} p_\theta(\mathbf{x}^{(b)} | \mathbf{x}^{(<b)})

  • 块间：自回归条件化，支持 KV 缓存
  • 块内：离散去噪扩散，支持并行生成

离散去噪扩散

每个块内部使用离散扩散过程：

前向过程：逐步 mask token q(xt(b)∣x0(b))=MaskSchedule(t)q(\mathbf{x}^{(b)}_t | \mathbf{x}^{(b)}_0) = \text{MaskSchedule}(t)

反向过程：逐步恢复 token pθ(xt−1(b)∣xt(b),x(<b))p_\theta(\mathbf{x}^{(b)}_{t-1} | \mathbf{x}^{(b)}_t, \mathbf{x}^{(<b)})

数据驱动噪声调度

传统扩散使用固定噪声调度。BD3-LM 提出数据驱动的噪声调度：

  • 估计不同时间步的梯度方差
  • 选择最小化训练方差的调度
  • 提高训练稳定性和收敛速度

任意长度生成

标准扩散语言模型（如 SEDD、MDLM）只能生成固定长度序列。BD3-LM 通过块级自回归结构自然支持任意长度：

  1. 生成第一个块（扩散）
  2. 条件化已生成块，生成下一个块
  3. 重复直到生成 EOS token 或达到最大长度

四、核心创新

创新点说明理论/实验依据
块级分解将序列分为块，块间 AR + 块内扩散连续插值 AR 和扩散
灵活长度生成突破扩散模型固定长度限制块级自回归结构
KV 缓存支持块间自回归天然支持 KV 缓存推理效率提升
数据驱动噪声调度最小化训练梯度方差收敛更快更稳定
梯度方差估计器高效估计训练方差理论分析
并行 token 采样块内 token 并行生成比纯 AR 更快

五、实验结果

语言建模基准

BD3-LM 在多个标准基准上达到扩散模型 SOTA：

数据集BD3-LMSEDDMDLMAR
OpenWebTextSOTA↓↓参考
PubMedSOTA↓↓参考
arXivSOTA↓↓参考

块大小权衡

块大小质量速度特性
K=1AR 级别串行纯自回归
K=8接近 AR更快平衡点
K=32略降更快更多并行
K=∞扩散级别最快(并行)纯扩散

生成质量

  • 困惑度 (Perplexity)：接近甚至匹配自回归模型
  • 生成速度：比纯自回归快（块内并行）
  • 可控性：支持 inpainting、条件生成等扩散特性

与现有方法对比

方法类型似然长度灵活KV 缓存并行生成
AR自回归高✅✅❌
SEDD扩散中❌❌✅
MDLM扩散中❌❌✅
BD3-LM混合高✅✅✅

六、相关工作整合

扩散语言模型

方法关键特点与 BD3-LM 关系
SEDDScore Entropy 离散扩散基线，固定长度
MDLMMasked 扩散语言模型基线，固定长度
SSD-LM半自回归 Simplex 扩散相关思路
LLaDA大规模离散扩散相关方向

自回归 + 扩散混合

方法关键特点与 BD3-LM 关系
Diffusion ForcingAR 结构 + 序列去噪相关思路
MDLM+AR混合方法类似动机

位置论文

方法关键特点与 BD3-LM 关系
Inference-time Scaling序列扩展 + 状态精炼两轴BD3-LM 是两轴结合的实例

七、代码与资源

代码仓库

代码结构

bd3lms/
├── main.py           # 训练和评估例程
├── noise_schedule.py # 噪声调度
├── diffusion.py      # 前向/反向扩散
├── dataloader.py     # 数据加载
├── utils.py          # LR 调度器、日志
└── models/           # 网络架构

预训练模型

  • HuggingFace: kuleshov-group/bd3-lms
  • 提供多个规模的预训练模型

基线实现

代码包含以下基线的完整实现：

  1. AR: 自回归模型
  2. SEDD: Score Entropy 离散扩散
  3. MDLM: Masked 扩散语言模型
  4. SSD-LM: 半自回归 Simplex 扩散（仅采样）

八、总结

核心贡献

  1. 概念创新：块级分解统一自回归和扩散语言模型
  2. 灵活插值：通过块大小在 AR 和扩散之间连续调节
  3. 实用价值：支持 KV 缓存 + 任意长度 + 并行生成
  4. 训练优化：数据驱动噪声调度最小化梯度方差
  5. SOTA 性能：扩散模型在语言建模基准上的新 SOTA
  6. 影响力：ICLR 2025 Oral，1000+ GitHub Stars

技术影响

  • 统一视角：打破了 AR 和扩散的二分法
  • 实用框架：为扩散语言模型提供了实用的工程方案
  • 研究方向：启发了更多混合生成范式的研究
  • 代码开源：完整的训练/评估代码和预训练模型

局限性

  • 块大小需要手动调节
  • 块内扩散仍有一定质量损失
  • 相比纯自回归，在极长序列上的优势有限

九、参考资源