arXiv:2409.19937cs.CV2024-09被引 11

MaskMamba融合Mamba与Transformer,加速图像生成并提升质量。

MaskMamba: A Hybrid Mamba-Transformer Model for Masked Image Generation

论文配图:MaskMamba: A Hybrid Mamba-Transformer Model for Masked Image Generation
图 1 · 摘自论文原文
  • 用标准卷积替换因果卷积,提升全局上下文建模能力。
  • 采用拼接替代乘法,推理速度比Transformer快54.44%(2048×2048)。
  • 支持文本/类别到图像生成,适合高效图像合成场景。

图像生成模型因依赖基于Transformer的骨干网络,面临可扩展性差和二次复杂度问题。本文提出MaskMamba,一种结合Mamba与Transformer架构的混合模型,采用掩码图像建模实现非自回归图像合成。通过两项关键改进重设计双向Mamba:(1) 以标准卷积替代因果卷积,更好捕捉全局上下文;(2) 使用拼接而非乘法,显著提升性能并加快推理速度。此外,研究了多种混合方案,包括串行与分组并行结构。引入上下文条件,使模型可同时完成类别到图像和文本到图像生成任务。实验表明,MaskMamba在生成质量上超越Mamba和Transformer基线模型,在2048×2048分辨率下推理速度提升54.44%。

原文摘要 · Abstract (English)

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.

图像生成Mamba混合模型加速推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。