Source-linked AI summary
Accelerating Unified Multimodal Models with Core-Expansion Routing and Unified Computation Scheduling
Wengyi Zhan, Chenqian Yan, Songwei Liu, Mingbao Lin, Rongrong Ji
TL;DR
Unified multimodal models waste computation across tokens, layers, and generation timesteps, while existing routing designs do not capture their structured cross-task redundancy. The paper introduces CE-Router with a shared core and generation-progress expansions, and UCS coordinates several acceleration operators using its routing signal. Across Show-o2 and BAGEL, the framework retains 98.03% of dense understanding performance with a 1.93× understanding speedup.
Problem
UMMs face redundant computation across visual tokens, model depth, and repeated denoising, while existing designs do not determine which importance structure is reusable across understanding and generation.
Method
CE-Router learns a task-shared core with progress-conditioned generation expansions, and UCS reuses its routing signal across token compaction, layer skipping, FFN pruning, cache reuse, and denoising-step early exit.
Results
1.93× understanding speedup is achieved while retaining 98.03% of dense performance on experiments with Show-o2 and BAGEL.
Takeaways & Limitations
The results support a core-expansion redundancy structure in which understanding and generation share a transferable importance component while generation adds progress-dependent corrections.
Takeaways & Limitations
The study focuses on transformer-based UMMs processing visual representations as token sequences, and practical speedups can vary across hardware and software stacks.
Abstract
from arXiv · showhide
Unified multimodal models jointly support understanding and generation, but incur substantial redundant computation across tokens, layers, and generation timesteps. Through token-importance probing, we identify an asymmetric core-expansion structure: understanding exhibits a stable importance component, while generation largely shares this component but requires progress-dependent corrections. We therefore propose CE-Router, which uses a task-shared core scorer and progress-conditioned generation expansions, optimized through generation decomposition and cross-task core alignment. At inference, CE-Router compacts token computation and supplies a learned routing signal to Unified Computation Scheduling, which coordinates layer skipping, FFN pruning, diffusion-head cache reuse, and denoising-step early exit. Experiments on two representative UMM architectures demonstrate consistent quality--efficiency improvements across both tasks, retaining 98.03\% of dense understanding performance with a 1.93$\times$ end-to-end inference speedup.
1 Introduction
Unified multimodal models share backbone computation across understanding and generation, creating redundant costs that vary by tokens, layers, and denoising progress. The paper identifies a shared importance core with generation-specific residuals, then uses CE-Router and UCS to accelerate both tasks.
- Motivation: UMMs incur redundant computation from long visual contexts, dense layer updates, and repeated generation processing across denoising steps.These costs affect attention, FFN computation, key–value-cache access, and trajectory latency.
- Motivation: Existing methods compress inputs, shorten trajectories, or route tasks separately, but do not derive a structured sharing boundary inside the coupled backbone.Fully shared routers miss generation-progress changes, whereas separate routers duplicate potentially reusable scoring structure.
- Core-expansion structure: Token-importance probes reveal a stable understanding component, while generation importance changes with denoising progress and transfers asymmetrically across tasks.The generation probe recovers much of the understanding target set, whereas understanding-to-generation coverage declines over time.
- Method: CE-Router combines a task-shared core scorer with generation-only, timestep-conditioned expansions that model progress-dependent residual variation.The core is input- and layer-dependent rather than a fixed token subset, preserving shared capacity while specializing generation.
- Method: UCS reuses CE-Router’s learned routing signal to coordinate layer skipping, FFN pruning, diffusion-head cache reuse, and denoising-step early exit.CE-Router compacts token computation by gathering retained tokens for attention and FFN operations.
- Results: 1.93× understanding speedup is achieved while retaining 98.03% of dense performance across experiments on Show-o2 and BAGEL.CE-Router provides the strongest task performance among accelerated baselines, while UCS further improves efficiency.
2 Related Work
Related work accelerates unified multimodal models through input compression, trajectory shortening, and within-backbone conditional computation. CE-Router instead learns a cross-task sharing boundary and reuses one task-aligned signal to coordinate computation.
- Boundary and trajectory methods: Input-level methods reduce visual tokens, while trajectory-level methods shorten generation through caching, distillation, or fewer sampling steps.These approaches act before or outside shared backbone computation.
- Within-backbone methods: FlashU combines multiple acceleration operators using training-free heuristics, whereas learned routing methods typically use independent understanding and generation routers.These designs either lack a common task-aligned signal or duplicate scoring structure across tasks.
- Paper's positioning: CE-Router learns the sharing boundary itself and reuses its routing signal to coordinate conditional computation across both tasks.This addresses the trade-off between compact full sharing and duplicated task-specific routing.
3 Methodology
Token-importance probing reveals stable understanding-side importance and progress-dependent generation variation. CE-Router captures this asymmetry with shared core scoring and generation expansions, then reuses routing decisions for token compaction and unified scheduling.
- 3.1 Probing Asymmetric Cross-Task Redundancy: Token probing finds persistent high-importance positions in understanding and increasing token replacement across generation denoising.Understanding-side attention bands persist across the backbone, while generation importance progressively departs from its initial selection.
- 3.1 Probing Asymmetric Cross-Task Redundancy: 84.45% of the understanding target set is recovered by a generation-trained probe, while generation coverage decreases from 83.67% to 78.40% during denoising.The transfer pattern indicates stronger generation-to-understanding reuse than reverse transfer at later generation stages.
- 3.2 Core-Expansion Scoring: CE-Router combines a task-shared core scorer with generation-only, timestep-selected expansion scorers.The core is input- and layer-dependent, while generation adds a segment-specific residual to model progress-dependent variation.
- 3.3 Core-Expansion Router Optimization: Two-stage optimization first decomposes generation importance across denoising segments, then aligns the shared core through joint understanding-generation training.The core is trained across segments and tasks, whereas each expansion specializes in its assigned generation segment.
- 3.4 Inference with Unified Computation Scheduling: At inference, retained-token routing compacts sequences and UCS reuses the decisions across layers, channels, diffusion heads, and denoising steps.Retained tokens are packed for attention and FFN computation, while dropped tokens follow a bypass path; UCS coordinates additional acceleration operators.
4 Experiments
Experiments evaluate CE-Router on two architecturally distinct UMM backbones across understanding and image-generation tasks. CE-Router achieves strong quality, while UCS further improves inference efficiency.
- Experimental setup: CE-Router is evaluated on the shared-Transformer Show-o2-1.5B and Mixture-of-Transformers BAGEL-7B-MoT backbones against FlashU, UniMoD, and dense inference.All methods use matched inference protocols.
- Experimental setup: UCS denotes Unified Computation Scheduling and is evaluated alongside performance and inference efficiency on multimodal understanding and image generation.The table marks best and second-best results with bold and underlining, respectively.
- Quantitative results: CE-Router preserves 99.36% and 99.96% of original understanding scores on Show-o2 and BAGEL, respectively, while attaining the highest GenEval and DPG-Bench scores.The comparison covers both architectures and both task types.
- Qualitative results: CE-Router produces cleaner object geometry, more consistent colors, richer details, and better spatial relations than FlashU and UniMoD in qualitative generation results.These advantages are largely retained after applying UCS.
Appendix D.
Ablations support the core-shared, generation-expanded router and the generation-first training strategy, while UCS improves efficiency with limited score degradation. The study also identifies a transformer-token-sequence scope boundary.
- Router architecture: Under the same parameter budget, understanding expansion improves Rel. Avg. by 0.16 points but trails the proposed design by 5.05 GenEval and 6.73 DPG-Bench points.The result supports allocating task-specific expansion primarily to generation.
- Generation segmentation: M = 4 delivers the best overall results with 185.15M parameters, whereas M = 8 saturates generation performance at 277.72M parameters.Increasing segmentation improves generation while keeping understanding stable, but increases router size.
- Training configuration: Generation-first training achieves the best overall balance, while joint training is suboptimal and UMM-first training reduces UMM performance to 95.61%.The comparison concerns alternative training configurations.
- Training configuration: λcon = 0.25 achieves the best overall performance, while removing the consistency loss causes substantial degradation on both tasks.Both smaller and larger consistency weights are less effective.
- Unified Computation Scheduling: UCS limits understanding and generation score drops to 1.33 and 1.01 points, compared with FlashU’s 18.74- and 3.28-point drops.Its trained unified routing signal enables larger computation and latency reductions than applying FlashU directly to CE-Router.
- Limitations: The study focuses on transformer-based unified multimodal models that process visual representations as token sequences.Different modality interfaces or generation paradigms may require architecture-specific adaptations.
A Experimental Details
The appendix provides implementation and measurement details rather than repeating the paper’s methodological derivations.
- Appendix scope: This section reports training configurations, baseline implementations, UCS operating points, and the efficiency measurement protocol.The main paper already describes the evaluated models, CE-Router, and the four UCS operators.
A.1 Evaluation and Efficiency Measurement Protocol
Evaluation covers two UMM backbones, multimodal understanding, and image generation under specified resolution and sampling settings. Efficiency is measured over complete end-to-end inference.
- Evaluation: The evaluation uses Show-o2-1.5B and BAGEL-7B-MoT for understanding and generation tasks.Table 4 summarizes the benchmarks, generation resolution, and denoising-step counts.
- Efficiency protocol: All generation training, evaluation, and efficiency measurements use 432 × 432 resolution, with TFLOPs and latency measured over the complete end-to-end inference path.Measurements are not restricted to a single layer or denoising step.
A.2 CE-Router Training Details
CE-Router training uses a two-stage procedure to learn a shared core scorer and generation-specific expansions, while UCS reuses routing signals across multiple acceleration operators. The appendix also specifies matched evaluation, implementation settings, and UniMoD reproduction details.
- CE-Router training: Two-stage optimization trains the task-shared core and generation expansions with task objectives and dense-execution consistency.Stage I trains generation across denoising segments; Stage II jointly trains understanding and generation with task-specific parameter updates.
- Training and evaluation settings: At r = 0.8, CE-Router and FlashU operate at approximately matched TFLOPs for performance comparison.Stage I fixes the keep ratio, while Stage II samples ratios from [0.6, 0.95] to learn a stable ranking across keep ratios.
- Unified Computation Scheduling: UCS maps four acceleration operators to CE-Router signals, using retained and dropped tokens to constrain potentially destructive operations.The operators are layer skipping, FFN pruning, diffusion-head cache reuse, and denoising-step early exit.
- Unified Computation Scheduling: UCS replaces fixed caching and step reduction with adaptive decisions controlled by routing stability, while enabling late-stage early exit for generation.The Show-o2 configuration enlarges the layer-skipping candidate budget; BAGEL generation additionally uses routing-guided cache reuse and late-stage early exit.
- Threshold selection: UCS thresholds are screened through one-time profiling and confirmed with a complete end-to-end inference evaluation.Profiling samples 20 examples per evaluation dataset and replays candidate thresholds offline; the reported procedure takes approximately 14.6 minutes.
- Baseline reproduction: UniMoD comparisons use an implementation based on the paper description because the official code has not been released.The reproduction covers Show-o2 and BAGEL architecture, training, and evaluation settings.
B.1 Received-Attention Analysis
Received attention is used as a proxy for understanding-side token importance, with within-layer normalization making persistent token-position patterns visible across the backbone. The analysis emphasizes relative importance rather than absolute attention mass across layers.
- Attention-based importance: Received attention averages head-level attention assigned to each token at each decoder layer.The analysis uses the dense, uncompressed model and aggregates attention over heads and 300 understanding examples.
- Observed persistence: Persistent vertical bands across most layers support a stable component in understanding-side token scoring.The paper cautions that received attention is only a proxy for routing importance, not a task-trained router score.
- Visualization: Log transformation, percentile clipping, and [0, 1] normalization make Figure 1(a) colors represent within-layer relative importance.The colors are not intended for direct comparison of absolute attention mass across layers.
- Probe analysis: Task-specific probes independently rank token logits within each layer and select top-k tokens for set-similarity analysis.The probes use independent parameters and do not implement CE-Router’s cross-task parameter sharing.
B.5 Cross-Task Probe Transfer
Cross-task transfer analysis compares directional top-k coverage between understanding and generation probes. Generation-probe transfer to understanding yields a single aggregate coverage of 84.45%, while understanding-probe transfer on generation varies by denoising step.
- Understanding-to-generation transfer: Understanding-probe coverage on generation samples varies with the generation timestep.The comparison is computed independently at each layer and evaluated timestep before averaging over samples and layers.
- Generation-to-understanding transfer: 84.45% is the aggregate top-k coverage when the generation probe transfers to understanding samples.Because understanding has no denoising trajectory, this transfer setting produces one result rather than timestep-specific values.
- Cross-task comparison: Figure 1(c) contrasts timestep-dependent understanding-probe transfer with the single aggregate generation-probe result.This directional comparison exposes asymmetry in how token-importance structure transfers across tasks.
C Detailed Experimental Results
Across two architecturally distinct UMM backbones, CE-Router preserves understanding quality and achieves the strongest accelerated generation results. UCS further reduces computation and latency while largely retaining these task-level and category-level advantages.
- Multimodal Understanding: 99.36% and 99.96% of original average understanding scores are retained on Show-o2 and BAGEL, respectively.After UCS, models further reduce TFLOPs and latency while maintaining the strongest overall understanding performance among accelerated baselines.
- Text-to-Image Generation: CE-Router achieves the highest overall GenEval and DPG-Bench scores among accelerated methods on both backbones.Category-level advantages are reported across object composition, attributes, spatial relations, and prompt-following.
- Text-to-Image Generation: UCS largely retains CE-Router’s category-level generation gains while continuing to outperform FlashU and UniMoD overall.The reported comparison states that runtime acceleration does not rely on disproportionate degradation in a single evaluation dimension.
D Additional Qualitative Results
Across both evaluated backbones and benchmarks, CE-Router delivers the most consistent visual quality among accelerated methods, preserving object details and spatial composition. Combined with UCS, it retains these qualitative advantages while improving end-to-end acceleration.
- CE-Router produces the most consistent visual quality across both backbones and benchmarks.It better preserves object identity, geometry, fine-grained details, and spatial composition.
- CE-Router with UCS largely retains CE-Router’s qualitative advantages while providing greater end-to-end acceleration.
- Figure 5 compares qualitative generations across prompts (a)–(p), including photographs, paintings, sketches, and composite scenes.