離散拡散のためのシンプレックス緩和

要旨

カテゴリカル生成のための離散拡散モデルは、中間状態空間とそれに対応する逆予測問題を決定する劣化カーネルによって定義される。我々は一様離散拡散を研究し、基礎となるカテゴリカル劣化過程を変更せずに、その訓練目的関数と逆遷移を拡張できるかを問う。我々はSimplaxを導入する。これは正確なディリクレ-カテゴリカル拡張であり、各劣化カテゴリカル状態を補助的なシンプレックス値変数と結合しつつ、元の一様拡散過程をそのカテゴリカル周辺分布として保持する。この拡張は、デノイザ入力を劣化カテゴリカル状態に保ったまま、扱いやすいラオ・ブラックウェル化された逆ブリッジ目的関数と、対応する確率的逆サンプラーをもたらす。実験的には、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.