Dimba: Transformer-Mamba Diffusion Models

TL;DR

Dimba introduces a hybrid Transformer-Mamba diffusion model for text-to-image generation, achieving comparable quality with improved efficiency and flexibility, and revealing promising architectural properties for large-scale diffusion models.

Contribution

This work presents the first hybrid Transformer-Mamba architecture for diffusion models, combining their strengths and optimizing for large-scale, resource-constrained image generation.

Findings

Dimba achieves comparable image quality to benchmarks.

The model offers higher throughput and lower memory usage.

It demonstrates flexible adaptation to resource constraints.

Abstract

This paper unveils Dimba, a new text-to-image diffusion model that employs a distinctive hybrid architecture combining Transformer and Mamba elements. Specifically, Dimba sequentially stacked blocks alternate between Transformer and Mamba layers, and integrate conditional information through the cross-attention layer, thus capitalizing on the advantages of both architectural paradigms. We investigate several optimization strategies, including quality tuning, resolution adaption, and identify critical configurations necessary for large-scale image generation. The model's flexible design supports scenarios that cater to specific resource constraints and objectives. When scaled appropriately, Dimba offers substantial throughput and a reduced memory footprint relative to conventional pure Transformers-based benchmarks. Extensive experiments indicate that Dimba achieves comparable performance compared with benchmarks in terms of image quality, artistic rendering, and semantic control. We also report several intriguing properties of architecture discovered during evaluation and release checkpoints in experiments. Our findings emphasize the promise of large-scale hybrid Transformer-Mamba architectures in the foundational stage of diffusion models, suggesting a bright future for text-to-image generation.