E-MMDiT: Revisiting Multimodal Diffusion Transformer Design for Fast Image Synthesis under Limited Resources

📄 arXiv: 2510.27135v1 📥 PDF

作者: Tong Shen, Jingai Yu, Dong Zhou, Dong Li, Emad Barsoum

分类: cs.CV

发布日期: 2025-10-31

🔗 代码/项目: GITHUB


💡 一句话要点

提出E-MMDiT,一种轻量高效的多模态扩散Transformer,用于资源受限下的快速图像合成。

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 扩散模型 Transformer 图像合成 轻量化 多模态 Token缩减 资源受限 生成式AI

📋 核心要点

  1. 现有扩散模型训练成本高昂,需要大量数据和算力,且模型结构复杂,推理延迟高,难以在资源受限场景下应用。
  2. E-MMDiT通过token缩减策略,结合高压缩视觉tokenizer、多路径压缩模块、位置强化和交替子区域注意力等技术,实现了轻量高效的设计。
  3. 实验表明,E-MMDiT仅用少量数据和算力即可训练,在图像生成质量上取得了有竞争力的结果,并提供了易于复现的基线。

📝 摘要(中文)

扩散模型在从文本提示生成高质量图像方面表现出强大的能力。然而,这些模型通常需要大规模的训练数据和大量的计算资源进行训练,或者存在结构复杂、延迟高等问题。为此,我们提出了高效多模态扩散Transformer(E-MMDiT),这是一个高效且轻量级的多模态扩散模型,仅有304M参数,用于在低训练资源下快速图像合成。我们提供了一个易于复现的基线,并取得了具有竞争力的结果。我们的512px生成模型,仅使用25M公共数据,在单个包含8个AMD MI300X GPU的节点上训练1.5天,在GenEval上达到0.66,并且通过一些后训练技术(如GRPO)可以轻松达到0.72。我们的设计理念以token缩减为中心,因为计算成本随着token数量的增加而显著增加。我们采用了一种高度压缩的视觉tokenizer来产生更紧凑的表示,并提出了一种新颖的多路径压缩模块来进一步压缩token。为了增强我们的设计,我们引入了位置强化(Position Reinforcement),它加强了位置信息以保持空间连贯性,以及交替子区域注意力(ASA),它在子区域内执行注意力以进一步降低计算成本。此外,我们提出了AdaLN-affine,一个高效的轻量级模块,用于计算Transformer块中的调制参数。我们的代码可在https://github.com/AMD-AGI/Nitro-E 获得,我们希望E-MMDiT能够成为未来研究的强大而实用的基线,并为生成式AI模型的普及做出贡献。

🔬 方法详解

问题定义:论文旨在解决扩散模型在资源受限场景下难以训练和部署的问题。现有扩散模型通常参数量巨大,需要大量的训练数据和计算资源,并且模型结构复杂,推理延迟高,这限制了它们在边缘设备或低成本环境中的应用。

核心思路:论文的核心思路是通过token缩减来降低计算成本。由于Transformer的计算复杂度与token数量呈平方关系,因此减少token数量可以显著降低计算量,从而实现轻量化和加速。

技术框架:E-MMDiT的整体架构基于扩散Transformer,主要包含以下模块:1) 高压缩视觉Tokenizer:将输入图像转换为更紧凑的token表示。2) 多路径压缩模块:进一步压缩token数量,降低计算量。3) 位置强化(Position Reinforcement):增强位置信息,保持空间连贯性。4) 交替子区域注意力(ASA):在子区域内执行注意力,降低计算复杂度。5) AdaLN-affine:高效计算Transformer块中的调制参数。

关键创新:论文的关键创新在于多方面的token缩减策略,包括高压缩视觉Tokenizer和多路径压缩模块,以及交替子区域注意力机制。这些创新共同作用,显著降低了模型的计算复杂度,使其能够在资源受限的环境下进行训练和推理。

关键设计:1) 高压缩视觉Tokenizer的具体压缩比例未知。2) 多路径压缩模块的具体实现方式未知。3) 位置强化通过在token中加入位置编码实现。4) 交替子区域注意力将图像划分为多个子区域,并在每个子区域内进行注意力计算,从而降低计算复杂度。5) AdaLN-affine是一种轻量级的自适应层归一化方法,用于计算Transformer块中的调制参数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

E-MMDiT仅使用304M参数,在单个包含8个AMD MI300X GPU的节点上,使用25M公共数据训练1.5天,即可在512px图像生成任务中,在GenEval上达到0.66。通过GRPO等后训练技术,可以轻松达到0.72。该结果表明E-MMDiT在资源受限的情况下,也能实现具有竞争力的图像生成质量。

🎯 应用场景

E-MMDiT具有广泛的应用前景,包括移动设备上的图像生成、低成本云服务、以及在资源受限的环境中进行AI研究和开发。该模型可以用于快速生成图像,例如个性化头像、艺术创作、以及数据增强等任务,促进生成式AI的普及。

📄 摘要(原文)

Diffusion models have shown strong capabilities in generating high-quality images from text prompts. However, these models often require large-scale training data and significant computational resources to train, or suffer from heavy structure with high latency. To this end, we propose Efficient Multimodal Diffusion Transformer (E-MMDiT), an efficient and lightweight multimodal diffusion model with only 304M parameters for fast image synthesis requiring low training resources. We provide an easily reproducible baseline with competitive results. Our model for 512px generation, trained with only 25M public data in 1.5 days on a single node of 8 AMD MI300X GPUs, achieves 0.66 on GenEval and easily reaches to 0.72 with some post-training techniques such as GRPO. Our design philosophy centers on token reduction as the computational cost scales significantly with the token count. We adopt a highly compressive visual tokenizer to produce a more compact representation and propose a novel multi-path compression module for further compression of tokens. To enhance our design, we introduce Position Reinforcement, which strengthens positional information to maintain spatial coherence, and Alternating Subregion Attention (ASA), which performs attention within subregions to further reduce computational cost. In addition, we propose AdaLN-affine, an efficient lightweight module for computing modulation parameters in transformer blocks. Our code is available at https://github.com/AMD-AGI/Nitro-E and we hope E-MMDiT serves as a strong and practical baseline for future research and contributes to democratization of generative AI models.