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, Theo X. Olausson, Devon Hjelm, Paul Dixon, Joao Monteiro, Pierre Ablin, Vishnu Banna, Arno Blaas, Nick Henderson, Kari Noriy, Dan Busbridge, Josh Susskind, Marco Cuturi, Irina Belousova, Luca Zappella, Russ Webb, Jason Ramapuram
Organizations: Apple · University of Cambridge · MIT · Google Deepmind (work done at Apple)
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
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 2Figure 3Figure 4Figure 5
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 7Figure 8
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.
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. ⋆ indicate the best hyperparameter for each metric. For Production Complexity, lower values indicate simpler audio and are considered as preferable.
Model
FID (Inception) ↓
FID (DINOv2) ↓
FAD ↓
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 ± 2.10
43.12 ± 1.45
21.52 ± 0.46
30.20 ± 0.39
52.01 ± 1.40
63.72 ± 0.99
52.20 ± 0.50
24.75 ± 1.09
67.19 ± 1.10
31.66 ± 1.36
Anti-mask
33.00 ± 2.10
36.97 ± 1.41
27.05 ± 0.50
32.86 ± 0.39
54.06 ± 1.40
66.04 ± 0.97
55.15 ± 0.50
26.27 ± 1.11
67.41 ± 1.09
34.90 ± 1.39
Table 2: Anti-masking results on the evaluation harness. We report mean accuracy ± standard deviation.
Appendix figures & tables19 assets
Supplementary material from the paper’s appendix.
Appendix
Model
# Codebooks
PESQ ↑
Content Enjoy ↑
Content Useful. ↑
Prod. Complex. ↑
Prob. Quality ↑
Down. Factor ↑
Higgs pretrained
8
3.168
5.670
6.192
1.559
6.510
960
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
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
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 ↓
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, β1,β2 , weight-decay, ϵ ) on a 350M model trained for 13B tokens. The plot shows cross-sections through the hyperparameter landscape through the best found hyperparameters.
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 × reduction in token count to achieve an equivalent loss.
Category
Module / Depth
LR
WD
α1
α2
ϵ
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 total depthblock 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.
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.
Figure 26: Samples generated by our model with different prompts. See Table 7 for the extensive list of prompts.
Figure 27: Samples generated by our model with different prompts.
Figure 28: Samples generated by our model with different prompts.
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.
(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 .
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.
Diffusion language models are a promising alternative to autoregressive models due to their potential for faster generation. Among discrete diffusion approaches, Masked diffusion currently dominates, largely driven by strong perplexity on language modeling benchmarks. In this work, we present the first scaling law study of uniform-state and interpolating discrete diffusion methods. We also show that Masked diffusion models can be made approximately 12% more FLOPs-efficient when trained with a simple cross-entropy objective. We find that perplexity is informative within a diffusion family but can be misleading across families, where models with worse likelihood scaling may be preferable due to faster and more practical sampling, as reflected by the speed-quality Pareto frontier. These results challenge the view that Masked diffusion is categorically the future of diffusion language modeling and that perplexity alone suffices for cross-algorithm comparison. Scaling all methods to 1.7B parameters, we show that uniform-state diffusion remains competitive on likelihood-based benchmarks and outperforms autoregressive and Masked diffusion models on GSM8K, despite worse validation perplexity. We provide the code, model checkpoints, and video tutorials on the project page: http://s-sahoo.github.io/scaling-dllms
Subham Sekhar Sahoo, Jean-Marie Lemercier, Zhihan Yang +4
NVIDIA, Santa Clara, USA · Department of Computer Science, Cornell Tech, NYC, USA · Department of Computer Science, Cornell University, Ithaca, USA +2
Masked Discrete Diffusion Models (MDMs) have achieved strong performance across a wide range of multimodal tasks, including image understanding, generation, and editing. However, their inference speed remains suboptimal due to the need to repeatedly process redundant masked tokens at every sampling step. In this work, we propose Sparse-LaViDa, a novel modeling framework that dynamically truncates unnecessary masked tokens at each inference step to accelerate MDM sampling. To preserve generation quality, we introduce specialized register tokens that serve as compact representations for the truncated tokens. Furthermore, to ensure consistency between training and inference, we design a specialized attention mask that faithfully matches the truncated sampling procedure during training. Built upon the state-of-the-art unified MDM LaViDa-O, Sparse-LaViDa achieves up to a 2x speedup across diverse tasks including text-to-image generation, image editing, and mathematical reasoning, while maintaining generation quality.