跳转至

OmniMamba: Efficient and Unified Multimodal Understanding and Generation via State Space Models

会议: ECCV2026
论文: ECCV 原文
代码: https://github.com/hustvl/OmniMamba
领域: 多模态 VLM
关键词: 状态空间模型、统一理解与生成、解耦词表、任务专属 LoRA、数据效率

一句话总结

OmniMamba 以 Mamba-2 为共享自回归主干,通过解耦视觉编码、输出词表和任务适配器,在作者所称约 2M 图文对的统一模型训练规模下兼顾理解与生成,在 COCO 上取得 5.50 FID-30K,但这一数据口径不包含继承的预训练组件成本。

研究背景与动机

统一视觉语言模型希望同一个系统既能回答图像问题,也能按文本生成图像,但这两件事对视觉表征的要求并不相同。 理解需要保留物体语义、关系和可用于回答问题的细节;生成还需要可还原空间结构与纹理的表示。 Show-o 将文本自回归与离散扩散结合,JanusFlow 将自回归与 rectified flow 结合,Emu3 则坚持统一的下一 token 预测。 这些路线说明,统一任务不等于只有一种生成范式,但多数仍使用 Transformer 主干。 随着视觉 token 和输出序列变长,注意力计算与缓存会成为成本来源;另一方面,紧耦合的表征和词表也可能增加有限数据下的学习负担。

本文关心的不是只把 Transformer 换成更快的模块,而是怎样让基础能力相对有限的 Mamba-2 在小规模多模态训练中同时学会两项任务。 如果图像和文本共用输出词表,模型除了预测内容,还得学会不能在图像序列中输出文本 token。 如果两项任务必须按固定比例联合预训练,理解与生成数据的不平衡又会限制各自的有效利用。 因此,计算效率与数据效率需要分别处理:状态空间主干负责前者,显式任务边界与分阶段训练负责后者。 论文所称的小数据训练是基于已有语言模型、视觉编码器和图像 tokenizer 的适配,并非从零训练全部组件。

OmniMamba 保留共享主干和统一交叉熵目标,却不要求所有接口都共享。 它让理解走连续视觉特征到文本的路径,让生成走文本到离散图像码的路径,并在共享层内保留任务专属的低秩更新。 这把一部分本来需要样本反复教会模型的模态规则,直接变成结构约束。 核心 idea:统一下一 token 预测与主干计算,同时解耦任务需要的视觉表示、适配参数和输出空间,以减少少数据训练中的跨任务干扰。

方法详解

整体框架

模型基于具有 48 层的 Mamba-2-1.3B,处理多模态理解(MMU)和文本到图像生成(T2I)。 MMU 的输入是图像与文本指令,输出为文本;T2I 的输入是文本条件,输出为可解码成图像的离散码序列。 两条路线依次经过任务解耦视觉编码、任务专属 LoRA 所在的共享主干,以及解耦输出词表。 其中,生成编码器用于训练时把真实图像转成监督 token,推理时并不要求额外输入一张目标图像。 推理生成的离散图像码由对应图像解码器恢复成像素;理解分支则直接输出文本 token。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
        Input["图像与文本条件"] --> Enc["任务解耦视觉编码"]
        Enc -->|"连续理解特征 / 训练图像码"| Backbone["任务专属 LoRA<br/>共享 Mamba-2 主干"]
        Text["T2I 推理文本条件"] --> Backbone
        Backbone --> Vocab["解耦输出词表"]
        Vocab -->|"MMU:文本 token"| Answer["文本回答"]
        Vocab -->|"T2I:图像 token"| Decoder["图像解码器 → 图像"]

图中展示的是表示与预测路径;两阶段训练控制哪些参数更新,不是推理时要额外执行的阶段。 Mamba-2 利用状态空间序列计算避免显式构造完整的 token 两两注意力关系,这是长序列效率的基础。 但生成仍然是逐 token 自回归,并不是把所有图像 token 一次性并行预测。 因此,线性主干带来的速度收益不能直接等同于更强的长上下文语义理解能力。

关键设计

1. 任务解耦视觉编码:让语义理解和像素重建使用合适的表示

对于理解任务,作者沿用 Prismatic VLMs 的思路,融合 DINOv2 与 SigLIP 的连续视觉特征。 DINOv2 提供空间结构相关信息,SigLIP 提供与语言联系紧密的语义信息,投影层把视觉表示接入语言主干。 主实验的理解输入分辨率为 384,重点是让主干获得回答问题所需的视觉条件。 这里没有强迫理解模型先把图像压缩成适合重建的离散码,因此不必让同一种压缩目标兼顾所有语义需求。

对于生成任务,作者采用 LlamaGen 训练的 VQVAE,把真实图像编码成离散图像 token,主实验生成分辨率为 256。 这些 token 是下一 token 预测的监督目标,推理得到相同码空间中的序列后再交给解码器。 理解分支的连续特征与生成分支的离散码不能互换,它们分别服务于条件感知和图像重建。 该 tokenizer 已在 ImageNet 上预训练,并在 50M LAION-COCO 与 10M 内部高美学质量数据上进一步微调。 因此,“生成训练只用 83K COCO 图像”应理解为本文生成模型适配阶段的规模,而不是视觉表示学习的全部数据成本。

2. 任务专属 LoRA:在共享状态空间计算前调节输入投影

每个 Mamba-2 层的输入线性投影都有两套 LoRA,分别对应 MMU 与 T2I。 执行理解时,仅激活 MMU LoRA;执行生成时,仅激活 T2I LoRA,另一条路径停用。 低秩更新与原输入投影共同作用,因此两项任务共享大部分参数,却能调整进入层内计算的特征。 作者的假设是,显式按任务参数化输入投影,可以帮助状态空间模型用更少数据适应不同任务需求。

这不是根据每个 token 内容学习专家路由的混合专家系统,而是由已知任务选择相应适配器。 路由把理解与生成的部分参数更新隔开,降低同一组轻量适配参数被两个目标反复拉扯的风险。 LoRA rank 设为 8,论文报告参数增量仅 0.65%。 该增量描述适配器的规模,不能推导为整个训练过程只更新这 0.65% 的参数;第二阶段更新范围另有说明。 表 5 的去 LoRA 对照让这一结构选择有了直接证据,而不只依赖理论上的任务干扰解释。

3. 解耦输出词表:把输出模态从学习问题变成接口约束

共享大词表要求模型同时判断“下一个内容是什么”和“它属于文本还是图像”。 OmniMamba 使用独立的文本词表与图像码本,并由任务选择相应输出头与词表。 MMU 的预测空间是文本 token,T2I 的预测空间是图像 token,因此可以从输出接口排除错误模态。 作者观察到统一词表配置有时在视觉生成过程中输出文本相关 token,需要额外后处理。 双词表并不自动提升图像语义正确性,但它减少了学习与纠正这类非法输出的负担。

两条路径仍采用同一种下一 token 预测范式,区别在于条件表示和合法目标集合。 这说明“统一”落在序列计算与训练目标上,而不是要求图像和文本必须共享所有离散索引。 输入格式通过 [MMU]、[T2I] 区分任务,通过 [SOI]、[EOI] 和 [SOT]、[EOT] 标记图像与文本边界。 理解序列先放图像条件再放文本,生成序列先放文本条件再预测图像码。 原文提取的完整 T2I 格式行有损坏,因此这里仅解释可核实的顺序与标记,不重建声称精确的完整字符串。

一个完整示例

以下是机制示例,不是论文新增的定量实验:先给模型一张街景图,要求描述画面,再另行输入一句街景描述要求生成图像。 在描述任务中,图像经过 DINOv2 与 SigLIP,连续特征经投影后成为 Mamba-2 的视觉条件。 模型激活 MMU LoRA,并通过文本词表依次预测回答;它不需要重建输入图像的像素。 在生成任务中,文本先进入共享主干,模型改用 T2I LoRA 和图像输出词表。 每次预测得到的是图像码而非自然语言词,完整码序列最终交给图像解码器。 训练这条路线时,真实图像会被 VQVAE 编码成目标码;推理时这些目标码尚不存在,必须由模型逐步生成。 这两次调用共享语言主干,却不意味着模型已经验证支持任意交错的图文长对话或图像编辑。

损失函数 / 训练策略

第 1 阶段分开进行任务专属预训练,核心 Mamba-2 保持冻结,以免一开始就让不同规模的数据争用主干更新。 MMU 训练视觉投影层与 MMU LoRA,停用 T2I 路径;T2I 训练其 LoRA 与图像预测端,停用 MMU 路径。 第 3.5 节在“image head”和“image decoder”之间用词不一致,因此仅凭正文无法确定后者是否特指 VQVAE 像素解码器。 本笔记不据此断言该预训练图像解码器一定被更新。 第 2 阶段进行统一微调,每条任务路径保留独立 LoRA,将两项损失相加后统一反向传播。 第 3.5 节首先说冻结视觉编码器、训练其他模块,随后又写通过“frozen backbone”共享参数,两种说法冲突;主干冻结状态需实现核实。

第 3.6 节说明两项任务都使用标准下一 token 交叉熵,不引入扩散去噪目标。 下面只是对这段文字目标的标准记号表达,不是原文编号公式;其中 \(y_i\) 是所属任务的目标 token,\(c\) 是输入条件。

\[ \mathcal{L}_{\mathrm{task}}=-\sum_i\log p_\theta(y_i\mid y_{<i},c),\qquad \mathcal{L}=\mathcal{L}_{\mathrm{MMU}}+\mathcal{L}_{\mathrm{T2I}}. \]

第 4.1 节列出理解预训练数据为 676K,包括 118K COCO 和 558K LLaVA-1.5 图文对;第 3.5 节的动机段却写 665K,存在数量不一致。 理解微调使用 665K LLaVA-1.5 多轮对话、220K LVIS-Instruct-4V 和 400K LRV-Instruct;生成使用 83K MS-COCO 2014 图像。 这些是跨阶段、跨数据集的不同样本口径,不能简单相加后当作唯一图文对数量,也不能把作者“少于 2M”的说法解释为全组件预训练总量。 所有阶段使用 AdamW,动量参数为 0.9 和 0.95,余弦学习率调度配合 warm-up,权重衰减为 0,梯度裁剪阈值为 1.0。 训练在 NVIDIA A800 上使用 BF16;速度测试则使用 NVIDIA 4090 和 FP16,二者不是相同条件。

实验关键数据

主实验

下表摘自原文表 1,第 10 页。均为统一模型,语言主干规模为 1.3B;指标越高越好,但分辨率、预训练主干和数据不同,并非只替换架构的受控比较。

模型 理解分辨率 POPE MME-P VQAv2 test GQA MMMU
Show-o 512 80.0 1097.2 69.4 58.0 26.7
JanusFlow 384 88.0 1333.1 79.8 60.3 29.3
OmniMamba 384 86.3 1290.6 77.7 60.8 30.6

OmniMamba 超过 Show-o 的所有上述指标,但并未全面超过 JanusFlow:GQA 和 MMMU 较高,POPE、MME-P 和 VQAv2 较低。 下表摘自原文表 2,第 11 页;FID-30K 使用 COCO 验证集随机选择的 30K 文本提示,数值越低越好。 “生成训练图像”是表中报告的任务数据,不包含全部预训练依赖,且 OmniMamba 与 U-ViT 在 COCO 训练集训练,其他方法并非同一训练协议。

模型 参数量 生成训练图像 FID-30K
Show-o 1.3B 35M 9.24
U-ViT 45M 83K COCO 5.95
OmniMamba 1.3B 83K COCO 5.50

因此,5.50 支持其 COCO 分布内生成能力,不能单凭此表认定其开放域文本遵循、审美或人脸生成全面领先。 表 2 的 Show-o 生成数据 35M 与正文讨论的统一模型数据 2B 不是同一口径,不能混用。

消融实验

下表摘自原文表 5,第 13 页。第 4.6 节称消融使用 CLIP 理解编码器、224 分辨率和相同的缩短训练步数;这里仅选取 130M 和 370M 行,不混入主实验的 1.3B 行。

实验 主干 解耦编码器 解耦词表 任务 LoRA POPE MME GQA FID-30K
1 Mamba2-130M 否 是 是 65.6 598 35.7 31.4
2 Mamba2-130M 是 是 是 80.6 930 46.1 23.7
5 Mamba2-370M 是 否 是 80.8 1036 53.6 19.1
6 Mamba2-370M 是 是 否 81.2 1003 54.0 14.4
7 Mamba2-370M 是 是 是 81.9 1100 55.3 10.3

实验 1 → 2 支持解耦编码器;实验 5 → 7 支持双词表;实验 6 → 7 支持任务 LoRA。 其中加入任务 LoRA 后,FID-30K 从 14.4 降到 10.3,下降 4.1,MME 从 1003 升到 1100。 原文词表消融段把比较编号写成实验 1 和 3,但表中这两行并非词表开关对照;本笔记按表格字段采用实验 5 和 7,并保留这一矛盾。

关键发现

  • 原文表 4,第 12 页:单张 NVIDIA 4090、FP16、batch size 16、256 分辨率下,OmniMamba 为 5.68 image/s,Show-o 为 0.81,JanusFlow 为 1.02。
  • 原文表 3,第 11 页:序列长度 16384 时,OmniMamba 为 248.1 token/s,使用 FlashAttention-2 与 KV cache 的 SmolLM2-1.7B 为 57.40 token/s。
  • 第 4.5 节的 119.2 倍加速发生在 16k 序列,相对 Show-o;63.0% 显存节省发生在 23k,二者不是同一个测量点,Show-o 此处也未使用 FlashAttention-2。
  • 长序列吞吐测试移除了输出长度限制,没有同时证明生成内容在超长上下文中仍然正确;速度、显存和任务质量应分开评价。

亮点与洞察

  • 数据效率不只来自减少参数,也来自缩小模型必须学习的规则集合。双词表预先排除了错误模态,是可以迁移到其他离散多模态系统的接口设计。
  • 共享主干与任务专属输入投影可以同时成立。LoRA 在进入状态空间计算的位置调节特征,比把两个完整模型简单拼接更节省重复参数。
  • 分开预训练、再统一微调让任务不必从随机接口开始就相互竞争。这一策略尤其适合样本量不对称、但已有可靠预训练组件的场景。

局限与展望

  • 作者在第 6 节承认,COCO 生成数据的覆盖限制了图像质量,特别是人脸仍较粗糙;更低 FID 不代表所有视觉内容都可靠。
  • Mamba-2 基础模型仅在最长 2048 token 的序列上训练。论文展示的是超长输出的计算可行性,不是超长上下文理解或 CoT 能力已经得到充分训练。
  • 全部组件的数据成本远大于生成适配阶段的 83K 图像。评价“数据节省”时需要拆开语言模型、视觉编码器、tokenizer 和统一模型的训练账目。
  • 两阶段训练本身缺少这里可核实的独立对照表;冻结范围、预训练样本数和词表消融编号的文字矛盾也妨碍严格复现。
  • 可继续在同等预训练依赖、相同训练预算和更广生成提示集下比较架构,把系统吞吐收益与表征及数据收益分离。

相关工作与启发

  • 与 Janus / JanusFlow 的关系:继承任务解耦视觉编码的方向,但把统一序列主干换成 Mamba-2,并使用纯自回归图像 token 预测而非 rectified flow。
  • 与 Show-o / Emu3 的关系:不把共享词表当作统一模型的必要条件。本文消融支持在有限训练预算下显式指定输出模态,但没有证明共享词表在大规模训练中必然更差。
  • 与 Cobra / Prismatic VLMs 的关系:前者为 Mamba 多模态理解及数据配方提供基础,后者启发 DINOv2 与 SigLIP 特征融合;OmniMamba 将理解进一步扩展到统一生成。
  • 与 LlamaGen 的关系:直接利用其图像 tokenizer 降低生成接口学习难度,也继承其预训练数据依赖;这不是从零学习视觉码本的实验。
  • 分类判断:核心任务是同一模型的多模态理解与生成,保留多模态 VLM 分类比仅归入图像生成更能覆盖论文贡献。

评分

  • 新颖性: 4/5。将线性状态空间主干与显式任务解耦结合,创新主要在统一系统设计。
  • 实验充分度: 3/5。覆盖理解、生成、速度与结构消融,但数据口径及训练策略对照仍不充分。
  • 写作质量: 3/5。主线清楚,冻结范围和消融编号等不一致影响复现判断。
  • 价值: 4/5。为依赖已有预训练组件的低预算统一模型研究提供了实用设计,但不应扩大解释其长序列能力。