Blocked Gibbs Meets Diffusion Transformers: Unsupervised Learning for Constraint Optimization¶

会议: NeurIPS 2026
arXiv: 2605.25129
代码: https://github.com/khalil-research/BloGDiT
领域: 优化/理论
关键词: 约束优化, 扩散Transformer, 块吉布斯采样, 无监督学习, 玻尔兹曼采样

一句话总结¶

针对标准扩散模型全局微小高斯去噪与离散约束优化局部靶向重修之间的根本失配，本文提出块吉布斯扩散 Transformer（BloGDiT），通过可退火掩码重采样与自适应掩码策略，以无监督玻尔兹曼采样实现了跨二值图优化与多值通用约束满足问题的强泛化求解。

研究背景与动机¶

约束满足问题（CSP）与约束优化问题（COP）是运筹学、组合数学与人工智能的基础核心，传统求解主要依赖回溯树搜索（如 Google OR-Tools）或手工设计的局部搜索算法（如大邻域搜索 LNS）。尽管这些经典求解器性能强劲，但它们通常针对每个实例从头求解，计算开销无法在问题实例之间分摊。近年来，基于图神经网络（GNN）的扩散模型被引入神经组合优化领域，试图学习可复用的迭代细化机制；然而，现有 GNN 扩散模型大多局限于二值图优化任务，依赖局部消息传递，难以建模稠密非二值变量（如数独中的九值分类变量）以及包含长程依赖的全局约束（如全互异约束 AllDifferent）。

为了克服 GNN 的局限，将具备全局注意力与优异可扩展性的 Diffusion Transformer（DiT）引入约束优化成为自然的构想。然而直接应用标准 DiT 的实验结果却令人大失所望，模型泛化极差且易陷入次优解。这种失败的根源在于标准连续扩散过程与离散约束优化逻辑之间的本质冲突：标准扩散在每一步对全量变量施加均质微小的连续高斯去噪；而离散约束优化在解接近可行时，其核心需求往往是保持绝大部分已满足约束的变量冻结不变，仅对少数违反约束或高不确定性的局部子集进行大幅度、破坏重构式的靶向修改（Surgical Edits）。全局均匀去噪不可避免地破坏了全局解的既有协同结构，导致反复震荡。

受到马尔可夫链蒙特卡洛（MCMC）中的块吉布斯采样（Blocked Gibbs Sampling）以及运筹学中大邻域搜索“破坏-修复”（Destroy-and-Repair）范式的启发，本文重新将扩散生成过程形式化为一系列条件重采样更新。核心 idea：提出块吉布斯扩散 Transformer（BloGDiT），通过引入辅助掩码变量将全变量联合去噪转化为基于块的条件高斯去噪，并在时间步上退火掩码块大小（前期大块探索、后期小块精确微调），配合无监督玻尔兹曼反向 KL 上界训练与自适应测试期变量选择，实现高效的端到端神经约束求解。

方法详解¶

整体框架¶

BloGDiT 将约束优化转化为从非归一化玻尔兹曼分布 $p_B(X) \propto \exp(-H(X)/\tau)$ 进行无监督采样的逆向扩散过程。问题实例中的离散变量被松弛到连续 logit 空间 $Z \in \mathbb{R}^{n \times K}$，并经行级 Softmax 投影为概率向量 $\tilde{X}$ 以评估平滑可微的约束惩罚能量 $H(\tilde{X})$。扩散前向过程对掩码选中的变量块注入高斯噪声，其余变量原样复制保持干净；逆向生成过程从高噪声先验出发，在每个时间步通过自适应或随机掩码选择待更新的变量块，利用带相对位置编码偏置的 Transformer 骨干网络，在给定未掩码“干净”上下文的条件下对该块变量进行去噪预测。掩码比例随时间单调退火，驱动解从全局粗粒度重组平滑过渡到局部细粒度修复。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["高斯基底先验 Z_T<br/>约束已知量固定 / 其余均匀随机"] --> B["时间退火掩码机制<br/>rho_t 从大到小动态调度"]
    B --> C["自适应掩码选择<br/>置信度边缘 / 违规度 / 关联约束"]
    C --> D["条件块扩散去噪<br/>未掩码变量冻结 / 掩码块经 DiT 预测均值方差"]
    D --> E["无监督玻尔兹曼目标优化<br/>最小化联合轨迹 KL 上界 (能量+熵+前向匹配)"]
    E --> F["离散解输出 X_0<br/>行级 argmax 获得最终可行解"]

关键设计¶

1. 掩码增强的块吉布斯扩散过程：从全量联合去噪解耦为条件块去噪

为了消除标准扩散对全变量更新引发的相互破坏，BloGDiT 在每个时间步 $t$ 引入辅助二值掩码向量 $m_t \in \{0, 1\}^n$。若第 $i$ 个变量的掩码位 $m_t^{(i)} = 1$，则该变量被纳入当前破坏修复块并被高斯扰动或去噪重采样；若 $m_t^{(i)} = 0$，则变量被严格冻结并恒等复制到下一时间步： $$q_\theta(z_{t-1}^{(i)} \mid Z_t, m_t) = \begin{cases} \delta(z_{t-1}^{(i)} - z_t^{(i)}) & \text{if } m_t^{(i)} = 0 \\ \mathcal{N}\left(\mu_\theta^{(i)}(Z_t, m_t), \mathrm{diag}((\sigma_\theta^{(i)})^2(Z_t, m_t))\right) & \text{if } m_t^{(i)} = 1 \end{cases}$$ 作者在数学上证明了掩码增强系统与原边际轨迹过程具有严格的采样等价性，即在联合分布 $q_\theta(Z_{0:T}, m_{1:T})$ 上边缘化掉所有掩码序列后，边际分布与条件转换核完全等价。这种设计赋予了 Transformer 聚焦于在已知部分有效赋值的强先验约束下重构局部子集的能力。

2. 掩码块大小的时间退火机制：平衡全局空间探索与局部精准修复

在离散搜索与优化中，初期需要大范围跳出局部极值，而后期接近可行解或极值解时则需要外科手术式的精准微调。BloGDiT 采用退火掩码率 $\rho_t \in [\rho_{\min}, \rho_{\max}]$。掩码率随时间步从 $t=T$ 到 $t=0$ 沿几何调度逐步衰减： $$\rho_t = \rho_{\min} \left(\frac{\rho_{\max}}{\rho_{\min}}\right)^{t/T}$$ 在扩散起始步，高掩码率（如 $\rho_{\max} = 0.9$）允许全图 90% 的变量重新采样，实现大尺度全局探索；随着时间步推移，掩码率逐步收缩至较低水平（如 $\rho_{\min} = 0.3$），使模型在保留绝大部分已建立有效模式的基础上，仅对少量冲突局部进行调整，完美拟合了大邻域搜索中逐步收敛邻域直径的直觉。

3. 测试期自适应掩码选择：引导模型对关键与违规变量做靶向修复

块吉布斯形式的显式掩码解耦带来了额外的推理期收益：逆向轨迹不再被强制绑定于随机掩码采样，而是可以根据中间状态的语义动态确定最需要重写的子集。BloGDiT 提出了三种互补的自适应掩码策略： - 置信度边缘策略（Top-Probability Margin）：衡量模型预测前两大候选值的 Softmax 概率差 $c_t^{(i)} = |q_\theta(x^{(i)}=j_1 \mid X_t) - q_\theta(x^{(i)}=j_2 \mid X_t)|$。概率差越小说明模型对该变量取值越摇摆，赋予其更高的掩码概率 $\rho_t^{(i)} \propto -c_t^{(i)}$。 - 违规严重度策略（Violation-Based）：依据当前离散化赋值统计每个变量参与违反约束的频次或权重 $v_i(X_t)$，采样概率与冲突数成正比 $\rho_t^{(i)} \propto v_i(X_t)$，精准狙击破坏约束的“病灶”。 - 约束关联网策略（Related Variables）：在超图或约束图上随机采样被违反的约束超边，直接将其涉及的所有关联变量成组激活作为破坏集，保留了约束内部的代数耦合性。

4. 结合约束图相对位置编码的连续 Logit 扩散 Transformer 架构

为了支持包含全局复杂约束的离散变量，BloGDiT 在变量连续松弛 logit 空间 $Z_t \in \mathbb{R}^{n \times K}$ 上操作。每个变量编码为其当前赋值的嵌入与绝对位置编码之和。针对约束依赖，模型构建二值约束图 $G=(V, E)$（只要变量 $i, j$ 共同参与任意约束则连边），并将其作为自注意力机制的相对位置编码偏置（RPE）： $$\mathbf{A}_{ij} = \frac{(\mathbf{W}^Q \mathbf{h}_i)^\top (\mathbf{W}^K \mathbf{h}_j)}{\sqrt{d}} + c \cdot \mathbb{I}[(i, j) \notin E], \quad c \le 0$$ 这一设计在保持 Transformer 全局跨变量交互能力的同时，对无直接约束关联的变量施加注意力负偏置，使网络在全局视野下依然能够高效感知局部的约束拓扑。

损失函数 / 训练策略¶

BloGDiT 采用完全无监督的方式进行训练，无需任何真实求解器标签。模型通过最小化逆向链边际与玻尔兹曼目标之间的逆向 KL 散度，推导并放缩出掩码增强的联合轨迹 KL 上界： $$\tau \mathrm{KL}\left(q_\theta(Z_{0:T}, m_{1:T}) \,\|\, p(Z_{0:T}, m_{1:T})\right) = \mathbb{E}_{q_\theta}[H(X_0)] - \tau \sum_{t=1}^T \mathbb{E}\left[h(q_\theta(Z_{t-1} \mid Z_t, m_t))\right] - \tau \sum_{t=1}^T \mathbb{E}\left[\mathbb{E}_{Z_{t-1}} \log p(Z_t \mid Z_{t-1}, m_t)\right] + C$$ 由于前向与逆向过程在训练阶段共享相同的掩码生成分布，掩码散度项严格抵消为零。目标由三项构成： 1. 最终能量项 $\mathbb{E}[H(X_0)]$：约束惩罚与优化目标的可微连续松弛，引导生成样本趋向零违规与高目标值； 2. 微分熵项 $h(q_\theta)$：限制在高斯分布的掩码坐标子空间上，具有简洁的闭式解 $\frac{1}{2}\sum_i m_t^{(i)} (\log(2\pi e) + \sum_k \log \sigma_{\theta, i, k}^2)$，防止扩散分布过早坍缩，保障生成多样性与探索能力； 3. 前向噪声匹配项：在高斯前向核下亦可闭式表达，惩罚预测均值和方差偏离实际加噪分布。

模型通过对 5 步逆向扩散反向传播端到端优化，使用 AdamW 优化器，学习率设为 $10^{-4}$，批大小为 512，并在训练周期中将温度 $\tau$ 从 1 线性退火至 0。

实验关键数据¶

主实验¶

论文在四类经典约束满足与组合优化基准上进行了评估：非二值高阶约束的数独（Sudoku）、非二值图约束满足的图着色（Graph Coloring）、二值图约束优化的最大独立集（MIS）与最大割（MaxCut）。测试涵盖分布内与严格的分布外（OOD）大尺寸图泛化。

任务 / 数据集	评测指标	BloGDiT (本文)	标准 DiT	最强学习基线	运筹求解器 (OR-Tools)	性能增益说明
Sudoku (SATNet / Easy)	解出率 (%) ↑	100.0	100.0	100.0 (ConsFormer)	100.0	达到与精确求解器相同的 100%
Sudoku (RRN / Hard OOD)	解出率 (%) ↑	94.1	48.6	77.7 (ConsFormer)	100.0	较标准 DiT 提升 +45.5%，超前 SOTA 16.4%
Coloring-5 ($n=50$)	解出率 (%) ↑	83.00	79.75	81.00 (ConsFormer)	83.08	逼近 OR-Tools 精确求解
Coloring-5 ($n=100$ OOD)	解出率 (%) ↑	54.67	48.91	47.33 (ConsFormer)	57.16	较标准 DiT 提升 +5.76%，显著超越基线
Coloring-10 ($n=100$)	解出率 (%) ↑	53.92	50.50	52.60 (ConsFormer)	52.41	超越 OR-Tools (+1.51%)
Coloring-10 ($n=200$ OOD)	解出率 (%) ↑	18.58	14.25	11.92 (ConsFormer)	10.25	较 OR-Tools 提升 +8.33% (近翻倍)
MIS (RB-small)	独立集大小 ↑	19.52	8.05	19.97 (RLSA) / 18.88 (DiffUCO)	20.10	较标准 DiT (+11.47) 彻底解决坍缩
MIS (RB-large)	独立集大小 ↑	37.05	0.29	40.19 (RLSA) / 38.10 (DiffUCO)	42.66	标准 DiT 在大图坍缩至 0.29，本文达 37.05
**MaxCut ($	V	=800$)**	相对最优差距 ↓	0.11	56.11	1.22 (ANYCSP) / 4.11 (DiffUCO)
**MaxCut ($	V	=1\text{K}$)**	相对最优差距 ↓	1.33	50.00	2.44 (ANYCSP) / 6.33 (DiffUCO)
**MaxCut ($	V	=2\text{K}$)**	相对最优差距 ↓	7.88	101.33	13.00 (RLSA) / 31.67 (DiffUCO)
**MaxCut ($	V	\ge 3\text{K}$)**	相对最优差距 ↓	32.75	162.00	51.63 (ANYCSP) / 116.75 (DiffUCO)

消融实验¶

论文针对自适应掩码选择（AMS）、目标函数中的微分熵项，以及掩码率退火区间 $[\rho_{\min}, \rho_{\max}]$ 进行了详尽消融分析。

测试配置 / 变体	Sudoku-Hard (%) ↑	Coloring-10 ($n=200$) (%) ↑	MIS RB-large (大小) ↑	MaxCut $\ge 3\text{K}$ (差距) ↓	说明
BloGDiT (完整模型)	94.12	18.58	37.05	32.75	自适应掩码 + 熵正则 + 退火调度
w/o 自适应掩码 (-AMS)	88.21	17.42	17.16	88.63	换为纯随机掩码，大图优化性能严重腰斩
w/o 微分熵 (-Entropy)	91.92	18.25	35.03	61.00	移除多样性熵项，优化易过早陷入局部极小
掩码区间消融 ($\rho_{\max}=0.9, \rho_{\min}=0.7$)	—	—	4.24	—	后期掩码过大，无法进行精细化微调修复
掩码区间消融 ($\rho_{\max}=0.5, \rho_{\min}=0.1$)	—	—	9.93	—	前期掩码过小，缺乏大范围全局空间跳跃能力

关键发现¶

局部靶向破坏修复的决定性价值：标准 DiT 在全局连续更新下，面对大尺寸图或高难度 OOD 约束问题时会出现严重的性能崩塌（如 MIS-RB-large 上平均独立集大小跌落至 0.29）。BloGDiT 通过掩码冻结绝大多数已建立的变量关系，使神经模型能集中容量处理局部冲突，一举将指标拉升至 37.05。
掩码退火的“两极约束”：消融数据显示，$\rho_{\max}$ 必须足够大（$\ge 0.7$）才能支持前期的跨模态空间探索；而 $\rho_{\min}$ 必须收敛至较小值（如 0.3），若保持大掩码（如 0.7）到末尾，由于每步破坏过剧，模型无法稳定在约束局部极值点上。
自适应变量优先级的增益放大：在图规模变大、约束复杂度剧增时（如 MaxCut $\ge 3\text{K}$ 和 MIS-RB-large），通过违规统计或不确定性边缘引导掩码，比盲目随机掩码的效果强两倍以上，证明了将 LNS 启发式与扩散结合的深层合理性。

亮点与洞察¶

将运筹学大邻域搜索与连续扩散 Transformer 完美统一：将经典 LNS 的“Destroy & Repair”思想严谨地形式化为带辅助掩码的扩散前向与逆向过程，消除了传统 Diffusion 均质全量扰动对组合逻辑结构的破坏。
无监督玻尔兹曼采样的通用泛化能力：摆脱了对 GNN 架构或预先标注的最优解训练集的依赖，借助 Transformer 强大的全局容量与统一松弛能量定义，直接在包含多值变量（九宫格分类、着色）和稠密高阶约束（AllDifferent）的广义 CSP/COP 任务上展现出媲美或超越传统图专用模型的能力。
跨领域的通用设计迁移启发：该架构的“掩码冻结 + 条件重采样 + 边缘置信度导向选择”机制，可以直接迁移到大语言模型的离散代码错误修复、逻辑推理链的回溯重试或分子图结构的局部基团替换优化中。

局限与展望¶

反向传播计算图展开受限：由于训练时需要对反向扩散链进行时间步展开以计算最终能量梯度，受显存限制训练通常仅能回传 5 个扩散步，未能充分利用超深扩散链的表达潜力。
连续 Logit 松弛的边际差距：本文是在连续松弛后的 logit 空间中执行高斯扩散，而非纯离散状态空间扩散（如 Masked Discrete Diffusion），在连续到离散的 argmax 投影阶段仍可能损失细微概率质量。
缺乏严格精确采样的理论保真度：虽然最小化逆向 KL 散度能生成高质量样本，但在有限步数下并不具备严格收敛至目标玻尔兹曼真实分布的理论保证。未来可在生成末端引入 Metropolis-Hastings 接受-拒绝修正以实现严格的 MCMC 保真采样。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将块吉布斯思想引入连续扩散 Transformer，建立了运筹 LNS 与扩散生成的严格理论纽带。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖数独、图着色、最大割、最大独立集四种不同拓扑与变量形态的基准，兼具分布外泛化与丰富消融。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严密自洽，问题切入直击痛点，图示与实验逻辑极其清晰。
价值: ⭐⭐⭐⭐⭐ 为通用离散约束满足与组合优化探索出一条超越 GNN 和手工局部搜索的崭新通用神经范式。