FlexUMM: Flexible Representation Learning for Unified Multimodal Models

Yinyi Luo; Haoyue Bai; Yi Ge; Nilton Cesar Rojas Vales; Marios Savvides; Jindong Wang

cs.AIOct 15, 2026

FlexUMM: Flexible Representation Learning for Unified Multimodal Models

Authors: Yinyi Luo, Haoyue Bai, Yi Ge, Nilton Cesar Rojas Vales, Marios Savvides, Jindong Wang

Organizations: Carnegie Mellon University · William & Mary · University 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, functional plug-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.

Figures & tables

Panel of Figure 1 — Figure 1: FlexUMM adapts both who computes and how much computation an input receives.

Panel of Figure 2 — Figure 2: Input-adaptive depth and quality–efficiency behavior.

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 selects Share , Decouple , or Exit for each token at each routed FFN. A separate candidate-depth router decides, at the sequence level, whether the remaining decoder layers are still needed.

Model	Params	Understanding					Generation
Model	Params	MMMU	MMBench-CN	MMBench-EN	MathVista	MMStar	GenEval	WISE
Bagel ( Deng et al., 2025 )	14B	51.9	80.94	82.65	71.3	63.20	78.81	0.4044
Janus-Pro ( Chen et al., 2025 )	7B	40.9	69.93	64.84	43.8	47.93	79.02	0.3900
Janus ( Wu et al., 2025a )	1.3B	27.3	37.19	27.69	26.6	26.67	40.04	0.2222
Janus Flow ( Ma et al., 2025 )	1.3B	29.0	60.31	64.16	34.8	39.13	49.99	0.2964
Show-O2 ( Xie et al., 2026 )	7B	47.9	76.68	77.25	51.5	54.53	59.87	0.3595

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 in bold .

Setting	$\eta_{\beta}$	$\bar{d}$	Rel. compute	MMMU	MMStar	MathVista
Aggressive	0.30	26.072	0.930	52.9	71.2	69.1
Balanced	0.50	27.841	0.993	54.0	70.9	73.4
Conservative	0.80	28.000	1.000	53.6	70.8	73.1

Table 2: Effect of the halting threshold

\eta_{\beta}

on Bagel.

Panel of Figure 4 — Figure 4: Quality-efficiency trade-off and adaptive exit distribution.

Variant	Understanding					Generation
Variant	MMMU	MMBench-CN	MMBench-EN	MathVista	MMStar	GenEval	WISE
Bagel	51.9	80.94	82.65	71.3	63.20	78.81	0.4044
Bagel+FlexUMM	53.6	85.87	86.81	73.1	70.80	82.03	0.4292
Capacity axis: which components are routed
Routed attention, shared FFN	52.1	82.26	81.95	71.2	62.30	81.21	0.4028
Routed attention, routed FFN	53.2	84.81	84.24	71.9	68.90	80.65	0.3995

Table 3: Ablation study on Bagel.

Method	Unison					Unified-Bench
Method	Int. Cons.	U.-G. Gen.	G.-G. Und.	Mut. Enh.	Overall	CLIP	DINOv2	DINOv3	LongCLIP	Overall
Bagel	80.72	72.09	33.68	35.87	55.59	89.47	78.77	72.40	93.21	83.46
+FlexUMM	83.21 $\uparrow$ 3.1%	72.96 $\uparrow$ 1.2%	37.08 $\uparrow$ 10.1%	46.08 $\uparrow$ 28.5%	59.83 $\uparrow$ 7.6%	90.91 $\uparrow$ 1.6%	79.88 $\uparrow$ 1.4%	74.53 $\uparrow$ 2.9%	94.87 $\uparrow$ 1.8%	85.05 $\uparrow$ 1.9%

Table 4: Unification test on Unison ( Liu et al., 2026 ) and Unified-Bench ( Yan et al., 2026 ) .

Model	C-Unsafe $\downarrow$	C-Reject $\uparrow$	I-Unsafe $\downarrow$	Both-Safe $\uparrow$
Bagel	43.72%	5.07%	49.20%	43.96%
+FlexUMM	42.35%	5.23%	49.06%	44.86%

Table 5: Uni-SafeBench safety results.

Backbone	Method	Understanding					Generation		Efficiency
Backbone	Method	MMMU	MMBench-CN	MMBench-EN	MathVista	MMStar	GenEval	WISE	Rel. compute
Bagel	Full backbone	51.9	80.94	82.65	71.3	63.20	78.81	0.4044	1.00
Bagel	FlexUMM	54.0 $\uparrow$ 4.0%	85.87 $\uparrow$ 6.1%	86.75 $\uparrow$ 5.0%	73.4 $\uparrow$ 2.9%	70.93 $\uparrow$ 12.2%	82.03 $\uparrow$ 4.1%	0.4292 $\uparrow$ 6.1%	0.96
Janus-Pro	Full backbone	40.89	64.84	69.93	43.8	47.93	79.02	0.3900	1.00
Janus-Pro	FlexUMM	46.11 $\uparrow$ 12.8%	68.47 $\uparrow$ 5.6%	73.11 $\uparrow$ 4.5%	51.0 $\uparrow$ 16.4%	52.33 $\uparrow$ 9.2%	81.03 $\uparrow$ 2.5%	0.4095 $\uparrow$ 5.0%	0.96
Harmon	Full backbone	35.0	39.43	48.94	22.0	36.53	67.81	0.3917	1.00

Table 6: Generalization across backbones with different sharing structures.

Panel of Figure 5 — Figure 5: Qualitative depth-selection progression for one generation prompt.

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Backbone	Params	Lang. layers	Candidate depths $\mathcal{D}$	Experts $E$	Trained params
Bagel	14.607B	28	$\{22,24,26,28\}$	2	0.056B (+0.381%)
Janus-Pro	7.420B	30	$\{21,24,27,30\}$	2	0.035B (+0.472%)
Harmon	1.5B	28	$\{22,24,26,28\}$	2	0.013B (+0.837%)

Appendix

Table 7: Backbone configurations for FlexUMM.

Regime	$\alpha$	$\alpha\lambda_{1}\to\alpha\lambda_{L}$	Description
weak-private	0.5	$0.075\to 1.325$	Half-amplitude opposing private target
medium-private	1.0	$0.150\to 2.650$	Unit-amplitude opposing private target
strong-private	1.5	$0.225\to 3.975$	1.5 $\times$ -amplitude opposing private target

Appendix

Table 8: Private-pressure settings for the controlled layerwise competition study.

Layer	Attention cosine	FFN cosine	FFN competition
1	1.000	0.955	0.000
2	1.000	0.796	0.000
3	1.000	0.439	0.000
4	1.000	0.034	0.004
5	1.000	-0.293	0.293
6	1.000	-0.520	0.520

Appendix

Table 9: Layerwise gradient alignment for the 8-layer model under the medium-private regime.

Model	Regime	Attention cosine	FFN cosine / competition
4L-8H	medium-private	$1.000\to 1.000$	$0.954\to-0.757$ / $0.000\to 0.757$
4L-8H	weak-private	$1.000\to 1.000$	$0.988\to-0.316$ / $0.000\to 0.316$
4L-8H	strong-private	$1.000\to 1.000$	$0.900\to-0.883$ / $0.000\to 0.883$
8L-8H	medium-private	$1.000\to 1.000$	$0.955\to-0.758$ / $0.000\to 0.758$
8L-8H	weak-private	$1.000\to 1.000$	$0.989\to-0.249$ / $0.000\to 0.249$
8L-8H	strong-private	$1.000\to 1.000$	$0.906\to-0.878$ / $0.000\to 0.878$

Appendix

Table 10: Aggregate layerwise trends from the first to the final layer.

Condition	Shared Attn.	Shared Norm	Attn. cosine	Norm cosine	FFN cosine
Shared Attn. + Shared Norm	$\checkmark$	$\checkmark$	$0.989\to 0.989$	$1.000\to 1.000$	$-1.000\to-1.000$
Shared Attn. + Private Norm	$\checkmark$	$\times$	$0.989\to 0.989$	N/A	$-1.000\to-1.000$
Private Attn. + Shared Norm	$\times$	$\checkmark$	N/A	$1.000\to 1.000$	$-1.000\to-1.000$
Private Attn. + Private Norm	$\times$	$\times$	N/A	N/A	$-1.000\to-1.000$

Appendix

Table 11: Shared normalization ablation in the controlled two-task setting.

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 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.

Selected depth	MME samples	Percentage
22	296	12.47%
23	18	0.76%
25	564	23.76%
26	1,146	48.27%
27	52	2.19%
28	298	12.55%

Appendix

Table 12: Selected-depth distribution for Bagel on MME.

Model	Difficulty	Unique best depths	Best-depth mean $\pm$ std	Oracle gain
4L	Uniform	4	$3.24\pm 0.98$	0.055
4L	Early-heavy	4	$2.97\pm 1.12$	0.054
4L	Late-heavy	4	$3.53\pm 0.82$	0.045
8L	Uniform	8	$6.23\pm 2.00$	0.106
8L	Early-heavy	8	$5.64\pm 2.34$	0.102
8L	Late-heavy	8	$6.89\pm 1.70$	0.088

Appendix

Table 13: Controlled progressive-depth results. Oracle gain is measured relative to fixed full-depth computation in BCE.

Variant	Rel. compute	MMMU	MMBench-CN	MMBench-EN	MathVista	MMStar
Graded mixing instead of discrete gate	0.987	54.8	82.4	85.2	70.2	69.6
No action-mixture target in $\mathcal{L}_{\mathrm{gate}}$	0.988	55.2	82.0	84.0	72.2	70.4
$E=4$ , $k=1$	0.958	54.0	85.87	86.75	73.4	70.93
$E=2,k=2$	0.978	54.4	82.0	84.8	72.4	70.2
$E=6,k=2$	0.951	55.2	86.6	87.1	73.6	74.4

Appendix

Table 14: Capacity-axis design ablations on Bagel. “Rel. compute” is the normalized route/action proxy defined in Section A.7 .

Suite	Case	Task	Chars.	Depth	Halt@26	Risk@26	Regret@26
MMBench-EN	493	MMBench-English	66	26	0.581	0.091	0.038
MMBench-EN	630	MMBench-English	41	26	0.640	0.081	0.030
MMMU	Biology_21	biology	218	26	0.511	0.098	0.045
MMMU	Biology_8	biology	726	26	0.506	0.088	0.035
MMMU	Chemistry_30	chemistry	257	26	0.505	0.090	0.037
MathVista	129	iconqa_integer	104	26	0.523	0.087	0.034

Appendix

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.

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.

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.

Group	$N$	Mean halt@26	Mean risk@26	Halt pass	Joint pass
early	10	0.546	0.092	100.0%	100.0%
full	740	0.424	0.132	5.0%	0.0%

Appendix

Table 16: Depth-26 gate-condition comparison. “Halt pass” uses the configured halt threshold 0.5; “Joint pass” requires both halt probability

\geq 0.5

and risk upper bound

\leq 0.1

Figure 10: Depth-26 gate signals for early-exit and full-depth inputs. The dashed lines show the configured halt-probability and risk thresholds.

Panel of Table 17 — 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.

Figure 11: Selected generation set. Each tile contains the FlexUMM image followed by the paired baseline image.

#	Prompt	FlexUMM result	Baseline failure mode
1	purple pizza	Pass	No purple pizza; red pizza detected
2	red dog	Pass	Brown dog rather than red
3	white scissors	Pass	Purple scissors rather than white
4	purple backpack	Pass	Suitcase/handbag detected, not backpack
5	chair left of a zebra	Pass	Bench detected; no chair
6	parking meter above a broccoli	Pass	Parking meter below the broccoli

Appendix

Table 18: Detector outcomes for the selected generation set. The rows are individual image-level wins from distinct prompts.

Task Interference

See all 12 papers

LatentUMM: Dual Latent Alignment for Unified Multimodal Models
May 18, 2026Yinyi Luo, Wenwen Wang, Hayes Bai +2Cross-Modal Learning Unified Multimodal Models
Carnegie Mellon University · William & Mary
Where a New Concept Must Enter: Entry Point Gates Cross-Task Usability in Unified Multimodal Models
Aug 18, 2026Zongyang Qiu, Yihan Wu, Kaixuan Fan +2Unified Multimodal Models Multimodal Generation
The Hong Kong University of Science and Technology (Guangzhou) · Columbia University · MMLab, CUHK +1
TorchUMM: A Unified Multimodal Model Codebase for Evaluation, Analysis, and Post-training
Apr 12, 2026Yinyi Luo, Wenwen Wang, Hayes Bai +6Unified Multimodal Models Multimodal Model Evaluation
Carnegie Mellon University · William & Mary · Auburn University +1

May 18, 2026J/K move · Enter open · S save

LatentUMM: Dual Latent Alignment for Unified Multimodal Models

Yinyi Luo, Wenwen Wang, Hayes Bai, Marios Savvides+1

Carnegie Mellon University · William & Mary

Cross-Modal Learning Unified Multimodal Models Cross-Modal Alignment Multimodal Foundation Models Latent Space Alignment Multimodal Understanding

Open paper AskPDF arXiv

Unified multimodal models (UMMs) achieve strong performance in both understanding and generation by learning a shared latent space, yet they often exhibit functional inconsistency between these two capabilities. We observe that this issue does not stem from a lack of shared representations, but from the absence of explicit alignment between the transformations that map into and out of the latent space. As a result, generation and re-encoding can follow inconsistent trajectories, leading to semantic drift under modality transitions. In this work, we propose LatentUMM, a framework that constructs an enhanced shared latent space to explicitly align these transformations and improve cross-modal consistency. LatentUMM consists of two stages. First, dual latent alignment enforces consistency at both the modality and capacity levels: cross-modal alignment uses a stronger embedding model to impose structured cross-modal semantics, while dual capacity alignment enforces bidirectional consistency under generation and re-encoding. Second, latent dynamics stabilization improves robustness via stochastic latent rollouts and preference optimization, favoring trajectories that better preserve semantic consistency. Experiments show that LatentUMM consistently improves multimodal consistency across diverse architectures. Code is available at: https://github.com/AIFrontierLab/TorchUMM/tree/main/src/umm/post_training/LatentUMM.

Abstract

Task Interference

Explore similar work

LatentUMM: Dual Latent Alignment for Unified Multimodal Models

Where a New Concept Must Enter: Entry Point Gates Cross-Task Usability in Unified Multimodal Models

TorchUMM: A Unified Multimodal Model Codebase for Evaluation, Analysis, and Post-training