# Soft-Di[M]O: Improving One-Step Discrete Image Generation with Soft Embeddings

arXiv: https://arxiv.org/abs/2509.22925 · HTML: https://mlsanity.com/paper/2509.22925/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 144 of its 161 references that are in the map.
- First posted on arXiv 26 Sep 2025.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 144)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://mlsanity.com/paper/2204.06125/) (2022)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis](https://mlsanity.com/paper/2307.01952/) (2023)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [DreamFusion: Text-to-3D using 2D Diffusion](https://mlsanity.com/paper/2209.14988/) (2022)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Large Scale GAN Training for High Fidelity Natural Image Synthesis](https://mlsanity.com/paper/1809.11096/) (2018)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [ProlificDreamer: High-Fidelity and Diverse Text-to-3D Generation with Variational Score Distillation](https://mlsanity.com/paper/2305.16213/) (2023)
- [Reproducible scaling laws for contrastive language-image learning](https://mlsanity.com/paper/2212.07143/) (2022)
- [MaskGIT: Masked Generative Image Transformer](https://mlsanity.com/paper/2202.04200/) (2022)
- [ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2304.05977/) (2023)
- [An overview of gradient descent optimization algorithms](https://mlsanity.com/paper/1609.04747/) (2016)
- [Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction](https://mlsanity.com/paper/2404.02905/) (2024)
- [One-step Diffusion with Distribution Matching Distillation](https://mlsanity.com/paper/2311.18828/) (2023)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Pick-a-Pic: An Open Dataset of User Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2305.01569/) (2023)
- [GenEval: An Object-Focused Framework for Evaluating Text-to-Image Alignment](https://mlsanity.com/paper/2310.11513/) (2023)
- [Improved Precision and Recall Metric for Assessing Generative Models](https://mlsanity.com/paper/1904.06991/) (2019)
- [On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes](https://mlsanity.com/paper/2306.13649/) (2023)
- [Large Language Diffusion Models](https://mlsanity.com/paper/2502.09992/) (2025)

## Built on by (most cited first, 6 of 6)

- [Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization](https://mlsanity.com/paper/2510.08233/) (2025)
- [VERTIGO: Visual Preference Optimization for Cinematic Camera Trajectory Generation](https://mlsanity.com/paper/2604.02467/) (2026)
- [One-step Diffusion Models with Bregman Density Ratio Matching](https://mlsanity.com/paper/2510.16983/) (2025)
- [Discrete Diffusion Bridges for Spatiotemporally Aligned Image Translation and Generation](https://mlsanity.com/paper/2608.29997/) (2026)
- [IDRF: Inverse-Distilled Reward Fine-tuning of Masked Discrete Diffusion Models](https://mlsanity.com/paper/2610.03641/) (2026)
- [Omni-Diffusion-Distill: Few-Step Distillation of Unified Multimodal Diffusion Large Language Models](https://mlsanity.com/paper/2610.10990/) (2026)

## Abstract (arXiv)

One-step generators distilled from Masked Diffusion Models (MDMs) compress multiple sampling steps into a single forward pass, enabling efficient text and image synthesis. However, they suffer two key limitations: they inherit modeling bias from the teacher, and their discrete token outputs block gradient flow, preventing post-distillation refinements such as adversarial training, reward-based fine-tuning, and Test-Time Embedding Optimization (TTEO). In this work, we introduce soft embeddings, a simple relaxation that replaces discrete tokens with the expected embeddings under the generator's output distribution. Soft embeddings preserve representation fidelity for one-step discrete generator while providing a fully differentiable continuous surrogate that is compatible with teacher backbones and tokenizer decoders. Integrating soft embeddings into the Di[M]O distillation framework (denoted Soft-Di[M]O) makes one-step generators end-to-end trainable and enables straightforward application of GAN-based refinement, differentiable reward fine-tuning, and TTEO. Empirically, across multiple MDM teachers (e.g., MaskBit, MaskGen), Soft-Di[M]O achieves state-of-the-art one-step results: improved class-to-image performance, a one-step FID of 1.56 on ImageNet-256 with GAN-based refinement, along with higher GenEval and HPS scores on text-to-image with reward fine-tuning, and further gains from TTEO.