The Design Space of Tri-Modal Masked Diffusion Models

Abstract: Discrete diffusion models have emerged as strong alternatives to autoregressive LLMs, 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.

• The paper proposes a unified tri-modal masked diffusion architecture that integrates text, image, and audio modalities using a bidirectional transformer and modality-specific tokens.
• The paper leverages CompleteP and SDE-based reparameterization to enable effective hyperparameter transfer and achieve batch-size invariance up to a critical threshold.
• The paper derives empirical scaling laws that reveal improved data efficiency per parameter and highlights modality-dependent inference regimes without cross-modal synergies.

Tri-modal Masked Diffusion Models: Architecture, Scaling, and Empirical Insights

Introduction and Motivation

The paper "The Design Space of Tri-Modal Masked Diffusion Models" (2602.21472) establishes a comprehensive empirical and theoretical analysis of discrete masked diffusion models (MDMs) extended to a unified tri-modal regime encompassing text, images, and audio. Traditional autoregressive (AR) paradigms have shown strong results in language modeling, but the causal, left-to-right constraint is poorly suited for flexible multimodal conditioning and infilling tasks. Discrete diffusion, by introducing order-agnostic iterative refinement, allows globally conditioned multimodal generation and infilling without explicit ordering or modality-specific architectural adaptations.

Unified Tri-modal MDM Architecture

The authors introduce a model integrating text, image, and audio modalities into a single discrete token stream handled by a bidirectional transformer backbone. Modality-specific boundary and mask tokens are utilized, and task tokens signal context-specific objectives such as text-to-image (T2I), text-to-speech (T2S), image captioning, and automatic speech recognition (ASR). Modality encapsulation is achieved via specialized tokenizers: SBER-MoVQGAN for images, Higgs Audio v2 for audio, and Tiktoken for text.

Figure 1: Tri-Modal masked diffusion model architecture with packed text and padded multimodal pairs for uniform sequence length.

This approach circumvents the need for modality-specific heads or adapters. Masked diffusion corrupts sequences by masking tokens independently per modality and then iteratively reconstructs the full sequence, enabling orderless infilling naturally suited to multimodal tasks.

Pretraining, Hyperparameter Transfer, and SDE-based Reparametrization

Extensive hyperparameter search at large scale is computationally prohibitive. The authors leverage CompleteP and SDE-based parameterization for robust hyperparameter transfer across model scales and for batch-size invariance. SDE reparameterization decouples physical batch size (chosen according to compute constraints) from logical batch size (required to maintain variance in stochastic optimization). Empirical analysis demonstrates that up to a critical batch size $B_\text{crit}$, training loss remains invariant.

Figure 2: Loss profile as a function of batch size, showing batch size invariance below the critical threshold $B_{\text{crit}}$.

Figure 3: The required iteration count $S_{\text{crit}}$ is nearly independent of model size, scaling sub-linearly with token horizon.

A key contribution is the formalization of the drift-horizon tradeoff, controlled via interpolation parameter $\gamma$, allowing optimal allocation between noise reduction (batch size) and trajectory length (iterations) during training.

Figure 4: The optimal drift-horizon balance for token horizon scaling lies at $\gamma^*\approx 0.44$, interpolating between noise reduction and deeper trajectory exploration.

Scaling Laws for Tri-modal Diffusion Models

The paper presents rigorous empirical scaling laws for tri-modal MDMs pretrained using CompleteP+SDE. Validation loss as a function of model size $N$ and token budget $D$ is best fit by a non-additive power law of the form $L(N, D) = E + (AN^{-a/b} + BD^{-1})^b$, which deviates from the standard additive Chinchilla form used for AR LLMs.

Figure 5: Training curves for tri-modal MDM, showing loss as a function of $(N, D)$ at fixed compute budgets.

Figure 6: Empirical scaling law fit ($R^2$=99.3%) for tri-modal MDMs, exhibiting high predictive fidelity.

The compute-optimal token per parameter ratio is derived to be $D^{\star}(N) \approx 7754 \cdot N^{0.84}$, implying tri-modal MDMs are more data-efficient per parameter than AR models. Notable is the finding that a 3B parameter model requires $\sim480$B tokens, markedly higher than AR scaling expectations.

Figure 7: Isoloss contours for tri-modal MDM models, illustrating directions for optimal hyperparameter transfer.

Empirical Results and Ablations

A unified 3B tri-modal MDM was trained on 6.4T tokens. The model achieves competitive results across text, image, and audio benchmarks, with performance detailed via multimodal radar plots using metrics such as FID, WER, Audiobox Aesthetics, and LM-harness scores.

Figure 8: Multi-modal evaluation summary for the 3B MDM: normalized radar and raw metric values across image, text, and audio.

A modality mixing ablation revealed no synergistic effects between modalities at the measured scale, with loss decreasing monotonically based on respective modality mixing ratios.

Figure 9: Loss contours as a function of tri-modal mixing ratios, showing competition for model capacity and no synergy.

Inference hyperparameter ablations demonstrate marked modality sensitivity: classifier-free guidance (CFG) scale, temperature, and top-p must be tuned independently for text, image, and audio.

Figure 10: Hyperparameter ablations for text-to-image generation, showing optimal FID across variable steps, CFG, temperature, and top-$p$.

Figure 11: Hyperparameter ablations for text-to-speech generation, showing metric-sensitive optimality for steps, CFG, temperature, and top-$p$.

Anti-masking interventions—where decorrelated masks are applied to each batch sample—yielded consistent improvements in training stability and downstream performance.

Theoretical Implications and Practical Impact

The unified architecture and unified token space streamline multimodal task conditioning and infilling. The SDE-based batch-size invariance offers pragmatic advantages for utilizing available compute and optimizing wall-clock times without performance sacrifice, up to a critical threshold. The scaling law results suggest that as multimodal MDM parameter count grows, data efficiency per parameter improves, contrasting the AR regime. The absence of synergistic cross-modal learning at moderate scale further underscores the need for scale or specialized regularization to exploit cross-modal information.

Future Directions

Potential extensions include:

• Scaling to larger parameter regimes to explore possible cross-modal synergies and emergent behaviors at high capacity.
• Refinement of multimodal data mixing strategies leveraging advanced mixture optimization frameworks.
• Integration of adaptive masking schedules per modality and dynamic task token conditioning.
• Acceleration of inference via parallel decoding protocols and optimized sampling heuristics for discrete diffusion (Wu et al., 28 May 2025).
• Analysis of compute/data scaling regimes for other modalities beyond text, image, and audio.

Conclusion

This paper delivers a rigorous empirical and theoretical exploration of the unified tri-modal masked diffusion model architecture, offering practical recipes for pretraining and hyperparameter transfer via CompleteP and SDE-based scaling. Strong numerical evidence supports batch-size invariance, prescriptive scaling laws, and modality-dependent inference regimes. The work sharply characterizes the challenges and opportunities present in multimodal diffusion pretraining and lays foundational principles for future research in multimodal generative modeling and order-agnostic sequence generation at scale.