Unrestrained Simplex Denoising for Discrete Data. A Non-Markovian Approach Applied to Graph Generation¶

会议: NeurIPS 2026
arXiv: 2603.28572
代码: https://github.com/yoboget/UNSID
领域: 图像生成
关键词: 概率单纯形, 非马尔可夫去噪, 离散图生成, 狄利克雷概率路径, 沃罗诺伊概率

一句话总结¶

针对传统离散扩散模型离散状态突变和标准连续流匹配误差累积的瓶颈，提出在概率单纯形上建模的无约束非马尔可夫去噪框架（UNSID），利用狄利克雷显式概率路径与一步重采样机制，在图与分子生成基准上实现了更快的采样速度与超越离散扩散/流匹配的生成质量。

研究背景与动机¶

离散结构（如分子图与通用网络拓扑）的生成建模近年来被离散扩散模型（如 DiGress）与离散流匹配（如 DeFoG）所主导。然而，在离散状态空间中直接进行跳转去噪存在固有的结构脆弱性：微小的去噪状态翻转（例如单条边的误删或突变）可能直接破坏图的连通性或拓扑不变量；同时离散图神经网络由于对称性受限，通常必须在每步去噪中重新计算昂贵的手工结构特征（如拉普拉斯特征向量、随机游走概率等）。为了平滑这一过程，部分工作尝试将离散类别松弛到无约束欧氏空间中进行扩散，但这不仅引入了冗余的自由度与高斯噪声，还破坏了概率单纯形的天然几何约束；而现有单纯形上的流匹配（Dirichlet Flow Matching）由于线性插值构造导致概率支集跳跃断裂，且反向动力学受限于连续轨迹约束，极易引发去噪误差的复利累积。

这种困境的核心矛盾在于：离散扩散模型由于硬状态跳跃而脆弱且依赖昂贵辅助特征，而现有的连续/单纯形去噪模型又普遍施加了严格的马尔可夫轨迹连续性约束（即要求反向去噪过程在小时间步内连续逼近前一噪声状态），导致局部预测误差在长采样链条中无法逃逸并不断放大。

本文的切入角度是打破去噪过程对前一噪声状态的强依赖，将离散数据直接映射为概率单纯形的顶点，并在单纯形上构建非马尔可夫的去噪机制。核心 idea：提出无约束单纯形去噪（UNSID），在概率单纯形上建立条件独立于历史噪声状态的显式狄利克雷概率路径，通过“后验预测-单纯形重采样”消除马尔可夫连续轨迹约束，从理论上实现总变差误差收缩，并彻底摆脱手工图特征。

方法详解¶

整体框架¶

UNSID 的整体流水线将多变量离散属性（图的节点类型与边类型）表示为多单纯形顶点上的单热向量。前向加噪过程在各个维度独立进行，通过显式狄利克雷分布将干净顶点平滑扩散到整个单纯形内部；去噪时，等变图神经网络接收包含连续单纯形权重的图结构作为输入，预测干净图属性的后验类别分布；反向采样不沿向量场积分，而是直接从预测后验中采样离散结构，并在前一时间步的狄利克雷路径上重新加噪采样，形成两阶段的收缩映射。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["先验噪声初始化<br/>均匀单纯形 / 边际混合"] --> B["当前单纯形状态 x_t"]
    B --> C["等变图去噪网络<br/>预测干净数据后验分布"]
    C --> D["无约束单纯形重采样<br/>后验采样结合狄利克雷重加噪"]
    D --> E["校正子采样与步进<br/>总变差误差逐步收缩"]
    E -->|迭代 T 步| F["离散图与分子输出<br/>顶点投影与化学有效性评估"]

关键设计¶

1. 显式狄利克雷概率路径与沃罗诺伊概率标定：解决单纯形插值的支集断裂与调度黑盒问题

传统的随机插值方法在单纯形上会造成严重的支集不连续（当噪声水平超过阈值时，后验分布崩溃到单顶点邻域）。UNSID 采用显式狄利克雷概率路径，对任意时间 $t \in [0, 1)$，节点或边的第 $i$ 维加噪分布定义为具有全单纯形支集的狄利克雷分布： $$q_t(\mathbf{x}^{(i)} \mid \mathbf{x}_1^{(i)}) = \text{Dir}(\mathbf{x}^{(i)}; \mathbf{1} + \alpha_t \mathbf{x}_1^{(i)})$$ 其中噪声调度满足 $\alpha_0 = 0$ 且 $\lim_{t \to 1} \alpha_t = \infty$，实践中采用 $\alpha_t = -a \log(1 - t)$。为了定量分析和校准该概率路径的混合速度，本文引入了几何工具“沃罗诺伊概率”（Voronoi Probability），即加噪状态 $\mathbf{x}_t$ 仍然落在其生成顶点 $\mathbf{e}_i$ 对应沃罗诺伊多胞形内的概率。针对狄利克雷路径，推导出了闭式解析解： $$P_{v_i}(\mathbf{x}_t \in \mathbb{V}_i \mid \mathbf{x}_1 = \mathbf{e}_i) = \sum_{k=0}^{K-1} \frac{(-1)^k \binom{K-1}{k}}{(k+1)^{\alpha_t + 1}}$$ 利用此解析表达式，研究者可以直观评估不同超参数 $a$ 在不同类别数 $K$ 下的噪声混合行为，消除了传统启发式设计调度曲线的不确定性。

2. 条件独立非马尔可夫去噪核：消除误差复利并实现总变差误差收缩

传统扩散与流匹配强行限制反向轨迹在时间极限下连续（即 $\lim_{h \to 0} p_{t+h \mid t}(\mathbf{x} \mid \mathbf{x}_t) = \delta(\mathbf{x} - \mathbf{x}_t)$），若当前状态处于去噪器预测不良的局部敏感区域，后续采样将被锁定在该区域引发误差累积。UNSID 引入非马尔可夫条件独立假设，即给定干净样本 $\mathbf{x}_1$ 时，任意不同时刻的噪声状态条件独立：$q_t(\mathbf{x}_t \mid \mathbf{x}_s, \mathbf{x}_1) = q_t(\mathbf{x}_t \mid \mathbf{x}_1)$。基于此假设，单步去噪核直接退化为简单的两阶段重采样过程： $$p_{t+h}(\mathbf{x}_{t+h} \mid \mathbf{x}_t) = \sum_{\mathbf{x}_1 \in \mathcal{X}} q_{t+h}(\mathbf{x}_{t+h} \mid \mathbf{x}_1) P_{1 \mid t}(\mathbf{x}_1 \mid \mathbf{x}_t) = \mathbb{E}_{\mathbf{x}_1 \sim P_{1 \mid t}(\cdot \mid \mathbf{x}_t)} \left[ q_{t+h}(\mathbf{x}_{t+h} \mid \mathbf{x}_1) \right]$$ 该设计使反向单步不仅避免了流匹配中计算昂贵向量场导数的开销，而且从理论上证明了误差收缩性质：流匹配在常微分方程流映射下只能保持已累积的总变差误差（$\|\Phi_\# \tilde{\pi}_t - \pi_{t+h}\|_{\text{TV}} = \|\tilde{\pi}_t - \pi_t\|_{\text{TV}}$），而 UNSID 的单步马尔可夫算子 $Q$ 能将总变差误差按因子至少严格压缩至 $(1 - \varepsilon_{t+h}) \|\tilde{\pi}_t - \pi_t\|_{\text{TV}}$，使采样轨迹能自发逃逸局部误差陷阱。

3. 内生区分度的单纯形等变去噪网络：摆脱手工结构特征

在图神经网络中，离散节点和边属性由于置换对称性受到 1-WL 测试表达能力上限的制约，迫使前人模型（如 DiGress、DeFoG）必须在每一步计算拉普拉斯谱特征或随机游走概率。UNSID 中的图加噪状态直接由连续概率向量构成，连续单纯形坐标天然作为各个节点和边的唯一特征标识符，理论上打破了对称性瓶颈并实现了全图通用表达。因此，UNSID 的去噪模型采用基于 GPS 图 Transformer 结构的等变 GNN，只需输入加噪的连续单纯形张量与图大小标识，彻底抛弃了每一步昂贵的谱分解与环结构统计特征，大幅精简了网络计算量。

损失函数 / 训练策略¶

后验分布 $P_{1 \mid t}(\mathbf{x}_1 \mid \mathbf{x}_t)$ 由参数化网络 $P_{1 \mid t}^\theta(\mathbf{x}_1 \mid \mathbf{x}_t) = \prod_{i=1}^L \text{Cat}(\mathbf{x}_1^{(i)}; \boldsymbol{\omega}^{(i)})$ 近似，其中 $\boldsymbol{\omega} = \text{softmax}(f_\theta(\mathbf{x}_t, t))$。模型通过最小化节点和边属性的加权负对数似然进行训练： $$\mathcal{L}_\theta = \mathbb{E}_{t, \mathcal{G}_1, \mathcal{G}_t} \left[ \frac{\gamma}{n} \sum_{i} \left[ -\log P_{1 \mid t}^\theta (\mathbf{v}_1^{(i)} \mid \mathcal{G}_t) \right] + \frac{1 - \gamma}{n(n-1)} \sum_{i \neq j} \left[ -\log P_{1 \mid t}^\theta (\mathbf{a}_1^{(i,j)} \mid \mathcal{G}_t) \right] \right]$$ 其中时间步 $t \sim \mathcal{U}(0, 1)$，加噪图 $\mathcal{G}_t \sim q_t(\mathcal{G}_t \mid \mathcal{G}_1)$，节点权重 $\gamma$ 平衡节点属性与边连接损失。在推理采样初始化时，除了均匀分布先验 $\text{Dir}(\mathbf{1})$ 之外，模型还支持基于边际类别频率构造的边际狄利克雷混合先验 $q_0^{\text{marg}}(\mathbf{x}) = \sum_{k=1}^K m_k \text{Dir}(\mathbf{x}; \mathbf{1} + \kappa \mathbf{e}_k)$，进一步提升稀疏图生成的初始化质量。

实验关键数据¶

主实验¶

在分子生成基准 ZINC250K（表 1）以及合成规则图基准 Planar 与 SBM（表 2）上，UNSID 展现出超越离散扩散与流匹配基准的生成质量。

数据集	模型	有效性 Valid (%) ↑	FCD / Deg. MMD (×10³) ↓	NSPDK / Orbit MMD (×10³) ↓	连通/聚类 Clust. (×10³) ↓
ZINC250K	训练集参考	100.00	1.13	0.10	-
	DiGress	94.99	3.48	2.10	-
	GruM	98.65	2.26	1.50	-
	GraBFN	99.22	2.12	1.30	-
	SID	99.50	2.06	2.01	-
	DiriFM (单纯形流匹配)	97.32	2.79	1.92	-
	UNSID (本文)	99.98	1.79	1.00	-
Planar	训练集参考	100.0	0.25	0.78	36.6
	DiGress	45.5	0.71	1.31	45.4
	GruM	91.0	0.38	6.42	40.7
	DeFoG	99.5	0.50	0.60	50.1
	DiriFM	0.0	6.44	45.04	196.6
	UNSID (本文)	100.0	0.36	0.78	39.9
SBM	训练集参考	93.5	1.57	37.0	50.1
	DiGress	51.0	1.28	39.6	51.5
	GruM	67.5	2.20	40.4	49.9
	DeFoG	90.0*	0.60*	55.6	51.7
	UNSID (本文)	78.5	1.74	52.1	49.9

*注：DeFoG 在 SBM 上的极高表现伴随 0% 的新颖率，存在严重训练集记忆过拟合现象。

消融实验¶

在 ZINC250K 数据集上针对采样函数评估步数（NFE）与狄利克雷噪声调度参数 $a$ 进行了系统消融（表 7、表 8）。

配置 / 超参	有效性 Valid (%) ↑	FCD ↓	NSPDK (×10³) ↓	说明
NFE = 16	94.29	6.54	7.97	极低步数下仍具备合理有效性
NFE = 64	99.64	3.26	2.79	仅用 64 步即达到 SOTA 级别的化学有效性
NFE = 256	99.91	2.17	1.59	性能逼近强基线
NFE = 1024 (默认)	99.98	1.71	0.87	达到最优 FCD 与结构距离指标
调度参数 a = 1	98.87	4.13	2.61	末期轨迹变化过陡，去噪难度骤增
调度参数 a = 3 (默认)	99.98	1.79	1.00	沃罗诺伊概率适中平滑，综合性能最优
调度参数 a = 6	99.92	1.66	0.82	后期路径平缓，保持极优分布相似度
调度参数 a = 10	99.91	1.95	1.25	后期矫正信号减弱，指标轻微回落

关键发现¶

去噪范式的本质决定性能差距：在相同网络结构下，比较单纯形流匹配（DiriFM）与 UNSID，DiriFM 在 Planar 上有效性甚至崩溃至 0%，在 ZINC250K 上的 FCD 为 2.79，而 UNSID 达到 1.79。这表明无约束采样核消除误差累积是算法制胜的关键，而非单纯引入单纯形表示。
采样效率与参数规模全面占优：在 Planar 数据集上，UNSID 的单次采样耗时仅 6.3 秒（参数量 5.4M），明显快于 DiGress 的 16.9 秒（8.9M）与 GruM 的 29.3 秒（7.1M），去噪过程中无需重算拉普拉斯特征带来了显著的时间优势。
步数泛化能力卓越：仅需 64 个 NFE，UNSID 生成分子的有效性就突破 99.6%，足以应对需要高通量分子生成的实际工业筛选场景。

亮点与洞察¶

沃罗诺伊概率解析解：论文巧妙地通过独立伽马变量的最大值比率，给出了高维单纯形顶点沃罗诺伊概率在狄利克雷加噪下的封闭求和解析式，为单纯形概率路径的动力学分析提供了扎实的理论量化工具。
跳出去噪历史的非马尔可夫重采样：摆脱了必须从 $t$ 步平滑连续移动到 $t+h$ 步的马尔可夫教条，将每一步转换为“全图预测 + 目标时间独立重新加噪”，从机制上根治了连续动力系统在敏感区域陷入误差复利循环的弊端。
端到端纯净架构：摒弃了离散图生成模型中必不可少的繁琐手工谱特征提取工程，使图神经网络能以纯净端到端的方式直接消费连续单纯形张量。

局限与展望¶

模态验证受限于图结构：虽然理论框架面向任意多变量离散数据，但文中仅在图与分子结构上完成了实证，在离散文本序列、代码、蛋白质序列等大规模语言/生物模态上的表现尚未验证。
稠密全连接图表示的可扩展性瓶颈：当前模型沿用全对邻接矩阵的连续表示，其显存与计算复杂度关于图节点数呈二次方增长，对包含数千节点的超大图尺度需要结合稀疏化或局部展开技术拓展。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次提出连续概率单纯形上的非马尔可夫无约束去噪框架，并给出了沃罗诺伊概率的解析解与误差收缩证明。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖分子图（ZINC250K、QM9H）、合成拓扑图（Planar、SBM）与真实酶蛋白图，配有严谨的同架构控制对照实验与消融。
写作质量: ⭐⭐⭐⭐⭐ 动机阐述清晰深刻，数学推导完备自洽，对单纯形病态行为与马尔可夫误差积累的理论剖析透彻。
价值: ⭐⭐⭐⭐⭐ 为离散数据（尤其图与分子）的生成建模开辟了极具竞争力的单纯形动力学路线，具有极高的理论与实用借鉴价值。