Nemotron-Labs-Diffusion-Uni: Elastic Self-Refining Unified Diffusion LLM with Cross-Paradigm On-Policy Distillation

1. Introduction

• Elastic-MoT presents three task-specific views of one model: an 8B understanding branch, a mixed 4B-understanding/8B-generation view, and a 16B interleaved view.Each view loads only the parameters needed for its task.Source: Figure 2 · p. 2
• The paper frames unified DLLMs as a simpler shared-objective alternative to AR-plus-diffusion systems, while noting their remaining performance gap to specialized models.Prior unified multimodal models commonly use discrete-token autoregressive modeling or combine autoregression for text with continuous diffusion for images.Source: §1 · pp. 1–2
• Cross-paradigm OPD transfers visual-understanding knowledge from an autoregressive VLM teacher and visual-generation knowledge from a continuous flow teacher to a unified DLLM.This approach addresses the lack of strong-performing MDM teachers.Source: §1 · p. 2
• NLD-Uni inserts intermediate states between mask and clean tokens, enabling multi-step transitions, finer inter-token interaction, and progressive self-correction for visual generation.This replaces the conventional one-step masked-to-clean transition.Source: §1 · pp. 2–3
• NLD-Uni achieves state-of-the-art performance among unified DLLMs across diverse multimodal tasks and greatly reduces the gap to specialized AR VLMs and continuous flow models.Evaluations include visual question answering, object grounding, multimodal reasoning, text-to-image generation, and image editing.Source: §1 · p. 3

2.1. Diffusion Language Models

• SBD gradually masks clean text blocks and predicts each clean block conditioned on previous clean blocks; inference samples blocks sequentially and tokens within each block in parallel.NLD-Uni applies this paradigm to multimodal modeling, treating each image as a resolution-dependent integral block.Source: §2.1 · p. 3

2.2. Masked Discrete Image Models

• Masked image models generate VQ-coded patches by unmasking them, but prior refinement approaches rely on heuristic confidence thresholds and retain a training-inference mismatch.Some methods prioritize less-correlated tokens or revise already-unmasked tokens.Source: §2.2 · p. 3

2.3. On-Policy Distillation (OPD)

• Prior OPD work transfers teacher supervision to student-generated sequences, while NLD-Uni extends the approach to jointly distill from an AR VLM and an image-generation teacher into a unified DLLM.The framework represents both continuous semantic embeddings and discrete VQ tokens for input images, while output images use VQ codes.Source: §2.3 · pp. 3–4, Figure 3 · p. 4

3.1. Model Architecture

• NLD-Uni uses an understanding expert tower and a generation expert tower initialized from the same diffusion language model, with joint attention between them.Images are encoded with both a semantic encoder and a VQ encoder.Source: §3.1 · p. 4

3.2. Discrete Diffusion with Progressive Self-Refinement

• Progressive self-refinement adds token age states between clean and masked, corrupting tokens through intermediate states rather than a direct jump.A per-token schedule sets the refinement interval; a ball tree supports progressively broader token substitutions in embedding space.Source: §3.2 · pp. 4–6, Figure 4 · p. 5

3.3. Distillation of Visual Understanding Ability

• For visual understanding, the student’s AR view receives teacher supervision, while its diffusion view is trained to match its own AR predictions on partially masked rollouts.The method uses a transitive objective because directly copying AR teacher scores to the diffusion view can produce ill-formed supervision under context mismatch.Source: §3.3 · pp. 6–7

3.4. Distillation of Visual Generation Ability

• For visual generation, a continuous-flow teacher re-denoises noisy student outputs into teacher-refined targets, which retain alignment with student generations while improving fidelity and source-image preservation.The student output is re-encoded in the teacher’s VAE; the paper trains on these targets with canonical SFT loss, adapting supervision to the student’s rollouts.Source: §3.4 · pp. 7–8, Figure 6 · p. 7

4.1. Main Results

• Training uses three stages, ending with teacher distillation for multimodal reasoning and image editing while retaining text-to-image SFT batches.No text-to-image OPD teacher is used because the base model outperforms the cited teacher candidates on GenEval and HPSv3; training runs on 64 H100 GPUs for 26 days.Source: §4.1 · p. 8
• NLD-Uni reports strong performance across visual understanding, multimodal reasoning, text-to-image generation, and image editing, with state-of-the-art results among diffusion language models on reasoning tasks.The reported comparisons include improvements over previous unified diffusion models on understanding, gains over the pre-OPD checkpoint on reasoning, and reduced gaps to specialized models on generation and editing.Source: §4.1 · pp. 8–9
• Cross-paradigm OPD improves performance across reasoning tasks, with the most notable gains on We-Math, Geometry3k, and MathVerse; naive OPD-LM-style score copying worsens the base model.Table 2 reports quantitative results on multimodal reasoning benchmarks and marks the naive teacher-score baseline with an asterisk.Source: §4.1 · p. 8, Table 2 · p. 9
• NLD-Uni outperforms existing unified diffusion models on GenEval, DPG, and HPSv3, and surpasses Flux.1-dev and Qwen-Image on HPSv3.Table 3 covers the three text-to-image benchmarks; the prose identifies HPSv3 as a visual-quality measure.Source: §4.1 · pp. 8–9, Table 3 · p. 9
• NLD-Uni outperforms discrete diffusion models on ImgEdit, improves by +0.36 over NLD-Uni-Base, and exceeds the off-policy setup by +0.27.It underperforms only GPT-4o and the Qwen-Image teacher among the broader model families discussed; Table 4 reports per-category and overall ImgEdit scores.Source: §4.1 · p. 9, Table 4 · p. 10

4.2. Ablation Studies

• Across a 20k-step run, self-refining diffusion performs best on average and consistently beats uniform diffusion on DPG and HPSv3, while remaining comparable on GenEval.All methods are evaluated with 64 diffusion steps; the largest improvements are on HPSv3, which the authors relate to visual fidelity.Source: §4.2 · p. 9, Figure 7 · p. 10

4.3. Qualitative Results

• The qualitative section presents additional text-to-image generation and image-editing examples as evidence of NLD-Uni’s generation capability.The examples are shown in Figures 8 and 9.Source: §4.3 · p. 10

5. Conclusion

• NLD-Uni combines cross-paradigm OPD and self-refining diffusion, and experiments demonstrate both approaches’ effectiveness in a state-of-the-art unified DLLM.The conclusion describes the model as capable of visual understanding and generation.Source: §5 · p. 10

A.1. Formulations of SBD Models

• SBD generates consecutive blocks sequentially, using within-block bidirectional attention to reduce the compute cost of fully bidirectional diffusion.Its training combines clean and noisy blocks with separate AR and diffusion cross-entropy losses.Source: §A.1 · p. 19, Figure 10 · p. 20

A.2. Formulations of Self-Refining Diffusion Process

• For image blocks, NLD-Uni uses intermediate mask states, with each token’s age increasing from clean state 0 toward mask state K.Text blocks use standard MDM, while the image process is extended to multi-block interleaved multimodal sequences.Source: §A.2 · p. 20, §A.2.1 · p. 20

A.2.1. Ball-Tree-Based Per-Token Corruption Process

• Ball-tree corruption progressively replaces a code with a uniformly sampled neighbor set of increasing radius, producing smooth semantic corruption in visualizations.The tree maps token age to progressively larger regions in the VQ embedding space; visual comparisons report smooth corruption for both this process and rectified flow.Source: §A.2.1 · pp. 21–22

A.2.2. Non-Markov Per-Token Jump Schedule

• A pre-sampled per-token transition interval enables intermediate states; its duration parameter ε→0 recovers canonical DNDM, which uses a single jump.The DNDM schedule pre-samples each token’s jump time and maintains it through the trajectory.Source: §A.2.2 · pp. 23–24

A.3. On-policy Distillation from AR VLMs

• SBD’s block-size-one AR view accepts standard AR OPD directly, while an added consistency loss supplies supervision for larger diffusion blocks.AR-only supervision would otherwise remove support for parallel generation with block sizes greater than one.Source: §A.3 · p. 25

A.3.1. Connection with OPD-LM

• NLD-Uni uses OPD within unified multi-objective training, unlike OPD-LM’s single-objective AR-to-diffusion language post-training.The unified run combines visual understanding, text-to-image, and image-editing objectives.Source: §A.3.1 · p. 25
• Small nonzero-temperature AR rollouts yield better diffusion performance, suggesting trajectory diversity can help despite relaxing the argmax property.At zero temperature, greedy tokens are AR conditional argmaxes, but that does not guarantee compatibility with the observed future context.Source: §A.3.1 · pp. 25–26
• Directly applying OPD-LM’s loss to heterogeneous teacher and student models performs poorly, and NLD-Uni instead uses the student’s own AR scores to supervise its diffusion view.The approach also samples rollouts from the AR view rather than the diffusion view.Source: §A.3.1 · pp. 25–26

A.4. On-policy Distillation from continuous flow models

• For image editing, the student generates an output that is decoded, re-encoded into the teacher’s continuous VAE, partially noised, and refined through multiple ODE steps.The noise strength r controls refinement, and the refined latent is converted back into discrete VQ codes for training.Source: §A.4 · pp. 26–27
• Teacher-refined outputs generally preserve the student’s output while adding detail, and can outperform pure teacher generation when the teacher over-edits.For the dining-table removal example, refinement constrained the teacher to a neighborhood of the student rollout, preserving extra chairs.Source: Figure 14 · p. 26, §A.4 · p. 27

B.1. Data Mixture

• The data mixture includes 127M filtered text-to-image images, visual-instruction and grounding data, and image-editing data; OPD uses 76.8k reasoning and 100k editing samples.The OPD subset is sampled from prompts with labels discarded.Source: §B.1 · p. 28

B.2. Optimization Setup

• Training initializes the understanding branch from NLD, copies it to the image-generation branch, then proceeds through image-text, instruction/editing, and OPD stages.Image-text pairs remain mixed into the final stage to preserve text-to-image ability; the stage-two checkpoint is NLD-Uni-Base and the final checkpoint is NLD-Uni.Source: §B.2 · p. 28

B.3. Infrastructure

• Training used 64 H100 GPUs across eight nodes for 26 days, including five days for the final OPD stage.The infrastructure passage gives the total duration and OPD-stage duration.Source: §B.3 · p. 28

C.1.1. Comparison with MDM-Prime

• MDM-Prime scores 0 on GenEval after 20k training steps and fails to converge at 1024 × 1024, though its outputs retain some object-like structure.Its visual-generation experiments are described as using 32 × 32 CIFAR-10 and ImageNet-32 images.Source: §C.1.1 · pp. 28–29

C.2. Effect of refinement steps

• At 64 inference steps, N=8 works best; N=1 favors small NFE, while N=16 adds no improvement and can hurt low-NFE performance.The study also reports that ε=1 works poorly at any tested total NFE.Source: §C.2 · p. 29, Table 6 · p. 30

C.3. Studies on distilling from AR VLMs

• The additional AR-VLM OPD ablations start from NLD-Uni-Base and post-train only multimodal reasoning, so their results are not directly comparable to the main joint-training result.This scope statement qualifies the appendix ablations rather than the main result.Source: §C.3 · p. 30

C.3.1. Training Objectives

• Only AR-OPD+Consistency and AR-OPD+SFT improve diffusion generation over baseline; other setups create mismatched supervision from mismatched context.Even student diffusion rollouts fail when the consistency target is the student AR score, which can favor tokens inconsistent with future diffusion-generated context; matching AR samples and targets mitigates this issue.Source: §C.3.1 · p. 31
• Consistency loss only collapses to 24.1 accuracy at 400 steps, as tying the student’s generation modes without external supervision quickly produces degenerate sequences.This run is outside the plotted range and is not shown in the figure.Source: §C.3.1 · p. 31

C.3.2. Sampling Temperature

• Greedy rollouts yield 76.9% AR accuracy, while temperature 0.4 yields 74.8% diffusion accuracy; strict per-token argmax alignment is unnecessary for effective diffusion transfer.Across tested temperatures, AR-OPD+Consistency has higher diffusion accuracy than alternatives such as AR-OPD+SFT.Source: §C.3.2 · p. 31
• Zero-temperature evaluation gives better accuracy even for models trained with non-zero rollout temperature.All evaluations use temperature 0 for consistency and reproducibility.Source: §C.3.2 · p. 31

C.3.3. Choice of Teachers

• Nemotron teacher outperforms the stronger Kimi teacher in this OPD setting, with the authors attributing the difference to Kimi’s larger teacher-student gap, different tokenizer, and different thinking-trace format.These early-exploration results use an earlier checkpoint and are not directly comparable to other ablations.Source: §C.3.3 · p. 31, Table 7 · p. 31

C.4. Studies on distilling from continuous flow models

• r=0.9 works best among tested noise strengths for image-editing OPD with the Qwen-Image teacher.The passage reports benchmark performance in Table 9.Source: §C.4 · pp. 31–32, Table 9 · p. 32
• Added Gaussian noise has a highly non-linear effect on the final image-editing output.Source: §C.4 · p. 32
• Even at r=0.9, generated images remain aligned with the student rollout and can differ substantially from a purely teacher-generated sample.Source: §C.4 · p. 32

C.5. Speed-Quality Tradeoff

• The speed-quality trade-off is measured using tokens per forward pass, throughput, and MathVista answer accuracy, and tuned via the unmasking confidence threshold.Source: §C.5 · p. 32, Table 10 · p. 32
• 552.58 tokens/s is reported for the vanilla setup, compared with 202.74 tokens/s for the AR baseline.Results are measured on H100 GPUs.Source: §C.5 · p. 32, Table 10 · p. 32
• Up to 997.95 tokens/s is achievable by relaxing the unmasking confidence threshold, with minor performance degradation.Source: §C.5 · p. 32, Table 10 · p. 32

D.1. Unified Architecture.

• Unified MLLMs use either a single dense transformer for image and text tokens or separate parameter sets for the modalities.Janus, MMaDa, and Show-O exemplify the first choice; BAGEL, MetaQueries, and LaViDa-O exemplify the second.Source: §D.1 · p. 32
• Elastic-MoT decouples attention layers in the final half of the model and enables joint attention only in early layers, extending an 8B base model to 16B.Source: §D.1 · p. 32
• Elastic-MoT loads 8B parameters for understanding, 8B generation plus 4B understanding parameters for generation, and all 16B for interleaved tasks.For understanding, this introduces no overhead compared with an 8B VLM.Source: §D.1 · pp. 32–33

• Generic multimodal understanding benchmarks are MMMU [42], RealWorldQA, ChartQA [43], DocVQA [96], InfoVQA, TextVQA [97], AI2D [45], RefCOCO, ScienceQA [98], ScreenSpot-Pro [47], and OCRBench-V2.Source: E. List of Benchmarks · p. 33
• Multimodal reasoning benchmarks include [100], [50], [101], DynaMath [102], Geometry3k [49], LogicVista [103], and We-Math.Source: E. List of Benchmarks · p. 33
• Multimodal generation benchmarks are GenEval [59], DPG [60], HPSv3 [61], and ImgEdit [65].Source: E. List of Benchmarks · p. 33

• Cross-paradigm OPD currently requires matching vocabulary spaces; the setup uses the Ministral 3 tokenizer, while cross-tokenizer methods are left for future work.Source: F. Limitations and Future Work · p. 33
• Further architecture optimization may be possible, but is outside this work’s focus.Source: F. Limitations and Future Work · p. 33
• The self-refining process uses a non-Markov formulation with predefined per-token corruption intervals and uniform progression through intermediate states.Source: F. Limitations and Future Work · p. 33
• Fine-grained text rendering is limited because such data is not widely present in the training mix; the authors plan to curate more of it.Source: F. Limitations and Future Work · p. 33