Soft-Di[M]O: Improving One-Step Discrete Image Generation with Soft Embeddings
Authors: Yuanzhi Zhu
Xi WANG
Stéphane Lathuilière
Vicky Kalogeiton
Venue: ICLR 2026 Poster (2026)
Keywords: Diffusion Distillation; Discrete Diffusion; Mask Diffusion; One-step Generation
Primary area: generative models

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 while cause minimum bias.
Integrating soft embeddings into the Di[M]O \citep{zhu2025di} 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 \citep{weber2024maskbit}, MaskGen \citep{kim2025democratizing}), 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 than teacher GenEval \citep{ghosh2023geneval} and HPS \citep{wu2023human} scores on text-to-image with reward fine-tuning, and further gains from TTEO.

Published as a conference paper at ICLR 2026
SOFT-DIMO: IMPROVING ONE-STEP DISCRETE IM- AGE GENERATION WITH SOFT EMBEDDINGS
Yuanzhi Zhu!
Xi Wang'
Stéphane Lathuilière
Vicky Kalogeiton
LIX, École Polytechnique, CNRS, IPP
Inria, Univ. Grenoble Alpes, CNRS, LJK
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 model- ing 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, simple relaxation that replaces discrete tokens with the expected embeddings under the generator's output distribution. Soft embed- dings preserve representation fidelity for one-step discrete generator while pro- viding a fully differentiable continuous surrogate that is compatible with teacher backbones and tokenizer decoders while cause minimum bias. Integrating soft embeddings into the Di[M (Zhu et al., 2025b) distillation framework (denoted Soft-Di(M makes one-step generators end-to-end trainable and enables straight- forward application of GAN-based refinement, differentiable reward fine-tuning, and TTEO. Empirically, across multiple MDM teachers (e.g. MaskBit (Weber et al. 2024), MaskGen (Kim et al. 2025)), Soft-DiMM]O achieves state-of-the-art one-step results: improved class-to-image performance, a one-step FID of 1 1.56 on ImageNet-256 with GAN-based refinement, along with higher than teacher GenEval (Ghosh et al. (202)) and HPS (Wu et al. 2023) scores on text-to-image with reward fine-tuning, and further gains from TTEO.
INTRODUCTION
Masked Diffusion Models (MDMs) generate discrete data by iteratively replacing masked tokens with predicted unmasked ones. They have recently achieved strong results across diverse domains, including text generation (Austin et al. 2021
Lou et al., 2023 Nie et al. 2024
2025 Lou et al.
2024
Gong et al. 2024 Deschenaux & Gulcehre 2024] Sahoo et al. [0255] Shi et al.
al. 2025). image synthesis (Chang et al. 2022 Gu al. 2022 Chang al. 2023]
Hu & Ommer
2024
Bai et al.] 2024] Patil et al.) 2024
Issenhuth et al. [021 Kim et al.) 2025 Weber et al,
2024
and multi-modal understanding and generation (Hu al. 2022 Li et al. 2024d| Xie al,
2024a
Wang et al. 2024a Yang et al. 2025 You et al.) 2025 Yu et al. 2025 Li et al. 2025 Wang
2025a Shi al.] (025b)) Building this success, recent work has demonstrated that these mul step procedures can be distilled into one-step processes, substantially accelerating inference
al.
and reducing computational cost while maintaining competitive quality (Zhu et al. 2025b Sahoo
et al.
2025 Yoo et al. 2025).
Despite their remarkable performance, current one-step discrete generators face two fundamental limitations: First, they inherit both modeling and aligument errors from the teacher: imperfections in the teacher's generative modeling and the under-alignment of the pre-trained teacher prior to reward fine-tuning, both error types of teacher are transmitted to the student during data-free distil- lation (Zhou et al. 2024b Luo et al. 2023bb Yin et al. 2024b). Second, their discrete token outputs block gradient flow. preventing the application of various key refinement techniques such as GAN training (Yin et al. 2024a] Zhou et al. (024b), differentiable reward fune-tuning (Luo) 2024 Luo
et al. 2024), and Test-Time Embedding Optimization (TTEO) (Eyring al. 2024). As result,

---

Published as a conference paper at ICLR 2026
Figure 1: Qualitative results produced by our one-step generators distilled from MaskBit et al. 2024) (top) and MaskGen-L Kim et al. 2025) (bottom).
Weber
rent distilled MDM generators have limited flexibility to leverage external knowledge, or incorporate user preferences better align with downstream objectives.
The gradient-blocking problem is common across discrete modeling tasks. Prior remedies include policy-gradient estimators such as REINFORCE (Williams 1992, Sutton et al. 1999 Ranganath eet al., 2014) and differentiable relaxations (e.g.. Gumbel-Soltmax, soft-argmax) (Gumbel 1954
Jang et al. 2016; Zhang et al., 2017 Gu et al. 2018). However, these methods often introduce biased or high-variance gradients and numerical instabilities (Jang al. [016 Liu al. (023a),
and they have not been systematically applied to one-step discrete generators.
To address these limitations, introduce minimal-change, decoder- relaxation of dis- crete outputs, called soft embeddings, tailored to one-step discrete generators. We define soft embed- ding as the expected token embedding under the generator's output probability distribution, serving directly as a continuous surrogate for discrete outputs. This simple construction: (i) maintains repre- sentation fidelity, when the model is confident and its output distribution is sharply concentrated, the soft-embers boft-embedding collapses to the corresponding discrete embedding, thereby mitigating the high- variance issue and stablizing training: (ii) is fully differentiable, enabling end-to-end gradient propa- gation through the generator logits; (iii) offers compatibility with the existing teacher backbone and tokenizer decoder; and (iv) provide small bias and very low variance as true loss surrogate.
Soft embeddings enable continuous training of one-step discrete generators, thus unlock refinements techniques that were previously inaccessible: (i) adversarial refinement (GAN training): discrimina- tors operate in the encoded space while gradients backpropagate through the teacher backbone into the generator; (ii) differentiable reward fine-tuning: the one-step generator can be fine-tuned with the gradient from the differentiable reward functions: (iii) TTEO: once trained, the generator's input embeddings corresponding to Tinit are optimized at inference time with respect to reward signals to maximize seward-based metrics.
2

---

Published as a conference paper at ICLR 2026
In this work, we combine the proposed soft embeddings with the prior state-of-the-arr one-step generator Di yielding Soft-Di(MJO. The continuous nature of soft embeddings enables Soft- DilMO to seamlessly incorporate the three post-distillation refinement techniques discussed above. In summary, our contributions are: (i) we introduce soft emheddings as a new representation for the outputs of one-step generators; (ii) demonstrate that this representation enables end-to-end differentiability for methods previously inaccessible to discrete oulputs, including; GAN training, differentiable reward fine-tuning, and test-time embedding optimization: (iii) we extend Di Mi to additional MDM teachers and visual tokenizers, and empirically achieve state-of-the ofr-the-rtt-tttt one-step performance on both to-sat and text-to-image tasks.
PRELIMINARY: MASKED DIFFUSION MODELS AND MM
Masked Diffusion Models (MDMs). MDMs (Chang et al. 2022 Shi et al. 2025a Hu & Ommer 2024 generate discrete data by learning to reverse a discrete, absorbing -ittiiiigggtteeeeee state diffusion process. For
image modeling, a pre-trained tokenizer (e.g., VQ-VAE (Van Den Oord et al. 2017 Esser et al.) 2021 is used: the encoder maps image into a sequence of discrete tokens To - where is the sequence length, and the decoder reconstructs the image from this token sequence.
In the forward process, the original token sequence Zy is gradually corrupted into a fully masked sequence I1. At a timestep E [0, 1], each token is independently masked with probability Tr. yielding intermediate sequence 2, containing both image and mask tokens.
The reverse process iteratively refines the sequence by filling the mask tokens. A model 6 predicts categorical distribution p=(x)/ =x))= softmax (2/g (x:)) for each masked token position where = (ar)) is the predicted logits. Logits (xx) has shape of [L. V). where is the sequence length and V = [V] is the size of the vocabulary. Starting from a fully masked sequence I1, masked tokens are replaced by tokens sampled from p (.)). while unmasked tokens remain unchanged, progressively refining the sequence back to Ty over multiple steps.
MDM is trained to minimize the cross-entropy loss for reconstructing the original tokens from their masked versions:
LMDM
& + g o |- logPolj (ooee(/o//re are, +)]]
(1)
where / is uniformly sampled and the inner expectation is over the forward process. This objective enables efficient, parallel token prediction.
O Distillation. To address the sampling inefficiency of multi-step MDMs, Zhu et al. 2025b) proposed Di which distills a multi-step MDM into a one-step generator. Built on on on-policy
distillation (Agarwal et al. 2024 Ross et al.] 2011 Ho al., 2016 Balakrishna et al. 2020 Arora let al. 2022] Schmidt] 2019). , the framework operates by first sampling a one-step prediction from the student. This output is then corrupted by the forward masking process produce an intermediate state S. The objective is to match the teacher's and student's output distributions conditioned on Zt-
Two key designs enable Di M O to work effectively in practice: (i) an auxiliary model is introduced to approximate the training objective and its gradient, following strategies in prior work (Poole et al.) 2022 Huang et al. 2024; Zhou et al. (024c); and (ii) an strategy is adopted to prevent student from mode collapsing when starting from the low-entropy. fully masked state. With these components, we obtain the following approximation to the gradient of the Di loss (o)) which is used to update the student parameters 0:
token initialization
Exunt E ( ( E
des (3rmi)
VaLun 2
V. Dan (p (P.))))))
(2)
d0
where Tini ~ Pinit is sampled from a designed initialization distribution Pinit 2g represents the stu- dent's output logits, H, denotes the intermediate state obtained by applying the forward mask kernel leju to the sampled tokens Ia, and (-l-) denotes the chosen divergence measure (summed over all masked token positions) between the distributions from the teacher (ø) and the auxiliary model (25). Class condition and text condition c omitted for simplicity. The scalar function w(t) balances the contribution of different mask ratios to the overall objective. The auxiliary model is trained with Eq. (I) with Tg as clean tokens. Notably, this loss gradient is applied directly to the generated logits, avoiding the need to backpropagate through the non-differentiable sampling operation.
3

---

Published as a conference paper at ICLR 2026
Disc.
Real Login Fake
softmax
M
Sinii
Emba (3a)
GT
27
Teacher
Student 0
One-step Generator
Logits
Tokens
Masked iz
Layer
Mao
Auxiliary 1
Embreee (2a)
Emb. C9
Figure 2: So Pipeline. Given a sampled initialization Zimis the one-step generator (student 0) outputs logits ze; discrete image tokens
{= sampled from these logits and used to form the Di distillation loss path). In parallel, the logits are converted into soft embeddings (purple path); these dif- ferentiable soft embeddings permit various supervision for post-training, such as GAN loss and reward loss, allowing the student to compete with or even surpass the teacher's performance. The right panel illustrates the soft-embedding construction for each position i: logits softmax p' and Σ, p2 E2 yield éz. We p} pa( aimm for simplicity. More details discriminator is visualized in Fig.
METHODS
While one-step generators distilled from MDMs can be fast and effective, existing approaches (e.g,, Di[MOO (Zhu et al. 2025b)) implicitly assume that the teacher fully captures the target distribution. Consequently, the student model inherits any modeling errors present in the teacher. Addition- ally, since the student produces discrete tokens, gradients cannot propagate through its outputs, preventing the direct application of standard post-training techniques such as GAN objectives and differentiable reward fine-tuning. To overcome these limitations, we build the idea of continuous relaxations (e.g.. Gumbel-Softmai (Gumbell 1954 Jang et al. (016)) see Appendix E and introduce soft embedding. defined the expected token embedding under the generator's output probability (Sec. (Sec.3.1) 3.1 The soft embedding is compatible with the teacher backbone and tokenizer decoder while preserving end-to-end differentiability. Leveraging this representation, the distilled one-step student can incorporate various post-training and refinement techniques (Sec. (.2). including GAN training (Sec. 3.2.1), differentiable reward fine-tuning (Sec. (.2.22). and Test-Time Embedding Optimization (TTEO, Sec. 3.2.3). The overall training pipeline is illustrated in Fig.
3.1 SOFT EMBEDDINGS
3.1.1
INTRODUCTION OF SOFT EMBEDDING
For each masked position, the soft embedding is computed from the model's predicted logits at that position. Taking masked position i example (right panel of Fig. (2). we begin with the model's predicted logits zz. These logits are first converted into a probability distribution over the entire vocabulary via the softmax function: pe (zt (iiis) = softmax(2z)) This distribution pe represents the model's confidence for each possible token filling the masked position. Instead of conventional discrete token sampling, the soft embedding ep is computed as the expectation of token embeddings under this probability distribution (= (init). Concretely, let E E RIM xd denote the embedding matrix (from either the teacher backbone or the tokenizer decoder), and E5 be the d-dimensional embedding vector for token j in the vocabulary V. The soft embedding is then computed as:
|
Emb(2f)
Σ Po =[ Eg E. po (If()mm))
In other words, the operator Emb(-) here maps logits (via softmax) to a convex combination of all the token embeddings. We use Emb, and EmbDee to denote this operation using the embedding layer from teacher & and the tokenizer decoder, respectively.
3.1.2 BENEFITS OF SOFT EMBEDDING
This design has three practical benefits. First, representation fidelity: high representation fidelity essential for soft embeddings serve effective approximation of tokens sampled from log-
4

---

Published as a conference paper at ICLR 2026
its. Empirically, one-step generators distilled from MDMs (e.g.. Di Zhu et al. 2025b and ReDi (Yoo et al., 2025)) often concentrate probability mass on only few candidate tokens per po- sition because they transfer randomness into the initial sequence (see Tab. 14 for more validation). Consequently, the soft embedding representation typically closely matches the embedding of the sampled token, with only small approximation error between the continuous and discrete repre- sentations. We provide empirical verification for this representation fidelity in the Appendix [.2 Second, differentiability: the mapping from logits soft embeddings is fully differentiable, al- lowing gradients from any downstream objective (e.g., a discriminator or differentiable reward) to propagate back the generator. This enables end-to-end optimisation for adversarial refinement and reward-based fine-tuning without resorting to non-differentiable sampling or gradient estima- tors. Third, compatibility: the soft embedding ég lies the same embedding space used by the teacher or tokenizer decoder. As result, soft embeddings can be fed directly into these components in the generation pipline without additional projection layers.
3.1.3 THEORETICAL JUSTIFICATION OF SOFT EMBEDDINGS.
The soft-ebbedding operator admits a principled theoretical justification as a surrogate for discrete token sampling. Let the ideal objective be the expectation over possible discrete tokens:
Coig (0) Ej-po [f(e)]]
(4)
where f denotes the differentiable downstream objective (e.g.. teacher backbone + discrimina-
head, or tokenizer decoder reward, introduced later in Sec. 3.2). Since sampling is non- differentiable, we optimize the surrogate Loont (0) = f()), where é E' pe is the soft embedding introduced in Sec. B.1.1 The validity of this relaxation follows from two theoretical observations.
Second-orde bias under concentrated logits (Lemma G.3
Let n == == [(ej (e, - =
be the embedding covariance under po. Lemma G.3 shows that under standard smoothness assump- tions,
(5) Thus, the soft-embedding surrogate incurs only a second-order bias the embedding variance. For one-step generators distilled from MDMs. the output logits are highly concentrated (empirically
( - E,[((e) s -/,ET
validated previous work Di (Zhu et al., 2025b) and ReDi (Yoo et al. 2025) and our Fig.
and Tab. 14h If most probability is concentrated then >> OEs), making the bias term (s) and therefore negligible. This explains why the continuous embedding& é closely matches the embedding of the sampled token.
token j with top top -1 probability ps(j*) = 1 1,
Gumbel-Softmax Beraight-Throughe (ST) exhibits first-order bias (Lemma G.4). Lemma G.4 shows that the gradient estimator of Gumbel-Softmax ttraight satisfies
LB [[VoLsT] VoLoriell < C Clll@naardd ésoft |I| + O(r) + O(EE))).
(6)
The dominant term arises from the mismatch between the hard forward token and the soft backward relaxation, producing a first-order bias that does not vanish even in the limit - 0 unless the logits are exactly one-hot. Moreover, the injected Gumbel noise yields high gradient variance, which further destabilizes adversarial or reward-based optimization.
Soft embedding therefore provides an exceptionally stable surrogate with: (i) second-order bias that shrinks as logits concentrate (Lemma G.3), (ii) deterministic, variance-free pathwise gradients, and (iii) no forward/backwant mismatch, unlike Gumbel-ST (Lemma G.4). These theoretical properties explain why soft embeddings a good candidate to enable the refinement techniques described next.
3.2
REFINEMENT TECHNIQUES UNLOCKED WITH SOFT EMBEDDINGS
3.2.1
ADVERSARIAL, FINE-TUNING
The proposed soft embeddings allow to feed generator outputs directly into frozen, pre-trained teacher backbone and to backpropagate discriminator gradients to the generator logits. Building on this, design& discriminator and corresponding training protocol tailored to soft embeddings.
5

---

Published as a conference paper at ICLR 2026
Diseriminator architecture. We reuse the pre-trained teacher as a frozen multi-scale feature ex- tractor and add lightweight discriminator heads on top of the extracted features (Sauer et al. 2024a Chen al., 2025b). Since the soft embedding is computed using the teacher's 's embedding layer, the remaining backbone layers can directly process it as input. For each selected transformer layer, add a convolutional head to aggregate spatial information; the resulting outputs concate- nated along the channel dimension and passed through a final MLP to produce the discriminator logit. Masked-embeddie
input augmentation. As the teacher model is pre-trained on masked sequences, to keep the discriminator inputs close to the teacher's 's distribution and to increase robust- ness as in diffusion GAN (Wang et al. 2022 Lin al. 2024. Yin al. 2024a), we apply random masking to the sequence of soft embeddings fed into the backbone when computing the adversarial features. Specifically, similar to the diffusion GAN where the discriminator received Gaussian- noised input, we replace random number of the computed soft embeddings or the ground truth embeddings with mask embeddings to predict fake and real logits, respectively. In this work, we use (-b to denote the soft embedding operation with the embedding layer of teacher model ø This operator also maps one-hot tokens from ground truth image token sequence to the correspond- ing token embeddings. We then use Emgg((Fg)) and Emb (29)r to denote the masked embedding sequences with mask ratio r as input to the teacher backbone.
GAN objective. Let PrOAN be the distribution over mask ratios used for discriminator augmentation and Pinit the generator's 's initialization distribution. The GAN training objective in our method is:
min max maxBBxar parra PMPPYCAN [og D,,((n, (Embs (1) + Exia PoPru [o, (Emb, (=e)r, (arrr)))) r))],
Go Dn
(7)
where D,, is the full discriminator (frozen backbone trainable heads). Class condition and text condition is omitted for simplicity.
For stable generator updates ct al. 2016 Shannon et al. 2020) for the generator:
the non-saturating GAN loss (Goodfellow al. 2014
Poole
LGAN (0) Exza P [- log log(D,, (D, (Embg(zo)r,
(8)
Thanks to the differentiability of operator Embs, (-), gradients from D, with respect to its input can propagate through the masking operator and the embedding matrix back the generator logits zg. enabling effective end-to-end adversarial training.
3.2.2
DIFFERENTIABL REWARD FINE-TUNING
followino and aesthetic quality are the two primary criteria in text-to-image generation (Luo)
2024 Luo et al. 2024 Wang et al., 2025b). To directly optimize these criteria, the research com-
munity has increasingly adopted learned reward models (Wu et al., 2023] Radford et al.) [021]
2023) and human- prena-eeeeeeenceeeeee signals (Wallace et al., [0224 Liu et al.] 2025a Kirstain ct al. 20033 Yang et al. 2024), which supply scalable, task-aligned supervision that complements likelihood- training. However, since reward functions typically operate decoded images, di- rectly optimizing such rewards is impractical for discrete generators because gradients cannot pass through the tokenizer decoders. By using soft embeddings make model outputs differentiable with respect to logits, we can now fine-tune the one-step generator end-to-end on differentiable reward functions. Concretely, we fine-tune with set of differentiable rewards R (e.g., a CLIP score for pooptt-magee alignment (Radford et al., 202T) and ImageReward for aesthetics (Xu et al. 2023 ):
Black et al.
Lreward (0) Σ
where A; denotes the weighting for reward model R,, Dec(.) is the tokenizer decoder, Emb (-) is the corresponding embedding operator, 2g is the model output logits and e is the prompt condition.
By mapping generator logits to soft embeddings using the embedding layer from the tokenizer, can feed the resulting embedding sequence into the tokenizer's decoder and obtain a differentiable image. A reward score is then computed construct loss and backward to the generator.
The final generator loss is: Le (0) (0) + WGGN + Wreward Freward (0),
(10)
where WGAN and Wreward balance the contribution of the GAN and reward terms against the base O distillation loss.

---

Published as a conference paper at ICLR 2026
LCM-SDXL
SDXI-DMD2
Sot=DiMoo
Qualitative comparison of one- and few-step distillation methods by using the same text prompts. Di and Soft-Di(M|O are both from MaskGen-L teacher. The number of inference steps is indicated by "s'.
Figure 3:
3.2.3 TEST-TIME EMBEDDING OPTIMIZATION
The soft embedding also enables Test-Time Scaling (TTS) (Eyring et al. 2024 Guo et al. 2024 Ma
et al. 2025) by optimizing the embedding input of the one-step student. TTS for one-step genera- tor involves spending extra inference compute to search over different initial inputs Tinits including through gradient-based optimization using the tokenizer decoder and through sampling-basee meth- ods like Best-of-N (BoN). Optimizing the discrete initial code sequence can be difficult and inef- ficient for MDMs; instead, since the embedding layer is fixed, initialize from the embeddings corresponding to Tini and directly optimize it. This optimization (TTEO) can be expressed as:
arg max MeaaxR(DD Gin
(11)
where ein denotes the input embedding be optimized, c the prompt text condition, and z (ein) the logits generated by the one-step generator with input e We apply the same technique for the tokenizer decoder that the entire test-time optimization remains fully differentiable, enabling a ReNO-style (Eyring et al. 2024) test-time scaling implementation.
EXPERIMENTS
In Sec. 4.1
we summarize the experimental setup. We evaluate our method on two representative class-to-image generation and text-to-image generation. Following prior work (Yin et al.
tasks:
2024a Zhou et al. 2024b; Luo 2024 Li et al. 2024b). we pair the base distillation algorithm M O with different post-traiming training techniques depending the task. For class-to-image we add adversarial linettuning [Di O + GAN, see Sec. (42) because FID our primary evaluation metric measures distribution-level fidelity, and common reward-based objectives risk driving samples

---

Published as a conference paper at ICLR 2026
away from the class-conditional target. For text-to-image, we instead use differentiable reward- based fine-tuning (Di[Mjo + Reward, see Sec. (.3) to directly optimize prompt adherence and aes- thetics using CaiP-based (CLIP score) and learned aesthetic reward (ImagcReward). In addition, we show that applying TTS to the one-step text-to-image generator can further improve prompt adherence and aesthetic quality. In Sec. 4.4 we ablate key design choices and hyperparameters.
4.1 EXPERIMENTAL. SETUP
Teacher Models and Datasets. Together with MaskGit (Chang ct al.
2022 and Meissonic (Bai
et al. 2024 reported in Di[MOO, we also include MaskBit (Weber et al.
2024 and MaskGin (Kim
al. 2025 stronger teachers for class-to-image and text-to-image generation, respectively. For
both MaskGit and MaskBit, we use the ImageNet-256 (Dnng et al. 2009) dataset. For text-to-image generation, since both Meissonic and MaskGen are SeTde on the aesthetic dataset, we conduct the corresponding experiment with synthetic images from the DALLE3- IM dataset (Egan et al.) 2024) when using GAN training. And we use the prompts from AAION-Aesthticss
dataset Cherti al. (202)) for Di[MO loss following the original implementation. Since our one-step student inherits many components and features from the teacher model, we list the key configurations including vocabulary sizes and output logits dimensions in Tab. 5
Evaluation. The metries we use for comparing ImageNet results are Fréchet Inception Distance (FID) (Heusel et al. (017) and Inception Score (IS) (Salimans et al. 2016). We also use precision (Prec.) and recall (Rec.) (Kynkäänniemi et al., 2019). density (Den.) and coverage (Cov.) (Naeem) al. 2020) to further evaluate the fidelity and diversity of generated images. For text-to-image generation, we follow Zhu et al. (2025b) to measure the HPSv2 (Wu et al.) 2023) for aesthetic quality and GenEval (Ghosh et al.) 2003) and CLIP score for prompt following ability. We also measure the FID on COCO-30K (Lin et al, 2014) and MJHQ-30K (Li et al. 2024a), denoted as ImageNet-222 2566, we calculate the FID using 5k generated images for ablations and 50k generated for benchmarking (50 images per class), comparing the statistics of the generated images with the reference statisties from ADM (Dhariwal & Nichol (021) using Clean-FID (Parmar et al. (202b)) On top of the quantitative metrics, we also provide visual generations of different methods qualitative comparison.
FIDcoco and FIDMHQQ respectively. In our experiments
4.2 CLASS-TO-IMAGE GENERATION
MaskGit teacher. For MaskGit teacher, we initialize the generator from the provided one-step generator in Di and use the same teacher model as Di[M]O. In Tab. 2 we present a quantitative performance comparison of our Soft-Di M with various acceleration methods, including the high- order sampler 0-trapezoidal (Ren et al. 2025), other distillation techniques (Hayakawa et al. 2024 Yoo al. 2025), and the teacher model MaskGit (Chang et al., 2022p. This result, a FID of 6.4 using only one forward pass, demonstrates that our method effectively improves Di[M] yielding superior metric scores that surpass all prior approaches using the same teacher model. These comparisons highlights the effectiveness of our method in matching the performance of multi-step teacher models while dramatically reducing the computational cost for inference.
Benchmarking with MaskBit teacher. MaskBit is a powerful MDM built on Lookup-Free Quanti- zation (LFQ) (Yu et al. 2023) tokenizer, with relatively fewer parameters. With this strong teacher, we compare our method to leading deep generative models, including GANs, diffusion models, their distilled counterparts, and various variations, focusing on generation quality. We first reimplement DilM on the MaskBit teacher, and we find that this baseline show limited generation diversity, as reflected by the recall (see Tab. D. After adding GAN loss to the training, we achieve a FID of 1.98, lower than 2, which can be considerred sufficient good for ImageNet-256 generation. With a further longer training with slightly change of hyperparameters. we got a state-of-the-art one-step FID of 1. 1,5 comparable the teacher model and beat the other GAN, diffusion models by large margin, demonstrate the goodness of our method.
4.3 TEXT-TO-IMAGE GENERATION
Result with Meissoicc teacher. Since the Meissonic model is fine-tuned with human preference data, we found that the teacher model distribution is already aligned with the reward encouraged dis-
8

---

Published as a conference paper at ICLR 2026
Table Class conditional ImageNet-255 results with MaskBit teacher.
Table 2: Quantitative results on class-conditionll ImageNet-2566 with MaskGit teacher. i We reproduced Halton Sampler on ImageNet-256 with MaskGit.
Madhed
Skepe
ISit)
2h11)
Beaff)
ADMd
Deantel de Nicholl 2031
250
59/M
1.59
196.7
0.53
Methed
FID 35
Rec
(it Coe
2124
6:M0
040
Patko NiC
230
61534
2.27
2782
0.33
Mavida
Bonky
2114
0,96
040
0097
2004
200
615M
2,96
277.3
0.83
1499
Maxbda
Bonkk
10:15
1923
022
Lix 45
Sen
b004
506
51B
2.18
263.3
0.81
0.55
MaxkGa
Bonkr
2133
$1.35
0.16
0.12
Mockot
MAR Mi NING
100
4005
1:98
201
84
6.70
MArrVItion Wett
p007
X17M
1.74
119-1
0 Oxpooikl
1O0
7.10
960
2663
0.90
0.38
1.47
0.95
Few-dtap Toon Sreira
Sing
Hakmn
Moree cl
X75
2012
0:38
1.46
1225
240.6
0:31
1097
675M
20.3
Moaioo dagstal.200 2023
618
3.45
SO1T
Nofewas Odierir
301
KM7
36.0
0.88
0:41
1.43
6364
2.58
SO1T
Xedeeaa Odlecin
3004
90.41
14.0
0.31
0:13
0.
2136
6.70
2012
696
N154
6:57
2136
Dharde Dollakka
335M
7.58
252.5
0.54
R
Werd
101
7.50
775.0
3.1
0:46
1.3
0.93
Dijt
(100
6.11
214.1
0.53
305
7:96
2000
3.0
0:44
131
0.97
DapE O- Misddtii
The
20055
335M
2.89
301E
O&T
11.35
1820
9:44
0.99
Soats Exje Masicii
2051
1.06
28144
0.61
D.5
691
2140)
0.38
0097
20M
273-2
0.91
Saa-Dero
6.40
0.55
1.21
8.4
Curdii Oerah unbokinng
DP4po IAAN od: Sal PMPP)
6.00
8.2
0.21
8.2
(C.65)
FID (50)
8.00
FID (SK)
8.0
FID (5k)
%.0
6.90
7.8
7.6
10000
20000 30000 40000 Training Steps
50000
10000 20000 30000 40000 50000 Tmining Steps
10000 20000 30000 40000 50000 Thaining Steps
(a) Choice of discriminator input
(b) GAN help training
(c) Embedding mask ratio Prexx
Figure 4: Ablation studies on ImageNet-256 using FID as the evaluation metric with MaskBit teacher.
tribution. We simply use reward loss to fine-tuning the one-step Meissonic generator from Di[M]O. We observe in Tab. 3 that this straightforward reward-only fine-tuning produces consistent improve- ments in both HPS and GenEval scores.
Result with MaskGen teacher. MaskGen is a strong MDM that uses semantic ID tokenizer Yu et al. 2024b) and is trained at 256 x 256 resolution. Although the provided MaskGen-L and Mask XL. differ in size, their performance is similar in our experiments, we MaskGen- I. as the teacher to train our one-step generator. We reimplemented DiMMO on the MaskGen-L. teacher to provide an initialization for reward fine-tuning. We observe that reward-only fine-tuning can induce reward hacking, producing over-saturated images (see Appendix [.10]: this motivates keeping the Di[MOO loss as a regularization to limit large divergence from the teacher distribution. By combining the reward loss with the DiMMOO loss we obtain higher GenEval and HPS scores than the teacher when sampling with 32 steps (see Tab. and more comparison with multi-step teacher in Appendix 1.12p. Visual comparisons appear in Fig.
TTS with MaskGen teacher. After the cewards-fne-tune one-step generator is obtained, we ap- ply TTEO with it. We use the SGD (Ruder, 2016) optimizer with a learning rate of 0.2 and no gradient clipping; other settings follow Eyring et al. (2024), e.g., we use the same combination of 4 rewards (CLIP score, ImageReward, HPS, and PickScore) as selection metric to pick the final samples. Our Soft- Soft-DDiMj DijM model further improves both GenEval and HPS V2.1 scores and substantially outperforms the multi-step step teacher (see Tab. B) Example outputs and qualitative comparisons can be found in Fig. 3 We provide corresponding ablation and cost-iy-quality analysis in Appendix|] I.8

---

Published as a conference paper at ICLR 2026
Table 3: Comparison of text-t--image generation methods across multiple metrics. Methods marked with are evaluated by running their publicly available checkpoints; all other numbers are taken from the corresponding papers as reported. f denotes the higher is better and & denotes the lower is better.
Mathal
Parans Fm
Geakvvll
HPS VOLL
Segle
Yhu
Coasting
Calors
Packiio
Coker All.
Oveaal
Andes
Concqet
Palating
Plate
Avrraged
0.92
0.27
B22
23.79
21.26
Ranack
0.94
0.56
049
B37
0.10
0.52
28.35
24:31
24.33
25.35
Podil
2.43
0.20
0.95
0.76
028
D35
0.13
0.58
12.44
33.36
33.86
27:48
0.31
8.2
033
B83
008
0.6
29.57
2003
0.43
33.44
0.313
8.97
0.78
B83
008
0.4
23.54
25.33
25.07
37.00
37.88
00221
19.02
25.00
20:36
ZAB
00342
0.99
022
D27
1412
28:27
30.91
25.45
SUXL-ONDDX
Thc
ZAB
01343
0.99
0..00
B39
28.21
30.39
20.09
2112
Medagee
4.4 ABLATIONS STUDY
In this section, we present ablation studies to validate several key design choice in our method. More ablations can be found in Appendices I.3
1.8
1.10
and
Choice of embedding representation. Because the distilled one-step generator yields highly con- centrated logits, using a straight-through
cuubel--softma estimator is not required and can intro- duce bias or high-variance errors in gradient estimation during training. As shown in Fig. 4a comparison between the Gumbel hard straight-through estimator and our soft-embedding method indicates that the soft embedding consistently improves FID throughout training.
Limitations of GAN-only distillation. Prior work suggests that GAN training alone can convert a
multi-step continuous diffusion model into a one-step generator (Xu et al. 2024 Lin et al. 2025
Sauer et al., 2024b a Chen et al., 2025a), but those approaches rely on sophisticated engineering tricks to work in practice. In our experiments, we first show that continued training with the Di M loss does not improve results (see Fig. 4b); conversely, using only the GAN loss (i.e., omitting the DilMO objective) fails to produce competitive generator and can even perform worse than the Oi baseline (see Fig. 4bb and Fig. I0). Moreover, when distilling from the MaskGen-L teacher with large initial mask ratio (e.g.. 0.95 as in our MaskGen setup), GAN-only distillation exhibits severe mode collapse: different prompts often yield identical outputs (see Fig. 13). Together, these findings indicate that Di[M O loss as a regularizer is crucial for our refinement algorithm.
Choice of GAN mask schedule. Analogous to denoising GANs, our mask-embedding GAN ob- jective not only leverages the teacher backbone's training input distribution but also acts form of data augmentation for the discriminator, helping to prevent discriminator overfitting. In Fig. 4c we evaluate several mask ranges, from no masking up range of [0,0.95]. The results show that using a GAN mask improves training, and that larger mask ranges yield the largest gains.
5 CONCLUSION
In this work, we advance visual MDM distillation by integrate additional post-trane training techniques to the O framework. Our approach tackles key challenges of discontinuity by proposing soft em- bedding. Through theoretical justification and extensive experiments across multiple tasks, includ- ing class-to-image generation and text-to-image generation, we show that with diverse pretrained models Soft-Di|Mj consistently surpasses both the base teacher MDMs and other distilled vari- ants. This work demonstrates the power of distribution matching methods on distillation for MDMs, contributes to the growing community of research exploring efficient generation of discrete data.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
This work was supported by Hi! PARIS and ANR/France 2030 program (ANR-23-IACI-00N)
WUR-22-CE23-0001 LNR-22-CE39-001E HilParis Grant and Chair, and was granted access the IDRIS High- Higg--Pefformacc Computing (HPC) resources under the allocation 2025-AD01 1014300R2 and 2025-AD01 1015894 made by GENCI. We sincerely thank Shilin Lu, Yuqing Wang and Yao Teng for their insightful discussions that contributed to this work. We are also grateful to Runlong Liao, Nizar Benbouchta, and Nicolas Dufour for their meticulous proofreading.
REFERENCES
Rishabh Agarwal, Nino Vieillard, Yongehao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem. On-policy distillation of language models: Learning from self- generated mistakes. In The Twelfth International Conference on Learning Representations, 2024.
Kushal Arora, Layla El Asri, Hareesh Bahuleyan, and Jackie Chi Kit Cheung. Why exposure bias matters: An imitation learning perspective of error accumulation in language generation. arXiv preprint arXiv:2204.01177. 2022.
Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Sub- ham Sekhar Sahoo, and Volodymyr Kuleshov. Block diffusion: Interpolating between autore- gressive and diffusion language models. arXiv preprint arXiv:2503.08 09577, 2025.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:1 191117911-99999 2021.
Jinbin Bai, Tian Ye, Wei Chow, Enxin Song. Qing- Guo Chen, Xiangtai Li, Zhen Dong, Lei Zhu, and Shuicheng Yan. Meissonic: Revitalizing masked generative transformers for efficient high- resolution text-to-image synthesis. arXiv preprint 1rXiv:2410.08261, 2024.
Ashwin Balakrishna, Brijen Thananjeyan, Jonathan Lee, Felix Li, Arsh Zahed, Joseph E Gonza- lez, and Ken Goldberg. On-policy robot imitation learning from a converging supervisor. In Conference on Robot Learning, pp. 21 PMLR, 2020.
Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 22669, 2023.
Heli Ben-Hamu, Itai Gat, Daniel Severo, Niklas Nolte, and Brian Karrer. Accelerated sampling from masked diffusion models via entropy bounded unmasking. arXiv preprint arXiv:2505.2488. 2025.
Victor Besnier and Mickael Chen. A pytorch reproduction masked generative image transformer. arXiv preprint arXiv:2310.14400 2023.
Victor Besnier, Mickael Chen, David Hurych, Eduardo Valle, and Matthieu Cord. Halton sched- uler for masked generative image transformer. Proceedings of the International Conference on Learning Representations (ICLR), 2025.
Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305. /3301, 2023.
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale GAN training for high fidelity 2018
natural image synthesis. arXiv preprint
Huiwen Chang, Han Zhang. Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11315- 11325, 2022.
Huiwen Chang. Han Zhang, Jarred Barber, AJ Maschinot, Jose Lezama, Lu Jiang, Ming Hsuan Yang, Kevin Murphy, William T Freeman, Michael Rubinstein, et al. Muse: Text-to-image gen- eration via masked generative transformers. arXiv preprint Xxiv:2301.007/00 2023.
11

---

Published as a conference paper at ICLR 2026
Tong Che, Yanran Li, Ruixiang Zhang. R Devon Hjelm, Wenjie Li, Yangqiu Song, and Yoshua Bengio. -ikeiioo- augmented discrete generative adversarial networks. arXiv preprint arXiv:1702.07983, 2017.
Dar-Yer Chen, Hmrishav Bandyopadhyay, Kai Zou, and Yi-Zhe Song. Nitrofusion: High-fidelity single-step diffusion through dynamic adversarial training. In Proceedings of the Computer Vision and Puttern Recognition Conference, pp. 7654-7663, 2025a.
Junsong Chen, Shuchen Xue, Yuyang Zhao, Jincheng Yu, Sayak Paul, Junyu Chen, Han Cai, Song Han, and Enze Xie. Sana-sprint: One-step diffusion with continuous-time consistency distillation. arXiv preprint arXiv:2503.09641 2025b.
Mehdi Cherti, Romain Beaumont, Ross Wightman, Mitchell Wortsman, Gabriel Ilharco, Cade Gor- don, Christoph Schuhmann, Ludwig Schmidt, and Jenia Jitsev. Reproducible scaling laws for contrastive language-iman image learning. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 2818-2829, 2023.
Trung Dao, Thuan Hoang Nguyen, Thanh Le, Due Vu, Khoi Nguyen, Cuong Pham, and Anh Tran. Swiftbrush v2: Make your one-step diffusion model better than its teacher. In European Confer- ence on Computer Vision, pp. 176-192. Springer, 2025.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. Icee, 2009.
Justin Deschenaux and Caglar Gulcehre. Beyond autoregression: Fast Ilms via self-distillation through time. arXiv preprint arXiv:2410.21035, 2024.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-878 8794, 2021.
Sander Dieleman. Guidance: cheat code for diffusion models, 2022. URL WKKLhtttps///beaanese
html
Ben Egan, Alex Redden, XWAVE, and SilentAntagonist. Dalle3 1 Million+ High Quality Captions, May 2024. URL http:////ggggggga.cccce/aaaaaaae
syrthetic- 1m dalle3 high
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recogni- tion, pp. 12873-12883, 2021.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning, 2024.
Luca Eyring, Shyamgopal Karthik, Karsten Roth, Alexey Dosovitskiy, and Zeynep Akata. Reno: Enhancing one-step text-to-image models through reward- based noise optimization. Advances in Neural Information Processing Systems, 37: 15447775111,,,, 2024.
William Fedus, Ian Goodfellow, and Andrew M Dai. Maskgan: better text generation via filling in the_. arXiv preprint 1rXiv:180..07736 2018.
Guhao Feng, Yihan Geng. Jian Guan, Wei Wu, Liwei Wang, and Di He. Theoretical benefit and limitation of diffusion language model. arXiv preprint arXiv:2502.09622. 2025.
Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One step diffusion via shortcut models. arXiv preprint trXiv:2410.12555 2024.
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky TQ Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching. Advances in Neural Information Processing Systems, 37: 2024.
12

---

Published as a conference paper at ICLR 2026
Zhengyang Geng, Mingyang Deng, Xingjian Bai, Zico Kolter, and Kaiming He. Mean flows for one-step generative modeling. arXiv preprint arXiv:2505.13447, 2025.
Dhruba Ghosh, Hannaneh Hajishirzi, and Ludwig Schmidt. Geneval: An object-focused framework for evaluating text-to-image alignment. Advances Neural Information Processing Systems, 36: 52132-52152, 2023.
Shansan Gong, Shivam Agarwal, Yizhe Zhang. Jiacheng Ye, Lin Zheng, Mukai Li, Chenxin An, Peilin Zhao, Wei Bi, Jiawei Han, et al. Scaling diffusion language models via adaptation from autoregressive models. arXiv preprint arXiv:2410. 17891, 2024.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
Jiatao Gu, Daniel Jiwoong Im, and Victor OK Li. Neural machine translation with gumbel-greede decoding. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
Jiatao Gu, Shuangfei Zhai, Yizhe Zhang, Lingjie Liu, and Joshua M Susskind. Boot: Data-free dis-
tillation of denoising diffusion models with bootstrapping. In ICMI. Workshop on Structured Probabilistic Inference {{&} Generative Modeling, 2023.
Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo. Vector quantized diffusion model for text-to-iman toi-tiggegeee synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10696-10706, 2022.
Emil Julius Gumbel. Statistical theory of extreme values and some practical upplications: a series of lectures, volume 33. US Government Printing Office, 1954.
Chuan Guo, Alexandre Sablayrolles, Hervé Jégou, and Douwe Kiela. Gradient-based adversarial attacks against text transformers. arXiv preprint arXiv:2104. 137333 2021.
Xiefan Guo, Jinlin Liu, Miaomiao Cui, Jiankai Li, Hongyu Yang, and Di Huang. Initno: Boosting text-to-image diffusion models via initial noise optimization. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9380-9389, 2024.
Satoshi Hayakawa, Yuhta Takida, Masaaki Imaizumi, Hiromi Wakaki, and Yuki Mitsufuji. Distil- lation of discrete diffusion through dimensional correlations. arXiv preprint arXiv:2410.08709, 2024.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho and Tim Salimans. arXiv:2207.12598. 2022.
Classifier-free
diffusion guidance.
arXiv preprint
Jonathan Ho, Jayesh Gupta, and Stefano Ermon. Model-free imitation learning with policy opti- mization. In International conference on machine learning, pp. 2760-2769. PMLR, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
Minghui Hu, Chuanxia Zheng, Heliang Zheng, Tat-Jen Cham, Chaoyue Wang, Zuopeng Yang, Dacheng Tao, and Ponnuthurai N Suganthan. Unified discrete diffusion for simultaneous vision- language generation. arXiv preprint arXiv:2211. 14842, 2022.
Vincent Tao Hu and Björn Ommer, [mask] is all you need. arXiv preprint arXiv:2412.0677 2024.
Zemin Huang. Zhengyang Geng, Weijian Luo, and Guo-jun Qi. Flow generator matching. arXiv preprint arXiv:2410.1933 19310, 2024.
Thibaut Issenhuth, Ugo Tanielian, Jérémie Mary, and David Picard. Edibert, a generative model for image editing. arXiv preprint 111.:2111.11264, 2021.
13

---

Published as a conference paper at ICLR 2026
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameteriization with gumbel-softnax. arXiv preprint arXiv:1611.01 144, 2016.
Minguk Kang, Jun-Jun--n Yan Zhu, Richard Zhang, Jaesik Park, Eli Shechtman, Sylvain Paris, and Taesung Park. Scaling up gans for text-to-image synthesis. In Proceedings of the IEEE/CVE conference computer vision and pattern recognition, pp. 20124-10134 2023.
Dongjun Kim, Chich- Hisin Lai, Wei- Hei-nn Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning proba- bility flow ode trajectory of diffusion. arXiv preprint (rXiv:23/0.0227)) 2023.
Dongwon Kim, Ju He, Qihang Yu, Chenglin Yang, Xiaohui Shen, Suha Kwak, and Liang-Chieh Chen. Democratizing text-to-image masked generative models with compact text-aware one- dimensional tokens. arXiv preprint arXiv:2501.07730, 2025.
Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, and Omer Levy. Pick- a-pic: An open dataset of user preferences for text-to-image generation. Advances neural information processing systems, 36:36652-36663, 2023.
Matt J Kusner and José Miguel Hernández-Lobate Gans for sequences of discrete elements with the gumbel-softmax distribution. arXiv preprint arXiv:1611.04033 2016.
Tuomas Kynkäääniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances in Neural Information Processing Systems, 32, 2019.
Daiging Li, Aleks Kamko, Ehsan Akhgari, Ali Sabet, Linmiao Xu, and Suhail Doshi. Playground v2. 5: Three insights towards enhancing aesthetic quality in text-to-image generation. arXiv preprint a0X.v: 17245, 2024a.
Jiachen Li, Weixi Feng, Wenhu Chen, and William Yang Wang. Reward guided latent consistency distillation. arXiv preprint arXiv:2403. 2024b..
Shufan Li, Konstantinos Kallidromitis, Hritik Bansal, Akash Gokul, Yusuke Kato, Kazuki Kozuka, Jason Kuen, Zhe Lin, Kai-Wei Chang, and Aditya Grover. Lavida: A large diffusion language model for multimodal understanding. arXiv preprint 17X. 2025.
Tianhong Li, Yonglong Tian, He Li, Mingyang Deng, and Kaiming He. Autoregressive image generation without vector quantization. Advances in Neural Information Processing Systems, 37: 56424-56445, 2024c.
Zijie Li, Henry Li, Yichun Shi, Amir Barati Farimani, Yuval Kluger, Linjie Yang, and Peng Wang. Dual diffusion for unified image generation and understanding. arXiv preprint arXiv:2501.00289 2024d.
Shanchuan Lin, Anran Wang, and Xiao Yang. Sdxl-lightning: Progressive adversarial diffusion
distillation. arXiv preprint arXiv:2402.1 13929, 2024.
Shanchuan Lin, Xin Xia, Yuxi Ren, Ceyuan Yang, Xuefeng Xiao, and Lu Jiang. Diffusion adversar- ial post-training for one-step video generation. arXiv preprint arXiv:2501.08311 2025.
Tsung-Yt Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C LLwrence Zitnick. Microsoft coco: Common objects context. In Computer vision-ECCV 2014: 13th European conference, zurich, Switzerland, September 6-12, 2014, pro- ceedings, part v 13, pp. 740-755. Springer, 2014.
Enshu Liu, Xuefei Ning. Yu Wang, and Zinan Lin. Distilled decoding 1: One-step sampling of image regr-segee models with flow matching. arXiv preprint arXiv:2412.17153 2024.
Jie Liu, Gongye Liu, Jiajun Liang, Ziyang Yuan, Xiaokun Liu, Mingwu Zheng, Xiele Wu, Qiulin Wang, Wenyu Qin, Menghan Xia, et al. Improving video generation with human feedback. arXiv preprint ar 25xiv:2501.139) 2025a.
14

---

Published as a conference paper at ICLR 2026
Liyuan Liu, Chengyu Dong, Xiaodong Liu, Bin Yu, and Jianfeng Gao. Bridging discrete and back- propagation: straight-througe and beyond. Advances in Neural Information Processing Systems, 36:12291-120 2023a.
Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209. 14577, 2022a.
Qiang Liu. On rectified flow and optimal coupling. preprint, 2022b.
Qiang Liu. Ieml tutorial on the blessing of flow. International conference on machine learning. 2025.
Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, and Qiang Liu. Instaflow: One step is enough for high-quality diffusion-based text-to-image generation. arXiv preprint arXiv:2309.06:: 06380, 2023b.
Zeyu Liu, Zanlin Ni, Yeguo Hua, Xin Deng, Xiao Ma, Cheng Zhong, and Gao Huang. Coda:
Repurposing continuous vaes for discrete tokenization. arXiv preprint arXiv:2503.17766, 2025b.
llya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arX1v:1711.05101,200 2017.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion language modeling by estimating the ratios of the data distribution. 2023.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. In Forty-first International Conference on Machine Learning, 2024.
Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02385 2021.
Simian Luo, Yiqin Tan, Longbo Huang, Jian Li, and Hang Zhao. Latent consistency models: Synthe- sizing high-resolution images with few-step inference. arXiv preprint 1xXiv22300004378. 2023a.
Weijian Luo. Diff-instruct++:: Training one-step text-to-image generator model to align with human preferences. arXiv preprint arXiv:2410. 18881, 2024.
Weijian Luo, Tianyang Hu, Shifeng Zhang, Jiacheng Sun, Zhenguo Li, and Zhihua Zhang. Diff- instruct: A universal approach for transferring knowledge from pre-trained diffusion models. Advances in Neural Information Processing Systems, 36:76525-76546, 2023b.
Weijian Luo, Colin Zhang, Debing Zhang, and Zhengyang Geng. David and goliath: Small one-step model beats large diffusion with score post-training. arXiv preprint arXiv:2410.20898, 2024.
Weijian Luo, Zemin Huang, Zhengyang Geng, J Zico Kolter, and Guo-jun Qi. One-step diffusion distillation through score implicit matching. Advances in Neural Information Processing Systems, 17:115377- 2025a.
Yihong Luo, Tianyang Hu, Weijian Luo, Kenji Kawaguchi, and Jing Tang. Reward- instruct: reward-centric approach fast photo-realistic image generation. arXiv preprint arXiv:2503.13070 13070, 2025b.
Yihong Luo, Tianyang Hu, Jiacheng Sun, Yujun Cai, and Jing Tang. Learning few-step diffusion models by trajectory distribution matching. arXiv preprint arXiv:2503.06674, 2025c.
Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Handen-Eijnden and Sain- ing Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. In European Conference on Computer Vision, pp. 23-40. Springer, 2024.
Nanye Ma, Shangyuan Tong, Haolin Jia, Hexiang Hu. Yu-Chuan Su, Mingda Zhang, Xuan Yang, Yandong Li, Tommi Jaakkola, Xubui Jia, et al. Inference-tine time scaling for diffusion models beyond scaling denoising steps. arXiv preprint 2iXiv:2501.0973 2025.
15

---

Published as a conference paper at ICLR 2026
Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 14297-14306, 2023.
Muhammad Ferjad Naeem, Seong Joon Oh, Youngjung Uh, Yunjey Choi, and Jaejun Yoo. Reliable fidelity and diversity metrics for generative models. In International conference machine learning, pp. 7176-7185. PMLR, 2020.
Thuan Hoang Nguyen and Anh Tran. Swiftbrush: One-step text-to-image diffusion model with variational score distillation. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 7807-7816, 2024.
Shen Nie, Fengqi Zhu, Chao Du, Tianyu Pang. Qian Liu, Guangtao Zeng, Min Lin, and Chongxuan Li. Scaling up masked diffusion models on text. arXiv preprint arXiv:2410. 18514, 2024.
Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. Large language diffusion models. arXiv preprint 2Sf2::2502.09992,, 2025.
Kong-Hyun Park, Chieh Chieh-Hsin Lai, Satoshi Hayakawa, Yuhta Takida, and Yuki Mitsufuji. Optimiz- ing sampling schedule of discrete diffusion models. CoRR, 2024.
Gaurav Parmar, Richard Zhang, and Jun-Yan Zhu. On aliased resizing and surprising subtleties in gan evaluation. In CVPR, 2022a.
Gaurav Parmar, Richard Zhang, and Jun-Yan Zhu. On aliased resizing and surprising subtleties in gan evaluation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11410- 110-114220 11420, 2022b.
Suraj Patil, William Berman, Robin Rombach, and Patrick von Platen. amused: An open muse reproduction. arXiv preprint arXiv:2401.01806 2024.
William Peebles and Saining Xie. Scalable diffuusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 4195-4205, 2023.
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhom, Jonas Müller, Joc Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.01992, 2023.
Ben Poole, Alexander A Alemi, Jaseha Sohl-Dickstein, and Anelia Angelova. Improved generator objectives for gans. arXiv preprint arXiv:1612.00777 2016.
Ben Poole, Ajay Jain, Jonathan T Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. arXiv preprint arXiv:2209.14900 14988, 2022.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference machine learning, pp- 8748-8763. PmLR, 2021.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint arXiv:2204.06125. 2022.
Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. In Artificial intelligence and statistics, pp. $12- PMLR, 2014.
Yinuo Ren, Haoxuan Chen, Yuchen Zhu, Wei Guo, Yongxin Chen, Grant M Rotskoff, Molei Tao, and Lexing Ying. Fast solvers for discrete diffusion models: Theory and applications of high- order algorithms. arXiv preprint vrXiv:2502.002:: 2025.
Steven Rennie, Etienne Marcheret, Youssef Mrouch, Jerret Ross, and Vaibhava Goel. Self-critical sequence training for image captioning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7008-7024, 2017.
16

---

Published as a conference paper at ICLR 2026
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVF Con- ference Computer Vision and Pattern Recognition, pp. 10684-10695, 2022.
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. reduction of imitation learning and struc- tured prediction to no-regret online learning. In Proceedings of the fourteenth international con- ference on artificial intelligence and statistics, pp. 627-635. JMLR Workshop and Conference Proceedings, 2011.
Sebastian Ruder. An overview of gradient descent optimization algorithms. arXiv preprint arXiv:1609.04747 2016.
Subham Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. Advances in Neural Information Processing Systems, 17:133013-130113 2024.
Subham Sekhar Sahoo, Justin Deschenaux, Aaron Gokaslan, Guanghan Wang. Justin Chiu, and Volodymyr Kuleshov. The diffusion duality. arXiv preprint arXiv:2506.1000 2025.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.005/2 2022.
Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
Tim Salimans, Thomas Mensink, Jonathan Heek, and Emiel Hoogeboom. Multistep distillation of diffusion models via moment matching, Advances in Neural Information Processing Systems, 37: (6046-30770, 2025.
Axel Sauer, Katja Schwarz, and Andreas Geiger. StyleGAN-XL: Scaling StyleGAN to large diverse datasets. In Special Interest Group on Computer Graphics and Interactive Techniques Conference Proceedings, pp. 1-10, 2022.
Axel Sauer, Frederic Boesel, Tim Dockhorn, Andreas Blattmann, Patrick Esser, and Robin Rom- bach. Fast high-resolution image synthesis with latent adversarial diffusion distillation. In SIG- GRAPH Asia 2024 Conference Papers, pp. 1-11, 2024a.
Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion dis- tillation. In European Conference on Computer Vision, pp. 87-103. Springer, 2024b.
Florian Schmidt. Generalization in generation: A closer look at exposure bias. arXiv preprint arXiv:1910.00292. 2019.
Matt Shannon, Ben Poole, Soroosh Mariooryad, Tom Bagby. Eric Battenberg, David Kao, Daisy Stanton, and RJ Skerry-Ryan. Non-saturating gan training as divergence minimization. arXiv preprint wtxiv:0000000029 2020.
Neta Shaul, Itai Gat, Marton Havasi, Daniel Severo, Anuroop Sriram, Peter Holderrieth, Brian Kar- rer, Yaron Lipman, and Ricky TQ Chen. Flow matching with general discrete paths: kinetic- optimal perspective. arXiv preprint arXiv:2412.03488) 2024.
Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis Titsias. Simplified and general- ized masked diffusion for discrete data. Advances in neural information processing systems, 37: 103131-103167, 2024.
Jiaxin Shi, Kehang Han, Zhe Wang, Arnand Doucet, and Michalis Titsias. Simplified and general- ized masked diffusion for discrete data. Advances in Neural Information Processing Systems, 37: 2025a.
Qingyu Shi, Jinbin Bai, Zhuoran Zhao, Wenhao Chai, Kaidong Yu, Jianzong Wu, Shuangyong Song, Yunhai Tong, Xiangtai Li. Xuelong Li, et al. Muddit: Liberating generation beyond text-to-image with unified discrete diffusion model. arXiv preprint arXiv:2505.224. 23600, 2025b.
17

---

Published as a conference paper at ICLR 2026
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. 2023.
Hannes Stark, Bowen Jing, Chenyu Wang, Gabriele Corso, Bonnie Berger, Regina Barzilay, and Tommi Jaakkola. Dirichlet flow matching with applications to dna sequence design. arXiv preprint arXiv:2402.0555 402.05841, 2024.
Peize Sun, Yi Jiang, Shoufa Chen, Shilong Zhang, Bingyue Peng, Ping Luo, and Zehuan Yuan. Autoregressive model beats diffusion: Llama for scalable image generation. arXiv preprint arXiv:2406.06522 06522, 2024.
Richard Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
Sophia Tang, Yinuo Zhang, Alexander Tong, and Pranam Chatterjee. Gumbel-softmax flow match- ing with straight-through guidance for controllable biological sequence generation. ArXiv, pp- arXiv-2503, 2025.
Keyu Tian, Yi Jiang, Zehuan Yuan, Bingyue Peng. and Liwei Wang. Visual autoregressive modeling: Scalable image generation via next-scale prediction. Advances in neural information processing systems, 37:84839-8448 2024.
Michael Tschannen, Cian Eastwood, and Fabian Mentzer. Givt: Generative vocaula-vocaaulary
transformers. In European Conference on Computer Vision, pp. 292-309. Springer, 2024.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Bram Wallace, Meihua Dang. Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong, Shafiq Joty, and Nikhil Naik. Diffusion model alignment using direct preference optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8228 8238, 2024.
Jin Wang, Yao Lai, Aoxue Li, Shifeng Zhang, Jiacheng Sun, Ning Kang, Chengyue Wu, Zhenguo Li, and Ping Luo. Fudoki: Discrete flow-based unified understanding and generation via kinetic- optimal velocities. arXiv preprint 2505.:2:00052.4477 2025a.
Shengquan Wang. A semi-supervised text generation framework combining deep transformer and a gan. arXiv preprint arXiv:2502.05937. 2025.
Xinyou Wang, Zaixiang Zheng, Fei Ye, Dongyu Xue, Shujian Huang, and Quanquan Gu. Dplm-2: A multimodal diffusion protein language model. arXiv preprint arXiv:2410.13788 2024a.
Yifei Wang, Weimin Bai, Colin Zhang, Debing Zhang, Weijian Luo, and He Sun. Uni-instruct: One-step diffusion model through unified diffusion divergence instruction. arXiv preprint arXiv:2505.20755. 2025b.
Yuqing Wang. Zhijie Lin, Yao Teng, Yuanzhi Zhu, Shuhuai Ren, Jiashi Feng, and Xihui Liu. Bridging continuous and discrete tokens for autoregressive visual generation. arXiv preprint 1643: 2025c.
Zhendong Wang, Huangjie Zheng, Pengcheng He, Weizhu Chen, and Mingyuan Zhou. Diffusion- gan: Training gans with diffusion. arXiv preprint arXiv:2206 2.V:22206.02262, 2022.
Zhengyi Wang, Cheng Lu, Yikai Wang, Fan Bao, Chongxuan Li, Hang Su, and Jun Zhu. Pro- lifiedreamer: High-fidelity and diverse text-to-3d generation with variational score distillation. Advances in Neural Information Processing Systems, 36, 2024b.
Mark Weber, Lijun Yu, Qihang Yu, Xueqing Deng, Xiaohui Shen, Daniel Cremers, and Liang- Chich Chen. Maskbit: Embedding image generation via bit tokens. arXiv preprint arXiv:2409.162) 2024.
Ronald J Williams.
Simple statistical learning. Machine learning, 8:229-256, 1992.
algorithms for connectionist reinforcement
18

---

Published as a conference paper at ICLR 2026
Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score v2: A solid benchmark for evaluating human preferences of text-to- image synthesis. arXiv preprint arXiv:2306.099: 2023.
Yecheng Wu, Junyu Chen, Zhuoyang Zhang, Enze Xie, Jincheng Yu, Junsong Chen, Jinyi Hu, Yao Lu, Song Han, and Han Cai. De-ar: Efficient masked autoregressive image generation with deep compression hybrid tokenizer. arXiv preprint urXiv:2507.04947 2025.
Jinheng Xie, Weijia Mao, Zechen Bai, David Junhao Zhang, Weihao Wang, Kevin Qinghong Lin,
Yuchao Gu, Zhijie Chen, Zhenheng Yang, and Mike Zheng Shou. Show-o: One single transformer to unify multimodal understanding and generation. arXiv preprint arXiv:2408. 12528, 2024a.
Sirui Xie, Zhisheng Xiao, Diederik P Kingma, Tingbo Hou, Ying Nian Wu, Kevin Patrick Murphy, Tim Salimans, Ben Poole, and Ruiqi Gao. Em distillation for one-step diffusion models. arXiv preprint ar 4rXvv:2405.16652, 2024b.
Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation. Advances in Neural Information Processing Systems, 36:1 15933-15935, 2023.
Yanwu Xu, Yang Zhao, Zhisheng Xiao, and Tingbo Hou. Ufogen: You forward once large scale text-to-image generation via diffusion gans. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8196-8206, 2024.
Hanshu Yan, Xingchao Liu, Jiachun Pan, Jun Hao Liew, Qiang Liu, and Jiashi Feng. Perflow: Piecewise rectified flow as universal plug-and-play accelerator. arXiv preprint urXiv:2405.07510, 2024.
Ling Yang. Ye Tian, Bowen Li, Xinchen Zhang, Ke Shen, Yunhai Tong. and Mengdi Wang. Mmada: Multimodal large diffusion language models. arXiv preprint arXiv:2505. 15809, 2025.
Shentao Yang. Tianqi Chen, and Mingyuan Zhou. A dense reward view on aligning text-to-image diffusion with preference. arXiv preprint arXiv:2402.08265 2024.
Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, and Bill Freeman. Improved distribution matching distillation for fast image synthesis. Advances in neural information processing systems, 37:47455-47487, 2024a.
Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Fredo Durand, William T Freeman, and Taesung Park. One-step diffusion with distribution matching distillation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6613-6623, 2024b.
Jachoon Yoo, Wonjung Kim, and Scunghoon Hong. Redi: Rectified discrete flow. arXiv preprint arXiv:2507.15897, 2025.
Zebin You, Shen Nie, Xiaolu Zhang, Jun Hu, Jun Zhou, Zhiwu Lu, Ji-Rong Wen, and Chongxuan Li. Llada-v: Large language diffusion models with visual instruction tuning. arXiv preprint arXiv:2505.16933 2025.
Lantao Yu, Weinan Zhang, Jun Wang, and Yong Yu. Seqgan: Sequence generative adversarial nets with policy gradient. In Proceedings of the AAAI conference on artificial intelligence, volume 31, 2017.
Lijun Yu, José Lezama, Nitesh B Gundavarapu, Luca Versari, Kihyuk Sohn, David Minnen, Yong Cheng, Vighnesh Birodkar, Agrim Gupta, Xiuye Gu, et al. Language model beats diffusion- tokenizer is key to visual generation. arXiv preprint arXiv:2310.0573:1 2023.
Qihang Yu, Ju He, Xueqing Deng, Xiaohui Shen, and Liang- Chhch Chen. Randomized autoregres- sive visual generation. arXiv preprint arXiv:2411.00777 2024a.
Qihang Yu, Mark Weber, Xueqing Deng, Xiaohui Shen, Daniel Cremers, and Liang Chigh Chen. An image is worth 32 tokens for reconstruction and generation. Advances in Neural Information Processing Systems, 177:128940-128968 2024b.
19

---

Published as a conference paper at ICLR 2026
Rumpeng Yu, Xinyin Ma, and Xinchao Wang. Dimple: Discrete diffusion multimodal large language model with parallel decoding. arXiv preprint arXiv:2505. 16990, 2025.
Yizhe Zhang, Zhe Gan, and Lawrence Carin. Generating text via adversarial training. In NIPS workshop Adversarial Training, volume 21, pp. 21-32. Academia. edu, 2016.
Yizhe Zhang, Zhe Gan, Kai Fan, Zhi Chen, Ricardo Henao, Dinghan Shen, and Lawrence Carin. Ad- versarial feature matching for text generation. In International conference on machine learning, pp. 4006 4015. PMLR, 2017.
Linqi Zhou, Stefano Ermon, and Jiaming Song. Inductive moment matching. arXiv preprint arXiv:2503.07565. 2025a.
Mingyuan Zhou, Zhendong Wang, Huangjie Zheng, and Hai Huang. Long and short guidance in score identity distillation for one-step text-to-image generation. arXiv preprint 2166:,2200.01661,, 2024a.
Mingyuan Zhou, Huangjie Zheng, Yi Gu, Zhendong Wang, and Hai Huang. Adversarial score identity distillation: Rapidly surpassing the teacher in one step. arXiv preprint 20.iv:2410.1499) 2024b.
Mingyuan Zhou, Huangjie Zheng. Zhendong Wang. Mingzhang Yin, and Hai Huang. Score identity distillation: Exponentially fast distillation of pretrained diffusion models for one-step generation. In Forty-first International Conference on Machine Learning, 2024c.
Mingyuan Zhou, Yi Gu, and Zhendong Wang. Few-step diffusion via score identity distillation. arXiv preprint arXiv:2505. 12674, 2025b.
Yuanzhi Zhu, Xingchao Liu, and Qiang Liu. Slimflow: Training smaller one-step diffusion models with rectified flow. In European Conference on Computer Vision, pp. 342-359, Springer, 2025a.
Yuanzhi Zhu, Xi Wang, Stéphane Lathuilière, and Vicky Kalogeiton. Di MMo: Distilling masked diffusion models into one-step generator. arXiv preprint (rXiv:2503.1545: 2025b.
Yuchen Zhu. Wei Guo, Jaemoo Choi, Guan-Horng Liu, Yongxin Chen, and Molei Tao. Mdns: Masked diffusion neural sampler via stochastic optimal control. arXiv preprint arXiv:2508. 10684, 2025c.
20

---

Published as a conference paper at ICLR 2026
Appendix for So
This appendix organized as follows:
Appendix A
Usage of large language models.
Appendix
Statement reproducibility.
Appendix
Broader impacts of this work.
Appendix
Limitation and future works.
Appendix
Discussion additional related works.
Appendix
Discussion on DLLM (Nie al. 2024
2025 Shi et 2024
Sahoo al.
2024).
Appendix
Theoretical justification of soft embedding.
Appendix
More experimental details.
Appendix
Additional experiments and corresponding findings.
Appendix
Failure cases of method.
Appendix
Additional visual results of one-step generations from our distilled models.
Appendix
List of all prompts used in this paper for text-to-image generation.
A USE OF LARGE LANGUAGE MODELS
We used large language models solely for text polishing and grammar correction during manuscript preparation. No LLMs were involved in the conception or design of the method, experiments, or analysis. All technical content, results, and conclusions have been independently verified and validated by the authors.
B REPRODUCIBILITY STATEMENT
We provide detailed training and evaluation protocols, including a full hyperparameter table, model and tokenizer specifications in the appendix. To enable reproduction and follow-up work, will publicly release the codes and model checkpoints (together with implementation notes and licensing information) upon publication.
C BROADER IMPACTS
Our work focuses distilling the multi-step generation process of MDMs into step, signif- icantly reducing inference time and computational costs, therefore lowering the carbon footprint during the inference. This advancement has the potential make high-quality generative models more accessible, facilitating applications in creative industries, content generation, and real-time systems. However, as with many generative modeling techniques, our method inherits biases from the teacher models. This could potentially lead to ethical concerns, including the generation of mis- leading or harmful content. Additionally, by enabling faster and more efficient content generation, our approach could lower the barrier to misuse, such as the creation of deepfakes or other deceptive media.
D LIMITATION AND FUTURE WORKS
Method Limitations. While our distillation method not limited by the teacher model, the teacher's performance still plays big role in the distillation, shown in text-to-image ex- periments. A stronger text-to-image teacher model can push the limit efficient discrete visual generation. In addition, because GAN operates in the latent space, the student's performance is bounded by the tokenizer's 's reconstruction FID (rFID), Notably, discrete tokenizers typically yield much higher rFID than continuous ones (Li al. 2024c). In the future, expect apply
21

---

Published as a conference paper at ICLR 2026
method to MDM trained with the adapted continuous tokenizer (w et al. 2025 Liu al. 2025b
Wang et al.) 2025c), such as TokenBridge (Wang et al.] (025c) to reach better generation perfor- mance. Besides, our model inherits the weakness of the teacher models, For example, like most MDM for image generation, our model can not generate images at varying resolutions.
Potential for Other Distillation Methods. Using soft embeddings not only enables gradient flow into the student's parameters for GAN training, it also makes a SiD-style distillation for MDMs feasible: both the teacher and auxiliary model can accept embedding inputs, allowing gradients to be backpropagated to update the student. Since Di is currently the only publicly available one- step discrete visual generative model (the ReDi model YYoo et al.) 2025) is not available), we leave fine-tuning one-step generators distilled with alternative distillation methods to future work.
Distillation for Multi-modal Discrete Model. While this work focuses on class-to-image and text-to-image generation, we believe that the proposed algorithm can also be applied to the visual
generation of multi-modal discrete diffusion models (Hu et al. [002]] Li et al., 2024d] Xie et al.) 2024a Wang et al.] [024a] Yang et al.) 2025; You et al. 2025 Yu et al., 2025 Li et al. 2025 Wang ct al. 2025a Shi et al. 2025b), such as show-o (Xie et al. 2024a) and Fudoki (Wang et al. 2025a).
Due to limited computational resources, we leave this direction to future work.
Discrete Diffusion beyond MDMs. For discrete diffusion schemes that employ a uniform (or uniform-like) kernel (Austin et al.) 2021 Gat et al., 2024; (Shaul et al.] 2024), our method natu- rally extends and avoids the out-of-distribution issues that can arise from the one-step generator's initialization. This stems from the fact that our distributional-mattoo
objective (including the orig- inal Di [M] loss) targets the generator's 's outputs regardless of its initialization distribution. We also acknowledge the related approach of Yoo al.] (2025), which first fine-tunes a masked diffusion model into a uniform-kernel discrete diffusion model and then applies ReFlow (Liu 2022a).
E RELATED WORK
Continuous Difflusion Distillation. Continuous diffusion models learn a Probability Flow ODE (PF-ODE) that connects the initial distribution to the target distribution. Sampling requires nu- merically solving this PF-ODE, often with many iterative steps. To accelerate generation, various
distillation techniques have been proposed, which can be broadly grouped into ODE-based and
distribution-based methods. ODE-based methods (Luhman & Luhman 2021: [ong et al., [023]
Liu 2022at Salimans & Ho 2022, Gu al., 2023 Meng al.] [2023] Yan et 2024 Zhu al.,
2025a) train a student model through regression objectives derived from the teacher's 's PF-ODE tra-
jectory. However, these approaches are difficult to apply to masked diffusion models (MDMs),
which lack an explicit PF-ODE formulation. Distribution-bast
methods (Luo et al.] 2023b Yin
let al. 2024b at Xu et al. 2024 Kim et al.] [023 Zhou et al.) [202c]|a /
Nguyen Tran 2024 Dao
eet al. 2025 Luo et al. 2025a Zhou et al. (025b) instead directly align the student's 's output dis-
tribution with the teacher's multi-step sampling distribution, typically via divergence minimization
or adversarial training (Sauer et al.) 2024b a). Variational Score Distillation (VSD) (Wang et al.
2024b] Yin et al. 22024b al Nguyen & Tran 2024; Dao et al. 2025| Luo et al.] 2023b] Xie et al.
2024b Salimans al. 2025 follows this paradigm by introducing an auxiliary model v that ap- proximates intermediate distributions of a one-step generator, serving as proxy to match the student and teacher.
Acceleration of Masked Diffusion Models. MDMs achieve strong generation performance, but their iterative nature makes inference both time- and compute-intensive. Consequently, improving MDM inference efficiency has become an active research topic. Prior work has approached this from several angles. Some methods focus on optimizing the sampling process: Park et al.) (2024) propose optimizing sampling schedules to reduce compounding decoding errors, while Ren et al. (2025) em- ploy higher-order solvers to allow larger sampling steps. Halton sampling (Besnier et al. 2025)) and entropy- -bouded sampling (Ben-Hamu et al. (023) have been proposed to accelerate generation by sampling multiple relatively independent tokens, leveraging spatial sparsity or entropy constraints. Other approaches explore distillation: [Haykkawa et al. (2024) adopt a strategy similar to Progres- sive Distillation (Salimans & Ho 2022), training a step student with a consistency loss that leverages intermediate states from real data, and Deschenaux & Gulcehre (024)) distill student to approximate the teacher's multi-step output distribution by minimizing divergence, which requires multiple teacher inferences and often multi-round distillation. Recently, (Zhu et al. 2025b) intro-
22

---

Published as a conference paper at ICLR 2026
duced Di[MOO, the first framework enabling effective one step generation for MDMs, significantly simplifying and accelerating inference. Finally, Yoo et al. (2025) propose fine-tuning an MDM with
stochastic initial states, followed by reflow stage with the simulated couplings (Liu 2022a b Liu
et al. 2023b). This idea of stochastic initialization is conceptually similar to that used in Di[MOO.
Adversarial Training for Discrete Modelling, Adversarial training for discrete data is challenging because the discriminator must compare generated sequences with ground truth, while the sampling process introduces non-differentiability Several strategies have been proposed to address this. A common approach is the REINFORCE estimator (Williams) 1992 Sutton et al.. 1999; Ranganath ct al. 2014), which treats generation
policy and the discrimmator score reward. It has been widely applied in discrete sequence modeling (Fedus et al. 2018; Rennie et al. 2017 Yu
et al. 2017 Che et al. 2017; Zhuu et al.] 2025c), but suffers from high-variance gradients that often destabilize training. Another line of work adopts differentiable relaxations of discrete tokens. The
Gumbel-Softmax trick (Gumbel| 1954: Jang et al. 2016 Gu al. 2018) generates a soft one-hot approximation: Punt softmax [[2g + (T))
9s ~ Gumbel 1), where T > 0 controls the smoothness. This method has been successfully applied in adversarial training for discrete data 2016) Guo et al.) [021 Wang. 2025 Tang et al., 2025). but suffers from biased gradients and potential numerical Instabilities (Jang et al. 2016 Liu et al., 2023a). Similarly, the soft-argmax operator (Zhang et al. 2016 2017) has been used to compute the soft embedding for LSTM- generators (Hochreiter & Schmidhuber 1997). and recently in discrete diffusion models (Stark al., 2024 Liu. 2025 Sahoo et al.) 2025 However, prior work has explored using soft embeddings for fast visual discrete generation.
Kusner& & Hernández-Lobar
F DISCUSSION ON DIFFUSION LARGE LANGUAGE MODELS (DLLM)
Most Diffusion Large Language Models (DLLMs) are instances of MDMs. Here, we discuss the challenges of directly applying Di(M O to language modelling for one-step generation with MDMs as an example. At first glance, one might expect that the only difference between image-based and text-based MDMs lies in the underlying data modality. However, this distinction turns out to be fun- damental for distillation. Visual Signals: Pseudo-Discrete with Redundancy. Visual signals are inherently continuous and highly redundant-humne perception remains unchanged even if many pixels are perturbed. This redundancy persists after discretization with tokenizers such as VQGAN. Even if some visual tokens are distorted, the decoded image is often perceptually similar to the orig- inal. Thus, visual discrete codes can be regarded as pseudo-discrete: although tokenized, they retain an underlying continuity. Language Tokens: Intrinsically Discrete. Language tokens, in contrast, intrinsically discrete. Every token choice has precise semantic and syntactic consequences, and small perturbations typically make a sentence ungrammatical or nonsensical.
This difference creates a fundamental challenge for MDM one-step generation. For existing discrete diffusion modelling, a one-step generator produces all logits for each token, then samples each of the tokens in parallel independently. To sample sentence correctly, the model's joint output distribution must collapse into an almost exact delta distribution over one sequence, which means the entropy for each prediction position must be as low as possible, otherwise the probability of generating valid sentence is essentially zero (Feng et al. (2025). This requirement for near- perrect precision the output distribution makes DLLM distillation significantly more challenging than visual diffusion model distillation.
G THEORETICAL JUSTIFICATION OF SOFT EMBEDDING
This section provides principled justification for using Soft Embedding as the surrogate for reward fine-tuning or adversarial training in one-step discrete generators. We establish (i) the ideal REIN- FORCE gradient reference, (ii) the bias and variance analysis of soft embedding and Gumbel- Softmax; and (iii why soft embedding is particularly suitable for one-step generators whose logits typically highly concentrated.
23

---

Published as a conference paper at ICLR 2026
G.1 PROBLEM SETUP
Consider the following ideal discrete objectivd
Lurig (0) = Ppj-p [f(e)
(12)
where is the predicted token distribution, e, is the embedding of token j. f denotes a determin- istic differentiable mapping (e.g.. teacher backbone + discriminator heads in adversarial training, or tokenizer decoder + reward function in reward fine-tuning).
The difficulty arises because sampling j ~ Pe is con-differentiable Common surrogates include:
1. REINFORCE, unbiased loss estimator but with high gradient variance.
2. Gumbel-Softmax (ST), biased and noisy.
3. Soft Embedding, smooth relaxation used in our method.
REINFORCE (ideal reference). Using the likelihood-ratio identity:
V VoLuni (0) Ey-pa [f(e) Vo log po]-
(13)
This estimator is unbiased, but for high-dimensional outputs (hundreds or thousands of tokens), the loss gradient variance is prohibitive.
Gumbel-Softmax Straight-Througge Gumbel-Softmax defines a soft relaxation:
Pn softmax ((=e + g)((T))
(14)
And the loss Lhard [[ ) uses hard sample j = arg max(fo) in the forward pass and soft sample in the backward pass to calculate the approximation dR
db deggrt:
(forward), Ehard = Earg max(fu)
êsoft = E (backward).
(15)
Soft Embedding (ours). Soft embedding replaces the discrete token by its embedding expectation:
= Ejrpo [ey] ET pe.
(16)
The surrogate objective becomes Lion (0) f(e). This yields fully pathwise gradient:
dzo
VoCuon (0) (E)(E) (diag(po) - p5 pp) po)
(17)
d0
G.2 BIAS ANALYSIS OF SOFT EMBEDDING AND GUMBEL-SOFTMAN
In this subsection, we provide formal analysis the bias introduced by Soft Embedding and
Gumbel-Softma when used as relaxations of the original objective in Eq. 112). We work under
standard smoothness assumptions and show that (i) the Soft-Embedding surrogate introduces only second-order bias in the embedding variance, while (ii) Gumbel--Softaa straight- trrighh-through produces a first-order biased estimator whose bias does not vanish even under concentrated logits.
G.2.1 ASSUMPTIONS
Assumption G.1 (Smoothness of f) The function f RRd - is twice continuously differentiable with Hessian bounded
[V for all e;,
(18)
Assumption G.2 (Embedding variance) Let se the expected embedding, define the embedding variance as:
: E= [ee #)(es - []
(19)
We denote its spectral norm by ||>||.
For simplicity, we consider only loss token since they sampled independently.
24

---

Published as a conference paper at ICLR 2026
G.2.2 SOFT EMBEDDING BIAS
We first show that the bias of Soft Embedding second-order the embedding variance.
Lemma G.3 (Second-order bias of Soft Embedding) Under Assumption G.I the soft-embedding surrogate satisfies:
T. |2|
(f()) - Ey-polf(ee) S
(20)
Thus the soft-embedding objective f(u) matches the true discrete objective to second-order term in | Furthermore, if the predicted distribution is concentrated on a mode j with
po(i)) =
Σ jØ1
Po(s)= =.
(21)
and if all embeddings satisfy lle,ll2 I B, then Eq. (20) is also (E))
|Σ||| O and thus the soft-embedding bias in
Proof:
(Second-order bias). Using the second-order Taylor expansion of f(e;) around
1
f(e,) = f(n) V 14) +
2 20, p1) 1)) V> - 14),
where Ey lies the line segment between ej and u. Taking expectation under po. the linear term vanishes sincellee, μm 0. Therefore,
B[f(e,)] f(e) + _ - 1) [ -w)] 2
By Assumption G.1
((e - V) - n1) S Lile, - 111?
Taking expectation:
L Elle, 1112 2 |||||
Bif(ee) - f(4) s
(Covariance under concentrated logits). Assume pa satisfies Eq. 21 and ||ey|l2 S B. Define the weighted average of non-max tokens
Σ j#1"
pa(ieej-
Then: (i) ||ll| < B and [|m]|2 < B (ii) [le, - Mll2 s 2B for all j (iii) since μ (1 E)e,- + Em,
|leg- /lI2 el|eg- ml12 < 2.B.
Now expand the covariance:
== (1- =)( -1- -00ee---llee
-1) + Σ
- m)(e; =
jØj
Using the bounds above,
(1 (1- -11/ < (1- E)
and
Σ poli)lle, m1l < =(2B)) jf1"
Hence,
|زJ|| < (1 -
O(E).
Substituting this into Eq.
20 proves that the soft-embedding bias is also O()).
25

---

Published as a conference paper at ICLR 2026
shows that the soft-embedding surrogate is intrinsically accurate when the model places most of its probability mass on a small number of candidate tokens. In the concentrated-logie regime (common for one-step generators distilled from diffusion models), the emhedding covariance scales as ||>|| = O(e), and therefore the bias between f(u) and the true dis- crete objective becomes negligible. This explains why soft emheddings provide faithful continuous approximation of discrete sampling in practice.
Remarks. Lemma G.3
G.2.3 BIAS OF GUMBRI--SFFMMAX STRAIGHT-
We now analyze the bias of the Gumbel-Softmax Straight-Througe (ST) estimator. Unlike soft em- bedding, Gumbel-Softmal ST exhibits firstits bias that persists even under concentrated logits.
Lemma G.4 (First-order bias of Gumbel-Softmat ST) Suppose Assumption G.1 holds, and in addition f LxpLipschitz on the convex hull of the embeddings. Let @nard (g) and Craft (g) denote the hard and soft embeddings produced by Gumbel-Softuax ST for given Gumbel noise g. Then the bias of the ST gradient estimator satisfies
VeLst V&Congll
2 C- [lénardd) ésoft Esoft + O(T) + O(1E2)l),
(22)
where C > is constant depending on L and sups ||0ésot/000) T is the Gumbel-Softme
lemperature, and Eg denotes the embedding covariance under Pg. Moreover whenever pe pa( (*) <
for the mode jr and fee-)) )Il 2 cu > 0, the bias term in Eq. (22 does not vanish as T - 0.
Proof: [Proof sketch] For fixed Gumbel noise g. the ST gradient can written as
Mesti (g)
gsT (0: g) = (chard (g))
00
where Jy is the Jacobian of f with respect to its input. Consider the "fully soft" reparameterized gradient
Isott (0:g) = J (Cron (g)) (g)
0e
Using the Lipschitz continuity of f and the chain rule, we obtain
Bg (0: g)] Eg [Sson((: [G]g)]]l < C Eg [[lénara(g) (n))
On the other hand, Eg [gaoft (0:g)] is the reparameterized gradient of a continuous relaxation objec- tive, whose bias with respect to the discrete objective Lorig is O(T) + O(IEE under Assump- tion G.I Combining the two bounds yields Eq.
Finally, in the concentrated regime where pa(j") ~ 1 but pa(i") < 1, the hard sample equals ej- with probability (j")), whereas the soft embedding stays strictly inside the convex hull. Conse- quently,
Ey [ (g) - ésoti (g)1l) Z (1 (o-po((i)) E,zr llle,- e,ll 0,
and if [[V )l 2 co this implies non-vanishing bias as +
Remarks. The key distinction is that Gumbel-Softmai ST incurs a first-order bias term propor- tional to the mismatch between forward and backward embeddings, | êgoftll. J. This bias: (i) does not vanish under temperature annealing when Po()*) 1; (ii) scales with Vfll: and (ii) persists even in the concentrated logit regime. In addition, the soft sample ésoft = F Po is a random variable whose expectation Ey [po] * po for finite r, creating additional bias due to Gumbel noise.
G.3 VARIANCE COMPARISON
In addition to the bias gap, the two methods exhibit fundamentally different variance behavior. Soft embedding produces fully deterministic forward pass and pathwise gradient Vof(E' (ET po), whose variance is identically zero given Zinit In contrast, Gumbel-Softmax introduces stochastic Gumbel noise in every forward pass and further mismatches the backward pass (soft) and forward pass (hard), yielding high-variance gradients of the form V vol(eg) where depends random perturbations.
26

---

Published as a conference paper at ICLR 2026
Discriminator
Real Logits
Lughs
Fint
aro
Student
Emba (#e)
One step
Logits
Teacher
Figure 5: Soft-Di(MjO Zoom details: the multi-scale discriminator in the pipeline Fig.
G.4 SUMMARY
Tab. summarizes the comparison. Soft-embedding yields (i) principled and variicrance surTo- gate objective, (ii) a bias that is provably second-order in embedding variance, and (iii) excellent stability in the highly concentrated logit regime characteristic of one-step generators. In contrast, Gumbel-Softmax straight-through combines significant bias with high variance, making it unstable for reward- or GAN-based fine-tuning.
Table 4: Comparison of gradient estimators for optimizing discrete generators.
Method
Blas
Variance
REINFORCE
None (unbiased)
Very high
Sombbe--offtmx
High
High
Soft Embedding (ours)
Small (2nd-order)
Very low
H MORE EXPERIMENTAL DETAILS
In this section, we provide further implementation and training details to facilitate reproducibility. The hyperparameters customized in this work and the model configurations are listed in Tab. 5 The complete algorithm with the input output dimension is presented Algorithm
H.1 IMPLEMENTATION
Goft-Di[M]O is implemented top of the O codebase. For MaskGit teacher models, the PyTorch implementation from (Besnier & Chen 2023), which relies on the original VQGAN
[Esser al., 2021); see discussion in (Sun et al. 2024 All teacher sampling procedures follow the
corresponding official implementations.
Unless otherwise specified, we use the following defaults:
Gradient clipping with norm 1.
Constant learning rate scheduler with linear warmup of 100 steps.
All model temperatures are fixed to during distillation (no top-p or top-k sampling).
The generator embedding layer is frozen during distillation improve stability.
Loss weight w(()) is used in Eq. 2 for simplicity.
H.2 DISTILLATION AND TEACHER SCHEDULES
For the Di loss we use the same mask schedule both for (i) obtaining the intermediate state d using generation with the one-step model and (ii) training the auxiliary model. Regarding the
27

---

Published as a conference paper at ICLR 2026
Table 5: Hyperparameterss and configurations used for distillation. Model architectures are the same as the teacher models. Largest Experiments in this work are conducted with 4 NVIDIA H100 GPUs, In the Output Logits Dimension, B is the batch size, V in MaskGit means mask token included, [B, L, 2, 128 in MaskBit represents the split group strategy where vocabulary
16384 2 is split into groups with 2 each.
Teacher Model
MaskGit
MaskBit (longer training)
Mcissonic
MaskGien
Batch Size
256
256
16
64
Training Steps
30k
50k(+50k)
10k
50k
Teacher CFG
2.0
1.5
4.0
4.0
Initial Mask Ratio
0.6
0.6
0.6
0.9
Di[MOO Divergence
Jeffreys -0.2
FKL
FKL
FKL
Fix Embedding Layer
Truc
True
Truc
Truc
Sampling Temperature
12
0.1
0.5
Generator Mask Schedule
arceos
arccos
cosine
arccos
Auxiliary Mask Schedule
arccos
cosine
GAN Loss Weight
100
25 (200)
100
GAN Mask Schedule
1400, 0.5]
0.95]
Discriminator Heads
4
Total #pprams in Discriminator
72M
85M
65M
Reward Loss Weight
2000
2000
CLIP Reward Relative Weight
0.2
0.2
Mixed Precision
bf16
bf16
bf16
bf16
Optimizer
Adam
Adam (AdamW)
Adam
Adam
Generator (81,85
(0.9,0.999)
(0.9.0.999)
(0.9,0.999)
(0.9,0.999)
Auxiliary Model (g),8g)
(0.9,0.999)
(0.9,0.999)
(0.9,0.999)
(0.9,0.999)
Discriminator (2)
(0.0.0.999)
(0.0,0.999)
(0.0,0.999)
(0.0,0.999)
Learning Rate
le-5
le-6 (5e-7)
le-6
1e-6
Weight Decay
0.0
0.0 (0.01)
0.0
FMA Ratio
0.9999
0,9999
0.9999
0.9999
Generation Image Resolution
256x226
256x256
1024x 1024
256x256
Sequence Length (I.)
256
256
4096
128
Downsampling Factor (Patch Size)
16x
16x16
16x16
512 (1D)
Vocabulary Size (V = IVD
1024
16384
8192
8192
Output Logits Dimension
L.
2,
L.
[,
divergence Dav in Di loss, follow Di for MaskGit and Meissonic, and Forward KL. (FKL) divergence for both MaskBit and MaskGen-L experiments. The same masking schedule is also used when training the teacher model. Both teacher models employ Classifier-Free Guidance (CFG) Ho & Salimans (2022). Since our generator distills the CFG output of the teacher models, its inference does not require CFG; this removes the extra sampling cost and improves sampling efficiency.
H.3 DISCRIMINATOR AND GAN TRAINING
When training the discriminator, we apply the same mask to both ground-truth embeddings and generated (fake) embeddings, and use the same prompts for real and fake samples. To balance gradient magnitudes, the GAN loss weight is set to 100 in most experiments so that GAN gradients are comparable to the Di[M gradients. For the longer runs, in MaskBit experiments. we increase this weight to 200. For MaskGit experiment, we use a mask ratio TGAN [0, 0.5). along with noise perturbation for the features from the teacher model backbone for stable training, which achieves slight better FID (FID 5k from 11.7 (only TGAN E [0, 0.95) as in MaskBit experiment) to 11.6). For the longer MaskBit training reported
Tab. [ we initialize the generator from the previous-stage best checkpoint. In addition, we halve the learning rate to 5e-7, increase the GAN loss weight to 200, and reduce the dropout rate to O. And we use an additional weight decay of 0.01 with AdamW V optimizer (Loshchilov & Hutter 2017). However,
find that either increasing the GAN loss weight or removing dropout (setting the dropout rate to 0) fails to improve the training of the one-step generator from the previous stage.
28

---

Published as a conference paper at ICLR 2026
Algorithm 1 Soft Distillation
Require: Pre-trained teacher model é, Pre-trained one-step generator 0o (optional), condition dataset De.
ground truth dataset Da. loss weights WGAN and Whewand
1: 0 <- copy Weightsss))))), -oopy (opyyeeigt))(B)) (optional), dh copy (opphesi))))))) // initialize
repeat
3: #*0 Generate logits 2e and tckens I6
4: Sample Zimi Pins. 6~ D. // with strategy in Di[M]O
5: Get generator logits 20 (TDini: c) E BBxxxxx11
6: 2o R sample po (20 (Tini) softmax(z2 c))
8 Update generator 8
Sample (0, 2up (Tiu,, c)) // Forward
9:
Calculate Po (zo (asl and Po (x0
10: # calculate Di[MJO loss
(D) -Bar (.nnttt==
((()) (Begol Da Dam(pollp)))) Emb((s) E' (2ol)amm)))
12:
Calculate soft embedding éØ
# calculate GAN loss
(0) - P F (Pm # ca alculate ate reward rd Toss
15:
16:
(rewna(0) x- E, (.e. (20),
17:
calculate total loss and update
Update using gradient of (0) = (O) + waNN + Wrewnnt Crewant (0)
19:
*** Update auxiliary modei v
20:
Sample t' ~ (0, 1]. žy (Z (r (Trini c))
21:
Update di with cross entropy loss (Eq. |
22:
888 Updato diseriminatos ")
23:
Sample [0, 0.95], calculate E," Emb(
24:
Sample real data (Ig, c) ~ Da, calculate Emb((g)a
25:
Update y with GAN objective (Eq.
26:
until convergence
27:
Return one-step generator
H.4 STRAIGHT-THROUGH GUMBEL IMPLEMENTATION
In frameworks such as PyTorch, setting ha hard=True in F. gumbel softmax implements a straight-through estimator that discretizes the forward pass using arg max while preserving tinuous gradients using Gumbel softmax in the backward pass Wang (025)..
H.5 ADDITIONAL EVALUATION DETAILS
For the evaluation of FIDcoco and CLIP score, we follow the exact evaluation setup as GigaGAN (Kang et al. 2023) and DMD2 Yin al. 2024a). More specifically, use 30K prompts from the COCO 2014 validation set to generate the corresponding images. The outputs are resized to 256x256 and compared with 40,504 real images from the same validation set using clean-FID (Par- mar et al. 2022a). We compute the CLIP score using the ViT-g-14-laion22 s vb2b b42k encoder. All MaskBit ablations are trained with batch size of 64 for 50k iterations. For ImageNet experiments, we evaluate final checkpoints across a sweep of generation temperatures and report the best metric for each checkpoint. For text-to-ing -octtto-mmage experiments, the generation temperature fixed at The FID, precision, recall, density and coverage in the imagenet experiments are all calculated with fea- tures extracted from the Inception V3 network. Note that the generation FID of our distilled MaskGit one-step generator outperforms that of the reconstruction FID of the tokenizer VQGAN. This is rea- sonable, as the reconstruction FID and generation FID for ImageNet are calculated against different reference data statistics, this is discussed in https://gituubcco//yyyeeaaa htapr//daaneeeooooooeeiiii/issse Yu et al. 2024a).
I ADDITIONAL EXPERIMENTS AND FINDINGS
In this section, we present additional experimental analysis and findings. In Appendix [.]] we vi- sualize and compare the tokenizers used in our work. In Appendices [.2 1.13 and 1.14 provide
29

---

Published as a conference paper at ICLR 2026
Figure 6: Comparison of reconstructions produced by different tokenizers. From top to bottom: (I) original image. (2) reconstruction using the MaskBit tokenizer, (3) reconstruction using the MaskGen TATTTok tok- enizer, and (4) reconstruction from Stable Diffusion v1.5 continuous VAE.
empirical evidence that the generator logits are highly concentrated. In Appendices I.3 I.4 and I.6 we show that our adversarial training algorithm robust a wide range of hyperparameters. In Ap- pendices I.5 and 1.7 we offer further comparisons between soft embedding and the Gumbel-Softmax Appendix [.8] we ablate the number of TTEO steps and analyze the
straight-througe estimator. corresponding cost cost- qusl-tyaityy tradeoff. In Appendix I.9 study the effect of different teacher CFG value for distillation. In Appendices [.10 and [.11 we report results obtained with individual re- ward terms and investigate the effect of different reward weights. In Appendíx 1.12, we present additional comparisons to multi-step teacher models under different evaluation setrings. Finally, in Appendix 1.15] we provided the result of soft embedding the Di OM] distillation framework.
1.1 RECONSTRUCTION ABILITY TOKENIZERS
We show reconstruction results produced by different tokenizers in Fig. 6 The figure compares the original image with reconstructions from the MaskBit tokenizer, the MaskGen TATTTok tokenizer, and the Stable Diffusion v1.5 continuous VAE. This comparison highlights the gap between the discrete image tokenizers and the continuous VAEs, illustrating the performance bottleneck of our method.
Figure 7: Comparison of decoded images produced from sampled discrete tokens (left) and from soft embed- dings (right).
30

---

Published as a conference paper at ICLR 2026
GAN Discriuninato Loss
GAN Generaorr Loss
9.0
Le0g-10.0
10
0.4
0.0
10000 20000
30000 40000 50000
10000
20000 30000 40000
50000
10000 20000
30000 40000 50000
Trining
Training
Training
(a) FIID
(h) Generator GAN loss
(c) Discriminator loss
Figure 8: Influence of different temperature in the calculation of sofumax
1.2 DECODED IMAGES FROM SAMPLED TOKENS vs. SOFT EMBEDDINGS
To validate the of soft embeddings for one-step generator, we compare images decoded from (i) sampled discrete tokens and (ii) the corresponding soft embeddings. Visually, the two sets of reconstructions are nearly identical (see Fig. 7 Quantitatively, over two batches of 64 images, we measure Mean Squared Error (MSE) of 6.02 x 10 between the sampled-token reconstructions and the soft-embedding reconstructions. For reference, the tokenizer reconstruction MSE is 3.84 10 10-6, These results support the practical usage of soft embeddings in our pipeline.
To further demonstrate the importance of representation fidelity, we vary the softmax temperature used to compute soft embeddings. In Fig. we show that both model performance and training sta- bility degrade as the temperature increases higher temperatures produce more diffuse probability distributions and thus larger approximation errors in the soft embeddings.
1.3 INFLUENCE OF GAN Loss WEIGHT
In Fig- 9 we conduct experiments with different GAN loss weights when applying our method to the MaskBit teacher and showcase that we need larger GAN loss weights to achieve lower FID. In our setting, the Di MJO loss weight for the generator is fixed to I, and no reward function is applied during the distillation. The loss weight for the auxiliary model and discriminator are both set to 1, as they are decoupled. We only need to balance the GAN loss with the Di[M]O loss for the generator. We also validate by comparing the loss gradients w.r.t logits 20 gradients for different GAN loss weights and observe that we need larger GAN loss weights to have comparable GAN loss gradients against Di MO loss gradient.
1.0
8.2
Gradion Norm Ratio
Mow!
Dizerimingtor Loss
(80)
8.0
Row inci
Stt
0.5
00.
191.
24
7.8
0.2
0.2
20000 Training Steps
Training Steps
Training Stepe
(a) FID
(b) Gradient Norm Ratio
(c)
Discriminator loss
Figure 9: Influence with different GAN loss weights WGAN. The GAN loss is not applied to the generator for the first 1000 steps.
1.4 GAN LOSS REDUCES DI'M SENSITIVITY TO HYPERPARAMETER
In addition to improving final performance, the GAN loss makes training more robust and easier to optimize. As shown in Fig. 10 when we initialize from a multistep MaskBit teacher and train with the Di loss (teacher CFG- 2) only, optimization rapidly converges to a local minimum. Introducing the GAN loss either alone or in combination with the Di[M]O objective-hhele model escape that local minimum and attain better FID. In practice, the best results obtained by combining both losses: the Di(M]O term preserves faithful distillation from the teacher, while the GAN loss increases robustness and improves final sample quality.
loss-either
the
31

---

Published as a conference paper at ICLR 2026
DIMIO CIC-d fer
Soft embedding
FID (5k)
GAN oely
Gumbel struight-rrough
GenEvel
0.48
12
0.4
8
10000 20000 30000 40000 50000
Traming Steps
10000
20000 30000 40000 50000 Training Steps
Figure 10: Ablation on the efg value and GAN training. The default CFG value for MaskBit is 1.5 (sec Tab. 5)
Figure 11: Ablation on the represen- tation used for reward fine- fine-tuning.
1.5 ADDITIONAL COMPARISON WITH GUMBEL STRAIGIT-TIRROGGII
In Fig. 12 provide comparison between soft-embedding representation to the Gumbel ttrought-through addition to Fig. 4a Training with soft embeddings noticeably more stable: the GAN generator loss is smoother and the real/fake discriminator logits show substantially less variance over the course of training.
Carble tbaghhrregh Brrigh
0.5
GIAN Genennee
1.5
Real Logits
1.5
Take Logits
1.0
1.0
0.5
-1.5
0.51
10000 20000 20000 40000 50600
10000 20000 30000 40000
50000
10000 2I0NO S00000 40000 50000
Training
Training
Triinin Steps
(a) Generator GAN loss
(b) Real logits
(c) Fake logits
Figure 12: Influence with different GAN loss weights WCAN The GAN loss is not applied to the generator for the first 1000 steps.
1.6 GAN ALONG TRAINING LEADS TO MODE COLLAPSE
In our MaskGen experiments, when we have a large initial mask ratio Fini 0.95, we observe mode collapse. As shown in Fig. [3]] for different prompts the generator will produce the same output.
Figure 13: Mode collapse with GAN alone distillation.
1.7 ABLATION SOFT EMBEDDING FOR MASKGEN-L TEACHER
We replicate the soft embedding ablation from Fig. 4a on the MaskGen-L. setup and present the training curves with GenEval as metric in Fig. This result shows that soft embeddings yield con- sistently higher GenEval scores during training, indicating improved prompt-following compared to using soft embeddings. fact, the cumbel--sfftma behaves similarly to the direct softmax used in soft-embedding approach when the injected Gumbel noise is small. However, the noise increases, the approximation enror the resulting soft embedding also grows and will cause a drop in performance.
32

---

Published as a conference paper at ICLR 2026
1.8 TEST- TIME EMBEDDING OPTIMIZATION (TTEO) WITH MULTI-START STRATEGY
To reduce computational cost while maintaining performance, we employ TTEO combined with multi-start Best-of-N (BoN) strategy, where multiple TTEO are executed in parallel for each prompt. In Fig. 14 we present the GenEval scores as function of both the number of BoN seeds N and the number of TTEO optimization iterations. Our results demonstrate that while the standard BoN baseline provides performance improvements, incorporating TTEO yields additional gains, thanks to soft embeddings. For all experiments reported in the main paper, we use N seeds and 4 TTEO iterations, which requires less inference time as pure BoN with N 64. We provide detailed runtime comparison in Tab. where the batch size is for each prompt.
0.60
GenEvel
SDXL
SD3-medium
0.55
RoN
BoN
TIEO iter
BoN TTEO iter=
16
24
32
40
48
56
64
BoN (N)
Figure 14: GenEval scores as function of BoN sampling iterations N. Curves compare the BoN baseline and BoN combined with TTEO (iter=2,4); horizontal dashed lines show multi-step diffusion reference performance
(SDXL (Podell al. 2023 SD3-medium (Esser al. 2024
Table 6: Runtime cost of different configuration in TTEO.
Contgaration
BoNcs
BcN2it
BcNL-1
BeNuss
BoNit TTTO mrn.?
BiAcs TTED Iur:?
BON-A TIDO man-
BeN: TTDX Hemo
Busins sped (njooepp:
1.46
2.64
4.00
8,57
2.52
5
9.58
1.9 CFG OF TEACHER MODEL.
The CFG mechanism for discrete models (e.g., AR and MDM) differs from the CFG commonly used in continuous diffusion/flow models (Dieleman 2022). In this work, across mask ratios as our default. We also experimented with several alternative CFG schedules (for example, linear schedules from 1.5 and cosine schedules from 1.5 2), but found that they do not improve Di(M performance.
constant CFG value
Tab. reports metrics for Di applied values. Consistent with Zhou et al. 22024)), we observe tradeoff between FID and CLIP score as the teacher CFG changes: a higher CFG produces better CLIP alignment but worse FID (for both COCO and MJHQ reference set). Importantly, further fine-tuning the distilled generator with a combination of the Di(M O objective and GAN loss reduces FID, with the largest improvements observed for FIDMIH We also attempted experiments with MaskGen- :I: however, the provided checkpoint cause unstable training under large CFG values and the model tend to collapse during distillation with Di[M in those settings. We use the generator distilled with CFG=4 MaskGen-L teacher as the initialization of our Soft-Di(MiO training.
MaskGen-L teacher using different teacher CFG
1.10 REWARD-ONLY FINE-TUNING
Tab. g we present results from fine-tuning the generator using different reward functions alone. We make the following observations. CLIP score is more strongly correlated with GenEval, reflect- ing improved prompt-following ability; in contrast, HPS primarily reflect aesthetic quality. Fine- tuning with CLIP reward alone improves prompt following (GenEval and CLIP score) but does not influence HPS score. On the other hand, optimizing HPS, ImageReward, or PickScore improves HPS and also yields silightly gains in GenEval (also hurt the FID different extends).
33

---

Published as a conference paper at ICLR 2026
Table 7: O performance MaskGen-L. with different CFG value. Models are trained with batch size of 64 and 100k steps
CPG value
2.0
2.0 +GAN
3.0
4.0
bsx4
+GAN
5.0
6.0
7.0
8.0
FIDama ())
6.79
5.87
8.52
9.73
9.68
6.58
10:45
10.96
11.35
11.57
FIDeoco (L)
22.74
21.26
23.99
24.15
24.56
21.46
24.56
24.60
24.72
24.77
Sutr-S (1)
0.296
0.294
0.298
0.299
0.299
0.298
0.298
0.298
0.298
0.298
GenEvaa (t)
0.432
0.424
0.430
0.429
0.426
0.427
0.439
0.425
IIPSV2.0 (f)
27.46
27.28
27.49
27.50
27.55
27.37
27.50
27.48
27.46
27.45
HPSv2.1 (f)
27.16
26.63
27.19
27.14
27.35
26.85
27.13
27.08
27.01
26.99
We further provide qualitative visual comparisons of generators fine-tuned with different single re-
wards in Fig. 15 The over-saturated appearance in some outputs is consistent with observations
reported in Li et al (2024b).
Table 8: Di(MOO model fine-tuned with different reward model alone, trained with 50k steps and bs=32. The similar effect of single reward alone fine-tuning is also reported in Reward-Instruee [uo et al. 2025b), but they did not report that the CLIP score can help with GenEval metric.
Rewsed Function
Wa
IragsRoowodd Xactal
Roaffird 2011
Brsciie
FEDyan
34,59
21.50
9.59
22.56
9.73
RiDeoco (1)
33.62
30.68
35:56
24.15
CLIP-Screc (
0.276
0.332
0.302
GesEval (T)
0.433
0.450
0.503
0.422
0.424
HPS-2 (1)
25.58
28.65
27.72
27.00
27.50
IPS:2 (1)
HPS
ImsgeRRewrd
CLI
PickSoore
Baseline
Figure 15: Visual comparison of models fine-tuned with individual reward functions. Each column shows outputs from a model fine-tuned only on the indicated reward (HPS, ImageReward, CLIP, PickScore) and the baseline (initial checkpoint used for fine-tuning). Prompts used are: A green dog and red cat and dynamic anime poster featuring blending realistic watercolors and modern anime. The palette is gloony and dramatic, dominated by dark crange and white. Weathercore aesthetic with swirling autumn leaves and mist, mysterious realism, evocative and emotional, cinematic composition --ar 69:128 --s 750 --niji 5.
collection of Asian woren in style
1.11 BALANCING CLIP AND IMAGEREWARD
Based on the results and analysis in Appendix 1.10 we use only the ImageReward and CLIP in our reward loss for the fine-tuning stage. In Tab. 9 we show results from fine-tuning the distilled MaskGen-L generator with ImageReward fixed at weight and varying the relative CLIP weight (overall reward loss weight 2000 in comparison to Di loss weight of 1). Increasing the CLIP weight consistently improves CLIP-Score and GenEval (better prompt alignment) but degrades Aes- thetic metrics (HPSv2.0/ / HPSv2.1). In practice, a moderate CLIP weight (roughly 0.2-0.5) offers a good trade-off between prompt adherence and image quality, and we select 0.2 as our final choice. A similar trend is observed for the Meissonic teacher as in Tab. TO only reward loss is used).
34

---

Published as a conference paper at ICLR 2026
Table 9: fine-tune Di(M model of MaskGen teacher with Di loss and reward, trained with 50K steps and bs=32. We use both ImageReward and CLIP as reward and ablate the relative loss weight of ImageReward. The Di(M loss weight I, the overall reward loss weight 2000, Inside the reward loss, the ImageReward loss weight 1., and we change the CLIP loss weight here.
fine-tune Di[MO model of Mcissonic teacher with Di[M loss and reward, trained with 10k steps and bs=16. We use both ImageReward and CLIP as reward and ablate the relative loss weight of CLIP compared to ImageReward. The Di[MjO loss weight is O, and the overall reward loss weight is 1000. Inside the reward loss, the ImageReward loss weight is and we change the CLIP loss weight here. This table is measured with generator tempera- ture=0.5.
Table 10:
CLIP compared
CLIP relativo weiglt
00
0.01
0.55
0.1
02
RM
1.0
FIDsmn (1)
12.11
11.36
10.48
9.74
10.36
9.0
10.43
CLIP relative weight
1.0
FilDoco (4)
24.24
22.29
CLI
0.305
GeaEval (?)
0.496
0.518
0.523
0.523
0.529
4.000
0:551
01:97
0:498
0 909
0.466
0.994
HPSx2.0 (f)
019/20111
10
28.15
27.00
HPSv2.J (f)
33.15
32,98
32.54
31,51
31.34
115x2.111
sttp
stpps
stopp
steps
steps
step
Figure 16: Ote-step generation from one-step generator compared to multi-step MaskGen-L teacher.
I.12 COMPARISON WITH MULTI-STEP TEACHER
In Tabs. to we compare our final one-step generator with the multi-step MaskGen-L teacher across all metries. The results show that one-step generator achieves both stronger prompt- following ability and higher aesthetic quality, without sacrificing FID. We provide a visual compar- ison in Fig. 16
35

---

Published as a conference paper at ICLR 2026
Table 11: Comparison of coDo, FIDM and CLIP- -Sooe for MaskGen-L (Kim et al. 2025 across varying generation steps and our one-step generator. The results are evaluated on MSCOCO val 30k dataset. The teacher sampling config is the same as the official ones. The asterisk item denotes the performance reported in the original paper.
Steps (4)
64
32
16
16
8
4
2
(owns)
FIDeoco (4)
22.60
22.70
22.78
22.64
22.89
24.84
44.21
88.70
23.44
FIDsmo (4)
7.35
7.43
7.74
7.59
8.32
11.58
29.92
77.57
10.59
CLIP- (f)
0.312
0.312
0.307
0.312
0.311
0.308
0.287
0.245
0.322
Table 12: Performance Geneval evaluation across varying generation steps for and XL., Each cell shows L/XL.
Steps (1)
64
32
16
(oanz)
Positin
0,000
0,053
0.00//
00000/00000
0.0M
Two Objects (t)
0,543/0.543
0.545/0.546
0,540 0.515
0,477 0L498
0.74
0.22//0.055
0.02D/0.000
0.59
Colors (*)
0.787:0.511
0.715/0.314
0,500 10543
0.76/ 10822
0.756
/
0,307/0.350
0,51
Color Arminttes (t)
0,148 0,340
0.133/0.15M
0,E31/0.150
0.010/00,00
0.18
Cousting (1)
0,384/0.400
0.572/0.403
0..35/
0.353 0.351
0.576
0.244/0.203
0.113/0.000
0.41
Single Otjject
0,992
0.942/0994
0,972 08947
0.978 IL947
0.966 0LA50
08813/0.444
0.525/0.563
0.98
Overall Scare (1)
/0.48)
0,479/0.467
0.482 70-499
0..24/0.509
0.178/0.165
0.509
Gamble strigdd doorgh
Top-1 Probubility
0.8
Sedt redsatingg
Entropy
0.4
10000 20000 30000 40000 50000 Trining Steps
20000 30000 50000 Training Steps
Figure 17: Entropy evolution during distillation.
Figure 18: Top-1 probability evolu- tion during distillation.
I.13 SAMPLING FROM ONE-STEP GENERATOR WITH arg max
To validate the observation from Di M Zhuu et al. 2025b) and ReDi (Yoo al. 2025) that output logits are highly concentrated with sharp predicted distributions, compare the FID of images sampled from the output distribution versus images obtained using Argmax in Tab. 14 As shown, both methods achieve nearly identical results, demonstrating that the output distribution is sharp enough for deterministic Argmax sampling (given the redundancy of the discrete latent space).
I.14 ENTROPY AND TOP- PROBABILITY EVOLUTION
In this subsection, we analyze the evolution of entropy and top-1 probability under different con- figurations. Both metrics are computed the mean across all tokens during distillation. We first present the evolution during Di distillation as a reference, where entropy decreases and top-1 probability increases over training. This Di MiO checkpoint is then used to initialize our Soft-Di M distillation. We compare the training dynamics of Soft-Di Di using our soft embedding approach versus Gumbel softmax straight-througs estimation. As shown in Figs. 17 and 18 our proposed soft embedding exhibits stronger concentration of the predicted logits compared to Gumbel softmax, achieving lower entropy and higher top-1 probability, corresponding Fig. 4a
36

---

Published as a conference paper at ICLR 2026
Table 13: Comparison of HPSV2.0 and HPSV2.1 scores across different generation steps for various visual styles. (L). Each cell shows v2.0 v2.1.
Stde
64
32
16
2
Fours HPSV2.1)
Arime
26.00/70 51
76 29.34
78.1//8.09
2604/26.72
27.79/2006
24.75/17.93 75.5/25778
29.51
Concegt Art
27/42/27.33
27.37//2200
27.3522554
2725/22:00
2704/26115
26.0/22.52
24.30/17.15
30.62
Paintings
27.8/227.3
27.44/22.99
27.9922332
2737/22555
27/2/25.00
25.93/2100
16443
20
Phono
2744525899
27:47/2887
2748/2578
25.36/2551
27.4//4700
2609/21.25
Average
27.64/27.70
27,62/22.00
27.56/22.43
27.52/27:4
27.5/26.11
26.17/2238
16.3
20.35
Table 14: Comparison of sampling methods showing nearly identical performance between sampling from logits and using arg max.
Method
FID-50K
IS
Precision
Rocall
Sample from logits: Categorical (logits) sample
1.56
273.2
0.81
0.60
Argmax: torch. argmax (logits,
dLm-
1.57
273.3
0.81
0.60
Table 15: Comparison of standard Di[M and Di using soft embedding.
Method
FID-50K
Precision
Recall
Standard Di
2.89
310.1
0.87
0.49
Soft embedding variant DiMO
3.33
315.9
0.88
0.47
I.15 USE SOFT EMBEDDING ALONE DI
While soft embedding is originally introduced to provide additional supervision beyond the teacher model, it can also be integrated directly into the O loss computation alone, without applying further downstream refinement losses. As illustrated in Fig. 19 the standard Di pipeline (top) first samples discrete tokens Fp, applies the forward mask-diffusin diate state H,, and then feeds it into the teacher á and auxiliary model d to construct the Di M loss. In the soft-embedding variant (bottom), we instead take the model logits za and immediately compute the corresponding soft embeddings using the teacher's embedding layer (the embedding layer of the teacher and the auxiliary model is the same and fixed all the time). Following the same mask schedule as in the standard diffusion process, a subset of these embeddings is then replaced by the mask-token embedding. This partially masked soft-embedding representation is subsequently passed to the teacher & and auxiliary & to form the Di[MOO loss. Compared to the original Di [M loss, the only change is replacing sampled discrete tokens with their soft-embedding counterparts in the forward pass. Importantly, in this variant, the Di loss does not backpropagate through the soft-embedding computation, the relaxation used only to stabilize the forward intermediate repre- sentation, not as differentiable surrogate. This contrasts with the refinement techniques introduced in the main method (GAN loss, reward tuning, TTEO), where gradients do propagate through the soft embedding, which is essential for enabling these refinement techniques on one-step MDMs.
process to obtain the interme-
In Tab. [5] we show that soft embedding variant of DiMO has a slight degradation in FID and recall compared to standard DiMO, likely due to the smoothing effect of soft embedding. However, the IS score and precision improves, potentially indicating that soft embedding helps generate higher- quality samples at the cost of diversity. This result also suggests that soft embedding serves good representation for the sampled tokens. In Fig. 20 we further provide the FID curve during distillation for both approaches, it shows that while the soft embedding variant (yellow curve) gives fast convergence, but the final converged performance of both approaches are similar with small gap. We leave it as future work to investigate this early convergence.
37

---

Published as a conference paper at ICLR 2026
Teacher d
Student 0
One-slep Generator
Logits
Tokens 2
Masked E,
Auxiliary y
Embby (=a)
Teacher +
Student
One-step Generator
Logits Se
Embo (Zp)
Auxiliary
Using soft embedding in D. Top: the original Di[M] O loss formulation. Bottom: the Di variant with soft embedding, where intermediate states represented using "partially masked" embeddings. The pink operator after embedding layer indicates that random subset soft embeddings is replaced with mask embeddings, following the same masking schedule as in the standard forward mask-diffusion process.
Figure 19:
Soft embbolding vainei DCMJO
12
FID (Sk)
10
10000
20000 30000 40000 50000
Training Steps
Figure 20: FID comparison of standard O and soft embedding variant Di OM The only difference of two experiments the way to construct the input teacher and auxiliary model, as illustrated in Fig. 19
J FAILURE CASES
Compared to the Di[MOO generator, we observe that our reward fine-tuned generator produces less diverse images for certain prompts (only a small portion in practice), as illustrated in Fig. 21
Figure 21: Visual examples of mode collapse of Soft- Pott-DijjO-MM MMsMoo-ea for certain prompts. Upper are
generated from Di generator and lower is generated with Soft- Boft-Di[M|O with MaskGen-L teacher. Prompts used Fantasie Antheia Crochet Beach Cover Up 6552 Small Twi Twilight and cou not found the image.
38

---

Published as a conference paper at ICLR 2026
Table 16: Class-conditional ImageNet-222 256 results with MaskBit teacher. This a more complete version of 1
Tab.
Farily
Methed
Stags (1) #Parmms
[DD())
15(1)
Prit)
Rec(")
BipGAN Brock
2015
112M
234.46
0,57
GipaDAN
Kang
A023
S69M
3.45
215.52
0.84
0.6
SyyeGANN
1504
ADM
Dhacwal & Nichol
2123
250
554M
4,59
156.00
0882
0.52
LDM-4G
Kontech et
250
400M
3.601
247
Diftaiio Flow
UMIT-
1.5)
Peckies Xic. 2023
250
67534
2.27
278.24
0,83
0.57
2.5
0.81
MisCCaf
237M
182.1
0.80
0.51
1emmaaa-A An Sux ot
2014
576
3.18
3.18
261.
088
0.58
MAR
2104:
4000
VAR-120
Tias al
2024
10
6004
2.57
302.6
0.8
0.56
TET
Yaetall
291.8
MAGVTB-2
Viret
64
907M
138
3194
MaABi
Weber et
2024
256
305M
1.52
325.6
256
MaABi
Wecker
2004
64
305M
1.66
319.90
0.8
0.60
Sengat
2003
675M
Shencor Frans al 2074
675M
10.00
675M
Few-Step from Scradh
2005a
67534
675M
1.59
Grng ca
2005 (XL/2)
678M
676M
LlaniGee L-co
Lio at 2004
326M
7.5%
237.5
0.54
0.37
Discore Disiilation
Di Hjo Makki Zhe et
2025b
305M
2.89
310.3
0.87
0.40
Sol- DUjjO (MaskBi (eacier)
305M
1.56
291.35
0.54
0.55
longor)
39

---

Published as a conference paper at ICLR 2026
K MORE QUALITATIVE RESULTS
In Figs. and we present randomly sampled ImageNet images generated single step by our distilled models with MaskBit teacher and MaskGit teacher, respectively. In Fig. 24 present image samples generated unconditionally in a single step by our distilled models with MaskGen-L teacher. Finally, in Figs. 25 and 26 we provide additional rext-to-image one-step generation results from our distilled model with the Meissonic teacher and MaskGen-L teacher.
I
Figure 22: Uncurated images from one-step generation with MaskBit teacher.
40

---

Published as a conference paper at ICLR 2026
Figure 23: Uncurated images from one-step generation with MaskGit teacher.
41

---

Published as a conference paper at ICLR 2026
Figure 24: One-step unconditional generation from reward fine-tuned one-step generator with MaskGen-L teacher.
42

---

Published as a conference paper at ICLR 2026
Figure 25: One-step generation from our reward fine-tuned one-step generator with MaskGen-L teacher.
43

---

Published as a conference paper at ICLR 2026
Figure 26: One-step generation from our reward fine-tune tuned one-step generator with Meissonic teacher.
44

---

Published as a conference paper at ICLR 2026
L
MISC.
Prompts
Below a collection of creative prompts we used to generate images in Figs. 7 16
25 and 26 (in the order from left to right, from top to bottom):
Fig.1
A pair old leather boots sprouting tiny photorealistic flowers and moss, blending nature and object.
A distinguished older gentleman in vintage study, surrounded by books and dim lighting, his face marked by wisdom and time. 8K, hyper-realistic, cinematic, post-produee post-production.
A cute pixel art of a cat in a sunny garden.
Design tranquil interiors that celebrate the calming power of sage green.
A floating castle in the sky, waterfalls cascading into the clouds, painted in whimsical Studio Ghibli style.
Ombre color scheme of neon yellow, neon blue, neon pink, neon green, neon orange, cute adorable alien insectoid Siberian musk deer, Chinese water deer creature with (Clawed Appendages) and colorful (Turquoise)- tailed alien wildlife photography, intricate biological detail, creature, insect focus, hu- mans, spray painted fluorescence on particularly skin, fluorescent, fluorescent dust, fluores- cent spray paint, glowing dust, subsurface scattering, Photorealistic, Hyper realistic, analog style, realistic, film photography, soft lighting, heavy shadow.
Horns, transparent nodules, highly de-
portrait in the style of polygonal painting in the style of Alexander Deineka, the face is divided into colored polygons of different sizes, mosaic, combination of complex cold and warm colors, a play with large color planes, modernism, cubism, sophistication, tenderness and expression. Light hair flutters in the wind, creating picturesque intersecting lines and planes.
A detailed isometric LEGO diorama of private plot landscape LEGO design on a slope, showcasing Tokyo Tower and nice parks with plants decorated ocean side. Design was made with LEGO blocks. The diorama includes visual elements like flowers, pathway, grill area, pond, plants and trees. The background is minimalist white, focusing the structured and secure architecture of the landscape design.
A metallic apple reflecting a distorted city skyline, hyper-detailed reflections.
Pasta & Pepper Salad.
Fig.33 3
bowl of strawberries with powdered sugar sprinkled on top.
schoolgirl with twin braids holding a sketchbook.
a pair glasses resting on an open notebook.
a dog toy lying on a tiled kitchen floor.
knight silver armor holding baner.
anime-style slime creature with cute round eyes.
Fig. 7
saharian landscape at sunset, 4k ultra realism, BY Anton Gorlin, trending on artstation, sharp focus, studio photo, intricate details, highly detailed, by greg rutkowski.
chinese red blouse, in the style of dreamy and romantic compositions, floral explosions -ar 24:37 -tyllize 750 6.
Digital 2D, Miyazaki's style, ultimate detailed, tiny finnest details, futuristic, sci-fi, magical dreamy landscape scenery, small cute girl living alone with plushified friendly big tanuki in the gigantism of wilderness, intricate round futuristic simple multilayered architecture, habitation cabin in the trees, dramatic soft lightning, rule of thirds, cinematic.
45

---

Published as a conference paper at ICLR 2026
poster art for the collection of the asian woman, in the style of gloomy, dark orange and white, dynamic anime, realistic watercolors, nintencore, weathercore, mysterious realism ar 69:128 -s 750 -niji 5.
A Yantasy-themeet portrait of female elf with golden hair and violet eyes, her attire shim- mering with iridescent colors, set in an enchanted forest. 8K, best quality, fine details.
'very beautiful girl in bright leggings, white short top, charismatic personality, professional photo, style of jessica drossin, super realistic photo, hyper detail, great attention to skin and eyes, professional photo.
(steampunk atmosphere, a stunning girl with a mecha musume aesthetic, adorned in- tricate cyber gogle,) digital art, fractal, 32k UHD high resolution, highres, professional photography, intricate details, masterpiece, perfect anatomy, cinematic angle cinematic lighting, (dynamic warrior pose:1)
(Pirate ship sailing into bioluminescence with galaxy in the sky), epic, 4k, ultra.
tshirt design, colourful, no background, yoda with sun glasses, dancing festival, ink splash, 8k.
Half-length head portrait of the goddess of autumn with wheat ears on her head, depicted dreamy and beautiful, by wlop.
Walter White dressed as medieval-style king.
A serene meadow with a tree, river, bridge, and mountains in the background under slightly overcast sunrise sky.
A closeup portrait of gray owl with spreaded wings attacking cinematic lighting, digital painting by Greg Rutkowski used as album cover art on Artstation.
A hyena fursona sitting in the grass in savannah at sunset.
Large dog looking at television show in living room.
A fantasy-themedd portrait female elf with golden hair and violet eyes, her attire shim- mering with iridescent colors, set in enchanted forest. 8K, best quality, fine details.
Fig. 16
a goldfish swimming in a round glass bowl.
a ceramic vase with dried flowers on shelf.
a notebook with doodles on the cover, anime style with cute screentones.
a koi fish swimming in a pond with rippling ink-drawn water. snow globe sitting on windowsill.
bench under tree with fallen leaves.
Fig. 25
bird perched on a windowsill with frost outside.
steaming cup of tea on a small saucer.
a shrine gate drawn with sharp black-and contrasts.
an anime-style vampire with elegant clothing and sharp eyes.
a street covered in fresh snow.
a commuter waiting at bus stop during sunset.
a street with puddles reflecting city lights.
basket of oranges on market stall.
bird's nest in the branches of tree.
basket of fresh bread kitchen counter.
jar of seasalt with cork lid beside it.
46

---

Published as a conference paper at ICLR 2026
cat curled up inside cardboard box.
train arriving at suburban station in the evening.
a wooden bench covered in frost.
a pirate with large hat feathered coat.
jar of honey with wooden dipper beside
sunflower field under a bright sky, anime style with vibrant colors.
fox mask hanging from a wooden wall, anime style with dramatic lighting. a bicycle parked under cherry blossom trees, petals falling in manga style. an empty parking lot at night.
collection of glass bottles lined up on a windowsill catching afternoon sunlight.
metal gate with peeling paint.
wind-up toy robot on a wooden table, anime style with cel shading.
an overgrown railway track disappearing into bushes.
an anime-style fairy with translucent wings glowing softly.
glass bottle lying on sidewalk.
jar filled with fireflies, anime style glowing in the dark.
anime-style samurai with a tattered cloak standing under moonlight.
a bottle of olive oil beside a loaf of bread.
a beach umbrella folded on the sand.
Fig. 26
Digital art of Prince of Roses.
A landscape featuring a Kyoto Animation-style building.
A path winding through forest depicted in digital art.
close-up portrait beautiful girl with an autumn leaves headdress and melting wax.
A neon-soaked cyberpunk alleyway with rain-drenched streets and futuristic holograms, gritty yet vibrant, hyper-realistic, ultra-detailed, cinematic scene.
A serene mountain landscape at sunrise, mist rolling over rugged peaks, ultra-detailed, photorcalistic, soft lighting, high-resolution, digital art.
A hyper-detailed closeup of a dew-covered insect on a vibrant leaf, extreme macro photog raphy style, ultra-realistic, high-resolution, intricate textures.
An anime-style magical girl in dynamic pose, vibrant colors, ultra-detailed costume and background, energetic, high-resolution, cinematic lighting.
An enchanted autumn forest with falling leaves and warm, glowing light, ultra-detailed, photorealistic, rich textures, digital art, serene mood.
An elegant Renaissance portrait of noble figure, detailed textures, soft natural lighting, ultra-detailed, classical, high-resolution, oil painting style.
A cybernetic humanoid robot portrait with metallic textures and neon accents, ultra- detailed, photorealistic, cinematic. futuristic digital art.
A vibrant street art mural on an urban wall, ultra-detailed, energetic, bold colors, high- resolution, digital painting, modern art style.
A dark fantasy warrior in intricately detailed armor standing in stormy battlefield, ultra- detailed, hyper-realistic, cinematic, dynamic action scene.
An intense lightning storm over vast desert landscape, ultra-detailed, dramatic, high- resolution, cinematic, digital art, atmospheric.
A detailed nature macro shot of a vibrant flower with dewdrops, ultra-detailed, photoreal- istic, high-resolution, digital painting, delicate textures.
47

---

Published as a conference paper at ICLR 2026
An ultra-realistic snowy mountain village under starry sky, ultra-detailed, atmospheric, cinematic, high-resolution, digital winter wonderland.
A humorous portrait cat dressed as a Victorian aristocrat, in vintage photorealism.
A photorealistic shot of bouquet of wildflowers in clear glass vase on sunlit win- dowsill.
A mecha jet fighter engages in an air battle with an explosion as a backdrop, set against a dark, starry sky in highly-detailed art piece by Stephan Martiniere.
A young black woman stands in front of ringed planet in space.
Digital art of a cherry tree overlooking valley with waterfall at sunset.
An astronaut in white futuristic cybernetic armor running on the surface the moon, fea- tured in artwork illustration on Artstation.
The image is headshot of happy girl with white hair in school uniform, illustrated by Ilya Kuvshinov.
A minimalistic heart drawing created using Adobe Illustrator.
The image is digital art headshot of owlfolk character with high detail and dramatic lighting.
Close up an eye with the Earth inside the pupil, inspired by Wes Anderson's art.
Landrover drives through rain-soaked forest in a highly-detailed digital artwork by Greg Rutkowski and Artgerm.
A snowy lake in Sweden captured in vibrant, cinematic style with intense detail and raytracing technology showcased on Artstation.
A heron silhouetted against beautiful sunrise, created by Greg Rutkowski.
A surreal portrait of woman with giant carnation face in flower field at sunset with colorful clouds and large sky, created by artist Simon Stålenhag.
48


Reviews:
[{"review_id": "BtdFhf6YXX", "reviewer": "Reviewer_rzsv", "summary": "The paper introduces Soft-DiMO, an extension to DiMO with soft-embeddings that allow differentiating through the discrete token sampling step for post-training use cases such as supervised finetuning and human-alignment / prompt-adherence. The paper shows the superiority of their method compared to other methods, like the Gumbel Softmax trick, which allows differentiating through discrete sampling steps. The paper shows a SoTA FID score on discrete image generation trained on ImageNet 256x256.", "questions": "1. On Line 44, the authors write \"First, they inherit modeling errors from the teacher\" as one of the challenges of current one-step discrete generators. Can the authors please clarify why this is a problem? Are the authors suggesting that the errors get distilled to the student in preference to the image statistics in a more prominent way than they occur in the teacher model? Given the teacher is a powerful model with minimal region for mistakes, wouldn't the student model also be minimal in the param region for mistakes (since zero errors is unavoidable)?\n2. My central question is around the arguments for why the proposed method works better than the others like Gumbel-Softmax other than a qualitative hand-wavy argument about high variance and high bias training signals. It is unclear why the proposed method alleviates such problems. A theoretical argument (despite any assumptions made) would help the reader to appreciate the idea proposed more than an empirical hack to circumvent post-training challenges of discrete sampling models.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The authors propose a simple way to backrprop through the discrete sampling step that unlocks a wide variety of potential post-training applications. \n- The authors demonstrate various post-training techniques using their method and show better performance\n- The authors achieve a SoTA FID score on ImageNet 256x256 \n- The paper also compares the presented method to a multiple existing works", "weaknesses": "Minor:\n- It is helpful for the reader to have the abstract be independent and include inline full citations of terms used like DiMO, MaskBit, MaskGen, HPS, etc. E.g.: (Author et al. 2022), in line with the abstract.\n\nMajor:\n- I've had issues with the readability of the paper especially in Section 2 and Section 3 where the problem statement is not well set up to be complete. For instance, what is the space of vocabulary V for this problem? What are the typical patch sizes for each token? What dimensions are the output probability logits for each token? \n- The central idea of the paper is to weight the output token classification probability with the corresponding token embedding to obtain a so-called \"soft-embedding\" that can be used to backprop to empirically make post-training work better than other techniques.\n-- However, the novelty is limited since the proposed soft-embedding is a weighted average w.r.t the word probability. \n-- Empirically, it is clear that it works better (at least in the hyperparam region used for experiments). But it is not clear *why* this works better, either intuitively and/or theoretically. Is this formulation really approximating backproping through the discrete sampling step better than the Gumbel-Softmax trick? If so, why? Is it purely due to optimization landscapes that are hard to theoretically prove or can some theoretical arguments be made to understand why this works better empirically?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "q74ngCfGqh", "reviewer": "Reviewer_jPZF", "summary": "The paper proposes Soft-Di[M]O, a one-step discrete image generator distilled from masked diffusion teachers (e.g., MaskBit, MaskGen). The key idea is to replace sampled discrete tokens with soft embeddings—the expectation of token embeddings under the model’s output distribution—so gradients can flow through the generator and decoder. This makes one-step students end-to-end trainable, enabling (i) GAN-based refinement, (ii) differentiable reward fine-tuning (e.g., CLIP/ImageReward).", "questions": "For TTEO, what selection metric do you use to pick the final image—e.g., CLIP score, ImageReward/HPS, discriminator logit for GenEval and HPS v2.1?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Soft embeddings (expected embedding; Eq. (3)) are a minimal change that preserves compatibility with teacher backbones/tokenizers while restoring differentiability.\n\n2. The framework cleanly integrates GAN loss, reward losses, broadening what one-step discrete students can optimize.\n\n3. The paper studies discriminator input choices and mask schedules.", "weaknesses": "1. It’s unclear how much improvement comes from the GAN refinement versus the soft-embedding relaxation itself. While Fig.-style ablations compare Di[M]O vs GAN only vs combined, it only reports 5k FID.\n\n2. TTEO (test-time scaling) increases inference compute; the paper could quantify cost-vs-quality more explicitly. For TTEO, what selection metric do you use to pick the final image—e.g., CLIP score, ImageReward/HPS, discriminator logit for GenEval and HPS v2.1?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "we2j0cPbeA", "reviewer": "Reviewer_UkJL", "summary": "This paper addresses a key limitation of one-step generative models distilled from Masked Diffusion Models: their discrete token blocks gradient flow. The authors propose soft embeddings, a simple and effective relaxation that replaces discrete tokens with the expected embedding under the generator's output probability distribution, which makes the one-step generator differentiable. By integrating this into the Di[M]O framework, the method unlocks the ability to apply refinement techniques, including GAN-based training, differentiable reward fine-tuning, and Test-Time Embedding Optimization. Empirically, it achieves state-of-the-art one-step results, notably a 1.56 FID on ImageNet-256 and improved alignment scores on text-to-image tasks.", "questions": "- What's the performance of only applying soft embeddings without GAN loss?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper identifies a significant and practical problem (blocked gradients in one-step discrete generators) and proposes an simple solution.\n- The differentiable feature enables many technique applications like GAN, reward fine-tuning.\n- The method achieves strong results on ImageNet ($1.56$ FID). And the MSE between samples decoded by soft embedding and the corresponding discrete embedding is very small.\n- The authors conduct comprehensive experiments validating the effectiveness and design choices.", "weaknesses": "- The idea of using the expectation of discrete embeddings is not a new technique in machine learning. So its novelty lies in the application to one-step MDMs.\n- The generation performance relies on the introduction of GAN loss, which potentially introduces training instability and conplexity of hyperparameter tuning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PHIndrW2p4", "reviewer": "Reviewer_S4vL", "summary": "This paper extends the Di[M]O framework by introducing a soft embedding relaxation for discrete diffusion distillation. Instead of hard token sampling, the model uses non-parametric expected embedding under the predicted token distribution, enabling differentiability throughout the training process. This modification makes the framework compatible with gradient-based extensions such as adversarial fine-tuning and preference optimization. Experiments evaluate both class-conditional image generation and text-to-image tasks.", "questions": "1. Does the relaxation introduce distributional bias or oversmoothing compared to discrete sampling?\n\n2. How does this compare against simpler continuous relaxations like Gumbel-Softmax or parametric estimators in the same setting?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. Clear and intuitive motivation: The paper articulates very well why enabling differentiability is valuable — it broadens the range of learning signals applicable to discrete diffusion models.\n\n2. Reasonable and meaningful novelty: While the relaxation itself (Gumbel–Softmax-style) is known, its application in the context of distilled discrete diffusion models is new and practically relevant. The contribution sits at a “methodological adaptation” level — modest but non-trivial.\n\n3. Clarity of presentation: The derivation of soft embeddings, training objective, and overall algorithm are explained cleanly and are easy to follow. The paper is well written and pedagogical.\n\n4. Empirical coverage across modalities: The inclusion of both class-conditional generation and text-to-image synthesis demonstrates generality of the approach beyond a single domain.", "weaknesses": "1. Overemphasis on auxiliary fine-tuning tasks: The paper devotes significant space to adversarial fine-tuning and differentiable reward fine-tuning, but these are not core contributions. They are applications or privileges enabled by the soft embedding, not new methods. This focus dilutes attention from the main relaxation technique, which deserves deeper theoretical and empirical exploration.\n\n2. Limited ablation on the method itself: The paper could better analyze the behavior of the soft embedding mechanism—e.g., temperature schedules, entropy evolution, convergence stability, and how it compares to direct Gumbel-Softmax or straight-through variants.\n\n3. Experimental scope modest: Experiments are conducted on moderate-resolution datasets and lack comparisons with alternative differentiable relaxations. Broader benchmarks or higher-resolution evaluation would strengthen the claims.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)