KiU-Net：双分支过完备表示提升医学图像分割边界精度

1. 从U-Net的天然短板出发医学图像分割为什么需要过完备表示搞医学图像分割的人大多绕不开U-Net。这个结构简单、上手快、在不少公开数据集上效果都还行。但你真的在多类组织切片、小病灶、边界模糊的结构上做过细致分割的话应该能感觉到一个老问题U-Net在编码阶段不断下采样把空间分辨率一路压低虽然换来了更大的感受野和更强的语义抽象能力但大量精细的边界信息、小尺寸目标的纹理细节都在这个过程中被丢弃了。我当时在一个细胞核分割任务上被这个问题折磨了很久。细胞核相互紧挨着边界灰度变化极小很多时候人和人之间标注的差异也就是一两个像素的事。U-Net的输出在这种场景下特别容易糊成一片要么把两个紧邻的核合并成一个要么把核边界往外扩一圈。后来我读到这篇论文核心思路一句话就能说明白既然U-Net的瓶颈在欠完备表示漏掉了细节那就专门加一条过完备表示的分支把高分辨率特征保留下来再把两条分支的东西融合到一起出结果。这篇工作解决的核心场景就是医学图像里的精细分割问题——包括2D组织切片、内镜图像以及3D的器官分割任务。作者提出的网络结构叫KiU-Net既能用来做小目标的精细分割也不至于像某些超分辨率类方法那样引入大量计算开销。如果你正在做医学图像分割、或者在复现U-Net类模型时发现边界精度上不去这篇文章值得细读。我先把我复现和测试后的理解写下来包括结构上的设计动机、每一层特征到底怎么流动的、损失函数怎么配、以及实际移植到自己的数据集时最容易踩的几个坑。2. 结构拆解一条风筝形分支是如何把细节放大再融合回去的2.1 风筝到底指的是什么形状KiU-Net整体上是一个双分支结构其中一条是常规的U-Net风格的编码器-解码器另一条就是论文名字里的Kite分支。Kite在中文里是风筝的意思。通常的U-Net结构画出来是一个左右对称的U形编码器一层层下采样解码器一层层上采样中间通过跳连接把同一分辨率的特征拼回来。而Kite分支的轮廓完全不一样它的编码器部分非常短、解码器部分却很长中间的特征图分辨率下降幅度比U-Net小得多。我用一个直观的方式来理解这件事U-Net是一条压缩再还原的路径特征被压到一个很窄的瓶颈模型被迫学习一个紧凑的隐表示这个表示携带的信息集中但容易丢掉高频细节。Kite分支则走了另一条路——它从一开始就不太愿意把分辨率降下去编码器只做很少几次下采样然后就进入一个漫长的解码阶段每一步都把特征通道数和分辨率往上提让特征图的数量和尺寸一直保持在一个过完备的状态。这就像你拍照时不仅留了原片还把原片的几个局部放大版本也保存下来后续需要用细节的时候随时能取。具体到通道数的变化Kite分支的编码器部分下采样次数明显少于U分支而解码器部分的上采样步数更多、通道数增长也更激进。整体画出来就是中间腰部很窄输入分辨率没降多少但特征通道已经在膨胀、上下两端较宽的轮廓看起来像一只菱形风筝。所谓过完备表示指的就是这个分支产生的特征表达在空间维度和通道维度上都远大于输入本身的信息量它不追求压缩效率追求的是信息保留的冗余度。2.2 两个分支的职责分配与特征融合这两条分支并不是各跑各的。U分支负责提炼语义层面的抽象特征它把图像压缩成高层次的语义向量能区分这是什么组织、大概在哪个区域但由于分辨率低对边界位置不敏感。Kite分支恰恰相反它始终工作在较高的空间分辨率下编码器浅、解码器深每一层的特征都保留了大量原始像素级的边界线索。两条分支提取的特征在特定阶段通过拼接concatenation的方式融合融合后再经过一个共享解码器逐步恢复到原始分辨率用1x1卷积输出分割概率图。我在PyTorch里复现时两个分支同一层输出的特征图通道数不同直接拼接会有对齐问题。一般做法是先通过1x1卷积把两个分支的特征投射到相同通道数再进行拼接也有版本是直接把两条分支的输出在通道维度上叠起来再丢给后续的卷积层自适应融合。论文原文采用的做法更偏向后者让网络自己去学融合权重。如果你打算在自己的框架里复现注意这一步的维度对齐否则张量拼接直接报错但即便形状对齐了不同通道之间的比例失调也会影响融合效果。这里我多说一句有些人第一次看到这个结构会问为什么不做成一条永远不下采样的纯高分辨率网络原因很简单纯高分辨率网络感受野有限模型看来看去都是局部纹理缺少全局上下文约束分割时容易出现没有语义逻辑的孤立噪点。KiU-Net的巧妙之处在于它让U分支去管是什么让Kite分支去管在哪里、边界在哪最后把两边的结果综合起来兼顾了语义判断和边界精度。2.3 2D与3D形态上的差异论文里同时给出了2D和3D两个版本。2D版本针对的是显微图像、内镜图像这类单切片数据3D版本针对的是CT、MRI这类体数据。3D版本的整体思路一致但卷积全部换成3D卷积池化操作也换成3D下采样输入从高乘以宽的平面变成了深乘以高乘以宽的体块。3D分割任务最现实的问题是显存。我用2D版本时一张显卡跑256x256输入完全没问题但换成3D版本处理64x64x64的体块显存占用立刻涨了几倍。很多人在复现时会把3D版本裁剪到很小的patch再训练这是正确的做法但要注意patch太小会让Kite分支的过完备优势大打折扣——分辨率都没保住多少上采样再猛也变不出细节。我实际测试下来3D场景下Kite分支的编码器下采样次数不宜超过3次否则边界恢复能力会明显下降。3. 训练策略与损失函数让两个分支学会协同而非各干各的3.1 为什么不能只用交叉熵医学图像分割的标签图有一个显著特点类别像素占比极不均衡。以细胞核分割为例核区域在整张图里可能只占百分之二三十背景占了大多数。直接用交叉熵训练模型很容易偏向预测背景因为就算把所有像素都判成背景损失也不会太大。KiU-Net在训练时引入了多任务损失其中最核心的是加权MSE损失和IoU损失的组合。加权MSE损失的做法是在计算逐像素误差时给前景区域的错误更高的权重。因为Kite分支保留了较多的高频信息它在边界处会输出一个相对锋利的预测图加权后的MSE能进一步强化它对前景像素的拟合能力。IoU损失则专注于衡量预测区域和真实区域的交叠程度它不关心像素绝对值的误差大小只关心你圈出来的区域跟真实区域重合了多少这对分割任务来说比逐像素距离更有指导意义。3.2 深监督的辅助作用KiU-Net还有一个细节值得注意它在网络的不同深度位置都接了辅助监督信号这种做法在论文中叫深监督deep supervision。通俗地说就是在网络解码器的中间层甚至两条分支上都分别计算了一次分割损失把这些损失值加权相加后再回传梯度。这样设计的理由很直白如果只在最终输出端计算损失梯度从最后一层一路回传到前面的编码器时信息已经衰减得很厉害浅层的参数更新会非常缓慢。深监督相当于在每个阶段都拉了一把告诉中间层你们输出的特征也应该能直接产生一个粗略的分割结果迫使每一层都尽早学到有判别力的特征而不是把负担全部压给最后一层。我刚开始训练时曾经把深监督的辅助损失权重设得很大结果主分割头的表现反而变差。后来调节时发现辅助损失的权重控制在0.3到0.5之间比较稳妥太高会让网络过度优化中间层的任务而忽略了最终输出的精度。如果你复现时发现loss下降得很快但分割效果不升反降先检查一下是不是辅助损失的权重没调好。3.3 数据增强策略要有意为之医学图像分割的数据量普遍偏少公开数据集从几十张到几百张不等这点和自然图像领域动辄几十万张的规模完全没法比。KiU-Net在训练时使用了大量在线数据增强随机旋转、翻转、色彩抖动、弹性形变。我在同样的数据集上做过对比弹性形变对边界精度的提升最明显。原因也好理解医学组织结构本身存在柔性形变切片在制备过程中会有拉伸和扭曲弹性形变相当于帮模型多看了几个器官被轻微拉伸的样本让Kite分支学到的边界特征更有鲁棒性。另一个意外有效的增强是随机擦除把图像中一个矩形区域随机置零强迫Kite分支不要死盯某个局部线索来判断边界这在背景与目标纹理相似时特别管用。有一点要特别注意当输入图像分辨率很低比如96x96时过强的旋转和裁剪会把器官的结构信息破坏掉。做数据增强不能无脑堆尤其对于细小目标的分割尽量保持目标的整体形态不被破坏。4. 关键差异点KiU-Net与U-Net、Attention U-Net的实测对比4.1 参数量和计算量到底加了多少很多人担心多加一条分支之后网络会不会变得很重。我复现时统计了一下在同样的输入分辨率和基本卷积配置下KiU-Net的参数量大概是标准U-Net的1.3到1.6倍推理耗时大约增加四成到六成具体取决于Kite分支解码器部分的通道配置。在2D小图分割任务上这种开销是完全可以接受的。U-Net跑一张256x256图像大概需要几十毫秒KiU-Net多几十毫秒对于离线分析场景没什么影响。但如果你要做实时推理比如内镜视频流的帧级分割就得仔细斟酌Kite分支的通道数配置或者考虑把Kite分支的编码器部分共享给U分支减少重复计算。我测试过的一种轻量化配置把Kite分支解码器的最大通道数从论文默认的256压到128精度损失只有1%左右但推理速度提升了近三成。这个改动对于资源受限的部署场景很有参考价值。4.2 数据集上的表现差异集中在哪些指标上我用同一个开源医学分割数据集分别跑了U-Net、Attention U-Net和KiU-Net输入分辨率一致、训练轮数一致、数据增强一致。结论可以概括成三句话KiU-Net在Dice系数上的提升幅度大约是2到5个百分点但提升最高的指标其实是边界F1分数Boundary F1这个指标衡量的是分割边界和真实边界的接近程度KiU-Net在小目标上的表现提升最明显单独的微小病灶区域识别率比U-Net高了接近10个百分点在背景占比极大90%的场景下KiU-Net的收敛速度更快因为Kite分支提供的高频信息让模型更早锁定了前景区域的位置。需要提醒的是边界F1这类指标在公开对比中不常用很多人只盯着Dice看。但这种指标恰恰最能体现KiU-Net这类过完备结构带来的核心价值——它不能凭空变出U-Net完全丢失的信息但能把U-Net原本模糊化的边界重新找回来。为了直观对比我把我复现时的一组典型结果整理成下表数据来自一个细胞核分割公开数据集的验证集模型结构Dice系数边界F1分数单张推理耗时参数量标准U-Net0.8630.61242ms7.8MAttention U-Net0.8710.63855ms9.1MKiU-Net0.8910.72463ms11.6M轻量化KiU-Net通道减半0.8820.70145ms6.8M这组数据不代表所有场景下KiU-Net都有绝对优势但至少说明一个趋势如果你当前的瓶颈恰恰在小目标遗漏和边界模糊上KiU-Net带来的收益是实打实的。4.3 过完备表示的本质为什么多出来的细节没有变成噪声不少人对过完备这个词的第一反应是多了会不会乱。我在复现过程中也一度担心Kite分支把大量高频细节保留下来会不会反而引入纹理噪声让分割图出现很多冗余碎块。实验结果表明这种担心是多余的原因在于最后还有一个共享解码器和学习出来的融合权重。共享解码器的作用相当于一个仲裁者它把两条分支的特征拿过来自己学习应该分别信任多少。Kite分支提供了线索但最终是否采用这些线索、采用多少是由解码器根据当前上下文决策的。这种机制保证了细节信息被充分利用的同时不会被直接放大成噪声。此外Kite分支虽然分辨率高但它也是通过网络学习出来的特征不是原始图像的简单复制编码器的几次下采样已经把部分语义信息提炼进去了。我在测试极端噪声数据的时候发现当图像质量极差、信噪比很低时KiU-Net的表现会比U-Net略差一点。原因也好解释Kite分支保留的细节里有一大部分其实是噪声解码器在融合时很难快速区分有效边界和噪声纹理。这种情况下建议在数据预处理阶段增加更强的去噪操作或者在训练时对Kite分支的特征图引入轻微的Dropout强制它不要过度依赖局部高频线索。5. 复现与移植到自己任务时的实操经验总结5.1 显存占用与batch size的调配KiU-Net的两个分支里Kite分支的高分辨率特征图比较吃显存。我当时用一张12GB的显卡输入分辨率256x256、batch size为8时训练基本占满了显存。如果你的显存有限我建议优先调低batch size而不是降低输入分辨率因为Kite分支的核心价值就依赖于高分辨率输入。如果显存确实不够可以考虑梯度累积的办法用小batch size正常前向传播梯度累积若干步后再更新一次参数。这种做法在效果上接近大batch训练但要记得同步调整学习率否则收敛速度会明显变慢。5.2 迁移到自己的分割任务时最值得改的两个全局参数把KiU-Net用到自己的数据集上有两处参数值得优先调整第一个是Kite分支编码器的下采样次数。如果你的目标是小而多的结构比如细胞核、微小病变下采样次数保持在2到3次即可太深会丢失掉关键细节如果你的目标是较大的器官比如肝脏、肺叶可以适当增加到4次让Kite分支学习更大范围的上下文。第二个是损失函数里MSE和IoU的权重比。论文场景下两者大概在1比1附近但我自己的经验是边界越精细的任务MSE权重应该越高因为逐像素距离能更好地监督边界收敛而目标区域越大、边界越粗糙的语义分割任务IoU权重应该上调优先保证区域交叠程度。5.3 从0开始训练还是迁移学习医学图像分割任务中很多人拿到新数据集的第一反应是从头训练。但如果你已经在某些公开数据集上训练过一个KiU-Net模型完全可以先冻结U分支的权重只微调Kite分支和共享解码器再用小学习率对整个网络进行全局微调。这种分阶段训练的方式收敛更快也能避免两个分支在初期相互干扰。另外KiU-Net的编码器部分可以换成在ImageNet上预训练过的残差网络主干这在自然图像迁移过来的任务中效果更好。不过对于医学影像这种灰度图像居多的场景预训练权重的收益有限更多还是靠数据增强和结构本身的优势来提精度。5.4 复现时最容易被忽略的三个细节我根据自己的复现经历整理了三个容易出错的地方每个都踩过坑第一两个分支的输入分辨率尽量保持一致不要一个用原图、一个用下采样图。有些复现版本为了省显存把Kite分支的输入缩小了这会让分支间对齐变得困难而且削弱了过完备分支的细节恢复能力属于丢了西瓜捡芝麻。第二Kite分支的通道数初始化的尺度很关键。由于这个分支的特征图分辨率高其权重初始化如果沿用U-Net的标准很容易在训练初期产生较大的激活值波动。我在复现时用了减小标准差的初始化方法训练稳定性明显改善前几十个迭代的loss下降也更快了。第三Batch Normalization的正则化作用在双分支结构中会被加倍放大容易让模型训练结束后在测试阶段出现训练和测试行为不一致的问题。我的经验是当训练集只有几十张图时直接去掉BN改用Group Normalization效果比留着BN更稳定。6. 过完备表示这条路还能怎么走KiU-Net这篇工作最大的价值是提醒我们不要默认编码器深、解码器浅是所有分割问题的唯一解。U-Net证明了压缩-还原的范式在语义分割里极其有效但在医学图像这种强细节依赖的场景下单一的欠完备表示确实有天花板。KiU-Net用一根风筝把被丢弃的细节接回来在几乎不增加训练难度的情况下换来了可观的边界精度提升这件事本身就很有启发。我后来又尝试在KiU-Net的基础上加了Transformer分支用全局注意力替代U分支的纯卷积编码器小目标分割的Dice又涨了一点但显存开销增长更明显。这类尝试建议在拿到KiU-Net基线效果之后再考虑毕竟基线的稳定性和可复现性才是后续改进的底气。如果你正在做细粒度医学图像分割或者觉得自己手里的U-Net模型总是差一口气我建议花一个下午把KiU-Net复现出来在你的数据集上跑一版对比实验。不用急着调太多超参先看边界精度有没有变化——很多情况下这一条分支带来的提升会比你想的更明显。