上下文至关重要:医学图像上下文分割中的支持集选择与失败检测

快速导读:研究支持集选择策略并训练分类器预测分割失败,以提升医学图像上下文分割的可靠性。

上下文学习(In-context learning, ICL)使分割模型能够在推理时通过少量标注示例适应新任务,无需重新训练,在医学图像分析中展现出巨大潜力。然而,其性能高度依赖作为唯一任务信号的支持集(support set)构成——不匹配的支持集会严重降低分割质量。现有研究多聚焦于模型架构设计,对支持集选择策略和分割失败预测的探索严重不足。

本文以MultiverSeg为上下文分割模型、DINOv3为图像编码器,在四个医学图像分割基准(EchoNet、WBC、HipXRay-Femur、HipXRay-Pelvis)上系统比较了随机采样与基于余弦相似度的支持集选择策略,并训练了一个Transformer分类器,仅以查询和支持图像的嵌入为输入,预测分割结果的IoU是否低于预设阈值,从而实现推理前的失败检测。

英文题目:Context Matters: Support Set Selection and Failure Detection for In-Context Medical Image Segmentation

论文出处:arXiv 每日论文精选 · arXiv:2608.05333

原始论文:PDF / 论文页面

这篇论文解决什么问题?

上下文分割的核心挑战在于:模型在推理时完全依赖支持集来理解任务,而支持集的质量直接影响分割结果。当支持图像与查询图像在解剖结构、成像模态或病理表现上不匹配时,模型可能产生严重错误的分割。这一问题在临床场景中尤为关键,因为错误分割可能导致误诊或不当治疗决策。

现有工作主要关注如何设计更强大的上下文分割模型,但对“如何选择支持集”这一上游问题缺乏系统研究。虽然Suo等人和Zhang等人分别探讨了视觉提示选择和有效示例构成,但均未在医学图像领域进行验证,也未将支持集选择与失败检测联合考虑。

此外,当前上下文分割系统缺乏在推理前预测分割失败的机制。临床部署要求系统不仅能分割,还要能主动识别“我可能做错了”的情况,以便触发人工审核或拒绝输出。这一可靠性筛查机制的缺失,构成了上下文分割从研究走向临床的关键障碍。

本文正是针对上述两个缺口展开:一方面量化不同支持集选择策略的性能差异,另一方面构建失败检测分类器,为安全部署提供双重保障。

核心创新

  • 首次系统比较了随机采样与基于相似度的支持集选择策略在医学图像上下文分割中的性能,并量化了不同支持集大小下的增益。
  • 提出了一种基于Transformer的分类器,仅利用查询和支持图像的嵌入(无需生成掩膜)即可预测分割失败,为临床部署提供了主动的可靠性筛查机制。

方法概览

1. 比较随机采样与基于相似度的支持集选择策略(使用DINOv3提取图像嵌入,通过余弦相似度检索最相似的K个支持样本)。2. 训练一个基于Transformer的分类器,仅以查询和支持图像的嵌入为输入,预测分割结果的IoU是否低于预设阈值,从而实现分割失败检测。

  • 支持集选择策略对比:随机采样从支持数据集中均匀选取K个样本,重复N=40次取平均IoU以消除随机性;相似度采样使用DINOv3提取图像嵌入,通过余弦相似度检索与查询最相似的K个支持样本,该策略为确定性方法,仅需单次采样。
  • 失败检测架构设计:采用Transformer编码器作为分类器主干。首先由预训练图像编码器(DINOv3)提取查询和支持图像的嵌入,然后为查询嵌入和学习到的“查询”角色嵌入相加,为支持嵌入和学习到的“支持”角色嵌入相加,以区分两者在上下文中的不同功能。
  • 角色嵌入与分类头:组合后的嵌入连同可学习的[cls] token一并输入标准Transformer编码器。编码器通过自注意力机制建模查询与支持图像之间的关系,最终由[cls] token经分类头输出二分类结果——预测分割IoU是否低于预设阈值τ。
  • 失败标签定义:以数据集特定的中位IoU作为阈值τ,将低于该阈值的分割结果标记为“失败”。这一相对阈值定义避免了绝对IoU标准在不同数据集间不可比的问题。
  • 训练数据构建:基于随机采样的分割结果构建失败检测数据集,确保训练数据覆盖多样化的成功与失败案例。分类器在训练时仅使用图像嵌入,无需访问分割掩膜或真实标注。
  • 评估指标:采用AUROC和AUPRC评估失败检测性能,这两个指标能全面反映分类器在不同阈值下的表现,尤其适用于类别不平衡的失败检测场景。

逐图理解论文

失败案例与直觉

失败案例与直觉
Fig. 3. In-context segmentation performance for random and similarity- based sampling. Box plots of IoU score distributions across four benchmarks (WBC, HipXRay-Femur, EchoNet, and HipXRay-Pelvis) as a function of support set size K ∈{1, 2, 4, 8, 16, 32} for random (green) and similarity-based (blue) selection. Perfor- mance improves with K under both strategies. Similarity-based selection consistently matches or outperforms random sampling, yielding the largest gains at small K where the support set is most constrained.

图3以箱线图形式展示了四种基准上两种采样策略在不同支持集大小下的IoU分布。蓝色为相似度采样,绿色为随机采样。可观察到两个一致趋势:性能随K对数增长,且相似度采样在K较小时优势明显,在K较大时两者趋同。

研究空白与核心贡献

研究空白与核心贡献
Fig. 2. Overview of the proposed in-context segmentation failure detection architecture. We use a transformer-based classifier to predict segmentation failure, defined as a binary indicator of whether the resulting segmentation falls below a pre- defined IoU threshold τ. First, a pre-trained image encoder (e.g., DINOv3) extracts embeddings for the query and support images. To distinguish between these roles, we add learned query and support embeddings to their respective image embeddings. These combined embeddings, alongside a learned class token ([cls]), serve as inputs to a standard transformer encoder. Finally, the [cls] output token passes through a classification head to predict the failure label.

图2呈现了失败检测分类器的完整架构。查询和支持图像首先经DINOv3编码器提取嵌入,然后分别加上可学习的查询角色嵌入和支持角色嵌入以区分功能。组合后的嵌入与[cls] token一同输入Transformer编码器,最终由分类头输出失败预测。

失败检测的可行性与局限

失败检测的可行性与局限
Table 1. Evaluation metrics for in-context segmentation failure detection on the test set. Overall classification metrics for the failure classifier across all four benchmarks, pooled across support set sizes. The IoU threshold (τ) denotes the dataset- specific median IoU used to define failure. The AUROC exceeds chance (0.50) on every dataset, demonstrating that segmentation failure carries a strong signal recoverable from the query and support embeddings alone.

表1汇总了四个基准上失败检测分类器的整体性能指标。AUROC在所有数据集上均超过0.50的随机水平,其中EchoNet表现最佳(AUROC: 0.80, AUPRC: 0.79),证明分割失败信号可从查询和支持嵌入中恢复。

实验如何设计?

  • 数据集:四个医学图像分割基准——EchoNet(心脏超声)、WBC(白细胞显微镜图像)、HipXRay-Femur(髋关节X光-股骨)和HipXRay-Pelvis(髋关节X光-骨盆),覆盖不同成像模态和解剖结构。
  • 上下文分割模型:MultiverSeg,一个专为医学图像设计的上下文分割模型,在推理时接收查询图像和支持集(图像-掩膜对),输出查询图像的分割掩膜。
  • 支持集大小:K∈{1,2,4,8,16,32},覆盖从极少示例到较丰富示例的完整范围,以观察性能随支持集大小的变化趋势。
  • 图像编码器:DINOv3,一个大规模预训练的视觉Transformer,用于提取查询和支持图像的嵌入表示,同时服务于相似度检索和失败检测分类器。
  • 失败检测分类器训练:使用随机采样策略在训练集上生成分割结果,以中位IoU为阈值构建二分类标签,训练Transformer分类器,在独立测试集上评估。

关键结果与论文证据

  • 相似度采样在低样本场景下优势显著:当K=1时,相似度采样相比随机采样在EchoNet上中位IoU提升0.16,在WBC上提升0.14(第7页)。这一增益在临床资源稀缺场景中具有重要价值。
  • 两种策略的性能差距随K增大而缩小:当K≥16时,相似度采样与随机采样的中位IoU趋于一致,表明在支持集足够大时,随机采样已能提供足够的任务信息(第6页,图3)。
  • 失败检测在所有基准上均超过随机水平:分类器的AUROC在四个数据集上均高于0.50,在EchoNet上达到最高(AUROC: 0.80, AUPRC: 0.79),证明分割失败信号可从图像嵌入中恢复(第7页,表1)。
  • 失败可预测性随支持集增大而提升:在EchoNet、WBC和HipXRay-Femur上,AUROC随K增大而上升,HipXRay-Femur从K=1时的0.60提升至K=32时的0.92(第8页,图4)。
  • HipXRay-Pelvis呈现异常模式:在该数据集上,失败检测AUROC在大支持集下反而下降至接近随机水平,表明方法的有效性可能依赖数据集特性(第8页,图4)。
  • 性能提升呈对数趋势:无论采用何种采样策略,中位IoU随K对数增长,四分位距收窄,说明增加支持集带来的边际收益递减(第6页,图3)。

阅读时需要注意

  • 模型与编码器单一性:仅评估了MultiverSeg和DINOv3的组合,未验证方法在其他上下文分割模型或图像编码器上的泛化性。
  • 未利用掩膜信息:支持集选择和失败检测均仅依赖图像嵌入,忽略了支持掩膜中蕴含的目标结构形状、大小等关键特征,可能限制性能上限。
  • 阈值定义简化:失败标签基于单一的中位IoU阈值,未考虑任务特定的临床可接受标准,且实验仅使用一次固定数据划分,未评估分布外偏移下的鲁棒性。

关联工作

  • MultiverSeg(第5页):本文使用的上下文分割模型,通过在支持集上条件化来适应新分割任务,代表了当前医学图像上下文分割的主流范式。
  • DINOv3(第5页):用于提取图像嵌入的预训练编码器,其自监督学习得到的表示在多个视觉任务上展现出强大的泛化能力,为相似度检索和失败检测提供了高质量特征基础。
  • Suo等人的视觉提示选择工作(第2页):研究了上下文分割中的提示选择策略,但未涉及医学图像领域,也未与失败检测结合。
  • Zhang等人的有效示例研究(第2页):探讨了视觉上下文学习中有效示例的构成要素,为本文的支持集选择研究提供了动机,但缺乏系统性的量化比较。

发表评论