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

2025/09/26 by Yuanzhi Zhu, Zhu, Yuanzhi, Xi Wang +5 · 2 citations
Computer Science · Engineering · #Advanced Vision and Imaging #Artificial Intelligence (cs.AI) #Computer Vision and Pattern Recognition (cs.CV) #FOS: Computer and information sciences #Image Processing Techniques and Applications #Machine Learning (cs.LG) #Medical Image Segmentation Techniques

paper · pdf · doi:10.48550/arxiv.2509.22925

openalex publication_date 2025/09/26 · openalex created_date 2025/10/19 · openalex updated_date 2026/07/28

Abstract

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.

Citations

• MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control
• ReDi: Rectified Discrete Flow
• DC-AR: Efficient Masked Autoregressive Image Generation with Deep Compression Hybrid Tokenizer
• The Diffusion Duality
• Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking
• Muddit: Liberating Generation Beyond Text-to-Image with a Unified Discrete Diffusion Model
• Uni-Instruct: One-step Diffusion Model through Unified Diffusion Divergence Instruction
• FUDOKI: Discrete Flow-based Unified Understanding and Generation via Kinetic-Optimal Velocities
• Dimple: Discrete Diffusion Multimodal Large Language Model with Parallel Decoding
• LaViDa: A Large Diffusion Language Model for Multimodal Understanding
• LLaDA-V: Large Language Diffusion Models with Visual Instruction Tuning
• MMaDA: Multimodal Large Diffusion Language Models
• Mean Flows for One-step Generative Modeling
• Few-Step Diffusion via Score identity Distillation
• CODA: Repurposing Continuous VAEs for Discrete Tokenization
• Halton Scheduler For Masked Generative Image Transformer
• Gumbel-Softmax Flow Matching with Straight-Through Guidance for Controllable Biological Sequence Generation
• Bridging Continuous and Discrete Tokens for Autoregressive Visual Generation
• Reward-Instruct: A Reward-Centric Approach to Fast Photo-Realistic Image Generation
• SANA-Sprint: One-Step Diffusion with Continuous-Time Consistency Distillation
• Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models
• Inductive Moment Matching
• Learning Few-Step Diffusion Models by Trajectory Distribution Matching
• Large Language Diffusion Models
• Theoretical Benefit and Limitation of Diffusion Language Model
• A Semi-Supervised Text Generation Framework Combining a Deep Transformer and a GAN
• Fast Solvers for Discrete Diffusion Models: Theory and Applications of High-Order Algorithms
• Improving Video Generation with Human Feedback
• Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps
• Diffusion Adversarial Post-Training for One-Step Video Generation
• Democratizing Text-to-Image Masked Generative Models with Compact Text-Aware One-Dimensional Tokens
• Dual Diffusion for Unified Image Generation and Understanding
• Distilled Decoding 1: One-step Sampling of Image Auto-regressive Models with Flow Matching
• Flow Matching with General Discrete Paths: A Kinetic-Optimal Perspective
• NitroFusion: High-Fidelity Single-Step Diffusion through Dynamic Adversarial Training
• Randomized Autoregressive Visual Generation
• Beyond Autoregression: Fast LLMs via Self-Distillation Through Time
• David and Goliath: Small One-step Model Beats Large Diffusion with Score Post-training
• Flow Generator Matching
• Scaling up Masked Diffusion Models on Text
• Diff-Instruct++: Training One-step Text-to-image Generator Model to Align with Human Preferences
• Scaling Diffusion Language Models via Adaptation from Autoregressive Models
• Adversarial Score identity Distillation: Rapidly Surpassing the Teacher in One Step
• DPLM-2: A Multimodal Diffusion Protein Language Model
• One Step Diffusion via Shortcut Models
• Distillation of Discrete Diffusion through Dimensional Correlations
• Jump Your Steps: Optimizing Sampling Schedule of Discrete Diffusion Models
• MaskBit: Embedding-free Image Generation via Bit Tokens
• SwiftBrush v2: Make Your One-step Diffusion Model Better Than Its Teacher
• Show-o: One Single Transformer to Unify Multimodal Understanding and Generation
• Discrete Flow Matching
• SlimFlow: Training Smaller One-Step Diffusion Models with Rectified Flow
• Autoregressive Image Generation without Vector Quantization
• An Image is Worth 32 Tokens for Reconstruction and Generation
• Simple and Effective Masked Diffusion Language Models
• Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation
• ReNO: Enhancing One-step Text-to-Image Models through Reward-based Noise Optimization
• Simplified and Generalized Masked Diffusion for Discrete Data
• Multistep Distillation of Diffusion Models via Moment Matching
• EM Distillation for One-step Diffusion Models
• Improved Distribution Matching Distillation for Fast Image Synthesis
• PeRFlow: Piecewise Rectified Flow as Universal Plug-and-Play Accelerator
• InitNO: Boosting Text-to-Image Diffusion Models via Initial Noise Optimization
• Score identity Distillation: Exponentially Fast Distillation of Pretrained Diffusion Models for One-Step Generation
• Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction
• Fast High-Resolution Image Synthesis with Latent Adversarial Diffusion Distillation
• Reward Guided Latent Consistency Distillation
• Scaling Rectified Flow Transformers for High-Resolution Image Synthesis
• Playground v2.5: Three Insights towards Enhancing Aesthetic Quality in Text-to-Image Generation
• SDXL-Lightning: Progressive Adversarial Diffusion Distillation
• A Dense Reward View on Aligning Text-to-Image Diffusion with Preference
• Dirichlet Flow Matching with Applications to DNA Sequence Design
• SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers
• aMUSEd: An Open MUSE Reproduction
• SwiftBrush: One-Step Text-to-Image Diffusion Model with Variational Score Distillation
• GIVT: Generative Infinite-Vocabulary Transformers
• One-step Diffusion with Distribution Matching Distillation
• Adversarial Diffusion Distillation
• Diffusion Model Alignment Using Direct Preference Optimization
• UFOGen: You Forward Once Large Scale Text-to-Image Generation via Diffusion GANs
• Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution
• A Pytorch Reproduction of Masked Generative Image Transformer
• GenEval: An Object-Focused Framework for Evaluating Text-to-Image Alignment
• Language Model Beats Diffusion -- Tokenizer is Key to Visual Generation
• Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference
• Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion
• InstaFlow: One Step is Enough for High-Quality Diffusion-Based Text-to-Image Generation
• SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis
• On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes
• Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis
• BOOT: Data-free Distillation of Denoising Diffusion Models with Bootstrapping
• Diff-Instruct: A Universal Approach for Transferring Knowledge From Pre-trained Diffusion Models
• ProlificDreamer: High-Fidelity and Diverse Text-to-3D Generation with Variational Score Distillation
• Training Diffusion Models with Reinforcement Learning
• Pick-a-Pic: An Open Dataset of User Preferences for Text-to-Image Generation
• Bridging Discrete and Backpropagation: Straight-Through and Beyond
• ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation
• Scaling up GANs for Text-to-Image Synthesis
• Consistency Models
• Muse: Text-To-Image Generation via Masked Generative Transformers
• Scalable Diffusion Models with Transformers
• Unified Discrete Diffusion for Simultaneous Vision-Language Generation
• On Distillation of Guided Diffusion Models
• DreamFusion: Text-to-3D using 2D Diffusion
• Rectified Flow: A Marginal Preserving Approach to Optimal Transport
• All are Worth Words: A ViT Backbone for Diffusion Models
• Classifier-Free Diffusion Guidance
• Diffusion-GAN: Training GANs with Diffusion
• Hierarchical Text-Conditional Image Generation with CLIP Latents
• Why Exposure Bias Matters: An Imitation Learning Perspective of Error Accumulation in Language Generation
• MaskGIT: Masked Generative Image Transformer
• StyleGAN-XL: Scaling StyleGAN to Large Diverse Datasets
• High-Resolution Image Synthesis with Latent Diffusion Models
• High-Resolution Image Synthesis with Latent Diffusion Models
• EdiBERT, a generative model for image editing
• Vector Quantized Diffusion Model for Text-to-Image Synthesis
• Structured Denoising Diffusion Models in Discrete State-Spaces
• Diffusion Models Beat GANs on Image Synthesis
• Gradient-based Adversarial Attacks against Text Transformers
• Learning Transferable Visual Models From Natural Language Supervision
• Knowledge Distillation in Iterative Generative Models for Improved Sampling Speed
• Taming Transformers for High-Resolution Image Synthesis
• Non-saturating GAN training as divergence minimization
• Reliable Fidelity and Diversity Metrics for Generative Models
• Generalization in Generation: A closer look at Exposure Bias
• Improved Precision and Recall Metric for Assessing Generative Models
• Large Scale GAN Training for High Fidelity Natural Image Synthesis
• MaskGAN: Better Text Generation via Filling in the______
• Neural Discrete Representation Learning
• Neural Machine Translation with Gumbel-Greedy Decoding
• Adversarial Feature Matching for Text Generation
• Maximum-Likelihood Augmented Discrete Generative Adversarial Networks
• Improved generator objectives for GANs
• Self-critical Sequence Training for Image Captioning
• GANS for Sequences of Discrete Elements with the Gumbel-softmax\n Distribution
• Categorical Reparameterization with Gumbel-Softmax
• SeqGAN: Sequence Generative Adversarial Nets with Policy Gradient
• An overview of gradient descent optimization algorithms
• Improved Techniques for Training GANs
• Microsoft COCO: Common Objects in Context
• Black Box Variational Inference
• A Reduction of Imitation Learning and Structured Prediction to No-Regret\n Online Learning
• Long Short-Term Memory
• Guided Score identity Distillation for Data-Free One-Step Text-to-Image Generation

Cited by

• One-step Diffusion Models with Bregman Density Ratio Matching
• Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization

Related

• Di\mathtt[M]O: Distilling Masked Diffusion Models into One-step Generator
• Nemotron-Labs-Diffusion-Image: Advancing Masked Discrete Diffusion for High-Resolution Image Synthesis
• Diffusion Models Beat GANs on Image Synthesis
• SSDD: Single-Step Diffusion Decoder for Efficient Image Tokenization
• Diffusion Guided Domain Adaptation of Image Generators