ReCo-KD: Region- and Context-Aware Knowledge Distillation for Efficient 3D Medical Image Segmentation

TL;DR

ReCo-KD tackles the deployment gap in 3D medical image segmentation by training a compact student to inherit both fine-grained anatomical detail and long-range context from a high-capacity teacher. It introduces two complementary modules, MS-SARD and MS-CA, to perform region-aware structure distillation and contextual alignment within the nnU-Net framework. The approach achieves near-teacher accuracy across BTCV, BraTS2021, Hippocampus, and a large 110-class brain dataset while drastically reducing parameters, FLOPs, and CPU latency, demonstrating strong practical impact for resource-constrained clinical settings. Its backbone-agnostic, training-time-only design makes it straightforward to integrate into existing workflows for efficient, accurate 3D segmentation.

Abstract

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.

Figures (6)

• Figure 1: Voxel distribution across background and foreground classes in three medical segmentation datasets. The left donut charts show the dominance of background voxels—BraTS 2021: 98.9%, Hippocampus: 94.7%, BTCV: 95.5%. The right bar charts reveal strong foreground imbalance, e.g. BTCV (Liver 2.45%; most other organs <1%). This imbalance suggests that equal voxel weighting in knowledge distillation may overlook small yet clinically critical structures.
• Figure 2: Visualization of spatial, channel, and contextual representations from the teacher and the student (encoder stage 1). The green square highlights the anchor point, and the surrounding maps show responses relative to this position. The discrepancies across all levels indicate the representational gap, motivating the need for knowledge distillation to align them.
• Figure 3: Overview of the proposed Region- and Context-aware Knowledge Distillation (ReCo-KD) framework for 3D medical image segmentation. Teacher and student share the same backbone, with the student using a channel-reduced width ($C/2^{t}$). Multi-scale feature maps from all encoder stages are distilled by two complementary modules: Multi-Scale Structure-Aware Region Distillation (MS-SARD), which highlights class-specific regions with scale-normalized weighting, and Multi-Scale Context Alignment (MS-CA), which aligns teacher–student affinity patterns to transfer long-range dependencies. Together these modules enable the compact student to achieve near-teacher accuracy with greatly reduced computation.
• Figure 4: Illustration of contextual alignment distillation using the 3D Global Context Block. Feature volumes from both teacher and student encoder are taken as inputs to align contextual representations across stages.
• Figure 5: Qualitative results on BraTS2021. Rows show axial, sagittal, and coronal views. The first column is the full slice with a red dashed box marking the region of interest; the others show the cropped region for Ground Truth, Teacher, Student, and our ReCo-KD. Zoom for the best view.