Skip to content

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

Conference: NeurIPS 2025 arXiv: 2510.22995 Code: https://github.com/SLDGroup/LoMix Area: Medical Image Segmentation Keywords: Multi-scale fusion, Logits mixing, Deep supervision, NAS, U-Net

TL;DR

LoMix introduces a Combinatorial Mutation Module (CMM) that generates "mutant" logits from multi-scale outputs via four fusion operators (addition / multiplication / concatenation / attention-weighted fusion) across all subset combinations, paired with NAS-style Softplus learnable weights for automatic contribution balancing. On Synapse 8-organ segmentation, Dice improves from 80.9% to 85.1% (+4.2%), and by +9.23% under 5% training data.

Background & Motivation

Background: U-shaped segmentation networks (UNet, TransUNet, etc.) produce multiple logits outputs at different scales across decoder layers. The standard practice is to use only the last layer (LL) output or apply uniform-weight deep supervision (DS).

Limitations of Prior Work: (a) Single-output strategies discard complementary information from intermediate layers; (b) uniform deep supervision weights are suboptimal across anatomical structures—small and large organs require information at different scales; (c) manual weight tuning is time-consuming and non-transferable.

Key Challenge: Multi-scale logits carry complementary information (coarse scales provide large-region context; fine scales capture precise boundaries), yet no effective mixing mechanism exists in prior work.

Goal: Automatically learn the optimal fusion strategy and weights for multi-scale logits.

Key Insight: Drawing on differentiable NAS search—parameterizing all possible fusion strategies and weights into a differentiable search space for end-to-end learning.

Core Idea: 4 fusion operators × all scale subset combinations = a "mutant" logits pool → Softplus learnable weights for end-to-end optimal mixing.

Method

Overall Architecture

A U-shaped network with \(L\) decoder layers produces logits \(P_i\) at each layer → CMM: generates "mutant" logits over all \(2^L - 1 - L\) multi-scale subsets × 4 fusion operators → each logit is assigned a Softplus weight \(w_u = \ln(1 + e^{\alpha_u})\) → weighted total loss \(\mathcal{L}_{total} = \sum w_i \mathcal{L}_i + \sum w_S^{(op)} \mathcal{L}_S^{(op)}\)

Key Designs

  1. Combinatorial Mutation Module (CMM):

    • Function: Generates a rich set of "mutant" logits from the original \(L\) scale logits.
    • Mechanism: Four fusion operators—addition \(\sum P_i\), multiplication \(\prod P_i\), concatenation followed by \(1\times1\) convolution, and attention-weighted fusion (AWF, softmax-normalized per-pixel). Applied to every non-empty multi-scale subset (\(2^L - 1 - L\) subsets) × 4 operators, yielding \(L + 4(2^L - 1 - L)\) logits in total.
    • Design Motivation: Different fusion operators capture different types of complementary information—addition performs averaging/smoothing, multiplication highlights intersections/high-confidence regions, and attention enables adaptive weighting.

  2. NAS-Style Softplus Learnable Weights:

    • Function: Automatically learns optimal weights for each logit (original and mutant).
    • Mechanism: \(w_u = \text{softplus}(\alpha_u) = \ln(1 + e^{\alpha_u})\), where \(\alpha_u\) is a learnable parameter. Softplus guarantees non-negativity and differentiability.
    • Design Motivation: More flexible than fixed weights and end-to-end optimizable; more stable than softmax normalization (allows certain logit weights to approach zero).

  3. Computational Feasibility for \(L \leq 5\):

    • Function: Prevents combinatorial explosion.
    • Mechanism: At \(L=5\), the total number of logits is \(5 + 4 \times 26 = 109\)—fully tractable in practice.
    • Design Motivation: Medical segmentation networks typically have 3–5 decoder layers, keeping the combination count manageable.

Loss & Training

  • Per-logit loss: \(\mathcal{L} = \beta \mathcal{L}_{CE} + \gamma \mathcal{L}_{DICE}\)
  • Total loss is the weighted sum over all logits.
  • Weights and network parameters are optimized jointly.

Key Experimental Results

Main Results

Dataset Backbone LoMix Dice Baseline (LL) Gain
Synapse 8-organ PVT-EMCAD-B2 85.1% 80.9% +4.2%
ACDC Cardiac PVT-EMCAD-B2 92.51% — —
Synapse (5% data) Various backbones +9.23% — Strong low-data regime

Ablation Study

Configuration Synapse Dice
Addition only 84.1%
Multiplication only 83.8%
All 4 operators 85.07%
Softplus vs. fixed weights +0.21–0.7%
Cross-backbone (UNet~PVT) Consistent +3.88–11.88%

Key Findings

  • The multiplication operator is most sensitive to Softplus weights (+0.7%)—attributed to its tendency to produce extreme values that necessitate weight regulation.
  • Gains are largest under low-data conditions (+9.23%)—multi-scale fusion acts as effective regularization.
  • Consistent improvements across 6 diverse backbones confirm true plug-and-play applicability.
  • HD95 also improves substantially (14.9 vs. 22.9), indicating enhanced boundary quality.

Highlights & Insights

  • The "mutant" logits concept is highly creative: Framing logits fusion as a combinatorial search space and applying NAS-style search is a natural and elegant transfer of ideas.
  • Strong advantage in low-data settings: Particularly practical for medical image segmentation where annotations are scarce.
  • Fully plug-and-play: Requires no modification to the backbone architecture—CMM and learnable weight layers are appended solely at the decoder end.

Limitations & Future Work

  • Validated only on 2D medical segmentation—3D volumetric or general dense prediction tasks remain untested.
  • Performance on very small organs or lesions is not evaluated.
  • The number of decoder layers \(L\) must remain \(\leq 5\) for computational tractability.
  • vs. Deep Supervision (DS): DS applies uniform weights; LoMix automatically learns heterogeneous weights.
  • vs. Auxiliary Losses: Auxiliary losses typically use fixed ratios; LoMix dynamically adjusts contributions.
  • Insights: The combinatorial mutation idea is transferable to other tasks requiring multi-scale feature fusion.

Rating

  • Novelty: ⭐⭐⭐⭐ Multi-scale logits combinatorial mutation combined with NAS-style weights is a novel design.
  • Experimental Thoroughness: ⭐⭐⭐⭐⭐ 6 backbones + 2 datasets + low-data regime + comprehensive ablation.
  • Writing Quality: ⭐⭐⭐⭐ Method description is clear and well-structured.
  • Value: ⭐⭐⭐⭐ A plug-and-play multi-scale fusion module, particularly suited for low-data scenarios.

Supplementary Method Notes

  • Computational cost of CMM: At \(L=4\), the total is \(4 + 4 \times 11 = 48\) logits—requiring only one CE+Dice loss computation per logit, keeping overhead manageable.
  • Convergence behavior of Softplus weights: Experiments show that in late training, most "mutant" logit weights approach zero while a few key combinations gain large weights—automatically discovering the most valuable scale combinations.
  • Relationship to and distinction from NAS: LoMix searches over loss weights rather than architectures, resulting in a much smaller search space (~50 continuous parameters vs. NAS's \(10^{18}\) discrete space) with no additional search budget required.
  • Clinical significance of HD95 improvement: A reduction from 22.9 to 14.9 indicates a substantial decrease in worst-case (top 5%) boundary errors—critical for applications such as surgical planning that demand precise delineation.
  • Interpretability of dynamic weights: Inspecting per-logit weights after training reveals which scale combinations are most informative, providing a window into how the network leverages multi-scale information.