医学图像半监督分割新范式:概率原型分类器应对数据不确定性

发布时间:2026/5/26 23:37:00

医学图像半监督分割新范式:概率原型分类器应对数据不确定性 1. 项目概述当医学图像分割遇上“不确定”的像素在医学影像分析这个领域我们这些做算法研发的每天都在和数据打交道。一个核心的痛点就是高质量的标注数据太稀缺了。请一位经验丰富的放射科医生来逐像素标注一张3D心脏MRI耗时耗力成本极高。因此半监督学习Semi-Supervised Learning, SSL自然成为了研究热点——如何用少量标注数据加上大量无标注数据训练出一个靠谱的模型。大家常用的方法比如基于一致性的Mean-Teacher框架或者基于伪标签的自训练确实在利用未标注数据上取得了不错的效果。后来为了应对医学图像中常见的类别不平衡问题比如心肌的像素远少于背景原型学习Prototype Learning被引入进来通过为每个类别学习一个“典型特征”原型来缓解模型对大类别的偏见。但是我和团队在实际研究和复现这些SOTA方法时发现了一个被长期忽视的“幽灵”数据不确定性。这可不是模型参数学得好不好的问题而是数据本身自带的“模糊性”。想象一下心脏左心房和周围组织的边界在MRI图像里常常是模糊的、渐变的没有一个清晰的界限。再比如CT或MRI图像中固有的噪声、伪影都会让某些像素点的特征变得“不可信”。传统的确定性方法把一个像素映射到特征空间的一个“点”上然后计算这个“点”到各个原型“点”的距离来做分类。这种做法对于这些处在模糊地带或带有噪声的“棘手像素”非常不友好模型很容易产生过分割或欠分割如图1所示。我们这篇发表在IEEE TMI上的工作核心就是直面这个“数据不确定性”的挑战。我们提出了一种全新的概率原型分类器其核心思想是不再把像素和原型看作特征空间里一个确定的点而是看作一个概率分布具体是多元高斯分布。这个分布不仅有均值表示最可能的特征还有方差表示对这个特征的不确定程度。一个边界模糊的像素其方差就会很大一个清晰的、典型的组织内部像素方差就会很小。这样一来整个分类过程就从“点对点”的距离比较升级为“分布对分布”的相似度度量。我们引入互似然分数来衡量一个像素分布属于某个原型分布的可能性这个分数同时考虑了特征距离和特征的不确定性。即使一个背景像素因为噪声等原因其特征均值意外地靠近了前景原型只要它的不确定性或者原型本身的不确定性足够高模型依然能凭借概率模型将其正确分类。简单来说我们给模型装上了一双能感知“模糊”和“不确定”的眼睛让它不再对图像中的噪声和模糊边界“视而不见”从而在标注数据极少的情况下依然能做出更鲁棒、更精确的分割。下面我就来拆解这个框架的每一个技术细节并分享我们在实现和调参过程中的实战经验。2. 核心框架设计从确定性到概率性的范式转变我们的整体框架建立在经典的Mean-Teacher架构之上这是一个在半监督学习中非常有效的“师生”框架。教师模型的权重是学生模型权重的指数移动平均因此更加稳定。教师模型对弱增强图像进行推理生成相对可靠的伪标签学生模型则在强增强图像上训练并同时接受真实标签和伪标签的监督。这种“弱-强”增强策略能有效提升模型的泛化能力。框架的创新核心在于分类头。我们并行使用了两个分类器线性分类器一个传统的卷积层主要职责是分割那些特征明显、确定性高的区域。概率原型分类器我们提出的新模块专门负责处理那些具有高不确定性的“棘手像素”。这种双分类器设计是一种巧妙的“分工协作”。线性分类器简单高效能快速处理大部分明确区域而概率原型分类器则像一位“疑难杂症专家”专注于解决边界模糊和噪声像素的分类问题。两者的预测通过损失函数进行融合共同指导学生模型的学习。2.1 确定性原型的局限与概率原型的优势在深入我们的概率原型之前有必要先看看传统确定性原型做了什么以及它在哪里遇到了瓶颈。在确定性原型方法如PCR中对于一个类别c其原型是通过简单平均该类所有像素的特征向量得到的p_c (1/N) * Σ z_i。这里z_i是像素x_i经过编码器网络后得到的确定性特征向量。分类时计算像素特征z_i与每个类别原型p_c的余弦相似度相似度最高的类别即为预测结果。这个过程的致命弱点在于“平等对待”。它假设所有属于类别c的像素贡献是均等的。但在现实中一个位于清晰心肌区域的像素和一个位于心内膜边界模糊地带的像素它们的特征可靠性能一样吗显然不能。后者带有更高的不确定性。在平均过程中这些高不确定性像素的“噪声”特征会被带入原型污染原型的代表性。更糟糕的是在分类时一个高不确定性的模糊像素其特征向量z_i可能因为噪声干扰而偏离其“本应”在的位置如果这个偏离恰好使它更靠近另一个类别的原型就会导致误分类。我们的概率原型分类器从根源上改变了这一范式像素表示为分布p(z|x_i) N(z; μ_i, diag(σ_i^2))。我们通过网络预测两个值均值μ_i最可能的特征和对角协方差矩阵diag(σ_i^2)各特征维度的不确定性。方差越大表示网络对这个像素的特征越“没把握”。原型更新为分布原型也不再是一个点而是一个分布ρ_c ~ N(z; û_c, diag(σ̂_c^2))。关键来了它的均值和方差不是简单平均而是基于贝叶斯估计的加权平均û_c Σ (σ̂_c^2 / σ_i^2) * μ_i1/σ̂_c^2 Σ 1/σ_i^2这个公式非常直观不确定性低方差小的像素在构建原型时权重更大不确定性高方差大的像素权重被自动降低。这就好比在专家评审中资深专家的意见权重更高一样。原型因此更能代表该类别的“可靠”特征。2.2 互似然分数概率匹配的“裁判”有了概率化的像素和原型如何衡量它们的匹配程度这里我们引入了互似然分数。它的思想很巧妙计算像素分布z_i和原型分布ρ_c相等的概率p(z_i ρ_c)。通过数学推导见原文公式5-6我们可以得到其对数形式的计算公式MLS(z_i, ρ_c) -1/2 * Σ_d [ (μ_i^(d) - û_c^(d))^2 / (σ_i^2(d) σ̂_c^2(d)) log(σ_i^2(d) σ̂_c^2(d)) ] C这个公式蕴含了两个强大的机制注意力机制第一项(μ_i^(d) - û_c^(d))^2 / (σ_i^2(d) σ̂_c^2(d))。它不再是简单的欧氏距离。分母是双方不确定性的和。如果某一维度上像素或原型的不确定性很大方差大那么即使均值差很大这一项的值也会被“稀释”变小。这意味着模型在匹配时会自动“忽略”那些不确定性高的、不可靠的特征维度而更关注那些双方都确定的维度。惩罚机制第二项log(σ_i^2(d) σ̂_c^2(d))。这一项直接惩罚高不确定性。只要像素或原型在某一维度上不确定性高这一项就会贡献一个很大的正值因为log函数在自变量大于1时单调增从而降低总的MLS分数。一个像素即使其均值非常靠近某个原型但只要它自身或该原型不确定性极高其MLS分数也会很低从而不会被轻易归入该类。实操心得理解MLS的物理意义你可以把MLS想象成一场“信任投票”。两个分布要匹配需要满足1. 它们的“中心思想”均值接近2. 它们对自己的“观点”都足够确信方差小。如果一方闪烁其词方差大即使它说的内容均值和另一方很像我们也不能完全相信它们是一致的。MLS通过数学公式完美地量化了这个“信任”过程。最后通过softmax函数将MLS分数转换为类别概率p_PPC(y_ic|x_i) exp(MLS(z_i, ρ_c)/τ) / Σ_c‘ exp(MLS(z_i, ρ_c’)/τ)其中τ是温度超参数用于控制概率分布的平滑程度。3. 实现细节与训练策略3.1 网络结构与实现要点我们的框架以经典的编码器-解码器网络如V-Net用于3DU-Net用于2D作为骨干特征提取器f。在此之上我们增加了两个预测头表征头 h(·)通常由几层卷积Conv、批归一化BatchNorm和激活函数ReLU堆叠而成负责从特征图中预测每个像素的D维均值向量μ。不确定性头 q(·)这是实现概率化的关键。它的输入同样是骨干网络提取的特征输出是D维的方差倒数1/σ^2为了训练稳定性我们预测方差的倒数。这个头通常设计得比表征头更深或更复杂一些以确保它能学习到有意义的、与特征相关的“不确定性”信号。在我们的实现中q(·)包含两个卷积操作后接两个在线性变换层作用在通道维度上和一个SoftPlus激活层确保输出为正。为什么预测方差倒数直接预测方差σ^2在训练中可能面临数值不稳定的问题因为方差需要始终为正且梯度可能爆炸。预测其倒数1/σ^2精度则更稳定。在计算时取其倒数即可得到方差。3.2 分布原型的更新策略贝叶斯估计原型不是一成不变的需要在训练过程中用新看到的数据不断更新。对于点原型常用动量更新p_c_new α * p_c_old (1-α) * p_c_current。但对于我们的分布原型简单的线性插值就不合适了。我们提出了基于贝叶斯估计的更新策略。假设我们已经有了基于前t次迭代数据累积得到的原型分布ρ_c_t ~ N(û_t, σ̂_t^2)以及基于最新第tT次迭代数据计算出的增量原型分布ρ_c_tT ~ N(û_tT, σ̂_tT^2)。那么融合了新旧信息的新原型分布为1/σ̂_new^2 1/σ̂_t^2 1/σ̂_tT^2û_new σ̂_new^2 * (û_t/σ̂_t^2 û_tT/σ̂_tT^2)这个公式的美感在于它的不确定性加权。新的均值û_new是旧均值和新增量均值的加权平均权重是各自方差的倒数即精度。这意味着更确定方差小的原型信息在更新中占据主导地位。如果新增量原型是基于一批非常不确定的像素计算的σ̂_tT^2很大那么它对最终新原型的贡献就会非常小从而防止噪声信息污染已建立好的可靠原型。3.3 采样规则与“懒惰更新”技巧用所有像素来更新原型在计算和内存上都是不现实的。更重要的是我们相信具有适当可靠性和合理模糊性的像素对构建有代表性的原型最有益。过于可靠的像素方差极低可能只代表一种非常特定的模式比如一片均匀的背景用它们更新原型会导致原型多样性不足。过于模糊的像素方差极高特征不可靠会引入噪声。因此我们设计了一个基于方差的采样规则。对于候选像素集合X我们计算所有像素方差的分布。我们只选择那些方差位于[δ_s, δ_w]百分位数之间的像素。δ_s是强阈值例如第10百分位过滤掉过于确定的像素δ_w是弱阈值例如第80百分位过滤掉过于不确定的像素。这样我们就能保留那些“既不太过确定以至于模式单一又不太过模糊以至于不可信”的像素用于原型更新。在训练初期我们发现如果每个迭代都更新原型T1会出现“更新瓶颈”累积原型的方差σ̂_t^2会迅速变得非常小因为每次更新都倾向于降低不确定性而新增量原型的方差σ̂_tT^2相对较大。根据贝叶斯更新公式σ̂_new^2会变得极小且û_new中来自新增量的权重σ̂_new^2/σ̂_tT^2会趋近于0导致新信息无法被有效纳入原型更新停滞。为了解决这个问题我们提出了一个经验技巧——“懒惰更新”。我们设置一个较长的更新周期T例如T200而不是每轮迭代都更新。这样做有两个好处拉长了更新间隔使得累积原型的方差σ̂_t^2不会下降得过快。经过较长时间的训练后模型处理数据的能力增强新增量原型的方差σ̂_tT^2也会降低。 这样两者方差的差距缩小新增量信息就能持续、有效地被融入到原型中。虽然叫“懒惰”但这是一种以退为进的策略保证了原型更新的持续性和稳定性。3.4 损失函数设计总损失函数由监督损失L_l和无监督损失L_u加权组成L L_l λ_u * L_u。λ_u是一个随训练轮数增加的sigmoid斜坡上升系数用于在训练初期侧重有标签数据后期逐步加大无标签数据的权重。监督损失针对有标签数据我们同时计算线性分类器和概率原型分类器的预测与真实标签的损失交叉熵损失和Dice损失的平均并求和。无监督损失针对无标签数据我们使用教师模型PPC分支生成的伪标签ŷ_i来监督学生模型。这里有一个关键的质量控制我们只使用那些不确定性低于一定阈值例如方差位于该图像所有像素方差的前κ%的像素预测来生成伪标签和计算损失。这确保了用于训练的伪标签是相对可靠的避免了错误标签的累积传播。4. 实验验证与结果分析我们在三个公开医学图像分割数据集上进行了全面评估左心房分割数据集LA、心脏多结构分割数据集ACDC和前列腺分割数据集PROMISE12。我们设置了严格的半监督场景例如在LA数据集中仅使用100个扫描中的4个或8个有标签数据进行训练。4.1 与SOTA方法的对比我们将方法分为三类进行对比基于一致性的方法如MT, UA-MT、自训练方法如ST以及混合方法如BCP, PCR。我们的方法属于混合方法。定量结果LA数据集在仅使用8个有标签样本8%标注率时我们的方法在Dice和Jaccard系数上取得了所有方法中的最佳性能在边界距离指标95HD, ASD上与最强的基线方法BCP相当。ACDC数据集在4个和7个有标签样本的设置下我们的方法在Dice和Jaccard系数上均显著优于所有对比方法包括BCP。特别是在极度稀缺的4标签设置下优势更为明显这证明了我们的方法在标签极度匮乏时处理数据不确定性的强大能力。PROMISE12数据集在20%有标签数据设置下我们的Dice分数比第二名PRCL高出超过3%证明了其在2D切片数据上的泛化能力。定性结果可视化 图5、6、7的视觉对比非常直观。在LA数据中其他方法在左心房边界模糊处容易出现过度分割或分割不足而我们的方法能更好地贴合真实的解剖形状。在ACDC数据中面对类内异质性如右心室内部纹理、类间同质性不同组织灰度相似和严重类别不平衡心肌较薄我们的方法能成功分割出小目标如第一行中的小结构并能更好地区分前景和背景。在PROMISE12数据中我们的方法对前列腺大小变化、边界模糊和强度变化等挑战表现出了更强的鲁棒性。4.2 消融实验为什么概率化有效概率表示 vs. 确定表示我们将我们的PPC与其确定性版本PCR将我们的概率模块替换为点原型和余弦距离进行对比。在相同的采样阈值下PPC几乎在所有设置下都优于PCR。这直接证明了将不确定性信息显式建模为概率分布的有效性。我们还测试了不同的采样阈值(δ_s, δ_w)发现同时过滤掉最不确定和最确定的像素如(10, 80)能取得最佳效果这印证了“适度模糊”的像素对原型学习最有价值。贝叶斯更新 vs. 动量更新用我们提出的贝叶斯估计更新策略替换标准的动量更新后模型性能在所有指标上均有提升尤其在标签数据极少时4标签提升更显著。这说明在更新时考虑信息的不确定性精度加权至关重要它能防止不可靠的增量信息污染已建立的良好原型。懒惰更新的影响我们测试了不同更新周期T。当T1每轮更新时性能最差出现了“更新瓶颈”。当T增大到100或200时性能显著提升。但当T300时性能又有所下降因为更新频率过低原型无法及时捕捉数据分布的变化。T200是一个较好的平衡点既能保证信息持续融入又能保持原型稳定。4.3 概率行为的深入探究为什么概率表示能提升性能对抗过自信深度神经网络常有“过自信”问题即使对错误预测也会给出很高的softmax分数。PCR确定性方法就受此困扰无法有效识别不可靠预测导致错误伪标签积累。而我们的方法预测的方差σ^2是一个更可靠的不确定性指标。如图11所示高方差区域亮黄色精准地定位在模糊边界上。利用这个指标过滤高不确定性预测能生成更高质量的伪标签从而提升最终分割结果。数据不确定性 vs. 模型不确定性我们还对比了另一种不确定性估计方法UA-MT使用蒙特卡洛Dropout估计模型不确定性。在有限标签下模型可能因数据不足而无法准确校准其估计的模型不确定性如预测熵可能不可靠。如图11(b)所示UA-MT对边界上的某些像素给出了低不确定性箭头所指。而我们的数据不确定性估计方差更稳健即使模型参数未完美校准它也能基于像素自身特征提供有意义的估计边界像素始终表现出高方差。概率表示和分布原型如何工作我们可视化了训练过程中原型和像素的均值、方差分布变化图12。发现了两个理想特性可区分性不同类别的原型均值在特征空间中是明显分离的并且与同类像素的均值分布一致。不确定性差异前景类如器官的原型方差通常高于背景类这符合事实——前景类包含更多边界模糊像素。在匹配时特性一保证了当像素特征距离能清晰判别时MLS有效特性二保证了当两个类别的像素特征距离相近时MLS能通过考虑不确定性方差来做出正确决策。而在确定性表示中这种情况只能随机分配。4.4 对噪声数据的鲁棒性除了边界模糊成像伪影也是数据不确定性的重要来源。我们在LA数据上模拟了散斑噪声和高斯模糊来测试模型鲁棒性。如图13所示随着噪声程度增加所有方法的性能都会下降但我们的方法下降幅度最小且在不同噪声水平下都保持了最高的Dice分数。这表明通过概率建模显式地捕捉像素级不确定性使模型对输入噪声具有更强的内在鲁棒性。5. 局限性与未来展望这项工作是将概率论思想引入像素级表示、以应对数据不确定性的初步尝试。虽然实验证明了其有效性但仍存在一些局限和可改进空间原型表示能力我们目前将每个类别的原型建模为单一的多元高斯分布。这可能过于简化无法充分捕捉一个类别内部可能存在的多个子模式例如心脏左心房的不同部位、不同形状。一个自然的扩展是使用高斯混合模型来表征每个类别的原型使其能建模更复杂的类别内部分布。互似然分数的边界问题MLS分数是无界的这迫使我们在softmax中引入温度参数τ来将其映射到合理的概率范围。τ的选择需要根据数据集进行调整并非易事。一个潜在的解决方案是绕过直接计算分布相等的概率采用蒙特卡洛采样的方法从像素分布和原型分布中分别采样若干特征向量然后计算这些样本之间的平均相似度如余弦相似度用这个经验估计值作为匹配分数。计算开销概率化表示和MLS计算不可避免地会增加一些计算量。尽管在原型更新中我们采用了采样和懒惰更新来缓解但在处理超高分辨率图像或实时应用场景时仍需进一步优化效率。从更广的视角看这项工作为半监督学习特别是医学图像分析开辟了一条新路径将感知不确定性作为模型的内在能力而不仅仅是事后的估计工具。未来我们可以探索将这种概率化思想与更先进的网络架构如Transformer、更复杂的分布族如归一化流相结合进一步提升模型在极端数据稀缺和高度不确定性环境下的性能。对于临床部署而言一个能“知道自己哪里不知道”的模型其预测结果将更具可信度和实用价值。

相关新闻