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

2026-01-13

AI Summary

This work addresses the challenge of deploying accurate yet efficient 3D medical image segmentation models in resource-constrained clinical settings, where existing lightweight architectures often suffer significant performance degradation. To this end, we propose ReCo-KD, a training-stage knowledge distillation framework that effectively transfers fine-grained anatomical structures and long-range contextual information from a teacher network to a compact student model. Our approach leverages multi-scale structure-aware region distillation (MS-SARD) and multi-scale context alignment (MS-CA), incorporating class-aware masks, scale-normalized weighting, and cross-level feature affinity alignment. Notably, ReCo-KD operates independently of the student backbone and integrates seamlessly with nnU-Net. Extensive experiments demonstrate that our method substantially reduces model parameters and inference latency across multiple public and complex aggregated 3D medical datasets while preserving segmentation accuracy close to that of the teacher model, highlighting its strong potential for clinical deployment.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: Segmentation
• Data Mining & Knowledge Management: Data Compression

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• 3D medical image segmentation
• model efficiency
• lightweight models
• clinical deployment
• performance degradation

Innovation

• Knowledge Distillation
• 3D Medical Image Segmentation
• Region-Aware
• Context Alignment
• Lightweight Model