We introduce Anchored Posterior Sampling (APS) for masked diffusion foundation models, built on two key innovations:

1. Quantized expectation — provides gradient-like guidance for discrete diffusion with a purely discrete embedding space.
2. Anchored remasking — enables adaptive decoding by preserving “anchor tokens” aligned with measurements.

APS supports a variety of linear and nonlinear inverse problems (super-resolution, deblurring, inpainting, HDR, nonlinear blur) as well as reference-guided stylization and text-guided editing.

APS achieves state-of-the-art performance among discrete samplers and remains competitive with continuous diffusion, while being more efficient at test time.

• [2026.02.03] Code for anchoring diffusion language model (Dream-7B-Instruct) released!
• [2025.10.02] Our paper is now on ArXiv!

We provide a standalone Jupyter notebook demonstrating APS on Dream-7B-Instruct:

APS_for_dLLM.ipynb - Complete implementation with:

• Reward-guided text generation using APS
• Baseline comparison (standard Dream diffusion)
• Visualization of reward progression
• Example prompts and outputs

# Create a conda environment
conda create -n aps python=3.10 -y
conda activate aps

# Install required packages
pip install torch transformers==4.57.5 tqdm matplotlib numpy
pip install accelerate

Note: Dream-7B requires transformers==4.57.5. Other versions may not compatible.

1. Open APS_for_dLLM.ipynb in Jupyter/Colab
2. Run all cells sequentially
3. Modify the prompt in cell "7.5 Baseline Generation" to test different inputs
4. Adjust APS hyperparameters in cell "2. Configuration":
   • T: Number of diffusion steps (default: 128)
   • M: Optimization steps per diffusion step (default: 3)
   • eta: Learning rate / guidance strength (default: 0.1)

• Quantized Expectation (QE): Enables gradient flow through discrete tokens
• Anchored Remasking: Preserves high-confidence tokens during generation
• Reward Model: Skywork-Reward-V2-Qwen3-0.6B (lightweight, efficient)

APS produces sharper textures and refined details compared to G2D2 and DPS.

APS generalizes to multiple tasks (motion blur, HDR, nonlinear blur) with large improvements in PSNR and LPIPS.

APS enables training-free stylization with a reference style image and prompt.

APS also supports text-guided block inpainting:

APS demonstrates better scaling than continuous diffusion samplers at high resolutions, achieving strong performance with only 15 steps at 1024×1024.

If you find this work useful, please cite:

@article{rout2025aps,
  title     = {Test-Time Anchoring for Discrete Diffusion Posterior Sampling},
  author    = {Rout, L. and Lugmayr, A. and Jafarian, Y. and Varadharajn, S. and Caramanis, C. and Shakkottai, S. and Shlizerman, I.},
  journal   = {arXiv preprint arXiv:2510.02291},
  year      = {2025},
  url       = {https://arxiv.org/abs/2510.02291}     
}