MaskMamba:一种用于掩码图像生成的混合 Mamba-Transformer 模型
MaskMamba: A Hybrid Mamba-Transformer Model for Masked Image Generation
AI总结:
本文提出混合 Mamba-Transformer 模型 MaskMamba,通过改进双向 Mamba 架构并引入上下文条件,在提升图像生成质量的同时实现推理速度大幅提升。
摘要:
图像生成模型由于主要依赖基于 Transformer 的主干网络,一直面临可扩展性和二次复杂度方面的挑战。在本研究中,我们提出了 MaskMamba,这是一种结合 Mamba 与 Transformer 架构的新型混合模型,利用掩码图像建模进行非自回归图像合成。我们通过两项关键修改精心重新设计了双向 Mamba 架构:(1)用标准卷积替换因果卷积,以更好地捕捉全局上下文;(2)使用拼接代替乘法,在显著提升性能的同时加快推理速度。此外,我们探索了 MaskMamba 的多种混合方案,包括串行和分组并行两种排列方式。进一步地,我们引入了上下文条件,使模型能够同时执行类别到图像和文本到图像的生成任务。我们的 MaskMamba 在生成质量上优于基于 Mamba 和基于 Transformer 的模型。值得注意的是,在 2048×2048 分辨率下,其推理速度相比 Transformer 实现了 54.44% 的显著提升。
英文摘要:
Image generation models have encountered challenges related to scalability and quadratic complexity, primarily due to the reliance on Transformer-based backbones. In this study, we introduce MaskMamba, a novel hybrid model that combines Mamba and Transformer architectures, utilizing Masked Image Modeling for non-autoregressive image synthesis. We meticulously redesign the bidirectional Mamba architecture by implementing two key modifications: (1) replacing causal convolutions with standard convolutions to better capture global context, and (2) utilizing concatenation instead of multiplication, which significantly boosts performance while accelerating inference speed. Additionally, we explore various hybrid schemes of MaskMamba, including both serial and grouped parallel arrangements. Furthermore, we incorporate an in-context condition that allows our model to perform both class-to-image and text-to-image generation tasks. Our MaskMamba outperforms Mamba-based and Transformer-based models in generation quality. Notably, it achieves a remarkable $54.44\%$ improvement in inference speed at a resolution of $2048\times 2048$ over Transformer.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。