Discrete Stochastic Localization

4.1 Three-step overview

4.2 Why the noisy token is enough

The cancellation that makes DSL work

Under the localization channel, the clean-token posterior has the form

$$p(x\mid z,\gamma)\propto p(x)\exp\left(z\cdot x-\frac{\gamma}{2}\|x\|^2\right).$$

The only term that explicitly contains the nominal SNR $\gamma$ is

$$-\frac{\gamma}{2}\|x\|^2.$$

If all token embeddings lie on the unit sphere, then $\|x\|^2$ is constant across candidate tokens. After posterior normalization, this term cancels:

$$p(x\mid z)\propto p(x)\exp(z\cdot x).$$

So the Bayes-optimal posterior mean becomes

$$\hat{x}(z)=\frac{\mathbb E_{P(x)}[x\exp(x\cdot z)]}{\mathbb E_{P(x)}[\exp(x\cdot z)]}.$$

The denoiser no longer needs to know which nominal SNR generated $z$. It only needs the noisy token itself.

DSL needs both ingredients: stochastic localization makes $z$ a posterior natural coordinate, and unit-norm geometry makes the SNR-dependent norm term constant. Either ingredient alone is not enough.

4.3 Arbitrary per-token SNR paths

A global diffusion clock is no longer required

Once the denoiser is time-invariant, different token positions do not need to share the same SNR.

DSL can assign every position its own SNR $\gamma_i$:

$$z_i=\gamma_i x_i+\sqrt{\gamma_i}\epsilon_i.$$

Each $\gamma_i$ can follow its own path. Some positions may remain masked, some may become clean, and some may sit in finite-SNR soft-evidence states.

This turns generation into path design.

4.4 Exact NLL views motivate the training support

The training states are not chosen as a sampler heuristic

DSL’s training support comes from two exact likelihood views of the same per-token SNR formalism.

4.5 Training recipe

Mixed-SNR support, one cross-entropy objective

DSL trains on a mixture of two state families:

Both branches use one token-level cross-entropy objective:

$$\mathcal L_{\mathrm{CE}}=-\mathbb E_{(s,z)\sim q_{\mathrm{train}}}\sum_{i=1}^{L}\log p_\theta(s_i\mid z).$$

For the endpoint branch, this is the natural categorical likelihood. For the continuous branch, cross-entropy acts as a posterior-matching surrogate: if the model learns the true token posterior, its posterior mean recovers the Bayes denoiser in embedding space.

Implementation note. In practice, DSL samples a training example from one of two branches. The ROAR branch samples a revealed subset and assigns endpoint SNRs. The continuous branch samples per-token SNRs from a continuous distribution. Both are passed through the same converter and backbone.

4.6 Architecture

Posterior-view converter

The backbone does not receive raw noisy embeddings directly. Each noisy token is converted into a soft mixture over the token table:

$$q_i^{\mathrm{conv}}(v\mid z_i)=\mathrm{softmax}\left(\beta\langle z_i,x_v\rangle+b_v\right),\qquad m_i^{\mathrm{conv}}=\sum_v q_i^{\mathrm{conv}}(v\mid z_i)x_v.$$

The converter makes every noisy token look like a posterior object:

The converter scale $\beta$ and bias $b_v$ are important parameters: they control how similarity to the unit-norm token table is turned into a soft token distribution. This also gives the backbone a bounded, token-like input representation.

4.7 Backbone and timestep handling

A standard Transformer/DiT backbone, queried without informative time labels

The sequence of converter outputs is fed to a standard Transformer encoder or DiT-style denoiser. DSL is designed to remain compatible with masked-diffusion backbones, while adding finite-SNR token evidence through the converter and training support.

For compatibility with pretrained masked-diffusion DiTs, the timestep embedding module can remain in place, but DSL feeds a constant value rather than an informative timestep. This matches the theory: the noisy token state already encodes the relevant corruption regime.

The paper makes this precise: DSL keeps the timestep module for compatibility but feeds $t=0$, since $z_i$ and $m_i^{\mathrm{conv}}$ already encode the corruption state.

4.8 Inference

Choose a path, not a new model

At inference time, DSL keeps the posterior estimator fixed and changes only the path through SNR space.

All of these paths query the same clean-token posterior estimator.