이산 확산을 위한 심플렉스 완화

초록

범주형 생성을 위한 이산 확산 모델은 오염 커널에 의해 정의되며, 이 커널은 중간 상태 공간과 관련된 역방향 예측 문제를 결정합니다. 우리는 균일 이산 확산을 연구하고, 기반이 되는 범주형 오염 과정을 변경하지 않고 학습 목적 함수와 역방향 전이를 확장할 수 있는지 묻습니다. 우리는 각 오염된 범주형 상태를 보조 심플렉스 값 변수와 결합하면서 원래의 균일 확산 과정을 범주형 주변 분포로 보존하는 정확한 디리클레-범주형 증강인 Simplax를 도입합니다. 이 증강은 노이즈 제거기 입력으로 오염된 범주형 상태를 유지하면서, 계산 가능한 라오-블랙웰 역방향 브리지 목적 함수와 그에 상응하는 확률적 역방향 샘플러를 제공합니다. 실험적으로, Simplax는 비조건부 OpenWebText 생성에서 생성 perplexity-엔트로피 트레이드오프를 개선합니다. 스도쿠에서는 오직 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.