用於離散擴散的單純形鬆弛

摘要

用於類別生成的離散擴散模型由一個損壞核所定義，該核決定了中間狀態空間及相關的反向預測問題。我們研究均勻離散擴散，並探討其訓練目標與反向轉換能否在不改變底層類別損壞過程的前提下得以增強。我們提出Simplax，這是一種精確的Dirichlet-類別增強方法，將每個受損的類別狀態與一個輔助的單純形值變量耦合，同時保留原始均勻擴散過程作為其類別邊緣分佈。該增強方法產生了易於處理的Rao-Blackwellized反向橋接目標及相應的隨機反向採樣器，同時保留受損的類別狀態作為去噪器的輸入。實驗上，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.