Source-linked AI summary
ReCo-KD: Region- and Context-Aware Knowledge Distillation for Efficient 3D Medical Image Segmentation
Qizhen Lan, Yu-Chun Hsu, Nida Saddaf Khan, Xiaoqian Jiang
TL;DR
3D medical segmentation needs accurate anatomical delineation, but high-capacity models can exceed the computing resources available in clinical settings and lightweight models may lose accuracy. ReCo-KD distills region-aware detail and multiscale contextual relationships into a compact student, achieving near-teacher performance while reducing computational demands. Its evaluation spans multiple public datasets and an aggregated dataset, with scope limited to a default nnU-Net student and homogeneous CNN-to-CNN distillation.
Problem
High-capacity 3D segmentation models can be impractical for resource-limited clinical deployment, while compact architectures typically trade accuracy for efficiency.
Method
ReCo-KD combines MS-SARD for critical-region supervision with MS-CA for multiscale teacher–student contextual alignment during training.
Results
ReCo-KD consistently narrows the teacher–student performance gap while reducing parameters and FLOPs by up to 94% and 93%, respectively, across multiple datasets.
Takeaways & Limitations
The training-only framework provides near-teacher segmentation accuracy with lower computational demands and no additional inference cost.
Takeaways & Limitations
Evaluation is restricted to the default nnU-Net student and homogeneous CNN-to-CNN distillation, leaving broader student designs and heterogeneous distillation for future work.
Abstract
from arXiv · showhide
Accurate 3D medical image segmentation is vital for diagnosis and treatment planning, but state-of-the-art models are often too large for clinics with limited computing resources. Lightweight architectures typically suffer significant performance loss. To address these deployment and speed constraints, we propose Region- and Context-aware Knowledge Distillation (ReCo-KD), a training-only framework that transfers both fine-grained anatomical detail and long-range contextual information from a high-capacity teacher to a compact student network. The framework integrates Multi-Scale Structure-Aware Region Distillation (MS-SARD), which applies class-aware masks and scale-normalized weighting to emphasize small but clinically important regions, and Multi-Scale Context Alignment (MS-CA), which aligns teacher-student affinity patterns across feature levels. Implemented on nnU-Net in a backbone-agnostic manner, ReCo-KD requires no custom student design and is easily adapted to other architectures. Experiments on multiple public 3D medical segmentation datasets and a challenging aggregated dataset show that the distilled lightweight model attains accuracy close to the teacher while markedly reducing parameters and inference latency, underscoring its practicality for clinical deployment.
I. INTRODUCTION
3D medical segmentation models must balance fine anatomical detail, long-range context, and deployment efficiency. ReCo-KD addresses these challenges by distilling region-aware and contextual representations from a high-capacity teacher into a lightweight student.
- Motivation: Compute- and memory-intensive models limit deployment on CPU-only, mobile, and point-of-care systems despite strong segmentation capacity.The introduction frames clinical adoption as constrained more by efficiency than algorithmic accuracy.
- Motivation: Small critical structures occupy less than 1% of voxels, while the largest-to-smallest class volume ratio exceeds 200:1.This imbalance biases distillation toward dominant structures and weakens supervision for rare regions.
- Motivation: 3D segmentation is difficult because volumetric networks must handle high spatial resolution, dense voxel dependencies, and memory constraints.
- Motivation: Representational discrepancies between teacher and student occur in spatial, channel, and contextual features, motivating alignment beyond final outputs.The framework targets the student’s intermediate feature hierarchy rather than relying only on output-level guidance.
- Proposed Framework: ReCo-KD jointly applies MS-SARD and MS-CA to transfer fine-grained structure and long-range anatomical dependencies during training.MS-SARD uses class-aware masks, scale-normalized weighting, and attention-enhanced matching; MS-CA aligns relational features across scales.
- Proposed Framework: The training-only framework integrates with state-of-the-art backbones, including nnU-Net, without adding inference cost.The implementation provides automatic configuration and multi-modal input support through nnU-Net.
- Results: Experiments across public datasets and an aggregated dataset report near-teacher accuracy with reduced parameters, FLOPs, and inference latency.
II. RELATED WORKS
Prior work spans CNN and Transformer segmentation architectures, lightweight model design, and knowledge distillation. ReCo-KD targets under-addressed 3D imbalance and anatomical-context challenges through region reweighting and multiscale relational alignment.
- 3D Medical Image Segmentation: CNN-based U-Net variants provide strong 3D segmentation performance, while Transformer methods model long-range dependencies but can increase computational cost.
- Efficient Architectures: Lightweight architectures reduce model size and computational cost through techniques such as early downsampling, factorized convolutions, and depthwise separable convolutions.Examples include ENet, ERFNet, MobileNet-family models, Mobile-UNet variants, and other efficient 3D segmentation designs.
- Knowledge Distillation: Knowledge distillation has expanded from softened classification outputs to feature transfer, structural relations, and class-wise affinity modeling for dense prediction.
- Knowledge Distillation: Medical-image KD has explored boundary-guided, region-wise, and multiscale structured distillation, but many methods remain focused on 2D settings.
- Knowledge Distillation: ReCo-KD reweights supervision toward rare critical voxels and aligns inter-voxel relations across feature levels to preserve global anatomical consistency without inference-time overhead.
III. METHODOLOGY
ReCo-KD distills multi-scale encoder features from a full teacher into a channel-reduced nnU-Net student using complementary region-aware and context-alignment objectives. Its training-only design addresses class imbalance, activation mismatch, and long-range anatomical dependencies without adding inference-time complexity.
- Student construction: The compact nnU-Net student uniformly reduces channel width while preserving depth, strides, skip connections, and input resolution.Students are parameterized by t ∈ {0, 1, 2, 3}, with a minimum channel value of 4 used under high compression.
- Distillation framework: ReCo-KD distills encoder outputs at multiple stages through feature alignment and complementary structure- and context-aware objectives.The teacher and student use the same backbone, while distillation operates only during training.
- Deployment: All distillation components are training-only, so inference complexity remains identical to that of the lightweight student.The teacher is used in the training pipeline, while deployment retains the compact student alone.
- MS-SARD: MS-SARD uses binary region masks and class-wise scale weighting to reduce background dominance and emphasize small, clinically critical structures.Large-volume regions receive smaller weights, whereas under-represented regions receive larger weights.
- MS-SARD: MS-SARD further reweights feature matching with teacher-derived spatial and channel activation masks, highlighting informative voxels and channels.The resulting loss combines region, scale, spatial-activation, and channel-activation masks across encoder stages.
- MS-CA: MS-CA complements region distillation by transferring long-range dependencies through teacher–student affinity alignment across feature levels.It uses a lightweight 3D global-context operator to preserve contextual coherence across the volume.
B. Multi-Scale Contextual Alignment (MS-CA)
MS-CA transfers holistic contextual patterns from teacher to student through a lightweight 3D global-context operator. It aligns contextual representations across multiple feature levels while preserving localized MS-SARD supervision.
- Motivation: MS-CA addresses the limitation that region-focused MS-SARD may underrepresent global dependencies across 3D volumes.Long-range interactions between spatially distant anatomical structures are important for structural completeness and global consistency.
- Method: The module transfers holistic contextual patterns from teacher to student using a lightweight 3D global-context operator adapted from GC-blocks.It is designed to provide contextual guidance without interfering with localized supervision from MS-SARD.
- Multi-scale alignment: The contextual alignment loss is computed across multiple stages to align teacher and student global contextual representations at different scales.The hyperparameter λ controls the contribution of this contextual alignment term.
- Context operator: The contextual operator aggregates global features through soft attention and refines them with a residual bottleneck that preserves spatial semantics.F denotes the 3D feature map, and Nv = D × H × W is the number of voxels.
C. Overall Loss
The overall objective combines three components, including the standard segmentation loss applied to the student prediction and ground-truth labels.
- Objective: The overall objective is defined as a combination of three components.
- Task loss: Ltask denotes the standard segmentation loss, such as Dice loss and cross-entropy, between student predictions and ground-truth labels.
A. Experimental Setups
The experiments evaluate ReCo-KD across three public 3D segmentation benchmarks and one more complex private brain-structure task using nnU-Net-based teacher and width-scaled student models. Evaluation follows nnU-Net preprocessing and a single-fold reporting protocol unless otherwise specified.
- Datasets: ReCo-KD is evaluated on four datasets: three public benchmarks and one private, more complex task.
- Datasets: BraTS2021 provides multiparametric MRI with tumor-region labels and uses an 80:20 training-validation split across 1251 cases.
- Datasets: MSD Hippocampus contains single-modality MRI with anterior and posterior hippocampus annotations across 263 training and 131 test volumes.
- Datasets: BTCV contains abdominal CT with 13 annotated organs and uses a 24/6 training-validation split from 30 labeled training volumes.
- Datasets: The private brain-structure dataset aggregates multiple cohorts into 110 anatomical categories, with 1,189 training subjects and evaluation on Mindboggle-101.
- Implementation: The implementation uses nnU-Net defaults for preprocessing, planning, deep supervision, augmentation, and sliding-window inference, with a residual-encoder teacher.
- Implementation: Students are formed by uniformly scaling channel width while keeping network depth, strides, patch size, and batch size unchanged.
- Implementation: A 1 × 1 × 1 adapter aligns features when teacher and student channel dimensions differ; temperature and loss weights are fixed to 0.5 and 1.
B. Main Results
Across abdominal, hippocampal, brain-tumor, and fine-grained brain-structure evaluations, ReCo-KD preserves or improves segmentation accuracy while substantially reducing model cost. The method is also competitive with existing knowledge-distillation approaches.
- BTCV: 85.01% mean Dice on BTCV surpasses SwinUNETRv2 at 81.26% and MedNeXt at 82.98%, with the largest gains on small or rare organs.Performance on large organs such as liver and spleen remains strong.
- Hippocampus: 88.93% mean Dice on Hippocampus is achieved with 1.57 M parameters and 9.17 GFLOPs at t=3, matching or exceeding lightweight baselines.The setting uses one-eighth of the teacher’s channel width.
- BraTS2021: 91.09% average Dice on BraTS2021 approaches the teacher’s 91.65%, with the largest improvement occurring on enhancing tumor at +2.21 Dice over the non-distilled student.Whole Tumor and Tumor Core performance remain stable.
- Large-scale brain structure dataset: With one-quarter of the teacher channels, the large-scale brain student retains 98.29% of teacher accuracy while reducing parameters by 93.72% and FLOPs by 93.48%.CPU inference time falls from 119 s to 34.6 s, a 70.92% reduction.
- KD comparison: On BTCV, ReCo-KD achieves 85.01% Dice and 84.29% NSD, while on BraTS2021 it reaches 91.09% Dice and 93.83% NSD among compared KD methods.
D. Qualitative Analysis
Qualitative comparisons show that ReCo-KD improves the compact student's segmentation appearance across BraTS2021 and BTCV, especially around boundaries and spurious regions.
- BraTS2021: ReCo-KD produces crisper BraTS2021 boundaries and better overlap across tumor subregions than the non-distilled Student.The Student tends to under-segment enhancing tumor and generate irregular boundaries near the core.
- BTCV: The qualitative BTCV comparison presents axial slices with CT images, ground truth, predictions from different methods, and ReCo-KD.Each row corresponds to one axial slice.
- BTCV: On BTCV, ReCo-KD shows crisper organ boundaries and fewer spurious regions than the compared methods.The comparison uses identical windowing and places ReCo-KD at the far right of each row.
E. Ablation Studies
Ablations on BTCV indicate that region-based distillation and contextual alignment each improve the compact baseline, while their combination gives the strongest overall result.
- Component ablations: All single-component settings improve Dice and NSD over baseline while reducing HD95.The tested settings include FG-distill, BG-distill, Mask-align, and MS-CA.
- Component ablations: The ablation evaluates Dice, NSD, and HD95 for variants using the non-distilled Student as the baseline.FG-distill uses teacher foreground masks, BG-distill uses background voxels, Mask-align removes region loss, and MS-CA performs contextual alignment only.
- Experimental setting: All student models use uniform channel reduction to one-quarter of the teacher's channels in the reported comparison.The teacher is the full-capacity nnU-Net with a residual encoder.
- Component ablations: +2.00% Dice is obtained by the full ReCo-KD combination relative to the non-distilled Student baseline.The combined MS-SARD and MS-CA setting attains the best overall results.
2) Efficiency Analysis of Channel Reduction Factor:
Channel reduction substantially lowers computational demands while preserving accuracy through t=2, making t=2 the adopted efficiency–accuracy trade-off.
- Efficiency results: 93.7% FLOPs and 67.2% peak memory are reduced at t=2 relative to t=0.The empirical scaling closely follows the expected width-reduction trend.
- Efficiency results: Inference latency falls from 10.38s to 3.40s, approximately 3.05× faster, under channel reduction.Latency is measured per case on BTCV using a single RTX 2080 Ti with AMP.
- Efficiency results: Accuracy remains effectively unchanged through t=2, with mDice decreasing from 85.64% to 85.01%.Accuracy degrades at t=3, so t=2 is adopted as the default trade-off.
- Qualitative context: The qualitative BraTS2021 and BTCV analyses compare teacher, student, ground truth, and ReCo-KD outputs across anatomical views and slices.BraTS2021 uses axial, sagittal, and coronal views; BTCV uses axial rows.
- Scope: ReCo-KD is evaluated with a default nnU-Net student and homogeneous CNN-to-CNN distillation.This setting restricts architectural search and may understate benefits from stronger lightweight students.