304M参数模型用少量数据快速生成高清图像,适合资源有限场景。
E-MMDiT: Revisiting Multimodal Diffusion Transformer Design for Fast Image Synthesis under Limited Resources
- 通过压缩视觉令牌和多路径压缩模块减少计算量
- 512像素图像生成仅需2500万数据,1.5天训练达GenEval 0.66分
- 轻量化设计适合低资源环境,代码开源助力普及
扩散模型在文本到图像生成中表现优异,但通常需要大量数据和算力。为此,我们提出高效多模态扩散Transformer(E-MMDiT),仅含304M参数,可在低资源下实现快速图像合成。该模型在单个8 AMD MI300X GPU节点上,仅用2500万公开数据训练1.5天,即可实现512px图像生成,GenEval得分为0.66,结合GRPO等后训练技术可达0.72。设计核心为降低令牌数量,采用高度压缩的视觉分词器,并引入新型多路径压缩模块进一步压缩令牌。为增强性能,提出位置强化机制以保持空间一致性,以及交替子区域注意力(ASA)以降低计算开销。此外,设计了轻量级的AdaLN-affine模块用于高效计算调制参数。代码已开源,旨在为后续研究提供实用基线,推动生成式AI的普惠化。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。