The Design Space of Tri-Modal Masked Diffusion Models

TL;DR

Apple’s first tri-modal Masked Diffusion Model (MDM) pretrained from scratch on text + image-text + audio-text, with a 3B-parameter / 6.4T token preliminary release. The contribution is less the model than the empirical design-space study around it: a SDE-based AdamW reparameterization that eliminates the “optimal batch size” search above a critical threshold, multimodal scaling laws fitted on 262 isoFLOP runs that show tri-modal MDMs are asymptotically more data-efficient per parameter than Chinchilla-style AR LLMs (a 3B model wants ~250B tokens, vs 60B for AR), and modality-dependent inference defaults (noise schedule, CFG, temperature differ for text vs image vs audio).

Key claims

A single transformer with a unified discrete vocabulary (text + SBER-MoVQGAN image tokens + Higgs Audio v2 tokens) and modality-specific boundary/mask tokens can be trained from scratch for tri-modal masked diffusion, supporting text gen, T2I, captioning, TTS, and ASR with no modality-specific heads or adapters [§3, §3.1].
Under a CompleteP + SDE reparameterization of AdamW, training loss is invariant to physical batch size up to a critical threshold $B_{\rm crit}$ , eliminating the optimal-batch-size search reported by prior work [§4.1, Fig. 5].
The SDE rescaling decouples physical batch size (chosen for GPU saturation / wall-clock) from logical batch size (chosen to balance gradient noise vs drift) — runs below $B_{\rm crit}$ are all FLOP-efficient for a given token budget; above it, runs trade extra tokens for faster wall-clock [§5.1].
The critical iteration count $T_{\rm crit}$ is empirically independent of model size $N$ at fixed token budget, and grows sub-linearly with token horizon $D$ [§5.1, Fig. 5].
A new “isonoise vs isohorizon” interpolation parameter $\gamma$ governs how extra tokens are allocated between reducing SDE drift ( $\gamma = 0$ ) and extending SDE horizon ( $\gamma = 1$ ); both endpoints in the prior literature are suboptimal, and the optimal value is empirically $\gamma^* \approx 0.50$ [§5.2, Eq. 4.1, Figs. 9-9].
The Kaplan-style multiplicative scaling-law form $L(N, D) = A/N^\alpha + B/D^\beta + \text{interaction}$ fits tri-modal MDM training curves at $R^2 = 99.3\%$ / MRE 0.5%, while the Chinchilla additive form fits poorly [§5.3, Eq. 5.1, Fig. 12].
Compute-optimal TPP decreases as $N$ grows for tri-modal MDM: the fitted frontier gives $D^* \propto N^{0.476}$ , so larger MDMs are asymptotically more data-efficient per parameter than AR LLMs (where Chinchilla gives $D^* \propto N$ ) [§5.3, Eq. 5.2, Fig. 3].
A 3B compute-optimal tri-modal MDM needs ~250B tokens, vs ~60B prescribed by Chinchilla for AR; the data-efficiency gap is maintained at all realistically reachable model sizes [§5.3, Fig. 3].
Modality-specific tokenizers (SBER-MoVQGAN for images, Higgs Audio v2 for audio, Tiktoken for text) plus cut-cross-entropy and z-loss regularization are used to manage the large unified vocabulary efficiently; tokenizer FLOPs are excluded from the scaling-law fit [§3.1, §5.3].
Optimal inference defaults (noise schedule, sampling temperature, CFG scale) differ significantly across modalities — there is no single “right” setting for a unified MDM [§1 contributions, §6 (truncated)].
The fitted scaling exponents are compatible with concurrent work on unimodal MDMs, validating that the data-efficiency advantage is a property of masked diffusion rather than of multimodality [§5.3].

Method

Each training sample belongs to one of three classes — text-only, image-text, or audio-text — wrapped with modality-specific boundary tokens and an explicit task token (T1/T2/T3) so the model knows which generation direction is being requested. Text-only sequences are packed to the fixed sequence length $L=2048$ ; multimodal sequences are right-padded. A continuous-time forward Bernoulli-masking process with schedule $\alpha_t$ (monotonic, $\alpha_0=1$ , $\alpha_1=0$ ) independently corrupts each position to its modality’s mask token. The denoising transformer is bidirectional with RMSNorm, SwiGLU, RoPE, QK-norm; the training loss is the standard ELBO weighting $\dot{\alpha}_t/(1-\alpha_t)$ implemented via cut-cross-entropy plus a z-loss regularizer for numerical stability of the large vocabulary. Inference is iterative unmasking on a linear schedule from a fully-masked target span, conditioned on the prompt and the target task token.

The headline methodology contribution is the CompleteP + SDE reparameterization of AdamW (extending Mlodozeniec et al.’s CompleteP): per-module multipliers on $\eta, \lambda, \beta_1, \beta_2$ are rescaled by a single factor $S(N_0/N, B_0/B; \gamma)$ that interpolates between conserving SDE drift ( $\gamma=0$ ) and conserving SDE horizon ( $\gamma=1$ ). Combined with the SDE rule, the physical batch size becomes a free knob below $B_{\rm crit}$ . The scaling-law experiment is 262 runs on 24 logarithmically-spaced isoFLOP curves between $10^{18}$ and $10^{21}$ FLOP with TPP ratios from 1 to 2000, fit by basin-hopping + L-BFGS with 90/10 bootstrap cross-validation. The preliminary 3B model trains for 1M steps at batch size 2048 on 6.4T tokens.

Results

Scaling-law fit: $R^2 = 99.3\%$ , MRE 0.5% across 262 isoFLOP runs; multiplicative Kaplan form preferred over additive Chinchilla form [§5.3, Fig. 12].
Compute-optimal frontier: $D^* \propto N^{0.476}$ , vs $D^* \propto N$ for AR Chinchilla — implies a 3B MDM wants ~250B tokens, ~4× more than AR’s 60B at the same parameter count [§5.3, Eq. 5.2, Fig. 3].
SDE invariance: identical final loss across batch sizes from small to $B_{\rm crit}$ at fixed token budget for a 320M / 13B-token setup; non-SDE runs show the classic U-shaped batch-size curve [§5.1, Fig. 5].
Drift–horizon tradeoff: optimal $\gamma^* \approx 0.50$ (fitted power law $a + b \cdot S^\beta$ ), beating both $\gamma=0$ (Kaplan default) and $\gamma=1$ (CompleteP default) [§5.2].
3B preliminary model: trained for 1M steps on 6.4T tokens with batch size 2048; reports “strong results” across text gen, T2I, and T2S/ASR tasks (the design-space study, not the leaderboard numbers, is the contribution) [§3.1, Abstract].

Why it’s interesting

A scaling-laws paper with an opinionated stance: pretraining-from-scratch multimodal MDMs change the compute-optimal recipe — they want proportionally more tokens than AR LLMs at the same parameter count, in the opposite direction of the recent push for under-trained sparse-MoE LLMs. That has direct consequences for the data side of any unified-generative-model effort at Luma. The SDE/CompleteP work is also the most aggressive “make batch size irrelevant” claim filed so far: it’s the natural successor to the Fitting-vs-µTransfer debate in Hyperparameter scaling laws (How to Set the Learning Rate for Large-Scale Pre-training? argues µP-style transfer is stale under modern QK-Norm; this paper argues a further SDE-based reparameterization removes batch-size search entirely). Finally, the tri-modal angle complements Joint audio-video generation from the other side — LTX-2 / MOVA / SkyReels-V4 are continuous diffusion dual-stream DiTs; this is a single discrete-token unified MDM stack with audio + text + image. Both recipes are now on the table for “one model, many modalities.”