TUBE: Tangent Upper Bound on Evidence for Discrete Diffusion Language Models¶
会议: NeurIPS 2026
arXiv: 2605.24292
论文: https://arxiv.org/abs/2605.24292
代码: https://github.com/kuleshov-group/bd3lms
领域: LLM 其他
关键词: 离散扩散模型 / 证据切线上界 / 对数似然评估 / 任意阶自回归 / 块扩散模型
一句话总结¶
针对离散扩散与任意阶自回归模型无法精确计算对数似然的问题,本文提出了基于对数函数切线展开的证据切线上界(TUBE),构建了严格无偏的有限样本蒙特卡洛上界估计量,并首次通过双向似然界定证实离散块扩散模型的真实似然依然显著落后于标准自回归基线。
研究背景与动机¶
自回归模型(ARM)凭借固定的从左至右因果分解,一直是现代大规模语言建模的主导范式,不仅易于扩展,还能提供精确可解的对数似然评估。然而,固定生成顺序在许多非序列任务中并非最优选择,促使研究界发展出任意阶自回归模型(AO-ARM)与离散掩码扩散模型(MDM)。这类模型将生成过程视为对潜在生成顺序或分组去掩码轨迹的概率平均,赋予了模型高度灵活的非单调生成与并行推理能力;而最新的块扩散模型(BM)则通过跨块自回归、块内扩散的方式进一步结合了 KV 缓存与局部并行生成的优势,在生成质量和下游指标上展现出与传统自回归模型相当的竞争力。
然而,对生成模型进行严格科学评估的关键指标是对数似然,而 AO-ARM 与 MDM 由于潜在顺序或去掩码路径的排列组合空间极其庞大,其边缘对数似然 \(\log p(x) = \log \mathbb{E}_\pi [p(x|\pi)]\) 在计算上完全不可行。以往文献普遍依赖证据下界(ELBO)或多样本紧化形式(\(\text{ELBO}_K\))作为替代指标,但下界只能衡量性能的保底水平,无法说明模型的真实对数似然究竟高于下界多少,导致学术界始终无法在对数似然这一根本度量上对扩散语言模型与自回归基线做出严谨、可靠的公允对比。
直接引入变分推断领域既有的对数似然上界估计量(如基于 Rényi 散度的 CUBO、热力学变分目标 TVO 以及重要性采样间隙界 IS-VG-B)也无法解决这一困境,因为这些方法均在蒙特卡洛均值外部复合了对数等非线性函数,在有限采样预算下受 Jensen 不等式或自归一化权重崩塌的影响而产生不可控的负向或正向偏差,破坏了数学上成立的上界性质。核心 idea:利用对数函数的凹性构造切线一阶 Taylor 主超平面,建立关于模型边缘分布线性的切线证据上界(TUBE),配合 tractable 代理分布实现有限样本下的严格无偏蒙特卡洛估计,从而与 ELBO 构成紧密的对数似然双向夹逼区间。
方法详解¶
整体框架¶
本文旨在为具备潜在生成顺序的离散潜变量语言模型(包括 AO-ARM、MDM 以及块扩散模型 BM)构建严格保界、有限样本无偏的对数似然上界估计。输入为离散文本序列,由于长序列联合概率尺度极小会导致全局代理函数失效,方法首先基于块扩散结构将全局序列分解为若干较短的块;在每个局部块内,分别从潜在顺序先验中独立抽取两组生成轨迹,一组用于构建自代理分布,另一组用于计算目标模型的无偏似然均值;最后通过对数切线不等式线性组合两部分,输出无偏的 \(\text{TUBE}\) 上界估计值,并与多样本下界 \(\text{ELBO}_K\) 联合输出真似然的紧致置信区间。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入离散文本序列 x"] --> B["块级序列分解与似然双向夹逼<br/>划分长度 L' 块,降低联合概率估计方差"]
B --> C["代理函数构建与自代理采样<br/>独立顺序多样本平均,紧贴真实模型密度"]
C --> D["切线证据上界与无偏蒙特卡洛估计<br/>对数切线一阶展开,保证有限样本无偏上界"]
D --> E["输出对数似然与困惑度双向夹逼区间<br/>ELBO_K ≤ log p(x) ≤ TUBE"]
关键设计¶
1. 块级序列分解与似然双向夹逼:解决长序列全局概率极小导致的数值与估计失效 对于长度为 \(L \sim 10^2 - 10^3\) 的完整自然语言文本,任意序列的边缘概率 \(p_{\text{model}}(x)\) 本质上是数百个词元概率的连乘,其绝对数值在计算机浮点表示中极度接近于零。这使得全局代理函数 \(\psi(x)\) 极难在绝对尺度上贴近真实边缘似然,一旦出现微小比值偏差便会导致上界严重松弛。为了攻克这一数值挑战,方法依托块生成模型(Block Models)的因果解耦结构,将全文划分为 \(B\) 个大小为 \(L'\) 的连续块 \(\mathcal{B}_1, \dots, \mathcal{B}_B\):
在此分解下,跨块之间保留严格的自回归因果结构,而块内条件分布 \(p(x^{\mathcal{B}_b} \mid x^{\mathcal{B}_{<b}})\) 则采用单词元逐步生成的 AO-ARM 潜顺序或多词元并行生成的 MDM 分组轨迹。每个局部块的维度大幅收缩至 \(L' \in \{4, 8, 16\}\),不仅使得 \(L'=4\) 时全部 \(4!=24\) 种全排列可以被解析枚举从而获得绝对准确的块似然真值,而且使得 \(L' \in \{8, 16\}\) 时局部概率的方差完全受控,为后续构造紧致的无偏切线上界提供了稳定的数值载体。
2. 代理函数构建与自代理采样:以无偏自蒸馏平均贴近真实边际似然 切线上界的紧致程度完全取决于变分代理密度 \(\psi(x)\) 与真实边缘分布 \(p_{\text{model}}(x)\) 的贴近程度(当 \(\psi(x) = p_{\text{model}}(x)\) 时切线恰好相切,上界达到紧致极值)。为此,论文探索了两类代理构造路径:第一类是外部自回归代理 \(\psi_{\text{ARM}}(x)\),利用同数据集训练的自回归模型提供快速解析概率,或用扩散模型生成的合成样本对其微调;第二类则是自代理(Self-Surrogate)构造。理论上,在模型最优点处,任意固定顺序下的条件似然 \(p_{\text{model}}(x \mid \pi)\) 都会逼近真实边缘似然 \(p_{\text{model}}(x)\)。自代理机制进一步采集 \(M\) 个独立的生成轨迹 \(\hat{\pi}^{(1)}, \dots, \hat{\pi}^{(M)}\) 并计算其算术均值:
为了彻底避免统计估计量之间的虚假关联,采样算法实施了严格的样本解耦隔离:用于构建自代理 \(\psi_M(x)\) 的 \(M\) 条潜顺序必须与后续用于评估似然估计量 \(\hat{p}_K(x)\) 的 \(K\) 条潜顺序相互独立。实验证实,自代理 \(\psi_M(x)\) 在点级上对扩散模型真实边际概率的逼近能力远超任何外部训练的自回归模型,随着采样数 \(M\) 的增加,上界与真实对数似然之间的间隙迅速收敛。
3. 切线证据上界与无偏蒙特卡洛估计:利用凸函数切线性质突破非线性变换带来的估计偏差 针对 CUBO、TVO 等现有变分上界在有限样本下因外层复合对数导致下沉破界的根本缺陷,TUBE 从对数函数的内凹性出发,利用 \(\log a\) 在点 \(b > 0\) 处的切线仿射变换:
令 \(a = p_{\text{model}}(x)\) 为目标模型的未知边缘似然,\(b = \psi(x)\) 为正值的代理函数,直接推导出针对对数似然的切线证据上界:
由于未知的真实边际概率 \(p_{\text{model}}(x) = \mathbb{E}_\pi [p_{\text{model}}(x \mid \pi)]\) 在上式中以严格线性的形式呈现,当采用 \(K\) 个独立同分布潜变量样本构建经验均值 \(\hat{p}_{\text{model}, K}(x) = \frac{1}{K}\sum_{k=1}^K p_{\text{model}}(x \mid \pi^{(k)})\) 时,期望算子可以直接穿透分子:
这一线性结构赋予了 TUBE 无与伦比的理论优势:对于任意有限的样本量 \(K \ge 1\),其实验估计量的期望在数学上严格大于等于真实对数似然,永不破界;且其估计方差为 \(\frac{1}{K \psi(x)^2} \operatorname{Var}_\pi [p_{\text{model}}(x \mid \pi)]\),具备标准的 \(\mathcal{O}(1/K)\) 快速方差收敛特性。
损失函数 / 训练策略¶
本文的方法核心聚焦于无偏推断与评估阶段,并不改变底座模型的预训练损失。评估对象沿用标准块掩码离散扩散模型(BD3-LM),其在预训练期间最小化吸收离散扩散的加权变分交叉熵损失,等价于不同时间步下去掩码分类交叉熵的积分。在似然界定评估时,将样本预算均分为两份(例如 \(K=M=12, 32, 64\) 分别对应块大小 \(L'=4, 8, 16\)),分别输入到 \(\psi_M\) 与 \(\hat{p}_K\) 中计算 \(\widehat{\text{TUBE}}\),并利用对数困惑度转换公式 \(\text{PPL} = \exp(-\log p(x))\),将对数似然的上下界映射为困惑度的双向界定区间 \([\text{TUBE}, \text{ELBO}_K]\)。
实验关键数据¶
主实验¶
主实验在 OpenWebText(OWT)与 One Billion Word(LM1B)基准数据集上对预训练块自回归/扩散语言模型(BD3-LM)进行全面评估。困惑度(PPL)数值越低代表模型似然越优。在 \(L'=4\) 且处于单词元全序列模式(AO-ARM)时,通过遍历全部 \(4!=24\) 种排列可获得精确无偏的真实似然基准(标记为 \(\star\))。
| 块大小 \(L'\) | 生成模式 / 步数 NFE | CUBO (\(\beta=2\)) ↓ | TVO\(_U\) ↓ | IS-VG-B ↓ | 本文 TUBE (PPL下界) ↓ | 参考 ELBO\(_K\) (PPL上界) ↓ | 真实 ARM 基线 | 与 ARM 差距 (|TUBE-ARM|) | |---|---|---|---|---|---|---|---|---| | \(L'=4\) | NFE = 1 | 97.78 | 97.78 | 97.78 | 97.78 | 97.78 | 17.54 | 80.24 | | \(L'=4\) | NFE = 2 | 24.34 | 27.68 | 29.11 | \(\ge 21.55 \pm 0.03\) | \(\le 27.40\) | 17.54 | 4.01 | | \(L'=4\) | NFE = 4 | 19.78 | 21.89 | 22.52 | \(\ge 19.16 \pm 0.02\) | \(\le 21.63\) | 17.54 | 1.62 | | \(L'=4\) | AO-ARM (全序) | 17.50 | 18.55 | 18.76 | \(\ge 17.74 \pm 0.01\) | 18.46\(^\star\) (精确值) | 17.54 | 0.20 | | \(L'=8\) | NFE = 2 | 29.04 | 33.41 | 34.95 | \(\ge 24.32 \pm 0.06\) | \(\le 32.63\) | 17.54 | 6.78 | | \(L'=8\) | NFE = 4 | 20.90 | 23.54 | 24.15 | \(\ge 20.23 \pm 0.03\) | \(\le 23.25\) | 17.54 | 2.69 | | \(L'=8\) | NFE = 8 | 18.85 | 20.83 | 21.17 | \(\ge 19.15 \pm 0.02\) | \(\le 20.67\) | 17.54 | 1.61 | | \(L'=8\) | AO-ARM (全序) | 17.95 | 19.31 | 19.46 | \(\ge 18.67 \pm 0.01\) | \(\le 19.24\) | 17.54 | 1.13 | | \(L'=16\) | NFE = 4 | 22.89 | 25.21 | 25.96 | \(\ge 19.03 \pm 0.08\) | \(\le 24.75\) | 17.54 | 1.49 | | \(L'=16\) | NFE = 16 | 18.44 | 19.98 | 20.26 | \(\ge 18.22 \pm 0.03\) | \(\le 19.83\) | 17.54 | 0.68 | | \(L'=16\) | AO-ARM (全序) | 18.08 | 19.39 | 19.56 | \(\ge 18.47 \pm 0.01\) | \(\le 19.30\) | 17.54 | 0.93 |
消融实验¶
消融实验深入对比了不同代理函数 \(\psi(x)\) 选型对 TUBE 上界收敛紧致性的影响(以 OWT 上 \(L'=4\) 与 \(L'=8\) 的 AO-ARM 困惑度估计为例,对比自代理 \(\psi_M\)、单顺序代理 \(\psi_\pi\)、外部标准自回归代理 \(\psi_{\text{ARM}}\) 以及用生成样本微调的自回归代理 \(\psi_{\text{ARM-FT}}\))。
| 代理函数配置 | 样本数 \(M\) | 估计性质 | OWT (\(L'=4\)) PPL下界 | OWT (\(L'=8\)) PPL下界 | 结论说明 |
|---|---|---|---|---|---|
| 自代理 \(\psi_M\) (本文默认) | 12 / 32 | 无偏且紧致 | 17.74 \(\pm\) 0.01 | 18.67 \(\pm\) 0.01 | 最贴近 ELBO\(_K\) 真值,界定区间最窄 |
| 自代理 \(\psi_M\) (小样本) | 2 | 无偏但方差略大 | 17.48 \(\pm\) 0.04 | 18.15 \(\pm\) 0.06 | 样本量增加使上界单调向真实似然紧缩 |
| 外部自回归 \(\psi_{\text{ARM}}\) | 1 (解析直通) | 无偏但极其松弛 | 14.82 \(\pm\) 0.02 | 15.10 \(\pm\) 0.03 | 外部模型密度分布偏移,导致切线上界松弛下沉(PPL偏小) |
| 微调自回归 \(\psi_{\text{ARM-FT}}\) | 1 (解析直通) | 微调后略有改善 | 15.12 \(\pm\) 0.02 | 15.44 \(\pm\) 0.03 | 合成数据微调仅轻微弥合密度差异,仍显著逊于自代理 |
| 单顺序自代理 \(\psi_\pi\) | 1 | 极小采样基线 | 16.92 \(\pm\) 0.05 | 17.65 \(\pm\) 0.08 | 单一顺序方差过大,边界紧致度不及均值聚合 |
关键发现¶
- 离散扩散与任意阶自回归的似然严格落后于传统 ARM:在 \(L'=4\) 能够精确穷举 24 种排列获得真值的设置下,AO-ARM 真实 PPL 确定为 18.46,而标准 ARM 达到 17.54;当扩展到 \(L'=8\) 和 \(L'=16\) 时,TUBE 提供的 PPL 最优下界分别为 18.67 与 18.47,依然严格大于 17.54。这首次以确凿无偏的数学证据终结了争议:现有离散扩散模型在纯分布拟合(对数似然)维度上尚未超越单向自回归模型。
- 现有变分上界估计量普遍存在结构性失效:TVO\(_U\) 与 IS-VG-B 在所有 \(\text{NFE} \ge 2\) 的设置下,其给出的 PPL 下界估计均荒谬地高于 \(\text{ELBO}_K\)(在对数似然域中低于下界),发生了上界穿透失效;而 CUBO 虽然在特定 \(\beta=2\) 下数值偏小,但由于其将对数放在蒙特卡洛平均之外,受 Jensen 不等式作用本质上引入了负向偏差,无法作为有效数学上界。
- 并行步数(NFE)与似然边界单调正相关:随着去掩码步数从粗粒度的 1 步逐步提升至精细的 16 步,模型的 TUBE-ARM 似然间隙由数百迅速收窄至 0.68 左右,表明离散扩散模型的并行推理加速是以显著的边缘似然牺牲为代价的。
亮点与洞察¶
- 将非线性函数求期望转化为切线仿射求期望:巧妙利用凹函数的切线不等式将对数运算移至确定性代理项,让不可求积的隐变量边际似然以严格线性的形式呈现,是蒙特卡洛估计中规避 Jensen 偏差的极简而强大的范式。
- 首次实现离散扩散模型似然的双向无偏夹逼:以往工作只能依靠单侧的 ELBO 自说自话,本文建立的 TUBE 与 ELBO 共同构成了封闭的区间估计,为学术界提供了终结“扩散生成质量媲美 ARM 是否等于分布拟合超越 ARM”争议的裁决工具。
- 可复用的解耦自代理策略:在没有现成高质量外部参照模型时,通过将模型自身的多个随机轨迹进行无偏独立拆分求平均,能够以极低实现成本构造出高质量变分代理,该思路可直接推广至连续扩散模型、离散 VAE 及分子序列生成等更广阔的潜变量模型中。
局限与展望¶
- 受限于块模型(Block Models)的局部假设:当前方法依赖将长文本截断为短块以保证数值稳定性,若直接应用于完全无因果解耦的千词元全局离散扩散模型,极小绝对似然将导致自代理估计出现严重的浮点下溢和方差爆炸。
- 计算开销随采样预算线性增长:获得足够紧密的双向夹逼区间需要分别评估 \(K\) 条与 \(M\) 条前向去掩码路径,评估一次验证集的计算量相当于标准自回归模型的数十倍,无法实时在线监测。
- 未来方向:探索更高效的自适应重要性采样顺序分布,以进一步降低 \(K\) 与 \(M\) 的需求;将 TUBE 的无偏边缘似然估计作为有限样本重要性权重,引入强化学习后训练(如 RLHF / DPO)中的策略比率(Policy Ratio)校正。
相关工作与启发¶
- vs CUBO & Rényi Variational Bound:CUBO 基于 Rényi 散度构建理论上界,但其实际估计量需要对蒙特卡洛均值取外层对数,Jensen 不等式直接消解了有限样本下的上界保证;TUBE 则保持对数在外、均值在内,理论与实验上均百分之百保界。
- vs TVO (Thermodynamic Variational Objective):TVO 利用热力学积分右黎曼和构建上界,但依赖自归一化重要性权重,在高 \(\beta\) 处权重极度塌缩带来 \(\mathcal{O}(1/K)\) 固有偏差与离散化残差;TUBE 无需连续路径积分,完全避免了权重崩塌。
- vs DUEL:DUEL 仅能在确定性去掩码调度下计算精确似然,限制了采样器的自由度;TUBE 兼容任意随机去掩码轨迹与分组生成策略。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 巧妙利用对数函数切线构造无偏上界,理论极其优雅利落。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖不同块大小、生成步数、全排列穷举基准及四种上界估计量对比,论证扎实无可辩驳。
- 写作质量: ⭐⭐⭐⭐⭐ 数学证明严密清晰,动机与实验洞察逻辑严丝合缝。
- 价值: ⭐⭐⭐⭐⭐ 澄清了离散扩散语言模型似然评估的核心争议,为该领域的评估体系树立了新基准。