Discrete Stochastic Localization for Non-autoregressive Generation

TL;DR

This paper introduces Discrete Stochastic Localization (DSL), a novel framework for discrete sequence generation that improves upon masked discrete diffusion models by supporting flexible SNR paths and enhancing distributional faithfulness.

Contribution

The authors propose DSL, a continuous-state framework with invariant denoising, enabling a single trained model to support various SNR paths and improve discrete sequence generation.

Findings

Fine-tuning with DSL improves distributional faithfulness (MAUVE) on OpenWebText.

A single checkpoint supports multiple sampling methods, including autoregressive and hybrid approaches.

The method achieves effective sampling with as few as 48 steps without retraining.

Abstract

Continuous diffusion is a natural framework for non-autoregressive generation but has generally lagged behind masked discrete diffusion models (MDMs) on discrete sequence generation. We argue that the bottleneck is not continuity itself, but a representation in which denoising depends on timestep-indexed noise regimes. We introduce \emph{Discrete Stochastic Localization} (DSL), a continuous-state framework with unit-sphere token embeddings whose Bayes-optimal denoiser is invariant to the nominal signal-to-noise ratio (SNR) under the localization channel. One trained network then supports an entire family of per-token SNR paths, with endpoint masked-diffusion paths as a special case. Fine-tuning a pretrained MDLM checkpoint with DSL substantially improves distributional faithfulness (MAUVE) on OpenWebText across all step budgets from $T{=}128$ to $T{=}1024$, and the same checkpoint supports random-order autoregressive sampling, as well as a hybrid continuous-then-discrete sampler using as few as T=48 total steps -- without distillation or retraining.