DiDi-Instruct: Discrete Diffusion Distillation

• DiDi-Instruct is a distillation framework that trains a compact student to emulating a multi-step teacher through integral KL-divergence minimization.
• It employs techniques like grouped reward normalization, intermediate-state matching, and reward-guided ancestral sampling to ensure robust, diverse generation.
• Empirical results show up to 128× speedup with competitive perplexity and entropy compared to standard autoregressive and diffusion models.

Discrete Diffusion Divergence Instruct (DiDi-Instruct) is a distillation framework designed to accelerate discrete diffusion-based LLMs by training a compact student model to emulate the behavior of a deeper, multi-step teacher via integral KL-divergence minimization. DiDi-Instruct leverages a combination of theoretical advances—integral KL-divergence (IKL) minimization, grouped reward normalization, intermediate-state matching, and a reward-guided ancestral sampler—to deliver ultrafast generation of discrete sequences, including text and protein, at quality levels matching or exceeding large masked diffusion teachers and standard autoregressive baselines with a minimal loss of entropy and diversity (Zheng et al., 29 Sep 2025).

1. Theoretical Foundation

At the core of DiDi-Instruct is the distillation of a high-quality multi-step masked discrete diffusion LLM (MDLM)—the teacher, parameterized by $\theta$—into a student with parameters $\nu$ running for only $K\ll 1024$ function evaluations (NFEs). This is operationalized via minimization of an integral KL-divergence (IKL) over the continuous diffusion time interval $t\in[0,1]$:

$$\text{DKL}_{\text{IKL}}(q_\nu\parallel q_\theta) := \int_0^1 \omega(t)\,\text{KL}( q_\nu(z_t, t) \parallel q_\theta(z_t, t) )\,dt.$$

Here, $q_{\theta}(z_t,t)$ is the marginal of the teacher under corrupting kernel $Q(\cdot|x)$, which defines a masked categorical (token-wise) forward process with schedule $\alpha_t$, and $q_\nu$ analogously for the student. Since the process is discrete and non-differentiable, gradients with respect to $\nu$ are estimated using the REINFORCE score-function identity. The learning signal is the reward

$\nu$0

which is estimated via an adversarial discriminator distinguishing student from teacher-generated states. The update direction is:

$\nu$1

where $\nu$2 and $\nu$3 are weighting schedules, and $\nu$4 is the student’s probability of generating a fully masked input at $\nu$5. The reward $\nu$6 is estimated with discriminator $\nu$7 and averaged over masked positions:

$\nu$8

Crucially, DiDi-Instruct introduces intermediate-state matching, requiring the student to match the teacher’s distributions not only at $\nu$9 (fully masked input) but at randomly sampled intermediate times $K\ll 1024$0, further decomposing the learning gradient and ensuring robustness.

2. Training and Inference Methodology

Distillation Loop

The DiDi-Instruct training pipeline alternates between discriminator and student updates:

1. Sampling: Times $K\ll 1024$1 are drawn for a mini-batch of size $K\ll 1024$2.
2. Corruption: Both student and teacher outputs are corrupted to $K\ll 1024$3 using the masked kernel $K\ll 1024$4 at $K\ll 1024$5.
3. Discriminator Update: Minimize

$K\ll 1024$6

1. Student Update: Compute normalized reward $K\ll 1024$7, then apply gradient updates using the mini-batch estimator.

Reward-Guided Ancestral Sampler (RGAS)

For inference, the student employs RGAS to maximize output quality:

• Early Steps (Gradient Tilting): Adjust prediction logits by the reward gradient ($K\ll 1024$8) for global structure shaping.
• Late Steps (Re-ranking): Draw multiple candidates $K\ll 1024$9 and select based on exponentiated rewards, enhancing diversity and sample quality.

RGAS empirically grants a 20–30% perplexity reduction over standard ancestral sampling for discrete diffusion models.

3. Distinctive Techniques and Innovations

Grouped Reward Normalization (GRN)

Raw reward signals $t\in[0,1]$0 can exhibit high variance, destabilizing updates. DiDi-Instruct normalizes them within each mini-batch:

$t\in[0,1]$1

where $t\in[0,1]$2 and $t\in[0,1]$3 are the batch mean and variance, closely analogous to advantage normalization in reinforcement learning. This normalization significantly stabilizes training.

Intermediate-State Matching

The algorithm samples intermediate time steps $t\in[0,1]$4 and matches both the student and teacher at those times, not just the initial/final diffusion states. The gradient is thus split, so the student is forced to denoise both fully masked and partially corrupted states, thereby preserving output entropy and preventing mode collapse.

Reward-Guided Ancestral Sampler (RGAS)

RGAS replaces conventional ancestral sampling by using reward gradients and candidate re-ranking to ensure high-probability, diverse sample coverage consistent with the teacher. Early steps nudge the trajectory via gradient tilting; later steps exploit re-ranking for robustness and diversity.

4. Empirical Performance and Benchmarking

Comprehensive evaluation on OpenWebText compared DiDi-Instruct to a 169M-parameter teacher (1024 steps) and GPT-2 Small. Metrics include perplexity (PPL; lower is better), entropy (Ent; higher means greater diversity), and speedup.

NFEs	GPT-2 Small	Teacher (1024)	DiDi-Instruct (PPL / Ent / Speedup)
8	92.1 / 5.12	38.5 / 5.22	62.2 / 5.17 / 128×
16	43.0 / 5.13	38.5 / 5.22	30.99 / 5.22 / 64×
32	32.3 / 5.09	38.5 / 5.22	23.60 / 5.20 / 32×
64	26.1 / 5.04	38.5 / 5.22	19.61 / 5.18 / 16×
128	21.4 / 5.02	38.5 / 5.22	17.50 / 5.17 / 8×

Even with 8 NFEs (128× faster than the teacher), DiDi-Instruct outperforms GPT-2 Small, and by 16 NFEs (64×), it surpasses the full-step teacher. Sample entropy remains within 1% of the teacher, suggesting no diversity collapse. Training requires approximately one GPU-hour on an H100, over 20× less than previous distillation approaches such as SDTT and DUO (Zheng et al., 29 Sep 2025).

5. Ablation Studies and Generalization

Cumulative ablation demonstrates the necessity of each novel component. For instance, omitting score decomposition or time coupling causes severe performance degradation (PPL > 30,000 or > 300). Final gains are delivered by guided inference, which reduces 8-NFE PPL from 803.9 (baseline) to 62.24, and 128-NFE PPL from 96.6 to 18.45.

Config	8 NFEs PPL	128 NFEs PPL
Baseline	803.9	96.6
+ Score Decompose	667.8	89.3
+ Coupled t	101.0	30.6
+ ω(t) Correction	94.96	20.98
+ π(t) Weighting	92.10	21.38
+ Regularization	88.27	18.33
+ Guided Inference	62.24	18.45

For larger model scales (424M parameters), DiDi-Instruct maintains >50% performance benefit up to 128 NFEs while preserving entropy. The methodology generalizes to discrete protein sequence generation, where, on UniRef50, DiDi-Instruct improves ESMFold pLDDT confidence by up to +10 points at 8 NFEs and achieves high foldability with ≤32 NFEs, compared to the teacher’s >128 required.

6. Impact and Scope

DiDi-Instruct constitutes the first single-stage, distribution-matching distillation method for discrete diffusion LLMs. Its design combines the IKL framework with robust stochastic and adversarial reward signals, intermediate matching, and efficient decoding, enabling state-of-the-art acceleration (up to 128×), quality equal to or exceeding much slower diffusion teachers, and minimal loss in sample diversity. The resulting framework delivers substantial decreases in wall-clock training cost (over 20× faster than prior methods) and demonstrates robustness across tasks, scales, and output domains such as protein sequence design (Zheng et al., 29 Sep 2025).