Diffusion模型内存爆炸？DiM的4个扫描方向黑科技实测报告

Diffusion模型内存爆炸？DiM的4个扫描方向黑科技实测报告

最近在部署一个高分辨率图像生成服务时，我又一次被显存不足的报错信息“教育”了。客户的需求很明确：生成1024x1024甚至更高清的营销素材图，但现有的主流扩散模型，无论是基于Transformer还是U-Net的架构，一旦分辨率超过512，显存占用就会呈指数级增长，推理速度也慢得让人心焦。这几乎是所有从事AIGC落地的工程师都会遇到的“天花板”。就在我准备妥协，建议客户降低分辨率要求时，一篇名为“DiM: Diffusion Mamba for Efficient High-Resolution Image Synthesis”的论文进入了视野。它提出的“多方向扫描”和一系列工程优化，号称能线性处理高分辨率图像，这听起来简直是为解决我们的痛点量身定做的。但论文归论文，实际效果如何？在真实的工程环境中，它的显存占用、生成质量和推理速度到底有没有宣传的那么神？带着这些疑问，我决定进行一次彻底的实测，从256x256到1024x1024，对比不同扫描模式，并分享在调参过程中关于“可学习填充令牌”的一些实战经验。这份报告，就是为所有同样被内存瓶颈困扰的同行准备的。

1. 高分辨率图像生成的困境与DiM的破局思路

当我们谈论高分辨率图像生成时，核心矛盾在于计算复杂度与硬件资源。传统的扩散模型骨干，尤其是Vision Transformer (ViT)，在处理图像时会将图像分割成一个个图像块（Patch），并将这些块视为一个序列。注意力机制的计算成本与序列长度的平方成正比。这意味着，当图像分辨率从256x256提升到1024x1024时，序列长度可能增加16倍，而注意力计算的开销则会暴增256倍。这直接导致了训练和推理时令人绝望的显存占用和计算时间。

Mamba 模型的出现提供了一种新思路。作为一种基于结构化状态空间模型（SSM）的架构，Mamba在处理长序列时具有线性复杂度的优势，这使其在语言建模等领域大放异彩。然而，图像是二维信号，Mamba原生的因果扫描（从前到后处理序列）模式会引入强烈的方向性偏见，这对于需要全局上下文的图像生成任务是致命的。

DiM（Diffusion Mamba）的创新之处，就在于它巧妙地解决了“如何让一维的Mamba理解二维图像”这个根本问题。其核心是一套组合拳：

• 多方向扫描：这是DiM的灵魂。它打破了单一方向的因果依赖，通过交替执行四种扫描路径（行主序、行逆序、列主序、列逆序），让每个图像块都能“看到”来自不同方向的上下文信息，从而构建起近似全局的感受野。
• 可学习填充令牌：在每行、每列的序列末尾插入特殊的、可训练的令牌。这就像一个“边界标识符”，帮助模型识别图像的空间结构，避免将图像边界误认为是普通的图像内容。
• 轻量级局部增强：在网络的输入和输出层加入轻量的深度可分离卷积。Mamba擅长捕捉长程依赖，而小卷积能有效补充模型对局部细节和纹理的建模能力，防止生成图像过于“平滑”或缺乏细节。
• “弱到强”训练策略：先在低分辨率（如256x256）上预训练，再向高分辨率（如512x512）微调。这大大加速了高分辨率模型的收敛，是一种非常实用的工程策略。

下面的表格对比了传统ViT-based Diffusion与DiM在处理高分辨率图像时的核心差异：