The Design Space of Tri-Modal Masked Diffusion Models

Louis Bethune; Victor Turrisi; Bruno Kacper Mlodozeniec; Pau Rodriguez Lopez; Lokesh Boominathan; Nikhil Bhendawade; Amitis Shidani; Joris Pelemans

cs.LGFeb 25, 2026

The Design Space of Tri-Modal Masked Diffusion Models

Authors: Louis Bethune, Victor Turrisi, Bruno Kacper Mlodozeniec, Pau Rodriguez Lopez, Lokesh Boominathan, Nikhil Bhendawade, Amitis Shidani, Joris Pelemans,

Organizations: Apple · University of Cambridge · MIT · Google Deepmind (work done at Apple)

Abstract

Discrete diffusion models have emerged as strong alternatives to autoregressive language models, with recent work initializing and fine-tuning a base unimodal model for bimodal generation. Diverging from previous approaches, we introduce the first tri-modal masked diffusion model pretrained from scratch on text, image-text, and audio-text data. We systematically analyze multimodal scaling laws, modality mixing ratios, noise schedules, and batch-size effects, and we provide optimized inference sampling defaults. Our batch-size analysis yields a novel stochastic differential equation (SDE)-based reparameterization that eliminates the need for tuning the optimal batch size as reported in recent work. This reparameterization decouples the physical batch size, often chosen based on compute constraints (GPU saturation, FLOP efficiency, wall-clock time), from the logical batch size, chosen to balance gradient variance during stochastic optimization. Finally, we pretrain a preliminary 3B-parameter tri-modal model on 6.4T tokens, demonstrating the capabilities of a unified design and achieving strong results in text generation, text-to-image tasks, and text-to-speech tasks. Our work represents the largest-scale systematic open study of multimodal discrete diffusion models conducted to date, providing insights into scaling behaviors across multiple modalities.

Figures & tables

Panel of Figure 1 — Figure 1: High-fidelity generation. The pretrain-only 3B MDM demonstrates strong prompt adherence alongside high-quality visual rendering of texture , lighting , and composition . Samples show: (a) natural daylight and depth of field ("egg in a field of crocuses"); (b) fine-grained fur texture in B&W ("lion’s face"); (c) soft, warm lighting with vintage color tones ("preparing bread dough"); and (d) complex multi-object arrangement ("noodle soup with toppings"). Extended generations in Appendix Appendix G .

Figure 10: Training curves for tri-modal MDM . We select 24 log-distributed isoFLOPS between 5e18 and 1e22 to cover the $(N,D)$ grid. The performance is dominated by the total compute budget $C$ , following formula in appendix H.1 of Busbridge et al. (2025) . The loss is computed on an independent validation set with identical mixture weights.

Figure 15: Tri-modal 3B overview across image, text, and audio. Left: radar summary (larger radius indicates better normalized score; vertex labels are raw values). Normalization is mixed by metric type: bounded metrics use natural bounds (GenEval/LM percentages to 100, WER to 1, Audiobox to 10), while unbounded Frechet metrics (FID-Incept, FID-DINOv2, FAD) use ECDF calibration from references, with lower-is-better inversion

s=1-F(x)

, where

F

is the ECDF. Right: raw unnormalized metrics.

Panel of Figure 16 — Figure 16: Loss contours for tri-modal mixture coefficients , taking $[1/3,1/3,1/3]$ as the reference point for the 0-level contour. We do not observe synergies between modalities at that model and data scale: they all compete for capacity and tokens.

Figure 17: Text-to-image hyperparameter ablations. We generate images from text prompts and compute FID against reference images on two datasets: CC12M in blue and train (eval seed) in orange. Top row (a–d) shows FID computed using DINOv2-L features, bottom row (e–h) shows FID using Inception-v3 features. Panels (a,e) vary generation steps; (b,f) vary CFG scale; (c,g) vary temperature; (d,h) vary top-

p

. Stars indicate optimal values for each metric-dataset combination.

Figure 18: Text-to-speech hyperparameter ablations. We evaluate audio generation from text prompts on two datasets (train set, LibriSpeech-PC) using ground-truth durations, with three metrics: FAD, WER, and Audiobox Aesthetics. Top row shows (a) steps sweep and (b) CFG sweep. Bottom row shows (c) temperature sweep and (d) top-

p

sweep.

\star

indicate the best hyperparameter for each metric. For Production Complexity, lower values indicate simpler audio and are considered as preferable.

Model	FID (Inception) $\downarrow$		FID (DINOv2) $\downarrow$		FAD $\downarrow$
Model	Train Data	CC12M	Train Data	CC12M	Train Data	LibriSpeech
Base	18.69	26.77	306.19	395.73	0.24	0.79
Anti-mask	17.81	21.04	302.61	361.00	0.22	0.55

Table 1: Multimodal ablation results comparing standard MDM training versus anti-masking.

	OpenBookQA	TruthfulQA	BBH	MMLU	Winogrande	ARC-Easy	HellaSwag	LogiQA2	PIQA	ARC-Challenge
Base	32.80 $\pm$ 2.10	43.12 $\pm$ 1.45	21.52 $\pm$ 0.46	30.20 $\pm$ 0.39	52.01 $\pm$ 1.40	63.72 $\pm$ 0.99	52.20 $\pm$ 0.50	24.75 $\pm$ 1.09	67.19 $\pm$ 1.10	31.66 $\pm$ 1.36
Anti-mask	33.00 $\pm$ 2.10	36.97 $\pm$ 1.41	27.05 $\pm$ 0.50	32.86 $\pm$ 0.39	54.06 $\pm$ 1.40	66.04 $\pm$ 0.97	55.15 $\pm$ 0.50	26.27 $\pm$ 1.11	67.41 $\pm$ 1.09	34.90 $\pm$ 1.39

Table 2: Anti-masking results on the evaluation harness. We report mean accuracy

\pm

standard deviation.

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Model	# Codebooks	PESQ $\uparrow$	Content Enjoy $\uparrow$	Content Useful. $\uparrow$	Prod. Complex. $\uparrow$	Prob. Quality $\uparrow$	Down. Factor $\uparrow$
Higgs pretrained	8	3.168	5.670	6.192	1.559	6.510	960
Higgs pretrained	4	2.544	5.577	6.087	1.561	6.390	960
DAC pretrained	9	3.658	5.555	5.994	1.567	6.255	512
DAC pretrained	4	2.396	5.069	5.415	1.603	5.597	512
DAC retrained	9	2.969	5.691	6.234	1.556	6.550	1024
DAC retrained	4	2.433	5.585	6.131	1.566	6.434	1024

Appendix

Table 3: Reconstruction metrics for different audio tokenizers.

Model	Type	Latent / Code Dim	Vocab Size	Tokens Per Image	rFID $\downarrow$
CC12M
cosmos-ci8x8 ( Agarwal et al., 2025 )	Continuous	16	-	1024	1.37
cosmos-di16x16 ( Agarwal et al., 2025 )	FSQ	6	65536	256	3.50
cosmos-di8x8-360p ( Agarwal et al., 2025 )	FSQ	6	65536	1024	1.80
ibq-262144 ( Shi et al., 2025 )	IBQ	256	262144	256	0.89
movqgan-270m ( ai-forever, 2023 )	MoVQ	256	16384	1024	0.50

Appendix

Table 4: Reconstruction FID for different image tokenizers. ImageNet consists of 50k validation examples and CC12M ( Changpinyo et al., 2021 ) consists of 50k samples from the full dataset. Best FID per section is shown in bold . When a continuous model achieves the best FID, the best discrete model is also bolded. *UniTok uses 8 categorical predictions of size 4096 (x8 underlying codes that are merged into a single token).

Figure 19: Small-scale hyper-parameter search for 0-shot transfer . Final average cross-entropy loss for a sweep of $5$ AdamW hyperparameters (learning rate, $\beta_{1},\beta_{2}$ , weight-decay, $\epsilon$ ) on a 350M model trained for 13B tokens. The plot shows cross-sections through the hyperparameter landscape through the best found hyperparameters.

Panel of Figure 20 — Figure 20: Performance as function of physical batch size. . In a) we see impact of batch size on total runtime, with the effect of the node-count clearly visible: more nodes reduces wall-clock time faster than simply saturating GPUs. In b) we see the runtime with 1, 2, 4, 8 and 16 nodes respectively. Doubling node count approximately halves the runtime. Finally, in c) , we see that saturated GPUs have much more efficient FLOP usage than non-saturated ones. Doubling node count never allow to recover training efficiency of fewer nodes, because of communication overhead.

Model
$N$ blocks	$24$
Dimension	$3072$
$N$ attention heads	$24$
QK Norm	Yes
Normalization	RMSNorm
Pre-norm	Yes

Appendix

Table 5: Model and training details for the 3B multimodal MDM .

Figure 21: Per-module hyperparameter search on a small scale model for 0-shot transfer . Average final ELBO for a sweep of AdamW’s per-module hyperparameters as well as initialization scales for a 350M model (including 80M non-embedding parameters) trained for 13B tokens, using a batch size of 256. Per-module tuning yields a 1.81

\times

reduction in token count to achieve an equivalent loss.

Category	Module / Depth	LR	WD	$\alpha_{1}$	$\alpha_{2}$	$\epsilon$	Init Scale
Standalone	Embedding weights (Audio)	2.192	1.009	0.962	1.493	1.494	1.826
	Embedding weights (Image)	1.013	0.864	2.108	0.685	0.734	0.554
	Embedding weight (Text)	3.937	1.593	1.421	1.791	0.317	0.379
	Unembedding weights (Audio)	1.633	1.510	3.442	0.594	0.742	3.422
	Unembedding weights (Image)	1.655	1.213	1.929	1.042	0.635	1.524
	Unembedding weights (Text)	3.008	0.737	1.346	0.955	1.206	0.341

Appendix

Table 6: Optimal per-module hyperparameter multipliers found for the LLaDA Multimodal model. Depth factors apply to all layers where the

\frac{\texttt{block depth}}{\texttt{total depth}}

falls within the highlighted fraction (as counted from the network input towards output).

Figure 22: Training curves for uni-modal text MDM models trained under CompleteP.

Panel of Figure 23 — Figure 23: Scaling laws for uni-modal text MDM s . a) Scaling law predictions for text-only MDM models, using CompleteP. b) Iso-FLOP curves for text-only MDM models under CompleteP parameterization (no SDE scaling).

Figure 24: Masking schedule ablation for audio generation across guidance scales. We evaluate four masking schedules (linear, cosine, polynomial, geometric) on ground-truth length audio generation quality using FAD, WER, and AudioBox Aesthetics metrics on our train mixture. Models are evaluated across CFG scales 1.0-10.0.

Figure 25: Masking schedule ablation across guidance scales. We evaluate the generation quality of four masking schedules (linear, cosine, polynomial, geometric) in CC12M and our train mixture (eval seed). Models are evaluated across CFG scales 1.0-10.0.

Panel of Figure 26 — Figure 26: Samples generated by our model with different prompts. See Table 7 for the extensive list of prompts.

Panel of Figure 27 — Figure 27: Samples generated by our model with different prompts.

Panel of Figure 28 — Figure 28: Samples generated by our model with different prompts.

Panel of Figure 29 — Figure 29: Samples generated by our model with different prompts.

(a) Lagerstroemia macrocarpa Wall, Queen’s flower, Queen’s crape myrtle, Lythraceae. Close-up of a vibrant pink flower with a yellow center in full bloom, sharply focused against a blurred green outdoor background. Numerous small green buds at different stages surround the flower on thin green stems.

(b) A cozy living room with a plush brown sofa and patterned pillows. A wooden coffee table sits in front, framed wall art above, rustic wooden ceiling, and a bed with a white comforter in the background. Decorative lighting and wall hangings add warmth.

(c) A photograph of a rainbow trout in a body of water. The trout is facing the left side of the image. The water is clear and there are some small rocks at the bottom of the water. There are also some small fish in the water.

(d) A close-up photograph of a painting with a variety of colors. The colors are bright and vibrant. The painting is made up of small squares of color that are all the same size.

(e) Watercolor drawing of lady with punk makeup.

(f) A teddy bear lies alone among the snow.

(g) A close-up photo of dried white flowers with brown centers, surrounded by purple and green dried flowers. The arrangement is set against a plain white wall, emphasizing texture and muted tones.

Appendix

Table 7: Prompts used to generate images in Figure 26 .

(u) A serene landscape with a snow-capped mountain beneath a clear blue sky. Tall golden grass fills the foreground, distant trees add contrast, and warm lighting suggests early morning or late afternoon.

(v) A close-up photo of a circular apple pie on a wooden surface, with apples arranged in a circular pattern and dusted with powdered sugar.

(w) A cartoon parrot wearing a blue-and-white sailor hat with an anchor, smiling while holding a guitar. The parrot has green and yellow feathers, an orange beak, blue eyes, and stands on a white background.

(x) Five bowls of chia pudding arranged in a grid on a white surface, topped with peanut butter, berries, kiwi with chocolate chips, nuts with honey, and banana with granola under bright lighting.

(y) A close-up photo of a hand holding a rolled up 20 Euro note. The background is out of focus, but there are green leaves and what appears to be a tree.

(z) A close-up photo of a small pizza topped with red sauce, cheese, and herbs on a white plate placed on a wooden table.

(aa) Fresh silver fish with pinkish tones displayed on crushed ice in a supermarket. The fish overlap slightly, with finely shaved white ice keeping them fresh and the background cropped tightly.

Appendix

Table 8: Prompts used to generate images in Figure 27 .

(ao) Skewers with charred meat, roasted golden potatoes, and roasted red tomatoes served on a white plate with a brown rim. The background is softly blurred with table elements visible.

(ap) Blossom fruit. beautiful spring.

(aq) An adult cow and calf lying on straw inside a wooden barn. The cow is brown with white patches, the calf mostly white with brown markings, lit by soft natural light.

(ar) Slow Down Tours on the train.

(as) A calm beach scene with a large wet boulder in the foreground, sandy shore, gentle blue waves, and a clear sky creating a tranquil atmosphere.

(at) A forest landscape with green, red, and yellow foliage in the foreground and a mountain covered in autumn leaves in the background.

(au) Gorgeous portrait of a blue peacock with silky blue feathers.

Appendix

Table 9: Prompts used to generate images in Figure 28 .

(bi) A wide shot of a valley surrounded by snow-covered mountains under a cloudy sky, evenly lit by natural daylight.

(bj) A white lighthouse on a rocky cliff surrounded by green vegetation, overlooking the ocean beneath a clear blue sky with scattered clouds.

(bk) A creeping thistle flower with a vibrant pink center on a sepia-toned old paper background, softly lit with blank space for text.

(bl) A dimly lit dining table set with plates of food and wine glasses. Three wine bottles are visible, chairs surround the table, and a wooden wall adds a cozy mood.

(bm) Streets of Kanazawa - japanese, japan, kanazawa, street, building, oriental, lantern.

(bn) A white bowl of soup with meat, carrots, and greens on a wooden table, with salt and pepper shakers, a tomato, and a beige napkin nearby.

(bo) A close-up view of stalactites hanging from a cave ceiling, with rugged textured walls and contrasting light and shadow highlighting natural formations.

Appendix

Table 10: Prompts used to generate images in Figure 29 .

Diffusion Models

Diffusion models learn to reverse gradual data corruption, enabling flexible generation, editing, forecasting, and inference.

See all 1,169 papers

Omni-Diffusion: Unified Multimodal Understanding and Generation with Masked Discrete Diffusion
Mar 6, 2026Lijiang Li, Zuwei Long, Yunhang Shen +6Masked Diffusion Models Unified Multimodal Models
Nanjing University · Tencent Youtu Lab · CASIA
Scaling Beyond Masked Diffusion Language Models
Feb 16, 2026Subham Sekhar Sahoo, Jean-Marie Lemercier, Zhihan Yang +4Diffusion Language Models Diffusion Models
NVIDIA, Santa Clara, USA · Department of Computer Science, Cornell Tech, NYC, USA · Department of Computer Science, Cornell University, Ithaca, USA +2
Sparse-LaViDa: Sparse Multimodal Discrete Diffusion Language Models
Dec 16, 2025Shufan Li, Jiuxiang Gu, Kangning Liu +4Masked Diffusion Models Diffusion Language Model Inference
Adobe · UCLA

Mar 6, 2026J/K move · Enter open · S save

Omni-Diffusion: Unified Multimodal Understanding and Generation with Masked Discrete Diffusion

Lijiang Li, Zuwei Long, Yunhang Shen, Heting Gao+5

Nanjing University · Tencent Youtu Lab · CASIA

Masked Diffusion Models Unified Multimodal Models Multimodal Generation Multimodal Foundation Models Multimodal Learning Diffusion Models+2

Open paper AskPDF arXiv

While recent multimodal large language models (MLLMs) have made impressive strides, they predominantly employ a conventional autoregressive architecture as their backbone, leaving significant room to explore effective and efficient alternatives in architectural design. Concurrently, recent studies have successfully applied discrete diffusion models to various domains, such as visual understanding and image generation, revealing their considerable potential as a promising backbone for multimodal systems. Drawing inspiration from these pioneering studies, we introduce Omni-Diffusion, the first any-to-any multimodal language model built entirely on mask-based discrete diffusion models, which unifies understanding and generation across text, speech, and images. Omni-Diffusion employs a unified mask-based discrete diffusion model to directly capture the joint distribution over discrete multimodal tokens. This approach supports not only bimodal tasks but also more complex scenarios involving multiple modalities. On a diverse set of benchmarks, our method outperforms or performs on par with existing multimodal systems that process two or more modalities, highlighting the significant promise of diffusion models in powering the next generation of multimodal foundation models. Project webpage: https://omni-diffusion.github.io.

Abstract

Diffusion Models

Explore similar work

Omni-Diffusion: Unified Multimodal Understanding and Generation with Masked Discrete Diffusion

Scaling Beyond Masked Diffusion Language Models

Sparse-LaViDa: Sparse Multimodal Discrete Diffusion Language Models