将扩散Transformer转为专家混合模型,60%参数激活量下保持性能
Dense2MoE: Restructuring Diffusion Transformer to MoE for Efficient Text-to-Image Generation
- 用专家混合(MoE)替换前馈层,激活参数减少62.5%
- 通过选择性激活模块提升稀疏性,激活参数总量降低60%
- 多步蒸馏优化确保转换后性能不降,适合大模型高效生成
扩散Transformer(DiT)在文生图任务中表现卓越,但参数量庞大导致推理开销高。现有压缩方法多依赖剪枝,激进剪枝常因模型容量下降引发性能显著退化。为此,本文首次将密集型DiT重构为专家混合(MoE)结构,实现有结构的稀疏化,在保留模型容量的同时减少激活参数。具体地,将DiT模块中的前馈网络(FFN)替换为MoE层,使FFN激活参数减少62.5%;进一步提出混合块(MoB),仅选择性激活部分DiT块以增强稀疏性。为确保从密集到MoE的有效转换,设计了多步蒸馏流程:基于泰勒指标的专家初始化、带负载均衡的知识蒸馏,以及用于MoB优化的分组特征损失。将大型扩散模型(如FLUX.1 [dev])重构为MoE结构后,激活参数减少60%,性能与原模型持平,并优于剪枝类方法。Dense2MoE为高效文生图生成建立了新范式。
原文摘要 · Abstract (English)
Diffusion Transformer (DiT) has demonstrated remarkable performance in text-to-image generation; however, its large parameter size results in substantial inference overhead. Existing parameter compression methods primarily focus on pruning, but aggressive pruning often leads to severe performance degradation due to reduced model capacity. To address this limitation, we pioneer the transformation of a dense DiT into a Mixture of Experts (MoE) for structured sparsification, reducing the number of activated parameters while preserving model capacity. Specifically, we replace the Feed-Forward Networks (FFNs) in DiT Blocks with MoE layers, reducing the number of activated parameters in the FFNs by 62.5\%. Furthermore, we propose the Mixture of Blocks (MoB) to selectively activate DiT blocks, thereby further enhancing sparsity. To ensure an effective dense-to-MoE conversion, we design a multi-step distillation pipeline, incorporating Taylor metric-based expert initialization, knowledge distillation with load balancing, and group feature loss for MoB optimization. We transform large diffusion transformers (e.g., FLUX.1 [dev]) into an MoE structure, reducing activated parameters by 60\% while maintaining original performance and surpassing pruning-based approaches in extensive experiments. Overall, Dense2MoE establishes a new paradigm for efficient text-to-image generation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。