FlexUMM adds 0.38% parameters to Bagel, lifting MMStar by 7.7 points while cutting compute 4%
FlexUMM: Flexible Representation Learning for Unified Multimodal Models
Yinyi Luo, Haoyue Bai, Yi Ge, Nilton Cesar Rojas Vales, Marios Savvides, Jindong Wang
cs.AI
2026-10-16
FlexUMM adds FFN routing and per-input early exit to unified multimodal models; on Bagel, MMStar rises from 63.20 to 70.93 at 0.96x relative compute and +0.381% parameters.
What problem this solves
Unified multimodal models (UMMs) try to do image understanding and image generation with one backbone. The two jobs pull the shared parameters in different directions, yet existing UMMs run every input and every capability through the same fixed path: same depth, same parameter sharing.
The paper points at two mismatches. First, depth. On Bagel, a learned router spreads the 2,374 MME questions over six exit depths between 22 and 28 layers, with a mean of 25.5. Inputs clearly do not need equal compute. Second, capacity. Which parts should be shared across understanding and generation, and which should be split, is not something a uniform design can answer.
Method
FlexUMM is a post-training plug-in. It trains only routers, halting heads and low-rank residual banks. The original attention, normalization and FFN weights stay as released.
- Shared attention. In a synthetic two-task study, attention gradients stay perfectly aligned across tasks (cosine 1.000 at every layer), while FFN gradient cosine falls from 0.955 at layer 1 to −0.758 at layer 8. So attention is kept as the shared communication channel and only the FFN gets specialized.
- Tri-state gate. At each routed FFN, every token chooses one of three actions: shared FFN only, shared FFN plus a low-rank residual, or skip the FFN. The residual's up-projection starts at zero, so inserting it leaves the pretrained model's behavior unchanged at initialization.
- Halting head. At candidate depths (22/24/26/28 for Bagel) a sequence-level head outputs a halt probability. If it passes 0.5, the model stops there and decodes with the existing LM or flow-matching head; the last candidate is the fallback.
- Training. Every candidate depth gets a task loss. The earliest candidate whose loss is within a margin of the full-depth one becomes the halting label, with extra penalties on token-level cost and expected depth.
Results
The main backbone is Bagel (14.6B), post-trained on ScienceQA, AI2D, TextVQA, MathV-360K and InstructPix2Pix.
| Metric | Bagel | +FlexUMM |
| MMMU | 51.9 | 54.0 |
| MMBench-CN | 80.94 | 85.87 |
| MMBench-EN | 82.65 | 86.75 |
| MathVista | 71.3 | 73.4 |
| MMStar | 63.20 | 70.93 |
| GenEval | 78.81 | 82.03 |
| WISE | 0.4044 | 0.4292 |
On other backbones, Janus-Pro's MMMU goes from 40.89 to 46.11 and Harmon's MathVista from 22.0 to 35.1, at relative compute of 0.96 and 0.89. The paper reports an average gain of 10.17% with 6.33% fewer FLOPs across the three backbones.
In ablations, a random router with identical parameters drops MMBench-EN to 74.34, below the 82.65 baseline, so the gain comes from learned allocation rather than extra parameters. Fixed depth 25 gives MathVista 48.2 and fixed depth 26 gives 67.2, against 73.1 for adaptive depth. Unison overall moves from 55.59 to 59.83 and Unified-Bench from 83.46 to 85.05. On Uni-SafeBench, Both-Safe rises from 43.96% to 44.86%, a small change.
Why it matters
If you already run a Bagel-class checkpoint, this is a recipe that adds under 0.4% parameters and leaves the backbone untouched. The halting threshold can be changed at inference time to trade accuracy for compute without retraining.
The main payoff is accuracy, not speed. Relative compute on Bagel only drops to 0.96.
Limitations
- Compute numbers are a proxy, not measured latency. The current wrapper still evaluates the dense base FFN before applying the action mask, so real savings come only from sequence-level early exit. Token-level skipping has no demonstrated benefit yet.
- The gradient-conflict evidence is synthetic and hand-designed. The authors state it does not measure real Bagel or Janus-Pro gradients.
- Some numbers disagree across tables: relative compute at threshold 0.5 is 0.993 in Table 2 but 0.96 in Table 6, and the FlexUMM MMMU score is 53.6 in the ablation versus 54.0 in the main table. The text claims best results on all five understanding benchmarks, yet Bagel+UniCot scores 83.0 on MathVista against 73.4.
- The gains are large (MMStar +7.7), but the paper gives no clearly matched plain fine-tuning control on the same post-training data.
- The early-exit case audit covers only 10 unique inputs, which the authors say is not a statistical estimate.
Terms
- UMM: a single network that handles both multimodal understanding and generation.
- FFN: the feed-forward block after attention in each Transformer layer, holding most parameters.
- Early exit: stopping at an intermediate layer and decoding from there instead of running all layers.
- Low-rank residual bank: a small adapter made of a down- and up-projection, far cheaper than a copied FFN.
- Flow matching: a generative training objective where the network predicts a velocity field from noise to image.
Source
What people are saying
Related papers
All paper explainers