PaperlayerView exact revision on arXiv

Source-linked AI summary

Omni-Diffusion-Distill: Few-Step Distillation of Unified Multimodal Diffusion Large Language Models

Hong Huang, Chenhongyi Yang, Junzhe Sun, Animesh Sinha, Wuyang Chen, Yifan Jiang

arXiv:2610.10990v1cs.CVcs.LG

TL;DR

Unified multimodal dLLMs can require many iterative decoding steps, and existing few-step distillation largely treats image and text separately. Omni-Diffusion-Distill uses a two-stage shared discrete-token framework with modality-specific remedies. It reports strong generation and understanding performance with fewer steps and substantial wall-clock speedups.

  • Problem

    Unified multimodal dLLMs require tens to hundreds of decoding steps, while existing few-step distillation largely treats image generation and text dLLMs separately.

    Source: §1 · p. 2

  • Method

    A shared discrete-token framework first replays cached teacher trajectories to teach step skipping, then refines on student rollouts with image entropy-matched guidance and a text collision penalty.

    Source: §1 · p. 3, §3.3 · p. 5, §3.4.1 · p. 5

  • Results

    At 8 image-generation and 64 understanding steps, it leads distilled baselines on the reported benchmarks, scoring 0.828 on GenEval and 57.2 on COCO captioning.

    Source: §1 · p. 3

  • Takeaways & Limitations

    One student serves image generation and multimodal understanding, with block-causal text attention allowing exact KV-cache reuse for understanding.

    Source: §4.4 · p. 10

    The method requires a shared discrete token space, iterative masked decoding, and classifier-free guidance for images.

    Source: §3.1 · p. 3

Abstract

from arXiv · show

Unified multimodal diffusion large language models (dLLMs) offer a single architecture for both image generation and multimodal understanding, but their iterative decoding requires tens to hundreds of forward passes. Existing few-step distillation methods largely focus on either image generation or text generation, making it unclear how to compress a fully discrete multimodal dLLM into a single efficient student while preserving both generation and understanding. We introduce Omni-Diffusion-Distill, a unified two-stage distillation framework that retains strong generation and understanding capabilities while substantially reducing the inference cost of a unified multimodal dLLM. Omni-Diffusion-Distill aligns the distillation of both generation and understanding, for both images and text, in the discrete token space. In the first stage, the student is trained to skip decoding steps by replaying cached teacher trajectories, and in the second stage the student is refined on intermediate states along its own rollouts. We further remedy two sources of degradation in unified distillation with a pairwise collision penalty that reduces repetition under parallel text decoding, and entropy-matched guidance that prevents entropy collapse caused by fitting the sharpened teacher distribution in image generation. Omni-Diffusion-Distill achieves state-of-the-art trade-offs between decoding efficiency and generation and understanding performance for multimodal dLLMs, reducing image generation from 128 to 8 decoding steps and multimodal understanding from 512 to 64, giving 18.2x and 21.2x wall-clock speedups. Under these budgets, it scores 0.828 on GenEval and 83.0 on DPG-Bench for text-to-image generation, while reaching GPT judge scores of 20.0 on MM-Vet and 57.2 on COCO captioning (twice the teacher's 28.4 at the same steps) for multimodal understanding.

1 INTRODUCTION

Omni-Diffusion-Distill addresses the cost of iterative decoding in unified multimodal dLLMs, where existing few-step distillation largely treats image and text separately. Its two-stage framework distills both modalities and reports strong generation and understanding performance at substantially fewer decoding steps.

1 INTRODUCTION

  • Unified multimodal dLLMs can require tens to hundreds of sequential decoding steps, while existing few-step distillation methods largely address image or text generation separately.Image and text can share a discrete token space in these models, but shortening both decoding processes at once raises modality-specific challenges.Source: §1 · p. 2
  • At 8 image-generation steps and 64 understanding steps, Omni-Diffusion-Distill leads distilled baselines on the reported benchmarks, with 0.828 on GenEval and 57.2 on COCO captioning.It is 18.2× faster than the teacher for text-to-image generation and 21.2× faster for understanding on one A100.Source: §1 · p. 3
  • The framework replays cached teacher trajectories to teach step skipping, then refines the student on its own rollouts with entropy-matched guidance for images and a collision penalty for parallel text decoding.The text penalty targets repeated tokens among positions decoded in parallel; the image guidance rescales the CFG target to match the conditional prediction's entropy.Source: §1 · p. 3

2 RELATED WORK

Prior few-step distillation methods are largely modality-specific, while Omni-Diffusion-Distill treats text and images as masked tokens of one model and refines on states generated by the student.

2 RELATED WORK

  • Few-step methods span image diffusion, discrete image generation, autoregressive LLMs, and text dLLMs, while unified-model acceleration has also used consistency distillation or speculative decoding with a continuous image branch.The paper distinguishes its approach by applying on-policy refinement to student-produced states and treating both modalities as masked tokens of one model.Source: §2 · p. 3

3 METHOD

Omni-Diffusion-Distill uses a shared discrete-token framework with off-policy trajectory replay followed by on-policy refinement. Image entropy matching and text-specific sequential targets and collision penalties address distinct degradation modes.

3.1 BACKGROUND AND PRELIMINARIES

  • A shared discrete token space, iterative masked decoding, and image classifier-free guidance define the method's stated requirements.The student uses block-causal attention for text and bidirectional attention for images; finished text blocks' KV cache can be reused exactly.Source: §3.1 · pp. 3–4

3.2 MOTIVATION FOR TWO-STAGE TRAINING

  • Stage 1 learns few-step decoding from teacher-trajectory states, but Stage 2 is needed because the student's own inference states can accumulate errors.Stage 1 trains on teacher states; Stage 2 trains on intermediate states from student rollouts and also addresses repetition and entropy collapse.Source: §3.2 · p. 4

3.3 STAGE 1: OFF-POLICY TRAJECTORY REPLAY FOR FEW-STEP DECODING

  • Stage 1 replays cached teacher trajectories so the student learns to make in one step commitments the teacher makes over several.Its four terms are trajectory matching, self-consistency, path training on reconstructed intermediate states, and a ground-truth masked cross-entropy anchor.Source: §3.3 · pp. 4–5

3.4 STAGE 2: ON-POLICY REFINEMENT ON STUDENT ROLLOUTS

  • Stage 2 rolls out the student, scores its intermediate states with teacher predictions, and trains on those states rather than relying only on teacher trajectories.The same on-policy loop covers both modalities, with task-specific states, teacher predictions, and matching terms.Source: Figure 3 · p. 4, §3.4 · p. 5, §3.4.2 · p. 6

3.4.1 IMAGE BRANCH: MATCHING THE DISTRIBUTION

  • For image generation, the student samples and re-masks an image from its rollout prediction, then matches an entropy-rescaled CFG teacher target at scored positions.The rescaled CFG target's entropy is set to a fixed fraction of the conditional prediction's entropy; an auxiliary network estimates the distribution of student-generated images.Source: §3.4.1 · p. 5

3.4.2 TEXT BRANCH: SEQUENTIAL TARGETS AND PENALTY ON PARALLEL COLLISIONS

  • For text, the teacher predicts scored positions sequentially so each target is conditioned on tokens already committed, while the student continues to decode in parallel.The text branch matches per-position predictions and adds a penalty for collisions between positions decoded in the same step.Source: §3.4.2 · p. 6
  • The pairwise collision penalty reduces but does not forbid repetition, and is applied only to text.It targets pairs of scored positions because independently sampled positions in the same decoding step can commit the same token.Source: §3.4.2 · p. 6

4 EXPERIMENTS

Omni-Diffusion-Distill is evaluated across text-to-image generation, image-to-image generation, and multimodal understanding at substantially reduced decoding budgets. Performance remains strong overall, though open-ended understanding shows larger losses than discriminative benchmarks.

4.1 EXPERIMENTAL SETUP

  • Experiments distill Lumina-DiMOO on approximately 51K examples and evaluate four tasks across 17 benchmarks, using 8/4/2 image steps and 64/32/16 understanding steps against teacher budgets of 128 and 512.Source: §4.1 · pp. 6–7

4.2 T2I: OUTPERFORMING EVERY DISTILLED BASELINE AND THE FEW-STEP TEACHER

  • 0.828 on GenEval and 83.0 on DPG-Bench at 8 steps; Omni is the best distilled method on all five T2I benchmarks.Source: §4.2 · p. 7, Table 1 · p. 7

4.3 I2I: EDITING DEGRADES LITTLE UNDER STEP REDUCTION

  • 2.6% on the I2I benchmark and 0.9% on Graph-200K: scores drop when the teacher goes from 128 to 8 steps, while Omni is on par with the best distilled baseline.Source: §4.3 · p. 8, Table 2 · p. 8
  • 3.32 final entropy and a 2.31 I2I benchmark score with entropy-matched guidance, versus 1.55 entropy and 2.04 without it.Source: §4.3 · p. 8
  • 3.26 on Add versus 2.52 for the best other distilled method; Omni leads on Add and Compose and ranks second on I2I Overall and Graph-200K.Source: §4.3 · p. 8

4.4 MMU: STEP REDUCTION PRIMARILY AFFECTS OPEN-ENDED GENERATION

  • Within half a point across all five discriminative scores when teacher steps fall from 512 to 64; the student is similarly stable from 64 to 16 steps.Source: §4.4 · p. 9, Table 3 · p. 9
  • 57.2 on COCO captioning versus 28.4 for the few-step teacher at 64 steps; Omni beats it on all five open-ended benchmarks but still trails the full-step teacher, most on DeCapBench.Source: §4.4 · p. 9, Table 4 · p. 10
  • At 8 image and 64 understanding steps, the student runs 18.2× and 21.2× faster than the teacher and outperforms distilled baselines on most benchmarks.Source: §4.4 · p. 10

A.3 DISTILLATION FOR FEW-STEP GENERATION

Prior few-step distillation methods largely address image generators, language models, or dLLMs separately; Omni-Diffusion-Distill applies one on-policy framework to both modalities of a unified discrete model. Its training combines cached teacher-trajectory replay with student-rollout refinement and modality-specific losses.

A.3 DISTILLATION FOR FEW-STEP GENERATION

  • Few-step distillation has developed in largely separate lines for image generators, autoregressive LLMs, and dLLMs, while existing unified-model methods use different generation or understanding setups.Table 5 compares policy, target/reference, discreteness, and modality across representative methods.Source: §A.3 · pp. 20, 22, Table 5 · p. 21
  • Omni-Diffusion-Distill uses one on-policy framework for both modalities, with modality-specific states, teacher predictions, and matching terms.A pairwise collision penalty targets repetition in parallel text decoding, while entropy-matched guidance addresses image degradation from step reduction.Source: §A.3 · p. 22

B ALGORITHMS AND LOSS DETAILS

  • Stage 1 replays cached teacher trajectories, using keyframes and sparse teacher predictions to train the student to skip decoding steps.The cache stores top-R teacher probabilities and one residual category; the student and teacher are compared over the same R+1 categories.Source: B ALGORITHMS AND LOSS DETAILS · pp. 22–23
  • Stage 1 also trains on rebuilt intermediate states and freshly masked outputs, while the image update uses a categorical surrogate because discrete sampling is not differentiable.The ground-truth term weights by 1/π and divides by output length, making it an unbiased estimate of full masked cross-entropy.Source: B ALGORITHMS AND LOSS DETAILS · p. 23
  • Entropy-matched guidance solves for a temperature at each scored position, clipping solutions outside [0.25, 8].The preference term trains the student to favor the teacher’s completion over a sample from a periodically refreshed student reference.Source: B ALGORITHMS AND LOSS DETAILS · p. 24
  • At inference, the student generates images without an unconditional forward pass and decodes understanding block by block with exact KV-cache reuse.The block-causal mask enables cache reuse; the teacher’s fully bidirectional attention does not.Source: B ALGORITHMS AND LOSS DETAILS · p. 24

C.1 TRAINING DATA AND TEACHER TRAJECTORIES

The training mixture spans text-to-image generation, image-to-image tasks, and multimodal understanding, with teacher trajectories cached once for reuse. The reported deployed-budget speedups differ across image generation and understanding, and baselines share the teacher, initialization, data, and update count.

C.1 TRAINING DATA AND TEACHER TRAJECTORIES

  • The 51K-example training mixture contains 11K text-to-image, 22K image-to-image, 13K multimodal-understanding, and 5K inpainting examples.Inpainting pairs are built from text-to-image images and used only as extra training data.Source: §C.1 · p. 24
  • Teacher trajectories are collected once over the mixture and reused for every Stage-1 run; image trajectories use 64 steps with classifier-free guidance and understanding uses 512 steps without guidance.The full-step image teacher in the main evaluation uses 128 steps.Source: §C.1 · p. 24, Table 7 · p. 25

C.2 OPTIMIZATION AND COST

  • The two training stages share the data mixture, batch size, and ground-truth anchor, but differ in losses and optimization settings.The collision penalty applies only to the text branch; the auxiliary network and student reference are used only in Stage 2 and discarded afterward.Source: §C.2 · p. 24, Table 6 · p. 25
  • Understanding speeds up 21.2× and text-to-image generation 18.2× at deployed budgets, with exact KV-cache reuse contributing to the understanding speedup.Without cache reuse, understanding is 6.2× faster; text-to-image speedup is below its 32× NFE reduction.Source: §C.2 · p. 24

C.3 BASELINES

  • Baselines share the Lumina-DiMOO teacher, initialization, training mixture, optimizer-update count, sampler, and decoding budgets, while retaining their published objectives.Di[M]O, CDLM, and T3D are each applied to the modality beyond their original design coverage.Source: §C.3 · p. 24

C.4 BENCHMARKS

The benchmark suite evaluates text-to-image generation, image-to-image editing and controllable generation, and multimodal understanding with discriminative and open-ended tasks.

C.4 BENCHMARKS

  • Text-to-image evaluation covers GenEval, DPG-Bench, UniGenBench, OneIG English prompts, and TIIF-Bench.Each benchmark uses its official metric.Source: §C.4 · p. 25
  • Image-to-image evaluation uses the I2I benchmark and Graph-200K, while multimodal understanding spans discriminative and open-ended benchmarks.MMU covers POPE, MME-P, MMBench, SEED, MMMU, MM-Vet, Vibe-Eval, LLaVA-Bench, COCO captioning, and DeCapBench.Source: §C.4 · p. 25
  • COCO captioning uses a paper-defined protocol rather than its official metric; the other listed benchmarks use their official metrics.The open-ended set includes COCO captioning alongside MM-Vet, Vibe-Eval, LLaVA-Bench, and DeCapBench.Source: §C.4 · p. 25

C.5 OPEN-ENDED JUDGING PROTOCOLS

Open-ended benchmarks are judged by GPT-5.5, with benchmark-specific protocols and a custom procedure for COCO captioning. The appendix extends the reported results to smaller decoding budgets and additional T2I categories and samples.

C.5 OPEN-ENDED JUDGING PROTOCOLS

  • All five open-ended benchmarks use the same judge; four retain their official grading prompt and aggregation, while COCO captioning uses a protocol defined by the authors.For COCO, the judge rates accuracy, completeness, and coherence from 1 to 5, with five reference captions provided as context, not ground truth.Source: §C.5 · pp. 25–26
  • MM-Vet uses its official soft-grading prompt, while Vibe-Eval rescales judge ratings from 1–5 to 0–100.LLaVA-Bench uses a pairwise relative score, and DeCapBench averages two fact-level F1 scores.Source: §C.5 · pp. 25–26
  • Discriminative benchmarks are parsed by standalone yes/no answers or option letters, making their scores insensitive to the rest of the output.Multiple-choice parsing falls back to matching the full option string when no standalone letter is found.Source: §C.5 · p. 26

D ADDITIONAL RESULTS

  • Additional results extend the main section with smaller decoding budgets, per-category T2I results, and more qualitative samples.Inference cost is timed on one A100-80GB at batch size 1 in bfloat16, over the token-space sampler only.Source: D ADDITIONAL RESULTS · p. 26, Table 8 · p. 26

D.1 RESULTS AT SMALLER DECODING BUDGETS

At smaller decoding budgets, Omni-Diffusion-Distill generally retains strong results across text-to-image generation, image-to-image generation, and multimodal understanding, though some task-specific exceptions emerge.

D.1 RESULTS AT SMALLER DECODING BUDGETS

  • On T2I, Omni-Diffusion-Distill leads every benchmark at every budget, with its margin generally growing as the budget shrinks.On discriminative MMU it is the best distilled method on all five benchmarks, with scores flat down to 16 steps.Source: §D.1 · p. 26
  • On open-ended MMU, scores fall for every method, but Omni-Diffusion-Distill keeps the largest COCO captioning margin while T3D leads LLaVA-Bench at 32 and 16 steps.On I2I, it leads Add at both budgets and Graph-200K at 4 steps, while competitors lead Overall by at most 0.07.Source: §D.1 · p. 26
  • At 2 steps on I2I, T3D leads Graph-200K, showing that Omni-Diffusion-Distill does not lead every benchmark at every budget.On the I2I benchmark Overall, T3D leads at both budgets and Di[M]O at 4 steps, by at most 0.07.Source: §D.1 · p. 26

D.2 RESULTS IN EACH CATEGORY OF THE T2I BENCHMARKS

The section reports category-level T2I results at 8 decoding steps, additional qualitative samples, and analysis of the distillation components and design choices.

D.2 RESULTS IN EACH CATEGORY OF THE T2I BENCHMARKS

  • The category results cover GenEval, DPG-Bench, UniGenBench, OneIG, and TIIF-Bench under each benchmark’s official protocol.Tables 13–17 report category scores; TIIF-Bench also compares short and long prompts.Source: §D.2 · pp. 26–27
  • At 8 decoding steps, Omni-Diffusion-Distill leads every overall score and has its largest margins in several categories requiring multiple prompt constraints.It matches or exceeds the full-step teacher in some categories, while text rendering and DPG-Bench Global remain the main gaps.Source: §D.2 · p. 27

D.3 ADDITIONAL QUALITATIVE SAMPLES

  • Additional examples are provided: Figure 8 presents more T2I samples in the same format as Figure 4, while Figure 9 presents controllable-generation samples.Source: §D.3 · p. 27

E ANALYSIS AND ABLATIONS

  • The analysis section checks whether the collision penalty and entropy-matched guidance address the two identified challenges and examines which design choices affect results.The cited analysis also includes smaller-budget results for discriminative and open-ended MMU.Source: E ANALYSIS AND ABLATIONS · p. 27, Table 11 · p. 28, Table 12 · p. 28

E.1 WHY THE COLLISION PENALTY AND ENTROPY-MATCHED GUIDANCE HELP

The analysis examines repetition in parallel text decoding and entropy collapse in image distillation, finding that the two targeted mechanisms address these distinct problems.

E.1 WHY THE COLLISION PENALTY AND ENTROPY-MATCHED GUIDANCE HELP

  • Loop@5 rises from 0.6% with full-step teacher decoding to 74.4% at the student budget, while Omni-Diffusion-Distill reduces it to 7.7%.Without the collision penalty, Loop@5 is 21.6%; the penalty directly reduces repetition under parallel text decoding.Source: §E.1 · p. 30
  • Without entropy-matched guidance, final student image entropy is 1.55, below the few-step teacher’s raw CFG target of 2.28; guidance raises it to 3.32.Entropy is measured at the final decoding step over masked image positions and samples.Source: §E.1 · pp. 29–30
  • Entropy-matched guidance tracks 0.7 times conditional entropy during decoding, whereas fixed temperatures fall below that target early and exceed it near the end.Fixed-temperature scaling also lowers benchmark scores.Source: §E.1 · p. 31
  • The collision penalty’s largest gain is on COCO captioning, while entropy-matched guidance’s largest gain is on the I2I benchmark, consistent with their distinct targets.The collision penalty acts on the probability that two scored positions commit the same token; entropy-matched guidance acts on teacher-prediction entropy.Source: §E.1 · p. 31

E.2 ABLATIONS

Ablations show that both training stages and the two targeted mechanisms matter, while also distinguishing useful design choices from alternatives that underperform.

E.2 ABLATIONS

  • Using Stage 1 as a cold start for Stage 2 improves every benchmark; without Stage 2, scores are 14.8 points lower on DPG-Bench and 21.1 lower on COCO captioning.The largest reported gain is 14.4% on OneIG; gains include MM-Vet rising from 19.4 to 20.0 and COCO captioning from 56.5 to 57.2.Source: §E.2 · pp. 32–33
  • With the pairwise collision penalty, scoring the sampler’s own masked positions improves MM-Vet and COCO captioning, matches Vibe-Eval, and lowers DeCapBench.The compared alternative re-masks tokens the student has already committed.Source: §E.2 · pp. 33–34
  • Sequential teacher conditioning raises Vibe-Eval from 6.78 to 8.09 and COCO captioning from 47.3 to 49.8, but the largest improvement comes from the pairwise collision penalty.The collision penalty raises COCO captioning to 57.2 and DeCapBench to 19.6; an adjacent-repeat penalty brings almost no improvement.Source: §E.2 · p. 34
  • The entropy fraction ρ = 0.7 performs best on seven of eight columns, with the largest gain on I2I, where the score rises from 2.04 to 2.31.At ρ = 0.9, GenEval falls to 0.176 and ImageReward is negative, indicating that further softening hurts generation quality.Source: §E.2 · p. 34
  • Fixed-temperature scaling scores below Off on every evaluated benchmark, with GenEval falling from 0.798 to 0.475 at ν = 1.4 and to 0.015 at ν = 2.0.At ν = 2.0, DPG-Bench is 26.2 and ImageReward is −2.11; the I2I scores are 2.03 and 1.98 versus 2.31 with entropy-matched guidance.Source: §E.2 · pp. 34–35

Related papers

Loading 2610.10990v1…