最近在做一个半监督医学图像分割项目卡在标注数据太少、模型总是过拟合的坎上试了不少对比学习方案效果都不理想。后来偶然想到一个思路既然大规模预训练的生成模型尤其是Stable Diffusion那套已经把视觉语义压缩得很好那能不能把它的VAE编码器直接借过来用来做医学图像的特征提取再配合一个双分布对齐的约束让少量标注数据和大量未标注数据在潜在空间里互相校准结果意外地顺最后在某个公开腹部器官数据集上只用了20%的标注量就追平了全监督基线。这篇文章就把这个“冻结SD-VAE双分布对齐”的生成式半监督分割方案完整拆开讲讲包括为什么这么设计、每一步怎么落地、损失怎么配以及我踩过的那些坑。1. 方案的整体设计思路1.1 为什么敢打SD-VAE的主意先说一个直觉。医学图像分割的痛点从来不是模型结构不够花哨而是标注成本极高。一张CT可能要请放射科医生花一两个小时勾轮廓一个batch的数据标注下来人力开销远超普通自然图像。所以半监督学习的核心诉求永远是怎么把那些没有标签的图像也充分利用起来。传统的半监督方案比如伪标签、一致性正则、对比学习思路都很好但都有一个共同的问题——需要从零开始或者从ImageNet预训练开始去学特征。而医学图像和自然图像差异巨大ImageNet的特征迁移过来常常水土不服。更麻烦的是医学图像本身的标注样本太少想靠这些样本把特征学好很容易陷入严重的过拟合。这时候SD-VAE的价值就出来了。Stable Diffusion的VAE编码器是在海量图文数据上训练过的它的潜在空间对纹理、结构、语义轮廓有非常强的表达能力。尽管它没见过医学图像但它学到的底层视觉结构——边缘、纹理基元、形状渐变——是通用的。更重要的是我们不需要微调它直接冻结把它当作一个固定且强大的特征提取器把医学图像映射到一个已经很好组织的潜在空间里。在这个空间里做分割相当于站在巨人的肩膀上干精细活。1.2 “双分布对齐”到底在干什么冻结SD-VAE只是第一步。真正让半监督跑起来的是“双分布对齐”这个设计。我把它叫作“双分布”是因为这里有两组分布需要对齐。第一组是有标注数据的特征分布和无标注数据的特征分布。理想情况下如果两者在潜在空间里分布一致那么从有标注数据学到的分割头就能自然地泛化到无标注数据上。第二组是特征分布和先验分布。因为SD-VAE的潜在空间本来就被约束成近似高斯分布所以我们可以让模型提取的特征也去贴近这个先验起到一个正则化的作用。这两组对齐合在一起就有了一种很微妙的效果有标注数据提供“锚点”无标注数据通过分布对齐被拉向同一个区域而先验约束又防止所有特征挤成一团或者漂移到无意义的位置。实验做下来这个组合的效果比单做任何一个都强不少。1.3 和常见半监督方案的差异化市面上常见的半监督分割无非几类。伪标签方案最直接但问题在于教师模型训得不够好时错的标签也会被当作“真理”学进去错误会累积。一致性正则方案比如在扰动下保持输出不变效果好但对扰动方式极其敏感医学图像里某些扰动比如小尺度弹性变形反而会破坏病变区域的结构语义。对比学习方案也很流行但它需要构造正负样本对负样本数不够或者构建不当模型很容易学到的只是“特征不同”而抓不到“同一结构语义”。我们这个方案就不太一样——它不依赖伪标签的置信度不需要手工设计复杂的扰动也不需要显式构造负样本。它靠的是生成式模型潜在空间里的分布形态来传递信息。一句话概括让先验替你说话让分布替你校正。2. 数据预处理与适配细节2.1 医学图像的“三通道化”SD-VAE的编码器是为三通道RGB图像设计的而医学图像通常是单通道灰度图CT、MRI。直接复制成三通道当然可以但我强烈不建议——因为SD-VAE对通道间的相关性是有刻板印象的比如自然图像里RGB三个通道的统计特征跟医学图像复制三份是完全不同的。我用的是一个更稳的做法把灰度图复制成三通道但分别做不同的窗宽窗位调整。比如CT图像我可以把软组织窗、骨窗、肺窗分别映射到R、G、B三个通道。这样不仅适配了SD-VAE的输入要求还顺带做了多窗信息融合一举两得。实测下来比直接复制三通道在分割任务上提升约2个点的Dice。注意并不是所有医学图像都需要窗宽窗位。如果是MRI更建议用T1、T2、ADC等不同序列的配准结果来填三通道信息量更大。2.2 尺寸与归一化的坑SD-VAE的原生输入分辨率是512x512但医学图像分割任务往往没有那么大常见256x256。这里关键问题不是分辨率而是归一化。自然图像的归一化用的是ImageNet均值和方差mean[0.485,0.456,0.406], std[0.229,0.224,0.225]但这是基于RGB三通道的统计量。如果直接把CT灰度图复制三通道然后套这个归一化相当于强迫SD-VAE看到一个它从未见过的输入分布编码出来的特征往往会退化。我的处理方式是先做医学图像本身的标准化比如CT的窗宽窗位处理min-max归一化然后再套ImageNet归一化把它作为第二层归一化。实操里我给这个流程起了个名字叫“双重归一化策略”核心是先让医学图像的数值分布回到[0,1区间和自然图像类似的量级再进SD-VAE标准的预处理逻辑。2.3 切片采样的策略医学图像本质是三维体数据但SD-VAE是二维的所以必须切片。这里有个取舍横断面还是冠状面单一方向会不会丢失空间信息我采用的是“2.5D方案”——对每个体素取横断面、冠状面、矢状面三个方向的切片分别过SD-VAE编码器然后把三个方向的潜在特征在通道维度上拼接再送进分割头。这样既避免了直接用3D模型带来的显存爆炸又保留了三维上下文。不过要提醒一下2.5D的代价是编码时间变成三倍推理也要三倍计算量。如果临床场景对时间敏感可以退化成单一横断面Dice的下降大约在1-3个点视器官而定。3. 核心方法实现3.1 架构总览整个方案分成三个部分冻结的SD-VAE编码器、一个轻量分割头、两个分布对齐模块。SD-VAE编码器负责把输入维度为 [B,3,H,W] 的图像映射成潜在特征 [B,4,H/8,W/8]这其实就是Stable Diffusion的VAE压缩后的特征图空间分辨率缩小8倍通道数变成4。这个特征就是分割网络的主干输入。分割头我选择的是一个轻量化的U-Net变体只不过它的第1层输入不再是原始图像而是SD-VAE出来的特征图。分割头内部包含跳连、注意力、上采样模块最终输出与输入分辨率相同的分割概率图。分布对齐模块分为两条标注分支把标注图像的潜在特征用分割头的编码器进一步提炼成高维语义特征和分割头的深层特征做对齐。无标注分支只前传特征不计算分割损失但计算特征与标注分支特征以及先验分布的分布距离。3.2 为什么要“冻结”这一步可能会有人质疑为什么不去微调SD-VAE编码器反正数据有这么多未标注的微调不是更好吗我一开始也是这么想的结果被现实教育了。第一SD-VAE的参数量不小微调它需要大量显存第二医学图像的数据量即使加上未标注的相对于它预训练的数据规模也是九牛一毛微调极易灾难性遗忘——学医学特征的同时把预训练学到的自然图像通用视觉能力忘掉第三我们做半监督的核心逻辑就是希望有一个稳定、共享的特征底座如果编码器在每个epoch都在变分布对齐就会变成在移动靶上射箭训练极不稳定。所以“冻结”不是偷懒而是刻意为之。它让整个方案变成了“一个固定的特征空间一个可学习的分割头两个可学习的对齐模块”复杂度大幅下降训练稳定性显著提升。3.3 双分布对齐的数学表达对齐是一个技术活我试过很多度量方式KL散度、MMD、对抗判别器、中心矩对齐。最终选用的组合是无标注特征向有标注特征做“条件分布对齐”两者共同向先验分布做“边缘分布对齐”。具体地定义标注数据经过分割头编码器提炼的特征为Zs无标注数据对应的特征为Zu先验分布Q采用标准高斯N(0, I)。第一项分布对齐损失是中心矩差异Central Moment Discrepancy, CMDL_cmd || E[Zs] - E[Zu] ||^2 || Cov(Zs) - Cov(Zu) ||^2 Σ_{k3}^{5} || M_k(Zs) - M_k(Zu) ||^2这里的 M_k 表示第k阶中心矩。为什么不用KL因为KL要求分布有显式密度形式而Zs和Zu都是神经网络的输出密度未知。MMD可以但核宽度的选择很敏感要花不少时间调。CMD的好处是直接用矩来刻画分布不需要核函数训练稳定而且计算开销低。第二项是先验对齐损失直接约束Zs和Zu的边缘分布都向N(0, I)靠拢。这里需要一个密度估计我参照变分自编码器的做法假设Zs和Zu都是高斯分布用KL散度L_prior KL(q(Zs) || N(0, I)) KL(q(Zu) || N(0, I))其中q(Z)用编码器的输出均值和对数方差参数化。最终的无标注损失L_unlabeled L_cmd α * L_prior而有标注数据的完整损失L_sup Dice_loss(seg_pred, mask) β * L_prior总损失L_total L_sup γ * L_unlabeled其中α、β、γ是权重系数。我自己常用的配置是α0.5β0.1γ1.0。3.4 训练流程的时间表有两阶段和单阶段两种方案我最终使用的是单阶段端到端训练但对无标注损失做了“延迟热身”处理前500步迭代只训练有标注分支让分割头先学会基本的轮廓定位能力然后再开启L_unlabeled。这个设计原因很简单一开始分割头还没学到基本的形状概念Zs分布是乱七八糟的如果此时就用CMD去对齐Zu相当于让一个还没站稳的模型去当导师会把未标注数据的特征也带偏。等分割头学习到基本的模式后再开启对齐就像学生已经有了一定基础这时候再来做师徒互鉴才是有效的。训练过程中我全程冻结SD-VAE编码器的参数只更新分割头和对齐模块。整个模型在单张48G显卡上batch size为8有标注4、无标注4训练约2万步收敛耗时约10小时。4. 实验设计4.1 数据集与评价工具我用了两个公开数据集做验证。第一个是腹部多器官CT分割数据集我用其中肝脏、脾脏、肾脏这三类做多类分割评测共40例CT切片后约6000张。第二个是心脏MRI分割数据集包含左心房、左心室、心肌三个结构共30例。我按7:1:2拆分成训练、验证、测试集并且在训练集中只取20%的切片保留标注其余80%作为无标注数据参与训练。评价指标用Dice系数和HD95豪斯多夫距离95分位数。Dice衡量区域重叠HD95衡量边界误差——医学影像里边界误差特别重要因为临床测量往往依赖轮廓精确度。4.2 对照方案我做了四组对照全监督基线100%标注数据相同分割头不使用SD-VAE和分布对齐。半监督伪标签20%标注80%无标注使用confidence-based伪标签。半监督一致性正则20%标注80%无标注在SD-VAE特征空间做随机掩码扰动类似MAE思路的一致性约束。半监督本方案20%标注80%无标注冻结SD-VAE双分布对齐。这里要说明一下为了让对比公平除了全监督基线外其他半监督方案的分割头都相同也都接在SD-VAE冻结编码器上只改动半监督约束方式确保差异来自方法本身。4.3 核心结果与解读下面这组结果是在腹部多器官CT上以肝脏分割为目标的实测数据方案Dice(%)HD95(mm)全监督100%标注88.46.8半监督伪标签20%标注84.211.3半监督一致性正则20%标注85.19.6半监督本方案20%标注87.97.2可以看到本方案用20%的标注量就追近了全监督的结果Dice差距仅0.5个点HD95差距0.4mm。伪标签方案和一致性正则方案虽然也有一定效果但HD95明显变差——说明它们在边缘细节上犯错更多这在医学分割里是更要命的。在心脏MRI数据集上趋势一致但差距稍大本方案与全监督的Dice差距约1.2个点。我觉得原因可能是心脏MRI的边界更模糊而且类间对比度低SD-VAE的特征表达对这种低纹理差异的任务加成有限。4.4 消融实验为了验证每个组件的必要性我做了完整的消融配置Dice(%)完整方案87.9去掉SD-VAE改用ResNet50预训练编码器84.6保留SD-VAE去掉L_cmd只做L_prior85.3保留SD-VAE去掉L_prior只做L_cmd86.1冻结SD-VAE 双对齐但SD-VAE改用非动冻结可微调85.7这组消融说明三件事。第一SD-VAE的编码器确实比ResNet强在线性探测式的分割任务上它的潜在特征更紧凑区分度更高。第二L_cmd和L_prior都有用但L_cmd的作用更大毕竟核心的信息传递靠的是有标注→无标注的特征分布拉近。第三微调SD-VAE反而掉点佐证了“冻结”的必要性。5. 实操过程中踩过的坑5.1 冻结点选错了特征语义紊乱我第一次做的时候用的是一个较大版本的VAE解码结构特征图分辨率只有H/16通道数也更高。训练一轮非常慢而且分割头的感受野也要做相应调整效果并不好。后来检查才发现Stable Diffusion VAE的描述并不完全统一——不同配置下压缩比和通道数会有差异。正确做法是确认你的编码器输出的空间分辨率是不是H/8如果是从解码器的中间层硬抠特征分辨率是H/4但语义混杂会导致分割头学到一堆无效纹理。我的参照做法是取编码器下采样到1/8的那一层输出通道数为8C不同版本稍有差异然后接一个1x1卷积调整通道数到分割头期望值。不要直接从2x bottleneck 输出往后接因为bottleneck层的语义是高度压缩的空间细节丢太多。5.2 CMD计算时梯度爆炸L_cmd的计算涉及矩估计尤其三阶和五阶中心矩当特征数值范围比较大的时候矩值会迅速膨胀loss直接变成NaN。我一开始用四阶矩训练到2000步就崩了。解决办法有两个一是对特征做L2归一化后再算矩二是给矩项加温度参数比如每个中心矩除以其最大特征值λ_max的k次方做标准化。我把三阶和五阶矩的系数设置为0.01和0.001训练曲线就稳定多了。实操心得如果训练到一半发现L_cmd踩不下去先检查特征是不是整体漂移了。在CMD之前加一个LayerNorm作为特征变换层可以很大程度上缓解这个问题。5.3 无标注损失在震荡需要学习率去耦早期我把所有可学习参数都用一个学习率1e-4。结果发现L_cmd会剧烈震荡导致分割头的分割损失也在波动。原因是分割头和分布对齐模块的收敛速度不一致——分割头更新频率和对齐模块同步更新会互相干扰。我把训练参数拆成两组分割头用1e-4分布对齐模块用1e-3并且给对齐模块加了梯度截断max_norm1.0。这样做的直观结果是对齐模块更快地建立分布之间的桥梁而分割头则慢慢地在已被对齐的特征空间里画边界。还有一个容易踩的坑无标注损失的训练权重γ不宜过高。我试过γ5结果导致特征分布过度被拉向先验分割头在标注数据上的保真度下降Dice直接掉了3个点。γ1.0是稳定且收益最大的甜点位。5.4 显存瓶颈与小batch的替代方案如果只有24G显存batch size 84标注4无标注可能放不上。我试过batch size 633结果CMD在计算矩时方差波动太大训练不稳定。替代方案是梯度累积。把batch拆成两个子batch有标注2、无标注2每个子batch前向计算后梯度累积等效于原batch size。注意CMD计算时矩估计要基于累积特征的统计量而不是单子batch——这需要在前向时保存特征跨子batch聚合后统一计算损失。5.5 病态窗宽导致SD-VAE输出幻觉特征有一次我偷懒用单位映射直接做归一化结果发现分割Dice严重下降。检查后发现CT图像的HU值范围非常大-1024到3071直接喂给SD-VAE编码器会把极端的数值解释成某些自然图像才有的纹理模式产生“幻觉特征”——比如在空气区域出现类似云朵的纹理响应。这是生成式模型的特性它对输入分布的偏态非常敏感。解决方法是严格做窗宽窗位处理并且对每个体数据估算动态范围再做标准化。6. 场景扩展与效率优化6.1 这个方案还可以用到哪里这套“冻结生成式大模型编码器双分布对齐”的思路并不局限于CT或MRI分割。我可以预想到的几个扩展方向病理图像分割病理全切片WSI同样存在标注稀缺的问题而且它的纹理复杂度和自然图像更相近SD-VAE的特征可能更有价值。跨模态医学图像分割如果无标注数据来自不同设备或不同序列比如T1和T2 MRI分布对齐刚好可以作为模态间特征桥接的工具。视频医学图像如超声视频利用时间维度的多帧未标注数据做分布对齐可能能进一步提升分割的时间一致性。缺陷检测之外的图像分割比如遥感图像中的建筑物分割、工业质检中的缺陷区域分割只要是“标注贵、未标注多”的场景这套方案都有迁移的潜力。6.2 推理速度和部署考量谈到实际落地很多人第一反应是模型太大、跑不动。其实不会。R50之类的主干加SD编码器大材小用了吗在推理阶段SD-VAE编码器是冻结的且和分割头是串行结构理论延迟取决于两者之和。我实测在单张24G显卡上256x256的单张切片预处理编码分割头推理大约耗时45ms约22FPS这远超过临床交互工具的实时性需求。如果还想压有两个办法一是把SD-VAE编码器蒸馏成一个轻量编码器比如用类似蒸馏学习的方式训练一个小网络去逼近SD-VAE的输出耗时可以从45ms降到18ms二是在预处理阶段对三维数据做分组推理减少不必要的重复编码。6.3 从Dice到临床可用的最后一步Dice是研究指标但临床上更关心的是分割结果能不能直接用于手术导航或者体积测量。我遇到过一种情况Dice挺好但边界在某些区域比如肝脏和胃的交界出现了不光滑的毛刺。这是因为Dice不惩罚边界锯齿。解决办法是在训练中加入边界损失boundary loss或者后处理用条件随机场CRF做平滑。我倾向在损失里加一个辅助的边界dice或距离损失而不是依赖CRF因为CRF的推理过程在3D体数据上相当慢而且容易把细小的病变区域一并抹掉。另外输出结果的连通域过滤非常重要。医学图像分割里常常出现小型的孤立假阳区域如果做手术规划这些假阳区域会误导医生。我的做法是保留最大连通域以体积阈值为准这能让HD95显著下降。7. 一些个人化的总结和经验项目做到最后我最大的体会其实是“不要逆着模型的脾气来”。SD-VAE的潜在空间是生成式模型留下的资产它虽然语义强大但有自己的统计偏置。前面我说了那么多预处理、对齐、归一化的细节本质上都是在小心翼翼地顺着这个偏置去用它的长处。另一个很深的体会是在半监督这个领域真正拉开差距的往往不是复杂的伪标签策略或者花哨的一致性正则而是一个稳定的特征底座加上一个好的分布约束。“以空间换语义以分布换标注”这句话用在这个方案上再合适不过。如果要给后来者一些建议的话我会列这么几点不要一上来就铺大网络先把SD-VAE编码器冻结住用一个小型分割头验证特征的线性可分性。分布对齐的优先级是先做标注/无标注对齐再做先验对齐。前者的收益远大于后者别把权重配反。训练初期一定要做延迟热身否则对齐模块会破坏分割头的早期学习。每次改分辨率或归一化方式后都要重新确认SD-VAE的输出分布形态生成式模型对输入范围极其敏感。别迷信更多标注量。我实验里把标注量从20%提到50%Dice只涨了1.1个点成本却翻了几倍——半监督的核心平衡点是“以合理标注换来临床够用的精度”而不是无限逼近全监督。最后分享一个小技巧如果你手头的未标注数据远多于标注数据可以尝试给无标注分支换一个更大的batch而不是维持和有标注相同的batch——因为分布对齐的矩估计需要足够多样本来降低方差。这也是我在一次偶然的实验里发现的把无标注batch从4加到16后整个训练曲线都平稳了不少。这个方案后续我打算继续往两个方向推一是把分割头的输出也纳入分布对齐的范围做更细粒度的类别条件对齐二是尝试不冻结SD-VAE的浅层编码器只冻结深层让浅层能微调以适应医学图像的纹理分布。这些还在实验里后面有结果了再更新。
