Source-linked AI summary
ZigMa: A DiT-style Zigzag Mamba Diffusion Model
Vincent Tao Hu, Stefan Andreas Baumann, Ming Gui, Olga Grebenkova, Pingchuan Ma, Johannes Schusterbauer, Björn Ommer
TL;DR
Diffusion models face scalability and quadratic-complexity challenges, while existing Mamba vision methods often overlook spatial continuity in their scans. The paper introduces zero-parameter Zigzag Mamba within Stochastic Interpolants and reports stronger performance than Mamba-based baselines across larger-scale visual settings.
Problem
Diffusion models face scalability and quadratic-complexity challenges, and existing Mamba vision methods often neglect spatial continuity when scanning visual tokens.
Method
The paper introduces Zigzag Mamba, a plug-and-play zero-parameter heterogeneous layerwise scan method, and applies it within the Stochastic Interpolant framework to images and videos.
Results
Zigzag Mamba outperforms related Mamba-based baselines and is evaluated on large-scale visual settings including 1024×1024 images and videos.
Takeaways & Limitations
Spatially continuous layerwise scans provide the paper’s supported route for applying Mamba to scalable visual generation across 2D and 3D data.
Takeaways & Limitations
The scan schemes are selected empirically and may be sub-optimal, while GPU constraints prevented exploration of longer training durations.
Abstract
from arXiv · showhide
The diffusion model has long been plagued by scalability and quadratic complexity issues, especially within transformer-based structures. In this study, we aim to leverage the long sequence modeling capability of a State-Space Model called Mamba to extend its applicability to visual data generation. Firstly, we identify a critical oversight in most current Mamba-based vision methods, namely the lack of consideration for spatial continuity in the scan scheme of Mamba. Secondly, building upon this insight, we introduce a simple, plug-and-play, zero-parameter method named Zigzag Mamba, which outperforms Mamba-based baselines and demonstrates improved speed and memory utilization compared to transformer-based baselines. Lastly, we integrate Zigzag Mamba with the Stochastic Interpolant framework to investigate the scalability of the model on large-resolution visual datasets, such as FacesHQ $1024\times 1024$ and UCF101, MultiModal-CelebA-HQ, and MS COCO $256\times 256$ . Code will be released at https://taohu.me/zigma/
1 Introduction
The paper targets diffusion-model scalability and quadratic complexity by adapting Mamba’s efficient long-sequence modeling to visual data. It identifies spatial continuity as a key gap in existing 2D/3D Mamba scans and proposes Zigzag Mamba within Stochastic Interpolants for larger-scale visual generation.
- 1 Introduction: Diffusion models face scalability challenges, and transformer-based approaches retain a quadratic-complexity bottleneck with sequence length despite efficiency techniques.This motivates investigating State-Space Models as an alternative backbone for scalable generation.
- 1 Introduction: Existing Mamba vision methods often flatten 2D tokens in row- or column-major order, neglecting spatial continuity, while multi-directional scans add parameters and GPU-memory burden.The paper frames spatial continuity as the central challenge in applying a 1D sequence model to 2D images and 3D videos.
- 1 Introduction: Stochastic Interpolants provide a generalized framework spanning normalizing flows, diffusion models, flow matching, and Schrödinger Bridges, motivating tests beyond prior 256×256 and 512×512 settings.The paper targets more complex scenarios, including 1024×1024 images and videos.
- 1 Introduction: Zigzag Mamba is a plug-and-play, zero-parameter layerwise scan paradigm designed to preserve spatial continuity when extending Mamba from sequences to images and videos.The method distributes scan complexity across layers and incorporates visual inductive bias without adding parameters.
- 1 Introduction: The authors report that Zigzag Mamba outperforms related Mamba-based baselines and explores Stochastic Interpolants on large-scale images and videos.The contribution summary identifies both the benchmark comparison and the expanded application scope.
2 Related Works
Related work spans State-Space Models, Mamba-based visual modeling, scan-curve designs, and diffusion-generation frameworks. The paper positions Zigzag Mamba as a simpler generative alternative focused on spatially continuous scans and distributed scan complexity.
- 2 Related Works: Mamba offers efficient long-sequence modeling through work-efficient parallel scanning and data-dependent innovations, but applying its 1D formulation to visual data remains challenging.Prior work includes applications across images, video, restoration, graphs, point clouds, and image generation.
- 2 Related Works: VisionMamba uses bidirectional SSMs for discriminative tasks, whereas Zigzag Mamba applies a simple alternative to generative diffusion models.The paper distinguishes its focus on backbone design and text-conditioned visual generation from related discriminative approaches.
- 2 Related Works: S4ND adds local convolution and Mamba-ND uses multiple scans within one block, while Zigzag Mamba distributes scan complexity across layers with zero parameter burden.This design aims to maximize the incorporation of visual inductive bias.
- 2 Related Works: The paper relates its approach to scan-curve research such as Hilbert curves, while reporting that preliminary Hilbert-curve results did not work well for its setting.The supplied passage presents this as an empirical contrast rather than a general impossibility.
- 2 Related Works: Diffusion research includes SDE- and ODE-based score-generation frameworks, while Stochastic Interpolants generalize across several generative-model families.The related-work context places Zigzag Mamba within broader generative-model development.
3 Method
ZigMa adapts Mamba to visual diffusion by preserving spatial continuity through layerwise token rearrangements, while extending the design to text-conditioned images and videos. Its scan structure retains linear sequence complexity and factorizes 3D spatial and temporal reasoning.
- Zigzag Mamba: ZigMa identifies spatial continuity as critical when generalizing Mamba from 1D sequences to 2D images and 3D videos.The method targets the mismatch between image neighborhoods and conventional flattened scan orders.
- DiT-Style Network: A DiT-style backbone uses single-scan Mamba blocks, with arrange-rearrange operations providing positional awareness across layers.The architecture follows the DiT framework and uses layer-specific rearrangement pairs around the forward scan.
- Zigzag Mamba: Eight space-filling zigzag schemes are cycled across layers as Ω_i = S_{i%8}, preserving continuous patch traversal without adding parameters.The input tokens are rearranged before scanning and restored afterward through the reverse permutation.
- Text Conditioning: Cross-attention and MLP conditioning extend ZigMa to timestep- and text-conditioned generation while retaining Mamba-based long-sequence modeling.The conditioning pathway separately modulates the Mamba scan and cross-attention mechanism.
- 3D Video Extension: For video, factorized 3D Zigzag applies separate spatial and temporal Mamba blocks because fully 3D zigzag scanning produced suboptimal optimization.The temporal component uses forward and backward scans, while spatial processing uses 2D zigzag scans.
- Complexity: Self-attention has quadratic sequence complexity, whereas Mamba and Zigzag Mamba are linear; Zigzag Mamba also removes the k-direction multiplier.The comparison uses sequence length M and fixed state parameter N = 16 by default.
4 Experiment
Experiments evaluate ZigMa across image and video datasets, scanning schemes, position embeddings, computational efficiency, and transformer comparisons. ZigMa improves with more zigzag scans, scales to high-resolution generation, and achieves competitive quality with favorable efficiency.
- Scan Scheme Ablation: Zigzag scanning consistently improves results over Sweep, with larger gains as the number of scans increases from 1 to 8.The relative gain is more prominent at 512 × 512 than at 256 × 256.
- Position Embedding Ablation: Learnable position embeddings perform best, while ZigMa remains nearly unchanged between sinusoidal and no position embeddings.The ablation reports marginal gains from learnable embeddings and stronger spatial inductive-bias incorporation than the baseline.
- Efficiency Ablation: ZigMa maintains GPU memory and FPS as Order Receptive Field increases, while Parallel, Bidirectional, and Vision Mamba lose FPS.ZigMa with an Order Receptive Field of 8 is faster without changing parameters.
- Transformer Comparison: Compared with transformer-based methods, ZigMa achieves comparable generation performance with significantly less memory consumption and fewer FLOPS.The comparison includes U-ViT and DiT, while Figure 6 also reports FPS and GPU memory across patch dimensions.
- Main Results: ZigMa outperforms the Bidirectional Mamba baseline on 1024×1024 FacesHQ and consistently outperforms the baseline on MS COCO.FacesHQ uses 4,096 latent-space tokens; the MS COCO comparison reports stronger performance for ZigMa, especially with 8 scans.
5 Conclusion
The conclusion presents Zigzag Mamba as a DiT-style diffusion model built within the Stochastic Interpolant framework. It addresses spatial continuity through heterogeneous layerwise scans and extends the design to 3D data through factorization.
- 5 Conclusion: Zigzag Mamba combines a DiT-style layout with heterogeneous layerwise scanning to use spatial inductive bias in 2D images.The model is developed within the Stochastic Interpolant framework.
- 5 Conclusion: The 3D Mamba design is factorized into 2D and 1D Zigzag Mamba components to facilitate optimization.
6 Limitations and Future Work
The authors identify empirical scan selection and limited training duration as limitations, while proposing broader applications of ZigMa as future work. They also anticipate applicability to other linear attention models.
- 6 Limitations and Future Work: The scan schemes are chosen empirically because the paper does not exhaustively enumerate spatially continuous zigzag paths, which may yield sub-optimal performance.
- 6 Limitations and Future Work: GPU resource constraints prevented experiments with longer training durations, although the authors anticipate similar conclusions.
- 6 Limitations and Future Work: Future work will explore ZigMa for long-sequence modeling across additional applications and domains.
- 6 Limitations and Future Work: The authors anticipate that the scan path may also suit linear attention models including RWKV, xLSTM, HGRN, and GLA.
7 Impact Statement
The impact statement frames ZigMa as a scalable approach for high-fidelity large-image generation and text-to-image conditioning. It also acknowledges the risk that such capabilities could enable harmful image generation.
- 7 Impact Statement: ZigMa targets scalable, high-fidelity generation of large images and supports text-to-image generation through cross-attention in the Mamba block.
- 7 Impact Statement: The authors acknowledge a risk of enabling harmful image generation as with other large-scale image synthesis systems.
8 Appendix
The appendix examines scan-path design, spatial continuity, model variants, and their effects on visual generation performance. It also reports that increasing spatial continuity improves FID.
- 8 Appendix: The appendix compares Hilbert and Zigzag scan paths under different Order Receptive Fields on unconditional MultiModal-CelebA256.Table 7 defines the comparison scope, while the accompanying figures illustrate Hilbert and Peano-related scan structures.
- 8 Appendix: The appendix lists text-to-image methods and ZigMa model variants, including Small, Base, Large, and XLarge configurations with optional cross-attention.Table 9 identifies the model sizes and defines CA as cross-attention for text-to-image conditioning.
- 8 Appendix: Visual results are reported for FacesHQ 1024 × 1024 and MS-COCO, with samples shown in Figures 15 and 16.
- 8.2 Spatial Continuity is Critical: Increasing spatial continuity through larger patch groups improves performance, while random shuffling produces notably inferior results.The analysis evaluates 2×2, 4×4, 8×8, and 16×16 patch grouping with Zigzag-8 at the group level.
- 8.2 Spatial Continuity is Critical: Spatial continuity is identified as a critical requirement when applying Mamba to 2D sequences.
8.3 Visualization
The visualization appendix presents high-fidelity samples across resolutions and analyzes Hilbert scanning as an alternative to Zigzag scanning. It reports that Hilbert’s complex structure is harder for the SSM to optimize.
- 8.3 Visualization: Generated samples maintain high visual fidelity across FacesHQ 1024 × 1024 and MultiModal-CelebA 512 × 512.Figure 12 visualizes results across these resolutions.
- 8.4 New Result about the Scanning Scheme: The appendix evaluates Hilbert scan variants with different angles and starting points while keeping parameters consistent.A Gilbert algorithm is used to preserve continuity across square sizes.
- 8.4 New Result about the Scanning Scheme: Hilbert scanning performs worse than Zigzag scanning because its complex structure appears to hinder SSM optimization on natural images.The authors hypothesize that structure may matter more than locality in generative tasks.
- 8.4 New Result about the Scanning Scheme: Zigzag scanning is presented as the simplest case of a Peano curve.
8.5 New Result of 2D visual data
This section studies ZigMa design choices and efficiency across 2D visual-generation settings. It reports stable speed and memory with larger receptive order, while larger patches worsen FID.
- 8.5 New Result of 2D visual data: The default ZigMa configuration is the Base model, while cross-attention is optional because it adds parameter and speed burdens.
- 8.5 New Result of 2D visual data: FID deteriorates as patch size increases, indicating that smaller patches are important for optimal performance.The ablation compares patch sizes 1, 2, 4, and 8.
- 8.5 New Result of 2D visual data: FPS varies with patch size in the appendix’s dedicated comparison.Figure 13 is labeled “FPS v.s. Patch Size.”
- 8.5 New Result of 2D visual data: Zigzag Mamba maintains GPU memory and FPS as receptive order increases, and Zigzag-8 can run faster without changing parameters.Parallel Mamba, Bidirectional Mamba, and Vision Mamba lose FPS as parameters increase.
- 8.5 New Result of 2D visual data: Zigzag Mamba consistently outperforms Sweep-1 and Sweep-2 in the loss and FID trends.Figure 14 compares the training loss and FID curves across resolutions on MultiModal-CelebA.
- 8.5 New Result of 2D visual data: In-context attention performs worse than cross-attention for text-to-image generation.The authors hypothesize that discontinuity between text tokens and image patch tokens contributes to this difference.
- 8.5 New Result of 2D visual data: The factorized 3D Zigzag Mamba uses the sst scheme, prioritizing spatial complexity over temporal information.The authors hypothesize that temporal redundancy exists and leave other s/t combinations for future work.
8.7 More related works
Related work covers multidimensional Mamba applications, token-shuffling methods, and patch-order studies. ZigMa differs by distributing scan complexity across network layers to preserve continuous spatial paths.
- 8.7 More related works: Prior work includes Mamba applications to medical imaging, restoration, graphs, NLP, tabular data, human motion, point clouds, and image generation.
- 8.7 More related works: ZigMa distributes scan complexity across every network layer, unlike Mamba-ND’s use of various scans within a single block.
- 8.7 More related works: Patch-order studies have focused on autoregressive transformers, whereas ZigMa focuses on Mamba-based structures.
- 8.7 More related works: ZigMa shuffles tokens to maintain a continuous spatial-filling scan path, with different orders across layers.This approach is related to prior spatial or channel shuffling methods.
8.8 More Details
This section details the model’s indexing, conditioning, training, evaluation, and dataset procedures. It also specifies the metrics, sampling protocols, and visualization resources used for assessment.
- Double-Indexing Issue for Ωi: Spatial-continuous Mamba reasoning uses arrange-and-rearrange indexing along the token dimension, while the authors note that indexing can be time-consuming for large token counts.The paper notes that torch.compile() can substantially ease this timing issue.
- Double-Indexing Issue for Ωi: The formulation assumes Mamba-based networks are permutation equivariant to token order and requires 50% fewer indexing operations.The supplied passages present this reduction as a comparison point for the indexing formulation.
- Evaluation Metrics: Evaluation uses FID, KID, and FDD for image fidelity, with FDD added because FID does not fully reflect human opinions, and FID/FVD for video fidelity.Image metrics use 5,000 real and 5,000 fake images, while video metrics use 200 sampled videos.
- More Details: The implementation details include EMA models for evaluation, optional learnable position encodings, and dataset-specific hyperparameters summarized in Table 10.The authors state that Mamba’s token order automatically encodes position, while extra learnable position encoding can nevertheless improve performance relative to the baseline.
- The conditioning of timestep and prompt: Text-conditioned experiments use MultiModal-CelebA and MS COCO text-image pairs, whose captions are encoded with a CLIP text encoder and supplied to the network as token sequences.COCO and MultiModal-CelebA typically provide 5 to 10 captions per image.
- The conditioning of timestep and prompt: Timestep and prompt conditions are incorporated by concatenating a condition token with image-patch tokens in the Mamba block.The conditioning process is described as being implemented according to Algorithm 1.