UniDiffuser: Unified Multi-Modal Diffusion

• UniDiffuser is a unified multi-modal generative framework that recasts all queries as noise prediction tasks with modality-specific timesteps.
• It employs a transformer-based (U-ViT) architecture to handle mixed image and text token streams, achieving near-specialist performance across diverse tasks.
• Empirical results demonstrate competitive FID and CLIP scores along with robust performance in applications from text-to-image generation to robotic control.

UniDiffuser is a unified multi-modal generative modeling framework based on diffusion processes, designed to parameterize and sample marginal, conditional, and joint distributions over diverse modalities such as images and text with a single transformer backbone. The principal innovation is the observation that all relevant generative queries—image generation, text generation, conditional and joint synthesis—are reducible to noise prediction tasks under modality-specific timesteps within a shared diffusion infrastructure. UniDiffuser's transformer-based architecture (U-ViT variant) natively accommodates mixed-modality token streams and achieves near-specialist performance, extensibility, and scalability across a suite of multi-modal generative and predictive tasks (Bao et al., 2023, Team et al., 30 Apr 2026).

1. Unified Multi-Modal Diffusion Formulation

The core of UniDiffuser is a multi-modal extension of denoising diffusion probabilistic models. Given $M$ modalities (e.g., image $X$ and text $Y$), each data pair $(x^X_0, x^Y_0)$ is processed via independent Gaussian perturbations: $$x^m_t = \alpha^m_t x^m_0 + \sigma^m_t \epsilon^m, \quad \epsilon^m \sim \mathcal{N}(0, I)$$ with independent timesteps $t^m$ for each modality, enabling flexible marginalization or conditioning. The joint forward process factorizes across modalities as products of Gaussian transitions.

The reverse process uses a single transformer-based noise predictor $\epsilon_\theta(\{x^m_t\}; t^1, ..., t^M)$ producing predicted perturbations $[\epsilon^1, ..., \epsilon^M]$. The model simultaneously learns to denoise all modalities under a mean-squared error objective: $$\mathcal{L}(\theta) = \mathbb{E}\| \epsilon_\theta(x^X_{t^X}, x^Y_{t^Y}, t^X, t^Y) - [\epsilon^X, \epsilon^Y] \|_2^2$$ Special cases—marginal, conditional, and joint distributions—are obtained by fixing or tying modality-specific timesteps. For example, setting $t^Y = T$ (max noise) marginalizes over $X$0 (unconditional image generation), whereas $X$1 (no noise) conditions on $X$2 (text-to-image) (Bao et al., 2023).

Continuous-time extensions apply in domains like video and action prediction, where diffusion ODEs are defined per modality and the model predicts velocity fields for video latents and action tokens, subject to SNR-based timestep sampling and flow-matching objectives (Team et al., 30 Apr 2026).

2. Transformer-Based Architecture

The UniDiffuser backbone employs a ViT-style transformer capable of handling mixed image and text token streams, as well as numeric timestep embeddings for each modality. The input sequence comprises (i) image latent tokens (e.g., 16 × 16 CLIP-VAE patches), (ii) text latent tokens (e.g., 77 GPT-2 prefix embeddings), and (iii) special tokens encoding timesteps $X$3, $X$4.

Key design features include:

• Modality- and position-specific token embeddings.
• Entirely self-attention-based cross-modal interaction (no explicit cross-attention modules).
• Post-layer-norm transformer blocks, multiple MLP layers, and long skip connections for improved optimization and gradient flow.
• Separate linear heads to generate per-modality noise predictions from partitioned final layer embeddings.

In the MotuBrain extension for world action modeling, a Mixture-of-Transformers (MoT) architecture is implemented with three streams (video, action, and text), cross-attending over joint tokenized sequences with structure controlled by "H-bridge" patterns (decoupled attention at shallow/deep layers, joint attention in the middle), and an independent text stream reserved for semantic conditioning without denoising heads (Team et al., 30 Apr 2026).

3. Training Paradigms and Datasets

UniDiffuser is trained end-to-end on large scale paired multi-modal corpora. Image-text training leverages up to 600 million high-quality pairs from LAION datasets using the following schedule:

• Stage 1: LAION2B-en, 256×256 resolution, 250K steps, large batch size, AdamW optimizer, 5K step warmup.
• Stage 2: Higher-resolution LAION subset, learning rate decay on plateau, 200K steps.
• Stage 3: LAION-aesthetics at 512×512, 220K steps, optimizer state resumed from Stage 2.

All tasks share a single training loss and network backbone, which supports seamless scalability and parameter-efficiency. Mixed precision training over 88 A100 (80 GB) GPUs reaches completion in approximately 28 days (Bao et al., 2023).

In robotic policy and world-modeling settings (MotuBrain), training involves multi-view video streams, instruction text, and action trajectories. Multiview video is processed via shared VAEs and 3D rotary positional embeddings; heterogeneous robot actions are unified by relative pose parameterization anchored to a reference frame and flattened into continuous tokens (Team et al., 30 Apr 2026).

4. Unified Inference Modes

UniDiffuser supports a range of query types at inference by manipulating modality-specific timesteps and clamping:

• Unconditional generation: Fix one modality's timestep at maximum (fully noised), sample the other.
• Conditional generation (e.g., text-to-image): Set the conditioning modality's timestep to zero (clean), sample the other.
• Joint image-text generation: Initialize both modalities at identical noise levels, denoise jointly.
• Image-to-text generation: Swap roles in the above.
• Variation and interpolation: Achieved through blocked Gibbs sampling or synchronous denoising.

The reverse diffusion is implemented with DDPM-style update equations, and classifier-free guidance is realized without extra training or null-token permutations by linearly combining predictions from different noise schedules (Bao et al., 2023). In MotuBrain, five operational modes are enabled (policy, world model, inverse dynamics, video generation, joint video–action prediction) by modulating which tokens are given or sampled, all with the same core parameters (Team et al., 30 Apr 2026).

Pseudocode Overview for Inference Modes

Task	Clamp/Noise	Key Parameter Setting
Unconditional image	$X$5 (text = noise)	Sample $X$6; denoise via reverse process
Text-to-image	$X$7 (text = data)	Sample $X$8 conditioned on $X$9
Image-to-text	$Y$0 (image = data)	Sample $Y$1 conditioned on $Y$2
Joint image+text	$Y$3	Denosing both jointly
Video/action/robotics modes	modulated clamping	Per-mode configuration (Team et al., 30 Apr 2026)

5. Empirical Results and Quantitative Performance

On MS-COCO text-to-image generation (10K zero-shot samples), UniDiffuser achieves:

• FID: 9.71 (vs. Stable Diffusion 8.59; Versatile Diffusion 10.09)
• CLIP-score (ViT-L/14): ≈ 0.26 at optimal guidance, superior to Versatile Diffusion (≈0.24)

For image-to-text generation, UniDiffuser achieves CLIP-score ≈ 0.515 (Vs. 0.495 for Versatile Diffusion). Model size is 952M parameters with a sampling runtime of 19.8s per 10 samples (25 steps) on an A100-80GB, peak memory 48GB (Bao et al., 2023).

In the robot control domain, MotuBrain attains average success rates of 95.8% (RoboTwin 2.0, clean) and 96.1% (randomized), leads in EWMScore (WorldArena benchmark), and can adapt to new embodiments with minimal demonstration data (Team et al., 30 Apr 2026).

6. Core Technical Insights

Several conceptual advances underpin UniDiffuser's efficacy:

1. Unified Diffusion Perspective: Marginal, conditional, and joint generation are recast as noise prediction under per-modality timestep scheduling, enabling a parameter-efficient architecture to support diverse tasks by adjusting clamping/noising.
2. Minimal Architectural Overhead: One backbone, a single loss (MSE over concatenated noise predictions), and unified forward/backward passes suffice for all generative and inferential settings.
3. Transformer Cross-Modal Handling: The U-ViT and MoT backbones leverage joint self-attention for modality fusion without explicit cross-attention modules, supporting arbitrary token mixtures.
4. Classifier-Free Guidance Intrinsic to Framework: By virtue of modeling all relevant distributional variants, classifier-free guidance mechanisms entail simply linear recombination of existing predictions, with no retraining or null-token requirement.
5. Scalability and Generality: Empirical performance on strict quantitative metrics is comparable to specialist architectures in vision and language domains, while uniquely supporting additional variation, interpolation, and multi-modal generation/translation tasks in a parameter-shared manner.
6. Real-World Application Generality: In robotics, the UniDiffuser formulation supports world modeling, policy, forward/inverse dynamics, and joint prediction within a single system, robust to heterogeneous multimodal data and capable of low-latency real-time deployment (Bao et al., 2023, Team et al., 30 Apr 2026).

Together, these attributes establish UniDiffuser as a foundational approach for scalable, unified multi-modal generative modeling and downstream control.