Discrete Bayesian Sample Inference for Graph Generation¶
会议: ICLR2026
OpenReview: https://openreview.net/forum?id=py2NPbAdvH
代码: 待确认
领域: 图生成 / 生成模型
关键词: 图生成, 贝叶斯样本推断, 离散扩散, 随机微分方程, 分子生成
一句话总结¶
本文提出 GraphBSI,把贝叶斯样本推断(BSI)从连续数据推广到离散类别数据,让模型在「概率单纯形上的分布参数空间」里迭代精化对图的信念而非直接演化离散图,并把这一过程写成一族可调噪声 \(\gamma\) 的 SDE,在 Moses、GuacaMol 分子生成基准上以一次性(one-shot)方式取得 SOTA。
研究背景与动机¶
领域现状:图(分子、知识图谱、社交网络)是离散、无序的结构化数据,传统为连续数据(如图像)设计的生成模型不好直接套用。近年主流是把图离散化后用离散扩散(DiGress、DisCo、Cometh)或流匹配(DeFoG)来生成,另一条线是贝叶斯流网络(BFN):它不在样本上加噪去噪,而是在「样本分布的参数」上演化——即便底层样本是离散的,分布参数也能连续平滑地变化,因此天然适配离散数据,代表工作 GraphBFN 在分子生成上很有竞争力。
现有痛点:BFN 虽然好用,但它的推导建立在信息论视角之上,整套框架(参数空间、信息累积、对贝叶斯更新取极限近似)门槛高、不直观,难以理解和改造。另一方面,离散扩散把图状态直接当作要去噪的对象,离散状态的跳变让采样轨迹不够平滑,在低函数调用预算下质量受限。
核心矛盾:要既保留「在连续参数空间上演化、平滑处理离散结构」这个好处,又要有一个简单、可解释、能自由调控采样随机性的框架。BFN 在「可解释性」和「可控性」上都欠缺。
本文目标:(1) 把贝叶斯样本推断(BSI,原本只针对连续数据)推广到类别数据,给离散生成一个比 BFN 更简洁的解释;(2) 把它形式化成 SDE,从而能引入一个可调噪声的采样器族;(3) 在此基础上做出一个高质量的一次性图生成模型。
切入角度:BSI 把「生成」看成一连串贝叶斯更新——从一个宽泛的先验信念出发,不断对未知样本做带噪测量并把信息融进信念,直到信念足够尖锐就采样。作者发现:只要把 BSI 里的连续高斯信念换成「单纯形上的类别信念(logits 上的高斯)」,整套贝叶斯更新依然成立,而且比 BFN 少了一层信息论包装。
核心 idea:在「类别分布参数(logits)」的连续空间里做贝叶斯样本推断,用一个神经网络反复重构未知图、再用带噪测量更新信念;并把这个迭代写成 SDE,引入噪声系数 \(\gamma\) 得到一族「边缘分布相同、随机性可调」的采样器。
方法详解¶
整体框架¶
GraphBSI 不直接对图本身去噪,而是维护一个「对要生成的图的信念」,这个信念是定义在节点/边类别 logits \(z\) 上的高斯,对应一个类别分布 \(\mathrm{Cat}(\mathrm{softmax}(z))\)。生成时从先验 \(z_0\sim\mathcal N(\mu_0,\beta_0 I)\) 出发,循环执行三步:① 用网络 \(f_\theta(z_t,t)\) 根据当前信念重构出对图的估计 \(\hat x\);② 围绕 \(\hat x\) 采一个带精度 \(\alpha\) 的噪声测量 \(y\sim\mathcal N(\hat x,\alpha^{-1}I)\);③ 用贝叶斯定理把测量融进信念 \(z\leftarrow z+\alpha y\)。随着时间从 \(t=0\) 到 \(t=1\),信念越来越尖锐,最后从 \(\mathrm{Cat}(\mathrm{softmax}(z_1))\) 采样(或直接量化重构结果)得到图。
把步数 \(k\to\infty\) 取极限,这个离散更新收敛成一个 SDE(Theorem 3):\(dz_t=\beta'(t)f_\theta(z_t,t)\,dt+\sqrt{\beta'(t)}\,dW_t\)。在此基础上作者推出一族保持边缘分布不变、只改随机性强弱的广义 SDE(噪声系数 \(\gamma\)),并设计了两套数值离散方案(Euler-Maruyama 与 Ornstein-Uhlenbeck)外加末步量化技巧来高效采样。整条流程如下:
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["先验信念 z0<br/>logits 高斯"] --> B["类别贝叶斯样本推断<br/>z ← z + αy 累积更新"]
B --> C["图适配<br/>节点/边独立类别 + 置换不变网络"]
C --> D["噪声可调 SDE 族<br/>γ 调控随机性、边缘不变"]
D --> E["OU 离散 + 末步量化<br/>低 NFE 高质量采样"]
E --> F["生成图 (X, A)"]
关键设计¶
1. 类别贝叶斯样本推断:把 logits 当信念,用贝叶斯更新代替直接去噪
原始 BSI 只处理连续数据,把信念建成各向同性高斯。本文要解决的痛点是:离散图没法直接套连续高斯信念。作者的做法是把样本放到概率单纯形上 \(x\in\Delta^n_{c-1}\),信念写成 \(p(x\mid z)=\mathrm{Cat}(x\mid\mathrm{softmax}(z))\),其中 \(z\in\mathbb R^{n\times c}\) 是类别 logits。关键结果是 Theorem 1:给定带精度 \(\alpha\) 的高斯测量 \(y\sim\mathcal N(x,\alpha^{-1}I)\),后验信念仍是同族类别分布,且 logits 更新极其简洁——
也就是说,logits 空间里的贝叶斯更新就是「按精度加权累加测量」。这点和连续 BSI 有本质区别:连续情形是把新信念与旧状态做插值,而类别情形是各分量直接累积。多步测量叠加后 \(z_k=z_0+\sum_i\alpha_i y_i\),配上单调递增的精度调度 \(\beta(t)\)(令 \(\alpha_i=\beta(t_{i+1})-\beta(t_i)\))就控制了到时刻 \(t\) 注入信念的总信息量。这样做的好处是绕开了 BFN 里对贝叶斯更新取极限的近似,给离散生成一个更干净的概率解释,并且能直接推出 ELBO 训练目标(Theorem 2),其简化形式就是一个加权重构损失 \(\mathbb E[\beta'(t)/2\cdot\lVert f_\theta(z,t)-x\rVert_2^2]\)——训练只需让网络从带噪信念里重构出真样本。
2. 图适配:节点与边各自独立类别,配置换不变网络生成变长图
有了类别 BSI,还要把它落到图上。图被表示为 \((X,A)\),\(X\in\Delta^N_{c_X-1}\) 是节点类别(原子类型)的 one-hot,\(A\in\Delta^{N\times N}_{c_A-1}\) 是边类别(键类型)的 one-hot,并特意把「边不存在」当作边的第一个类别。作者把每个节点、每条边都当作类别信念的一个独立分量,于是上面的类别 BSI 框架可以逐分量地直接套用;节点和边之间的依赖关系不在信念里建模,而是交给重构网络 \(f_\theta\) 去捕捉。为保证生成分布对节点编号顺序无关,\(f_\theta\) 选用置换不变的图 Transformer(沿用 DiGress/DeFoG 的架构,输入含类别概率、熵、随机游走特征、时间步正弦嵌入),在各向同性噪声下整个生成模型即为置换不变。对于变长图,先从数据集的节点数边缘分布里采一个节点数 \(N\),再对节点/边生成,实现上通过对多余节点和边做掩码完成。这套表示让同一个离散框架既能生成分子图,也能(如附录所示)生成 DNA 序列等一般离散数据。
3. 噪声可调 SDE 族:一个 \(\gamma\) 在确定性 ODE 与高随机采样器之间连续滑动,且边缘分布不变
离散多步更新在 \(\Delta t\to0\) 极限下收敛为 SDE \(dz_t=\beta'(t)f_\theta dt+\sqrt{\beta'(t)}dW_t\)(Theorem 3)。但固定的随机性未必最优,作者借鉴 Karras et al. (2022) 的思路,推出一族保持边缘密度 \(p_t(z_t)\) 不变、只用 \(\gamma\) 调随机性的广义 SDE(Theorem 4):
\(\gamma=0\) 退化成确定性的概率流 ODE;\(\gamma=1\) 恢复原始 SDE;\(\gamma\) 越大轨迹越随机,极限 \(\gamma\to\infty\) 时每步都用最新预测覆盖当前状态。理论上所有 \(\gamma\) 边缘分布相同,但实践中分数项 \(\nabla_{z_t}\log p_t(z_t)\) 没有闭式解只能近似,高随机性能纠正前几步的错误却需要更细的离散。关键是这个分数函数不用额外训练:Theorem 5 证明 BSI 损失本身就是一个分数匹配损失,分数模型可直接由重构网络给出 \(s_\theta(z,t)=\big(\mu_0+\beta(t)f_\theta(z,t)-z\big)/\big(\beta(t)+\beta_0\big)\)。和 BFN 不同,CatBSI 是从先验信念 \(p(z\mid t=0)\) 采样而非固定先验,天然避开了 \(t=0\) 处的不连续。
4. OU 离散 + 末步量化:低函数调用预算下也能高质量采样
SDE 没有闭式解,需数值采样。简单的 Euler-Maruyama(EM)在足够细的时间网格上表现不错,但在高 \(\gamma\)、少步数下会不稳定。作者改为在每个小区间内冻结重构 \(\hat x=f_\theta(z_t,t)\),使广义 SDE 在该区间退化为 Ornstein-Uhlenbeck 过程从而可解析求解(Theorem 6),得到带精确边缘的更新 \(z_{t+\Delta t}\sim m+(z_t-m)e^{-\kappa\Delta t}+\sqrt{\tfrac{\gamma\beta'}{2\kappa}(1-e^{-2\kappa\Delta t})}\,\mathcal N(0,I)\),其中 \(\kappa=\tfrac{(\gamma-1)\beta'}{2(\beta_0+\beta)}\)。当 \(\Delta t\to0\) 时 OU 收敛回 EM。另一个提效技巧是「量化代替采样」:当 \(t=1\) 附近信念已足够尖锐时,最后一步直接采样会引入多余噪声,但此时信念里的信息已够网络做近乎完美的重构,于是提前在较低精度停下、把重构结果投影(量化)回样本空间即可。再加上非均匀时间网格参数 \(\rho\)(\(t_i=(i/k)^\rho\))把函数调用更密集地放在更需要的时间段,整套方案让 GraphBSI 在仅 50 次函数评估下就有竞争力,500 次时进一步刷新 SOTA。
损失函数 / 训练策略¶
训练目标是类别 BSI 的 ELBO(Theorem 2)。\(k\to\infty\) 时其可优化部分化简为加权重构损失(Eq. 5):\(\mathcal L=\mathbb E_{x,t,z}\big[\beta'(t)/2\cdot\lVert f_\theta(z,t)-x\rVert_2^2\big]\),对应 Algorithm 2——采样本 \(x\)、采先验 \(z_0\)、采时刻 \(t\)、按精度做一次带噪测量得到信念 \(z\)、让网络从 \(z\) 重构 \(x\) 并做 MSE。当 \(\beta_0\to0\)(先验在 \(t=0\) 退化为 Dirac)时,该损失与连续时间类别 BFN 损失只差一个常数。实践上采用指数精度调度、对节点/边用数据集边缘分布加小方差(1.0)的高斯先验,并用预条件让网络预测信念与真值之差:\(f_\theta(z,t)=\mathrm{softmax}\big(z+\tilde f_\theta(z,t)\big)\)。
实验关键数据¶
主实验¶
在 GuacaMol 和 Moses 两个分子生成基准上,对比扩散/流匹配 SOTA(DiGress、DisCo、Cometh、DeFoG、GraphBFN),报告 EM 与 OU 两种离散、50 与 500 步两种预算。
| 数据集 | 配置 | 关键指标 | GraphBSI | 之前最好 |
|---|---|---|---|---|
| GuacaMol | OU, 500 步 | Validity ↑ | 99.6 | 99.0 (DeFoG) |
| GuacaMol | EM, 500 步 | FCD ↑ | 82.6 | 73.8 (DeFoG) |
| Moses | EM, 500 步 | FCD ↓ | 0.72 | 1.07 (GraphBFN) |
| Moses/Guaca | OU, 50 步 | Validity ↑ | 99.7 / 99.2 | 83.9 / 91.7 (DeFoG) |
仅 50 步时,GraphBSI(两种离散)在除 Moses novelty 外的所有指标上都超过 DeFoG;500 步时 OU 离散在 GuacaMol 上全指标超过所有现有模型并把 validity 打满,EM 离散把 Moses 的 FCD 从 1.07 压到 0.72。
合成图基准¶
在 planar / tree / SBM 合成图上(Table 2,每次生成 40 图、重复 5 次):
| 配置 | 数据集 | V.U. ↑ | 说明 |
|---|---|---|---|
| GraphBSI (OU) 1000 | Planar | 100.0 | validity 打满 |
| GraphBSI (EM) 1000 | Tree | 96.5 | Ratio 1.3,分布相似度最好之一 |
| GraphBSI (OU) 1000 | SBM | 77.5 | SBM 较难,validity 适中 |
作者指出合成集仅 128 张图、评测方差大,mean ratio 需谨慎解读。
消融实验¶
| 配置 | 现象 | 结论 |
|---|---|---|
| \(\gamma=0\)(概率流 ODE) | 所有指标都很差 | 纯确定性采样不行,噪声不可或缺 |
| \(\gamma\) 中等(20~100) | FCD 最优 | FCD 偏好中等随机性 |
| \(\gamma\) 增大 | SNN、Filters 提升但 novelty 下降 | 高噪声让样本更贴近训练分布 |
| EM vs OU(高 \(\gamma\)) | EM 在 \(\gamma\Delta t\) 过大时失稳 | OU 在各预算/噪声下都稳,普遍持平或优于 EM |
| 非均匀网格 \(\rho<1\) | FCD 在 50 步下改善 | 后段时刻细化离散更划算 |
关键发现¶
- 噪声调控是性能主驱动:Fig. 8 显示优化 \(\gamma\) 是性能增益的关键来源;\(\gamma=0\) 时全面崩坏,说明把 BSI 写成可调噪声 SDE 这一步是涨点核心。
- OU 离散在低预算下尤为关键:局部线性化让少步数采样仍稳定高质,是 50 步即 SOTA 的工程基础。
- 指标间存在权衡:高随机性同时提升相似度类指标、压低 novelty——生成越像训练集,新颖性越低。
亮点与洞察¶
- 统一视角:把贝叶斯样本推断推广到类别数据后,理论分析直接打通了 BFN 与扩散模型的联系——logits 上的累积贝叶斯更新、SDE、分数匹配损失三者在一个框架里自洽,比信息论包装的 BFN 更易懂。
- 一个网络两用:Theorem 5 证明 BSI 重构损失本身就是分数匹配损失,于是分数函数无需额外训练即可从重构网络读出,省掉一套分数网络。
- 可迁移技巧:「边缘不变、\(\gamma\) 调随机性」的 SDE 族(借鉴 Karras)+「末步量化代替采样」是通用采样加速思路,可迁移到其他离散/连续生成器,尤其在低 NFE 场景。
- 离散即类别:把「边不存在」编码成边的一个类别、节点/边各自独立类别,是处理图离散结构干净利落的表示,让同一框架顺手覆盖 DNA 序列等一般离散数据。
局限与展望¶
- 节点和边在信念层被建成相互独立的类别,所有结构依赖都压给重构网络 \(f_\theta\) 承担,复杂高阶依赖是否被充分建模存疑。
- 分数项 \(\nabla_{z_t}\log p_t(z_t)\) 只能近似,导致理论上「所有 \(\gamma\) 边缘相同」在实践中不成立,必须靠网格搜索 \(\gamma\)/步数/\(\rho\) 才能拿到最优,调参成本不低。
- 合成图(仅 128 图)评测方差大、SBM 上 validity 仅 77.5,说明在结构更复杂的非分子图上仍有差距。
- novelty 与相似度/有效性的权衡未被根治:高噪声涨 FCD/SNN 却掉 novelty,如何同时拿到高质量与高新颖性仍是开放问题。
相关工作与启发¶
- vs 贝叶斯流网络(GraphBFN):BFN 在信息论视角下操作分布参数、对贝叶斯更新取极限近似,门槛高;本文用 BSI 视角把生成看成一串贝叶斯更新,避开极限近似、给出更简洁解释,并从先验信念采样而非固定先验,避开 \(t=0\) 不连续。Moses FCD 从 1.07(GraphBFN)降到 0.72。
- vs 离散扩散(DiGress / DisCo / Cometh):它们在离散图状态上去噪;本文在连续 logits 空间演化信念、离散性自然由 softmax 处理,采样轨迹更平滑,低步数下质量更高。
- vs 流匹配(DeFoG):DeFoG 是当前一次性图生成强基线;GraphBSI 在 50 步即在多数指标上超过它,500 步进一步全面领先,差异主要来自可调噪声 SDE 与 OU 离散。
- vs 概率流 ODE(Xue et al. 2024):\(\gamma=0\) 即退化为该确定性方案,但实验表明纯 ODE 表现最差,凸显本文「引入并调控随机性」相对纯确定性采样的价值。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 把 BSI 推广到类别数据并统一 BFN/扩散/分数匹配,框架层面有真贡献。
- 实验充分度: ⭐⭐⭐⭐ 分子基准 SOTA 充分,合成图与噪声/网格消融到位;非分子复杂图覆盖偏弱。
- 写作质量: ⭐⭐⭐⭐ 理论推导(6 个定理)清晰、图示直观,但符号密度较高。
- 价值: ⭐⭐⭐⭐ 给离散图生成一个更可解释、可控的框架,采样加速技巧有通用迁移价值。