离散扩散的单纯形松弛

摘要

用于类别生成的离散扩散模型由一个损坏核定义，该核决定了中间状态空间以及相应的反向预测问题。我们研究均匀离散扩散，并探讨能否在不改变底层类别损坏过程的前提下，丰富其训练目标和反向转移。我们提出 Simplax，一种精确的 Dirichlet-类别增强方法，它将每个被损坏的类别状态与一个辅助的单纯形值变量耦合，同时保持原始均匀扩散过程作为其类别边际。这种增强产生了一个可解的 Rao-Blackwell 化反向桥目标函数，以及相应的随机反向采样器，同时保留被损坏的类别状态作为去噪器输入。实验上，Simplax 在无条件 OpenWebText 生成上改善了生成困惑度-熵的权衡。在数独任务上，仅在 30 线索谜题上训练的模型在所有评估的线索密度中取得了对比方法中的最高准确率，包括最小唯一可解的 17 线索情形，并且在无条件生成中也实现了最高的有效性。

English

Discrete diffusion models for categorical generation are defined by a corruption kernel, which determines the intermediate state space and the associated reverse prediction problem. We study uniform discrete diffusion and ask whether its training objective and reverse transitions can be enriched without changing the underlying categorical corruption process. We introduce Simplax, an exact Dirichlet--categorical augmentation that couples each corrupted categorical state with an auxiliary simplex-valued variable while preserving the original uniform diffusion process as its categorical marginal. This augmentation yields a tractable Rao--Blackwellized reverse-bridge objective and a corresponding stochastic reverse sampler, while retaining the corrupted categorical state as the denoiser input. Empirically, Simplax improves the generative perplexity--entropy tradeoff on unconditional OpenWebText generation. On Sudoku, a model trained exclusively on 30-clue puzzles achieves the highest accuracy among the compared methods across all evaluated clue densities, including the minimum uniquely solvable 17-clue regime, and also achieves the highest validity in unconditional generation.