PaperlayerView exact revision on arXiv

Source-linked AI summary

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, 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

arXiv:2602.21472v1cs.LG

TL;DR

The paper addresses the underexplored design space of native tri-modal masked diffusion pretraining across text, image, and audio. It introduces a unified from-scratch model and systematically studies scaling, batch-size, modality-mixing, noise-schedule, and inference choices. The resulting analysis provides empirical scaling guidance, removes optimal-batch-size tuning through SDE reparameterization, and identifies modality-specific inference tradeoffs.

  • Problem

    Native multimodal masked diffusion pretraining remains underexplored compared with adapting pretrained unimodal models, despite design choices that can strongly affect stability and compute efficiency.

  • Method

    The paper pretrains a unified tri-modal masked diffusion model from scratch and analyzes multimodal scaling laws, SDE-based batch-size reparameterization, and modality-specific inference settings.

  • Results

    The study reports empirical scaling guidance, eliminates the need to search for an optimal batch size up to a critical threshold, and finds substantially different optimal inference settings across modalities.

  • Takeaways & Limitations

    A single discrete diffusion formulation can support flexible conditioning across text, image, and audio while requiring modality-specific inference choices.

  • Takeaways & Limitations

    Beyond a critical batch size, additional tokens have diminishing marginal utility or increase wall-clock time, so over-trained models cannot be both FLOP-efficient and fast to train.

Abstract

from arXiv · show

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.

1 Introduction

The paper studies native tri-modal masked diffusion pretraining across text, image, and audio, focusing on the design choices that govern scaling, stability, efficiency, and inference. It introduces a unified model and reports modality-dependent inference tradeoffs, SDE-based batch-size scaling, and empirical multimodal scaling laws.

  • Motivation: Existing multimodal masked diffusion work largely adapts pretrained unimodal models, leaving the pretraining design space underexplored.The paper emphasizes that choices affecting stability and compute efficiency can dominate at pretraining scale.
  • Contributions: A unified tri-modal MDM represents text, image, and audio in one discrete token space and supports conditioning in multiple directions.The model uses modality-specific boundary and mask tokens while sharing a transformer backbone and vocabulary.
  • Contributions: The model supports text-to-image generation, image captioning, text-to-speech, and automatic speech recognition through one architecture.
  • Contributions: SDE-based reparameterization makes training loss invariant to batch size up to a critical threshold, removing the need to search for an optimal batch size.The paper also evaluates per-module hyperparameter scaling within the multimodal MDM setting.
  • Contributions: The fitted scaling law is L(N, D) = E+ (AN^-a/b + BD^-1)^b, and the compute-optimal frontier is D★(N) ≈7754 · N^0.84.The paper reports this form as a better fit than the additive alternative and finds asymptotically greater data efficiency than its autoregressive comparison.
  • Contributions: Optimal inference noise schedules and sampling parameters differ significantly across text, image, and audio generation.

2 Background and Related Work

Masked diffusion models reconstruct masked discrete tokens through iterative denoising and have been extended from text to image and audio applications. Prior multimodal systems mainly cover one or two modalities, motivating a from-scratch tri-modal model with a unified token stream.

  • Masked diffusion background: Masked diffusion progressively replaces tokens with [MASK] and trains a reverse model to reconstruct the original sequence.Continuous-time formulations represent the forward and reverse processes as continuous-time Markov chains.
  • Text: Text diffusion has progressed from small datasets to large language MDMs, including an 8B-parameter model trained on 2.3T tokens.The cited literature reports strong performance on MMLU and GSM8K for LLaDA.
  • Image: Image MDMs have modeled pixels directly at low resolution or used pretrained tokenizers to represent images as discrete token grids.
  • Audio: Audio diffusion research is comparatively sparse, with prior work combining discrete diffusion or masking-unmasking with audio tokenizers.
  • Multimodal MDMs: Existing multimodal MDMs unify language and image tasks, but typically remain bimodal, initialize from pretrained weights, or operate at smaller scale.This paper adds audio, pretrains from scratch, accounts for total token budget, and jointly optimizes representations across modalities.
  • Multimodal MDMs: Tri-modal architecture diagrams pack pure text and pad image-caption or audio-transcription pairs to a maximum length while ignoring padding in attention and loss.

3 Method

The method concatenates discrete text, image, and audio tokens into a unified sequence, applies modality-specific masking, and trains a bidirectional denoiser to reconstruct masked positions. Generation begins from a fully masked target and iteratively reveals tokens under modality constraints.

  • Data representation: Training uses text-only, image-text, and audio-text samples, with each modality represented by tokens from a modality-specific vocabulary.
  • Vocabulary and sequences: A shared vocabulary combines disjoint text, audio, and image vocabularies with modality-specific special tokens.Task tokens signal which generation task the model should perform.
  • Vocabulary and sequences: Text-only sequences are packed to the maximum length, while shorter mixed-modality sequences are right-padded after EOStext.Padding is used to match the fixed sequence length L★.
  • Forward process: The continuous-time forward process independently masks positions with probability β_t, using dedicated mask tokens for each modality.The corruption level increases monotonically from the original sequence at t=0 to a fully masked sequence at t=1, and task tokens are never masked.
  • Denoising objective: The reverse model maps a corrupted sequence to logits over the unified vocabulary at every position and computes loss over masked, non-padding tokens.Cut-cross-entropy improves memory efficiency, while w(t)=1/t yields an unbiased ELBO estimator under Bernoulli masking.
  • Inference: Inference starts from a fully masked sequence and repeatedly samples modality-constrained candidate tokens before unmasking selected positions.A predefined linear schedule determines which masked positions are revealed at each step until no masks remain.
  • Model architecture: The 3B tri-modal model uses a bidirectional transformer with RMSNorm, SwiGLU MLPs, rotary positional embeddings, and QK-norm.It is pretrained from scratch for 1M steps with batch size 3072 and sequence length 3256.

4 Hyperparameter Transfer

The section develops hyperparameter-transfer rules based on CompleteP and SDE reparameterization, including a batch-size-invariant training scheme and scaling of AdamW hyperparameters. It also reports empirical SDE-regime behavior that identifies critical batch and iteration thresholds.

  • Transfer rules: Hyperparameter transfer rules allow settings found at small scale to be applied to larger models, avoiding infeasible large-scale grid searches.The study considers μP, depth-μP, u-μP, and CompleteP as candidate transfer rules.
  • SDE parametrization: AdamW hyperparameters are reparameterized through an SDE view because their effects are redundant with batch size.The reparameterization balances stochastic-gradient noise, drift, and the SDE horizon when batch size changes.
  • SDE parametrization: SDE parametrization keeps training performance similar across batch sizes up to the critical threshold Bcrit, eliminating the need to search for an optimal batch size Bopt.This behavior contrasts with the U-curve reported for non-SDE parametrization.
  • Scaling rules: The SDE-scaling factor κ interpolates between conserving stochastic drift and conserving the SDE horizon as token budget increases.γ=0 defines isonoise scaling, γ=1 defines isohorizon scaling, and intermediate γ values interpolate between them.
  • Critical batch size: Below Bcrit, larger batches reduce the required iteration count while maintaining constant loss; above Bcrit, SDE discretization breaks and training becomes FLOP-inefficient.Under the SDE regime, the critical iteration count Scrit is reported as constant with respect to model size.

5 Scaling Behavior of MDM under the SDE Transfer Rule

The SDE transfer rule clarifies how critical batch size scales with token horizon and model size, while enabling compute-flexible batch selection below a critical threshold. It also identifies a nontrivial drift–horizon allocation and derives scaling laws for compute-optimal tri-modal MDM training.

  • Critical batch-size scaling: Above Scrit, asymptotic loss depends mainly on model size N and token budget D, whereas too few integration steps break the SDE approximation.The critical batch size is Bcrit = D/(L Scrit), with L denoting sequence length.
  • Critical batch-size scaling: SDE reparameterization makes training loss invariant to batch size up to the critical batch size Bcrit, removing the need to tune an optimal batch size Bopt.The physical batch size can be selected for available compute, provided it remains below Bcrit.
  • Critical batch-size scaling: Scrit is independent of model size, while Bcrit is unaffected by model size under SDE scaling at equal token horizon.A batch size below Bcrit for a smaller model is therefore also safe for a larger model at the same token horizon.
  • Critical batch-size scaling: Scrit and the corresponding Bcrit increase sub-linearly with token horizon D, creating a tradeoff between wall-clock time and FLOP efficiency.Scaling GPUs proportionally with D eventually risks reaching Bcrit, after which marginal token utility drops or wall-clock time increases.
  • Optimal drift–horizon tradeoffs: The drift–horizon parameter γ allocates additional tokens between reducing stochastic drift through virtual batch size and extending the SDE horizon through virtual iterations.Neither γ = 0 nor γ = 1, corresponding to two literature settings, is optimal; the best allocation is approximately γ* ≈ 0.44.
  • Scaling laws for tri-modal MDM: A 3B tri-modal MDM requires at least 480B tokens, compared with 60B tokens reported for autoregressive language models under Chinchilla.The paper reports that this gap persists at realistically reachable model sizes, while noting that cross-family comparisons require caution.

6 Data

The experiments use a fixed 33% text, 33% image-text, and 33% audio-text mixture, with modality-specific datasets assembled from filtered and licensed sources.

  • 6 Data: All experiments except the modality-mixing ablation use equal 33% proportions of text, image-text, and audio-text data.The token horizon is smaller than the total dataset size, corresponding to a single global epoch.
  • 6 Data: The text corpus combines web, academic, technical, licensed, and synthetic reasoning data with additional cleaning and PII filtering.
  • 6 Data: The audio-text corpus contains 2M hours of web audio transcribed by Whisper and filtered for privacy, speech activity, dialogue, production quality, and complexity.
  • 6 Data: The image-text corpus aggregates CC3M, CC12M, COYO, recaptioned text datasets, and other licensed sources, with PII filtering applied.

7 Results

The results examine unified multimodal performance, modality mixing, inference hyperparameters, and anti-masking. They show modality-specific design tradeoffs and consistent anti-masking improvements across evaluated modalities.

  • 7.1 Unified Design: 10.06 FID-Inception and 107.61 FID-DINOv2 are reported on CC12M, while GenEval Overall reaches 48.89.
  • 7.1 Unified Design: The 3B tri-modal model is evaluated across text, image, and audio settings using modality-specific benchmarks and metrics.Image evaluation uses FID with DINOv2-L and Inception-v3 features on CC12M and training data.
  • 7.2 Modality Mixing: Modality-mixing experiments vary text, image-text, and audio-text weights across 15 experiments using a 320M model and 13B tokens.The loss-contour analysis uses an equal-mixture reference and reports no observed cross-modal synergies at that scale.
  • 7.3 Best Generation Hyperparameters: Text-to-image evaluation varies steps, CFG, temperature, and top-p across CC12M and training data, with FID improving with more steps but showing diminishing returns.Intermediate values are generally optimal for CFG and temperature, while top-p behaves differently.
  • 7.3 Best Generation Hyperparameters: Text-to-speech quality improves with more steps but exhibits metric-dependent tradeoffs, including stronger transcription accuracy and worse audio fidelity at higher CFG.Trends broadly generalize from the training evaluation set to LibriSpeech-PC.
  • 7.4 Anti-Masking: Anti-masking duplicates each sample with complementary masking patterns, while compute-matched baselines process unique samples across two epochs.
  • 7.4 Anti-Masking: Anti-masking improves multimodal performance, with the largest gains in audio generation quality, and produces consistent improvements across most text tasks.The comparison reports FID for images and FAD for audio, alongside mean accuracy and standard deviation for text evaluation.
  • 7.4 Anti-Masking: Anti-masking improves benchmark performance in the multiple-epochs setup without additional computational cost.

8 Conclusion

The paper presents tri-modal generation as unified iterative refinement over shared discrete tokens and maps the training and inference choices that govern efficiency and quality.

  • 8 Conclusion: Text, images, and audio share one token stream and transformer backbone, supporting captioning, text-to-image, ASR, and TTS without modality-specific heads.The study also reports SDE-based reparameterization, empirical scaling behavior, modality-dependent sampling, and compute-matched anti-masking gains.

Appendices

The appendices formalize masked diffusion, its loss weighting, and the relationship between masking corruption and timestep contributions. They show that inverse-probability weighting corrects the sampling bias introduced by masking.

  • Forward Process: The forward process progressively masks tokens through a Markov chain governed by timestep-specific transition matrices.
  • Forward Process: The cumulative probability of a token remaining unmasked through time t is ᾱ_t, while its probability of having been masked is 1−ᾱ_t.
  • Loss Weighting: The weighting function w(t) balances timestep contributions according to the masking schedule, which determines how quickly information is destroyed.
  • Loss Weighting: Early timesteps retain most context and are easier to reconstruct, whereas late timesteps mask most tokens and require inference from minimal context.
  • Unbiasedness: At a fixed timestep, averaging reconstruction loss over masked positions randomly subsamples token positions.
  • Unbiasedness: Multiplying by 1/t corrects this subsampling so each token contributes equally in expectation across timesteps.The appendix identifies this as inverse-probability weighting analogous to time-dependent diffusion loss weighting.

B.1 Audio Tokenizer Ablations

The tokenizer ablations compare audio and image tokenizers under sequence-length constraints, while documenting the training and runtime settings used for the broader model analysis.

  • Audio tokenizer selection: A fixed context length of 3256 and clips up to 30 seconds require a low-rate audio codec that preserves perceptual quality.The study compares pretrained DAC, Higgs Audio v2, and a same-data DAC-style tokenizer.
  • Audio tokenizer selection: Higgs Audio v2 with 4 codebooks is selected as the default audio tokenizer for its rate–distortion trade-off.More codebooks improve reconstruction but quickly become impractical under the sequence-length constraint.
  • Image tokenizer selection: Image tokenizer comparisons evaluate reconstruction FID on ImageNet at 256 and 512 resolution and on CC12M.The compared discrete tokenizers include Cosmos, IBQ, OpenMagVIT, Unitok, and MoVQGAN.
  • Training and search setup: The global hyperparameter search uses 2900 small-scale trial runs on a 320M-parameter model trained for 13B tokens.The resulting optimum initializes the per-module multiplier search.
  • Training and search setup: SDE parametrization permits a wide range of batch sizes, while runtime falls sharply with node count and FLOP efficiency declines under sub-linear scaling.The runtime study varies per-GPU batch size and node count while tracking GPU saturation and wall-clock time.
  • Training and search setup: The 3B multimodal MDM’s model and training details are collected in Table 5.The passage identifies the table as the location of the complete training specification.

D MDM with Per-module Hyperparameters

The paper extends global hyperparameter transfer with structured per-module AdamW multipliers, finding that different parameter groups and depths favor distinct optimization settings.

  • Method: Per-module tuning assigns unique AdamW multipliers while sharing depth and module-type factors across repeated blocks.The searched parameters include learning rate, weight decay, β1, β2, and 𝜖.
  • Results: 1.81× fewer tokens are required to reach equivalent loss with per-module tuning on the 350M model.The experiment uses a 350M model, including 80M non-embedding parameters, trained for 13B tokens with batch size 256.
  • Results: Embedding and unembedding weights favor effective learning rates up to ∼4× larger than baseline module settings.Attention projections and MLP gates are tuned more conservatively, often with increased 𝜖 for numerical damping.
  • Results: Later blocks use smaller steps and stronger stabilization according to the learned depth factors.The resulting multipliers are described as highly structured rather than uniform.
  • Results: Table 6 lists the optimal per-module multipliers and the depth-factor convention used for the LLaDA Multimodal model.Depth factors apply across layers according to block depth.

E Extended Scaling Laws Results

The extended scaling-law analysis defines FLOP accounting and applies CompleteP-based scaling to text-only models, with training curves and predictions reported across model sizes and budgets.

  • FLOP accounting: FLOP calculations exclude input and output embeddings because the 117k multimodal vocabulary makes them larger than the transformer backbone in small models.The reported scaling-law model size N therefore uses non-embedding parameters.
  • Scaling-law derivation: The compute-optimal derivation imposes C = 6ND and minimizes a power-law objective under that constraint.The stated minimizers follow from monotonicity of the constrained objective.
  • Uni-modal text scaling: Uni-modal text MDM scaling runs use CompleteP without SDE scaling rules, with batch sizes varied to maximize GPU occupancy.Models use sequence length 4096 with packing and truncation and no padding.
  • Uni-modal text scaling: Training curves and scaling-law predictions for text-only MDMs are presented in Figures 22 and 23.The figures cover curves as functions of model size, token budget, and FLOP budget, including iso-FLOP curves.

F Masking Schedules for Image and Audio Generation

The masking-schedule study compares linear, cosine, polynomial, and geometric schedules for image and audio generation across guidance scales, with polynomial performing best in the reported evaluations.

  • Experimental setup: Four masking schedules are trained in matched 1B-model experiments for 100k steps while holding other hyperparameters fixed.The schedules are linear, cosine, polynomial, and geometric, using continuous-time ELBO weighting.
  • Image generation: Polynomial masking consistently achieves the best image quality across FID-Inception and FID-DINOv2 on CC12M and the train mixture.The best image performance occurs at CFG scales from 7 to 9 using 1024 diffusion steps at 256×256 resolution.
  • Audio generation: Polynomial masking also consistently achieves the best audio quality across FAD, WER, and AudioBox Aesthetics.Audio evaluation uses ground-truth durations, 512 diffusion steps, and 10,000 generated samples.
  • Qualitative examples: Extended image generations are shown in Figures 26–29 with prompts listed in Tables 7–10.The prompts were selected from synthetic captions using quality filtering and diversity.
  • Evaluation figures: Figures 25 and 24 organize schedule quality comparisons across CFG scales from 1.0 to 10.0 for image and audio generation.The image and audio figures report the corresponding metric-based evaluations described in the text.
Loading 2602.21472v1…