Conservation Laws for Diffusion Models¶

会议: NeurIPS 2026
arXiv: 2607.10067
代码: https://github.com/zivaharoni/conservation-laws-diffusion-models
领域: 优化/理论
关键词: 扩散模型, 守恒律, 广义外信息传递 (GEXIT), 失配交叉熵, 有限容量效应

一句话总结¶

基于编码理论中的广义外信息传递（GEXIT）函数建立统一的不匹配守恒律，证明离散与连续无记忆信道下扩散模型的数据-模型交叉熵可精确表示为沿噪声路径对局部后验边缘导数的积分，将训练简化为基于信道倾斜的外先验分类负对数似然学习，并在马尔可夫信源、text8 和 CIFAR-10 上揭示了模型有限容量引起的显著信道依赖差异。

研究背景与动机¶

自回归（AR）模型通过概率链式法则逐坐标分解数据分布，直接最小化精确的负对数似然（即数据与模型的交叉熵），但其串行生成机制在长序列展开时容易放大分布偏移误差。相比之下，扩散模型沿噪声退化轴定义前向腐蚀过程并用学得的去噪器逆向还原，在图像生成上极为成功，也日益受到掩码扩散语言模型、离散状态空间扩散（D3PM）和连续嵌入空间文本扩散的青睐。然而，扩散模型在训练时通常只优化局部去噪目标或变分下界：连续高斯扩散常借助加权均方误差与变分界的等价性，而离散扩散中去噪损失与精确序列似然之间的数学联系往往是间接的。

近期已有研究尝试利用沿加噪路径的信息论导数来刻画扩散似然，例如高斯信道下的互信息-最小均方误差（I-MMSE）恒等式，以及针对掩码信道或泊松信道的专门似然公式。但这些分析各自局限在特定噪声族内，尚缺乏一个能够同时统摄掩码、均匀替换、连续高斯等广泛无记忆退化信道的单一理论框架；同时，已有信息论恒等式多假设去噪器能完美恢复真实后验（匹配设定），而实际神经网络学到的是有偏近似后验（失配设定），此时真正需要精确刻画的量是数据分布与模型分布之间的交叉熵（CE），且尚不清楚在失配设定下是否仍能仅凭单坐标边缘后验算得全局序列似然。

本文从编码理论分析迭代译码器的广义外信息传递（GEXIT）面积定理出发，将变量轴上的概率链式法则推广为噪声退化轴上的信息耗散积分。核心 idea：针对任意无记忆退化信道族推导失配 GEXIT 守恒律，将数据-模型交叉熵精确表示为局部信道得分加权边缘后验导数沿噪声路径的积分，从而把各类离散与连续扩散训练统一简化为结合已知信道似然倾斜的 token 级分类交叉熵优化，并从理论上将路径无关的精确守恒律与有限容量去噪器对不同信道族的近似偏差解耦。

方法详解¶

整体框架¶

本文提出的统一理论与实现框架将扩散模型的似然评估与后验训练建立在无记忆退化信道族 \(\mathcal{W}=\{W_t\}_{t\in\mathcal{T}}\) 之上。给定来自有限字母表 \(\mathcal{X}^n\) 的干净序列 \(X=(X_1,\dots,X_n)\)，前向过程沿连续可微路径 \(\gamma(\tau)\)（\(\tau\in(0,1)\) 连接干净数据端点与纯噪声无信息端点）独立腐蚀各坐标得到观测 \(Y(\mathbf{t})\)。在建模端，神经网络无需从零拟合完整联合后验，而是利用外信息分解仅输出各位置的外先验（extrinsic prior）logits，再结合已知的前向信道核做解析贝叶斯倾斜得到边缘后验 \(P^\theta_{X_i\mid Y}\)。训练阶段在所有容许信道上共享统一的 token 级交叉熵目标；评估阶段则将学得的边缘后验代入各信道特化的失配 GEXIT 导数公式，沿噪声路径数值积分守恒律曲线下面积，从而直接获得序列级的数据-模型交叉熵（BPC / bpd），并可进一步利用该曲线构造等面积（等立信息）逆向采样调度。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["干净离散序列 X 与无记忆退化信道路径 γ(τ)<br/>生成含噪观测 Y(t)"] --> B["信道感知外先验倾斜建模<br/>网络输出外先验 logits + 已知信道核解析倾斜"]
    B --> C["跨信道共享边缘交叉熵训练<br/>高斯信道在离散支撑集上取期望避免嵌入坍缩"]
    C --> D["失配 GEXIT 守恒律与局部性<br/>用边缘后验与信道得分计算局部信息导数"]
    D --> E["离散与连续信道特化面积积分<br/>掩码 / 均匀替换 / 高斯 SNR 积分输出精确似然与等面积采样"]

关键设计¶

1. 信道感知外先验倾斜建模：利用已知腐蚀核将去噪分解为外先验预测与解析贝叶斯融合

在任意无记忆容许信道中，单坐标前向转移核 \(W_{t_i}(y_i\mid x_i)\) 是完全已知的设计先验，如果让神经网络直接回归完整边缘后验 \(P^\theta_{X_i\mid Y}(x_i\mid y)\)，不仅浪费模型容量去重复记忆已知的局部信道似然，还容易破坏贝叶斯因子分解的一致性。为此，根据贝叶斯公式，第 \(i\) 个坐标在给定全部含噪观测 \(y\) 下的真实边缘后验严格正比于仅依赖其余位置观测 \(y_{-i}\) 的外先验（extrinsic prior）与局部信道似然之积。神经网络只需参数化外先验项并输出类别 logits \(\ell_i^\theta(x_i;y)\)，然后显式加上已知信道对数似然做倾斜（tilting）：

\[P^\theta_{X_i\mid Y}(x_i\mid y) = \mathrm{Softmax}_{x_i\in\mathcal{X}}\!\left(\ell_i^\theta(x_i; y) + \log W_{t_i}(y_i\mid x_i)\right)\]

在理想情况下外先验只依赖 \(y_{-i}\)（例如在合成马尔可夫实验中通过自注意力掩码排除当前位置自身输入）；在标准双向 Transformer 或 U-Net 实现中，为了在单次前向传播中并行预测所有位置，网络以完整 \(y\) 为输入，但显式的 \(\log W_{t_i}(y_i\mid x_i)\) 倾斜项会引导网络专注学习来自上下文 \(y_{-i}\) 的外部先验信息。

2. 跨信道共享边缘交叉熵训练：以分类后验均值替代连续嵌入回归以消除嵌入坍缩

由于后续守恒律仅依赖单坐标边缘后验，无论底层噪声是掩码、均匀替换还是连续高斯，训练目标都可以统一归结为沿加噪路径 \(\gamma(\tau)\) 最小化各位置边缘后验的负对数似然（即共享的 token 级交叉熵损失 \(\mathcal{L}_{\mathrm{CE}}(\theta)\)）。这一设计在将高斯扩散应用于离散数据（通过单射嵌入 \(f:\mathcal{X}\to\mathbb{R}^d\) 映射到连续空间后加噪 \(Y_i(\rho)=f(X_i)+\rho^{-1/2}Z_i\)）时尤为关键：传统连续文本扩散通常训练回归器直接最小化嵌入空间均方误差（MSE），当嵌入映射与去噪器同时学习时，极易发生把所有嵌入向量压向零点的嵌入坍缩（embedding collapse）。本文依靠外先验倾斜机制，让高斯模型依然在离散字母表 \(\mathcal{X}\) 上输出分类分布并结合高斯核似然 \(\log W_\rho(y_i\mid x_i)\) 做分类交叉熵训练，仅在评估 I-MMSE 曲线时通过该分类后验在离散支撑集上求期望来计算嵌入空间后验均值 \(m_{Q,i}(y)=\mathbb{E}_{X'\sim Q_{X_i\mid Y(\rho)}}[f(X')]\)，从而彻底避免了目标嵌入与预测均值相互追逐导致的退化坍缩。

3. 失配 GEXIT 守恒律与局部性：仅用单坐标边缘后验精确刻画全局序列交叉熵

现有信息论扩散分析大多局限于模型后验等于真实后验 \(P_{X\mid Y}\) 的匹配假设，或者需要难以处理的联合分布 \(Q_{X\mid Y}\)。本文利用无记忆信道下 \(Y_i\) 给定 \(X_i\) 条件独立于 \((X_{-i}, Y_{-i})\) 的马尔可夫性质，证明了联合条件交叉熵对局部噪声参数 \(t_i\) 的偏导数严格等于第 \(i\) 个坐标边缘交叉熵的偏导数（GEXIT 局部性）：\(\frac{\partial}{\partial t_i}\mathrm{CE}(P_{X\mid Y(\mathbf{t})}, Q_{X\mid Y(\mathbf{t})}) = \frac{\partial}{\partial t_i}\mathrm{CE}(P_{X_i\mid Y(\mathbf{t})}, Q_{X_i\mid Y(\mathbf{t})})\)。进一步，定义信道得分函数 \(S_{t_i}(x_i,y_i)=\frac{\partial}{\partial t_i}\log W_{t_i}(y_i\mid x_i)\)，结合外先验与 \(t_i\) 无关的性质，可推导出失配 GEXIT 导数定理（Theorem 3.1）：

\[\frac{\partial}{\partial t_i}\mathrm{CE}\!\left(P_{X_i\mid Y}, Q_{X_i\mid Y}\right) = \mathbb{E}\!\left[-S_{t_i}(X_i,Y_i)\log Q_{X_i\mid Y}(X_i\mid Y) + \mathbb{E}_{X'\sim Q_{X_i\mid Y}(\cdot\mid Y)}\!\left[S_{t_i}(X',Y_i)\mid Y\right]\right]\]

上式第一项是信道得分加权的对数似然项，第二项是后验失配校正项（当 \(Q=P\) 时校正项期望为零，自然退化为经典匹配 GEXIT 公式）。沿连接干净端点（\(\tau\downarrow 0\) 时条件交叉熵为 \(0\)）到无信息纯噪声端点（\(\tau\uparrow 1\) 时恢复无条件交叉熵 \(\mathrm{CE}(P_X, Q_X)\)）的路径 \(\gamma(\tau)\) 对上述局部导数求和并积分，即得到全局失配 GEXIT 守恒律。这意味着在无限步连续时间极限下，无需计算高维联合后验，仅凭单坐标边缘后验与已知信道得分就能精确重构完整序列的交叉熵。

4. 离散与连续信道特化面积积分：统一掩码、均匀替换与高斯 I-MMSE 似然评估

将通用失配 GEXIT 导数代入常见扩散信道，可导出极其简洁的闭合评估曲线，把不同扩散范式置于同一信息论标尺下：对于擦除率为 \(t_i\in[0,1]\) 的掩码扩散（Masked diffusion），失配校正项恰好相互抵消为零，局部交叉熵导数简化为该位置被掩码与被观测时的局部交叉熵之差 \(\mathrm{CE}_i^{\mathrm{mask}}(\mathbf{t}) - \mathrm{CE}_i^{\mathrm{obs}}(\mathbf{t})\)。对于退化进度为 \(s_i\in[0,1]\)、保持概率 \(a_i=1-s_i(1-1/|\mathcal{X}|)\) 的均匀替换信道（Uniform channel），局部导数由错误观测交叉熵 \(\mathrm{CE}_i^{\mathrm{err}}\)、正确观测交叉熵 \(\mathrm{CE}_i^{\mathrm{cor}}\) 以及模型对当前观测符号赋予的后验概率期望 \(\mathbb{E}[Q_{X_i\mid Y}(Y_i\mid Y)]\) 共同决定。对于以信噪比 \(\rho>0\) 参数化的高斯扩散（Gaussian diffusion），利用高斯向量 Stein 恒等式（分部积分）将信道得分转化为对观测的后验梯度与散度，对数似然项与后验失配校正项中的二阶矩精确抵消并配成完全平方，从而证明沿原生 SNR 轴的交叉熵面积恰好等于嵌入空间失配均方误差积分的一半（即广义失配 I-MMSE 公式）：

\[\mathrm{CE}(P_X, Q_X) = \frac{1}{2}\int_0^\infty \mathbb{E}_P\!\left[\sum_{i=1}^n \left\|f(X_i) - m_{Q,i}(Y(\rho))\right\|^2\right] d\rho\]

此外，有了验证集上的守恒律导数曲线，还可以直接按累积面积等分噪声区间来构造等立信息调度（equal-information schedule），在信息耗散剧烈的噪声区自动加密采样步、在平坦区放大步长，使逆向扩散每一步承担相等的似然积分贡献。

损失函数 / 训练策略¶

所有信道模型均通过最小化沿加噪路径 \(\gamma(\tau)\) 的共享边缘交叉熵目标进行训练：

\[\mathcal{L}_{\mathrm{CE}}(\theta) = \int_0^1 \mathbb{E}_{X, Y(\gamma(\tau))}\!\left[-\frac{1}{n}\sum_{i=1}^n \log P^\theta_{X_i\mid Y(\gamma(\tau))}\!\left(X_i\mid Y(\gamma(\tau))\right)\right] d\tau\]

训练时随机采样路径参数 \(\tau\)（掩码信道在 \([0,1]\) 均匀采样，均匀替换信道在 \([1/256, 0.9999]\) 采样，高斯信道在预设 SNR 区间如 text8 的 \([10^{-6},1]\) 或 CIFAR-10 的 \([10^{-6},10]\) 上对数均匀采样），并通过式 (14) 的信道倾斜后验计算交叉熵损失。守恒律积分不作为单独的训练损失，而是作为训练集/验证集/测试集上的精确似然尺度评估代理（在测试集上取 100 个噪声网格点数值积分求面积）。

实验关键数据¶

主实验¶

表 1：字符级语言建模 text8 上的扩散模型似然对比（指标为每字符比特数 BPC \(\downarrow\)，越低越好）

模型 / 噪声信道	是否使用外部数据	BPC (\(\downarrow\))	参数量与骨干配置
Masked (本文)	否	1.424	128.6M (12层 DiT, 隐藏维768)
Uniform (本文)	否	1.513	128.6M (12层 DiT, 隐藏维768)
Gaussian (本文)	否	2.087	129.2M (12层 DiT, 隐藏维768)
D3PM mask/absorbing	否	\(\le 1.45\pm 0.02\)	标准离散扩散基线
SEDD Absorb	否	\(\le 1.39\)	得分熵离散扩散
MDLM	是	\(\le 1.40\)	掩码扩散语言模型
MD4	是	\(\le 1.37\)	广义掩码离散扩散
EDLM	是	\(\le 1.24\)	基于能量的扩散语言模型

表 2：CIFAR-10 图像生成质量（FID \(\downarrow\) / IS \(\uparrow\)，1000 步采样）与似然评估（bpd \(\downarrow\)）对比

模型 / 噪声信道	采样调度 (Schedule)	FID (\(\downarrow\))	IS (\(\uparrow\))	似然 bpd (\(\downarrow\))
Uniform (本文)	cosine	59.51	\(6.21\pm 0.06\)	3.66
Masked (本文)	equal-info	30.69	\(6.87\pm 0.08\)	3.31
Gaussian (本文)	log-uniform	26.66	\(7.34\pm 0.08\)	3.24
D3PM uniform	cosine	\(51.27\pm 2.15\)	\(5.99\pm 0.14\)	-
D3PM absorbing	mutual-info	\(41.28\pm 0.65\)	\(6.26\pm 0.10\)	-
D3PM Gauss + logistic	linear	\(7.34\pm 0.19\)	\(8.56\pm 0.10\)	\(\le 3.435\pm 0.007\)
DDPM	linear-\(\beta\)	3.17	9.46	\(\le 3.70\)
Improved DDPM	-	-	-	\(\le 3.40\)
Score SDE / ML Score SDE	geometric-\(\sigma\)	2.20	9.89	\(\le 2.99\)

消融实验¶

本文的核心对照与消融围绕不同信道族在有限模型容量下的近似误差以及字母表规模对守恒律估计的影响展开（表 3 汇总自合成一阶马尔可夫信源实验、BCJR 精确后验对比与跨模态信道消融）：

实验维度 / 配置	评估设定	关键表现与观测结果	说明
因果自回归 Transformer	合成马尔可夫信源 (\(K\in[8,512]\))	紧贴真实熵率 \(H(X)\)	变量轴链式法则参考基准
Masked GEXIT 估计	合成马尔可夫信源 (\(K\in[8,512]\))	全字母表范围逼近真实熵率與 BCJR 精确曲线	离散擦除后验最易被有限容量网络拟合
Uniform GEXIT 估计	合成马尔可夫信源 (\(K\in[8,512]\))	全字母表范围紧贴真实熵率與 BCJR 精确曲线	离散均匀替换在小到中等词表下保持高准度
Gaussian GEXIT 估计	合成马尔可夫信源 (\(K\in[8,512]\))	始终高于真实熵率，且偏差随 \(K\) 增大而显著扩大	连续嵌入空间去噪在高维离散词表下存在显著有限容量损耗
信息不确定性重参数化对齐	text8 统一横轴 \(H(X\mid Y)\)	\(t=0.1\) 掩码对应 \(0.555\) nats；\(s=0.1\) 均匀对应 \(0.879\) nats；\(\rho=0.1\) 高斯对应 \(\approx 0.009\) nats	将不同原生噪声轴映射到统一条件熵标尺以对比局部耗散

关键发现¶

理论路径不变性与实际有限容量偏差的鲜明对比：在合成一阶马尔可夫信源（字母表大小 \(K\) 从 8 扩展到 512，利用 BCJR 前向-后向算法计算精确参考曲线）上，尽管理论守恒律保证任何无记忆信道下的精确后验积分面积都严格等于真实熵率，但固定容量（8 层、隐藏维 256）的去噪 Transformer 对不同信道诱导的后验族逼近能力截然不同：掩码和均匀信道的 GEXIT 面积估计紧贴真实熵率，而高斯信道的估计值明显偏高（即存在额外的失配 KL 散度），且误差随字母表大小 \(K\) 增加而单调恶化。
没有单一噪声信道能跨模态通吃：在相同的架构与共享交叉熵训练协议下，最适合的噪声信道高度依赖数据模态的内在几何结构。在离散符号属性极强的字符级语言建模（text8）上，掩码扩散取得了最优的 \(1.424\) BPC（优于同设置下均匀信道的 \(1.513\) BPC 和高斯信道的 \(2.087\) BPC，且超越了不使用外部数据的基线 D3PM absorbing 的 \(1.45\) BPC）；而在具有天然连续光度序关系的图像数据（CIFAR-10，表示为 3072 个 8-bit 子像素 token）上，高斯信道反超离散信道，取得了三者中最佳的似然（\(3.24\) bpd vs. 掩码 \(3.31\) bpd、均匀 \(3.66\) bpd）和最佳生成质量（FID \(26.66\) vs. 掩码 \(30.69\)、均匀 \(59.51\)）。
守恒律曲线可直接指导等立信息采样：通过计算 CIFAR-10 验证集上的 GEXIT 守恒律曲线并按等累积面积划分采样点（equal-information schedule），掩码扩散在 1000 步采样下达到了 \(30.69\) FID 和 \(6.87\) IS，显著优于原版 D3PM absorbing 结合互信息调度的 \(41.28\) FID 和 \(6.26\) IS。

亮点与洞察¶

把扩散模型诠释为“噪声轴上的概率链式法则”：自回归模型无论按什么变量顺序展开条件概率，总熵都守恒；本文漂亮地证明了扩散模型无论沿什么无记忆退化信道耗散信息，沿噪声路径的局部信息导数积分也同样守恒。这一视角将编码理论中的 GEXIT 面积定理引入生成建模，在一套数学语言下同时推导出了掩码差值公式、均匀替换公式和高斯失配 I-MMSE 公式。
最让人“啊哈”的局部性推论（Infinitesimal Locality）：序列生成模型最核心的困难在于联合分布的高维耦合，但本文证明在连续退化极限下，评估完整序列的交叉熵根本不需要联合后验 \(Q_{X\mid Y}\)，单坐标边缘后验 \(Q_{X_i\mid Y}\) 就包含了重构全局似然所需的全部一阶导数信息。这为“为什么扩散模型只用逐 token 独立分类交叉熵训练就能学好全局序列分布”提供了迄今最透彻的信息论解释。
外先验信道倾斜（Logits Tilting）与抗坍缩技巧具备极高可迁移性：式 (14) 将网络职责剥离为纯粹预测上下文外先验 \(\ell_i^\theta(x_i;y)\)，再直接把已知前向转移对数概率 \(\log W_{t_i}(y_i\mid x_i)\) 加到 logits 上；同时在高斯扩散中只用分类分布对嵌入求期望而不直接回归连续向量。这两个轻量级技巧完全不挑骨干网络，可直接迁移到离散蛋白质序列生成、代码扩散模型或音频 token 扩散中，既能免除模型学习已知噪声核的负担，又能根治连续嵌入扩散的坍缩顽疾。

局限与展望¶

无穷小局部性与有限步采样之间的鸿沟：作者明确指出，边缘后验的充分性（locality）仅在无穷小退化步长（无限步连续极限）下严格成立；在实际有限步（如几十到几百步）逆向采样中，每一步转移跨越较大的噪声区间，此时各坐标之间的联合后验相关性不可忽略，仅靠边缘后验无法唯一确定最优的有限步联合采样器。此外，学习到的边缘后验集合在实现中并未强制要求可积（即未约束其严格对应某个全局自洽的联合分布 \(Q_X\)）。
大词表下高斯离散支撑集期望的计算瓶颈与实验规模限制：高斯信道实现需要在整个离散字母表 \(\mathcal{X}\) 上对嵌入向量求分类后验期望 \(m_{Q,i}(y)\)，这在字节级（\(|\mathcal{X}|=256\)）或小词表下开销很小，但如果直接扩展到数万词表（如 BPE 32k/128k）的大语言模型，每步对全词表嵌入求和并计算 Stein 散度将带来显著显存与计算开销。实验层面目前仅覆盖了合成马尔可夫链、text8 和 CIFAR-10（且大型实验仅使用单随机种子 2026），在 CIFAR-10 上统一子像素离散化架构的 FID（\(26.66\)）与专门针对连续像素优化的 DDPM/Score SDE（FID \(2.20\sim 3.17\)）仍有较大差距。
未来可改进方向：利用本文覆盖非均匀噪声路径 \(\mathbf{t}=(t_1,\dots,t_n)\) 的通用公式（式 (5) 与式 (9)），未来可以不再局限于所有位置共享单一标量噪声调度 \(\gamma_i(\tau)=\tau\)，而是根据不同空间位置或不同语义 token 的局部 GEXIT 导数陡峭程度，自适应优化按位置异构的加噪/去噪路径；针对大词表高斯扩散，可引入稀疏 Top-\(k\) 后验截断或分层 softmax 来降低条件均值计算成本。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 将编码理论中的 GEXIT 面积定理推广到失配交叉熵设定，首次在统一微分/积分框架下打通了掩码、均匀替换和连续高斯扩散的精确似然表征。
实验充分度: ⭐⭐⭐⭐ 在带解析真值的一阶马尔可夫信源（BCJR 基准）、字符级语言模型 text8 和图像 CIFAR-10 上系统验证了有限容量效应，但受算力限制未扩展到大词表 LLM 与高分辨率图像。
写作质量: ⭐⭐⭐⭐⭐ 数学推导极其严谨流畅，从匹配熵裂项一路自然过渡到失配 GEXIT 定理与三大信道特例，附录中 Stein 恒等式与 Verdú 公式的互证尤为清晰。
价值: ⭐⭐⭐⭐⭐ 不仅为离散与连续扩散提供了统一的似然评估标尺和无坍缩分类训练范式，更把“前向噪声信道与调度路径的选择”确立为可量化优化的模型设计自由度。