DPRD 用 ROI 掩码与位移向量对齐蒸馏,让仅占教师约5%参数的学生在 AMOS 上以 85.46% Dice 超过 MedNeXt 教师
核心概要
该工作提出位移保持关系蒸馏(DPRD),在 nnU-Net 中对 ROI 掩码后的病例级嵌入做批量成对位移向量对齐并配合关系尺度归一化,在 ISLES 2022 与 AMOS 2022 上以约5%教师参数和约3% FLOPs 的学生取得优于 Logits KD、FitNet、RKD、CIRKD 的分割结果,AMOS 上 Dice 达 85.46%、HD95 降至 11.72 mm。
Fig. 1. Overview of DPRD. RAFM obtains ROI-aware case-level embeddings, and DPRA aligns pairwise displacement relations between teacher and student representa- tions.
· 第 3 页深度剖析
提出位移保持关系蒸馏,用批量成对位移向量加关系尺度归一化,施加尺度不变的局部关系约束,而非匹配绝对激活。 相较把样本关系压缩为标量距离或角度的 RKD,以及为 2D 图像构建稠密像素级关系的 CIRKD,DPRD 在紧凑的病例级嵌入上建模高维位移方向与相对尺度。 方法在 ISLES 2022 与 AMOS 2022 上与 Logits KD、FitNet、RKD、CIRKD 在同一 nnU-Net 训练设置下对比,并给出消融。
设计跨架构、ROI 感知的多阶段蒸馏方案,对 3D 体数据内存友好,并显式缓解背景主导的监督稀释。 ROI 感知特征掩码(RAFM)用解剖先验把对齐聚焦到任务相关结构,掩码仅在训练阶段使用,推理时移除教师与投影层。 消融显示去掉 RAFM 后 AMOS Dice 从 85.46% 降至 81.68%、HD95 升至 21.65 mm。
在 ISLES 2022 与 AMOS 2022 上一致提升紧凑学生模型,边界敏感指标改善尤为明显。 ISLES 上 DPRD 取得最佳 Dice 75.16%、NSD 85.92%、HD95 13.08 mm,HD95 优于 MedNeXt 教师的 14.12 mm;AMOS 上取得 85.46% Dice、85.29% NSD、11.72 mm HD95,Dice 略高于教师的 85.37%,HD95 由 14.52 mm 降至 11.72 mm。 结果基于 nnU-Net 标准验证划分,ISLES 为 200 训练/50 验证病例,AMOS 为 240 训练/60 验证病例,指标含均值与标准差。
消融表明方向轨迹重建与相对距离一致性提供互补监督。 分别去掉轨迹映射或距离正则化都会使 HD95 升至 15.03 mm 与 17.07 mm,完整框架取得最佳 11.72 mm。 消融在 AMOS 2022 的 MobileUNetV3 学生上进行,逐项开关 RAFM、Traj.、Dist. 三个组件。
启示与展望
该结果面向需要在算力受限环境中部署 3D 分割器的场景,例如图像引导流程中的目标与风险结构定位;方法以 nnU-Net 为载体,教师为冻结的 MedNeXt,学生为 MobileUNetV3 与 PlainConvUNet,推理时移除教师与投影层,因此部署模型参数与 FLOPs 与学生一致。适用对象是具备训练期体素级标注的 MRI 卒中病灶与腹部多器官分割任务,作者也指出 RAFM 依赖真值标签掩码,未来拟探索教师预测 ROI 或不确定性引导掩码以放宽该要求。
读者仍需留意:ROI 掩码依赖训练期真值标签,在弱监督或标签稀缺场景的适用性尚未验证;评估限于 ISLES 2022 与 AMOS 2022 两个基准,跨中心与跨扫描仪的外观漂移是否被完全覆盖仍待更多数据检验;批量位移约束是随采样病例、增强块、编码器阶段、迭代与轮次累积的随机局部监督,其在不同批量规模下的稳定性值得关注;此外,本文为全文解析,但公式与部分表格在文本中呈现不完整,若需复现具体归一化与损失细节应查阅原文与代码。
