LoMix: Learnable Weighted Multi-Scale Logits Mixing for Medical Image Segmentation

TL;DR

LoMix addresses the underutilization of multi‑scale logits in U‑shaped medical segmentation networks by introducing a differentiable, NAS‑inspired framework that learns how to mix coarse and fine predictions. The Combinatorial Mutation Module (CMM) creates a rich set of mutant logits through four fusion operators across all decoder subsets, while NAS‑inspired softplus weights $w= ext{softplus}( heta)$ learn how strongly each logit should guide training, yielding an end‑to‑end trainable loss ensemble with zero inference overhead. Empirically, LoMix achieves state‑of‑the‑art or competitive results across Synapse, ACDC, BUSI, and ISIC datasets, with particularly large gains under data scarcity (up to +9.23% DICE with 5% labeled data) and across diverse backbones. The method provides a transparent, data‑efficient, plug‑and‑play improvement that enhances boundary capture and context integration without modifying inference pipelines, and it lays a foundation for extending mixed‑scale supervision to 3D and other dense prediction tasks.

Abstract

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.

Figures (13)

• Figure 1: The proposed LoMix supervision strategy during training. (a-b) An input image is processed by a U–shaped network, producing original multi‑scale logits $P_{orig}$ from different stage of decoders. (c) The Combinatorial Mutation Module synthesizes additional predictions by applying four fusion operators to every non‑trivial original prediction subset: Addition, Concatenation, Multiplication, and a learnable Attention-Weighted Fusion (awf). (d) NAS‑inspired weight‑learning produces differentiable softplus‑transformed weights $w$ for every original logit in $P_{orig}$ or mutated logit in $P_{mut}$, optimized jointly with network parameters via back‑propagation. (e) All original and mutated logits are supervised by a loss objective weighted by $\{w\}$.
• Figure 2: Limited data evaluation on Synapse 8‑organ segmentation with PVT‑EMCAD‑B2. LoMix uses NAS-inspired softplus weighting. $\uparrow$ denotes higher is better. Results are averaged over five runs. Detailed results are provided in Table \ref{['tab:synapse_limited_data_ablation']}.
• Figure 3: Comparison of different fusion operation combinations using NAS-inspired Softplus weights and PVT-EMCAD-B2 model for Synapse 8-organ segmentation. $\uparrow$ indicates higher is better.
• Figure 4: Effect of NAS‑inspired Softplus weight learning on Synapse 8‑organ segmentation with PVT‑EMCAD‑B2. Fixed = equal loss weights, Learned = NAS‑inspired softplus weights.
• Figure 5: Comparison of different supervision schemes on Synapse 8-organ segmentation across five backbones. LoMix uses NAS-inspired softplus weighting. Sup.: Supervision.