面向高分辨率图像生成的扩散 Transformer-to-Mamba 蒸馏
计算机视觉与模式识别
2025-06-25 v1
摘要
自注意力机制在扩散变换器（DiT）中的二次计算复杂度，在高分辨率图像生成中引入了巨大的计算成本。虽然线性复杂度的 Mamba 模型作为一种潜在替代方案，但直接训练 Mamba 仍然在实践中存在挑战。为此，本文提出一种扩散 Transformer-to-Mamba 蒸馏（T2MD）方法，构建了一个高效的训练流程，促进从基于自注意力的变换器向线性复杂度状态空间模型 Mamba 的过渡。我们建立了一个融合自注意力和 Mamba 的混合模型，同时实现了效率和全局依赖性。通过提出的层级教师强制和基于特征的知识蒸馏，T2MD 缓解了从头训练状态空间模型的难度和高昂成本。我们从提炼出的 512×512 分辨率基模型开始，通过轻量级适配和高分辨率微调将生成推进至 2048×2048 图像。实验表明，我们的训练路径在低开销但高质量文本到图像生成方面表现出色。重要的是，我们的结果也证明了使用顺序和因果 Mamba 模型生成非因果视觉输出的可行性，为未来的探索指明了潜力。
引用
@article{arxiv.2506.18999,
title = {Diffusion Transformer-to-Mamba Distillation for High-Resolution Image Generation},
author = {Yuan Yao and Yicong Hong and Difan Liu and Long Mai and Feng Liu and Jiebo Luo},
journal= {arXiv preprint arXiv:2506.18999},
year = {2025}
}