Test-Time Anchoring for Discrete Diffusion Posterior Sampling

TL;DR

APS is a training-free posterior sampler for pretrained masked discrete-diffusion foundation models (MMaDA, Dream, VQ-Diffusion), aimed at the same applications that continuous-diffusion samplers like DPS / PSLD / RB-Modulation handle: linear and nonlinear inverse problems, reference-based stylization, and text-guided block inpainting. Two ideas: (i) quantized expectation — tilt the unconditional categorical distribution by the measurement likelihood, treat the tilted probabilities as coefficients over the codebook embeddings, take the expectation through the decoder via a straight-through estimator, and quantize back to a token; (ii) anchored remasking — when picking which tokens to keep at the next step, compute confidence under the tilted posterior instead of the unconditional prior, so structurally informative “anchor” tokens (often near-zero prior mass) survive instead of being overwritten by easy background tokens. The result is the first discrete-diffusion sampler that exceeds both prior discrete baselines (G2D2, SGDD, SVDD-PM) and competitive continuous baselines (DPS, PSLD) on FFHQ and ImageNet at 256² and 512², while scaling to 1024² in 15 steps where PSLD takes ~12 min.

Key claims

Standard confidence-based remasking in masked discrete diffusion has two failure modes for posterior sampling: sampling bias (the model commits to a spurious token under the unconditional prior, locking it in for all future steps) and independent token-wise confidence (informative foreground tokens remain masked because background tokens have higher per-token confidence) [§3.1, Fig. 6].
Quantized expectation replaces direct sampling with a straight-through-quantized expectation in the span of the codebook embeddings, treating tilted-posterior probabilities as coefficients; the resulting token can have zero prior probability mass yet strong measurement consistency, which the paper calls an “anchor token” [§3.1, Algorithm 1].
Anchored remasking computes per-token confidence under the tilted posterior x_φ(z_t) rather than the unconditional prior x_θ(z_t), so anchor tokens (low prior mass, high posterior mass) survive instead of being remasked in favor of high-prior background tokens [§3.2].
Theoretical contribution: Theorem 3.1 derives a training upper bound L_DDPS that integrates measurements into the reverse diffusion process; Theorem 3.2 derives a test-time bound L_APS that lets a lightweight network do the posterior tilt without retraining the (potentially 8B-parameter) foundation denoiser [§3, Theorems 3.1–3.2].
On ImageNet super-resolution (4×), APS reports up to 35.82% LPIPS gain and 10.94% PSNR gain over the prior discrete state-of-the-art (G2D2); on FFHQ SR(4×) the gains are 31.36% LPIPS and 7.05% PSNR [§4 Abstract; Table on SR/Gaussian deblurring].
APS matches or beats strong continuous baselines (DPS pixel-space, PSLD latent-space) on linear inverse problems (super-resolution, Gaussian deblurring, random inpainting, motion deblurring) and on nonlinear inverse problems (HDR, nonlinear blur) on FFHQ [§4 General Inverse Problems].
Reference-based stylization works as a posterior-sampling instance: APS on top of MMaDA conditions on a single reference style image plus a class prompt and outperforms continuous IP-Adapter and StyleAligned on ImageReward and CLIP-T despite a weaker generative prior [§4 Stylization].
Text-guided block inpainting at 1024×512: APS produces prompt-aligned garment completions where Imagen 3, FLUX, and HD-Painter visibly miss prompt details (untucked oversized T-shirt with vibrant colors, fully buttoned shirt) [§4 Text-Guided Block Inpainting, Fig.].
Question answering: APS with a lightweight Qwen3-0.6B reward model (joint sampling instead of Dream-7B-Instruct’s independent-marginal sampling) yields +21.99% relative gain on a Qwen3-8B held-out oracle judge [§4 QA, Table 4, Appendix C.1].
Scaling behavior: APS-L runs at 1024² in 15 steps / 484s on A100, where continuous methods need 500–1000 steps; PSLD takes ~738s at 512² and ~12 min for high-res — discrete diffusion’s “better test-time scaling” is the explicit thesis of the sampling-efficiency table [§4 Sampling Efficiency].
The method generalizes prior work in the same lineage: G2D2 (Murata et al., 2024) extends RB-Modulation’s proximal sampler to VQ-Diffusion with star-shaped noising and Gumbel-Softmax dequantization, but its gradient guidance depends on continuous relaxations that APS shows are unnecessary if the tilt is done directly in the discrete embedding span [§2 Related Work].

Method

APS works on a pretrained masked discrete-diffusion foundation model (MMaDA at 512², VQ-Diffusion at 256², Dream-7B for QA). At each reverse-diffusion step:

Compute the unconditional categorical distribution x_θ(z_t) over codebook tokens at each masked position from the frozen denoiser.
Quantized expectation: Tilt x_θ(z_t) by the measurement likelihood log q(y | x_φ(Z_t)) to obtain a posterior estimate x_φ(z_t). Treat the tilted probabilities as coefficients over the codebook basis vectors; pass the expectation through the decoder with a straight-through estimator; quantize the resulting embedding back to the nearest valid token. This produces an “anchor token” that minimizes the measurement reconstruction error and may have arbitrarily low mass under the prior.
Anchored remasking: Decide which tokens to keep for the next step using confidence under x_φ(z_t) instead of x_θ(z_t). Anchor tokens (low prior mass, high posterior mass) survive; high-prior background tokens that disagree with the measurement get remasked.
Iterate until all tokens are unmasked.

The training-free version (APS) keeps the foundation model’s weights frozen and only optimizes a lightweight tilt at test time, justified by Theorem 3.2. A retrained variant (L_DDPS-minimizing) also exists but is reserved for cases where retraining is feasible. The whole method is summarized in Algorithm 1 of the paper.

Results

ImageNet SR (4×): up to 35.82% LPIPS and 10.94% PSNR over G2D2; consistently better than SGDD and SVDD-PM on Tables 1–2.
FFHQ SR (4×): 31.36% LPIPS and 7.05% PSNR gain over the prior discrete SOTA.
General inverse problems (FFHQ): APS and APS-L recover high-fidelity images from severely degraded inputs across random inpainting, motion deblurring, HDR, and nonlinear blur, outperforming DPS and PSLD on linear tasks and competitive on nonlinear tasks.
Reference-based stylization: beats IP-Adapter and StyleAligned on ImageReward and CLIP-T using only the MMaDA prior + APS.
Text-guided block inpainting (1024×512): prompt-aligned completions where Imagen 3 / FLUX / HD-Painter fail on the prompt-specified details.
Question answering: +21.99% relative improvement over Dream-7B-Instruct on a Qwen3-8B-judged held-out benchmark, via a Qwen3-0.6B reward model that enables joint sampling.
Sampling efficiency: APS at 1024² in 15 steps / 484s on A100 vs PSLD at 512² / ~12 min — claims a Pareto-frontier shift on the cost-vs-resolution axis at the high end.

Why it’s interesting

Posterior sampling has been a continuous-diffusion story (DPS, ΠGDM, PSLD, ReSample, RB-Modulation) — APS is the first non-toy demonstration that masked discrete diffusion is the natural substrate for it, because the model’s own categorical distribution is the tilt target. This complements the wiki’s Diffusion Language Models thread, which has so far treated discrete diffusion as a serving/throughput story (Mercury, WeDLM, I-DLM) on the text side; APS demonstrates the same architectural choice pays off on the image side, on a different downstream task (posterior sampling vs throughput) and for a different reason (training-free guidance vs serving efficiency). The Litu Rout lineage also includes RB-Modulation on the continuous side, so this is the same author moving from continuous to discrete by the same proximal-sampler design. Adjacent on the wiki: the Classifier-Free Guidance variants cluster has been about training-time fixes to a runtime miscalibration (SoFlow, Ostris, CAFM); APS is a different framing — guidance comes from the measurement in posterior sampling, not from text — but reuses the discrete-diffusion “tilt the categorical” mechanism that doesn’t exist in continuous flow models.