通过可学习的多尺度概率混合提升医学图像分割精度,无需额外计算开销。
LoMix: Learnable Weighted Multi-Scale Logits Mixing for Medical Image Segmentation
- 设计可微分模块LoMix,自动融合多尺度预测结果
- 在8器官数据集上比单输出监督提升4.2%的DICE分数
- 特别适合标注数据少的场景,数据效率高且推理无负担
U型网络在多个空间尺度输出概率图,分别捕获粗粒度上下文与细粒度细节。但训练时仍孤立处理这些输出——要么仅监督最高分辨率输出,要么对所有尺度施加相同权重的深层监督,未探索多尺度组合。为此,我们提出LoMix(概率混合),一种受神经架构搜索启发的可微分即插即用模块,生成新的多尺度混合输出,并学习每种混合如何指导训练过程。具体而言,LoMix使用四种轻量级融合算子(加法、乘法、拼接、注意力加权融合)生成丰富合成变异图,每个原始或变异图分配一个通过软正切函数优化的损失权重,实现网络参数与权重的联合优化,模拟一步架构搜索,自动发现最优尺度、组合与算子。将LoMix嵌入近期U型架构(如基于PVT-V2-B2主干和EMCAD解码器)在Synapse 8器官数据集上,相比单输出监督提升4.2% DICE,相比深层监督提升2.2%,相比等权重加法融合提升1.5%,且推理零开销。当训练数据稀缺(如仅1或2个标注扫描)时,优势扩大至+9.23%,凸显其数据高效性。在四个基准和多种U型网络上,LoMix相比单输出监督最多提升13.5% DICE,验证了其通用性、数据高效性、完全可解释性及推理无开销特性。代码已开源。
原文摘要 · Abstract (English)
U-shaped networks output logits at multiple spatial scales, each capturing a different blend of coarse context and fine detail. Yet, training still treats these logits in isolation - either supervising only the final, highest-resolution logits or applying deep supervision with identical loss weights at every scale - without exploring mixed-scale combinations. Consequently, the decoder output misses the complementary cues that arise only when coarse and fine predictions are fused. To address this issue, we introduce LoMix (Logits Mixing), a NAS-inspired, differentiable plug-and-play module that generates new mixed-scale outputs and learns how exactly each of them should guide the training process. More precisely, LoMix mixes the multi-scale decoder logits with four lightweight fusion operators: addition, multiplication, concatenation, and attention-based weighted fusion, yielding a rich set of synthetic mutant maps. Every original or mutant map is given a softplus loss weight that is co-optimized with network parameters, mimicking a one-step architecture search that automatically discovers the most useful scales, mixtures, and operators. Plugging LoMix into recent U-shaped architectures (i.e., PVT-V2-B2 backbone with EMCAD decoder) on Synapse 8-organ dataset improves DICE by +4.2% over single-output supervision, +2.2% over deep supervision, and +1.5% over equally weighted additive fusion, all with zero inference overhead. When training data are scarce (e.g., one or two labeled scans), the advantage grows to +9.23%, underscoring LoMix's data efficiency. Across four benchmarks and diverse U-shaped networks, LoMiX improves DICE by up to +13.5% over single-output supervision, confirming that learnable weighted mixed-scale fusion generalizes broadly while remaining data efficient, fully interpretable, and overhead-free at inference. Our code is available at https://github.com/SLDGroup/LoMix.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。