Diffusion Mamba Transformer

• Diffusion Mamba Transformer is a class of neural architectures that replaces or augments traditional self-attention with efficient Mamba SSM layers.
• It achieves linear or near-linear complexity by using spatially-aware scanning and dynamic masking to preserve key structural inductive biases.
• Hybrid models combining Mamba SSMs with self-attention demonstrate significant computational savings and improved performance across image, language, and multimodal tasks.

A Diffusion Mamba Transformer is a class of neural architectures for diffusion models that employs the Mamba state-space model (SSM) as a principal or hybrid backbone, with the goal of replacing or augmenting the traditional Transformer self-attention mechanism. This design enables more efficient modeling—especially for image, language, music, and multimodal generative tasks—by leveraging the linear or near-linear time complexity of the Mamba SSM, often combined with carefully engineered attention, fusion, or scanning schemes to preserve essential structural inductive biases. The approach is motivated by the need to overcome the quadratic scaling constraints of self-attention in high-resolution, long-sequence, or resource-constrained diffusion generative models across a range of domains.

1. State-Space Mamba: Mathematical Formalism and Advantages

The core of a Diffusion Mamba Transformer is the use of the Mamba SSM layer, which models feature sequences with a discretized linear recurrence: $$h_t = A\,h_{t-1} + B\,x_t,\qquad y_t = C\,h_t + D\,x_t$$ where $x_t$ is the input embedding at position $t$, $h_t$ the latent state, $y_t$ the output, and $A, B, C, D$ are learned or input-dependent matrices. In practice, Mamba layers adopt selective scanning mechanisms with input-dependent parameters, allowing each token or patch to dynamically adjust the state-space recurrence behavior (Wang et al., 2024, Mo et al., 2024, Fei et al., 2024).

This linear recurrence admits fast, parallel implementations using prefix-scan or convolutional methods, yielding per-layer compute and memory costs that grow linearly with sequence length ($O(N)$), in stark contrast to the $O(N^2)$ complexity of self-attention (Mo et al., 2024, Teng et al., 2024, Hu et al., 2024).

2. Network-Level Architectures: Pure, Hybrid, and Hierarchical Variants

Several principal architectural paradigms have emerged:

• Pure Mamba Backbones: All blocks are Mamba SSMs (possibly with bidirectional or multi-directional recurrences). This achieves maximal linearity and efficiency, as in DiM for image/video (Mo et al., 2024, Mo, 2024), DiffMa for CT-to-MRI (Wang et al., 2024), and DiffuApriel for language modeling (Singh et al., 19 Nov 2025). Performance often matches or exceeds Transformer/U-Net baselines with a fraction of GFLOPs and memory, especially for high resolution or long sequences.
• Hybrid Mamba–Transformer Models: Self-attention (for explicit pairwise/global context) and Mamba SSMs (for efficient long-range propagation) are interleaved or fused at various granularities. Examples include block-level alternation (Fei et al., 2024), sparse attention injection for global context (Singh et al., 19 Nov 2025), local windowed attention plus SSM (Fu et al., 2024), and globally-shared Transformer modules (Phung et al., 2024). Hybridization enables models to trade off speed and global context adaptively.
• U-Net/Hierarchical Integration: Within encoder–decoder hierarchies, Mamba blocks replace convolution/attention blocks for both patch-level and latent-level sequence modeling (USM (Ergasti et al., 18 Apr 2025), LaMamba-Diff (Fu et al., 2024)). This allows for progressive reduction/restoration of sequence length while maintaining state-space propagation of global context.

A summary of representative architectures:

Model	Attention Mechanism	Mamba Integration	Scaling	Notable Benchmarks
DiffMa (Wang et al., 2024)	Spiral cross-sequence (soft-masked)	Pure SSM stack	Linear time	CT→MRI (SSIM↑/PSNR↑)
DiM (Mo et al., 2024)	None	Bidirectional SSM	Linear time	ImageNet (FID, IS)
Dimba (Fei et al., 2024)	Block-alternating cross-attn	Hybrid Transformer–SSM	Hybrid	COCO, User study
LaMamba-Diff (Fu et al., 2024)	Local (windowed), no global	U-Net w/ local attn+SSM	Linear time	ImageNet, FID/IS
DiMSUM (Phung et al., 2024)	Cross-attn fusion + periodic GST	SSM + periodic Transform	Hybrid	CelebA, LSUN
SMDIM (Yuan et al., 27 Jul 2025)	Periodic sparse (MFA block)	SSM + self-attn	Near-linear	Symbolic music

3. Inductive Biases: Scanning, Masking, and Structural Continuity

To address issues inherent in flattening multi-dimensional signals for SSM processing, Diffusion Mamba Transformers incorporate inductive bias schemes:

• Spatially-Aware Scanning: Sequentialization strategies such as spiral-scan (Wang et al., 2024), zigzag or multiple scan directions (Hu et al., 2024, Teng et al., 2024), and scan-switching (Lu et al., 15 Oct 2025) preserve local neighborhood continuity and enhance the modeling of spatial correlations otherwise lost in naïve rasterization.
• Attention Masking and Dynamic Importance: Soft-masked cross-sequence attention (as in DiffMa (Wang et al., 2024)) employs a learned mask to modulate attention weights, emphasizing diagnostically or semantically important patches/tokens during denoising.
• Wavelet and Frequency Fusion: Models such as DiMSUM (Phung et al., 2024) and Proffusion-WM (Zhang et al., 6 May 2025) combine classic SSM/Mamba with wavelet transforms and frequency-domain decomposition, augmenting spatial SSM propagation with multi-resolution or frequency-selective information.

These schemes are empirically shown to improve convergence rate, generation fidelity, and stability, especially for high-dimensional or structured modalities.

4. Hybridization with Self-Attention: Design, Trade-Offs, and Empirical Performance

Hybrid Diffusion Mamba Transformer models clarify the contexts in which linear SSMs suffice and when attention is needed:

• Interleaved Attention: By alternating Mamba and self-attention layers, models such as Dimba (Fei et al., 2024) and DiffuApriel-H (Singh et al., 19 Nov 2025) recover the global context coverage of pure Transformer models while gaining significant throughput and memory advantages (up to 4.4× in language, 20–30% memory reduction in text-to-image).
• Local Attention Integration: Windowed self-attention within Mamba blocks, as in LaMamba-Diff (Fu et al., 2024), captures local detail without incurring the quadratic cost of global attention. This configuration retains linear or near-linear scaling while achieving SOTA FID for high-resolution images at a fraction of GFLOPs and parameters compared to DiT.
• Transformer-Driven Distillation: Teacher–student training via blockwise teacher forcing from a DiT teacher to a Mamba student (T2MD (Yao et al., 23 Jun 2025)) enables high-fidelity, high-resolution image synthesis even up to 4K resolution, with improved sample efficiency and convergence.

A plausible implication is that SSM-based models can, with minimal attention injection and the appropriate structural biases or training strategies, rival or surpass pure transformers in both efficiency and generative quality for many modalities.

5. Domain-Specific Adaptations and Applications

Diffusion Mamba Transformer architectures are operationalized across diverse domains:

• Medical Imaging: DiffMa (Wang et al., 2024) and MD-Dose (Fu et al., 2024) utilize Mamba SSM blocks for CT→MRI and radiation dose prediction, offering superior SSIM, PSNR, and MSE compared to ViT/U-Net at lower computational cost.
• High-Resolution Image/Video: DiM (Mo et al., 2024, Teng et al., 2024) demonstrates linear time/space scaling for large-scale images and videos, supporting efficient training, inference, and training-free upsampling.
• Language and Music: DiffuApriel (Singh et al., 19 Nov 2025) (text), SMDIM (Yuan et al., 27 Jul 2025) and Proffusion-WM (Zhang et al., 6 May 2025) (symbolic music) show that bidirectional or hybrid SSM+attention stacks outperform transformers alone in long-range sequence fidelity, with drastic reductions in parameter and memory footprints.
• Autonomous Driving and Policy Learning: Pi-DiMT (Zhou et al., 31 Jan 2026) interleaves Mamba with self-attention and physics-inspired modules to produce reliable and physically plausible motion planning. GMF-Drive (Wang et al., 8 Aug 2025) replaces transformer fusion with spatially-aware SSM modeling in BEV spaces, achieving new SOTA on NAVSIM.

6. Computational Complexity, Training Strategies, and Benchmarks

The defining feature of all variants is favorable scaling:

• Complexity:
  • Mamba SSM layer: $O(ND)$ per layer, where $N$ is sequence length.
  • Pure transformer (self-attention): $O(N^2D)$ per layer.
  • Hybrid: Interleaving Mamba and attention reduces net cost to $O(ND + N^2D/K)$ for K:1 Mamba:Attention.
  • Empirically, models achieve 10–80% reduction in GFLOPs, 1.5–4.4× speedups, and flexible scaling to “extreme” sequence lengths ($L=2\times2048$ for image, $L=65\,\mathrm{K}$ for text).
• Training Protocols:
  • Weak-to-strong curriculum (low- to high-resolution) (Teng et al., 2024, Yao et al., 23 Jun 2025).
  • Teacher-forcing/feature distillation to stabilize Mamba layer training (Yao et al., 23 Jun 2025).
  • Joint diffusion and auxiliary (wavelet/score entropy/contrastive) losses to align inductive biases and empirical gradients (Wang et al., 2024, Lu et al., 15 Oct 2025).
• Benchmarks and Empirical Gains:
  • FID, IS, SSIM, PSNR improvements across image synthesis.
  • Substantial memory and wall-clock savings in sample and batch throughput.
  • Enhanced multimodal and reasoning abilities in hybrid and unified architectures (Lu et al., 15 Oct 2025).

7. Current Limitations and Future Directions

Despite the empirical benefits, challenges remain:

• Context-Length/Model-Size Scaling: Pure SSMs may struggle with non-causal dependencies in purely spatial signals, though bidirectionality and hybridization mitigate this (Mo et al., 2024). For tasks requiring complex compositionality, sparse or adaptive attention may be required.
• Implementation Complexity: Efficient parameterization of dynamic SSM transition matrices and the design of multi-directional scans adds some overhead compared to standard transformer blocks.
• Theory and Generalization: Deeper theoretical understanding of SSM-attention trade-offs, optimal scan/inductive bias design, and non-causal-causal model distillation remains open. Further study is also necessary for real-world deployment in unstructured, multi-modal, and long-context settings (Mo et al., 2024, Zhou et al., 31 Jan 2026).

Ongoing research explores adaptive hybrid scheduling (Fei et al., 2024), learned scan patterns (Phung et al., 2024), large-scale multi-modal fusion (Lu et al., 15 Oct 2025), and the extension of linear-complexity generative modeling to video, 3D, and beyond (Mo, 2024, Mo et al., 2024, Lu et al., 15 Oct 2025).

---

References:

• Soft Masked Mamba Diffusion Model for CT to MRI Conversion (Wang et al., 2024)
• Scaling Diffusion Mamba with Bidirectional SSMs for Efficient Image and Video Generation (Mo et al., 2024)
• Dimba: Transformer-Mamba Diffusion Models (Fei et al., 2024)
• LaMamba-Diff: Linear-Time High-Fidelity Diffusion Models Based on Local Attention and Mamba (Fu et al., 2024)
• DiMSUM: Diffusion Mamba -- A Scalable and Unified Spatial-Frequency Method for Image Generation (Phung et al., 2024)
• Diffusion Mamba for Efficient High-Resolution Image Synthesis (Teng et al., 2024)
• DiffuApriel: High-Throughput Diffusion LMs with Mamba Backbone (Singh et al., 19 Nov 2025)
• Diffusion Transformer-to-Mamba Distillation for High-Resolution Image Generation (Yao et al., 23 Jun 2025)
• Symbolic Music Diffusion with Mamba (Yuan et al., 27 Jul 2025)
• Mamba-Diffusion Model with Learnable Wavelet for Controllable Symbolic Music Generation (Zhang et al., 6 May 2025)
• U-Shape Mamba: State Space Model for faster diffusion (Ergasti et al., 18 Apr 2025)
• ZigMa: A DiT-style Zigzag Mamba Diffusion Model (Hu et al., 2024)
• End-to-End Multi-Modal Diffusion Mamba (Lu et al., 15 Oct 2025)
• Physics-informed Diffusion Mamba Transformer for Real-world Driving (Zhou et al., 31 Jan 2026)
• GMF-Drive: Gated Mamba Fusion with Spatial-Aware BEV Representation for End-to-End Autonomous Driving (Wang et al., 8 Aug 2025)
• Efficient 3D Shape Generation via Diffusion Mamba with Bidirectional SSMs (Mo, 2024)