跳转至

Bernoulli Flow Models: Self-Consistent Generative Modeling for Binary Data

会议: NeurIPS 2026
arXiv: 2610.11362
代码: https://github.com/minimo6666/BFM_NeurIPS2026
领域: 图像生成
关键词: 二值流模型, 离散生成建模, 二值潜空间图像生成, 免训练少步采样, 自一致性跨步后验

一句话总结

本文提出伯努利流模型(BFM),通过在二值参数空间预定义统一的连续全局伯努利概率流路径并推导任意时间跨度的闭式后验转移分布,彻底消除了离散扩散在少步跨步采样时的训练-推理结构失配,在无需蒸馏或重训的前提下将 LSUN Churches \(256\times 256\) 的 16 步采样 FID 从基线 BLD 的 204.10 大幅降至 9.22。

研究背景与动机

二值变量的生成建模在二值图像生成、高分辨率图像的二值潜空间表示(如 Binary Latent Diffusion, BLD)以及统计物理自旋系统(如 Ising 模型)中具有基础性地位。虽然近年来面向二值空间的扩散模型已经能够在百步乃至千步采样下取得优秀的生成质量,但高昂的函数评估次数(NFE)严重制约了其实际推理效率。在连续空间扩散模型中,DDIM 和 DPM-Solver 等免训练加速采样器可以直接在稀疏推理时间网格上重评估同一连续动力学,从而在不重新训练或蒸馏的情况下大幅削减采样步数;然而在二值离散状态空间中,这类免训练少步加速却难以直接奏效。

造成这一困境的核心矛盾在于:现有二值扩散模型(如 BLD、D3PM 等)大多建立在固定单步离散马尔可夫链之上,先逐级定义离散单步前向加噪转移,再据此推导局部单步逆向后验。当推理阶段为了加速而将采样网格变粗(例如将 256 步训练的模型压缩至 16 步推理)时,每一次逆向更新都必须跨越多个中间训练时间步。然而现有方法在跨步采样时往往沿用局部单步似然转移(仅使用单个端点的噪声率参数)去近似真实的多步跨区间似然转移,导致训练时的生成过程与少步推理时的跨步过程发生严重的结构性失配,在低 NFE 下生成图像迅速退化为模糊噪声。

本文的切入角度是跳出“逐级定义单步马尔可夫链”的传统范式,直接从连续时间的全局概率路径出发构建二值生成动力学。只要边缘分布和任意区间转移率都能写成连续时间坐标的解析函数,改变推理步数就不再是对跳过离散步的粗略近似,而仅仅是在新的时间网格上重新代入同一个解析后验公式。核心 idea:在伯努利参数空间中定义连接真实二值数据与均匀白噪声的连续全局伯努利概率流路径,并利用对称二值组合算子严格推导出任意时间间隔下的闭式跨步后验分布,从而实现天然具备跨步自一致性的免训练少步二值生成。

方法详解

整体框架

BFM 将二值生成建模表述为连续时间区间 \(t_{\tau}\in[0,1]\) 上的概率流逆转过程(其中 \(t_0=1\) 对应干净二值数据 \(X_{t_0}\in\{0,1\}^d\),\(t_T=0\) 对应均匀伯努利白噪声先验 \(X_{t_T}\sim\mathcal{B}(0.5)\))。对于高分辨率图像(如 \(256\times 256\) 的 LSUN 和 FFHQ),模型首先利用预训练的二值自编码器(BAE)将图像编码为 \(16\times 16\times 64\) 的紧凑二值潜张量;对于原生二值数据(如 Binarized MNIST 和二维 Ising 自旋晶格),则直接在原始二值空间上操作。在训练阶段,模型沿连续全局伯努利概率流路径随机采样时间步并一步生成加噪二值状态,送入 Transformer 骨干网络直接预测干净二值数据 \(X_{t_0}\);在推理阶段,从均匀伯努利白噪声出发,根据目标 NFE 预算划分任意粗细的时间网格,将网络预测的干净数据概率代入自一致性跨步解析后验公式中,逐步采样恢复出最终的二值样本并解码为图像。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入:干净二值数据或BAE二值潜表示<br/>与均匀伯努利先验 B(0.5)"] --> B["连续全局伯努利概率流路径<br/>参数空间线性插值与对称算子构造"]
    B --> C["干净数据预测参数化与误差有界性<br/>Transformer预测x0与温度缩放"]
    C --> D["自一致性跨步解析后验<br/>任意推理网格步长闭式贝叶斯采样"]
    D --> E["输出:生成二值样本<br/>或经BAE解码生成256×256图像"]

关键设计

1. 连续全局伯努利概率流路径:在参数空间直接建立数据与白噪声的最优传输插值

为了彻底摆脱对固定离散时间步的依赖,BFM 仿照连续流匹配中的最优传输线性插值,直接在伯努利分布的参数空间上定义从干净数据 \(X_0\sim\pi_0\)(时间 \(t_0=1\))到二值白噪声先验 \(X_1\sim\mathcal{B}(0.5)\)(时间 \(t_T=0\))的连续时间概率路径:\(X_{t_{\tau}}\sim\mathcal{B}(t_{\tau}X_0+(1-t_{\tau})X_1)\)。对均匀先验 \(X_1\sim\mathcal{B}(0.5)\) 求边缘期望后,任意时刻 \(t_{\tau}\in[0,1]\) 的边缘分布都可以通过作者定义的对称二值组合算子 \(g(a,b)=a(1-b)+b(1-a)\) 写成紧凑的解析形式:

\[ q(X_{t_{\tau}}\mid X_{t_0}) = \mathcal{B}(X_{t_{\tau}}; \alpha_{t_{\tau}}), \quad \alpha_{t_{\tau}} = g\left(X_{t_0}, \frac{1-t_{\tau}}{2}\right) = \frac{1+(2X_{t_0}-1)t_{\tau}}{2} \]

基于这一全局边缘分布和全概率公式,作者进一步证明了任意相邻时刻从 \(t_{\tau-1}\) 到 \(t_{\tau}\)(\(0\le t_{\tau} < t_{\tau-1}\le 1\))的前向马尔可夫转移概率为 \(q(X_{t_{\tau}}\mid X_{t_{\tau-1}})=\mathcal{B}(X_{t_{\tau}}; g(X_{t_{\tau-1}}, \gamma_{t_{\tau}}))\),在推导过程中数据依赖项 \((2X_{t_0}-1)\) 被精确约去,使得比特翻转率 \(\gamma_{t_{\tau}} = \frac{0.5(t_{\tau-1}-t_{\tau})}{t_{\tau-1}}\) 仅取决于两个端点的时间坐标。这种先定全局路径、再导出无数据依赖翻转率的构造,使前向过程在连续时间轴上任意两个时刻之间都具备严格的解析定义。

2. 自一致性跨步解析后验:通过任意时间跨度闭式解消除跳步采样近似误差

针对传统二值扩散模型在少步推理(低 NFE)时因用单步转移近似跨步转移而崩溃的痛点,BFM 利用上述连续时间翻转率性质,将单步贝叶斯后验直接推广至任意跨步区间 \(0\le t_{\tau} < t_{\tau-k}\le 1\)(\(k>1\))。由于任意跨度下的前向翻转率均可精确写为 \(\gamma_{t_{\tau}\mid t_{\tau-k}} = \frac{0.5(t_{\tau-k}-t_{\tau})}{t_{\tau-k}}\),结合对称算子 \(g(\cdot,\cdot)\) 对分母的折叠化简,真实跨步后验分布(Theorem 1)具有如下严格的闭式解:

\[ q(X_{t_{\tau-k}}\mid X_{t_{\tau}}, X_{t_0}) = \mathcal{B}\left(X_{t_{\tau-k}}; \frac{g(X_{t_{\tau}}, \gamma_{t_{\tau}\mid t_{\tau-k}})\alpha_{t_{\tau-k}}}{g\big(g(X_{t_{\tau}}, \gamma_{t_{\tau}\mid t_{\tau-k}}), 1-\alpha_{t_{\tau-k}}\big)}\right) \]

这一性质被称为跨步采样下的“自一致性(Self-Consistency)”:当推理预算从训练时的 \(T\) 步(如 256 步)大幅缩减为 \(K\) 步(如 16 步)时,模型无需任何重训练或分布蒸馏,只需将新的稀疏时间网格代入 \(\gamma_{t_{\tau}\mid t_{\tau-k}}\) 和 \(\alpha_{t_{\tau-k}}\) 重新计算后验参数即可。相比之下,BLD 官方实现在跨步采样时仅用单步噪声参数 \(\beta^{t-m}\) 代替多步累积乘积 \(\prod_{j=t-m+1}^t(1-\beta^j)\),在步长大时产生巨大偏差,而 BFM 的解析闭式解在任意网格粒度下都与原始全局流路径保持严格一致。

3. 干净数据预测参数化与误差有界性:以二值交叉熵直接约束逆向转移误差

在真实推理时干净数据 \(X_{t_0}\) 未知,若直接让神经网络回归整个后验概率,网络需要同时拟合随噪声状态和时间步长剧烈变化的复杂插值关系。BFM 观察到在闭式跨步后验中,\(X_{t_{\tau}}\) 与时间转移率 \(\gamma_{t_{\tau}\mid t_{\tau-k}}\) 均已知,唯一未知量只有 \(X_{t_0}\)。因此,BFM 采用结构化的 \(X_{t_0}\)-预测参数化:用 Transformer 网络 \(\mathcal{T}_{\theta}\) 配合温度参数 \(\kappa\) 输出干净数据为 1 的概率预测 \(f_{\theta}(X_{t_{\tau}}, t_{\tau}) = \sigma(\mathcal{T}_{\theta}(X_{t_{\tau}}, t_{\tau})/\kappa)\),并将其直接替换闭式后验中的 \(X_{t_0}\)(即用 \(g(f_{\theta}(X_{t_{\tau}}, t_{\tau}), \frac{1-t_{\tau-k}}{2})\) 替换 \(\alpha_{t_{\tau-k}}\))来构造参数化逆向转移 \(p_{\theta}(X_{t_{\tau-k}}\mid X_{t_{\tau}})\)。

为了从理论上支撑这一参数化,作者严格推导了学习到的逆向采样误差界(Proposition 1 & 2):对于任意逆向区间 \(0\le t < s < 1\),单步/跨步真实后验与模型后验之间的全变差(TV)距离满足:

\[ D_{\mathrm{TV}}\big(q(X_s\mid X_t, X_{t_0}), p_{\theta}(X_s\mid X_t)\big) \le \frac{s^2-t^2}{s(1-t)^2}\left|f_{\theta}(X_t, t) - X_{t_0}\right| \]

同时,两者的坐标级 KL 散度被证明受二值交叉熵(BCE)损失乘以区间常数 \(C_{s,t}\) 的上界控制,且 \(K\) 步低 NFE 采样的累积分布 TV 误差受各步软 \(\ell_1\) 预测误差 \(\varepsilon_j\) 的加权和控制。这从理论上打通了“最小化训练期 BCE 损失”与“最小化任意跨步网格下逆向生成分布误差”之间的因果链条。

损失函数 / 训练策略

训练目标采用标准的逐坐标二值交叉熵损失(BCE),直接监督网络对干净二值状态 \(X_{t_0}\) 的恢复能力:

\[ \mathcal{L}_{\mathrm{BCE}}(\theta) := -\mathbb{E}_{t_{\tau}, X_{t_0}, X_{t_{\tau}}}\Big[X_{t_0}\log f_{\theta}(X_{t_{\tau}}, t_{\tau}) + (1-X_{t_0})\log\big(1-f_{\theta}(X_{t_{\tau}}, t_{\tau})\big)\Big] \]

在训练策略与超参数配置上,BFM 保持与 BLD 完全一致的 24 层、12 注意力头、768 维嵌入的 Transformer 骨干网络;高分辨率图像任务(LSUN Bedrooms、LSUN Churches、FFHQ \(256\times 256\))复用 BLD 在 6 亿张 LAION 图像上预训练的通用 C64 二值自编码器(BAE),将图像映射为 \(16\times 16\times 64\)(序列长度 256)的二值潜空间表示,训练批大小为 96;Binarized MNIST(序列长度 28)和 \(20\times 20\) 二维 Ising 模型(序列长度 20)直接在原生二值空间上训练,批大小为 512。优化器统一采用 Adam(\(\beta_1=0.9, \beta_2=0.99\),无权重衰减),经过 10,000 步线性 warmup 达到峰值学习率 \(2\times 10^{-4}\);默认采样温度设为 \(\kappa=0.9\),概率流路径默认采用线性(Linear)调度策略。

实验关键数据

主实验

表 1 展示了在 \(256\times 256\) 高分辨率图像生成任务(LSUN Bedrooms、LSUN Churches、FFHQ)上,BFM 与连续生成模型及离散生成模型在标准步数(800K 训练迭代,64 步采样,50K 生成样本)下的全面对比(其中 BLD* 为相同软硬件与训练配置下的严格复现结果):

数据集 方法类型 方法 采样步数 (NFE) FID \(\downarrow\) Precision \(\uparrow\) Recall \(\uparrow\)
LSUN Bedrooms \(256\times 256\) 连续生成模型 LDM-4 200 2.95 0.66 0.48
LSUN Bedrooms \(256\times 256\) 连续生成模型 Patch-DM 50 6.04 0.56 0.44
LSUN Bedrooms \(256\times 256\) 离散生成模型 D3PM 200 6.60 0.60 0.35
LSUN Bedrooms \(256\times 256\) 离散生成模型 VQ-Diffusion 200 7.19 0.54 0.37
LSUN Bedrooms \(256\times 256\) 离散生成模型 BLD* (复现基线) 64 6.22 0.69 0.36
LSUN Bedrooms \(256\times 256\) 离散生成模型 BFM (本文) 64 5.86 0.69 0.37
LSUN Churches \(256\times 256\) 离散生成模型 D3PM 200 6.02 0.68 0.39
LSUN Churches \(256\times 256\) 离散生成模型 VQ-Diffusion 200 6.88 0.72 0.37
LSUN Churches \(256\times 256\) 离散生成模型 BLD* (复现基线) 64 5.48 0.66 0.41
LSUN Churches \(256\times 256\) 离散生成模型 BFM (本文) 64 5.32 0.66 0.41
FFHQ \(256\times 256\) 离散生成模型 D3PM 200 9.49 0.71 0.41
FFHQ \(256\times 256\) 离散生成模型 VQ-Diffusion 200 8.79 0.70 0.43
FFHQ \(256\times 256\) 离散生成模型 BLD* (复现基线) 64 11.46 0.68 0.48
FFHQ \(256\times 256\) 离散生成模型 BFM (本文) 64 10.87 0.68 0.49

此外,在原生二值流形 Binarized MNIST 生成任务上,BFM 取得了 0.42 的最佳 FID,显著优于离散流匹配方法 SFM(4.62)、\(\alpha\)-Flow(5.02)、基于分数的离散模型 DMPM(2.89)以及 CoVAE(0.58)。在 64 步单样本推理资源测试中,Image-256(含 BAE 解码)单张生成耗时仅 0.3837 秒、峰值显存 1.0743 GB;Binarized MNIST 和 Ising 模型耗时约 0.35 秒、峰值显存仅 0.65 GB。

消融实验

表 2 汇总了 LSUN Churches \(256\times 256\) 上不同伯努利概率流路径调度器(100K 训练迭代)的消融对比,以及同一 256 步训练模型在不同推理 NFE 下的跨步采样表现(含原始 BLD 与附录 D 中作者修复跨步转移公式后的 Corrected BLD 对比):

实验维度 配置 / 采样步数 BFM FID \(\downarrow\) 对比基线 FID \(\downarrow\) 说明
调度器消融 (100K iter) Linear (线性调度,默认) 9.22 9.55 (BLD) 线性概率路径取得最佳生成质量
调度器消融 (100K iter) Half-cosine (半余弦) 9.41 9.53 (BLD) 优于 BLD,但略逊于线性调度
调度器消融 (100K iter) Reverse half-cos (反半余弦) 9.27 9.54 (BLD) 三种调度下 BFM 均稳定优于 BLD
免训练少步采样 (100K iter) 256 步训练 \(\rightarrow\) 16 NFE 9.22 204.10 (原始 BLD) 原始 BLD 因单步近似崩溃,BFM 保持自一致
修正跨步公式对比 NFE = 16 10.41 10.49 (Corrected BLD) 相对降低 0.76% FID
修正跨步公式对比 NFE = 32 8.28 8.37 (Corrected BLD) 相对降低 1.08% FID
修正跨步公式对比 NFE = 64 7.78 7.98 (Corrected BLD) 相对降低 2.51% FID
修正跨步公式对比 NFE = 128 7.70 7.85 (Corrected BLD) 相对降低 1.91% FID
修正跨步公式对比 NFE = 256 7.59 7.79 (Corrected BLD) 相对降低 2.57% FID

关键发现

  • 跨步自一致性彻底解决了低 NFE 采样崩溃问题:当把 256 步训练的检查点直接用于 16 步采样时,原始 BLD 的 FID 暴跌至 204.10,生成图像完全失去结构;而 BFM 凭借任意区间的闭式跨步后验,在 16 步下依然保持 9.22 的低 FID。即使作者在附录 D 中主动帮 BLD 修复了多步噪声累乘公式(Corrected BLD),BFM 在 16 至 256 的所有 NFE 档位上依然一致取得 0.76%–2.57% 的 FID 相对降低。
  • 全局一致的概率路径带来了更快的训练收敛速度:在 LSUN Bedrooms 的 100K 到 400K 迭代训练曲线中,BFM 在完全相同的网络与优化器下全程保持比 BLD 更低的 FID,因为不同采样时刻的监督信号都源自同一条全局闭式路径,时间步之间的训练信号更加协调。
  • 跨领域物理规律保真度高:在附录 C 的 \(20\times 20\) 二维 Ising 自旋系统生成实验中,BFM 在 \(T\in[1.50, 3.10]\) 连续温度区间内生成的自旋构型不仅视觉上准确复现了从有序到无序的相变,其宏观绝对磁化强度曲线 \(M(T)\) 也与 Wolff 团簇蒙特卡洛算法的真实物理模拟结果高度吻合。

亮点与洞察

  • 对称二值组合算子 \(g(a,b)=a(1-b)+b(1-a)\) 的代数化简极为精妙:二值概率推导通常容易陷入繁琐的分情况讨论或复杂的分子分母展开,而本文引入对称算子 \(g(a,b)\) 后,不仅边缘分布和前向翻转分布可以统一写成嵌套形式,连贝叶斯后验的分母也能被精确折叠为 \(g(g(X_{t_{\tau}}, \gamma), 1-\alpha)\),且前向翻转率推导中自然约去了数据项 \((2X_{t_0}-1)\),这是整个闭式跨步框架得以成立的代数支点。
  • 最让人“啊哈”的洞察在于揪出了离散扩散少步崩溃的真凶:过去人们普遍认为离散扩散在低 NFE 下效果差是因为缺乏类似连续 ODE 的高阶数值求解器,而本文不仅提出 BFM 解决了该问题,还在附录 D 中通过严谨的理论对照指出,BLD 在低 NFE 下崩溃的根本原因是其跨步采样用单步参数 \(\beta^{t-m}\) 替代了多步累积项,这一诊断为整个二值生成社区厘清了离散马尔可夫链与连续概率流之间的本质联系。
  • 设计思路可无缝迁移至其他二值/比特级生成任务:BFM 对现有二值扩散代码库的改动极小(仅替换前向加噪公式和逆向后验采样公式,网络架构和 BCE 损失完全不变),因此可以直接作为即插即用的生成底座迁移到二值神经压缩、离散语义哈希、组合优化求解以及量子/统计物理二值态模拟等任务中。

局限与展望

  • 极少步(1–4 步)生成仍受限于神经网络预测误差:正如理论误差界 Proposition 1 和 2 所示,虽然给定真实 \(X_{t_0}\) 时的理想(Oracle)跨步后验是完全精确的,但实际推理使用网络预测值 \(f_{\theta}(X_t, t)\) 代替 \(X_{t_0}\),当在接近纯噪声的较大区间上做单步或两步极大跨越时,网络预测误差仍会被区间系数放大;未来可将 BFM 的解析流路径与一致性训练(Consistency Training)或渐进式蒸馏结合,以实现高质量单步生成。
  • 高分辨率图像生成受限于预训练二值自编码器(BAE)的天花板:在 LSUN 和 FFHQ 上,BFM 在 64 步下的 Recall(0.37–0.49)和 FID(5.32–10.87)虽然全面优于同配置的 BLD*,但与连续潜扩散模型 LDM(FID 2.95–4.98)相比仍有一定差距,这主要受限于冻结的 C64 二值自编码器在离散化压缩时的表征信息损失。
  • 各坐标条件独立假设:当前理论推导与采样均假设给定干净状态或当前状态时各二值比特位之间条件独立(乘积伯努利分布),在强空间耦合的极少步跳跃中可能忽略同一时间步内比特间的联合相关性,未来可探索结合自回归局部校正或耦合采样策略。

相关工作与启发

  • vs Binary Latent Diffusion (BLD, CVPR 2023):BLD 在二值潜空间中定义了固定步数的离散马尔可夫扩散链,其跨步采样存在近似失配,导致 16 步采样时 FID 退化至 204.10;本文复用与 BLD 相同的 BAE 与 Transformer 架构,但改用连续全局伯努利概率流路径与闭式跨步后验,不仅在 64 步标准推理下全面超越 BLD,更实现了 16 步下 FID 9.22 的免训练稳定加速。
  • vs Analog Bits (Chen et al., 2022):Analog Bits 将离散二值比特强行松弛为实数域上的连续模拟比特(Analog Bits)并套用连续高斯扩散,推理结束时再通过阈值截断恢复二值;本文则始终严格停留在原生伯努利二值概率空间 \(\mathcal{X}=\{0,1\}^d\) 内建模,避免了连续松弛带来的几何流形失配。
  • vs 通用离散流匹配 (Discrete Flow Matching / Categorical Flow Matching, NeurIPS 2024):通用离散流匹配面向一般类别型(Categorical)状态空间构建连续时间生成流或单纯形上的黎曼流,通常仍需借助欧拉离散化近似积分;本文专注于伯努利二值变量这一特殊结构,利用对称算子直接推导出无需数值积分近似的任意时间跨度精确闭式后验。

评分

  • 新颖性: ⭐⭐⭐⭐ 巧妙利用对称二值算子在连续时间参数空间推导出任意步长的无数据依赖翻转率与闭式跨步后验,理论推导优雅且切中痛点。
  • 实验充分度: ⭐⭐⭐⭐ 在 LSUN、FFHQ、Binarized MNIST 及统计物理 Ising 模型上进行了全面验证,并附带了修正 BLD 基线对照、延迟显存实测与 Bootstrap 置信区间分析。
  • 写作质量: ⭐⭐⭐⭐⭐ 痛点剖析一针见血,主文定理与附录证明衔接严密,对 BLD 跨步近似缺陷的剖析尤为透彻。
  • 价值: ⭐⭐⭐⭐ 可作为二值潜空间生成与科学二值数据建模的直接替代底座,在不改网络结构和损失函数的前提下提供可靠的免训练少步采样能力。