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

AI summary

Overview

Research area: Multimodal machine learning — specifically unified discrete diffusion large language models (DLLMs) that handle both visual understanding and visual generation in a single model.

Technical level: Advanced. The paper assumes familiarity with masked diffusion, vector quantization, block-wise semi-autoregressive decoding, and on-policy distillation.

Scope: The paper introduces NLD-Uni, a unified diffusion LLM trained via cross-paradigm on-policy distillation from autoregressive vision-language and continuous flow teachers, plus a progressive self-refining masking process for image generation.

What This Paper Is About

Unified diffusion language models offer a single training objective and parallel token generation for both text and images, but they lag well behind specialized autoregressive VLMs on understanding and behind continuous flow models on image generation. The authors' goal is to close that gap without abandoning the diffusion paradigm, by importing knowledge from strong teachers that operate under entirely different generation paradigms. A secondary goal is to fix a specific weakness of masked diffusion for images: once a token is unmasked it can never be revised, even though tokens in the same image are highly correlated.

Key Contributions

1. Cross-paradigm on-policy distillation (OPD). A framework that lets a unified DLLM learn from teachers with incompatible generation paradigms — an autoregressive VLM teacher for visual understanding and a continuous flow model teacher for image generation — despite there being no strong DLLM teachers to distill from.

2. A transitive distillation objective for understanding. Rather than copying autoregressive teacher scores directly onto the student's diffusion predictions (which the authors show produces ill-formed targets), the method supervises the student's autoregressive view with the teacher and then supervises the student's diffusion view with its own autoregressive view via a consistency loss.

3. A self-refining masked diffusion process. The state space is extended from "clean token or mask" to a token vocabulary crossed with a discrete token age, with intermediate states between clean and masked. During inference, tokens can be committed intermediately and progressively refined, enabling inter-token interaction that standard masked diffusion lacks.

4. The NLD-Uni model itself. A two-tower Elastic Mixture-of-Transformer with three elastic views of one model: an 8B understanding branch for the VLM view, a 4B understanding plus 8B generation combination for text-to-image, and both full branches (16B) for interleaved generation.

Main Findings

• Strongest unified DLLM on understanding: NLD-Uni reaches an average of 78.3 across eleven image-understanding benchmarks in Table 1, a +17.1 improvement over the previous diffusion state of the art. Individual gains include +10.3 on MMMU, +8.4 on InfoVQA, +7.2 on ScienceQA, +48.8 on ScreenSpot-Pro, and +23.2 on OCRBenchV2.

• Competitive with autoregressive VLMs: That 78.3 average edges past Qwen3-VL-8B's 78.1 on the same benchmark set, though the authors attribute the understanding gains mostly to scaled training data rather than to OPD.

• OPD substantially lifts multimodal reasoning: Against the pre-OPD NLD-Uni-Base checkpoint, cross-paradigm OPD adds +5.5 on MathVista, +6.4 on DynaMath, +9.6 on MathVision, +9.8 on MathVerse, +13.5 on Geometry3k, +3.5 on LogicVista, and +15.1 on We-Math, yielding 77.0, 74.7, 45.4, 45.3, 67.9, 55.9, and 74.9 respectively.

• Naive OPD actively harms performance: Applying an OPD-LM-style loss that copies teacher scores straight onto the diffusion view drops MathVista from 71.5 to 47.7 and MathVision from 35.8 to 17.9, confirming the context-mismatch problem the transitive objective is designed to solve.

• Text-to-image quality close to continuous models: NLD-Uni scores 0.89 overall on GenEval, 86.6 on DPG, and 11.12 on HPSv3. The HPSv3 score exceeds Qwen-Image's 10.91 and FLUX.1-dev's 10.88, and beats the hybrid LLaDa-2-Uni's 9.47.

• OPD improves image editing: On ImgEdit, NLD-Uni reaches 4.17 overall, a +0.36 gain over NLD-Uni-Base (3.81) and +0.27 over an off-policy knowledge-distillation baseline that uses independent teacher outputs (3.90). It trails only Qwen-Image (4.27) and GPT-4o (4.20).

• Self-refinement beats alternative masking schemes: Over a 20k-step run at 64 diffusion steps, the self-refining process is best on average, is comparable to uniform diffusion on GenEval, consistently better on DPG and HPSv3, and shows its largest advantage on HPSv3, suggesting the benefit is concentrated in visual fidelity.

• Ball-tree corruption is a memory necessity: With 131k distinct VQ codes, a full pairwise distance matrix would need 32GB of VRAM; the 18-depth binary tree storing the same structure takes 524 KB as a single 1D array.

Methodology in Plain English

The model is split into two "towers" that talk to each other through joint attention: one for understanding and one for generating images. Both start from the same diffusion language model weights. Images coming in are encoded twice — once as continuous semantic embeddings, once as discrete VQ tokens — while images going out use only VQ codes.

For the self-refinement idea, instead of a token being either clean or completely masked, each token gets an "age" from 0 (clean) to K (fully masked). During training, a token is corrupted in several steps, jumping to progressively more distant tokens in the codebook. Those jumps are organized by a ball tree over the embedding space, so early jumps land on tokens that are near neighbors and later jumps wander further away. A separate continuous-time schedule decides when each token performs its jumps, controlled by a hyperparameter ε that sets how long the intermediate period lasts; when ε shrinks to zero the method reduces to prior work (DNDM), and when ε equals 1 all tokens degrade at the same rate.

For distillation on understanding tasks, the block-wise diffusion formulation collapses to ordinary autoregressive prediction when the block size is 1. That gives the student an autoregressive "view" that can be scored by an autoregressive VLM teacher. The naive approach of also copying those teacher scores onto the diffusion view fails because the contexts differ — the paper's example is a teacher confidently predicting "New York" after "Jason visited," which would push the student toward the nonsense "New Francisco" given the partial state "[m] Francisco." The fix is transitive: the teacher supervises the student's autoregressive output, and the student's own autoregressive output (with gradients stopped) supervises its diffusion output.

For generation, distillation is harder because flow models live in continuous latent space. The authors take a student's generated image, re-encode it with the teacher's VAE, add noise, and re-denoise with the teacher conditioned on the original image and instruction. The result — a "teacher-refined student output" — stays aligned with what the student produced but has fewer artifacts and better source-image preservation, making it a usable training target under a standard supervised loss. The authors note this is only loosely "on-policy" in the conventional sense.

Training proceeds in three stages: text-image pairs for basic alignment and generation, then instruction following, editing, grounding, and interleaved data, then a final OPD stage using Nemotron-3 Nano Omni for reasoning and Qwen-Image for editing. No OPD teacher is used for text-to-image because the base model already beats Qwen-Image and Flux on GenEval and HPSv3. The full run takes 64 H100 GPUs over 26 days, with the OPD stage accounting for 5 of those days.

Why This Matters

The work suggests diffusion LLMs need not be permanently second-tier. By showing that a DLLM can absorb capabilities from teachers in different paradigms, it opens a route to improving unified diffusion models without waiting for a strong DLLM teacher to exist. The self-refining diffusion process also addresses a structural limitation — the independence assumption across simultaneously sampled tokens — that has constrained masked image models' fidelity.

Real-world applications implied by the evaluated tasks:

• Instruction-based photo editing: the ImgEdit results cover adding, adjusting, extracting, replacing, removing, background changes, style transfer, and hybrid edits.
• GUI automation agents: strong ScreenSpot-Pro scores (49.2, versus 0.0–0.4 for several diffusion baselines) support locating interface elements from screenshots.
• Visual document and chart assistants: high scores on ChartQA (89.9), DocVQA (94.2), InfoVQA (84.0), and OCRBenchV2 (62.1) map onto document parsing and data extraction.
• STEM tutoring from diagrams and figures: We-Math 74.9, Geometry3k 67.9, MathVista 77.0, and MathVision 45.4 correspond to reasoning over math imagery.

Industry relevance: A single model with elastic views (8B, effectively 4B+8B, and 16B) that loads only the parameters a task needs is attractive for serving cost, and the faster parallel decoding of diffusion is a direct efficiency argument against autoregressive serving pipelines. The distillation recipe is also a practical playbook for teams with strong proprietary teachers but constrained student architectures.

Future Directions

• Better data mixture and architecture optimization. The authors explicitly acknowledge that the model can be improved along these lines, framing current results as a starting point rather than a ceiling.

• Closing the remaining gap to continuous flow models on fine-grained generation. NLD-Uni leads on HPSv3 but still trails Qwen-Image on DPG (86.6 versus 88.3) and ImgEdit overall (4.17 versus 4.27), so the fidelity and instruction-following gap is narrowed, not eliminated.

• Extending cross-paradigm OPD to settings where no teacher exists. The authors deliberately skipped an OPD teacher for text-to-image because no candidate beat the base model on GenEval and HPSv3, leaving open how to improve when teachers are weaker than the student.

• Scaling the self-refining process beyond the tested regime. The ablation covers a 20k-step run at 64 diffusion steps; the paper notes prior partial-masking work (MDM-prime) only demonstrated results on 32x32 images and did not scale to this setting, which invites further study of how token-age depth K and the ε schedule behave at larger scale.

Target Audience

Researchers and engineers working on diffusion language models, unified multimodal models, or knowledge distillation. It is most useful to readers already comfortable with masked diffusion formulations and multimodal tokenization, who want a concrete recipe for transferring capabilities between mismatched model paradigms and a technically detailed alternative to standard absorning-mask diffusion for image generation. Practitioners evaluating whether to build a unified model rather than an autoregressive-plus-diffusion hybrid will find the head-to-head benchmark tables directly actionable.

Authors’ abstract

We propose Nemotron-Labs-Diffusion-Uni (NLD-Uni), a unified discrete diffusion large language model (DLLM) for visual understanding and generation tasks. Compared with prior unified diffusion models, NLD-Uni introduces two major innovations. First, we propose a unified cross-paradigm on-policy distillation (OPD) framework that successfully transfers the visual understanding ability of autoregressive VLMs and visual generation ability of continuous flow models, to a unified DLLM. Second, NLD-Uni incorporates a novel progressive self-refining diffusion process that gradually corrupts each token across multiple steps during training, as opposed to adopting the conventional masked diffusion formulation that directly corrupts a clean image token to an uninformative mask token. At inference, this design allows NLD-Uni to make intermediate commit to a token, enabling progressive self-refinement for better visual quality. Through extensive experiments, we show that NLD-Uni achieves state-of-the-art performance among unified DLLMs on diverse tasks including visual question-answering, multi-modal reasoning, object grounding, text-to-image generation, and image editing.