arXiv:2411.04168v4 公告类型：替换-交叉 摘要：我们介绍了一种新颖的状态空间架构，有效地利用空间和频率信息，增强扩散模型在图像生成任务中对输入图像局部特征的归纳偏置。尽管状态空间网络，包括Mamba，这是一种革命性的循环神经网络进步，通常从左到右扫描输入序列，但在处理图像数据时，它们在设计有效的扫描策略方面面临困难。我们的方法表明，将小波变换整合到Mamba中，增强了对视觉输入局部结构的意识，并通过将频率解耦成小波子带，分别表示低频和高频分量，更好地捕捉了频率的长期关系。这些基于小波的输出随后通过交叉注意力融合层与原始Mamba输出无缝融合，结合了空间和频率信息以优化状态空间模型的顺序意识，这是图像生成中细节和整体质量的关键。此外，我们引入了一个全局共享变换器来增强Mamba的性能，利用其捕捉全局关系的独特能力。通过在标准基准上的广泛实验，我们的方法在DiT和DIFFUSSM上都表现出了优越的结果，实现了更快的训练收敛并提供了高质量的输出。代码和预训练模型可在https://github.com/VinAIResearch/DiMSUM.git上获得。