Mixture of States: Routing Token-Level Dynamics for Multimodal Generation (2511.12207v1)

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.

• The paper introduces a dynamic, token-level router that adapts cross-modal fusion in diffusion models for improved alignment and efficiency.
• It employs a separation of understanding and generation towers, achieving state-of-the-art results with minimal computational overhead.
• Empirical results show that the adaptive routing significantly enhances text-to-image synthesis and instruction-based image editing tasks.

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

Introduction and Motivation

"Mixture of States: Routing Token-Level Dynamics for Multimodal Generation" (2511.12207) addresses fundamental limitations in contemporary fusion paradigms for multimodal diffusion models, particularly in text-to-image generation and instruction-based image editing. Traditional approaches such as cross-attention and self-attention utilize static or rigid alignment between representations from different modalities, often leading to suboptimal conditioning given the dynamic, denoising-driven nature of diffusion models. Additionally, prior layer-wise fusion strategies—like Mixture-of-Transformers (MoT)—require strictly symmetric model architectures, impeding flexibility and scalability.

This paper introduces Mixture of States (MoS), a framework underpinned by a learnable, token-wise router that facilitates time- and input-adaptive aggregation of hidden states across modalities. This sparse and dynamic routing mechanism enables flexible fusion between asymmetric transformers, supporting optimal cross-modal alignment at every denoising step with negligible computational overhead.

Technical Approach: The MoS Architecture

Central to MoS is the separation of the model into an understanding tower ($\mathcal{U}$, typically a frozen transformer for text or text+image) and a generation tower ($\mathcal{G}$, a diffusion backbone). These modules are linked via a lightweight, token-wise router ($\mathcal{R}$), which dynamically determines (for each token at each generation block and timestep) which hidden states from the understanding tower to access and how to weight them.

Figure 2: MoS introduces a learnable router facilitating token-level, sparse, and dynamic connections between asymmetric transformer blocks for flexible multimodal fusion.

Unlike prior fixed-fusion approaches, the router's operation considers:

• The context embedding ($c$)
• The noised image latent at timestep $t$ ($z_t$)
• The current denoising timestep ($t$)

All are projected into a shared latent space and used to predict a routing logit matrix $\mathcal{W}\in \mathbb{R}^{m\times n}$, where $m$ and $n$ are the number of layers in $\mathcal{U}$ and $\mathcal{G}$, respectively. Routing is executed independently for each prompt token using a sparse top-$k$ selection strategy with $\epsilon$-greedy exploration during training, promoting both efficiency and avoidance of local optima.

The MoS paradigm generalizes cross-modality fusion: any visual token, at any denoising step and transformer block, can attend to any understanding tower layer for any prompt token, enabling rich, adaptive, and semantically grounded interactions.

Figure 3: Contrasts MoS routing scheme with cross-attention, self-attention, and MoT, highlighting the flexibility and granularity enabled by the learnable router.

Empirical Results and Ablations

Main Results

MoS achieves state-of-the-art (SoTA) performance across several challenging multimodal benchmarks, despite having only 3–5B learnable parameters—comparable or superior to approaches requiring up to 4× larger model sizes, such as Qwen-Image (20B parameters). Specifically, on GenEval, DPG, and WISE benchmarks, MoS-L (5B) closely matches or surpasses the strongest known open models.

Figure 1: MoS exhibits superior sample quality and alignment compared to MoT baseline across all training steps as measured by GenEval and DPG.

Figure 6: MoS’s router adds negligible latency, allowing for high-throughput image generation compared to existing models on identical hardware.

Key Ablations

Router Conditioning: Incorporating timestep and image latent information into the router's input—beyond static prompt embeddings—correlates with consistent improvements in FID and CLIP-based metrics, confirming the importance of dynamic, step-aware fusion.

Token-specific Routing: Predicting routing matrices for individual prompt tokens—rather than sample-wise global routing matrices—yields measurable gains, reinforcing the necessity of fine-grained, token-level conditioning.

Figure 4: Comparing router output strategies, token-specific prediction yields demonstrably better FID and CLIP scores.

Adaptive Layer Selection: MoS’s adaptive, learned state selection surpasses hand-crafted, fixed-layer routing as well as cross-attention and MoT baselines, even under equal compute and parameter budgets. Empirical results provide strong evidence against the assumed necessity of symmetric, layer-aligned fusion.

Router Efficiency and Sparsity: The lightweight transformer-based router contributes minimal compute overhead (on the order of milliseconds per $1024^2$ image step) and the sparse top-$k$ selection with $\epsilon$-greedy training further accelerates convergence and stabilizes training.

Visual and Qualitative Comparisons

On both text-to-image synthesis and instruction-based image editing tasks (including complex, multi-entity, and textual scenes), MoS-L outputs highly photorealistic and semantically aligned outputs with controllable style and content.

Figure 5: Examples of high-fidelity, photorealistic MoS-Image generations (left) and precise, instruction-aligned MoS-Edit results (right).

Figure 8: Visual comparison on challenging text-to-image prompts. MoS outperforms or matches larger competitor models on precise object arrangement and text rendering.

Figure 10: Instruction-based image editing—MoS-L/S produces precise, high-quality edits for hybrid and textual instruction cases, outperforming baselines.

Implications and Future Research

Theoretical and Practical Impact

MoS demonstrates that token-level, sparse, dynamic routing between asymmetric transformers is both tractable and highly effective, obviating the need for tedious hand-crafted fusion or symmetric architecture alignment in multimodal diffusion settings. This advances both practical training—reducing hardware and data barriers—and theoretical understanding, as the strict final-layer or layer-to-layer correspondence implicit in prior work is empirically refuted.

Architecturally, the use of a frozen, scalable understanding tower with a lightweight, dynamic-generation tower enables efficient scaling and inference. The adaptive router design is extensible to more general forms of conditional generative modeling, beyond text and images.

Limitations and Prospects

While MoS excels in the as-tested setting, further exploration is warranted in:

• Dual-way fusion: Extending one-way adaptive routing to permit bidirectional inter-tower interaction.
• Preference alignment: Leveraging reinforcement learning and advanced instruction finetuning for human-aligned outputs.
• Further efficiency: Exploring quantization, distillation, and caching optimizations within the MoS design.
• Explainability: Using router output patterns as interpretable explanations for cross-modal flow.
• Generalization: Application to other modalities (e.g., video, audio) and more complex task regimes.

Conclusion

Mixture of States (MoS) establishes a flexible and efficient foundation for scaling multimodal diffusion models. Through token-adaptive and time-aware routing between asymmetric transformers, it achieves state-of-the-art performance with substantially reduced compute and parameter budgets, and demonstrates robust capabilities for both generation and editing tasks. Empirical evidence from both quantitative benchmarks and qualitative visualizations confirms the superiority of dynamic, state-routing fusion over prior hand-crafted approaches. MoS is a promising platform for developing the next generation of scalable, high-fidelity, and interpretable multimodal AI systems.