Source-linked AI summary
Mixture of States: Routing Token-Level Dynamics for Multimodal Generation
Haozhe Liu, Ding Liu, Mingchen Zhuge, Zijian Zhou, Tian Xie, Sen He, Yukang Yang, Shuming Liu, Yuren Cong, Jiadong Guo, Hongyu Xu, Ke Xu, Kam-Woh Ng, Juan C. Pérez, Juan-Manuel Pérez-Rúa, Tao Xiang, Wei Liu, Shikun Liu, Jürgen Schmidhuber
TL;DR
Text-to-image generation and editing must align heterogeneous text and visual representations with the changing denoising process. MoS routes between hidden states dynamically by token and timestep; its models reach state-of-the-art results and match or surpass a model four times larger. The paper leaves human-preference alignment and router explainability for future work.
Problem
Text-to-image generation and editing must align heterogeneous text and visual representations with distinct modeling objectives and a dynamic denoising process.
Source: §1 · p. 2Method
MoS uses a learnable, token-wise router to select and aggregate text hidden states according to the input and denoising step.
Source: §1 · p. 2, §3.2 · p. 5Results
MoS achieves state-of-the-art performance across image generation and editing, and its 5B-parameter model matches or surpasses the 20B-parameter (4× larger) model (Wu et al., 2025a).
Source: §1 · p. 3, §5 · p. 11Takeaways & Limitations
MoS provides a flexible fusion approach for multimodal diffusion models, supporting asymmetric text and visual backbones.
Source: §1 · p. 3, §5 · p. 11The paper primarily uses supervised fine-tuning and leaves human-preference alignment techniques, including CoT or GRPO, for future work.
Source: §6 · p. 11
Abstract
from arXiv · showhide
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 $\epsilon$-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.
1 Introduction
Text-to-image generation and editing require aligning heterogeneous text and visual representations with a changing diffusion process. MoS uses adaptive, token-wise routing across hidden states and denoising steps, and its 3–5B-parameter models achieve state-of-the-art performance on image generation and editing.
1 Introduction
- Text–visual alignment must bridge heterogeneous representations and distinct modeling objectives, while prior fusion methods impose fixed or costly interaction patterns.Self-attention can provide deeper, bidirectional fusion but its cost scales quadratically with combined sequence length; MoT requires symmetric backbones with one-to-one block correspondence.Source: §1 · p. 2
- The authors identify adaptive layer selection, timestep-dependent conditioning, and token-specific features as principles for aligning text and visual representations.Their ablations find fixed-layer selection suboptimal, no experimental support for symmetric one-to-one alignment, and token-wise layer sourcing more effective than one shared layer embedding.Source: §1 · p. 2
- MoS gives the vision branch access to text hidden states across layers and uses a learnable token-wise router to aggregate features dynamically at each denoising step.Any visual token can attend to text tokens from any encoder layer within any visual transformer block.Source: §1 · p. 2
- MoS-Image and MoS-Edit support image generation and editing, and the approach is designed to merge asymmetric text and visual backbones.The paper presents these 3–5B-parameter models as a blueprint for scaling MoS.Source: §1 · pp. 2–3
- MoS-based models achieve state-of-the-art image-generation and editing performance at similar parameter scale; the 5B model matches or surpasses the 20B-parameter (4× larger) model (Wu et al., 2025a).The models include MoS-Image and MoS-Edit.Source: §1 · p. 3
2 Related Work
Related work spans diffusion-model conditioning, unified multimodal training, and dynamic computation. MoS differs from prior fusion designs by routing across models rather than only within a model, and by relaxing MoT’s symmetry constraints.
2 Related Work
- Text-to-image diffusion models typically condition on frozen text-encoder embeddings through cross-attention, self-attention, or layer-wise attention, leaving an initial static signal for a dynamic denoising process.MoS instead uses a learnable router to select and aggregate conditioning based on the prompt, denoising step, and noised image.Source: §2 · p. 3
- Unlike common joint-training approaches for unified tasks, MoS adopts multi-stage training and freezes the text branch to focus computation on generation components.The staged approach is described as avoiding mixed-batch throughput bottlenecks and difficulty balancing competing objectives across modalities.Source: §2 · pp. 3–4
- Dynamic neural networks typically route computation within a model; MoS extends sparse, input-conditioned computation to collaboration between models.The cited examples include MoE, MoD, and MoR.Source: §2 · p. 4
- MoT connects modality-specific transformers through shared attention but requires identical hidden dimensions and one-to-one block correspondence; MoS replaces this with sparse routing for asymmetric transformers.MoS is designed to remove those symmetry constraints.Source: §2 · p. 4
3 Mixture of States: Unifying Asymmetric Transformers
MoS uses a learned router to mediate hidden-state interactions between understanding and generation towers, adapting routing to context and denoising step. The paper applies this design to image generation and editing and trains the models through a progressive schedule.
3.1 Design Overview
- MoS uses understanding and generation towers, with full layer-level hidden states as interaction units and a learned router connecting the towers.The design supports state transfer between heterogeneous transformers with different depths.Source: §3.1 · p. 4
3.2 MoS Design Details
- The router takes timestep, noised image embedding, and context as inputs, predicts token-specific layer affinities, and selects top-k states with ε-greedy exploration during training.A lightweight router processes tokenized inputs with two bidirectional-attention transformer blocks before projecting context tokens to routing logits.Source: §3.2 · pp. 5–6
3.3 MoS for Image Generation
- For image generation, MoS routes selected context-conditioned hidden states to generation blocks as weighted, projected in-context features; inference uses deterministic top-k selection.The router predicts weights conditioned on context, timestep, and noisy latent; inference sets ε = 0.Source: §3.3 · p. 6
3.4 MoS for Image Editing
- For image editing, the understanding tower processes the source image and instruction together, and MoS routes their combined representation to the generation tower.At inference, the generation tower refines a latent using Gaussian noise and the clean reference image, then decodes the edited image.Source: §3.4 · pp. 6–7
3.5 Training MoS at Scale
- MoS-Image training progresses from 512 × 512 to 1024 × 1024, then uses aesthetic tuning and 2048 × 2048 super-resolution fine-tuning.The four stages use progressively curated data, with the final stage targeting 2K-resolution generation.Source: §3.5 · p. 7
- MoS-Image training requires ∼3,000 A100 days, compared with 6,250 A100 days reported for Stable Diffusion v1.5.The MoS-Image training schedule comprises four progressive stages.Source: §3.5 · p. 7
4 Experiments
MoS reports strong results across text-to-image generation, image editing, and image understanding, while its ablations favor adaptive routing over fixed alternatives. The experiments also report low router overhead and lower end-to-end latency than Qwen-Image and Bagel.
4.1 Systematic Ablations on Router Designs
- Ablations use a randomly selected pretraining subset and MJHQ; the default setup pairs LLaMA3.2-3B understanding with a 1B generation tower.The MoS–MoT comparison is an exception, using symmetric PLM-1B towers.Source: §4.1 · pp. 7–8
- The full router condition—prompt, noised latent, and timestep—achieves optimal performance, while token-wise prediction consistently beats prompt-wide prediction.The tested conditions include prompt alone, prompt plus latent, and prompt plus latent and timestep; Figure 5 depicts prompt-wide and per-token output configurations.Source: §4.1 · p. 8, Table 1 · p. 8, Table 2 · p. 8
- MoS surpasses hand-crafted static routing on both FID and CLIP, supporting adaptive layer selection over predefined connections.Source: §4.1 · p. 9, Table 3 · p. 9
- Under identical parameters, data, and compute, MoS consistently outperforms MoT across training stages and beats a 5B cross-attention baseline on GenEval and DPG.The cross-attention baseline scores 0.74 GenEval and 83.40 DPG; MoS scores 0.79 GenEval and 85.61 DPG.Source: §4.1 · p. 9, Figure 6 · p. 9, Table 4 · p. 9
- With a 3B generation tower, the router adds 0.008s per iteration; MoS has lower end-to-end latency than Qwen-Image and Bagel.The latency comparison uses a single A100 GPU for 1024 × 1024 image generation.Source: §4.1 · p. 9, Figure 7 · p. 10
4.2 Results on Text-to-Image Generation
- MoS-L consistently outperforms existing approaches across GenEval, DPG, WISE, and oneIG-EN.The reported comparison includes Qwen-Image at 20B parameters versus MoS at 5B; challenging examples with more than four entities and dense text also show improved prompt alignment.Source: §4.2 · pp. 9–10, Table 5 · p. 10
4.3 Results on Instruction-based Image Editing
- The 5B-parameter MoS model achieves state-of-the-art performance on ImgEdit and GEdit.Evaluations use GPT-4o and cover object-level, scene-level, text-editing, and hybrid cases.Source: §4.3 · pp. 10–11
4.4 Results on Image Understanding
- With self-CoT, WISE rises from 0.54 to 0.65 for MoS-L and from 0.47 to 0.55 for MoS-S.The understanding tower is frozen; self-CoT uses its reasoning-based captions as contextual guidance for generation, while its standalone understanding performance remains unchanged.Source: §4.4 · p. 11
5 Conclusion
MoS uses a learnable router to select hidden states dynamically across modalities, tokens, and denoising steps, supporting asymmetric understanding and generation towers. The paper reports state-of-the-art benchmark performance while rivaling or surpassing models up to four times larger at reduced computational cost.
5 Conclusion
- A learnable router selects hidden states dynamically across modalities, tokens, and denoising steps, avoiding rigid synchronization in prior architectures.Source: §5 · p. 11
- MoS supports asymmetric understanding and generation towers, executing the larger understanding model once while the generation module synthesizes visual detail.Source: §5 · p. 11
- Across multiple benchmarks, MoS rivals or surpasses models up to four times larger at markedly reduced computational cost.Source: §5 · p. 11
6 Limitation and Future Studies
The paper identifies five open directions, spanning early-fusion validation, human-preference alignment, further efficiency gains, explainability analysis, and small-object artifacts.
One-Way to Dual-Way Setting
- MoS's effectiveness in early-fusion settings remains unvalidated; a bidirectional router extension is deferred because of computational and data constraints.The proposed extension would use multiple projection layers to establish bidirectional transformer connections.Source: §6 · p. 11
Human Preference Alignment
- The models primarily use SFT post-training, leaving CoT- or GRPO-based human-preference alignment for future work.The authors state that the model can likewise benefit from such techniques.Source: §6 · p. 11
Efficiency Improvement
- Low-precision quantization, model distillation, and feature caching are possible further acceleration directions, but remain future work.The paper describes the models as relatively smaller than prior state-of-the-art models and naturally more efficient.Source: §6 · p. 11
Explainability
- Analysis of the router's relative-importance predictions for explainability is outside this work's scope.The authors identify those predictions as a potential basis for interpreting cross-modal interactions.Source: §6 · p. 11
Visual Artifacts
- The model can produce artifacts when generated objects are very small, an issue also seen in other DiT or unified models.Source: §6 · p. 12
A Implementation Details
The appendix reports implementation choices for stable, efficient training and examines router designs; its operation-space ablation favors global hidden states with a shared projection.
A Implementation Details
- Implementation choices include QK-Norm, separate modality-specific normalization, and four auxiliary learnable token registers to support stability or performance.The token registers have no explicit training objectives.Source: A Implementation Details · p. 19
- The generation tower drops timestep conditioning because the authors found negligible benefit alongside an approximately 20% parameter overhead.Source: A Implementation Details · p. 19
- Training uses FSDP and activation checkpointing, mixed precision by module, and custom Triton RMSNorm and fused-FFN kernels.The VAE remains in float32, the understanding tower uses bfloat16, and generation-tower and router collectives use bfloat16 all-gather with float32 gradient reduce-scatter.Source: A Implementation Details · p. 19
- Dynamic-resolution training assigns samples to resolution-and-aspect-ratio buckets, while dataset mixing ratios are grid-searched and diffusion-step sampling changes with resolution.Logit-normal sampling is used for 512×512 pretraining and later replaced by mode sampling for 1024×1024 and 2048 × 2048 training.Source: A Implementation Details · p. 19
- The reported setup uses AdamW at 1 × 10−4, a 4k-step warm-up, cosine decay to 1.5 × 10−5, and a top-k router with k = 2 and ϵ = 0.05.Pretraining runs 400k–1200k steps by stage, with 50k HQ fine-tuning steps.Source: A Implementation Details · p. 19
B Additional Discussions on Router Designs
- The router-design comparison considers head- or state-projected global attention and independent- or shared-projection global hidden states.Global attention transfers keys and values across branches, while the hidden-state scheme concatenates branch representations before transformer blocks.Source: B Additional Discussions on Router… · p. 20, Figure 8 · p. 20
B.1 Router Operation Space
- Global hidden states with a shared projection layer yield the optimal configuration in the operation-space ablation, which excludes the MoS router.Source: §B.1 · p. 20
B.2 Router Architecture Design
The router ablations favor modality-specific input normalization, ϵ-greedy exploration, and consolidating two hidden states.
B.2 Router Architecture Design
- ϵ-greedy exploration notably accelerates convergence across training steps.Source: Figure 10 · p. 21
- The ablation identifies k = 2 as the best-performing sparsity setting.Source: Figure 10 · p. 21
- Separate RMSNorm layers align modality representation scales and outperform shared RMSNorm across all evaluation metrics.Source: §B.2 · pp. 20–21, Figure 9 · p. 21
B.3 ϵ-greedy Strategy and Sparsity Design
The router ablation supports ϵ-greedy exploration and top-2 selection, while explaining the trade-offs at smaller and larger k.
B.3 ϵ-greedy Strategy and Sparsity Design
- ϵ-greedy exploration notably accelerates convergence across training steps.The tested setting uses ϵ = 0.05, randomly selecting a layer with 5% probability; this value was informed by grid search.Source: §B.3 · p. 21
- k = 2 consistently outperforms the other tested candidates for forming the final guidance feature.Source: §B.3 · p. 21
- At k = 1 the router may overfit to a single layer, while larger k can dilute representations and degrade performance.Source: §B.3 · p. 21
B.4 Scalability of MoS
The section examines scaling the understanding tower and how MoS handles image-editing context. It reports gains from a larger understanding tower and the highest average editing score when both towers receive the reference image.
B.4 Scalability of MoS
- MoS interaction consistently outperforms hand-crafted designs and also benefits from scaling up the understanding tower.Source: Figure 11 · p. 22
- Enlarging the understanding tower yields consistent, stable improvements, and its producer–consumer embeddings make scaling cost-efficient in the authors’ empirical analysis.The tower need not be updated across all training stages.Source: §B.4 · p. 22
B.5 MoS for Image Editing
- Providing source images to both the understanding and generation towers achieves the highest average GEdit-Bench score among the three tested input configurations.The authors hypothesize the towers contribute semantic information and low-level visual features, respectively.Source: §B.5 · p. 22, Figure 12 · p. 22
B.6 Ablation Study on Inference Strategy
The ablation tests common inference enhancements on MoS-S using GenEval. More inference steps improve generation quality, while scheduling and rescaling add slight gains.
B.6 Ablation Study on Inference Strategy
- The evaluation uses an intermediate MoS-S checkpoint trained for 400k steps at 512 × 512 resolution.Source: §B.6 · p. 22
- On GenEval, more inference steps consistently improve generation quality; CFG works in the typical 5.0–7.5 range, while linear–quadratic scheduling and rescaling add slight gains.Source: §B.6 · p. 23, Figure 13 · p. 23
B.7 MoS Router Visualization Analysis
The router visualizations show routing changes across denoising steps and tokens. They also report no evidence that final-layer features are effective or that MoT has a consistent layer-to-layer correspondence.
B.7 MoS Router Visualization Analysis
- Early denoising sparsely selects influential features from layers of varied depths; later, middle-layer weights rise and distributions smooth, as Fig. 14 visualizes layer contributions averaged across blocks and tokens.Source: §B.7 · p. 23, Figure 14 · p. 24
- At a fixed denoising step, individual tokens show distinct routing patterns, consistent with the ablation finding that token-wise prediction outperforms sample-wise prediction.Source: §B.7 · p. 23
- The analysis finds no evidence that final-layer embeddings are effective and no consistent pattern of strict layer-to-layer correspondence in MoT.Source: §B.7 · p. 23