跳转至

LARGO: Low-Rank Hypernetwork for Handling Missing Modalities

会议: NeurIPS 2026
arXiv: 2605.06086
代码: https://github.com/nielsvyncke/LARGO
领域: 多模态VLM
关键词: 缺失模态, 超网络, 张量分解, CP分解, 多模态分割

一句话总结

LARGO 摒弃传统的特征空间插补策略,在权重空间利用典范多元(CP)张量分解对所有 \(2^N-1\) 个模态组合的专用卷积权重进行联合低秩重参数化,以仅相当于单模型的参数预算实现了对任意缺失模态子集的极致自适应和前沿分割性能。

研究背景与动机

多模态医学图像(如脑肿瘤 MRI 中的 FLAIR、T1、T1c、T2 以及缺血性卒中成像中的 DWI、ADC、FLAIR)能够为临床诊断提供高度互补的病理与解剖特征。然而在实际临床场景中,由于采集时间限制、扫描设备缺失、患者身体状况或协议差异,完整获取全部模态往往不切实际,导致测试阶段普遍面临部分模态缺失的严峻挑战。

现有应对缺失模态的方法大多在特征空间进行设计,但在“专用性”(针对特定模态组合的极致分割精度)与“可扩展性”(支持全部 \(2^N-1\) 种模态子集)之间陷入两难困境。单模型零填充或模态丢弃训练迫使固定的网络参数同时适配所有模态组合,严重牺牲了特定组合的最优表现;特征解耦与多分支融合结构复杂,极难直接迁移至不同骨干网络或新数据集;基于生成对抗网络或扩散模型的插补方法不仅依赖多阶段训练,更存在幻觉生成虚假病灶的临床风险;而若为每种模态子集独立训练 \(2^N-1\) 个专用模型,参数量与存储开销随模态数呈指数爆炸,完全无法工程部署。

这一困境的核心矛盾在于:不同模态子集对应的专用网络在功能上存在巨大的参数重叠与知识相关性,但现有架构无法在不显著膨胀参数量的前提下显式容纳这组指数规模的模型族。本文的核心 idea 是:将全体 \(2^N-1\) 个缺失模态专用模型的卷积层参数沿新增的“模型维度”堆叠为高阶张量,并通过典范多元(CP)分解将其联合压缩为一个低秩超网络,仅需存储一组紧凑的因子矩阵并在推理时即时重构目标专用权重,在单模型参数开销下兼得专用模型的精度与全组合支持。

方法详解

整体框架

LARGO 的核心架构围绕“权重空间跨模型联合低秩分解”展开。针对包含 \(N\) 个模态的多模态系统,共有 \(M = 2^N - 1\) 种非空输入组合。LARGO 不对特征层进行繁重的多分支特征对齐或模态补齐,而是以标准骨干(如 3D nnU-Net)为主体,将各层卷积核参数定义为一个多维权重流形。网络分为两部分处理:第一部分为包含具体输入通道适配的浅层 Stem 与最终分类输出的 Head,由于其参数量占比极小(在 BraTS 上仅占总参数的 0.131%),直接保留 \(M\) 个独立的未压缩轻量副本;第二部分为占据绝对主导参数量的内部卷积和转置卷积层,全部采用低秩卷积层(LRConv)重参数化。

在推理阶段,当系统接收到某种模态子集输入时,超网络直接根据该模态组合索引 \(m\),利用预先优化好的紧凑因子矩阵进行张量缩并,瞬时恢复出该模态组合专用的卷积核参数,随后执行标准前向计算。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入模态子集<br/>X_m 与组合索引 m ∈ {1,...,M}"] --> Stem["模态适配 Stem 层<br/>独立通道输入卷积"]
    Stem --> Body["U-Net 主体网络<br/>由因子矩阵即时重构卷积权重"]
    subgraph Core["超网络重参数化机制"]
        direction TB
        Factors["跨模型因子矩阵<br/>A (模型), B (Cin), C (Cout), D (K)"] --> Recon["张量缩并即时重构<br/>按组合索引 m 提取切片"]
    end
    Core -.-> Body
    Body --> Head["任务预测 Head 层<br/>独立输出分割图"]
    Head --> Out["最终病灶分割结果 Y_hat"]

关键设计

1. 跨模型维度的四阶卷积张量构建与 CP 分解:消除指数模型族的冗余 传统张量分解通常用于单个神经网络的内部压缩(分解单个层的通道或空间维度)。LARGO 首次引入“模型维度”,将 \(M = 2^N - 1\) 个专用模型第 \(l\) 层的卷积核权重 \(\mathcal{V}^{[m]} \in \mathbb{R}^{C_{\text{in}} \times C_{\text{out}} \times K}\)(其中 \(K\) 表示展平后的空间卷积核尺寸,如 3D 卷积下 \(K = K_d \times K_h \times K_w = 27\))堆叠为一个四阶总权重张量 \(\mathcal{W} \in \mathbb{R}^{M \times C_{\text{in}} \times C_{\text{out}} \times K}\)。随后,利用典范多元(Canonical Polyadic, CP)张量分解将其表示为 \(R\) 个一阶秩一分量的外积和: $\(\mathcal{W} = \sum_{r=1}^{R} \bm{a}_r \circ \bm{b}_r \circ \bm{c}_r \circ \bm{d}_r\)$ 其中 \(\bm{a}_r \in \mathbb{R}^M\)、\(\bm{b}_r \in \mathbb{R}^{C_{\text{in}}}\)、\(\bm{c}_r \in \mathbb{R}^{C_{\text{out}}}\)、\(\bm{d}_r \in \mathbb{R}^K\) 分别为因子矩阵 \(\bm{A}, \bm{B}, \bm{C}, \bm{D}\) 的第 \(r\) 列。因子矩阵 \(\bm{A}\) 专门用于索引 \(M\) 种模态组合,捕获了跨模态组合之间的参数继承与共享关系,而 \(\bm{B}, \bm{C}, \bm{D}\) 则抽象了通用的输入通道、输出通道与空间几何滤波模式。这一设计直接打破了传统超网络必须学习复杂非线性权重生成网络的常规套路,以严谨的多线性代数结构直接编码权重流形。

2. 恒定参数预算导向的秩自适应选取准则:实现单模型参数等价 为了使超网络总参数量严格对齐单一基线骨干网络的参数规模,LARGO 推导了一套通用的层自适应秩选取准则。对于原始单个卷积层,参数量为 \(C_{\text{in}} C_{\text{out}} K\);而在 CP 分解超网络下,重参数化层包含 4 个因子矩阵,参数量为 \((M + C_{\text{in}} + C_{\text{out}} + K) R\)。为了达到严格 \(1/M\) 的压缩比(即整个超网络容纳 \(M\) 个模型后的总参数量等同于 1 个传统模型),每层的秩 \(R\) 被精确设定为: $\(R = \left\lfloor \frac{C_{\text{in}} C_{\text{out}} K}{M + C_{\text{in}} + C_{\text{out}} + K} \right\rfloor\)$ 每一层还为各模型保留独立的偏置向量 \(\bm{b} \in \mathbb{R}^{C_{\text{out}} \times M}\)。借助该准则,无论网络深度和层通道数如何变化,模型在每层都能自适应确定既保持单网络参数约束、又最大化模型表达能力的张量秩。

3. 条件按需切片的前向重构与数值稳定归一化:兼顾推理高效与训练稳定 在推理执行期间,若先计算完整的四阶全局张量 \(\mathcal{W}\) 再进行切片,会带来巨大的计算冗余与峰值显存占用。LARGO 采用高效的局部张量缩并,针对给定的模态组合索引 \(m\),仅利用因子矩阵 \(\bm{A}\) 在第 \(m\) 行的标量元素 \(A_{m,r}\),将权重重构简化为: $\(\mathcal{W}^{(m)} = \sum_{r=1}^{R} A_{m,r} \left( \bm{b}_r \circ \bm{c}_r \circ \bm{d}_r \right)\)$ 从而完全避免了为未使用的模态组合生成参数。在训练稳定性方面,由于四阶外积对梯度的缩放极为敏感,LARGO 在反向传播中对对应于输入通道、输出通道和空间核尺寸的基向量(\(\bm{b}_r, \bm{c}_r, \bm{d}_r\))进行 \(\ell_2\) 范数归一化,将所有尺度自由度完全吸纳进模型维度原子 \(\bm{a}_r\) 中。同时,\(\bm{B}, \bm{C}, \bm{D}\) 初始化为 Kaiming 正态分布,\(\bm{A}\) 初始设为全 1 向量,偏置设为 0,确保了从零端到端平稳收敛,无需任何预训练权重或先验 CPD 拟合。

损失函数 / 训练策略

LARGO 采用结合随机模态丢弃与全模态引导的双分支联合监督策略。基础分割损失结合了软 Dice 损失与交叉熵损失:\(\mathcal{L}_{\text{seg}} = \mathcal{L}_{\text{Dice}} + \mathcal{L}_{\text{CE}}\)。在每个训练迭代步中,算法从离散均匀分布中随机采样一个缺失模态组合索引 \(m \sim \mathcal{U}(1, M)\),构造子集输入 \(\bm{\mathcal{X}}_m\),同时保留全模态输入 \(\bm{\mathcal{X}}\)(对应索引 \(M\)),总损失定义为两者等权重的平均: $\(\mathcal{L}_{\text{total}} = \frac{1}{2} \left( \mathcal{L}_{\text{seg}}\left(\bm{\mathcal{H}}_{\bm{\varphi}}(\bm{\mathcal{X}}_m; m), \bm{\mathcal{Y}}\right) + \mathcal{L}_{\text{seg}}\left(\bm{\mathcal{H}}_{\bm{\varphi}}(\bm{\mathcal{X}}; M), \bm{\mathcal{Y}}\right) \right)\)$ 其中第一项迫使超网络全面适应所有可能的子模态输入空间,第二项则持续提供高信噪比的全模态梯度引导,防止子模型权重在极端单模态欠拟合下偏离最优解。

实验关键数据

主实验

在 BraTS 2018(4 模态,共 15 种缺失场景)与 ISLES 2022(3 模态,共 7 种缺失场景)两个权威医学多模态分割基准上,LARGO 与包括前沿 Transformer、掩码预训练及特征解耦方法(mmFormer、M³AE、ShaSpec、SimMLM)进行了 5 折交叉验证评测。下表展示了 BraTS 2018 上 15 种模态组合的整体区域平均分割表现(Whole Tumor, WT; Tumor Core, TC; Enhancing Tumor, ET)以及全场景综合均值对比。

数据集 / 场景配置 评价指标 本文 (LARGO) mmFormer M³AE ShaSpec SimMLM
BraTS 2018 (WT 均值) Dice (%) ↑ 87.05 84.90 85.40 86.08 85.85
BraTS 2018 (TC 均值) Dice (%) ↑ 79.35 76.24 77.07 79.01 76.82
BraTS 2018 (ET 均值) Dice (%) ↑ 62.83 58.71 60.63 62.08 61.09
BraTS 2018 全场景均值 Overall Avg Dice (%) ↑ 76.41 73.28 74.37 75.73 74.59
ISLES 2022 全场景均值 Stroke Lesion Dice (%) ↑ 62.68 56.88 60.14 50.92 60.02
avMNIST 异构分类均值 Test Accuracy (%) ↑ 97.00 96.56 95.89 96.11 96.00

在 BraTS 2018 的 15 种模态配置下,LARGO 在全部 15 种组合的 WT 上均位列第 1;在 TC 上 11 种位列第 1,其余 4 种位列第 2;在 ET 上 14 种位列第 1,仅 1 种位列第 2。在全指标共 52 组评测场景中(BraTS 45 项 + ISLES 7 项),LARGO 夺得 47 项第一。特别是在具有严苛信息缺失的单模态极端场景下(如 BraTS 单 T1 模态从基线的 68.38% 跃升至 78.41%,ISLES 单 FLAIR 从 38.60% 跃升至 41.09%),优势极为显著。

消融实验

消融实验重点验证了张量分解秩 \(R\) 的取值对模型表征能力与容量权衡的影响,以及 CP 分解与 Tucker 分解在等参数量下的表现对比(基于 ISLES 2022 5 折交叉验证)。

配置 / 实验设定 模态组合配置 卒中病灶 Dice (%) ↑ 说明与观察
秩消融 \(R/4\) 全场景平均 61.07 极致压缩,但低秩约束过紧导致容量受限
秩消融 \(R/2\) 全场景平均 62.09 性能随秩提升平稳上升
默认基准 \(R\) (1/M 压缩) 全场景平均 62.68 严格匹配单模型参数量,兼顾性能与效率
秩扩展 \(2R\) 全场景平均 63.11 参数量增加,性能进一步提升
秩扩展 \(7R\) 全场景平均 63.36 参数规模与独立训练专用模型族相当
独立专用模型族 (Dedicated) 全场景平均 62.90 参数量与 \(7R\) 相当,但平均指标甚至低于 \(7R\) 与接近 \(R\)
结构消融:Tucker 分解 全场景平均 61.97 同等参数量下较 CP 分解落后 0.71%,单模态缺失受损更明显
结构消融:CP 分解 (LARGO) 全场景平均 62.68 因子化形式对跨模型流形建模表达更优

关键发现

  • 低秩超网络性能超越独立训练的专用模型:当秩扩展至 \(7R\) 时,LARGO 凭借因子共享带来的正则化效应与全模态知识迁移,取得了 63.36% 的均值,甚至超过了耗费海量算力完全独立训练的各专用模型(Dedicated,62.90%)。
  • 极端单模态场景下的鲁棒性提升最为显著:传统特征对齐与插补方法在仅有单一较弱模态时(如缺乏增强造影剂的 T1 或 FLAIR)极易发生严重表征崩溃,而 LARGO 依靠权重空间的低秩流形插值,保留了强大的单模态特征提取路径。
  • 计算与显存开销适中:在 BraTS 2018 3D nnU-Net 上,LARGO 仅增加 +0.1% 的参数量(22.59M vs 22.57M),动态重构使得推理 GFLOPs 增加 24.6%,峰值显存增加 25.1%,但带来了高达 +3.13% 的全场景 Dice 提升,远优于参数暴增至 187.7M 的 ShaSpec。

亮点与洞察

  • 从“特征空间调和”转向“权重空间流形压缩”:既往工作执着于在特征表征阶段填补缺失或解耦公共/特有特征,而 LARGO 洞察到不同模态子集的模型权重点具有极高的内在相关性,首次将跨模型的网络族压缩问题抽象为张量多线性代数分解,范式新颖。
  • 无须辅助损失与复杂训练轮次:不依赖知识蒸馏、教师模型、对抗判别器或多阶段掩码预训练,仅通过单阶段端到端训练与标准的分割 Loss,便自发实现了多模态子集的协同收敛。
  • 架构无关性与跨模态通用性:该低秩卷积与全连接超网络设计完全独立于上层任务逻辑,不仅无缝嵌入 3D nnU-Net,并在视听多模态分类任务(avMNIST)中展现了出色的泛化潜力,易于扩展至通用多模态骨干中。

局限与展望

  • 模态数量扩展上限受组合数制约:虽然内层卷积实现了低秩共享,但未压缩的 Stem 和 Head 层参数量仍与模态组合数 \(M = 2^N - 1\) 呈指数关系。当模态数 \(N \le 10\) 时完全可控(占比千分之几),但若面对数百种极度海量的输入模态,Stem/Head 仍需引入统一通道映射或稀疏化处理。
  • 目前主要面向 CNN 架构:核心机制针对卷积张量维度(\(M \times C_{\text{in}} \times C_{\text{out}} \times K\))定制,未来尚需探索在多头自注意力(MHA)权重矩阵以及大型 Transformer 骨干中的通用张量化适配形式。
  • 线性超网络假定的表达天花板:CP 分解本质上是多线性映射,对于模态间极度剧烈的非线性权重跳变,引入微型非线性调制网络可能会进一步释放超网络潜能。

相关工作与启发

  • vs mmFormer / M³AE: 后两者依赖 Transformer 交叉注意力或掩码自编码器进行特征维度的跨模态补全,计算图复杂且调参敏感;LARGO 在权重空间进行张量重参数化,推理时只运行极简的重构骨干,精度更高且结构稳健。
  • vs ShaSpec: ShaSpec 通过显式分离模态共享特征与专有特征提升鲁棒性,但参数量膨胀至 187.7M;LARGO 仅用 22.6M 参数(减少 88% 空间)便取得了显著更优的平均分割精度(76.41% vs 75.73%)。
  • vs 传统张量网络压缩 (Kim et al., Lebedev et al.): 传统方法均聚焦于“模型内部(intra-model)”对单层通道/空间降维压缩;LARGO 另辟蹊径,首次探索“模型家族间(inter-model)”的联合低秩参数化,为模型集成与动态自适应网络提供了全新视角。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次跳出特征空间,在权重空间利用张量分解联合压缩指数级缺失模态专用模型族。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 3D 医学分割两大数据集(BraTS 2018、ISLES 2022)共 22 种组合 52 项指标评测,辅以跨域 avMNIST 验证与完备的秩/结构消融。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰严密,实验设置公平严谨,动机与模型流形逻辑高度自洽。
  • 价值: ⭐⭐⭐⭐⭐ 为临床多模态医学影像落地缺失模态难题提供了极具工程优雅性与极佳精度的单模型解决方案。