PaperlayerView exact revision on arXiv

Source-linked AI summary

Uni-Light: An Ultra-Lightweight Framework via Uncertainty-Aware Knowledge Distillation for Brain Tumour Segmentation

Libing Kuang, Soren Salehi, Ziling Wu, Ahmad P. Tafti, Armaghan Moemeni

arXiv:2609.06729v2cs.CV

TL;DR

3D brain tumour segmentation faces a trade-off between computational cost and fidelity in complex regions. Uni-Light pairs a compact multi-scale student with uncertainty-aware distillation and boundary loss; reported results show strong segmentation performance alongside reduced model and inference costs.

  • Problem

    Brain tumour segmentation methods can be computationally demanding, while lightweight models may lack fidelity in complex tumour regions.

    Source: §1 · p. 2

  • Method

    Uni-Light combines a compact multi-scale student with teacher-uncertainty-weighted distillation and boundary-aware optimisation.

    Source: §1 · p. 2, §3.1 · p. 3

  • Results

    Uni-Light achieves average DSC scores of 89.97% on BraTS2023-GLI and 84.40% on MSD-BTS, with reported SOTA performance and a higher HD95 trade-off.

    Source: Comparison with SOTA Methods · p. 6

  • Takeaways & Limitations

    Teacher uncertainty can reprioritise training samples without additional annotation, while inference efficiency depends solely on the student architecture.

    Source: §5 · pp. 8–9

    Reliance on a single teacher constrains uncertainty estimation, and RCE's emphasis on volumetric accuracy can increase boundary errors in challenging cases.

    Source: §5 · p. 9

Abstract

from arXiv · show

Accurate 3D brain tumour segmentation from multi-modal Magnetic Resonance Imaging (MRI) is essential for clinical diagnosis and treatment planning. Existing brain tumour segmentation methods often suffer from heavy computational demands, while current lightweight architectures frequently lack the capacity to maintain segmentation fidelity in complex tumour regions. To address these issues, we propose a novel ultra-lightweight framework (Uni-Light) that achieves high-fidelity segmentation with substantially reduced computational overhead. It combines multi-scale convolutions with an uncertainty-aware knowledge distillation scheme that directs the student model toward hard-to-classify regions, complemented by a Signed Distance Field boundary loss for geometric constraints. Experimental results on BraTS2023-GLI and MSD-BTS datasets demonstrate that Uni-Light reduces parameters by 97.56%, floating-point operations (FLOPs) by 73.03%, and inference memory footprint by 81.58%, while surpassing the state-of-the-art model by an average of 1.47% in Dice score, offering a highly competitive trade-off between segmentation accuracy and computational efficiency in resource-constrained clinical settings. This work also advances data engineering for medical imaging by demonstrating that teacher model uncertainty can be exploited as a data-driven supervisory signal, re-prioritising the training data distribution without requiring additional annotation.

1 Introduction

3D brain tumour segmentation is important for surgical planning and treatment monitoring, but computationally intensive models and spatially uniform distillation leave efficiency and difficult tumour regions as challenges. Uni-Light combines a compact multi-scale architecture with uncertainty-aware distillation and boundary-aware optimisation.

1 Introduction

  • Transformer and Mamba-based [9] models improve contextual modelling but add prohibitive computational overhead, motivating lightweight designs.Examples of lightweight approaches include depthwise convolutions [15] and multi-scale feature fusion.Source: §1 · p. 2
  • Uniform distillation overlooks spatially varying tumour difficulty, misallocating supervision in ambiguous boundaries and enhancing components.The cited passage identifies this as a limitation of existing distillation frameworks.Source: §1 · p. 2
  • Uni-Light combines compact multi-scale architecture, teacher-uncertainty-weighted supervision, and boundary-aware optimisation for 3D brain tumour segmentation.The framework is evaluated on BraTS and MSD [26] benchmarks.Source: §1 · p. 2

2 Related Work

Prior work reduces 3D segmentation costs through efficient CNN and attention-based designs, while knowledge distillation compresses models using several forms of transferred supervision. Uni-Light focuses on combining architectural efficiency with uncertainty-aware distillation and boundary-aware supervision.

2 Related Work

  • Efficient CNNs use parallel convolutions and parameter aggregation, while lightweight Transformers use axial or hybrid attention to reduce computational overhead.DiffBTS investigates lightweight diffusion models for 3D volumes.Source: §2 · p. 2
  • Knowledge distillation has progressed from softened probabilities to feature hints, attention transfer, and structured alignment for model compression.Medical segmentation has also used distillation to compress 3D convolutional networks and improve lightweight architectures.Source: §2 · p. 3
  • Uniform distillation can overlook ambiguity in tumour boundaries and enhancing subregions, while integration of signed-distance-field losses with spatially adaptive distillation remains underexplored.Uni-Light addresses this gap with a lightweight framework intended to maintain high-fidelity segmentation under computational constraints.Source: §2 · p. 3

3 Methodology: Uni-Light Framework

Uni-Light combines an efficient student, uncertainty-aware distillation, and joint losses for boundary delineation. Its student uses a compact multi-scale backbone and deep supervision; training weights teacher guidance by voxel uncertainty and combines task, distillation, RCE, and boundary losses.

3 Methodology: Uni-Light Framework

  • Uni-Light integrates an efficient student architecture, uncertainty-aware distillation, and a joint loss function for precise boundary delineation.These are the framework's three core components.Source: §3 · p. 3

3.1 Lightweight Student Architecture

  • The student adopts LATUP-Net's Parallel Convolutional Module in a shallower, narrower 3D U-shaped network, removes attention, and adds decoder deep supervision.The PCM extracts multi-scale features through parallel kernel paths; auxiliary heads predict at one-half and one-quarter resolution during training.Source: §3.1 · pp. 3–4

3.2 Dynamic Weighted Knowledge Distillation

  • DWKD derives voxel uncertainty from the teacher's softened predictions and uses it to assign dynamic supervision weights, with uncertainty often highest at tumour boundaries or complex tissues.Weights are normalized by the maximum entropy in the current batch, with a small constant preventing a zero denominator.Source: §3.2 · pp. 4–5

3.3 Loss Function

  • The composite objective balances task, DWKD, RCE, and boundary losses; task loss spans tumour regions and deep-supervision levels, with extra penalty weight for enhancing tumour.The loss coefficients are set to λ1 = 0.6, λ2 = 0.2, λ3 = 0.1, and λ4 = 0.1; region-specific penalties are v = {1, 1, 3}.Source: §3.3 · p. 5
  • RCE uses a confidence-thresholded teacher mask to ignore selected voxels, suppress confident background regions, and focus learning on uncertain areas.DWKD weights the KL divergence between teacher and student predictions voxel by voxel.Source: §3.3 · p. 5
  • The SDF boundary loss supplies continuous geometric supervision, penalizing background false positives and rewarding correct predictions inside the tumour.The signed distance map assigns negative distances inside the tumour and positive distances outside it.Source: §3.3 · pp. 5–6

4 Experiments

Experiments evaluate Uni-Light on BraTS2023-GLI and MSD-BTS, measuring segmentation accuracy and computational efficiency. They report benchmark results, implementation details, and ablations of the model’s components.

Datasets.

  • The evaluation uses 1,251 BraTS2023-GLI cases and 484 MSD-BTS cases, each with four MRI sequences and annotations for three tumour sub-regions.Results are reported for ET, TC, and WT.Source: Datasets · p. 6

Evaluation Metrics.

  • DSC and HD95 measure segmentation accuracy, while parameters, FLOPs, latency, inference memory, and training memory assess computational efficiency.HD95 evaluates boundary delineation using the 95th percentile of distances between predicted and ground-truth surface points.Source: Evaluation Metrics · p. 6

Implementation Details.

  • Training uses AdamW on a single NVIDIA T4 GPU with batch size 2 and five-fold cross-validation, with on-the-fly flipping, rotations, intensity scaling, and Gaussian noise.Volumes are resampled to 1283; the supplied augmentation passage lists random flipping and 90° rotations.Source: Implementation Details · p. 6

Comparison with SOTA Methods.

  • 89.97% average DSC on BraTS2023-GLI and 84.40% on MSD-BTS; Uni-Light surpasses BraTS-UMamba in ET and TC and has 3.5× and 4.7× lower latency than SwinUNETR and SegMamba, respectively.The reported average DSC scores are on Table 1’s two datasets; latency comparisons are reported in the accompanying results.Source: Comparison with SOTA Methods · p. 6, Table 1 · p. 7, Table 2 · p. 7
  • The cited passage says Uni-Light’s training footprint in MB avoids OOM issues typical of heavy architectures on standard hardware, while Figure 3 comparisons show superior segmentation fidelity and sharper boundaries.Source: Comparison with SOTA Methods · p. 7

Ablation Study.

  • 89.95% DSC with RCE (Ours) accompanies 6.29 mm HD95, versus 3.88 mm for S3; SDF boundary loss reduces HD95 from 5.14 to 4.23 mm.The authors attribute the higher final HD95 to sparse boundary outliers from RCE and describe this effect as orthogonal to boundary precision.Source: Table 3 · p. 8, Ablation Study · p. 8
  • DWDK (S1) weights supervision by teacher uncertainty, SDF loss (S2) refines contours, and PCM (S3) adds multi-scale context.The ablation evaluates these components on BraTS2023.Source: Ablation Study · p. 8, Table 3 · p. 8

5 Conclusion and Discussion

Uni-Light combines multi-scale convolutions, uncertainty-aware distillation, and boundary-aware optimization; experiments on BraTS2023 and MSD-BTS report maintained segmentation fidelity alongside lower parameter and computational costs. Teacher uncertainty also serves as a data-driven signal for reprioritizing training samples without additional annotation.

5 Conclusion and Discussion

  • On BraTS2023 and MSD-BTS, Uni-Light maintains high segmentation fidelity while reducing parameters and computational cost; inference efficiency depends only on the student architecture.The teacher model influences the student’s performance ceiling.Source: §5 · p. 8
  • Teacher uncertainty can reprioritize training samples as a data-driven supervisory signal without additional annotation.Source: §5 · p. 9
  • A single teacher limits uncertainty quality, and RCE’s volumetric emphasis can increase boundary errors in challenging cases.The authors suggest postprocessing as a possible mitigation and plan to explore stronger attention-based teachers and validation on additional medical imaging tasks.Source: §5 · p. 9

Related papers

Loading 2609.06729v2…