扩散Transformer到Mamba蒸馏用于高分辨率图像生成
Diffusion Transformer-to-Mamba Distillation for High-Resolution Image Generation
AI总结：
本文提出扩散Transformer到Mamba蒸馏（T2MD），通过混合模型和知识蒸馏，实现从DiT到Mamba的高效训练，支持高分辨率图像生成，验证了因果Mamba用于非因果视觉任务的可行性。
摘要：
扩散Transformer（DiT）中自注意力的二次计算复杂度在高分辨率图像生成中引入了大量的计算成本。虽然线性复杂度的Mamba模型成为一种潜在的替代方案，但直接训练Mamba在经验上仍然具有挑战性。为了解决这个问题，本文提出了扩散Transformer到Mamba蒸馏（T2MD），形成了一种高效的训练流程，促进了从基于自注意力的Transformer到线性复杂度状态空间模型Mamba的过渡。我们建立了一个扩散自注意力和Mamba混合模型，同时实现了效率和全局依赖。通过提出的层级教师强制和基于特征的知识蒸馏，T2MD缓解了从零开始训练状态空间模型的难度和高成本。从蒸馏的512×512分辨率基础模型开始，我们通过轻量级适应和高分辨率微调将生成推向2048×2048图像。实验表明，我们的训练路径导致了低开销但高质量的文本到图像生成。重要的是，我们的结果也证明了使用顺序和因果Mamba模型生成非因果视觉输出的可行性，暗示了未来探索的潜力。
英文摘要：
The quadratic computational complexity of self-attention in diffusion transformers (DiT) introduces substantial computational costs in high-resolution image generation. While the linear-complexity Mamba model emerges as a potential alternative, direct Mamba training remains empirically challenging. To address this issue, this paper introduces diffusion transformer-to-mamba distillation (T2MD), forming an efficient training pipeline that facilitates the transition from the self-attention-based transformer to the linear complexity state-space model Mamba. We establish a diffusion self-attention and Mamba hybrid model that simultaneously achieves efficiency and global dependencies. With the proposed layer-level teacher forcing and feature-based knowledge distillation, T2MD alleviates the training difficulty and high cost of a state space model from scratch. Starting from the distilled 512$\times$512 resolution base model, we push the generation towards 2048$\times$2048 images via lightweight adaptation and high-resolution fine-tuning. Experiments demonstrate that our training path leads to low overhead but high-quality text-to-image generation. Importantly, our results also justify the feasibility of using sequential and causal Mamba models for generating non-causal visual output, suggesting the potential for future exploration.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
