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

Shufan Li; Hanrong Ye; Greg Heinrich; Yonggan Fu; Haotian Zhang; Sifei Liu; Aditya Grover; Pavlo Molchanov

cs.CVOct 27, 2026

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

Authors: Shufan Li, Hanrong Ye, Greg Heinrich, Yonggan Fu, Haotian Zhang, Sifei Liu, Aditya Grover, Pavlo Molchanov

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.

Figures & tables

Figure 1: (Left) NLD-Uni introduces a new paradigm for building unified Diffusion LLMs through on-policy distillation from cross-paradigm teachers: VLMs for visual understanding and continuous diffusion models for visual generation. (Right) NLD-Uni achieves highly competitive visual understanding and generation performance. Visual understanding performance is averaged across tasks in Table 1 . Text-to-image and image editing reports GenEval and ImgEdit Score. Multimodal reasoning is measured on MathVista. Axes are normalized by the best model’s performance on each task.

Figure 2: NLD-Uni supports diverse multimodal tasks by adopting an Elastic-MoT architecture with three elastic views of a single model. The VLM view uses the 8B understanding branch; the T2I view combines 4B understanding parameters with the 8B generation branch; and the interleaved view uses both full branches (16B). Each view loads only the parameters needed for its task. Yellow and green blocks denote understanding and generation layers, gray blocks indicate omitted layers, and blue arrows denote joint attention. Representative examples demonstrate visual reasoning, object grounding, text-to-image synthesis, and image editing guided by reasoning and localization.

Figure 3: Overall Framework of NLD-Uni. NLD-Uni is a unified diffusion model whose input and output are both multimodal. Images in the input stream are represented by both continuous embeddings of a semantic encoder and discrete tokens of a VQ encoder. Images in the output stream are represented solely using VQ codes. NLD-Uni is trained to predict the clean token distribution given partially masked inputs. For image generation, we incorporate an intermediate masking state.

Figure 4: Design of Self-Refinement Process. (Left) We extend the state space of mask diffusion with token age

a

with

a=0

being the clean state and

a=K

being the masked state. (Middle) Instead of one-step jump, each token is progressively masked through the interval

(\tau,\tau^{\prime})

with length

\epsilon

. Colors encode different tokens. (Right) We implement the progression from age

a=0

a=K

through random jumps based on a ball-tree in the embedding space. Token distributions are diffused from a one-hot vector to balls with growing sizes. This example has tree depth

D=4

Figure 5: On-Policy Distilling from Autoregressive VLMs. Top: Given a prompt and clean prefix blocks, we generate the current block with the student’s AR view by temporarily adopting a block size of 1. Left: AR on-policy distillation loss aligns the student’s AR distribution with a frozen AR teacher. Right: Consistency loss aligns the student’s diffusion view to its own AR view. The diffusion prediction is obtained by evaluating the student on randomly masked AR rollouts. We apply stop-grad to the AR target.

Figure 6: On-Policy Distilling from Continuous Flow Models. Given a student output, we generate a teacher-refined output by adding noise to the student sample and re-denoise it with the teacher model. For reference, we also show an independent teacher generation. We observe that the teacher’s refinements are closely aligned to the student’s output while having better visual fidelity and source image preservation.

Model	MMMU	RealWorldQA	ChartQA	DocVQA	InfoVQA	TextVQA	AI2D	RefCOCO	Sci.QA	ScreenSpot-Pro	OCRBenchV2	Avg
Split	val	test	test	val	val	val	test	val	test (img)	test	test (EN)	Avg
Autoregressive Models
LLaVA-1.6-7B [ 52 ]	35.1	-	64.9	74.4	37.1	-	66.6	-	73.2	-	-	-
Qwen2.5-VL-7B [ 53 ]	58.6	-	84.9	82.6	-	-	83.9	-	-	-	-	-
Qwen3-VL-8B [ 54 ]	69.6	71.5	89.6	95.6	83.4	82.9	84.8	87.3	75.2	54.6	64.1	78.1
Cosmos3-Nano [ 10 ]	60.9	72.2	85.5	94.5	81.5	83.4	85.0	79.2	93.0	56.1	65.9	77.9

Table 1: Quantitative results on image-level understanding tasks. NLD-Uni is evaluated using a confidence-based diffusion sampler following [ 51 ] .

Model	MathVista	DynaMath	MathVision	MathVerse	Geometry3k	LogicVista	We-Math
Autoregressive Models
Qwen2.5-VL-7B [ 53 ]	68.2	55.7	22.4	39.0	37.1	42.4	33.3
Qwen3-VL-8B [ 54 ]	77.2	62.2	37.9	67.7	57.5	50.1	69.0
BAGEL [ 5 ]	73.1	-	24.1	-	-	-	45.8
Discrete Diffusion Models
Fast-dVLM [ 57 ]	57.3	51.3	20.2	33.7	27.0	37.1	41.0

Table 2: Quantitative results on multi-modal reasoning benchmarks. *Denotes naively applying teacher score to student’s diffusion view using OPD-LM’s loss.

			GenEval							DPG	HPSv3
Model	Family	Params	Single Obj. $\uparrow$	Two Obj. $\uparrow$	Counting $\uparrow$	Colors $\uparrow$	Position $\uparrow$	Color Attri. $\uparrow$	Overall $\uparrow$	Overall $\uparrow$	Score $\uparrow$
Gen. Only
SD3-Medium [ 63 ]	Cont.	2B	0.99	0.94	0.72	0.89	0.33	0.60	0.74	84.1	-
FLUX.1-dev [ 62 ]	Cont.	12B	0.98	0.81	0.74	0.79	0.22	0.45	0.66	83.8	10.88
Qwen-Image [ 39 ]	Cont.	20B	0.99	0.92	0.89	0.88	0.76	0.77	0.87	88.3	10.91
Meissonic [ 25 ]	Disc.	1B	0.99	0.66	0.42	0.86	0.10	0.22	0.54	-	-

Table 3: Text-to-Image Generation Performance on GenEval, DPG, and HPSv3.

\dagger

LLaDa-2-Uni incorporates a pretrained 6B continuous diffusion decoder. *Half of the understanding layers are not loaded in T2I tasks due to the Elastic-MoT design.

Model	Add	Adjust	Extract	Replace	Remove	Action	Background	Style	Hybrid	Overall
Other Methods
Qwen-Image [ 39 ]	4.38	4.16	3.43	4.66	4.14	4.69	4.38	4.81	3.82	4.27
GPT-4o [ 1 ]	4.61	4.33	2.90	4.35	3.66	4.89	4.57	4.93	3.96	4.20
LLaDa2-Uni [ 58 ]	3.76	4.16	2.40	4.04	3.82	4.42	4.07	4.60	3.97	3.92
Qwen2.5VL+Flux [ 67 ]	4.07	3.79	2.04	4.13	3.89	4.52	3.90	4.84	3.04	3.80
FluxKontext dev [ 68 ]	3.76	3.45	2.15	3.98	2.94	4.26	3.78	4.38	2.96	3.52

Table 4: Per-Category and overall scores on ImgEdit benchmark. *Off-policy indicates we use an independent teacher output as the target, as opposed to teacher-refined student output.

Figure 7: Ablation Studies on the Self-Refining Diffusion Process. We report performance on GenEval (a) DPG (b) and HPSv3 (c) benchmarks at different training steps using different paradigm.

Figure 8: Qualitative Results on Text-to-Image Tasks . Results are sampled at $1024\times 1024$ resolution.

Figure 9: Qualitative Results on Image Editing Tasks . For each input-output pair, we present the input image on the left and the output image on the right.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 10: The design of SBD formulation. (Left) SBD adopts a two-stream training process that concatenates clean blocks and noisy blocks in the sequence dimension. It applies AR CE loss on the clean blocks, and diffusion CE loss on the noisy blocks. (Right) SBD models use a specialized attention mask with a causal mask in the clean blocks and block-bidirectional mask in the noisy blocks. Each noisy block can attend to all other tokens in itself and all tokens in preceding clean blocks.

Figure 11: Visualization of Ball-Tree-Based Corruption Process. (Right) Ball-Tree is a balanced binary tree whose leaf nodes correspond to unique tokens in the vocabulary. The forward corruption process

U_{k|k-1}

increases the token age

a

and decreases the node depth

d

. (Left) Each node has a geometric representation which is a high-dimensional ball in the VQ embedding space. Decreasing depth

d

corresponds to increasing ball radius. In this visualization, the depth

D=4

and maximum token age

K=3

Figure 12: Visual Reconstruction Across Different Token Ages. To investigate whether the proposed age-based corruption process is semantically meaningful for visual generation tasks, we show reconstructed images at different token ages and compare them with the rectified flow process in the continuous latent space of Qwen-VAE. Results indicate that changes in the token age represent a smooth corruption process in the pixel space. In this visualization, all tokens are corrupted at the same rate, which corresponds to $\epsilon=1$ .

Figure 13: Image PSNR Across Different Token Ages. To investigate whether the proposed age-based corruption process is semantically meaningful for visual generation tasks, we plot the PSNR of reconstructed images at different token ages and compare them with the rectified flow process in the continuous latent space of Qwen-VAE. Results indicate that changes in the token age represent a smooth corruption process in the pixel space measured by PSNR. In the clean state, Qwen-VAE has better PSNR than NLD-Uni, suggesting its continuous VAE has better reconstruction capability.

Figure 14: Additional Qualitative Illustrations of the Teacher Refinement Process. We show additional qualitative samples of teacher-refined student outputs. They generally align with the student roll-outs but have better visual fidelity and source image preservation. In some examples where the teacher over-edits, the refined student outputs are preferable to pure teacher generation.

	Stage 1	Stage 2	Stage 3
Learning Rate	$1\times 10^{-5}$	$1\times 10^{-5}$	$2\times 10^{-6}$
Steps	400k	100k $\times$ 2	10k
$\beta_{1}$	0.99	0.99	0.99
$\beta_{2}$	0.999	0.999	0.999
optimizer	AdamW	AdamW	AdamW
Dataset Used	A	A,B,C	B and subset of A, C

Appendix

Table 5: Training configurations across three stages. We use letters A-C to represent different dataset.

Figure 15: Comparison with MDM-Prime. Progression of Generated Images at Different Training Steps. The prompt is "Highly stylized, whimsical illustration of a fox wearing a tiny monocle and top hat, reading a miniature book in a cozy, mushroom-filled forest glade. Children’s book art, gentle colors, charming, adorable."

Actual NFE	$\boldsymbol{N=1}$	$\boldsymbol{N=4}$	$\boldsymbol{N=8}$	$\boldsymbol{N=16}$
4	7.3700	-2.4895	—	—
8	9.6214	9.0038	0.0834	—
16	10.5449	10.8205	10.3044	3.0348
32	10.8921	11.0137	10.9894	10.7620
48	10.9378	11.0369	11.1167	11.0861
64	11.0051	11.0287	11.1050	11.1010

Appendix

Table 6: Effect of Refinement Steps Per Token. We report HPSv3 scores across actual NFE and number of refinement stages

N

Figure 16: Ablation Studies on distilling from AR VLMs . We report MathVista Acc throughout training steps.

	AI2D	ChartQA	DocVQA	TextVQA	OCRBench	MMMU	MathVista
	Test	Test	Val	Val		Dev/Val	Mini
Teacher Performance
Nemotron Teacher	88.44	89.80	95.73	85.06	88.20	70.67	82.70
Kimi Teacher	85.14	94.72	97.46	83.93	92.70	77.44	88.90
NLD-Uni w/ OPD
OPD w/ Nemotron	85.98	90.28	95.01	77.82	85.30	61.11	78.10

Appendix

Table 7: Ablation study of teacher choices. Results are sampled in AR view with a block size of one.

Temperature	Diffusion Performance	AR Performance
0.0	73.9	76.9
0.4	74.8	76.0
0.6	74.6	75.4
1.0	73.8	75.1

Appendix

Table 8: Effect of OPD rollout temperature. We vary the sampling temperature of student AR rollouts during OPD and report MathVista accuracy.

Method	$\mathbf{r=0.5}$	$\mathbf{r=0.8}$	$\mathbf{r=0.9}$	$\mathbf{r=1.0}$
ImgEdit	4.06	4.09	4.17	3.90

Appendix

Table 9: Ablation on Noise Ratio . Image Editing Results across

r

values

Mode / Threshold	Tokens/Forward	Tokens/s/GPU	MathVista (%)
AR	1.00	202.74	76.1
Diffusion 0.95	3.94	463.84	75.2
Diffusion 0.90	4.64	552.58	75.0
Diffusion 0.60	8.31	997.95	73.4
Diffusion 0.50	7.23	734.03	70.8

Appendix

Table 10: Speed-Quality Tradeoff on MathVista Benchmark. Results are measured on H100 GPUs.

Diffusion VLMs

See all 16 papers

LLaDA2.0-Uni: Unifying Multimodal Understanding and Generation with Diffusion Large Language Model
Apr 22, 2026Inclusion AI, Tiwei Bie, Haoxing Chen +15Unified Multimodal Models Multimodal Generation
AGI Research Center, Inclusion AI
Omni-Diffusion-Distill: Few-Step Distillation of Unified Multimodal Diffusion Large Language Models
Oct 7, 2026Hong Huang, Chenhongyi Yang, Junzhe Sun +3Diffusion Language Model Distillation Diffusion Model Inference Acceleration
Simon Fraser University · Meta
Ask, Solve, Generate: Self-Evolving Unified Multimodal Understanding and Generation via Self-Consistency Rewards
Jun 25, 2026Ritesh Thawkar, Shravan Venkatraman, Omkar Thawakar +5Unified Multimodal Models Multimodal Generation
1Mohamed bin Zayed University of Artificial Intelligence · 4Linköping University · 3Australian National University +1

Apr 22, 2026J/K move · Enter open · S save

LLaDA2.0-Uni: Unifying Multimodal Understanding and Generation with Diffusion Large Language Model

Inclusion AI, Tiwei Bie, Haoxing Chen, Tieyuan Chen+14

AGI Research Center, Inclusion AI

Unified Multimodal Models Multimodal Generation Diffusion-Based Image Generation Multimodal Foundation Models Diffusion Language Models Image Editing+1

Open paper AskPDF arXiv

We present LLaDA2.0-Uni, a unified discrete diffusion large language model (dLLM) that supports multimodal understanding and generation within a natively integrated framework. Its architecture combines a fully semantic discrete tokenizer, a MoE-based dLLM backbone, and a diffusion decoder. By discretizing continuous visual inputs via SigLIP-VQ, the model enables block-level masked diffusion for both text and vision inputs within the backbone, while the decoder reconstructs visual tokens into high-fidelity images. Inference efficiency is enhanced beyond parallel decoding through prefix-aware optimizations in the backbone and few-step distillation in the decoder. Supported by carefully curated large-scale data and a tailored multi-stage training pipeline, LLaDA2.0-Uni matches specialized VLMs in multimodal understanding while delivering strong performance in image generation and editing. Its native support for interleaved generation and reasoning establishes a promising and scalable paradigm for next-generation unified foundation models. Codes and models are available at https://github.com/inclusionAI/LLaDA2.0-Uni.

Abstract

Diffusion VLMs

Explore similar work

LLaDA2.0-Uni: Unifying Multimodal Understanding and Generation with Diffusion Large Language Model

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

Ask, Solve, Generate: Self-Evolving Unified Multimodal Understanding and Generation via Self-Consistency Rewards