DiMSUM：扩散Mamba——一种用于图像生成的可扩展统一空间-频率方法

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