トライモーダルマスク拡散モデルの設計空間

The Design Space of Tri-Modal Masked Diffusion Models

February 25, 2026

著者: 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

要旨

離散拡散モデルは、自己回帰言語モデルの強力な代替として登場し、近年では単一モダリティの基盤モデルを初期化し、双モダリティ生成に向けてファインチューニングする手法が提案されている。従来のアプローチと異なり、本研究ではテキスト、画像-テキスト、音声-テキストデータからスクラッチで事前学習された初の三モダリティマスク拡散モデルを提案する。我々は、マルチモダリティにおけるスケーリング則、モダリティ混合比率、ノイズスケジュール、バッチサイズ影響を体系的に分析し、最適化された推論サンプリングのデフォルト値を提供する。バッチサイズ分析からは、最近の研究で報告されていた最適バッチサイズ調整の必要性を排除する、確率微分方程式（SDE）に基づく新たな再パラメータ化手法を導出した。この再パラメータ化は、計算制約（GPU飽和度、FLOP効率、実時間）に基づいて選択される物理的バッチサイズと、確率的最適化における勾配分散のバランスを取るために選択される論理的バッチサイズを分離する。最後に、6.4トークンのデータで予備的な30億パラメータ三モダリティモデルを事前学習し、統一設計の有効性を実証するとともに、テキスト生成、テキストから画像への変換、テキストから音声への変換タスクで強力な結果を達成した。本研究は、複数モダリティにわたるスケーリング挙動に関する知見を提供し、マルチモダリティ離散拡散モデルにおける最大規模の体系的な公開研究を代表するものである。

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.