Mixture of States: Routing Token-Level Dynamics for Multimodal Generation

AI summary

Overview

Research area: Multimodal generative modeling — specifically the fusion mechanism that connects a language/vision understanding model to a diffusion-based image generator. The paper sits at the intersection of diffusion models, mixture-of-experts-style dynamic routing, and vision-language architectures.

Technical level: Advanced. The paper assumes familiarity with diffusion/flow matching, transformer blocks, cross-attention and self-attention fusion, and mixture-of-transformers designs.

Scope: The paper proposes Mixture of States (MoS), a learnable token-wise router that dynamically selects hidden states from a frozen understanding tower to condition a trainable generation tower, and validates it on text-to-image generation (MoS-Image) and instruction-based image editing (MoS-Edit).

What This Paper Is About

Multimodal diffusion models need to inject text (or text plus a reference image) into an image generator, but the two sides are built with different objectives and often different architectures, so aligning them is hard. Existing fusion designs are rigid: cross-attention uses only the final text-encoder layer, self-attention concatenates everything at quadratic cost, and Mixture-of-Transformers (MoT) forces a strict one-to-one layer correspondence between symmetrically sized towers. The goal of this paper is a fusion mechanism that lets any visual token, at any denoising step, pull features from any layer of an arbitrarily sized understanding model.

Key Contributions

1. Mixture of States (MoS): A flexible fusion mechanism whose interactions are dynamic (conditioned on the denoising timestep), sparse (top-k selection over layers), and state-based (full layer-level hidden states rather than same-layer key-value vectors), enabling asymmetric understanding and generation towers.
2. Two model families as a scaling blueprint: MoS-Image and MoS-Edit, spanning 3B to 5B learnable parameters, which use a 100M-parameter router and are built on a frozen understanding tower plus a from-scratch generation tower.
3. State-of-the-art results at small parameter scale: MoS models match or surpass counterparts up to 4 times larger across GenEval, DPG, WISE, oneIG, ImgEdit and GEdit, with lower end-to-end latency than Qwen-Image and Bagel on a single A100 GPU.
4. A systematic ablation of router design principles: adaptive layer selection versus fixed layers, timestep-dependent versus static conditioning, and token-specific versus sample-wise routing.

Main Findings

• Dynamic conditioning wins: Routing on the prompt alone gave FID 21.12 and CLIP 21.40 on MJHQ; adding the noised latent gave FID 21.89 and CLIP 21.53; adding the timestep on top of both gave the best result at FID 20.15 and CLIP 21.74.
• Token-specific beats sample-wise: A single shared routing solution predicted from a [CLS] token gave FID 21.66 and CLIP 21.48, while per-token routing gave FID 20.17 and CLIP 21.63.
• Learned routing beats hand-crafted routing: Evenly skipping text-encoder layers gave FID 21.51 and CLIP 22.04, versus FID 17.77 and CLIP 22.91 for MoS.
• MoS beats cross-attention: Under the same data and parameter budget, a 5B cross-attention baseline scored 0.74 GenEval and 83.40 DPG, while MoS scored 0.79 GenEval and 85.61 DPG.
• MoS beats MoT consistently: Under identical parameters, data and compute, MoS outperformed MoT across all training stages on GenEval and DPG.
• Leading benchmark results at 3B and 5B: On GenEval, MoS-S scored 0.89 and MoS-L 0.90, versus 0.87 for the 20B Qwen-Image, 0.81 for the 4.8B SANA-1.5, 0.66 for the 12B FLUX.1[Dev], and 0.88 for the 14B Bagel. On DPG, MoS-S scored 86.33 and MoS-L 87.01, versus 88.32 for Qwen-Image and 84.70 for SANA-1.5.
• Strong editing scores: On GEdit, MoS-S scored 7.41 and MoS-L 7.86, versus 7.56 for Qwen-Image, 6.52 for Bagel, on a scale of 0 to 10. On ImgEdit, MoS-S scored 4.17 and MoS-L 4.33, versus 4.27 for Qwen-Image and 3.20 for Bagel, on a scale of 0 to 5.
• WISE benchmark: MoS-S scored 0.47 and MoS-L 0.54; both trailed Qwen-Image's 0.62 but exceeded Bagel's 0.52, FLUX.1[Dev]'s 0.50, and SANA-1.5's unpublished result (not reported).
• Frozen understanding tower still helps: Applying self-CoT improved WISE from 0.54 to 0.65 for MoS-L and from 0.47 to 0.55 for MoS-S without additional training, and the paper notes the frozen understanding tower's own understanding performance is unchanged and therefore omitted.
• Negligible router cost: With a 3B generation tower generating a 1024×1024 image on a single A100 GPU, the router adds only 0.008s per iteration.
• Training efficiency: The four-stage MoS-Image pipeline required roughly 3,000 A100 days in total, compared with the 6,250 A100 days reported for Stable Diffusion v1.5.

Methodology in Plain English

The model has two towers. An understanding tower (frozen during training) reads the context — the text prompt for generation, or text plus a reference image for editing — and produces hidden states from every one of its layers. A generation tower (trained from scratch) performs the diffusion/rectified-flow process that turns noise into an image. Standing between them is a small router, the paper's main invention.

At every denoising step, the router looks at three things: the prompt, the current noisy latent, and a sinusoidal embedding of the timestep. It is a lightweight transformer of about 100M parameters that outputs, for every context token, a matrix of weights describing how much each understanding-tower layer should contribute to each generation-tower layer. Each generation block independently picks the top-k contributing hidden states, reweights them, projects them to its own dimensionality, and concatenates them with the visual tokens for in-context processing. During training, an epsilon-greedy strategy occasionally picks k random layers instead of the top-k ones to avoid a bad local optimum; at inference, selection is deterministic top-k with epsilon set to 0.

Training runs in four progressive stages: 512×512 low resolution, 1024×1024 high resolution, then aesthetic tuning, then 2048×2048 super-resolution tuning. The editing model is an extension requiring additional training on paired editing data. Ablations were run at a smaller scale using LLaMA3.2-3B as the understanding tower and a 1B generation tower, with the MoT comparison using a symmetric 1B/1B setup because MoT requires matched depths.

Why This Matters

Research impact: The paper challenges the assumption that text and visual backbones must be symmetric and layer-aligned. It shows that a sparse, learned, timestep-aware router can bridge architecturally mismatched models, which opens the door to pairing any pretrained understanding model with any diffusion generator. It also reframes mixture-of-experts-style routing as inter-model collaboration rather than intra-model adaptivity.

Real-world applications (as demonstrated or directly implied by the paper's tasks):

• Text-to-image generation with precise multi-entity and attribute following, including scenes with more than four entities and dense visual text.
• Instruction-based image editing, such as modifying a subject's attributes while preserving the reference image.
• 2K-resolution image synthesis via the super-resolution fine-tuning stage.
• Cost-sensitive deployment where a large frozen understanding model is run only once and a small generation model handles synthesis.

Industry relevance: The efficiency case is the headline. A 5B model matching or surpassing a 20B model (4 times larger) on most benchmarks, with lower end-to-end latency and higher throughput than Qwen-Image and Bagel on the same A100 platform, matters for serving costs. The roughly 3,000 A100-day training budget, against 6,250 A100 days for Stable Diffusion v1.5, is also a concrete argument for cheaper model development.

Future Directions

• From one-way to dual-way interaction: The paper states that MoT has shown strong scalability under early-fusion training, while MoS's effectiveness in early-fusion settings remains to be validated. The suggested principled extension is to give the router multiple projection layers to enable bidirectional interaction (the limitation section is truncated at this point).
• Router architecture exploration: The paper defers a deeper dissection of router architecture and of which features are most beneficial to route to the appendix.
• Scaling behavior: The paper confirms MoS adheres to standard scaling laws, but only two variants (MoS-S and MoS-L) are released, leaving larger-scale extrapolation as an open question.
• Inference-time optimizations: The paper reports that MoS benefits from common inference-time optimizations, but how far these can be pushed alongside routing is not fully explored.

Target Audience

This paper is most useful for researchers and engineers working on diffusion-based generative models, multimodal fusion architectures, and efficient large-scale training. Practitioners deciding how to wire a pretrained language or vision-language model into an image generator — especially those constrained by asymmetric model sizes or inference cost — will find the router design and ablation results directly actionable. Readers new to diffusion models or mixture-of-experts routing will find the architecture details dense, since the paper assumes familiarity with flow matching, transformer blocks, and existing fusion baselines.

Authors’ abstract

We introduce MoS (Mixture of States), a novel fusion paradigm for multimodal diffusion models that merges modalities using flexible, state-based interactions. The core of MoS is a learnable, token-wise router that creates denoising timestep- and input-dependent interactions between modalities' hidden states, precisely aligning token-level features with the diffusion trajectory. This router sparsely selects the top-$k$ hidden states and is trained with an $ε$-greedy strategy, efficiently selecting contextual features with minimal learnable parameters and negligible computational overhead. We validate our design with text-to-image generation (MoS-Image) and editing (MoS-Editing), which achieve state-of-the-art results. With only 3B to 5B parameters, our models match or surpass counterparts up to $4\times$ larger. These findings establish MoS as a flexible and compute-efficient paradigm for scaling multimodal diffusion models.