Bagel 仅加 0.38% 参数做路由与早退,MMStar 涨 7.7 分、算力降 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 在统一多模态模型上加 FFN 路由与按输入早退,Bagel 的 MMStar 从 63.20 升到 70.93,相对算力降到 0.96,参数仅增 0.381%。

这篇在解决什么

统一多模态模型(UMM)想用一个骨干同时做图像理解和图像生成。问题是两种能力对同一套参数的要求会打架,而现有 UMM 的计算路径基本固定:每个输入、每种能力都走同样的层数,共享同样的参数。

论文指出两处错配。一是深度:Bagel 在 MME 的 2,374 道题上,学出来的路由让输入分散在 22 到 28 层共六种深度,均值 25.5,说明不同输入需要的计算量不同。二是容量:理解和生成该共享哪些部件、该分开哪些部件,一刀切都不合适。

方法

FlexUMM 是挂在预训练 checkpoint 上的后训练插件,只训练路由器、停止头和低秩残差库,原始注意力、归一化和 FFN 权重不动。

共享注意力。作者用合成的双任务实验观察梯度:注意力层的跨任务梯度余弦始终是 1.000,FFN 层从第 1 层的 0.955 掉到第 8 层的 −0.758。结论是注意力当通信通道共享,FFN 才需要分化。
Tri-state Gate(三态门)。每个 token 在每个被路由的 FFN 处三选一:只走共享 FFN、共享 FFN 加一个低秩残差(specialized)、或跳过 FFN(bypass)。残差库的上投影零初始化,所以插入瞬间模型行为和原版一致。
Halting Head(停止头)。在候选深度(Bagel 为 22/24/26/28 层)上给整条序列算一个停止概率,超过阈值 0.5 就在该层输出,否则继续,最后一个候选兜底。输出直接复用原来的 LM head 和 flow-matching head。
训练。对每个候选深度同时算任务损失,把「不比满深度差太多」的最早候选当作停止标签,再加上容量成本和期望深度两项正则。

结果

主骨干是 Bagel(14.6B),后训练数据包括 ScienceQA、AI2D、TextVQA、MathV-360K 和 InstructPix2Pix。

指标	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

跨骨干:Janus-Pro 的 MMMU 从 40.89 到 46.11,Harmon 的 MathVista 从 22.0 到 35.1,相对算力分别为 0.96 和 0.89。论文称三个骨干平均提升 10.17%,FLOPs 降 6.33%。

消融里,随机路由在参数量不变的情况下,MMBench-EN 掉到 74.34,低于原版 82.65,说明收益来自学到的分配,不是多出来的参数。固定深度 d=25 时 MathVista 只有 48.2,固定 d=26 为 67.2,都远低于自适应的 73.1。统一性基准 Unison 总分 55.59 到 59.83,Unified-Bench 83.46 到 85.05。Uni-SafeBench 上 Both-Safe 从 43.96% 到 44.86%,幅度很小。

为什么重要

对手上有 Bagel 一类 checkpoint 的人,这是一个参数开销不到 0.4%、不改骨干的后训练配方,精度和算力两头都有账可算。阈值 η 可以在推理时调,不用重训就能换精度和算力的比例。

它的主要价值在精度,不在省算力。Bagel 上相对算力只降到 0.96。

局限与存疑

算力数字是代理指标,不是实测延迟。当前实现仍会先算完整个基础 FFN 再套动作掩码,真正的节省只来自序列级早退,token 级的 bypass 与 decouple 在这个实现里没有实测收益。
梯度冲突的证据来自人工构造的合成实验,作者自己说明它不代表真实 Bagel 或 Janus-Pro 的梯度,真实骨干上的逐组件梯度测量没有做。
数字前后不完全一致:Table 2 里 η=0.5 的相对算力是 0.993,Table 6 写的是 0.96;Table 3 的 FlexUMM 行 MMMU 是 53.6,Table 1 是 54.0。正文称在五项理解基准全部最优,但 MathVista 上 Bagel+UniCot 为 83.0,高于 73.4。
增益幅度不小(MMStar +7.7),但后训练用了理解类和编辑类数据,论文未给出同数据量的普通微调(Bagel+SFT 的对比设置也没写清)作为严格对照。
早退案例审计只覆盖 10 个唯一输入,作者承认不能当统计结论。

术语

UMM:Unified Multimodal Model,同一网络同时做多模态理解与生成。
FFN:Transformer 里每层注意力之后的前馈网络,承担大部分参数与非线性变换。
早退(early exit):推理时不跑完所有层,在中间层就用输出头给结果。
低秩残差库:用两个小矩阵(降维再升维)组成的轻量适配器,比复制整个 FFN 便宜得多。
flow matching:一类生成训练方式,让网络预测从噪声到图像的速度场。

原文与代码

社区讨论

FlexUMM 论文：统一多模态模型表征空间其实并不统一 jindong_wang92
FlexUMM 论文:统一多模态模型的嵌入空间其实并不统一 jindong_wang92