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

TL;DR

MoS (Mixture of States) is a fusion paradigm for multimodal diffusion transformers that replaces fixed cross-attention / MoT-style modality coupling with a small learnable, token-wise router. For each image token at each denoising step, the router sparsely picks top-k layers of the text encoder’s hidden states to read from, so the conditional signal is per-token, per-timestep, and adaptive in layer depth rather than a single fixed final-layer text embedding shared by all tokens. Trained with an ε-greedy exploration schedule, MoS-Image and MoS-Editing reach SOTA on text-to-image generation and editing at 3B–5B params, matching or surpassing models up to 4× larger.

Key claims

Using a single fixed layer of the text encoder (typically the final layer) or rigid one-to-one layer alignment between text and image streams is suboptimal — diffusion models do not consume language features in a strictly sequential or layer-aligned manner [Community note, §Adaptive Layer Selection].
Standard text-to-image systems encode the text once and keep it static throughout denoising, which creates an information mismatch with the evolving diffusion trajectory; conditional signals should adapt to the noise level / denoising step [Community note, §Input-Dependent Conditional Signals].
Each token should be allowed to source its representation adaptively from different text-encoder layers, rather than sharing a single layer embedding across all tokens [Community note, §Token-Specific Conditioning].
The router is token-wise, denoising-timestep-dependent, and input-dependent — it selects which text hidden states each image token attends to as a function of both the token and the diffusion step [Abstract].
The router uses top-k sparse selection over text-encoder hidden states and is trained with an ε-greedy exploration strategy, with minimal learnable parameters and negligible compute overhead [Abstract].
The design is presented as “orthogonal” to the prior fusion families — cross-attention, self-attention (concat-and-attend), and Mixture-of-Transformers — rather than as a refinement of them [Community note].
MoS-Image (text-to-image) and MoS-Editing (image editing) at 3B–5B parameters match or surpass counterparts up to 4× larger on the reported benchmarks [Abstract].

Method

MoS sits at the modality-fusion layer of a multimodal diffusion transformer. Where MMDiT-style architectures concatenate text and image tokens and run joint self-attention, and cross-attention architectures pull a fixed text feature (usually the final-layer embedding) into the image stream, MoS makes the read different along three axes at once: per image token, per denoising timestep, and per text-encoder layer.

Concretely, the text encoder produces a stack of hidden states across its layers. For each image (query) token at each denoising step, a small learnable router emits a sparse top-k selection over those layers’ hidden states; the selected hidden states are then aggregated as the conditional context for that token. The router is conditioned on the token, on the input, and on the diffusion timestep, so the same image token can read from different text layers at noise-heavy steps than at near-clean steps, and different image tokens at the same step can read from different layers. Training uses an ε-greedy schedule so the router explores layer choices early and exploits them later.

The router itself is parameter- and compute-light (the paper emphasizes “negligible computational overhead”), which is what lets a 3B–5B model match much larger ones — the gain is from how the existing text features are routed, not from extra capacity.

Results

MoS-Image and MoS-Editing achieve top scores on text-to-image generation and editing at 3B–5B parameters, matching or surpassing counterparts up to 4× larger [Abstract].
The HuggingFace paper page lists no quantitative benchmark numbers in the abstract; specific FID / GenEval / editing-benchmark numbers are not extracted here (arxiv PDF was rate-limited at filing time).

Why it’s interesting

MoS is doing something different from the dual-stream MMDiT recipe the wiki has been accumulating (Dual-stream diffusion transformer): instead of two parallel transformer streams coupled by bidirectional cross-attention at every block, it keeps a single image-side denoiser and makes the text read the adaptive part — token, layer, and timestep at once. That makes it the third distinct multimodal-fusion paradigm in the wiki, alongside cross-attention (LTX-2, SkyReels-V4) and MoT/concat (MMDiT lineage, Exploring the Deep Fusion of Large Language Models and Diffusion Transformers for Text-to-Image Synthesis). The “diffusion models don’t consume text features in a strictly layer-aligned manner” claim is the headline — if it holds, it argues that the standard practice of taking the last-layer text embedding leaves capability on the table for free. The mechanism also rhymes with the routing-redesign cluster in MoE Routing Design: it is a router over states, not a router over experts — a sibling primitive to Raven’s “router over recurrent slots” (Raven Part 1 — Memory as a Set of Slots), suggesting a broader trend of token-wise sparse routers being repurposed across architectural layers (FFN-experts → recurrent slots → multimodal hidden states).