End-to-End Multi-Modal Diffusion Mamba

AI summary

Overview

Research area: End-to-end multi-modal generative modeling — architectures that handle images and text within a single model, combining diffusion-based generation with Mamba state-space sequence modeling.

Technical level: Advanced. The paper assumes familiarity with diffusion models, variational autoencoders, state-space models, and multi-modal benchmarks.

Scope: The paper introduces MDM (Multi-modal Diffusion Mamba), a 7-billion-parameter unified architecture that encodes and decodes images and text through one shared VAE and a Mamba-based multi-step selection diffusion decoder, and evaluates it on image generation, captioning, VQA, and text reasoning tasks.

What This Paper Is About

Most large multi-modal systems bolt together separate encoders and decoders for each modality, which the authors argue prevents the model from learning a genuinely shared joint representation across images and text, and slows inference. This paper proposes a single unified pipeline — MDM — in which one variational autoencoder maps both images and text into a shared noisy latent space, and a Mamba-based multi-step selection diffusion decoder reconstructs either modality from that space using one objective. The goal is to generate high-resolution images and long text sequences efficiently without sacrificing multi-modal understanding.

Key Contributions

1. A unified end-to-end architecture (MDM) with a computational complexity of O(MLN²), which the authors contrast against the O(ML²N/G) complexity of prior end-to-end models such as MonoFormer. This is presented as the enabler for efficient long-sequence text and high-resolution image generation.

2. A multi-step selection diffusion model that merges autoregressive and diffusion-based generative paradigms into a single unified learning objective, so that multi-modal data can be generated simultaneously within one diffusion process.

3. Empirical validation across many tasks, including image generation on ImageNet and COCO, image captioning on Flickr30K and COCO, VQA on VQAv2, VizWiz and OKVQA, text comprehension and reasoning on seven datasets, and math/world-knowledge tasks on GSM8k, MATH, and MMLU.

4. An efficient decoder design built on Mamba-2 blocks with image and text "scan switches" (four directions for images, two for text, following the settings of DiM), plus a Score Entropy Loss objective that replaces Markov-chain-based denoising updates.

Main Findings

• Image generation is competitive but not first overall. On ImageNet 256×256, MDM (7B, diffusion) reaches FID 2.49, IS 281.4, Precision 0.86, and Recall 0.59. The paper states this is second in FID, IS, and Precision and third in Recall across the compared models. For reference, MonoFormer (1.1B) reports FID 2.57, IS 272.6, Precision 0.84, Recall 0.56; DiT-XL/2 (675M) reports FID 2.27, IS 278.2.

• Text-to-image results. MDM reports FID 5.91 and Gen Eval 0.68 on COCO 256×256. The paper claims MDM reached the top three on these metrics and achieved state-of-the-art on Gen Eval. Comparable reported Gen Eval scores include Transfusion 0.63, Dual-DiT 0.65, JanusFlow 0.70, SD-3 0.68, and Chameleon 0.39.

• Captioning performance. The MDM baseline scores CIDEr 62.4 on Flickr 30K and 109.6 on COCO; the instruction-tuned variant, InstructMDM, scores 75.2 on Flickr 30K and 122.1 on COCO. The paper states MDM ranked second among all models on captioning.

• VQA is competitive but behind the leaders. MDM scores 60.3 on VQAv2, 39.8 on VizWiz, and 47.1 on OKVQA; InstructMDM raises these to 66.7, 46.3, and 51.6. The paper notes MDM surpasses several traditional models including InstructBLIP, mPLUG-Owl, DREAMLLM, and Emu, while trailing the top-performing models.

• Text-only tasks favor traditional models. The authors report that MDM and other end-to-end multi-modal models perform worse than well-known traditional models on text-to-text tasks, which they attribute to dropping separate language encoders, visual encoders, and multi-modal fusion encoders. Against other end-to-end models, MDM outperforms MonoFormer and surpasses Chameleon on seven out of ten datasets.

• Instruction tuning helps across the board. InstructMDM (7B) exceeds the MDM baseline on every metric listed in the text and reasoning table, including HellaSwag (74.8 vs 70.6), MMLU (59.2 vs 54.4), PIQA (83.7 vs 79.9), and GSM8k (46.0 vs 40.5).

• Ablation results. Reducing image/text scan switches from six (①②③④/①②) to three (①②/①) raises FPS from 1.357 to 1.405 but worsens FID from 2.49 to 3.96. Replacing Mamba with a Transformer raises FPS further to 1.914 but degrades FID to 6.72.

• Long-sequence efficiency. The paper reports MDM outperforming other end-to-end multi-modal models on inference speed for sequences exceeding 1280 tokens, and that MDM generates image and caption data simultaneously. Some qualitative outputs are described as exceeding GPT-4V.

Methodology in Plain English

The authors build a single network that treats images and text the same way at the latent level. Images are split into patches, text is tokenized with SentencePiece using Unigram BPE and embedded, and both are pushed through one shared variational autoencoder that produces Gaussian parameters (mean and variance). Noise is then added to samples from that distribution, which is the forward diffusion step.

The decoder is the novel part. Instead of a Transformer, it uses 49 Mamba-2 blocks, each with dimension 2048, forming a 7-billion-parameter model. Two mechanisms drive it. First, a multi-step selection diffusion process gradually removes noise, trained under a Score Entropy Loss rather than a Markov-chain objective — the authors argue this stabilizes denoising and improves sampling quality on discrete data. Second, "scan switches" read the latent sequence in several different temporal directions (four for image patches, two for text), and the Mamba state-space structure then decides which items to attend to and which to ignore at each denoising step, based on a derived condition that the score ratio should approximate the true data density ratio. The Step is written as a second-order numerical update derived from DPM-Solver.

Three kinds of learnable padding tokens (time, class, pad) are inserted to mark the diffusion step, the data category, and sequence boundaries. Training combines an L2 reconstruction loss for images, cross-entropy for text, a KL divergence term on the latent space, and the Score Entropy Loss. Optimization used the AdamW optimizer without weight decay, a constant learning rate of 0.0001, and an EMA of model weights with coefficient 0.9999, with a DDPM scheduler and DPM-Solver for sampling. Baseline training used ImageNet, JourneyDB, and UltraChat; the instruction variant was tuned on COCO, GQA, OCR-VQA, TextVQA, and VisualGenome.

Why This Matters

If a single model can encode, decode, and generate images and text through one latent space and one objective, it removes the interface engineering that dominates today's multi-modal stacks and could reduce inference cost for long inputs. The paper's complexity argument — linear-ish scaling through Mamba rather than the quadratic cost of Transformers — targets exactly the regime where current end-to-end models struggle: high-resolution images and long text sequences.

Real-world applications suggested by the tasks evaluated:

• Simultaneous image-and-caption generation for content creation and accessibility tooling (the paper demonstrates generating an image of a dog alongside its description).
• Visual question answering for assistive and information-retrieval systems, evaluated on VQAv2, VizWiz, and OKVQA.
• Image captioning for photo organization, media indexing, and dataset annotation, evaluated on Flickr30K and COCO.
• General text comprehension, reasoning, and math/world-knowledge assistance, evaluated on HellaSwag, OpenBookQA, Wino-Grande, ARCEasy, ARCChallenge, BoolQ, PIQA, GSM8k, MATH, and MMLU.

Industry relevance: The paper positions MDM against GPT-4V, Gemini Pro, and Mistral, and reports gains over end-to-end baselines such as MonoFormer, LlamaGen, and Chameleon. Its efficiency claim is the main industrial hook, since serving high-resolution image and long-context text generation is where inference cost dominates. The caveat the authors themselves raise is that on text-only tasks, end-to-end models still lag behind traditional multi-encoder systems.

Future Directions

• Closing the text-only gap. The authors observe that end-to-end models underperform traditional models on text comprehension and reasoning and speculate this stems from abandoning dedicated language, visual, and fusion encoders. Whether the unified latent space can be improved to remove that penalty is unresolved.

• Scaling and modality expansion. The model is evaluated at 7 billion parameters with 49 Mamba blocks and dimension 2048; the paper does not report scaling behavior, nor does it extend beyond images and text to video or audio.

• Pushing the efficiency claims further. The ablation shows a clear speed-versus-quality tradeoff between six and three scan switches, and between Mamba and Transformer blocks. Finding configurations that retain FID 2.49-level quality at higher FPS remains open.

• Strengthening instruction tuning. The instruction-tuned variant beats the baseline on every reported metric, but the paper does not explore how far that tuning can be pushed or whether it closes the gap to top VQA systems such as GPT-4V and Gemini Pro on VizWiz and OKVQA.

Target Audience

Researchers and engineers working on multi-modal generative models, diffusion architectures, and state-space models, particularly those interested in replacing Transformer backbones or unifying modality pipelines. It is also relevant to practitioners evaluating the tradeoff between end-to-end simplicity and benchmark performance, and to readers tracking whether Mamba-based designs can compete with Transformers at scale. The mathematical sections on Score Entropy and the selection theorem make it more suitable for readers with a background in probabilistic generative modeling than for general audiences.

Authors’ abstract

Current end-to-end multi-modal models utilize different encoders and decoders to process input and output information. This separation hinders the joint representation learning of various modalities. To unify multi-modal processing, we propose a novel architecture called MDM (Multi-modal Diffusion Mamba). MDM utilizes a Mamba-based multi-step selection diffusion model to progressively generate and refine modality-specific information through a unified variational autoencoder for both encoding and decoding. This innovative approach allows MDM to achieve superior performance when processing high-dimensional data, particularly in generating high-resolution images and extended text sequences simultaneously. Our evaluations in areas such as image generation, image captioning, visual question answering, text comprehension, and reasoning tasks demonstrate that MDM significantly outperforms existing end-to-end models (MonoFormer, LlamaGen, and Chameleon etc.) and competes effectively with SOTA models like GPT-4V, Gemini Pro, and Mistral. Our results validate MDM's effectiveness in unifying multi-modal processes while maintaining computational efficiency, establishing a new direction for end-to-end multi-modal architectures.