跳转至

SAM+D: Parameter-Efficient Dimensional Lifting of SAM-Family Models via Depth-Routed LoRA and Depth Shifting

会议: ECCV 2026
论文: ECCV 原文
代码: https://github.com/JerrySongCST/SAM-Plus-D
领域: 医学图像
关键词: 基础模型微调, 参数高效微调(PEFT), 维度提升, 3D/4D分割, SAM

一句话总结

SAM+D 提出了一种通用的参数高效升维框架,通过在冻结的 Transformer 块中插入深度路由 LoRA 专家(DRLoRA)与零参数跨切片深度位移模块(DSM),仅训练约 2.8%~3.7% 的参数即可将 2D SAM 升维至 3D 体积分割,并首次将视频 SAM2 无缝升维至 4D 时空生物细胞追踪。

研究背景与动机

以 Segment Anything Model(SAM)为代表的二维视觉基础模型在自然图像分割上展现出惊人的零样本泛化能力,而 SAM2 进一步通过记忆机制将该能力扩展至 2D 视频。然而在现实世界的医学断层扫描(如 3D CT/MRI)以及显微时间序列成像(如 4D 活细胞追踪)中,解剖组织或细胞在深度方向上呈现出强烈的连续性与结构重叠,原生的 2D 切片独立推理完全忽视了切片间(inter-slice)的上下文依赖,极易产生切片间断裂和假阳性;同时现有的 3D 专用大模型(如 nnU-Net、SAM-Med3D、SegVol)要么从头训练需要庞大的 3D 标注体量,要么完全破坏了 2D 预训练的大规模先验。

目前将 2D 基础模型适配至 3D 任务的微调方案存在三个显著技术瓶颈:其一,传统低秩适应(LoRA)对体积内所有切片应用无差别的均匀参数变换,忽略了解剖结构在深度轴上的非均匀性(例如器官顶底部的锐利边缘与内部均质区域需要截然不同的表征调优),缺乏位置感知路由机制;其二,现有建模层间依赖的方法大多依赖计算密集的 3D 卷积或层间交叉注意力,参数量与显存开销激增,尚无零额外计算代价的层间通信方案;其三,现有方法通常割裂地针对 SAM 或 SAM2 设计特定适配层,缺乏统一的通用架构来同时实现 SAM 的 2D→3D 升维以及 SAM2 的 (2D+T)→(3D+T) 时空升维。

针对上述瓶颈,本文提出了一种即插即用、结构极简的升维范式。核心 idea:将切片间通信与深度空间自适应解耦,利用无参数的深度通道位移(DSM)在注意力前实现切片特征无损交互,并基于归一化深度坐标路由轻量 LoRA 专家(DRLoRA)进行非均匀空间更新,在冻结绝大部分主干参数的前提下实现 SAM 家族从 2D 到 3D/4D 的参数高效维度提升。

方法详解

整体框架

SAM+D 统一支持 3D 体积分割与 4D 时空分割。对于 3D 体积数据,输入体积经沿深度轴均匀采样并缩放后,以多切片构成的批次形式送入冻结的 ViT-B 编码器;在编码器的每个 Transformer Block 中,特征先经过 DSM 进行相邻切片间的通道平移,随后在自注意力计算中注入经归一化深度坐标调制的 DRLoRA 增量更新;提取的 3D 特征体结合带相对深度编码的提示向量,送入轻量适配的解码器和 3D 卷积后处理头输出稠密分割。对于 4D 时空设置,相同的 DRLoRA 与 DSM 插入到 SAM2 的 Hiera 骨干中,通过对记忆注意力模块注入 LoRA 捕捉三维时空匹配模式,实现端到端的时空目标追踪。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入 3D 体积 / 4D 显微序列<br/>沿深度采样切片批次 D_s"] --> B["阶段 1:深度位移模块 (DSM)<br/>相邻切片双向通道平移 (零额外参数/MAC)"]
    B --> C["阶段 2:深度路由 LoRA (DRLoRA)<br/>归一化深度 z 驱动轻量 MLP 动态加权 E 个专家"]
    C --> D["阶段 3:冻结注意力与特征提取<br/>更新 Q/V 投影并保持 K 与主干参数冻结"]
    D --> E["阶段 4:提示相对深度嵌入与解码适配<br/>注入 z_embed 调制切片距离并由 3D 头平滑输出"]
    E -->|4D 序列任务| F["阶段 5:SAM2 记忆注意力 LoRA 适配<br/>跨时间帧体积记忆匹配与连续追踪"]

关键设计

1. 深度位移模块 (DSM):零参数开销的跨切片特征交换 独立处理各切片会丢失深度轴拓扑连续性,而引入 3D 卷积或切片注意力又会显著拉升显存和 FLOPs。针对这一矛盾,DSM 将视频领域的时序位移机制(TSM)迁移至空间深度轴。对于特征张量 \(\mathbf{X} \in \mathbb{R}^{D_s \times H \times W \times C}\),DSM 沿通道维度按位移比例 \(\alpha\)(默认 25%)将其切分为三组:前 \(\alpha C\) 个通道由前一切片 \(i-1\) 向前复制(forward shift),次 \(\alpha C\) 个通道由后一切片 \(i+1\) 向后复制(backward shift),其余 \((1-2\alpha)C\) 个通道保持原地不变: $\(\mathbf{X}' = \text{DSM}(\mathbf{X}) = \text{concat}(\mathbf{X}^{\text{fwd}}, \mathbf{X}^{\text{bwd}}, \mathbf{X}^{\text{id}})\)$ 首尾边界切片则复制自身特征以防引入人工伪影。DSM 严格置于每个 Transformer 块的自注意力计算之前,这意味着进入 \(Q, K, V\) 投影的每个空间 token 已经预先融合了上下切片的上下文信息。由于纯粹基于步长切片内存拷贝,DSM 引入了 0 个新增参数与 0 个乘加计算(MAC),在 RTX 5090 上的单次拷贝延迟仅为 0.34 ms(仅占单个编码块前向耗时的 1.7%),却使网络在多层堆叠后能逐层累积宽广的深度感受野。

2. 深度路由 LoRA 专家 (DRLoRA):基于空间深度坐标的自适应低秩更新 传统 LoRA 在整个 3D 体积内对所有切片应用相同投影增量,但解剖组织沿体轴差异剧烈(如肝脏顶部边缘与实质内部特征迥异),单一低秩矩阵表征容量有限。DRLoRA 引入 \(E\) 个并行的 LoRA 专家 \(\{ \mathcal{E}_1, \dots, \mathcal{E}_E \}\),每个专家维护独立的降维与升维权重 \(\mathbf{W}_{\text{down}}^{(i)} \in \mathbb{R}^{r \times d}, \mathbf{W}_{\text{up}}^{(i)} \in \mathbb{R}^{d \times r}\)\(W_{\text{up}}\) 初始化为 0 以确保训练起始等价于恒等映射)。与自然语言中基于 token 内容路由的 MoE 不同,DRLoRA 创新性地采用标量切片深度位置 \(z = d / D_s \in [0, 1]\) 作为路由依据: $\(\mathbf{w}(z) = \text{softmax}(\mathcal{R}(z)), \quad \mathcal{R}: \mathbb{R} \xrightarrow{\text{Linear}} \mathbb{R}^h \xrightarrow{\text{GELU}} \mathbb{R}^h \xrightarrow{\text{Linear}} \mathbb{R}^E\)$ 最终加权混合更新量为 \(\text{DRLoRA}(\mathbf{x}, z) = \sum_{i=1}^E w_i(z) \cdot \mathbf{W}_{\text{up}}^{(i)} \mathbf{W}_{\text{down}}^{(i)} \mathbf{x}\)。该设计具有多重精妙优势:第一,切片深度天然均匀分布在 \([0, 1]\) 之间,模型在训练中绝对不会发生传统 MoE 的专家坍缩(Expert Collapse),因此完全无需额外的负载均衡辅助损失;第二,路由网络仅为极小的两层 MLP,参数量仅 \(\sim 0.14\text{K}\)(相较于基于特征内容的路由网络降低了数十倍);第三,遵循低秩微调经验,DRLoRA 仅注入到自注意力的 \(Q\)\(V\) 投影矩阵(\(K\) 保持冻结),使得自适应更新能够动态调控关注何处(\(Q\) 控制注意力模式)以及传递什么信息(\(V\) 控制信息载荷),同时冻结的 \(K\) 维持了预训练特征空间的全局对齐基准。

3. 相对深度提示嵌入与轻量 3D 头部解耦精炼 在单点提示(single-point prompt)交互分割中,原始 SAM 解码器缺乏对 3D 绝对与相对空间距离的感知。为此,论文引入了一个两层 MLP 构成的相对深度编码器 \(z_{\text{embed}}\)。若用户在切片 \(z_p\) 处给出一个点击,该模块计算每个切片 \(d\) 距提示切片的相对空间偏移 \((z_p - d) / D_s\),生成与密集特征维度对齐的偏置并直接叠加到稠密提示向量上: $\(\mathbf{D}'_d = \mathbf{D}_{\text{prompt}} + z_{\text{embed}}\left( \frac{z_p - d}{D_s} \right)\)$ 这使得解码器能自适应地感知当前切片是位于提示点所在层,还是渐进衰减的远端邻层。随后在解码器的注意力层同样嵌入 LoRA,输出各切片的 2D logits 掩码堆叠成 3D 体积后,流经一个紧凑的 3D 卷积头完成切片间的平滑过渡与多分类决策。

4. 4D 时空升维与记忆注意力 LoRA 协同适配 在针对活细胞时空追踪的 4D 任务中,模型必须同时应对 3D 空间体积和第四维时间轴。SAM+D 保持模块设计的高度统一性:将相同的 DSM 和 DRLoRA 插入 SAM2 的层次化 Hiera 编码器中,但严格限定 DSM 和 DRLoRA 仅在 3D 空间深度轴上生效,不跨时间帧混合。时间一致性的维系完全交由 SAM2 预训练的记忆网络接管;由于原始 SAM2 记忆注意力模块是基于 2D 自然视频训练的,直接用于 3D 体积会导致时空关联偏差,因此论文在所有 4 层 Memory Attention 的交叉注意力投影中注入 LoRA(\(r=16\)),以极低代价重塑跨帧体积记忆检索能力。训练过程仅使用逐帧 DiceCE 损失,无需显式的追踪图优化损失,即可自然涌现出稳定的四维细胞追踪与分割能力。

损失函数 / 训练策略

训练采用标准的逐切片与逐帧混合 Dice-CrossEntropy 损失(DiceCE)。在 3D 设置下,输入裁剪为中心位于提示点的 \(3 \times 128 \times 128 \times 128\) 体积,采用 AdamW 优化器,学习率设为 \(4 \times 10^{-4}\),配合余弦退火调度训练 500 个 epoch。在 4D 设置下,以目标细胞为中心裁剪 \(3 \times 32 \times 256 \times 256\) 体积块,时序长度 \(T=16\),学习率设为 \(1 \times 10^{-4}\) 训练 500 个 epoch。实验全程在 3 张 NVIDIA RTX 6000 Pro GPU 上使用 BF16 混合精度与梯度检查点训练,3D 体积推理单例耗时仅约 228.5 ms。

实验关键数据

主实验

在 3D CT 肿瘤分割(KiTS 肾脏、Pancreas 胰腺、LiTS 肝脏、Colon 结肠)以及 4D 细胞追踪挑战赛(Fluo-N3DH-SIM+)上,SAM+D 在单点提示设置下全面超越或比肩全监督 3D 网络及现有多参数适配方法,且可训练参数量显著更少。

表 1:四大 3D CT 肿瘤分割基准对比(单点提示 1 pt/vol,报告 Dice (%) 与 NSD (%), tolerance = 2 mm)

方法类别 方法 可训练参数量 KiTS (肾) Dice / NSD Pancreas (胰) Dice / NSD LiTS (肝) Dice / NSD Colon (结肠) Dice / NSD
全监督 3D nnU-Net 30.76 M 73.09 / 77.42 41.70 / 62.92 60.27 / 75.60 44.38 / 53.66
全监督 3D Swin-UNETR 62.19 M 66.48 / 72.44 39.66 / 58.33 52.16 / 65.71 33.53 / 41.16
交互式基线 Visual Sampler 33.67 M 72.47 / 78.25 49.48 / 69.49 52.51 / 64.29 53.49 / 68.05
SAM 适配 SAM-B (逐切片 1 pt/slice) 0 M 36.30 / 29.86 24.01 / 26.74 6.71 / 7.63 28.83 / 33.63
SAM 适配 3DSAM-adapter 25.46 M 80.16 / 87.40 53.69 / 76.44 58.02 / 71.55 59.29 / 76.10
SAM 适配 Med-SA 13.00 M 70.73 / 82.41 43.59 / 73.35 51.61 / 72.04 53.56 / 74.85
SAM 适配 MA-SAM (无提示) 63.04 M 69.25 / 65.84 34.72 / 47.72 56.28 / 62.87 50.32 / 56.37
本文方法 SAM+D (SAM 解码器) 2.57 M 84.74 / 92.11 59.68 / 79.05 63.33 / 73.52 63.70 / 77.69

表 2:4D 细胞追踪挑战赛基准对比 (Fluo-N3DH-SIM+ Seq 02 测试集,报告 DET, SEG, TRA 与 OPCTB (%))

方法 分割来源 可训练参数量 SEG ↑ DET ↑ TRA ↑ OPCTB ↑
BGU-IL CellPose 127.70 M 53.80 85.20 84.20 69.00
Ultrack CellPose 0 M 54.51 85.43 85.17 69.84
TrackStra CellPose 6.90 M 54.15 84.86 84.16 69.16
SAM2 (逐切片) CellPose 0 M 54.19 84.84 84.14 69.16
SAM2+D (本文) CellPose 3.11 M 56.79 87.32 87.15 71.97
SAM2+D (本文) GT 3.11 M - - 99.31 99.47

消融实验

论文在 LiTS 肝脏肿瘤数据集上(单点提示,ViT-B 骨干)进行了细致的模块贡献与超参数消融分析。

表 3:LiTS 验证集消融分析(默认配置带 *)

实验组别 具体设置 Dice (%) NSD (%) 说明
(a) 核心模块累加 DRLoRA only 56.41 65.28 仅深度路由 LoRA
(a) 核心模块累加 + DSM 59.17 68.74 加入深度位移,切片间交互提升 +2.76% Dice
(a) 核心模块累加 + z-embed 61.05 71.03 加入提示相对深度嵌入,提升 +1.88% Dice
(a) 核心模块累加 + Dec. LoRA* (完整模型) 63.33 73.52 解码器注入 LoRA,最终完整模型
(b) DSM 位移比例 \(\alpha\) \(\alpha = 0.125\) 61.22 71.46 比例偏小,切片交互不足
(b) DSM 位移比例 \(\alpha\) \(\alpha = 0.25\)* 63.33 73.52 最佳平衡点,保留 50% 本切片特征
(b) DSM 位移比例 \(\alpha\) \(\alpha = 0.5\) 62.08 72.15 过度平移冲淡本切片本征信息
(c) 专家数 \(K\) \(K = 1\) (标准 LoRA) 58.86 68.21 无深度路由,相比 4 专家落后 4.47% Dice
(c) 专家数 \(K\) \(K = 2\) 61.47 71.30 双专家初步解耦上/下切片
(c) 专家数 \(K\) \(K = 4\)* 63.33 73.52 最佳专家细分颗粒度
(c) 专家数 \(K\) \(K = 8\) 63.10 73.18 参数冗余导致性能轻微饱和回落
(d) LoRA 秩 \(r\) \(r = 4\) / \(r = 8\) 60.72 / 62.85 70.43 / 73.01 低秩表达能力受限
(d) LoRA 秩 \(r\) \(r = 16\)* / \(r = 32\) 63.33 / 62.19 73.52 / 72.34 \(r=16\) 达到最优性价比

关键发现

  • DRLoRA 彻底超越均匀 LoRA 与内容 MoE:将 \(K=1\) 的标准全局均匀 LoRA(58.86%)替换为以归一化深度为条件的 \(K=4\) 专家系统,Dice 瞬间飞跃至 63.33%(净增 +4.47%)。同时补充材料证实,在同等参数预算下,深度标量路由显著领先于基于 token 向量内容的 MoLoRA/MixLoRA 达 +7~9% Dice,印证了解剖结构在深度坐标上的演变属于宏观全局先验,标量路由直接、干净且永不坍缩。
  • DSM 以零成本换取强交互:仅靠前后各平移 25% 通道的内存浅拷贝,不增加任何参数和 MAC 乘加计算,即为模型带来 +2.76% Dice 增益,超越了计算代价成倍增加的 3D 卷积与层间交叉注意力方案。
  • 4D 时空分割突破检测器瓶颈:在 Cellpose 初始检测设定下,所有主流时空算法的分割精度受制于检测器上限死卡在 \(\sim 54\%\) SEG,而 SAM2+D 借助 Memory Attention LoRA 的长程时空记忆线索反哺空间分割,一举突破检测瓶颈达到 56.79% SEG 和 71.97% OPCTB。

亮点与洞察

  • 将 TSM 升维至空间深度轴的极简创新:原先用于视频动作识别的 Temporal Shift Module 第一次被证明可以直接应用于 3D 体积医学图像的深度切片交互。它以 0 参数、0 MAC、0.34ms 的硬件级低开销解决了层间孤立问题,兼具优异的性价比与优雅感。
  • 解耦「空间深度变化」与「时序记忆传递」:在 4D 扩展中,作者没有粗暴地把 DSM/DRLoRA 扩展到时间维度造成混淆,而是保持 DSM 专注空间深度切片融合,将时间维度完全留给 SAM2 预训练的 Memory Bank。这种模块职责的纯粹分离让基础模型预训练能力得以最大化保留。
  • 以宏观深度标量代替高维 Token 内容做 MoE 路由:避开了传统视觉/语言 MoE 中计算复杂、易退化坍缩、必须加 auxiliary loss 调平的通病。标量坐标 \(z \in [0, 1]\) 作为天然且连续的均匀分布输入,让两层 MLP 路由网络稳定高效。

局限与展望

  • 单解剖轴假定:当前 3D 流程默认沿固定扫描体轴(Z 轴)展开,对冠状位、矢状位、横断面三向各向同性体积未做三轴融合预测;未来可探索多轴切片集成与全向 DSM 平移。
  • 4D 复杂分裂事件下的提示依赖:在生物细胞追踪中,频繁的细胞分裂和新视野侵入仍需依赖外部检测器触发新点提示,无法完全摆脱辅助检测器实现端到端闭环。
  • 多模态与超大模型覆盖:目前尚未在最新的多模态概念分割模型 SAM3 上进行验证,其文本-视觉跨模态匹配能否直接兼容 DRLoRA 升维仍待探索。

相关工作与启发

  • vs 3DSAM-adapter: 3DSAM-adapter 引入了较为厚重的 3D 深度可分离卷积适配器,可训练参数高达 25.46 M;SAM+D 仅需 2.57 M 参数即可在 KiTS 上取得 84.74% vs 80.16% 的显著优势,性能与效率双赢。
  • vs Med-SA (SD-Trans): Med-SA 提出空间-深度双分支注意力,深度分支使得自注意力计算量翻倍且需 13.0 M 参数;SAM+D 的 DSM 模块 0 参数 0 MAC 即可实现同等乃至更优的跨切片信息融合。
  • vs 传统 MoE-LoRA (MoLoRA / MixLoRA): 传统方案基于高维特征做路由,需复杂的平衡损失且难以收敛;DRLoRA 基于归一化深度坐标路由,参数微小(0.14K)且从根源上免除了专家坍缩问题。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 巧妙重构 TSM 与深度标量路由 MoE,开创性地以统一架构极低成本实现 SAM 家族从 2D 到 3D 与 4D 的参数高效升维。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 4 个标准 3D CT 肿瘤分割任务和 1 个挑战性的 4D 显微细胞追踪基准,包含详尽消融与深度对比。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严谨清晰,图表完备,消融设计逻辑链极强。
  • 价值: ⭐⭐⭐⭐⭐ 为医学 3D 与生物 4D 领域复用自然图像 2D 视觉基础模型提供了极高参考价值的标准化范式。