FlexUMM：统一多模态模型的灵活表示学习

# 1 Introduction
Source: [https://arxiv.org/html/2610.22682](https://arxiv.org/html/2610.22682)
![[Uncaptioned image]](https://arxiv.org/html/2610.22682v1/cmu-logo.png)![[Uncaptioned image]](https://arxiv.org/html/2610.22682v1/logos/wm-logo.png)![[Uncaptioned image]](https://arxiv.org/html/2610.22682v1/logos/umd-logo.png)

FlexUMM: Flexible Representation Learning for Unified Multimodal Models

Yinyi Luo1111Contact: yinyil@andrew\.cmu\.edu, Haoyue Bai2, Yi Ge3, Nilton Cesar Rojas Vales2, Marios Savvides1, Jindong Wang2222Corresponding authors: jdw@wm\.edu

1Carnegie Mellon University2William & Mary3University of Maryland

###### Abstract

Unified multimodal models \(UMMs\) can integrate understanding and generation within a single backbone\. However, it is commonly known that these two capabilities can cause conflicting representations, leading to a central question: how representation and computation should be shared across functionalities? Existing UMMs often rely on a largely fixed structure, where different inputs and functionalities follow the same depth and parameter\-sharing pattern\. We argue that this assumption is too restrictive: more effective unification requires sharing the computation that supports cross\-functional communication while specializing the capacity that becomes task\- or input\-specific\. We introduce FlexUMM, a flexible UMM framework that dynamically decides when to preserve shared computation, when to route specialized capacity, and when to stop further processing\. Evidence shows that this flexible allocation improves UMMs along two complementary axes\. It improves accuracy by reducing harmful competition between functionalities, and it improves efficiency by avoiding unnecessary computation once an input representation is sufficiently mature\. Beyond simple architectural unification from the pre\-trained models, FlexUMM works as a model\-agnostic,functionalplug\-in to make existing UMMs more unified, stronger, and efficient\. Code is available at:[https://github\.com/AIFrontierLab/TorchUMM/tree/main/src/umm/post\_training/flexumm](https://github.com/AIFrontierLab/TorchUMM/tree/main/src/umm/post_training/flexumm)\.

Unified multimodal models \(UMMs\) are built around an appealing idea: understanding and generation can use one backbone instead of being served by two separate models\([Wang et al\., 2022](https://arxiv.org/html/2610.22682#bib.bib16);[Lu et al\., 2022a](https://arxiv.org/html/2610.22682#bib.bib17);[Sun et al\., 2024](https://arxiv.org/html/2610.22682#bib.bib30);[Deng et al\., 2025](https://arxiv.org/html/2610.22682#bib.bib31)\)\. Using the same shared architecture, one UMM can do both generation and understanding tasks with the same network\. This unification, however, creates a fundamental challenge in knowledge representation and computation allocation\([Wu et al\., 2026c](https://arxiv.org/html/2610.22682#bib.bib61);[Kil and Kim, 2026](https://arxiv.org/html/2610.22682#bib.bib57)\)\. A shared backbone must simultaneously support different capabilities, yet not every input requires the same computation, and not every component of the representation should be shared across capabilities\([Han et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib55);[Kang et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib56)\)\.

\(a\)Shared communication versus routed FFN capacity\.\(b\)Cross\-functional gradient cosine\.
Figure 1:FlexUMM adapts both who computes and how much computation an input receives\.Existing UMMs largely rely on a fixed computational structure \([Figure1\(a\)](https://arxiv.org/html/2610.22682#S1.F1.sf1)\), leaving both the allocation of capacity across capabilities and depth across inputs predetermined\([Kil and Kim, 2026](https://arxiv.org/html/2610.22682#bib.bib57);[Qiu et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib59)\)\. Such a rigid allocation can lead to two forms of mismatch: computation may be spent on inputs or capabilities that do not require it, while shared components may be overused even when their computation is not equally beneficial to different capabilities\([He et al\., 2025](https://arxiv.org/html/2610.22682#bib.bib18);[Wu et al\., 2026c](https://arxiv.org/html/2610.22682#bib.bib61);[Ye et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib19)\)\. This limitation raises a central question:how should a UMM adapt its representation and computation to both capabilities more adaptively?Addressing this question requires resolving two closely related challenges\.

First, computation depth should adapt to the input and capability\.Not every input requires the full depth of a deep network\([Graves, 2016](https://arxiv.org/html/2610.22682#bib.bib28);[Elbayad et al\., 2019](https://arxiv.org/html/2610.22682#bib.bib29)\)\. In a UMM, however, the required depth depends not only on the input but also on the capability being executed\([Kang et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib56);[Su et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib60)\)\. The depth at which a representation becomes sufficient for understanding should not coincide with the depth required for image generation, even though both capabilities operate on the same backbone\. A fixed execution depth therefore imposes the same computational budget on inputs with different requirements\. It is evident even within a single functionality: using a learned policy on Bagel, all MME\([Fu et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib53)\)inputs are distributed across six distinct depths \([Figure2\(a\)](https://arxiv.org/html/2610.22682#S1.F2.sf1)\)\. Moreover, the quality–efficiency behavior in[Figure2\(b\)](https://arxiv.org/html/2610.22682#S1.F2.sf2)suggests that computation beyond the sufficient depth can introduce unnecessary interference and reduce accuracy, rather than simply incurring redundant FLOPs\. These observations motivate input\-adaptive depth allocation that determines whether further computation remains useful\.

\(a\)Selected depths\.
\(b\)Quality–efficiency behavior\.
Figure 2:Input\-adaptive depth and quality–efficiency behavior\.Second, computational capacity must adapt to the capability\.While depth determines how far computation proceeds, it does not determine which components should be shared or specialized\. Understanding and generation rely on a common backbone, but their objectives can impose conflicting demands on the same parameters\([Han et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib55);[Wu et al\., 2026c](https://arxiv.org/html/2610.22682#bib.bib61)\)\. Separating the capabilities entirely would reduce such interference, but would also discard the cross\-capability communication that motivates unified modeling\([Kang et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib56);[Su et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib60)\)\. The challenge is therefore to preserve shared computation where the capabilities benefit from interaction while allowing specialization where their objectives diverge\([Peng et al\., 2026a](https://arxiv.org/html/2610.22682#bib.bib58);[Qiu et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib59)\)\. Our controlled diagnostic, under the task construction described in the appendix, suggests a structured sharing pattern: attention gradients remain aligned across functionalities throughout the network, whereas FFN gradients increasingly diverge with depth \([Figure1\(b\)](https://arxiv.org/html/2610.22682#S1.F1.sf2), Appendix[B\.1](https://arxiv.org/html/2610.22682#A2.SS1)\)\. This is a design hypothesis rather than a direct measurement of every released backbone, and it suggests that communication and transformation capacity should not be treated uniformly\. Communication benefits from sharing while transformation capacity requires adaptive specialization\. These interactions also have implications beyond accuracy and efficiency\. If shared computation is adapted without accounting for capability\-specific demands, improvements in one functionality may come with safety regressions in another\.

We introduceFlexUMM, a flexible computation allocation method that adapts both capacity and depth on a shared UMM backbone\. FlexUMM uses capability\-adaptive allocation to preserve shared attention and selectively specialize FFN capacity through per\-token routing, and depth\-adaptive allocation to determine whether further computation is needed through per\-sequence early exit\. Together, these mechanisms form a progressiveshare, specialize, and stoppolicy that dynamically allocates computation across capabilities and inputs \([Figure1\(a\)](https://arxiv.org/html/2610.22682#S1.F1.sf1)\)\. FlexUMM is implemented as a lightweight adaptation on top of the pretrained checkpoint, adding only routers, halting heads and low\-rank residual banks, with only 0\.381% increase in model size\.

Our contributions are summarized as follows:

- •We identify two complementary sources of conflicts in UMMs: input\-dependent depth and capability\-dependent capacity\. Our analysis shows that attention provides a shared communication pathway, while FFN capacity exhibits increasing functional divergence across layers\.
- •We propose FlexUMM, which jointly enables token\-level capacity allocation and sequence\-level depth adaptation within a unified backbone, allowing computation to be shared, specialized, or terminated according to the input and capability\.
- •FlexUMM improves both multimodal understanding and generation while providing an adjustable accuracy\-compute trade\-off at inference time, with negligible extra computation\.

## 2 Related Work

Unified multimodal models\.UMMs combine multimodal understanding and generation within a single architecture, with different designs for shared representations, tokenizers, and computation paths\([Xie et al\., 2025](https://arxiv.org/html/2610.22682#bib.bib20);[Chen et al\., 2025](https://arxiv.org/html/2610.22682#bib.bib21);[Deng et al\., 2025](https://arxiv.org/html/2610.22682#bib.bib31);[Peng et al\., 2026a](https://arxiv.org/html/2610.22682#bib.bib58)\)\. Recent studies further examine how architectural choices affect the interaction between functionalities\([Kang et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib56);[Qiu et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib59);[Wu et al\., 2026c](https://arxiv.org/html/2610.22682#bib.bib61);[Luo et al\., 2026b](https://arxiv.org/html/2610.22682#bib.bib2);[Luo et al\., 2026a](https://arxiv.org/html/2610.22682#bib.bib3);[Bai et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib4)\)\. These models validate unification, but their computation paths remain largely fixed, with sharing patterns and depth predetermined before inference\.

Adaptive sharing and specialization\.Multi\-task models mitigate task interference by mixing task\-specific representations, routing through shared and private experts, or learning layerwise sharing policies\([Misra et al\., 2016](https://arxiv.org/html/2610.22682#bib.bib23);[Ma et al\., 2018](https://arxiv.org/html/2610.22682#bib.bib24);[Tang et al\., 2020](https://arxiv.org/html/2610.22682#bib.bib32);[Sun et al\., 2020](https://arxiv.org/html/2610.22682#bib.bib25);[Rahimian et al\., 2023](https://arxiv.org/html/2610.22682#bib.bib26)\)\. UMMs have also introduced function\-specific computation to reduce conflicts: Janus decouples visual encoding for understanding and generation, while more recent work specializes conflicting visual computation while preserving semantic interaction\([Wu et al\., 2025a](https://arxiv.org/html/2610.22682#bib.bib22);[Wu et al\., 2026c](https://arxiv.org/html/2610.22682#bib.bib61)\)\. These approaches mainly target conventional multi\-task learning or rely on predefined specialization, rather than dynamically allocating shared and specialized computation\.

Adaptive computation\.Adaptive computation methods reduce visual processing through token selection, pruning, or merging\([Ryoo et al\., 2021](https://arxiv.org/html/2610.22682#bib.bib33);[Rao et al\., 2021](https://arxiv.org/html/2610.22682#bib.bib27);[Bolya et al\., 2022](https://arxiv.org/html/2610.22682#bib.bib34)\)\. PruMerge, FastV, SparseVLM, DyVTE, and HiDrop\([Shang et al\., 2025](https://arxiv.org/html/2610.22682#bib.bib35);[Chen et al\., 2024a](https://arxiv.org/html/2610.22682#bib.bib36);[Zhang et al\., 2024](https://arxiv.org/html/2610.22682#bib.bib37);[Wu et al\., 2026d](https://arxiv.org/html/2610.22682#bib.bib6);[Wu et al\., 2026b](https://arxiv.org/html/2610.22682#bib.bib7)\)exploit visual\-token redundancy or vary token participation across layers, while MuE, Mixture\-of\-Depths, and MoR adapt executed depth or recursive computation\([Tang et al\., 2023](https://arxiv.org/html/2610.22682#bib.bib5);[Raposo et al\., 2024](https://arxiv.org/html/2610.22682#bib.bib39);[Bae et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib8)\)\. These methods adapt token participation or executed depth; FlexUMM additionally adapts whether computation is shared or specialized\.

## 3 Method

Figure 3:FlexUMM separates three responsibilities in the implementation\.The shared decoder stack supplies attention and normalization for every route\. A tri\-state capacity gate selectsShare,Decouple, orExitfor each token at each routed FFN\. A separate candidate\-depth router decides, at the sequence level, whether the remaining decoder layers are still needed\.As illustrated in[Figure3](https://arxiv.org/html/2610.22682#S3.F3), FlexUMM introduces an input\-adaptive computation method for UMMs\. The model preserves a shared backbone for multimodal communication while adapting both the amount of capacity allocated to individual tokens and the computation depth required by each input\. It consists of three stages: 1\) Shared Fusion: All tokens undergo the same attention, normalization, and FFN \(feed\-forward network\) to establish a common multimodal representation\. 2\) Adaptive Specialization: The shared backbone remains intact, while a Tri\-state Gate determines whether each token uses the shared FFN alone, receives an additional specialized residual transformation, or bypasses the FFN\. 3\) Adaptive Depth: A sequence\-level Halting Head determines whether the current representation is sufficient, allowing different inputs to terminate computation at different depths\. Thus, FlexUMM adapts computation along two complementary dimensions: capacity at the token level and depth at the sequence level, while retaining a shared multimodal communication pathway\.

### 3\.1Shared Communication and Adaptive Capacity

Let𝒉\(ℓ−1\)∈ℝn×d\{\\bm\{h\}\}^\{\(\\ell\-1\)\}\\in\\mathbb\{R\}^\{n\\times d\}denote the hidden states entering layerℓ\\ell, wherennis the sequence length andddis the hidden dimension\. Each layer follows the standard Transformer backbone with multi\-head attention and normalization shared across modalities:

𝒉~\(ℓ\)=𝒉\(ℓ−1\)\+Attn\(ℓ\)​\(Norm1\(ℓ\)​\(𝒉\(ℓ−1\)\)\),𝒖\(ℓ\)=Norm2\(ℓ\)​\(𝒉~\(ℓ\)\),\\tilde\{\{\\bm\{h\}\}\}^\{\(\\ell\)\}=\{\\bm\{h\}\}^\{\(\\ell\-1\)\}\+\\mathrm\{Attn\}^\{\(\\ell\)\}\\left\(\\mathrm\{Norm\}\_\{1\}^\{\(\\ell\)\}\\left\(\{\\bm\{h\}\}^\{\(\\ell\-1\)\}\\right\)\\right\),\\qquad\{\\bm\{u\}\}^\{\(\\ell\)\}=\\mathrm\{Norm\}\_\{2\}^\{\(\\ell\)\}\\left\(\\tilde\{\{\\bm\{h\}\}\}^\{\(\\ell\)\}\\right\),\(1\)where𝒖k\(ℓ\)∈ℝd\{\\bm\{u\}\}\_\{k\}^\{\(\\ell\)\}\\in\\mathbb\{R\}^\{d\}and𝒉~k\(ℓ\)∈ℝd\\tilde\{\{\\bm\{h\}\}\}\_\{k\}^\{\(\\ell\)\}\\in\\mathbb\{R\}^\{d\}represent the normalized and residual state for tokenkk, respectively\.

Adaptive capacity allocation\.At a given layerℓ\\ell, a Tri\-state Gate dynamically selects an actionak\(ℓ\)∈𝒜=\{SHARED,SPECIALIZED,BYPASS\}a\_\{k\}^\{\(\\ell\)\}\\in\\mathcal\{A\}=\\\{\\mathrm\{SHARED\},\\mathrm\{SPECIALIZED\},\\mathrm\{BYPASS\}\\\}to allocate FFN capacity for tokenkk\. Specifically,SHARED\\mathrm\{SHARED\}routes the token through the shared base FFN,SPECIALIZED\\mathrm\{SPECIALIZED\}augments the shared FFN with a learned decoupled residual subset, andBYPASS\\mathrm\{BYPASS\}skips the current FFN update\. The gate predicts the action based on the token representation and routing context, enabling input\-adaptive capacity allocation without requiring manually annotated token attributes\.

Letsk,a\(ℓ\)s\_\{k,a\}^\{\(\\ell\)\}denote the gate logit for actionaa\. The action probability and inference decision are:

πk,a\(ℓ\)=exp⁡\(sk,a\(ℓ\)/τ\)∑a′∈𝒜exp⁡\(sk,a′\(ℓ\)/τ\),ak\(ℓ\)=arg⁡maxa∈𝒜⁡sk,a\(ℓ\)\.\\pi\_\{k,a\}^\{\(\\ell\)\}=\\frac\{\\exp\\left\(s\_\{k,a\}^\{\(\\ell\)\}/\\tau\\right\)\}\{\\sum\_\{a^\{\\prime\}\\in\\mathcal\{A\}\}\\exp\\left\(s\_\{k,a^\{\\prime\}\}^\{\(\\ell\)\}/\\tau\\right\)\},\\qquad a\_\{k\}^\{\(\\ell\)\}=\\arg\\max\_\{a\\in\\mathcal\{A\}\}s\_\{k,a\}^\{\(\\ell\)\}\.\(2\)
For theSPECIALIZED\\mathrm\{SPECIALIZED\}action, a learned routing head selects a residual\-bank indexzk\(ℓ\)z\_\{k\}^\{\(\\ell\)\}from the token representation and routing context\. The index identifies a learned residual subset rather than a manually assigned token label\. The overall output state is𝒉k=𝒉~k\+Δ​𝒉k\{\\bm\{h\}\}\_\{k\}=\\tilde\{\{\\bm\{h\}\}\}\_\{k\}\+\\Delta\{\\bm\{h\}\}\_\{k\}, where the token\-level FFN updateΔ​𝒉k\\Delta\{\\bm\{h\}\}\_\{k\}is defined as

Δ​𝒉k=\{F⁡\(𝒖k\),ak\(ℓ\)=SHARED,F⁡\(𝒖k\)\+Rzk\(ℓ\)​\(𝒖k\),ak\(ℓ\)=SPECIALIZED,𝟎,ak\(ℓ\)=BYPASS\.\\Delta\{\\bm\{h\}\}\_\{k\}=\\begin\{cases\}F\\left\(\{\\bm\{u\}\}\_\{k\}\\right\),&a\_\{k\}^\{\(\\ell\)\}=\\mathrm\{SHARED\},\\\\\[2\.0pt\] F\\left\(\{\\bm\{u\}\}\_\{k\}\\right\)\+R\_\{z\_\{k\}^\{\(\\ell\)\}\}\\left\(\{\\bm\{u\}\}\_\{k\}\\right\),&a\_\{k\}^\{\(\\ell\)\}=\\mathrm\{SPECIALIZED\},\\\\\[2\.0pt\] \\mathbf\{0\},&a\_\{k\}^\{\(\\ell\)\}=\\mathrm\{BYPASS\}\.\\end\{cases\}\(3\)Here,FFdenotes the shared base FFN, whileRzk\(ℓ\)R\_\{z\_\{k\}^\{\(\\ell\)\}\}denotes the selected learned residual subset\. The residual subsets are optimized jointly with the routing mechanism during post\-training, providing additional input\-adaptive capacity while preserving the shared FFN as the common foundation\.

### 3\.2Adaptive Computation Depth

While the capacity gate dynamically allocates computation within each layer, the effective computation depth may vary across inputs\. We therefore introduce a depth router that selects from a predefined set of candidate exit depths𝒟=\{d1,…,dM\}\\mathcal\{D\}=\\\{d\_\{1\},\\ldots,d\_\{M\}\\\}, whered1<⋯<dMd\_\{1\}<\\cdots<d\_\{M\}\. Here,mmcorresponds to a configured candidate depth rather than a continuously learned computation depth\.

At candidate depthdmd\_\{m\}, the depth router receives the current shared representation and produces a halting logitrmr\_\{m\}\. The corresponding halting probability is defined asphalt,m=σ⁡\(rmτdepth\),p\_\{\\mathrm\{halt\},m\}=\\sigma\\left\(\\frac\{r\_\{m\}\}\{\\tau\_\{\\mathrm\{depth\}\}\}\\right\),whereτdepth\\tau\_\{\\mathrm\{depth\}\}denotes the temperature of the depth router andσ⁡\(⋅\)\\sigma\(\\cdot\)is the sigmoid function\. We distinguishτdepth\\tau\_\{\\mathrm\{depth\}\}from the temperature used by the capacity gate, as the two temperatures govern different routing decisions\. In our implementation,τdepth=1\.0\\tau\_\{\\mathrm\{depth\}\}=1\.0\.

Sequential Exit Distribution\.To jointly optimize the candidate exits, we define a sequential exit distribution based on the halting probabilities:

qm=\{phalt,m​∏r=1m−1\(1−phalt,r\),m<M,∏r=1M−1\(1−phalt,r\),m=M\.q\_\{m\}=\\begin\{cases\}p\_\{\\mathrm\{halt\},m\}\\displaystyle\\prod\_\{r=1\}^\{m\-1\}\(1\-p\_\{\\mathrm\{halt\},r\}\),&m<M,\\\\\[6\.0pt\] \\displaystyle\\prod\_\{r=1\}^\{M\-1\}\(1\-p\_\{\\mathrm\{halt\},r\}\),&m=M\.\\end\{cases\}\(4\)The final candidate serves as a mandatory fallback, ensuring that∑m=1Mqm=1\\sum\_\{m=1\}^\{M\}q\_\{m\}=1\. The resulting expected computation depth is𝔼⁡\[𝒟\]=∑m=1Mqm​dm\.\\mathbb\{E\}\[\\mathcal\{D\}\]=\\sum\_\{m=1\}^\{M\}q\_\{m\}d\_\{m\}\.During training, the full\-depth candidate provides the reference computation for assessing the quality of earlier exits\. Each candidate is evaluated according to its task objective, and candidates satisfying a predefined quality constraint relative to the full\-depth reference are considered eligible for early termination\. The depth router is optimized to favor earlier eligible candidates while retaining the full\-depth candidate as a fallback\. This training procedure encourages input\-adaptive computation allocation without requiring the model to learn an unrestricted continuous exit depth\.

Inference\-Time Exit Policy\.At inference, the model evaluates the configured candidate depths in ascending order and terminates at the first candidate whose halting probability exceeds a threshold:

m⋆=min⁡\{m≥mmin\|phalt,m≥ηβ\},m^\{\\star\}=\\min\\left\\\{m\\geq m\_\{\\min\}\\;\\middle\|\\;p\_\{\\mathrm\{halt\},m\}\\geq\\eta\_\{\\beta\}\\right\\\},\(5\)wheremminm\_\{\\min\}specifies the earliest eligible exit andηβ\\eta\_\{\\beta\}is an inference\-time threshold controlling the computation\-quality trade\-off\. If no earlier candidate satisfies the criterion, the model falls back to the final candidatem=Mm=M\. In our experiments, we set the inference threshold toηβ=0\.50\\eta\_\{\\beta\}=0\.50\.

### 3\.3Training and Inference

FlexUMM jointly optimizes task performance and computational efficiency through token\-level capacity allocation and sequence\-level depth adaptation\. The overall objective is defined as

ℒ=ℒtask\+λcap​ℒcap\+λdepth​ℒdepth,\\mathcal\{L\}=\\mathcal\{L\}\_\{\\mathrm\{task\}\}\+\\lambda\_\{\\mathrm\{cap\}\}\\mathcal\{L\}\_\{\\mathrm\{cap\}\}\+\\lambda\_\{\\mathrm\{depth\}\}\\mathcal\{L\}\_\{\\mathrm\{depth\}\},\(6\)whereλcap\\lambda\_\{\\mathrm\{cap\}\}andλdepth\\lambda\_\{\\mathrm\{depth\}\}are hyperparameters\.

Task Objective\.To preserve the native task interfaces of the unified model, we reuse the original output heads at each candidate depth rather than introducing task\-specific prediction heads\. Specifically, the language\-model head is used during the autoregressive decoding process for understanding, whereas the flow\-matching head is used during the generation decoding process\. The task loss at candidate depthdmd\_\{m\}is defined as

ℓmtask=\{ℒCE​\(LMHead⁡\(𝒈m\),y\),understanding,‖FlowHead⁡\(𝒈m\)−𝐯‖22,generation,\\ell\_\{m\}^\{\\mathrm\{task\}\}=\\begin\{cases\}\\mathcal\{L\}\_\{\\mathrm\{CE\}\}\\left\(\\mathrm\{LMHead\}\(\{\\bm\{g\}\}\_\{m\}\),y\\right\),&\\text\{understanding\},\\\\\[4\.0pt\] \\left\\\|\\mathrm\{FlowHead\}\(\{\\bm\{g\}\}\_\{m\}\)\-\{\\mathbf\{v\}\}\\right\\\|\_\{2\}^\{2\},&\\text\{generation\},\\end\{cases\}\(7\)where𝐯\{\\mathbf\{v\}\}denotes the target flow\-matching velocity andℒCE\\mathcal\{L\}\_\{\\mathrm\{CE\}\}denotes the cross\-entropy loss for the understanding task\. The expected task loss across candidate exits is then formulated as

ℒtask=∑m=1Mqm​ℓmtask,\\mathcal\{L\}\_\{\\mathrm\{task\}\}=\\sum\_\{m=1\}^\{M\}q\_\{m\}\\ell\_\{m\}^\{\\mathrm\{task\}\},\(8\)whereqmq\_\{m\}is the sequential exit probability defined in Eq\. \([4](https://arxiv.org/html/2610.22682#S3.E4)\)\. This formulation allows the depth router to account for task quality across candidate computation depths while retaining the native prediction interfaces of the underlying model\.

Capacity Regularization\.The capacity regularizer encourages efficient allocation of token\-level computation across the adaptive layers\. Let𝒜=\{SHARED,SPECIALIZED,BYPASS\}\\mathcal\{A\}=\\\{\\mathrm\{SHARED\},\\mathrm\{SPECIALIZED\},\\mathrm\{BYPASS\}\\\}denote the available actions, and letcac\_\{a\}represent the normalized relative computational cost associated with actionaa\. We define

ℒcap=1L​n​∑ℓ=1L∑k=1n∑a∈𝒜ca​πk,a\(ℓ\),\\mathcal\{L\}\_\{\\mathrm\{cap\}\}=\\frac\{1\}\{Ln\}\\sum\_\{\\ell=1\}^\{L\}\\sum\_\{k=1\}^\{n\}\\sum\_\{a\\in\\mathcal\{A\}\}c\_\{a\}\\,\\pi\_\{k,a\}^\{\(\\ell\)\},\(9\)whereπk,a\(ℓ\)\\pi\_\{k,a\}^\{\(\\ell\)\}denotes the action probability at layerℓ\\ellfor tokenkk, as defined in Eq\. \([2](https://arxiv.org/html/2610.22682#S3.E2)\)\. The relative costs are determined according to the computation associated with each action, withcBYPASS≤cSHARED≤cSPECIALIZEDc\_\{\\mathrm\{BYPASS\}\}\\leq c\_\{\\mathrm\{SHARED\}\}\\leq c\_\{\\mathrm\{SPECIALIZED\}\}\. The cost values are normalized with respect to a common reference to ensure that the regularizer reflects the relative computational burden of the available routes\.

##### Depth Regularization\.

To promote input\-adaptive early termination, we penalize the expected computation depth normalized by the maximum candidate depth:

ℒdepth=𝔼⁡\[𝒟\]dM=1dM​∑m=1Mqm​dm\.\\mathcal\{L\}\_\{\\mathrm\{depth\}\}=\\frac\{\\mathbb\{E\}\[\\mathcal\{D\}\]\}\{d\_\{M\}\}=\\frac\{1\}\{d\_\{M\}\}\\sum\_\{m=1\}^\{M\}q\_\{m\}d\_\{m\}\.\(10\)This regularization encourages the router to favor earlier candidate exits when they provide sufficient task quality, while retaining the full\-depth candidate as a fallback\.

##### Inference and Computational Efficiency\.

During inference, FlexUMM evaluates candidate depths in ascending order and terminates execution at the first eligible candidate according to the halting policy in Eq\. \([5](https://arxiv.org/html/2610.22682#S3.E5)\)\. Unlike uniformly deep execution, this adaptive policy allows samples to exit at earlier depths, reducing unnecessary decoder computation and improving inference efficiency\.

By jointly adapting token\-level capacity and sequence\-level depth, FlexUMM reduces redundant computation while preserving the shared multimodal processing pathway\. Token\-level capacity allocation adjusts the computational route within each layer, while sequence\-level early exiting reduces the number of executed layers, enabling more efficient inference\.

## 4 Experiments

We perform extensive experiments on effectiveness, efficiency, and explanations of FlexUMM\.

### 4\.1Experimental Setup

Backbone models\.We evaluate FlexUMM on three UMM backbones with different sharing structures: Bagel\([Deng et al\., 2025](https://arxiv.org/html/2610.22682#bib.bib31)\), Janus\-Pro\([Chen et al\., 2025](https://arxiv.org/html/2610.22682#bib.bib21)\), and Harmon\([Wu et al\., 2025b](https://arxiv.org/html/2610.22682#bib.bib40)\), with Bagel serving as the primary backbone\.

Post\-training data\.FlexUMM is trained applied to each released checkpoint using ScienceQA\([Lu et al\., 2022b](https://arxiv.org/html/2610.22682#bib.bib41)\), AI2D\([Kembhavi et al\., 2016](https://arxiv.org/html/2610.22682#bib.bib42)\), TextVQA\([Singh et al\., 2019](https://arxiv.org/html/2610.22682#bib.bib43)\), MathV\-360K\([Shi et al\., 2024](https://arxiv.org/html/2610.22682#bib.bib44)\), and InstructPix2Pix\([Brooks et al\., 2023](https://arxiv.org/html/2610.22682#bib.bib45)\)\.

Evaluation benchmarks and metrics\.For understanding, we use MMMU\([Yue et al\., 2024](https://arxiv.org/html/2610.22682#bib.bib46)\), MathVista\([Lu et al\., 2024](https://arxiv.org/html/2610.22682#bib.bib47)\), MMStar\([Chen et al\., 2024b](https://arxiv.org/html/2610.22682#bib.bib48)\), and MMBench\([Liu et al\., 2024](https://arxiv.org/html/2610.22682#bib.bib49)\)\. For generation, GenEval\([Ghosh et al\., 2023](https://arxiv.org/html/2610.22682#bib.bib50)\)and WISE\([Niu et al\., 2025](https://arxiv.org/html/2610.22682#bib.bib51)\)\. For unified evaluation, Unison\([Liu et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib52)\)and Unified\-Bench\([Yan et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib54)\)\.

Implementation details\.Only the routers, halting heads, and low\-rank residual banks are trained; the backbone configurations and trainable parameter counts are summarized in Appendix[Table7](https://arxiv.org/html/2610.22682#A1.T7)\. All experiments are conducted using TorchUMM\([Luo et al\., 2026c](https://arxiv.org/html/2610.22682#bib.bib1)\)for fair comparison\. More details in Appendix[A](https://arxiv.org/html/2610.22682#A1)\.

### 4\.2Main Results

FlexUMM improves understanding accuracy over existing UMMs\.[Table1](https://arxiv.org/html/2610.22682#S4.T1)compares FlexUMM with existing unified models across understanding and generation benchmarks\. FlexUMM achieves the best results on all five understanding benchmarks and consistently improves both generation metrics over its Bagel backbone, while remaining behind generation\-specialized models\. These gains demonstrate the benefit of adaptive computation allocation without increasing model scale\.

Table 1:Main results\. FlexUMM is applied on top of the Bagel backbone; all other rows are released unified models evaluated under the same protocol\. Best results are inbold\.ModelParamsUnderstandingGenerationMMMUMMBench\-CNMMBench\-ENMathVistaMMStarGenEvalWISEBagel\([Deng et al\., 2025](https://arxiv.org/html/2610.22682#bib.bib31)\)14B51\.980\.9482\.6571\.363\.2078\.810\.4044Janus\-Pro\([Chen et al\., 2025](https://arxiv.org/html/2610.22682#bib.bib21)\)7B40\.969\.9364\.8443\.847\.9379\.020\.3900Janus\([Wu et al\., 2025a](https://arxiv.org/html/2610.22682#bib.bib22)\)1\.3B27\.337\.1927\.6926\.626\.6740\.040\.2222Janus Flow\([Ma et al\., 2025](https://arxiv.org/html/2610.22682#bib.bib15)\)1\.3B29\.060\.3164\.1634\.839\.1349\.990\.2964Show\-O2\([Xie et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib14)\)7B47\.976\.6877\.2551\.554\.5359\.870\.3595Show\-O2\([Xie et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib14)\)1\.5B37\.161\.0465\.0937\.943\.8755\.490\.3349Show\-O\([Xie et al\., 2025](https://arxiv.org/html/2610.22682#bib.bib20)\)1\.3B26\.111\.0143\.4829\.027\.8765\.060\.3037Emu3\([Wang et al\., 2024](https://arxiv.org/html/2610.22682#bib.bib9)\)8B31\.445\.6659\.5344\.943\.6045\.760\.3373Emu3\.5\([Cui et al\., 2025](https://arxiv.org/html/2610.22682#bib.bib13)\)34B29\.218\.1317\.5641\.730\.2781\.830\.6331Omnigen2\([Wu et al\., 2026a](https://arxiv.org/html/2610.22682#bib.bib10)\)7B46\.075\.4376\.3663\.552\.4778\.530\.4029Harmon\([Wu et al\., 2025b](https://arxiv.org/html/2610.22682#bib.bib40)\)1\.5B35\.039\.4348\.9422\.036\.5367\.810\.3917Bagel\+SFT–52\.680\.8882\.6573\.168\.2778\.030\.2274Bagel\+UniCot\([Qin et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib12)\)–53\.180\.9983\.1283\.070\.0077\.910\.4037Bagel\+IRG\([Huang et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib11)\)–48\.057\.3560\.4768\.061\.5372\.060\.3842Bagel\+FlexUMM–54\.085\.8786\.7573\.470\.9382\.030\.4292

One model serves multiple budgets\.The inference\-time thresholdηβ\\eta\_\{\\beta\}traces an accuracy\-compute curve without retraining\.[Table2](https://arxiv.org/html/2610.22682#S4.T2)shows that adaptive halting consistently outperforms static truncation across the operating range, with the optimal depth varying across benchmarks\. This supports input\- and benchmark\-dependent computation allocation rather than a single fixed depth\.

Table 2:Effect of the halting thresholdηβ\\eta\_\{\\beta\}on Bagel\.Settingηβ\\eta\_\{\\beta\}d¯\\bar\{d\}Rel\. computeMMMUMMStarMathVistaAggressive0\.3026\.0720\.93052\.971\.269\.1Balanced0\.5027\.8410\.99354\.070\.973\.4Conservative0\.8028\.0001\.00053\.670\.873\.1\(a\)Quality–efficiency trade\-off\.\(b\)Exit\-layer distribution\.
Figure 4:Quality\-efficiency trade\-off and adaptive exit distribution\.[Figure4\(a\)](https://arxiv.org/html/2610.22682#S4.F4.sf1)further shows the accuracy\-efficiency trade\-off: the efficiency\-safe policy preserves full\-depth quality while reducing the layer\-token compute proxy, whereas more aggressive early exit yields larger savings at some quality cost\.

Low adaptation overhead\.The added routers, halting heads, and low\-rank residual banks increase the checkpoint size by only 0\.381% \(Appendix[Table7](https://arxiv.org/html/2610.22682#A1.T7)\)\. Early exit reduces the layer\-token compute proxy, while Decouple activates only a selected low\-rank residual branch\. We report these as compute proxy savings rather than direct wall\-clock speedups, since the current dense implementation still evaluates the base FFN \([Section4\.7](https://arxiv.org/html/2610.22682#S4.SS7)\)\.

### 4\.3Ablation Study

Our ablations are run on Bagel in[Table3](https://arxiv.org/html/2610.22682#S4.T3)\.

First, FlexUMM specializes FFN capacity while keeping attention shared, following the gradient\-alignment analysis in[SectionB\.1](https://arxiv.org/html/2610.22682#A2.SS1)\. The first block of[Table3](https://arxiv.org/html/2610.22682#S4.T3)confirms this choice: routing attention alone underperforms the backbone on most understanding benchmarks, while routing both sublayers recovers only part of the gain\. Specialization is therefore most effective when applied to the FFN while preserving shared attention\.

Second, FlexUMM uses input\-adaptive depth rather than a fixed truncation\. The second block of[Table3](https://arxiv.org/html/2610.22682#S4.T3)shows that fixed depths around FlexUMM’s mean exit depth are consistently suboptimal: changing the fixed depth shifts which benchmarks degrade without eliminating the loss\. This supports the need for input\-dependent depth allocation\.

Third, we test whether gains come from learned routing rather than added parameters\. The last block of[Table3](https://arxiv.org/html/2610.22682#S4.T3)keeps the architecture and parameter count fixed but replaces learned routing with random decisions\. Random routing loses FlexUMM’s gains and falls below the backbone on several benchmarks, confirming that performance depends on effective computation allocation rather than parameter addition alone\. The remaining design\-choice ablations, including the capacity and stopping variants, are reported in Appendix[C\.1](https://arxiv.org/html/2610.22682#A3.SS1)\.

Table 3:Ablation study on Bagel\.VariantUnderstandingGenerationMMMUMMBench\-CNMMBench\-ENMathVistaMMStarGenEvalWISEBagel51\.980\.9482\.6571\.363\.2078\.810\.4044Bagel\+FlexUMM53\.685\.8786\.8173\.170\.8082\.030\.4292*Capacity axis: which components are routed*Routed attention, shared FFN52\.182\.2681\.9571\.262\.3081\.210\.4028Routed attention, routed FFN53\.284\.8184\.2471\.968\.9080\.650\.3995*Depth axis: fixed instead of input\-dependent*Fixed depthd=25d=2551\.381\.6183\.5848\.268\.8077\.190\.3974Fixed depthd=26d=2652\.181\.0184\.1367\.267\.9076\.870\.4039*Routing policy*Random router52\.480\.4874\.3471\.064\.5376\.380\.3849

Table 4:Unification test on Unison\([Liu et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib52)\)and Unified\-Bench\([Yan et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib54)\)\.MethodUnisonUnified\-BenchInt\. Cons\.U\.\-G\. Gen\.G\.\-G\. Und\.Mut\. Enh\.OverallCLIPDINOv2DINOv3LongCLIPOverallBagel80\.7272\.0933\.6835\.8755\.5989\.4778\.7772\.4093\.2183\.46\+FlexUMM83\.21↑\\uparrow3\.1%72\.96↑\\uparrow1\.2%37\.08↑\\uparrow10\.1%46\.08↑\\uparrow28\.5%59\.83↑\\uparrow7\.6%90\.91↑\\uparrow1\.6%79\.88↑\\uparrow1\.4%74\.53↑\\uparrow2\.9%94\.87↑\\uparrow1\.8%85\.05↑\\uparrow1\.9%

### 4\.4Evaluation on UMM Unification

We evaluate FlexUMM on two unified\-model benchmarks, Unison\([Liu et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib52)\)and Unified\-Bench\([Yan et al\., 2026](https://arxiv.org/html/2610.22682#bib.bib54)\)\. Effective unification requires a model to support both understanding and generation while maintaining consistent performance across these complementary capabilities\.[Table4](https://arxiv.org/html/2610.22682#S4.T4)shows consistent improvements over the baseline on both unification benchmarks, with gains across all evaluated dimensions and overall scores\. These results suggest that adaptive computation and input\-dependent depth selection help allocate computation more flexibly across different capabilities, improving the balance between understanding and generation\. Overall, FlexUMM enhances the unified performance of the model across both evaluated benchmarks\.

### 4\.5Evaluation on Safety

Table 5:Uni\-SafeBench safety results\.ModelC\-Unsafe↓\\downarrowC\-Reject↑\\uparrowI\-Unsafe↓\\downarrowBoth\-Safe↑\\uparrowBagel43\.72%5\.07%49\.20%43\.96%\+FlexUMM42\.35%5\.23%49\.06%44\.86%

We further assess safety using Uni\-SafeBench\([Peng et al\., 2026b](https://arxiv.org/html/2610.22682#bib.bib62)\), which is important for UMMs that combine multimodal understanding with generation, requiring both accurate safety recognition and appropriate response behavior\. As shown in[Table5](https://arxiv.org/html/2610.22682#S4.T5), FlexUMM lowers both unsafe rates while increasing the rejection rate and the rate of outputs that are safe under both criteria\. These improvements suggest that adaptive depth selection provides more flexible computation for processing safety\-relevant cues and making appropriate response decisions\.

### 4\.6Generalization Across Backbones

FlexUMM is agnostic to UMM architectures\.[Table6](https://arxiv.org/html/2610.22682#S4.T6)evaluates FlexUMM on two additional backbones, Janus\-Pro and Harmon\. Generalization across backbones is important for assessing whether the benefits of adaptive computation extend beyond a specific architecture or sharing structure\. FlexUMM consistently improves understanding and generation performance across all three backbones while reducing computation, demonstrating robust generalization across unified\-model architectures with different sharing structures\. The results suggest that the proposed share\-specialize\-stop formulation can adapt to diverse backbone designs without relying on architecture\-specific modifications\.

### 4\.7Effectiveness Analysis

We analyze the computation structure underlying FlexUMM: attention gradients remain aligned across functionalities, while FFN gradients become increasingly competitive with depth \([SectionB\.1](https://arxiv.org/html/2610.22682#A2.SS1)\)\. Depth requirements also vary substantially across inputs \([SectionB\.2](https://arxiv.org/html/2610.22682#A2.SS2)\), supporting FlexUMM’s capability\-adaptive capacity allocation and input\-adaptive computation depth\.

Exit\-layer distribution is consistent across benchmarks\.[Figure4\(b\)](https://arxiv.org/html/2610.22682#S4.F4.sf2)shows that across benchmarks, the router systematically exits at earlier candidate depths\. Consequently, the sequence\-depth component contributes a modest reduction in relative computation across benchmarks\. This consistent behavior indicates that adaptive depth selectively allocates less computation to easier inputs, rather than relying heavily on early exit for only specific benchmarks\.

### 4\.8Case Study

We inspect representative outputs under the same generation prompt and seed\.

Understanding\.The details are provided in Appendix[AppendixD](https://arxiv.org/html/2610.22682#A4)which cover medical terminology, geometric reasoning, person identification, counting, and color recognition\.

Generation\.We use this case study to illustrate how input\-adaptive depth selection affects generation quality under matched evaluation conditions\. All outputs are produced with the same checkpoint, prompt, random seed, and decoding settings; only the available candidate\-depth configurations differ\. The Bagel baseline executes the full backbone, whereas the 2\-, 3\-, and 4\-exit configurations allow the policy to select from progressively finer\-grained computation depths\. As shown in[Figure5](https://arxiv.org/html/2610.22682#S4.F5), the baseline fails to satisfy the requested spatial relation, while adaptive depth selection progressively recovers the relation and improves visual consistency\. The example is illustrative rather than a standalone benchmark result: Appendix[AppendixD](https://arxiv.org/html/2610.22682#A4)reports ten strict paired understanding wins and six detector\-verified generation wins, while the aggregate benchmark results quantify the overall accuracy\-efficient benefit\.

Table 6:Generalization across backbones with different sharing structures\.BackboneMethodUnderstandingGenerationEfficiencyMMMUMMBench\-CNMMBench\-ENMathVistaMMStarGenEvalWISERel\. computeBagelFull backbone51\.980\.9482\.6571\.363\.2078\.810\.40441\.00FlexUMM54\.0↑\\uparrow4\.0%85\.87↑\\uparrow6\.1%86\.75↑\\uparrow5\.0%73\.4↑\\uparrow2\.9%70\.93↑\\uparrow12\.2%82\.03↑\\uparrow4\.1%0\.4292↑\\uparrow6\.1%0\.96Janus\-ProFull backbone40\.8964\.8469\.9343\.847\.9379\.020\.39001\.00FlexUMM46\.11↑\\uparrow12\.8%68\.47↑\\uparrow5\.6%73\.11↑\\uparrow4\.5%51\.0↑\\uparrow16\.4%52\.33↑\\uparrow9\.2%81\.03↑\\uparrow2\.5%0\.4095↑\\uparrow5\.0%0\.96HarmonFull backbone35\.039\.4348\.9422\.036\.5367\.810\.39171\.00FlexUMM38\.6↑\\uparrow10\.3%42\.6↑\\uparrow8\.0%53\.09↑\\uparrow8\.5%35\.1↑\\uparrow59\.5%43\.07↑\\uparrow17\.9%71\.71↑\\uparrow5\.8%0\.4197↑\\uparrow7\.1%0\.89

*a photo of a chair left of a zebra*Baseline2\-depth selection3\-depth selectionFlexUMM \(4 exits\)![Refer to caption](https://arxiv.org/html/2610.22682v1/chair_left_baseline.png)![Refer to caption](https://arxiv.org/html/2610.22682v1/chair_left_2exits.png)![Refer to caption](https://arxiv.org/html/2610.22682v1/chair_left_3exits.png)![Refer to caption](https://arxiv.org/html/2610.22682v1/chair_left_flexumm_v8.png)Figure 5:Qualitative depth\-selection progression for one generation prompt\.

## 5 Conclusion

We study adaptive computation in unified multimodal models and propose FlexUMM, a unified framework that dynamically allocates computation according to both capability and input complexity\. By adapting computation within and across layers while preserving a shared multimodal backbone, FlexUMM enables the model to allocate computation where it is most useful rather than applying uniform computation to all tokens and inputs\. Across three unified multimodal backbones, FlexUMM consistently improves overall performance, achieving an average improvement of10\.17%while reducing FLOPs by6\.33%\. These results show that adaptive computation provides an effective principle for improving both the capability and efficiency of unified multimodal models\. We discuss limitations in Appendix[E](https://arxiv.org/html/2610.22682#A5)\.

### AI use statement

We used generative AI tools to assist with manuscript revision, literature search, and code development\. The experimental design and methodological choices were made by the authors\.

### Reproducibility statement

Dataset protocols, preprocessing, model pools, and training hyperparameters are provided in Appendix[A](https://arxiv.org/html/2610.22682#A1)\.

## References

- S\. Bae, Y\. Kim, R\. Bayat, S\. Kim, J\. Ha, T\. Schuster, A\. Fisch, H\. Harutyunyan, Z\. Ji, A\. Courville,et al\.Mixture\-of\-recursions: learning dynamic recursive depths for adaptive token\-level computation\.Advances in Neural Information Processing Systems38,pp\. 96572–96617\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p3.1)\.
- Baiet al\.\(2026\)H\. Bai, Y\. Luo, W\. Wang, Q\. Wen, and J\. WangUniPath: adaptive coordination of understanding and generation for unified multimodal reasoning\.arXiv preprint arXiv:2605\.11400\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p1.1)\.
- Bolyaet al\.\(2022\)D\. Bolya, C\. Fu, X\. Dai, P\. Zhang, C\. Feichtenhofer, and J\. HoffmanToken merging: your vit but faster\.arXiv preprint arXiv:2210\.09461\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p3.1)\.
- Brookset al\.\(2023\)T\. Brooks, A\. Holynski, and A\. A\. EfrosInstructpix2pix: learning to follow image editing instructions\.In2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition \(CVPR\),pp\. 18392–18402\.Cited by:[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p2.1)\.
- Chenet al\.\(2024a\)L\. Chen, H\. Zhao, T\. Liu, S\. Bai, J\. Lin, C\. Zhou, and B\. ChangAn image is worth 1/2 tokens after layer 2: plug\-and\-play inference acceleration for large vision\-language models\.InEuropean Conference on Computer Vision,pp\. 19–35\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p3.1)\.
- Chenet al\.\(2024b\)L\. Chen, J\. Li, X\. Dong, P\. Zhang, Y\. Zang, Z\. Chen, H\. Duan, J\. Wang, Y\. Qiao, D\. Lin,et al\.Are we on the right way for evaluating large vision\-language models?\.Advances in Neural Information Processing Systems37,pp\. 27056–27087\.Cited by:[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p3.1)\.
- Chenet al\.\(2025\)X\. Chen, Z\. Wu, X\. Liu, Z\. Pan, W\. Liu, Z\. Xie, X\. Yu, and C\. RuanJanus\-pro: unified multimodal understanding and generation with data and model scaling\.arXiv preprint arXiv:2501\.17811\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p1.1),[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p1.1),[Table 1](https://arxiv.org/html/2610.22682#S4.T1.6.1.4.1)\.
- Cuiet al\.\(2025\)Y\. Cui, H\. Chen, H\. Deng, X\. Huang, X\. Li, J\. Liu, Y\. Liu, Z\. Luo, J\. Wang, W\. Wang,et al\.Emu3\. 5: native multimodal models are world learners\.arXiv preprint arXiv:2510\.26583\.Cited by:[Table 1](https://arxiv.org/html/2610.22682#S4.T1.6.1.11.1)\.
- Denget al\.\(2025\)C\. Deng, D\. Zhu, K\. Li, C\. Gou, F\. Li, Z\. Wang, S\. Zhong, W\. Yu, X\. Nie, Z\. Song,et al\.Emerging properties in unified multimodal pretraining\.arXiv preprint arXiv:2505\.14683\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p1.1),[§2](https://arxiv.org/html/2610.22682#S2.p1.1),[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p1.1),[Table 1](https://arxiv.org/html/2610.22682#S4.T1.6.1.3.1)\.
- Elbayadet al\.\(2019\)M\. Elbayad, J\. Gu, E\. Grave, and M\. AuliDepth\-adaptive transformer\.arXiv preprint arXiv:1910\.10073\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p3.1)\.
- Fuet al\.\(2026\)C\. Fu, P\. Chen, Y\. Shen, Y\. Qin, M\. Zhang, X\. Lin, J\. Yang, X\. Zheng, K\. Li, X\. Sun,et al\.Mme: a comprehensive evaluation benchmark for multimodal large language models\.Advances in Neural Information Processing Systems38\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p3.1)\.
- Ghoshet al\.\(2023\)D\. Ghosh, H\. Hajishirzi, and L\. SchmidtGeneval: an object\-focused framework for evaluating text\-to\-image alignment\.Advances in Neural Information Processing Systems36,pp\. 52132–52152\.Cited by:[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p3.1)\.
- Graves \(2016\)A\. GravesAdaptive computation time for recurrent neural networks\.arXiv preprint arXiv:1603\.08983\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p3.1)\.
- Hanet al\.\(2026\)Y\. Han, H\. Chen, A\. Han, Z\. Wang, X\. Liu, S\. Zhang, D\. Zou,et al\.Turning internal gap into self\-improvement: promoting the generation\-understanding unification in mllms\.InInternational Conference on Learning Representations,Vol\.2026,pp\. 136501–136534\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p1.1),[§1](https://arxiv.org/html/2610.22682#S1.p4.1)\.
- Heet al\.\(2025\)S\. He, C\. Deng, A\. Li, and S\. YanUnderstanding and harnessing sparsity in unified multimodal models\.arXiv preprint arXiv:2512\.02351\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p2.1)\.
- Huanget al\.\(2026\)W\. Huang, S\. Chen, Z\. Xie, S\. Cao, S\. Tang, Y\. Shen, Q\. Yin, W\. Hu, X\. Wang, Y\. Tang,et al\.Interleaving reasoning for better text\-to\-image generation\.InInternational Conference on Learning Representations,Vol\.2026,pp\. 106153–106182\.Cited by:[Table 1](https://arxiv.org/html/2610.22682#S4.T1.6.1.16.1)\.
- Kanget al\.\(2026\)J\. Kang, H\. Yoon, J\. Jung, J\. Min, M\. Jeon, B\. Hwang, S\. Jung, and S\. KimTransferability between understanding and generation in unified multimodal models\.InEuropean Conference on Computer Vision,pp\. 626–647\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p1.1),[§1](https://arxiv.org/html/2610.22682#S1.p3.1),[§1](https://arxiv.org/html/2610.22682#S1.p4.1),[§2](https://arxiv.org/html/2610.22682#S2.p1.1)\.
- Kembhaviet al\.\(2016\)A\. Kembhavi, M\. Salvato, E\. Kolve, M\. Seo, H\. Hajishirzi, and A\. FarhadiA diagram is worth a dozen images\.InEuropean conference on computer vision,pp\. 235–251\.Cited by:[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p2.1)\.
- Kil and Kim \(2026\)D\. Kil and S\. KimEfficient unified multimodal understanding \(eumu\): winning solution for the mumu track at the 8th lsvos challenge\.arXiv preprint arXiv:2609\.19451\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p1.1),[§1](https://arxiv.org/html/2610.22682#S1.p2.1)\.
- Liuet al\.\(2026\)J\. Liu, X\. Shuai, H\. Ding, and Y\. JiangUnison: benchmarking unified multimodal models via synergistic understanding and generation\.arXiv preprint arXiv:2606\.26984\.Cited by:[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p3.1),[§4\.4](https://arxiv.org/html/2610.22682#S4.SS4.p1.1),[Table 4](https://arxiv.org/html/2610.22682#S4.T4),[Table 4](https://arxiv.org/html/2610.22682#S4.T4.4)\.
- Liuet al\.\(2024\)Y\. Liu, H\. Duan, Y\. Zhang, B\. Li, S\. Zhang, W\. Zhao, Y\. Yuan, J\. Wang, C\. He, Z\. Liu,et al\.Mmbench: is your multi\-modal model an all\-around player?\.InEuropean conference on computer vision,pp\. 216–233\.Cited by:[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p3.1)\.
- Luet al\.\(2022a\)J\. Lu, C\. Clark, R\. Zellers, R\. Mottaghi, and A\. KembhaviUnified\-io: a unified model for vision, language, and multi\-modal tasks\.arXiv preprint arXiv:2206\.08916\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p1.1)\.
- Luet al\.\(2024\)P\. Lu, H\. Bansal, T\. Xia, J\. Liu, C\. Li, H\. Hajishirzi, H\. Cheng, K\. Chang, M\. Galley, and J\. GaoMathvista: evaluating mathematical reasoning of foundation models in visual contexts\.InInternational Conference on Learning Representations,Vol\.2024,pp\. 23439–23554\.Cited by:[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p3.1)\.
- Luet al\.\(2022b\)P\. Lu, S\. Mishra, T\. Xia, L\. Qiu, K\. Chang, S\. Zhu, O\. Tafjord, P\. Clark, and A\. KalyanLearn to explain: multimodal reasoning via thought chains for science question answering\.Advances in neural information processing systems35,pp\. 2507–2521\.Cited by:[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p2.1)\.
- Luoet al\.\(2026a\)Y\. Luo, H\. Gokhale, M\. Savvides, J\. Wang, and S\. HeSelf\-corrected image generation with explainable latent rewards\.arXiv preprint arXiv:2603\.24965\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p1.1)\.
- Luoet al\.\(2026b\)Y\. Luo, W\. Wang, H\. Bai, M\. Savvides, and J\. WangLatentUMM: dual latent alignment for unified multimodal models\.arXiv preprint arXiv:2605\.17766\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p1.1)\.
- Luoet al\.\(2026c\)Y\. Luo, W\. Wang, H\. Bai, H\. Zhu, H\. Chen, P\. He, M\. Savvides, S\. Li, and J\. WangTorchumm: a unified multimodal model codebase for evaluation, analysis, and post\-training\.arXiv preprint arXiv:2604\.10784\.Cited by:[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p4.1)\.
- Maet al\.\(2018\)J\. Ma, Z\. Zhao, X\. Yi, J\. Chen, L\. Hong, and E\. H\. ChiModeling task relationships in multi\-task learning with multi\-gate mixture\-of\-experts\.InProceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining,pp\. 1930–1939\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p2.1)\.
- Maet al\.\(2025\)Y\. Ma, X\. Liu, X\. Chen, W\. Liu, C\. Wu, Z\. Wu, Z\. Pan, Z\. Xie, H\. Zhang, X\. Yu,et al\.Janusflow: harmonizing autoregression and rectified flow for unified multimodal understanding and generation\.In2025 IEEE/CVF Conference on Computer Vision and Pattern Recognition \(CVPR\),pp\. 7739–7751\.Cited by:[Table 1](https://arxiv.org/html/2610.22682#S4.T1.6.1.6.1)\.
- Misraet al\.\(2016\)I\. Misra, A\. Shrivastava, A\. Gupta, and M\. HebertCross\-stitch networks for multi\-task learning\.InProceedings of the IEEE conference on computer vision and pattern recognition,pp\. 3994–4003\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p2.1)\.
- Molchanovet al\.\(2019\)P\. Molchanov, A\. Mallya, S\. Tyree, I\. Frosio, and J\. KautzImportance estimation for neural network pruning\.In2019 IEEE/CVF conference on computer vision and pattern recognition \(CVPR\),pp\. 11256–11264\.Cited by:[§B\.1\.3](https://arxiv.org/html/2610.22682#A2.SS1.SSS3.Px1.p1.2)\.
- Niuet al\.\(2025\)Y\. Niu, M\. Ning, M\. Zheng, W\. Jin, B\. Lin, P\. Jin, J\. Liao, C\. Feng, F\. Meng, K\. Ning,et al\.Wise: a world knowledge\-informed semantic evaluation for text\-to\-image generation\.arXiv preprint arXiv:2503\.07265\.Cited by:[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p3.1)\.
- Penget al\.\(2026a\)W\. Peng, L\. Meng, Y\. Cai, X\. Zhuang, Y\. Yang, R\. Fang, C\. Wu, J\. Lin, Z\. Wu, and S\. BaiUnified multimodal autoregressive modeling with shared context\-visual tokenizer is key to unification\.arXiv preprint arXiv:2606\.18249\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p4.1),[§2](https://arxiv.org/html/2610.22682#S2.p1.1)\.
- Penget al\.\(2026b\)Z\. Peng, Y\. Xu, Q\. Zhang, J\. Shen, Y\. Zhang, H\. Xu, Y\. Wang, and G\. GouDoes unification come at a cost? uni\-safebench: a safety benchmark for unified multimodal large models\.arXiv preprint arXiv:2604\.00547\.Cited by:[§4\.5](https://arxiv.org/html/2610.22682#S4.SS5.p1.1)\.
- Qinet al\.\(2026\)L\. Qin, G\. JIA, Y\. Sun, T\. Li, H\. Pan, M\. Yang, X\. Yang, C\. Qu, Z\. Tan, and H\. LiUni\-cot: towards unified chain\-of\-thought reasoning across text and vision\.InInternational Conference on Learning Representations,Vol\.2026,pp\. 15999–16028\.Cited by:[Table 1](https://arxiv.org/html/2610.22682#S4.T1.6.1.15.1)\.
- Qiuet al\.\(2026\)Z\. Qiu, Y\. Wu, K\. Fan, B\. Li, and H\. XiongWhere a new concept must enter: entry point gates cross\-task usability in unified multimodal models\.arXiv preprint arXiv:2608\.17564\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p2.1),[§1](https://arxiv.org/html/2610.22682#S1.p4.1),[§2](https://arxiv.org/html/2610.22682#S2.p1.1)\.
- Rahimianet al\.\(2023\)E\. Rahimian, G\. Javadi, F\. Tung, and G\. OliveiraDynaShare: task and instance conditioned parameter sharing for multi\-task learning\.In2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops \(CVPRW\),pp\. 4535–4543\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p2.1)\.
- Raoet al\.\(2021\)Y\. Rao, W\. Zhao, B\. Liu, J\. Lu, J\. Zhou, and C\. HsiehDynamicvit: efficient vision transformers with dynamic token sparsification\.Advances in neural information processing systems34,pp\. 13937–13949\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p3.1)\.
- Raposoet al\.\(2024\)D\. Raposo, S\. Ritter, B\. Richards, T\. Lillicrap, P\. C\. Humphreys, and A\. SantoroMixture\-of\-depths: dynamically allocating compute in transformer\-based language models\.arXiv preprint arXiv:2404\.02258\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p3.1)\.
- Ryooet al\.\(2021\)M\. S\. Ryoo, A\. Piergiovanni, A\. Arnab, M\. Dehghani, and A\. AngelovaTokenlearner: what can 8 learned tokens do for images and videos?\.arXiv preprint arXiv:2106\.11297\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p3.1)\.
- Shanget al\.\(2025\)Y\. Shang, M\. Cai, B\. Xu, Y\. J\. Lee, and Y\. YanLlava\-prumerge: adaptive token reduction for efficient large multimodal models\.In2025 IEEE/CVF International Conference on Computer Vision \(ICCV\),pp\. 22857–22867\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p3.1)\.
- Shiet al\.\(2024\)W\. Shi, Z\. Hu, Y\. Bin, J\. Liu, Y\. Yang, S\. K\. Ng, L\. Bing, and R\. K\. LeeMath\-llava: bootstrapping mathematical reasoning for multimodal large language models\.InFindings of the Association for Computational Linguistics: EMNLP 2024,pp\. 4663–4680\.Cited by:[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p2.1)\.
- Singhet al\.\(2019\)A\. Singh, V\. Natarajan, M\. Shah, Y\. Jiang, X\. Chen, D\. Batra, D\. Parikh, and M\. RohrbachTowards vqa models that can read\.In2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition \(CVPR\),pp\. 8309–8318\.Cited by:[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p2.1)\.
- Suet al\.\(2026\)Z\. Su, H\. Wei, K\. Cen, Y\. Wang, G\. Chen, C\. Yuan, and X\. ChuGeneration enhances understanding in unified multimodal models via multi\-representation generation\.arXiv preprint arXiv:2601\.21406\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p3.1),[§1](https://arxiv.org/html/2610.22682#S1.p4.1)\.
- Sunet al\.\(2024\)Q\. Sun, Y\. Cui, X\. Zhang, F\. Zhang, Q\. Yu, Y\. Wang, Y\. Rao, J\. Liu, T\. Huang, and X\. WangGenerative multimodal models are in\-context learners\.InProceedings of the IEEE/CVF conference on computer vision and pattern recognition,pp\. 14398–14409\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p1.1)\.
- Sunet al\.\(2020\)X\. Sun, R\. Panda, R\. Feris, and K\. SaenkoAdashare: learning what to share for efficient deep multi\-task learning\.Advances in Neural Information Processing Systems33,pp\. 8728–8740\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p2.1)\.
- Tanget al\.\(2020\)H\. Tang, J\. Liu, M\. Zhao, and X\. GongProgressive layered extraction \(ple\): a novel multi\-task learning \(mtl\) model for personalized recommendations\.InProceedings of the 14th ACM conference on recommender systems,pp\. 269–278\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p2.1)\.
- Tanget al\.\(2023\)S\. Tang, Y\. Wang, Z\. Kong, T\. Zhang, Y\. Li, C\. Ding, Y\. Wang, Y\. Liang, and D\. XuYou need multiple exiting: dynamic early exiting for accelerating unified vision language model\.In2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition \(CVPR\),pp\. 10781–10791\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p3.1)\.
- Wanget al\.\(2022\)P\. Wang, A\. Yang, R\. Men, J\. Lin, S\. Bai, Z\. Li, J\. Ma, C\. Zhou, J\. Zhou, and H\. YangOfa: unifying architectures, tasks, and modalities through a simple sequence\-to\-sequence learning framework\.InInternational conference on machine learning,pp\. 23318–23340\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p1.1)\.
- Wanget al\.\(2024\)X\. Wang, X\. Zhang, Z\. Luo, Q\. Sun, Y\. Cui, J\. Wang, F\. Zhang, Y\. Wang, Z\. Li, Q\. Yu,et al\.Emu3: next\-token prediction is all you need\.arXiv preprint arXiv:2409\.18869\.Cited by:[Table 1](https://arxiv.org/html/2610.22682#S4.T1.6.1.10.1)\.
- Wuet al\.\(2025a\)C\. Wu, X\. Chen, Z\. Wu, Y\. Ma, X\. Liu, Z\. Pan, W\. Liu, Z\. Xie, X\. Yu, C\. Ruan,et al\.Janus: decoupling visual encoding for unified multimodal understanding and generation\.In2025 IEEE/CVF Conference on Computer Vision and Pattern Recognition \(CVPR\),pp\. 12966–12977\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p2.1),[Table 1](https://arxiv.org/html/2610.22682#S4.T1.6.1.5.1)\.
- Wuet al\.\(2026a\)C\. Wu, J\. Wang, P\. Zheng, R\. Yan, S\. Xiao, X\. Luo, Y\. Wang, W\. Li, X\. Jiang, Y\. Liu,et al\.Omnigen2: towards instruction\-aligned multimodal generation\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,pp\. 21964–21975\.Cited by:[Table 1](https://arxiv.org/html/2610.22682#S4.T1.6.1.12.1)\.
- Wuet al\.\(2026b\)H\. Wu, Y\. Fan, D\. Jinyang, J\. Tong, Y\. Ma, and X\. ShenHidrop: hierarchical vision token reduction in mllms via late injection, concave pyramid pruning, and early exit\.InInternational Conference on Learning Representations,Vol\.2026,pp\. 97296–97316\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p3.1)\.
- Wuet al\.\(2026c\)P\. Wu, H\. Diao, W\. Fan, L\. Lu, D\. Lin, and Z\. LiuUncovering understanding\-generation synergy in native unified multimodal models: from representation, task to system\.arXiv preprint arXiv:2609\.01607\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p1.1),[§1](https://arxiv.org/html/2610.22682#S1.p2.1),[§1](https://arxiv.org/html/2610.22682#S1.p4.1),[§2](https://arxiv.org/html/2610.22682#S2.p1.1),[§2](https://arxiv.org/html/2610.22682#S2.p2.1)\.
- Wuet al\.\(2026d\)Q\. Wu, W\. Lin, Y\. Zhou, W\. Ye, Z\. Zeng, X\. Sun, and R\. JiAccelerating multimodal large language models via dynamic visual\-token exit and the empirical findings\.Advances in Neural Information Processing Systems38,pp\. 168378–168403\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p3.1)\.
- Wuet al\.\(2025b\)S\. Wu, W\. Zhang, L\. Xu, S\. Jin, Z\. Wu, Q\. Tao, W\. Liu, W\. Li, and C\. C\. LoyHarmonizing visual representations for unified multimodal understanding and generation\.In2025 IEEE/CVF International Conference on Computer Vision \(ICCV\),pp\. 17739–17750\.Cited by:[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p1.1),[Table 1](https://arxiv.org/html/2610.22682#S4.T1.6.1.13.1)\.
- Xieet al\.\(2025\)J\. Xie, W\. Mao, Z\. Bai, D\. J\. Zhang, W\. Wang, K\. Q\. Lin, Y\. Gu, Z\. Chen, Z\. Yang, and M\. Z\. ShouShow\-o: one single transformer to unify multimodal understanding and generation\.InInternational Conference on Learning Representations,Vol\.2025,pp\. 28240–28264\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p1.1),[Table 1](https://arxiv.org/html/2610.22682#S4.T1.6.1.9.1)\.
- Xieet al\.\(2026\)J\. Xie, Z\. Yang, and M\. Z\. ShouShow\-o2: improved native unified multimodal models\.Advances in Neural Information Processing Systems38,pp\. 47490–47518\.Cited by:[Table 1](https://arxiv.org/html/2610.22682#S4.T1.6.1.7.1),[Table 1](https://arxiv.org/html/2610.22682#S4.T1.6.1.8.1)\.
- Yanet al\.\(2026\)Z\. Yan, K\. Lin, Z\. Li, J\. Ye, H\. Han, H\. Wang, Z\. Wang, B\. Lin, H\. Li, X\. Xiao,et al\.Unified multimodal models as auto\-encoders\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,pp\. 41903–41912\.Cited by:[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p3.1),[§4\.4](https://arxiv.org/html/2610.22682#S4.SS4.p1.1),[Table 4](https://arxiv.org/html/2610.22682#S4.T4),[Table 4](https://arxiv.org/html/2610.22682#S4.T4.4)\.
- Yeet al\.\(2026\)S\. Ye, M\. Xu, S\. Gu, D\. He, L\. Wang, and W\. HuUnderstanding vs\. generation: navigating optimization dilemma in multimodal models\.InInternational Conference on Learning Representations,Vol\.2026,pp\. 25118–25141\.Cited by:[§1](https://arxiv.org/html/2610.22682#S1.p2.1)\.
- Yueet al\.\(2024\)X\. Yue, Y\. Ni, K\. Zhang, T\. Zheng, R\. Liu, G\. Zhang, S\. Stevens, D\. Jiang, W\. Ren, Y\. Sun,et al\.Mmmu: a massive multi\-discipline multimodal understanding and reasoning benchmark for expert agi\.InProceedings of the IEEE/CVF conference on computer vision and pattern recognition,pp\. 9556–9567\.Cited by:[§4\.1](https://arxiv.org/html/2610.22682#S4.SS1.p3.1)\.
- Zhanget al\.\(2024\)Y\. Zhang, C\. Fan, J\. Ma, W\. Zheng, T\. Huang, K\. Cheng, D\. Gudovskiy, T\. Okuno, Y\. Nakata, K\. Keutzer,et al\.Sparsevlm: visual token sparsification for efficient vision\-language model inference\.arXiv preprint arXiv:2410\.04417\.Cited by:[§2](https://arxiv.org/html/2610.22682#S2.p3.1)\.

Appendix FlexUMM: FlexUMM

## Contents

## Appendix AConfiguration

This section describes the runtime configuration of FlexUMM\. It focuses on how the gates, residual banks, and depth controller are inserted into the shared decoder and how they interact during a forward pass\. It intentionally does not describe parameter selection, tuning schedules, or the post\-training procedure\.

### A\.1Backbone configurations

The backbone settings and the number of trained parameters used in our experiments are summarized in[Table7](https://arxiv.org/html/2610.22682#A1.T7)\.

Table 7:Backbone configurations for FlexUMM\.BackboneParamsLang\. layersCandidate depths𝒟\\mathcal\{D\}ExpertsEETrained paramsBagel14\.607B28\{22,24,26,28\}\\\{22,24,26,28\\\}20\.056B \(\+0\.381%\)Janus\-Pro7\.420B30\{21,24,27,30\}\\\{21,24,27,30\\\}20\.035B \(\+0\.472%\)Harmon1\.5B28\{22,24,26,28\}\\\{22,24,26,28\\\}20\.013B \(\+0\.837%\)
### A\.2Insertion into the shared decoder

The runtime configuration attaches FlexUMM to the language\-decoder stack of the UMM rather than constructing a second model\. Each decoder layer keeps its pretrained attention, normalization, and base FFN modules\. For a layerℓ\\ell, the shared part of the computation is therefore the ordinary decoder operation

𝒉~\(ℓ\)=𝒉\(ℓ−1\)\+Attn\(ℓ\)⁡\(Norm1\(ℓ\)⁡\(𝒉\(ℓ−1\)\)\),𝒖\(ℓ\)=Norm2\(ℓ\)⁡\(𝒉~\(ℓ\)\)\.\\tilde\{\{\\bm\{h\}\}\}^\{\(\\ell\)\}=\{\\bm\{h\}\}^\{\(\\ell\-1\)\}\+\\operatorname\{Attn\}^\{\(\\ell\)\}\\\!\\left\(\\operatorname\{Norm\}\_\{1\}^\{\(\\ell\)\}\\\!\\left\(\{\\bm\{h\}\}^\{\(\\ell\-1\)\}\\right\)\\right\),\\qquad\{\\bm\{u\}\}^\{\(\\ell\)\}=\\operatorname\{Norm\}\_\{2\}^\{\(\\ell\)\}\\\!\\left\(\\tilde\{\{\\bm\{h\}\}\}^\{\(\\ell\)\}\\right\)\.\(11\)The TAFE module is placed at the FFN call\. Thus, attention remains the shared communication channel, while the FFN residual is the component whose capacity can be shared, specialized, or suppressed\. Understanding and generation have separate TAFE controllers and residual banks, even when both token streams occur in one packed transformer forward\.

The runtime preserves the original names and parameter objects of the base FFN projections\. A checkpoint can therefore load the pretrained attention, normalization, and FFN weights directly; FlexUMM adds only the controller modules and residual banks around those weights\. Packed batches carry token indexes, sample lengths, and task metadata\. Before a gate is called, the runtime maps each routed token back to its packed sample so that the correct task embedding and modality\-specific controller are used\.

### A\.3TAFE controller

For tokenkkat layerℓ\\ell, the gate forms a route feature by combining the current hidden state with task, layer, route\-kind, and remaining\-budget information:

𝒛k\(ℓ\)=Ws​LN⁡\(𝒖k\(ℓ\)\)\+𝒆task⁡\(k\)\+𝒆layer⁡\(ℓ\)\+𝒆kind\+Wb​bk\(ℓ\)\.\{\\bm\{z\}\}\_\{k\}^\{\(\\ell\)\}=W\_\{s\}\\,\\operatorname\{LN\}\(\{\\bm\{u\}\}\_\{k\}^\{\(\\ell\)\}\)\+\{\\bm\{e\}\}\_\{\\mathrm\{task\}\(k\)\}\+\{\\bm\{e\}\}\_\{\\mathrm\{layer\}\(\\ell\)\}\+\{\\bm\{e\}\}\_\{\\mathrm\{kind\}\}\+W\_\{b\}b^\{\(\\ell\)\}\_\{k\}\.\(12\)Two small MLP heads read this feature\. The utility head produces three values, one for each action in\{Share,Decouple,Exit\}\\\{\\textsc\{Share\},\\textsc\{Decouple\},\\textsc\{Exit\}\\\}, and the subset head produces one logit for each residual bank\. The action decision uses the cost\-adjusted utility in Eq\. \([2](https://arxiv.org/html/2610.22682#S3.E2)\); the subset decision is the argmax bank index \(or the top\-kkbank indexes when multiple residuals are enabled\)\. The subset index is predicted from the hidden state and metadata; it is not a human\-supplied task label\.

The hard inference path is implemented tokenwise\. For aSharetoken, the base FFN output is retained\. For aDecoupletoken, one selected residual bank is evaluated and added to the base output\. For anExittoken, the FFN output is replaced by zero, so the decoder’s outer residual connection carries the post\-attention state forward unchanged by the FFN\. If top\-kkrouting is used, the selected residual outputs are averaged before being added\. The resulting FFN update is exactly the piecewise operation in Eq\. \([3](https://arxiv.org/html/2610.22682#S3.E3)\)\.

During differentiable routing, the hard action is replaced by the soft mixtureπ\\pifrom Eq\. \([2](https://arxiv.org/html/2610.22682#S3.E2)\)\. Withρk,e\\rho\_\{k,e\}denoting the soft bank weights, the routed FFN output is

Δ​𝒉k\(ℓ\)=\(πk,Share\(ℓ\)\+πk,Decouple\(ℓ\)\)​F​\(𝒖k\(ℓ\)\)\+πk,Decouple\(ℓ\)​∑e∈𝒯kρk,e\(ℓ\)​Re​\(𝒖k\(ℓ\)\),\\Delta\{\\bm\{h\}\}\_\{k\}^\{\(\\ell\)\}=\\left\(\\pi\_\{k,\\textsc\{Share\}\}^\{\(\\ell\)\}\+\\pi\_\{k,\\textsc\{Decouple\}\}^\{\(\\ell\)\}\\right\)F\(\{\\bm\{u\}\}\_\{k\}^\{\(\\ell\)\}\)\+\\pi\_\{k,\\textsc\{Decouple\}\}^\{\(\\ell\)\}\\sum\_\{e\\in\\mathcal\{T\}\_\{k\}\}\\rho\_\{k,e\}^\{\(\\ell\)\}R\_\{e\}\(\{\\bm\{u\}\}\_\{k\}^\{\(\\ell\)\}\),\(13\)where𝒯k\\mathcal\{T\}\_\{k\}is the selected top\-kkbank set\. TheExitprobability contributes neither the shared FFN output nor a specialized residual\. Evaluation switches back to the discrete argmax route\.

### A\.4Low\-rank residual banks

Each configured specialized bank is a residual adapter rather than a copied FFN block\. For bankee, the configured residual block computes

Re​\(𝒖\)=Ue​SiLU⁡\(De​RMSNorm⁡\(𝒖\)\),De∈ℝr×d,Ue∈ℝd×r\.R\_\{e\}\(\{\\bm\{u\}\}\)=U\_\{e\}\\,\\operatorname\{SiLU\}\\\!\\left\(D\_\{e\}\\,\\operatorname\{RMSNorm\}\(\{\\bm\{u\}\}\)\\right\),\\qquad D\_\{e\}\\in\\mathbb\{R\}^\{r\\times d\},\\quad U\_\{e\}\\in\\mathbb\{R\}^\{d\\times r\}\.\(14\)The down projection is initialized with a small random matrix and the up projection is initialized to zero\. Consequently, inserting the banks does not change the released backbone’s initial forward function, while each bank can learn a different functionality\- or modality\-conditioned residual\. The dense wrapper first evaluates the pretrained base FFNF⁡\(𝒖\)F\(\{\\bm\{u\}\}\)for every routed token\. The selected action is then applied to that result: aSharetoken retainsF⁡\(𝒖\)F\(\{\\bm\{u\}\}\), aDecoupletoken adds one selected low\-rank residualRe​\(𝒖\)R\_\{e\}\(\{\\bm\{u\}\}\), and anExittoken masks the FFN update after the base result has been formed\. ThusDecoupleadds only a low\-rank residual bank \(roughly the two rank\-rrprojections\) rather than a second full FFN, butShareorExitshould not be interpreted as removing the dense base\-FFN kernel in this runtime\. The action\-level compute proxy credits the low\-rank branch only to selected tokens; direct execution savings come from terminating the sequence before later decoder layers\.

### A\.5Candidate\-depth execution

The depth controller is separate from the tokenwise TAFE gate\. It is a small MLP shared across the configured candidate exits\. At candidate depthdmd\_\{m\}, the controller pools the current shared hidden state for the routed sequence and combines that pooled state with the candidate\-depth embedding, route\-kind embedding, task embedding, and generation timestep when one is available\. An optional previous\-candidate state supplies a representation\-change feature\. The controller returns a halt logit and, when enabled, a predicted harmful regret and uncertainty scale\.

The runtime evaluates layers only up to the largest configured candidate\. At each candidate it computesphalt,mp\_\{\\mathrm\{halt\},m\}and applies the eligibility threshold from Eq\. \([5](https://arxiv.org/html/2610.22682#S3.E5)\)\. The first eligible candidate is selected; the last candidate is always the fallback\. If all active token routes takeExitat an earlier layer, the decoder loop can terminate immediately, before the next candidate\-depth check\. Thus the two controls have different granularities: TAFE acts on token–layer pairs, while the depth controller terminates the remaining sequence\-level computation\.

#### A\.5\.1Quality\-constrained router training

The depth router is trained from all configured candidate depths in one forward pass\. If the candidate depths ared1<⋯<dMd\_\{1\}<\\cdots<d\_\{M\}, the runtime stores the hidden statehi,mh\_\{i,m\}for sampleiiat everydmd\_\{m\}while executing the backbone todMd\_\{M\}\. It then applies the existing task head to every candidate: the language\-model head produces a teacher\-forced cross\-entropy loss for understanding, while the flow\-to\-VAE head produces a flow\-matching MSE against the same generation target used at the final depth\. Hence the router is not trained from a single early\-exit head or from a sampled candidate; every candidate contributes a task\-cost valueci,mc\_\{i,m\}on every training example\.

The quality constraint is defined relative to the final candidate\. Letδ\\deltabe the configured quality margin and letai,ma\_\{i,m\}denote any task\-specific validity mask \(for example, the label\-correct target used by the understanding router\)\. The candidate is eligible when

si,m=\[ci,m≤ci,M\+δ\]ai,m=1,si,M=1\.s\_\{i,m\}=\\mathbf\{1\}\\\!\\left\[c\_\{i,m\}\\leq c\_\{i,M\}\+\\delta\\right\]a\_\{i,m\}=1,\\qquad s\_\{i,M\}=1\.\(15\)The final\-candidate override guarantees a valid target even when no earlier candidate satisfies the quality constraint\. The default supervised target is the first eligible candidate,

ti=min⁡\{m:si,m=1\},t\_\{i\}=\\min\\\{m:s\_\{i,m\}=1\\\},\(16\)although the implementation also supports selecting the lowest\-cost or best candidate target when that strategy is requested\. The candidate costs, eligibility masks, and teacher\-regret labels are detached before they supervise the router; they are labels for allocation, not an additional gradient path through the router into the backbone\.

The sequential probabilitiesqi,mq\_\{i,m\}in Eq\. \([4](https://arxiv.org/html/2610.22682#S3.E4)\) are the differentiable training surrogate for the first\-halt policy\. The router\-side objective combines the expected detached candidate cost, a normalized expected depth, entropy regularization, and supervised/calibration terms:

ℒrouter=λtask​∑mqi,m​sg⁡\(ci,m\)\+λdepth​∑mqi,m​dmdM−λent​H​\(qi\)\+λsup​ℒsup​\(qi,ti\)\+λcal​ℒcal​\(qi,si\)\+λrisk​ℒrisk\.\\begin\{split\}\\mathcal\{L\}\_\{\\mathrm\{router\}\}=\{\}&\\lambda\_\{\\mathrm\{task\}\}\\sum\_\{m\}q\_\{i,m\}\\,\\operatorname\{sg\}\(c\_\{i,m\}\)\+\\lambda\_\{\\mathrm\{depth\}\}\\sum\_\{m\}q\_\{i,m\}\\frac\{d\_\{m\}\}\{d\_\{M\}\}\-\\lambda\_\{\\mathrm\{ent\}\}\\,H\(q\_\{i\}\)\\\\ &\+\\lambda\_\{\\mathrm\{sup\}\}\\mathcal\{L\}\_\{\\mathrm\{sup\}\}\(q\_\{i\},t\_\{i\}\)\+\\lambda\_\{\\mathrm\{cal\}\}\\mathcal\{L\}\_\{\\mathrm\{cal\}\}\(q\_\{i\},s\_\{i\}\)\+\\lambda\_\{\\mathrm\{risk\}\}\\mathcal\{L\}\_\{\\mathrm\{risk\}\}\.\\end\{split\}\(17\)Heresg\\operatorname\{sg\}denotes stop\-gradient\. The supervised term is a masked binary\-cross\-entropy loss on the sequential halt hazards: candidates beforetit\_\{i\}are trained to continue, candidatetit\_\{i\}to halt, and later candidates are not treated as competing positive exits\. The calibration term aligns the predicted halt behavior with the eligible\-candidate mask\. The optional risk head predicts both a harmful\-regret probabilityr^i,m\\hat\{r\}\_\{i,m\}and a positive uncertainty scaleu^i,m\\hat\{u\}\_\{i,m\}\. Its loss is a binary cross\-entropy term for the task\-specific risk target plus a weighted smooth\-L1L\_\{1\}term that calibrates the uncertainty scale\. For understanding, the harmful\- regret target marks a candidate only when it is wrong while the final\-candidate teacher is correct, so correcting a teacher mistake is not penalized\. For generation, the risk target is built from normalized excess flow\-MSE relative to the final candidate \(together with the configured candidate\-cost strategy\), rather than from label correctness\.

For clarity, the complete training sequence is: \(1\) run to the deepest candidate and collect every candidate hidden state; \(2\) evaluate every native task head and form the per\-candidate task costs; \(3\) construct quality, label\-correctness, and harmful\-regret targets from detached candidate and final candidate predictions; \(4\) formqi,mq\_\{i,m\}and the router objective above; and \(5\) backpropagate the native task loss, multi\-exit task losses, final\-depth hidden\-state distillation, router loss, and TAFE loss through the configured trainable adaptation modules\. The router input features are detached as well, so the router cannot reshape the backbone representation merely to satisfy its allocation labels\.

At inference, the router scans candidates in increasing depth\. It halts at the first candidate satisfying the minimum\-depth condition, the halt probability threshold, and when the risk head is enabled, the upper confidence bound

r^i,m\+κ​u^i,m≤τrisk\.\\hat\{r\}\_\{i,m\}\+\\kappa\\hat\{u\}\_\{i,m\}\\leq\\tau\_\{\\mathrm\{risk\}\}\.\(18\)The final candidate remains an unconditional fallback\. Thusqi,mq\_\{i,m\}is a smooth surrogate used during training, whereas deployment uses the first eligible candidate with fixed validation\-time thresholds; it does not inspect test labels or recompute the quality constraint at inference\.

### A\.6Forward\-pass sequence

For one packed forward pass, the configured computation proceeds as follows:

1. 1\.The runtime identifies the understanding and generation token streams, records their packed positions, and associates each sample with its route kind and task metadata\.
2. 2\.At each decoder layer, the shared attention and normalization run on the packed sequence\. The corresponding TAFE gate is then called on each active stream, and the selected FFN residual is written back to the same packed positions\.
3. 3\.After the layer update, an all\-routeExitcondition may stop the decoder loop\. Otherwise, at a configured candidate depth, the depth controller decides whether the sequence continues\.
4. 4\.The selected hidden state is passed to the backbone’s existing output head: the language\-model head for understanding or the flow\-to\-VAE head for generation\. No independent prediction head is introduced for an intermediate depth\.

### A\.7Metrics

We report three efficiency quantities\.*Relative compute*is a normalized compute proxy,𝔼⁡\[𝒞proxy\]/𝒞max\\mathbb\{E\}\[\\mathcal\{C\}\_\{\\mathrm\{proxy\}\}\]/\\mathcal\{C\}\_\{\\max\}, where the reference is the same backbone run fully shared to full depth\. The proxy has two components: \(i\) layer–token work, which decreases when the depth policy stops a sequence and skips the remaining layers, and \(ii\) the action\-conditioned FFN term, which charges a full shared\-FFN unit toShareand only the selected low\-rank residual\-bank work toDecouple\. In the routing telemetry this action term uses nominal relative costs of1\.01\.0,0\.750\.75, and0\.00\.0forShare,Decouple, andExit, respectively; these weights describe allocation, not hardware\-measured FLOPs\. The latter is important because the dense base FFN is evaluated before the tokenwise action mask in the current wrapper\. Thus the direct and robust runtime saving is the sequence\-level skip; the realized latency benefit of the low\-rank term depends on sparse or fused kernel support and is not measured by this paper’s proxy\. Consequently, the relative values in Table[6](https://arxiv.org/html/2610.22682#S4.T6)should not be read as early\-exit percentages or as wall\-clock speedups: they combine the depth and action\-level accounting terms, while launch, packing, and memory effects are outside the proxy\.*Mean exit depth*d¯=∑ℓ∈𝒟q\(ℓ\)​ℓ\\bar\{d\}=\\sum\_\{\\ell\\in\\mathcal\{D\}\}q^\{\(\\ell\)\}\\ellsummarizes when the model stops computing\.*Specialization rate*α¯=1N​L∑i=1N∑ℓ=1L𝟏\[ai\(ℓ\)=DECOUPLE\]\\bar\{\\alpha\}=\\frac\{1\}\{NL\}\\sum\_\{i=1\}^\{N\}\\sum\_\{\\ell=1\}^\{L\}\\mathbf\{1\}\\big\[a\_\{i\}^\{\(\\ell\)\}=\\mathrm\{DECOUPLE\}\\big\]is the fraction of theN×LN\{\\times\}Ltoken–layer pairs that take theDecoupleaction of[Equation3](https://arxiv.org/html/2610.22682#S3.E3), and summarizes how much capacity specializes\.

## Appendix BFindings

### B\.1Sharing Is Beneficial Early, But Not Uniformly Throughout Depth

In early layers, different functionalities can benefit from processing a common representation\. Shared attention provides a mechanism for information exchange across modalities and functionalities, allowing different pathways to use complementary contextual information\. However, as representations become increasingly task\-specific, forcing all functionalities through the same representational capacity can induce competition\. This is particularly problematic for the FFN, which provides much of the transformer’s representational capacity\.

#### B\.1\.1Controlled \(Synthetic\) Layerwise Competition Diagnostic

This study is designed to answer one mechanistic question: when two functionalities share the same transformer component, do their learning signals update that component in compatible or conflicting directions? We use a synthetic two\-task setting so that the amount of shared structure and task\-private pressure can be controlled directly\. The shared target, private targets, and depth\-increasing private weight are deliberately constructed interventions\. They are useful for isolating a mechanism, but they are not measurements of a released Bagel, Janus\-Pro, Harmon, or Show\-o backbone\.

Each input sequence is associated with two tasks, denoted taskAAand taskBB\. Both tasks depend on a shared targetyshy^\{\\mathrm\{sh\}\}, but each task also has its own private target,yAy^\{A\}oryBy^\{B\}\. The shared target creates a reason for the two tasks to benefit from common computation, while the private targets create task\-specific pressure\. For layerℓ\\ell, we write the task losses as

ℒℓA\\displaystyle\\mathcal\{L\}^\{A\}\_\{\\ell\}=ℒℓsh\+λℓ​ℒℓA,priv,\\displaystyle=\\mathcal\{L\}^\{\\mathrm\{sh\}\}\_\{\\ell\}\+\\lambda\_\{\\ell\}\\mathcal\{L\}^\{A,\\mathrm\{priv\}\}\_\{\\ell\},\(19\)ℒℓB\\displaystyle\\mathcal\{L\}^\{B\}\_\{\\ell\}=ℒℓsh\+λℓ​ℒℓB,priv\.\\displaystyle=\\mathcal\{L\}^\{\\mathrm\{sh\}\}\_\{\\ell\}\+\\lambda\_\{\\ell\}\\mathcal\{L\}^\{B,\\mathrm\{priv\}\}\_\{\\ell\}\.\(20\)Here,ℒℓsh\\mathcal\{L\}^\{\\mathrm\{sh\}\}\_\{\\ell\}is the shared objective at layerℓ\\ell,ℒℓA,priv\\mathcal\{L\}^\{A,\\mathrm\{priv\}\}\_\{\\ell\}andℒℓB,priv\\mathcal\{L\}^\{B,\\mathrm\{priv\}\}\_\{\\ell\}are task\-private objectives, andλℓ\\lambda\_\{\\ell\}controls the strength of the private objective at depthℓ\\ell\. Increasingλℓ\\lambda\_\{\\ell\}with depth lets us isolate how growing task\-specific demands affect shared computation\.

For a transformer componentccat layerℓ\\ell, such as attention, normalization, or FFN, letθℓc\\theta^\{c\}\_\{\\ell\}denote its parameters\. We compute the task\-specific gradients

gℓA,c\\displaystyle g^\{A,c\}\_\{\\ell\}=∇θℓcℒℓA,\\displaystyle=\\nabla\_\{\\theta^\{c\}\_\{\\ell\}\}\\mathcal\{L\}^\{A\}\_\{\\ell\},\(21\)gℓB,c\\displaystyle g^\{B,c\}\_\{\\ell\}=∇θℓcℒℓB\.\\displaystyle=\\nabla\_\{\\theta^\{c\}\_\{\\ell\}\}\\mathcal\{L\}^\{B\}\_\{\\ell\}\.\(22\)The alignment between the two tasks for componentccat layerℓ\\ellis measured by gradient cosine similarity:

cosℓc=⟨gℓA,c,gℓB,c⟩‖gℓA,c‖2​‖gℓB,c‖2\.\\operatorname\{cos\}^\{c\}\_\{\\ell\}=\\frac\{\\langle g^\{A,c\}\_\{\\ell\},g^\{B,c\}\_\{\\ell\}\\rangle\}\{\\\|g^\{A,c\}\_\{\\ell\}\\\|\_\{2\}\\\|g^\{B,c\}\_\{\\ell\}\\\|\_\{2\}\}\.\(23\)Positive cosine means the two tasks update the shared component in compatible directions\. Negative cosine means the tasks compete: improving one task pushes the shared component in a direction that hurts the other\. We report FFN competition as

compℓFFN=max⁡\(0,−cosℓFFN\)\.\\operatorname\{comp\}^\{\\mathrm\{FFN\}\}\_\{\\ell\}=\\max\(0,\-\\operatorname\{cos\}^\{\\mathrm\{FFN\}\}\_\{\\ell\}\)\.\(24\)Thus, a competition value of 0 means no negative alignment, while larger values indicate stronger conflict\.

For a model withLLlayers, the depth\-dependent private weight is

λℓ=0\.15\+2\.5​\(ℓ−1L−1\)1\.35,ℓ∈\{1,…,L\}\.\\lambda\_\{\\ell\}=0\.15\+2\.5\\left\(\\frac\{\\ell\-1\}\{L\-1\}\\right\)^\{1\.35\},\\qquad\\ell\\in\\\{1,\\ldots,L\\\}\.\(25\)Thus,λ1=0\.15\\lambda\_\{1\}=0\.15in the first layer andλL=2\.65\\lambda\_\{L\}=2\.65in the final layer for every model depth\. The three regimes use the sameλℓ\\lambda\_\{\\ell\}schedule and differ by the private\-target amplitudeα\\alpha, which scales the opposing task\-private target before it entersℒℓA,priv\\mathcal\{L\}^\{A,\\mathrm\{priv\}\}\_\{\\ell\}andℒℓB,priv\\mathcal\{L\}^\{B,\\mathrm\{priv\}\}\_\{\\ell\}\. For readability, Table[8](https://arxiv.org/html/2610.22682#A2.T8)also reports the endpoint range of the amplitude\-weighted private pressure,α​λ1→α​λL\\alpha\\lambda\_\{1\}\\to\\alpha\\lambda\_\{L\}\.

Table 8:Private\-pressure settings for the controlled layerwise competition study\.Regimeα\\alphaα​λ1→α​λL\\alpha\\lambda\_\{1\}\\to\\alpha\\lambda\_\{L\}Descriptionweak\-private0\.50\.075→1\.3250\.075\\to 1\.325Half\-amplitude opposing private targetmedium\-private1\.00\.150→2\.6500\.150\\to 2\.650Unit\-amplitude opposing private targetstrong\-private1\.50\.225→3\.9750\.225\\to 3\.9751\.5×\\times\-amplitude opposing private targetWe evaluate three model configurations \(4, 8, and 12 layers\), these three private\-objective regimes, eight random seeds per condition, and 16,384 gradient samples per condition\. Theweak\-privateregime keeps private pressure small, thestrong\-privateregime makes private pressure dominant, andmedium\-privateuses an intermediate setting where the tasks still share context but develop increasingly task\-specific objectives with depth\.

Table 9:Layerwise gradient alignment for the 8\-layer model under themedium\-privateregime\.LayerAttention cosineFFN cosineFFN competition11\.0000\.9550\.00021\.0000\.7960\.00031\.0000\.4390\.00041\.0000\.0340\.00451\.000\-0\.2930\.29361\.000\-0\.5200\.52071\.000\-0\.6630\.66381\.000\-0\.7580\.758Within this controlled family of synthetic conditions, we observe a consistent asymmetric pattern: shared attention maintains strong alignment across depth, while FFN alignment progressively deteriorates and becomes increasingly competitive in deeper layers\. This conditional result is a design hypothesis and stress test, not evidence that the same gradient trajectory must occur in every production backbone\. For example, in the 8\-layermedium\-privatesetting, FFN cosine drops from0\.9550\.955in layer 1 to−0\.758\-0\.758in layer 8, while attention cosine remains1\.0001\.000throughout\.

Table 10:Aggregate layerwise trends from the first to the final layer\.ModelRegimeAttention cosineFFN cosine / competition4L\-8Hmedium\-private1\.000→1\.0001\.000\\to 1\.0000\.954→−0\.7570\.954\\to\-0\.757/0\.000→0\.7570\.000\\to 0\.7574L\-8Hweak\-private1\.000→1\.0001\.000\\to 1\.0000\.988→−0\.3160\.988\\to\-0\.316/0\.000→0\.3160\.000\\to 0\.3164L\-8Hstrong\-private1\.000→1\.0001\.000\\to 1\.0000\.900→−0\.8830\.900\\to\-0\.883/0\.000→0\.8830\.000\\to 0\.8838L\-8Hmedium\-private1\.000→1\.0001\.000\\to 1\.0000\.955→−0\.7580\.955\\to\-0\.758/0\.000→0\.7580\.000\\to 0\.7588L\-8Hweak\-private1\.000→1\.0001\.000\\to 1\.0000\.989→−0\.2490\.989\\to\-0\.249/0\.000→0\.2490\.000\\to 0\.2498L\-8Hstrong\-private1\.000→1\.0001\.000\\to 1\.0000\.906→−0\.8780\.906\\to\-0\.878/0\.000→0\.8780\.000\\to 0\.87812L\-16Hmedium\-private1\.000→1\.0001\.000\\to 1\.0000\.956→−0\.7510\.956\\to\-0\.751/0\.000→0\.7510\.000\\to 0\.75112L\-16Hweak\-private1\.000→1\.0001\.000\\to 1\.0000\.989→−0\.2870\.989\\to\-0\.287/0\.000→0\.2870\.000\\to 0\.28712L\-16Hstrong\-private1\.000→1\.0001\.000\\to 1\.0000\.903→−0\.8830\.903\\to\-0\.883/0\.000→0\.8830\.000\\to 0\.883
#### B\.1\.2Shared Normalization Ablation

We further isolate the role of shared normalization using the same controlled two\-task setting\. This ablation separates three possible sources of alignment: shared attention, shared normalization, and shared FFN capacity\. The objective is to determine whether positive cross\-task alignment comes only from attention, or whether normalization can also provide a stabilizing shared pathway\.

The notation is the same as in the layerwise study\. “Shared Attn\.” indicates whether taskAAand taskBBuse the same attention parametersθℓattn\\theta^\{\\mathrm\{attn\}\}\_\{\\ell\}\. “Shared Norm” indicates whether the tasks use the same normalization parametersθℓnorm\\theta^\{\\mathrm\{norm\}\}\_\{\\ell\}\. When a component is private, each task has its own copy of the parameters, so there is no shared component gradient cosine to report for that component\. This is marked as N/A in the table\.

The FFN objective in this ablation is deliberately constructed to oppose the two tasks, so the FFN cosine is expected to be negative\. This controlled design lets us ask whether shared attention and shared normalization can remain aligned even when the capacity\-bearing FFN objective is maximally competitive\.

Table 11:Shared normalization ablation in the controlled two\-task setting\.ConditionShared Attn\.Shared NormAttn\. cosineNorm cosineFFN cosineShared Attn\. \+ Shared Norm✓\\checkmark✓\\checkmark0\.989→0\.9890\.989\\to 0\.9891\.000→1\.0001\.000\\to 1\.000−1\.000→−1\.000\-1\.000\\to\-1\.000Shared Attn\. \+ Private Norm✓\\checkmark×\\times0\.989→0\.9890\.989\\to 0\.989N/A−1\.000→−1\.000\-1\.000\\to\-1\.000Private Attn\. \+ Shared Norm×\\times✓\\checkmarkN/A1\.000→1\.0001\.000\\to 1\.000−1\.000→−1\.000\-1\.000\\to\-1\.000Private Attn\. \+ Private Norm×\\times×\\timesN/AN/A−1\.000→−1\.000\-1\.000\\to\-1\.000

The results reveal an important asymmetry in unified computation: communication and capacity do not have the same optimal sharing pattern\. The key observation is not simply that competition increases with depth, but that different components exhibit fundamentally different sharing profiles\. Attention remains aligned while FFN capacity becomes increasingly task\-specific\. This motivates the first design principle of FlexUMM: preserve shared communication while selectively specializing task\-specific capacity\.

##### Evidence boundary\.

The paper therefore separates the mechanistic diagnostic from the real\-model evidence\. The gradient cosines above are synthetic and human\-designed; they show that the proposed sharing asymmetry can arise from increasing private pressure, but they do not establish its magnitude in a production backbone\. The empirical tests on Bagel and Janus\-Pro \(and the additional Harmon benchmark\) measure task quality, routing behavior, and compute proxies\. The separate Show\-o2 token\-importance experiment is a real Show\-o\-family analysis of functionality\-specific importance, not a direct layerwise gradient\-cosine measurement\. These real\-model results test whether the resulting flexible allocation policy is useful; they should not be read as a direct validation of the exact synthetic gradient values\. Direct component\-gradient measurements on each released backbone would be a separate experiment\.

#### B\.1\.3Layerwise Token\-Importance Sharing Across Functionalities

We further examine whether a fixed degree of sharing between functionalities is appropriate throughout computation in a real UMM\. We use Show\-o2 and estimate token importance separately for understanding and generation, then compare the resulting functionality\-specific importance patterns on the same held\-out inputs\.

##### Token\-importance estimation\.

For tokeniiat layerℓ\\ell, we define the layer\-induced representation update as

Δ​hℓ,i=hℓ,iout−hℓ,iin\.\\Delta h\_\{\\ell,i\}=h^\{\\mathrm\{out\}\}\_\{\\ell,i\}\-h^\{\\mathrm\{in\}\}\_\{\\ell,i\}\.\(26\)Following first\-order Taylor importance estimation\[[Molchanov et al\., 2019](https://arxiv.org/html/2610.22682#bib.bib38)\], we assign each token the score

sℓ,i=\|∇hℓ,ioutℒ⊤​Δ​hℓ,i\|,s\_\{\\ell,i\}=\\left\|\\nabla\_\{h^\{\\mathrm\{out\}\}\_\{\\ell,i\}\}\\mathcal\{L\}^\{\\top\}\\Delta h\_\{\\ell,i\}\\right\|,\(27\)which measures the alignment between the layer update and the gradient of the task loss\.

For understanding,ℒ\\mathcal\{L\}is the cross\-entropy loss on the expected answer tokens under teacher forcing\. For generation,ℒ\\mathcal\{L\}is the flow\-matching MSE loss computed along the generation trajectory using a Show\-o2\-generated reference latent\. Show\-o2 remains frozen when collecting these scores\.

##### Functionality\-specific predictors\.

We train two lightweight predictors, denoted MLP\-U and MLP\-G, to estimate the log Gradient×\\timesUpdate score from the token hidden state and layer identity\. MLP\-U is trained only on understanding samples, while MLP\-G is trained only on generation samples\.

The understanding set contains four task categories: counting, spatial relations, attribute combinations, and visual\-state recognition, with six samples per category\. For each category, four samples are used for training, one for validation, and one for testing\. The generation set contains 12 samples, split into 8 for training, 2 for validation, and 2 for testing, with all 49 generation steps recorded\. Each predictor is trained with two random seeds, and the evaluation results are averaged across seeds\.

As a basic check that the predictors capture their corresponding importance signals, we compare their predicted top\-20% tokens with those selected directly by the Gradient×\\timesUpdate scores\. MLP\-U reaches64\.8%64\.8\\%overlap on understanding test samples, while MLP\-G reaches34\.3%34\.3\\%on generation test samples\. Cross\-functionality transfer is lower, with23\.9%23\.9\\%for MLP\-U on generation and17\.9%17\.9\\%for MLP\-G on understanding\. For independent random top\-20% selections, the expected normalized overlap is approximately 20%, since each selected token has a 20% chance of appearing in the other set\.

##### Layerwise sharing\.

To compare the two functionality\-specific importance patterns under the same input, we apply MLP\-U and MLP\-G to identical held\-out samples\. At each layer, each predictor selects its topKKtokens, whereK=⌈0\.2​N⌉K=\\lceil 0\.2N\\rceilandNNis the number of valid tokens\. We compute

Overlapℓ=\|𝒯ℓU∩𝒯ℓG\|K,\\mathrm\{Overlap\}\_\{\\ell\}=\\frac\{\\left\|\\mathcal\{T\}^\{U\}\_\{\\ell\}\\cap\\mathcal\{T\}^\{G\}\_\{\\ell\}\\right\|\}\{K\},\(28\)where𝒯ℓU\\mathcal\{T\}^\{U\}\_\{\\ell\}and𝒯ℓG\\mathcal\{T\}^\{G\}\_\{\\ell\}are the top\-KKsets predicted by MLP\-U and MLP\-G, respectively\.

Figure 6:Layerwise top\-20% token overlap between MLP\-U and MLP\-G predictions on held\-out Show\-o2 understanding and generation samples\. Generation results are averaged over generation steps\.Figure[6](https://arxiv.org/html/2610.22682#A2.F6)shows a depth\-dependent pattern across both functionalities\. On understanding inputs, the overlap is generally above the random baseline in earlier layers and becomes lower in deeper layers\. On generation inputs, the trend is more pronounced: overlap is substantially above random in early layers and progressively decreases with depth, falling below the random baseline in the final layers\. These results are consistent with stronger sharing in earlier computation and increasing functionality\-specific specialization at greater depth\.

##### Variation across generation steps\.

We further examine whether the relationship between the two functionality\-specific predictors changes over the generation trajectory\. At each generation step, we average their top\-20% token overlap across layers\.

Figure 7:Top\-20% token overlap between MLP\-U and MLP\-G predictions across generation steps on held\-out Show\-o2 generation samples, averaged over layers\.As shown in Figure[7](https://arxiv.org/html/2610.22682#A2.F7), the overlap generally decreases over the generation trajectory and approaches the random baseline at later steps\. The final step is an exception to this overall trend\. Together, the layerwise and generation\-step analyses suggest that a fixed sharing pattern may be too restrictive, as functionality\-specific token\-importance patterns vary across both model depth and generation progress\.

### B\.2Computation Depth Should Be Input\-Dependent

Even after accounting for parameter competition, a second limitation remains: UMMs generally require every input to traverse the same number of layers\. Different inputs can reach a sufficiently mature representation at different depths\. A simple understanding query may require relatively little computation, whereas a difficult reasoning problem may benefit from deeper processing\. Generation inputs can also differ substantially in the amount of refinement they require\.

#### B\.2\.1Bagel Selected\-Depth Distribution on MME

We first investigate this question using Bagel with a learned FlexUMM routing policy on the full 2,374\-question MME benchmark\. The goal is to test whether inputs within the same model and the same functionality prefer different amounts of computation\.

Letxix\_\{i\}denote theii\-th MME input\. Bagel contains 28 language layers, and the routing policy evaluates a candidate exit\-depth set

𝒟=\{22,23,25,26,27,28\}\.\\mathcal\{D\}=\\\{22,23,25,26,27,28\\\}\.\(29\)For each inputxix\_\{i\}, the router selects one depth

di=rϕ​\(xi\),di∈𝒟,d\_\{i\}=r\_\{\\phi\}\(x\_\{i\}\),\\qquad d\_\{i\}\\in\\mathcal\{D\},\(30\)whererϕr\_\{\\phi\}is the learned routing policy anddid\_\{i\}is the selected number of language layers used for that input\. A smallerdid\_\{i\}means the model exits earlier and saves computation; a largerdid\_\{i\}means the model continues deeper for additional refinement\.

Table[12](https://arxiv.org/html/2610.22682#A2.T12)reports how often each depth is selected\. “MME samples” is the number of benchmark examples assigned to a given depth, and “Percentage” is this count divided by the full 2,374\-question benchmark size\.

Table 12:Selected\-depth distribution for Bagel on MME\.Selected depthMME samplesPercentage2229612\.47%23180\.76%2556423\.76%261,14648\.27%27522\.19%2829812\.55%Total2,374100\.00%The selected depth ranges from 22 to 28 layers, with a mean selected depth of 25\.514, a minimum of 22, a maximum of 28, and six distinct selected depths\. Nearly half of the examples select depth 26, but substantial mass remains at both extremes: 12\.47% exit at depth 22 and 12\.55% run the full 28 layers\. This variation occurs within a single model and a single benchmark, rather than simply reflecting differences between understanding and generation tasks; it is genuine input\-level variation in how much computation an example requires\. The result provides real\-model evidence that computation depth should not necessarily be treated as a fixed architectural constant\.

#### B\.2\.2Controlled Progressive\-Depth Study

To isolate the underlying phenomenon, we construct a controlled progressive\-depth setting\. The objective is to test whether variable optimal depth can arise even when the data\-generating process is known and controlled, rather than being an artifact of a particular real benchmark\.

Each samplexix\_\{i\}has a binary labelyi∈\{0,1\}y\_\{i\}\\in\\\{0,1\\\}and a hidden required depthLi∗L\_\{i\}^\{\*\}\. The required depth represents the first layer at which the sample contains enough evidence for a reliable prediction\. For a model withLLlayers, each intermediate layerℓ∈\{1,…,L\}\\ell\\in\\\{1,\\ldots,L\\\}produces a prediction probability

pi,ℓ=fℓ​\(xi\),p\_\{i,\\ell\}=f\_\{\\ell\}\(x\_\{i\}\),\(31\)wherefℓf\_\{\\ell\}denotes the model’s prediction head applied at depthℓ\\ell\.

The synthetic evidence is organized in three phases:

- •Before the required depth,ℓ<Li∗\\ell<L\_\{i\}^\{\*\}, the sample provides only weak evidence, so early predictions are uncertain\.
- •At the required depth,ℓ=Li∗\\ell=L\_\{i\}^\{\*\}, the sample provides a strong identifying signal for the label\.
- •After the required depth,ℓ\>Li∗\\ell\>L\_\{i\}^\{\*\}, additional layers receive zero\-mean refinement noise, so deeper computation is not guaranteed to improve the prediction\.

For each intermediate prediction, we compute binary cross\-entropy

BCE⁡\(yi,pi,ℓ\)=−yi​log⁡pi,ℓ−\(1−yi\)​log⁡\(1−pi,ℓ\)\.\\operatorname\{BCE\}\(y\_\{i\},p\_\{i,\\ell\}\)=\-y\_\{i\}\\log p\_\{i,\\ell\}\-\(1\-y\_\{i\}\)\\log\(1\-p\_\{i,\\ell\}\)\.\(32\)An oracle chooses the best exit depth for each sample by minimizing this loss:

d^i=arg⁡minℓ∈\{1,…,L\}⁡BCE⁡\(yi,pi,ℓ\)\.\\hat\{d\}\_\{i\}=\\arg\\min\_\{\\ell\\in\\\{1,\\ldots,L\\\}\}\\operatorname\{BCE\}\(y\_\{i\},p\_\{i,\\ell\}\)\.\(33\)Here,d^i\\hat\{d\}\_\{i\}is the empirically best depth for sampleii, which may differ from the hidden required depthLi∗L\_\{i\}^\{\*\}because the model’s intermediate predictions can be affected by noise and imperfect learning\.

We compare this oracle variable\-depth policy with fixed full\-depth computation\. LetNNbe the number of samples\. The full\-depth loss is

ℒfull=1N​∑i=1NBCE⁡\(yi,pi,L\),\\mathcal\{L\}\_\{\\mathrm\{full\}\}=\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}\\operatorname\{BCE\}\(y\_\{i\},p\_\{i,L\}\),\(34\)and the oracle variable\-depth loss is

ℒoracle=1N​∑i=1NBCE⁡\(yi,pi,d^i\)\.\\mathcal\{L\}\_\{\\mathrm\{oracle\}\}=\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}\\operatorname\{BCE\}\(y\_\{i\},p\_\{i,\\hat\{d\}\_\{i\}\}\)\.\(35\)We report oracle gain as

Gain=ℒfull−ℒoracle\.\\operatorname\{Gain\}=\\mathcal\{L\}\_\{\\mathrm\{full\}\}\-\\mathcal\{L\}\_\{\\mathrm\{oracle\}\}\.\(36\)A positive gain means that choosing the best depth per sample improves prediction quality relative to always using the final layer\.

We evaluate 4\-, 8\-, and 12\-layer models; three difficulty distributions; 32,768 samples per condition; and nine total model/difficulty conditions\. The uniform difficulty distribution spreads required depths across the full model, the early\-heavy distribution places more mass on shallower required depths, and the late\-heavy distribution places more mass on deeper required depths\.

In Table[13](https://arxiv.org/html/2610.22682#A2.T13), “Unique best depths” counts how many distinct values appear among the oracle\-selected depths\{d^i\}i=1N\\\{\\hat\{d\}\_\{i\}\\\}\_\{i=1\}^\{N\}\. “Best\-depth mean±\\pmstd” reports the mean and standard deviation of the oracle\-selected depth distribution\. “Oracle gain” is the BCE improvement of the oracle variable\-depth policy over fixed full\-depth computation\.

Table 13:Controlled progressive\-depth results\. Oracle gain is measured relative to fixed full\-depth computation in BCE\.ModelDifficultyUnique best depthsBest\-depth mean±\\pmstdOracle gain4LUniform43\.24±0\.983\.24\\pm 0\.980\.0554LEarly\-heavy42\.97±1\.122\.97\\pm 1\.120\.0544LLate\-heavy43\.53±0\.823\.53\\pm 0\.820\.0458LUniform86\.23±2\.006\.23\\pm 2\.000\.1068LEarly\-heavy85\.64±2\.345\.64\\pm 2\.340\.1028LLate\-heavy86\.89±1\.706\.89\\pm 1\.700\.08812LUniform129\.24±2\.999\.24\\pm 2\.990\.14312LEarly\-heavy128\.33±3\.548\.33\\pm 3\.540\.14112LLate\-heavy1210\.23±2\.5810\.23\\pm 2\.580\.121Every condition contains multiple distinct optimal depths\. For the 12\-layer uniform\-difficulty model, all 12 candidate depths are selected by at least some samples andStd⁡\(L∗\)=2\.99\\mathrm\{Std\}\(L^\{\*\}\)=2\.99\. The oracle variable\-depth policy consistently improves over fixed full\-depth computation, with gains ranging from 0\.045 to 0\.143 BCE\.

Across model depths and difficulty regimes, different samples have different optimal computation depths, and adapting depth can improve both computation efficiency and prediction quality\. This motivates the second design principle of FlexUMM: computation depth should be determined by the maturity of each input rather than fixed globally\.

## Appendix CAdditional Ablation and Policy Analyses

### C\.1Remaining Design Choices

##### Discrete versus graded specialization\.

[Equation2](https://arxiv.org/html/2610.22682#S3.E2)switches discretely between shared and specialized computation rather than interpolating between them\.[Table14](https://arxiv.org/html/2610.22682#A3.T14)compares this discrete gate with a graded alternative that mixes the shared and specialized outputs by a continuous weight, and also ablates the action\-mixture target insideℒgate\\mathcal\{L\}\_\{\\mathrm\{gate\}\}, the expert countEE, and the routing widthkk\. Removing the action\-mixture target tests the failure mode it is meant to prevent: gate saturation, where the shared expert is abandoned and the model degenerates into a fully specialized backbone that has lost cross\-functional transfer\.

##### Expert count and routing width\.

Widening the pool toE=6E=6experts edges out the deployedE=4,k=2E=4,k=2configuration on every benchmark in[Table14](https://arxiv.org/html/2610.22682#A3.T14), but FlexUMM is designed to be a lightweight adapter, andE=6E=6adds proportionally more trained parameters thanE=4E=4’s0\.381%0\.381\\%overhead \([Table7](https://arxiv.org/html/2610.22682#A1.T7)\) for a marginal accuracy gain\. We keepE=4E=4as the deployed setting and read theE=6E=6result as headroom rather than a reason to grow the added parameter budget\. Shrinking the pool toE=2,k=2E=2,k=2instead removes the shared\-versus\-specialized choice at each routed layer, which costs more than widening does, confirming that a nontrivial expert pool is needed for the router to specialize meaningfully\.

Table 14:Capacity\-axis design ablations on Bagel\. “Rel\. compute” is the normalized route/action proxy defined in[SectionA\.7](https://arxiv.org/html/2610.22682#A1.SS7)\.VariantRel\. computeMMMUMMBench\-CNMMBench\-ENMathVistaMMStarGraded mixing instead of discrete gate0\.98754\.882\.485\.270\.269\.6No action\-mixture target inℒgate\\mathcal\{L\}\_\{\\mathrm\{gate\}\}0\.98855\.282\.084\.072\.270\.4E=4E=4,k=1k=10\.95854\.085\.8786\.7573\.470\.93E=2,k=2E=2,k=20\.97854\.482\.084\.872\.470\.2E=6,k=2E=6,k=20\.95155\.286\.687\.173\.674\.4

### C\.2Early\-Exit Subset Audit

The main paper already reports the overall selected\-depth distribution, so we do not repeat that plot here\. Instead, we analyze the rare early\-exit cases in the same four\-GPU lightweight screen\. The screen produced 16 raw routes with a selected depth below 28\. Four\-way circular choice evaluation creates duplicate MMBench routes, leaving 10 unique inputs after deduplication: three MMMU inputs, five MathVista inputs, and two MMBench\-English inputs\.

Every unique early case selected depth 26; no case selected depth 22 or 24\. This is a targeted audit of the early\-exit tail, not a representative estimate of the full benchmark distribution\. The evaluators do not provide a common per\-example correctness field across all three suites, so the analysis focuses on the input and router signals rather than claiming early\-exit accuracy\.

#### C\.2\.1What receives an early exit?

Table[15](https://arxiv.org/html/2610.22682#A3.T15)lists every unique early case\. The subset is concentrated in a few visually grounded or short\-answer task families: MathVista icon\-style questions account for four of the five MathVista cases, while the MMMU cases come from biology and chemistry\. Question length is not uniform: the subset contains very short 21\-character questions as well as a 726\-character biology question\. Thus the observed early exits cannot be explained by a simple short\-question rule\.

Table 15:Unique early\-exit cases in the lightweight screen\. MMBench circular\-choice copies are deduplicated\. The router selected depth 26 in every case; the signals are recorded at the selected depth\.SuiteCaseTaskChars\.DepthHalt@26Risk@26Regret@26MMBench\-EN493MMBench\-English66260\.5810\.0910\.038MMBench\-EN630MMBench\-English41260\.6400\.0810\.030MMMUBiology\_21biology218260\.5110\.0980\.045MMMUBiology\_8biology726260\.5060\.0880\.035MMMUChemistry\_30chemistry257260\.5050\.0900\.037MathVista129iconqa\_integer104260\.5230\.0870\.034MathVista157iconqa\_integer21260\.5240\.0970\.044MathVista163iconqa\_integer21260\.5350\.0960\.043MathVista187iconqa\_integer21260\.5350\.0970\.043MathVista92vqa2\_0\_text60260\.6020\.0930\.040Figure[8](https://arxiv.org/html/2610.22682#A3.F8)summarizes the subset composition and plots question length against the halt probability at the selected depth\. All points lie at or above the router’s 0\.5 threshold, while their marker size encodes the residual\-risk upper bound\. The small sample and repeated task families make this an audit of behavior, not a statistical generalization\.

Figure 8:Early\-exit subset composition and router signal\. Each point is a unique input selected at depth 26; marker size is proportional to the residual\-risk upper bound\.
#### C\.2\.2Router traces around the early decision

Figure[9](https://arxiv.org/html/2610.22682#A3.F9)follows the complete router trace for each unique early case\. The individual traces show a consistent pattern: the risk upper bound falls from depths 22 to 26, while halt probability rises and crosses the 0\.5 decision threshold at depth 26\. The router therefore exits when its own risk signal has decreased sufficiently, rather than because the input was assigned a fixed shallow depth in advance\. Since these traces stop once the policy exits, there is no depth\-28 point for an early case\.

Figure 9:Candidate\-depth router traces for the unique early\-exit subset\. Faint lines are individual cases and the black line is their mean at each evaluated candidate depth\.
#### C\.2\.3Why the gate passes for early cases

The implementation requires both a halt probability of at least0\.50\.5and a risk upper bound of at most0\.10\.1before exiting\. We compare these two gate signals at depth 26 for every deduplicated input in the screen\. As shown in Figure[10](https://arxiv.org/html/2610.22682#A3.F10)and Table[16](https://arxiv.org/html/2610.22682#A3.T16), all 10 early cases satisfy both conditions, whereas none of the 740 full\-depth cases does\. Only 5\.0% of full\-depth cases cross the halt threshold alone, and their risk signal remains too high\. This provides a more mechanistic explanation for the rare early exits than question length or task identity by itself\.

Table 16:Depth\-26 gate\-condition comparison\. “Halt pass” uses the configured halt threshold 0\.5; “Joint pass” requires both halt probability≥0\.5\\geq 0\.5and risk upper bound≤0\.1\\leq 0\.1\.GroupNNMean halt@26Mean risk@26Halt passJoint passearly100\.5460\.092100\.0%100\.0%full7400\.4240\.1325\.0%0\.0%Figure 10:Depth\-26 gate signals for early\-exit and full\-depth inputs\. The dashed lines show the configured halt\-probability and risk thresholds\.##### Interpretation\.

The early\-exit subset provides concrete evidence that the learned policy is active on individual understanding inputs, but it also clarifies the current checkpoint’s operating point: early exits are rare and all occur at depth 26\. The result supports input\-conditioned routing as a mechanism while motivating future calibration or threshold studies if broader compute savings are desired\.

## Appendix DExtended Case Study

This appendix gives the full selected qualitative set referenced in Section[4\.8](https://arxiv.org/html/2610.22682#S4.SS8)\. The examples are strict paired wins: the FlexUMM answer or image is correct under the benchmark evaluator and the baseline comparator is not\. They are selected for auditability and diversity, and are not intended to replace aggregate benchmark results\.

### D\.1Understanding cases

Table[17](https://arxiv.org/html/2610.22682#A4.T17)lists all ten paired understanding examples\. The MMBench entries are representative displayed rows from the four\-way circular\-group evaluation used in the paired output analysis\.

Table 17:Complete selected understanding set\. Each row includes the benchmark input image and the paired answers\. The FlexUMM answer agrees with the ground truth and the baseline answer does not\.\#Input imageFlexUMMBaselineGT1![[Uncaptioned image]](https://arxiv.org/html/2610.22682v1/figures/case_study/understanding_01_mmmu_music.png)ABA2![[Uncaptioned image]](https://arxiv.org/html/2610.22682v1/figures/case_study/understanding_02_mmmu_clinical.png)DBD3![[Uncaptioned image]](https://arxiv.org/html/2610.22682v1/figures/case_study/understanding_03_mathvista_28.png)8\.57\.58\.54![[Uncaptioned image]](https://arxiv.org/html/2610.22682v1/figures/case_study/understanding_04_mathvista_270.png)100∘100^\{\\circ\}130∘130^\{\\circ\}100∘100^\{\\circ\}5![[Uncaptioned image]](https://arxiv.org/html/2610.22682v1/figures/case_study/understanding_05_mmbench_en_1443.png)CBC6![[Uncaptioned image]](https://arxiv.org/html/2610.22682v1/figures/case_study/understanding_06_mmbench_en_596.png)ABA7![[Uncaptioned image]](https://arxiv.org/html/2610.22682v1/figures/case_study/understanding_07_mmbench_cn_350.png)BAB8![[Uncaptioned image]](https://arxiv.org/html/2610.22682v1/figures/case_study/understanding_08_mmbench_cn_3560.png)BCB9![[Uncaptioned image]](https://arxiv.org/html/2610.22682v1/figures/case_study/understanding_09_mmstar_352.png)DCD10![[Uncaptioned image]](https://arxiv.org/html/2610.22682v1/figures/case_study/understanding_10_mmstar_276.png)BCB
### D\.2Generation cases

Figure[11](https://arxiv.org/html/2610.22682#A4.F11)shows six distinct selected generation wins spanning color, object identity, and spatial\-relation constraints\. The set includes one sample for each of the repeated red\-dog and white\-scissors prompts, while omitting the purple\-bear samples\. Within each tile, the FlexUMM image is on the left and the baseline comparator is on the right\. The pairs were generated with identical prompts and evaluation settings, and the detector outcome is summarized in Table[18](https://arxiv.org/html/2610.22682#A4.T18)\.

![Refer to caption](https://arxiv.org/html/2610.22682v1/figures/case_study/generation_contact_sheet.png)Figure 11:Selected generation set\. Each tile contains the FlexUMM image followed by the paired baseline image\.Table 18:Detector outcomes for the selected generation set\. The rows are individual image\-level wins from distinct prompts\.\#PromptFlexUMM resultBaseline failure mode1purple pizzaPassNo purple pizza; red pizza detected2red dogPassBrown dog rather than red3white scissorsPassPurple scissors rather than white4purple backpackPassSuitcase/handbag detected, not backpack5chair left of a zebraPassBench detected; no chair6parking meter above a broccoliPassParking meter below the broccoli

## Appendix ELimitations

Our evaluation considers a finite set of candidate depths for each backbone, which limits the granularity of the available quality\-compute operating points\. In addition, our experiments focus on a limited set of unified multimodal backbones and model scales\. Extending the evaluation to finer\-grained depth configurations and a broader range of architectures could further characterize the flexibility and applicability of the proposed framework\.