Lo Spazio di Progettazione dei Modelli di Diffusione con Maschera Tri-Modale

The Design Space of Tri-Modal Masked Diffusion Models

February 25, 2026

Autori: 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

cs.AI

Abstract

I modelli di diffusione discreta sono emersi come valide alternative ai modelli linguistici autoregressivi, con lavori recenti che inizializzano e mettono a punto un modello base unimodale per la generazione bimodale. Diversamente dagli approcci precedenti, introduciamo il primo modello di diffusione mascherata tri-modale addestrato da zero su dati di testo, immagine-testo e audio-testo. Analizziamo sistematicamente le leggi di scaling multimodali, i rapporti di miscelazione delle modalità, le pianificazioni del rumore e gli effetti della dimensione del batch, e forniamo impostazioni predefinite ottimizzate per il campionamento in fase di inferenza. La nostra analisi sulla dimensione del batch produce una nuova riformulazione basata su equazioni differenziali stocastiche (SDE) che elimina la necessità di ottimizzare la dimensione ottimale del batch come riportato in lavori recenti. Questa riformulazione disaccoppia la dimensione fisica del batch, spesso scelta in base ai vincoli computazionali (saturazione della GPU, efficienza dei FLOP, tempo di esecuzione), dalla dimensione logica del batch, scelta per bilanciare la varianza del gradiente durante l'ottimizzazione stocastica. Infine, addestriamo preliminarmente un modello tri-modale da 3 miliardi di parametri su 6,4 mila miliardi di token, dimostrando le capacità di un design unificato e ottenendo risultati solidi nella generazione di testo, in compiti di text-to-image e di text-to-speech. Il nostro lavoro rappresenta il più ampio studio sistematico e aperto sui modelli di diffusione discreta multimodale condotto fino ad oggi, fornendo intuizioni sui comportamenti di scaling attraverso multiple modalità.

English

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.