Discrete Stochastic Localization for Non-autoregressive Generation

Yunshu Wu; Jiayi Cheng; Longxuan Yu; Partha Thakuria; Rob Brekelmans; Evangelos E. Papalexakis; Greg Ver Steeg

cs.LGFeb 18, 2026

Discrete Stochastic Localization for Non-autoregressive Generation

Authors: Yunshu Wu, Jiayi Cheng, Longxuan Yu, Partha Thakuria, Rob Brekelmans, Evangelos E. Papalexakis, Greg Ver Steeg

Organizations: University of California Riverside · New York University

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 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.

Figures & tables

Panel of Figure 1 — Figure 1: Discrete Stochastic Localization (DSL) dynamically “localizes” to sample discrete tokens. (Left) DSL dynamics on a cyclic toy sequence (ABCDE, BCDEA, …); details in Appendix D . (Right) State summary at one snapshot: each token occupies a different SNR regime, transitioning from masked to soft to hard decoding as SNR increases.

Panel of Figure 2 — Figure 2: The DSL posterior decomposes into direction and norm axes. (a) t-SNE projection of converter outputs for 25 probe tokens spanning five semantic classes. Trajectories sweep $\bm{z}_{i}=\gamma\,{\bm{e}}_{v}$ from [MASK] across $\gamma\in[10^{-2},80]$ and are colored by SNR. (b) Mean top-1 token recovery and target-token probability as $\gamma$ increases.

Method	$T=128$	$T=256$	$T=512$	$T=1024$
Data (reference)	1.000	1.000	1.000	1.000
AR (reference)	–	–	–	0.760
SEDD (absorb)	0.007	0.008	0.009	0.008
MDLM ‡	0.015	0.023	0.031	0.042
MDLM+FB ‡	0.064	0.086	0.103	0.133
MDLM+DFM ‡	0.041	0.098	0.168	0.254

Table 1: OWT (

L=1024

) unconditional generation. We report MAUVE

\uparrow

as the primary metric; the full table with MAUVE / GenPPL / SentEnt is deferred to Appendix F . Bold marks the best MAUVE in each column among non-reference rows. Rows marked

\ddagger

are taken from ( Wang et al., 2025 ) ; rows marked

\lx@sectionsign

from ( Rout et al., 2025 ) . All DSL-FT rows use the same DSL-finetuned checkpoint. DSL-FT + ROAR (naive) and DSL-FT + MDLM are non-refinement decoder ablations: tokens are committed once and never remasked. The ReMDM-family rows add iterative remasking refinement.

$T_{\rm cont}$	$T_{\rm MDM}$	NFE	$\sigma_{\rm switch}$	MAUVE $\uparrow$
16	16	32	0.49	0.501
16	32	48	0.49	0.662
16	48	64	0.49	0.702
32	16	48	0.46	0.587
32	32	64	0.46	0.589
32	48	80	0.46	0.724

Table 2: Few-step OWT generation with the DSL hybrid continuous-then-MDM sampler. All rows use the same DSL-finetuned checkpoint. For each

(T_{\rm cont},T_{\rm MDM})

, we report a balanced

\sigma_{\rm switch}

from a small handoff sweep. The full sweep with MAUVE / GenPPL / SentEnt is in Appendix F .

Method	BPC ( $\downarrow$ )
Continuous Diffusion
Plaid (MD4 impl.)	$\leq$ 1.48
DSL (Ours)	$\leq$ 1.45
Discrete Diffusion
Mult. Diffusion	$\leq$ 1.72
D3PM Uniform	$\leq$ 1.61

Table 3: Bits Per Character (BPC) on Text8 test set. Continuous diffusion entries (Plaid, DSL) are ELBO upper bounds; discrete diffusion and AR entries are exact NLL — direct numerical comparison across these categories is not meaningful

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Description
$i=1,\ldots,L$	I ndex for sequences of length $L$
$t\in[0,\infty)$	Index for continuous t ime dynamics, equal to SNR
$\gamma_{i}\in[0,\infty)$	Per token SNR, defined on a contour or as function of $t$
$s_{i}\in\mathcal{V}$	$i$ -th s ymbol/token from the vocabulary $\mathcal{V}$
${\bm{x}}_{i}=\enc(s_{i})\in\mathbb{R}^{d}$	E mbed token as vector on unit hyper-sphere surface ( ${\\|{\bm{x}}\\|}=1$ )
$\bm{z}_{i}\in\mathbb{R}^{d}$	Noisy embedding of the $i$ -th token at noise level $t$

Appendix

Table 4: Summary of notation and key relations.

Aspect	Prior continuous diffusion LMs	DSL
Intermediate states	Use continuous finite-SNR noisy embeddings.	Also uses continuous finite-SNR noisy embeddings.
Posterior input	The token posterior is naturally a function of the pair $({\bm{y}}_{\tau},\tau)$ .	The token posterior is a function of the realized state $\bm{z}$ alone.
Denoiser form	Typically learns a timestep-conditioned map $f_{\theta}({\bm{y}}_{\tau},\tau)$ .	Learns one time-agnostic map $f_{\theta}(\bm{z})$ .
Role of SNR / time	The nominal timestep helps the model interpret the noisy embedding.	SNR is absorbed into the posterior natural coordinate.
Interpretation across noise levels	Different noise levels behave like a family of timestep-indexed denoising tasks.	Different noise levels are regions of one posterior-state space.
Token geometry	Usually does not enforce the unit-sphere geometry needed to cancel token-dependent norm terms.	Unit-sphere token embeddings make the Bayes posterior SNR-invariant.

Appendix

Table 5: Prior continuous diffusion LMs vs. DSL. Both use continuous finite-SNR states. The difference is how those states are parameterized, interpreted, and exposed to the denoiser.

Step budget $T$	128	256	512	1024
$\eta_{\mathrm{cap}}(T)$	0.010	0.008	0.008	0.002
top- $p$ nucleus	0.9
Tokenizer / $L$	GPT-2 BPE / 1024
MAUVE embed / $K$	GPT-2 Large / 500
GenPPL oracle	GPT-2 Large
SNR schedule	linear, $\texttt{snr\_min}=10^{-3}$ , $\texttt{snr\_max}=100$ , $\rho=7$ , $\gamma=4$

Appendix

Table 6: Decoding hyperparameters for OWT experiments under our step-budgeted ReMDM-loop protocol. All settings are identical across

T

except

\eta_{\mathrm{cap}}(T)

$T$	$\eta_{\mathrm{cap}}(T)$	mean rewrites-per-token $\left(\mathbb{E}[R_{i}]\right)$
128	0.010	0.579834
256	0.008	0.924500
512	0.008	1.610600
1024	0.002	1.848940

Appendix

Table 7: Rewrite statistics for the step-budgeted ReMDM-loop protocol (computed from Eq. 64 , averaged over 5k samples).

$T$	step $k$	$t_{k}$	mean_max_prob $\uparrow$	nucleus size (top- $p{=}0.9$ ) $\downarrow$
128	0	1.0000	0.0398	10268.59
128	56	0.5625	0.9391	12.17
128	71	0.4453	0.9585	8.72
128	127	0.0078	0.9928	1.74
256	0	1.0000	0.0398	10268.59
256	113	0.5586	0.9455	8.72

Appendix

Table 8: Logged posterior sharpness checkpoints from the sampler logs (“ReMDM logits stats” and “ReMDM nucleus size”).

Panel of Figure 3 — Figure 3: DSL correction under masked and garbled inputs. The input contains both masked positions and visible-but-wrong tokens. DSL can assign zero SNR to masked tokens and small positive SNR to uncertain visible tokens, allowing the same refinement dynamics to both fill missing values and correct garbled ones.

Panel of Figure 4 — Figure 4: Sampling diagnostics under a fixed step budget. (a) Masking and reveal schedule. (b) Remasking intensity and realized rewrites per token. (c) Posterior sharpening measured by mean max-probability and top- $p$ nucleus size.

Panel of Figure 5 — Figure 5: Endpoint smoothing improves near-clean calibration. We compare atomic ROAR endpoints ( $\gamma\in\{0,\gamma_{\max}\}$ ) to smoothed endpoint ranges. Smoothing reduces ECE at large SNR and yields reliability curves closer to the diagonal near the clean limit.

Panel of Figure 6 — Figure 6: Endpoint smoothing improves the step–quality tradeoff under fixed decoding. Using the same ReMDM-style sampler with identical schedules, smoothed-endpoint checkpoints achieve stronger MAUVE across step budgets while maintaining comparable or better GenPPL.

Method	$T=128$	$T=256$	$T=512$	$T=1024$
Data (reference)	1.000 / 14.8 / 5.44	1.000 / 14.8 / 5.44	1.000 / 14.8 / 5.44	1.000 / 14.8 / 5.44
AR (reference)	–	–	–	0.760 / 12.1 / 5.22
SEDD (absorb)	0.007 / 119.2 / 5.65	0.008 / 112.9 / 5.63	0.009 / 107.8 / 5.62	0.008 / 104.7 / 5.62
MDLM ‡	0.015 / 61.5 / 5.52	0.023 / 55.8 / 5.49	0.031 / 53.0 / 5.48	0.042 / 51.3 / 5.46
MDLM+FB ‡	0.064 / 42.8 / 5.44	0.086 / 39.0 / 5.41	0.103 / 37.0 / 5.40	0.133 / 33.8 / 5.35
MDLM+DFM ‡	0.041 / 37.9 / 5.31	0.098 / 31.2 / 5.30	0.168 / 24.2 / 5.26	0.254 / 21.7 / 5.20

Appendix

Table 9: OWT (

L=1024

) unconditional generation under masked-refinement decoders. Each entry is MAUVE

\uparrow

/ GenPPL

\downarrow

/ SentEnt

\uparrow

; bold marks per-column best MAUVE among non-reference rows. Rows marked

\ddagger

are taken from [ 31 ] ; rows marked

\lx@sectionsign

from [ 21 ] . All other rows use our protocol. The two DSL rows use the same DSL-finetuned checkpoint with different ReMDM-family remasking schedules.

$T_{\rm cont}$	$T_{\rm MDM}$	NFE	$\sigma_{\rm switch}$	GenPPL $\downarrow$	MAUVE $\uparrow$	SentEnt $\uparrow$
16	16	32	0.49	60.3545	0.5008	5.0521
16	32	48	0.49	52.4820	0.6619	5.0664
16	48	64	0.49	49.0116	0.7017	5.0604
32	16	48	0.45	62.3669	0.5872	5.1147
32	32	64	0.46	53.7585	0.5891	5.1061
32	48	80	0.46	51.4659	0.7240	5.1033

Appendix

Table 10: Full diagnostics for the selected DSL hybrid continuous-then-MDM sampler configurations on OWT. All rows use the same DSL-finetuned checkpoint. The main text reports MAUVE only; here we include GenPPL and sentence entropy diagnostics for the selected handoff configurations.

Asset	License / terms of use
OpenWebText	Public research dataset used for open-ended text generation evaluation; we cite the original release and follow its terms.
Text8	Public benchmark dataset used for character-level likelihood evaluation.
LLaDA / MDLM checkpoints	Pretrained checkpoints used for initialization and/or baseline comparison; we cite the original releases and follow their model-use terms.
MAUVE	Evaluation metric for open-ended text generation; we cite the original paper/package and follow the package license.
GPT-2 tokenizer / evaluator	Used for tokenization and GenPPL evaluation; we cite the original release and follow the corresponding terms.

Appendix

Table 11: Existing assets used in this work.

Non-Autoregressive Text Generation

Non-autoregressive text generation produces multiple tokens in parallel to reduce decoding time while maintaining coherent, useful outputs.

See all 31 papers

Discrete Stochastic Localization for Non-autoregressive Generation
May 13, 2026Yunshu Wu, Jiayi Cheng, Longxuan Yu +4Non-Autoregressive Text Generation Masked Diffusion Models
University of California Riverside · New York University
DSL-LLaDA: Scaling Continuous Denoising to 8B Masked Diffusion LMs
May 31, 2026Longxuan Yu, Yunshu Wu, Yu Fu +5Masked Diffusion Models Diffusion Language Model Decoding
University of California, Riverside · Georgia Institute of Technology · Microsoft
Clock Diffusion: Efficient Semi-Autoregressive Continuous Diffusion Language Models
Oct 1, 2026Yair Schiff, Omer Belhasin, Roy Uziel +7Diffusion Language Model Inference Language Modeling
NVIDIA · Cornell University

May 13, 2026J/K move · Enter open · S save