跳到正文
Apple Machine Learning Research·· 19 小时前AI 评分59

Apple 发布三模态掩码扩散模型设计空间研究

The Design Space of Tri-Modal Masked Diffusion Models

AI 导读

Apple 团队发布首个从零预训练的三模态掩码扩散模型(MDM)研究,在文本、图文和音频文本数据上训练。论文系统分析了多模态缩放规律、模态混合比例、噪声调度和 batch size 影响,提出基于 SDE 的重参数化方法免去调优最优 batch size,并在 3B 模型规模(6.4T tokens)上验证了文本生成、T2I 和 T2S 任务的表现。

正文

AuthorsLouis Béthune, Victor Turrisi, Bruno Kacper Mlodozeniec†**, Pau Rodriguez Lopez, Lokesh Boominathan, Nikhil Bhendawade, Amitis Shidani, Joris Pelemans, Theo X. Olausson‡**, Devon Hjelm, Paul Dixon, João 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§**

Discrete diffusion models have emerged as strong alternatives to autoregressive language models, with recent work initializing and finetuning a base unimodal model for bi-modal generation. Diverging from previous approaches, we introduce the first tri-modal Masked Diffusion Models (MDM) pretrained from scratch on text, image-text, and audio-text data. We systematically analyze multimodal scaling laws, modality mixing ratios, noise schedules, batch-size effects and provide optimized inference sampling defaults. Our batch-size analysis yields a novel stochastic differential equation (SDE) based reparameterization, eliminating 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 (GPU saturation, FLOP-efficiency, wall-clock time) from the logical batch size, chosen to balance the variance of gradients during stochastic optimization. Finally, we pretrain a preliminary model showcasing the capabilities of a unified design, achieving strong results at 3B model scale (6.4T tokens), in both text generation, T2I tasks, and T2S tasks. Our work represents the largest scale systematic open study of multimodal discrete diffusion models conducted to date, providing valuable insights into scaling behaviors across multiple modalities.

  • † University of Cambridge
  • ‡ MIT
  • § Google DeepMind
  • ** Work done while at Apple

Related readings and updates.

Discrete diffusion is a promising framework for modeling and generating discrete data. In this work, we present Target Concrete Score Matching (TCSM), a novel and versatile objective for training and fine-tuning discrete diffusion models. TCSM provides a general framework with broad applicability. It supports pre-training discrete diffusion models directly from data samples, and many existing discrete diffusion approaches naturally emerge as…

Read more

*Equal Contributors

Preserving training dynamics across batch sizes is an important tool for practical machine learning as it enables the trade-off between batch size and wall-clock time. This trade-off is typically enabled by a scaling rule; for example, in stochastic gradient descent, one should scale the learning rate linearly with the batch size. Another important machine learning tool is the model EMA, a functional copy of a target model…

Read more

来源:Apple Machine Learning Research · machinelearning.apple.com