PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Medical Image Segmentation Based on Knowledge Distillation

Dian Qin, Jiajun Bu, Zhe Liu, Xin Shen, Sheng Zhou, Jingjun Gu, Zhijua Wang, Lei Wu, Huifen Dai

arXiv:2108.09987v1eess.IVcs.CV

TL;DR

Medical image segmentation models can achieve strong accuracy but often require computational resources and storage that hinder practical deployment. The paper distills knowledge from cumbersome teachers into lightweight students and introduces RAD to transfer semantic region information through inter-class contrasts. Across LiTS and KiTS19 experiments, the approach substantially improves lightweight segmentation performance while preserving portability and narrowing the teacher–student gap.

  • Problem

    Existing medical image segmentation methods often rely on large computational complexity and storage, creating a practical efficiency challenge.

  • Method

    The paper trains lightweight student networks through knowledge distillation and uses RAD to transfer semantic region information via teacher-derived region contrast maps.

  • Results

    32.6% improvement raised tumor-segmentation Dice from 0.516 to 0.684, while the method narrowed the teacher–student gap from 0.229 to 0.061.

  • Takeaways & Limitations

    The distilled lightweight network provides a portable alternative for medical image segmentation scenarios requiring higher operating speed and lower storage usage.

Abstract

from arXiv · show

Recent advances have been made in applying convolutional neural networks to achieve more precise prediction results for medical image segmentation problems. However, the success of existing methods has highly relied on huge computational complexity and massive storage, which is impractical in the real-world scenario. To deal with this problem, we propose an efficient architecture by distilling knowledge from well-trained medical image segmentation networks to train another lightweight network. This architecture empowers the lightweight network to get a significant improvement on segmentation capability while retaining its runtime efficiency. We further devise a novel distillation module tailored for medical image segmentation to transfer semantic region information from teacher to student network. It forces the student network to mimic the extent of difference of representations calculated from different tissue regions. This module avoids the ambiguous boundary problem encountered when dealing with medical imaging but instead encodes the internal information of each semantic region for transferring. Benefited from our module, the lightweight network could receive an improvement of up to 32.6% in our experiment while maintaining its portability in the inference phase. The entire structure has been verified on two widely accepted public CT datasets LiTS17 and KiTS19. We demonstrate that a lightweight network distilled by our method has non-negligible value in the scenario which requires relatively high operating speed and low storage usage.

I. INTRODUCTION

Medical image segmentation requires accurate pixel-level predictions despite challenging tissue appearances and locations, while high-performing models remain costly to deploy. The paper addresses this trade-off with knowledge distillation and a region-aware module that transfers segmentation capability to lightweight networks.

  • Medical image segmentation is difficult because organs and tumors exhibit diverse appearances, irregular sizes, unpredictable locations, and contrast-agent variations.
  • Performance tends to decline when segmentation models are simplified for faster speed.
  • The proposed architecture distills segmentation capability from well-trained medical image networks into lightweight student networks.
  • RAD transfers semantic region information by making students mimic inter-class contrasts computed from intermediate feature maps using ground-truth masks.
  • RAD avoids ambiguous boundaries by encoding internal information within each semantic region, helping correct subtle student segmentation errors.
  • 32.6% improvement raised tumor-segmentation Dice from 0.516 to 0.684, while the student network remained 21 times smaller than its teacher.
  • The method was evaluated on the public LiTS and KiTS19 datasets with ablation considerations.

II. RELATED WORK

Prior medical segmentation research emphasizes accurate architectures, whereas practical deployment also requires lightweight models. Knowledge distillation provides a route to transfer capability from computationally demanding teachers to compact students without affecting efficiency.

  • Medical Image Segmentation: Medical segmentation research has developed UNet variants, GAN-based methods, PSPNet, and Deeplab architectures for challenging imaging tasks.
  • Medical Image Segmentation: Many established segmentation models perform well but require relatively high storage and computational resources, making them suitable as teachers.
  • Medical Image Segmentation: No widely accepted lightweight model dedicated to medical image segmentation had emerged, although lightweight networks had been explored for specific applications.
  • Knowledge Distillation: Knowledge distillation transfers knowledge from powerful cumbersome networks to lightweight models to improve performance without affecting efficiency.
  • Knowledge Distillation: Existing distillation research includes output-logit, intermediate-feature, attention-map, similarity-preserving, detection, and segmentation approaches.

III. METHODOLOGY

The architecture distills prediction and intermediate feature information from a teacher into a lightweight student through complementary modules. IMD aligns feature-map scales and transfers importance maps, while RAD transfers region relationships and PMD transfers pixel-level predictions.

  • The architecture uses IMD, RAD, and PMD to transfer intermediate and prediction knowledge from teacher to student networks.IMD and RAD transfer intermediate information through importance and region affinity maps, while PMD transfers predictive capability from the teacher’s segmentation output.
  • PMD treats segmentation as pixel-level classification and compares teacher and student probabilities at corresponding spatial positions.The module calculates a loss for pixel pairs at the same spatial locations using Kullback-Leibler divergence.
  • IMD rescales student feature maps to the teacher’s spatial size before constructing comparable importance maps.The rescaling uses unpooling, pooling, or no operation according to the relative spatial sizes of the feature maps.
  • Importance maps sum absolute feature activations across channels, and their distillation loss uses normalized differences between corresponding maps.The l1 norm is used because importance maps are relatively sparse in medical image segmentation scenarios.
  • The rescaling design enables distillation between teacher and student feature maps with different spatial sizes.This removes the requirement for strictly identical feature-map spatial dimensions.

C. Region Affinity Distillation

RAD transfers relationships among semantic regions rather than relying on ambiguous tumor boundaries. It extracts class-specific feature information with labeled masks and trains the student to match the teacher’s region contrasts for binary or multi-class segmentation.

  • RAD transfers relationship information between semantic regions from the teacher network to the student network.The module is designed to encode regional knowledge and semantic differences among regions for medical image segmentation.
  • RAD uses resized labeled masks to extract class-specific region information from intermediate feature maps.Region contrast values are computed by measuring similarity among the extracted class regions.
  • The region affinity loss compares student and teacher region contrast values or vectors using a selected norm.The scalar form averages similarities across class pairs, while the vector form retains all pairwise similarities for numerous semantic classes.
  • Binary segmentation matches contrast between target and background, whereas multi-class segmentation matches contrasts across all class pairs.The multi-class contrasts form a graph representing relationships among the semantic regions.

D. Training Process

The modules are integrated into an end-to-end objective combining segmentation and distillation losses. After training, the teacher and distillation modules are removed, leaving the lightweight student for inference.

  • The proposed modules are integrated and optimized end-to-end with a total loss that includes the general segmentation loss.The segmentation term can be cross entropy or Dice loss.
  • The experiments set α to 0.1 and β1 and β2 to 0.9, with matched β values reported as insensitive to individual fluctuations.The stated insensitivity is supported by the corresponding ablation results.
  • The training procedure updates only the student using a pretrained teacher and representative low-level and high-level feature pairs.Using two to four feature pairs with IMD and RAD is reported as the most efficient choice.
  • The teacher and distillation modules are discarded during inference after training.The resulting student retains its parameter count while receiving segmentation improvements from distillation.

A. Setup

The evaluation uses established teacher and lightweight student architectures on LiTS and KiTS19 CT datasets under a standardized training setup. Models are trained with fixed preprocessing, cross-validation, and convergence procedures.

  • RA-UNet serves as a teacher architecture, while ENet and other open-source lightweight networks serve as students.The experiments follow the official network structures and hyperparameters when training the architectures independently.
  • The experiments use 512 × 512 CT inputs, dataset-specific HU windowing, PyTorch implementations, and an NVIDIA GeForce RTX 3090 GPU.The liver window is −40 to 160 HU and the kidney window is −200 to 300 HU.
  • Training runs for up to 60 epochs with five-fold cross-validation, and reported scores are ranges derived from the last 20 epochs of each fold.
  • LiTS contains 201 CT scans collected with different scanners and acquisition protocols, with diverse liver-tumor presentations and image resolutions.The dataset includes varying axial and longitudinal resolutions, slice counts, and tumor sizes.
  • KiTS19 contains 210 abdominal CT scans with manual kidney and kidney-tumor masks and slice thicknesses ranging from 1 mm to 5 mm.The dataset also spans varying axial resolutions, fields of view, and tumor volumes.

C. Evaluation Metric

The experiments use per-case Dice coefficient as the primary segmentation metric, with VOE and RVD provided as complementary error measures. The evaluation also spans multiple teacher–student networks and public segmentation benchmarks.

  • Evaluation metrics: Dice coefficient is the chief evaluation metric, reported uniformly per case for volumetric segmentation.The metric compares prediction and ground-truth volumetric tumor masks.
  • Evaluation metrics: VOE and RVD complement Dice as error metrics whose values, or absolute values for RVD, should be minimized.Dice instead rewards larger values.
  • Experimental setup: The ablation evaluates distillation across multiple teacher and lightweight student network combinations.Teachers include RA-UNet, PSPNet, and UNet++, while students include ENet, MobileNetV2, and ResNet-18.
  • Experimental setup: The experiments cover LiTS17 and KiTS19 and include comparisons with contemporary segmentation networks and alternative distillation methods.The supplied table captions identify cross-experiments on LiTS and KiTS19 and comparisons with other knowledge-distillation methods.

1) Primary Results:

Across LiTS and KiTS19, the distillation architecture improves lightweight students using stronger teachers, with especially large gains for kidney-tumor segmentation. Students can approach teacher performance, and the method often accelerates convergence.

  • Primary results: 13.6% (0.078) is the largest reported liver-tumor Dice improvement, achieved by ENet among the evaluated lightweight students.MobileNetV2 and ResNet-18 improve by 10% (0.055) and 9.5% (0.044), respectively.
  • Primary results: 32.6% is the reported kidney-tumor segmentation improvement, while MobileNetV2 gains a 0.168 Dice-score increase with RA-UNet.The 0.168 increase is identified as the most visible promotion value in the experiments.
  • Primary results: 0.967 is the kidney-segmentation Dice score reached by ENet after learning from PSPNet.The paper identifies ENet as the strongest student for kidney segmentation in this result.
  • Primary results: Some students approach teacher-level performance across all four segmentation tasks.The reported tasks include liver tumor, liver, kidney tumor, and kidney segmentation.
  • Primary results: The method often accelerates convergence, bringing students that train poorly toward their teacher’s level during training.Validation trends are coordinated by epochs and evaluated with Dice coefficient.

2) Contemporary Rank:

The contemporary-method comparison positions distilled ENet as a low-parameter model with competitive tumor-segmentation performance. It ranks highly for liver tumor and outperforms most listed methods for kidney tumor.

  • Contemporary rank: The comparison orders all networks by ascending parameter count and identifies RA-UNet as ENet’s teacher.FLOPs are also listed for a constant 384×384 input.
  • Contemporary rank: ENet distilled from RA-UNet ranks fourth in liver-tumor Dice and surpasses PSPNet and DeeplabV3+.The comparison includes contemporary segmentation models arranged by parameter count.
  • Contemporary rank: 0.676 is ENet’s kidney-tumor Dice score, exceeded only by RA-UNet among the compared models.The ranking compares methods while retaining the distilled student’s small model size.

3) Comparison with Other Knowledge Distillation Methods:

EMKD outperforms the compared knowledge-distillation methods on both evaluated tasks, with particularly strong kidney-tumor results. Component ablations show positive contributions, while gains are limited when teacher–student performance gaps are small.

  • Distillation comparison: EMKD achieves the strongest results among the compared distillation methods on both evaluated tasks.The comparison fixes the teacher as RA-UNet and the student as ENet.
  • Distillation comparison: The advantage of EMKD is especially pronounced for kidney-tumor segmentation.The paper reports this qualitative comparison across the knowledge-distillation methods.
  • Component ablation: PMD, IMD, and RAD each positively affect student-network performance in the component ablations.The components are evaluated on LiTS and KiTS19 using Dice, VOE, and RVD.
  • Component ablation: Performance gains are theoretically difficult to make remarkable when the teacher–student performance gap is tiny.The paper gives liver segmentation as the example of this limitation.
  • Hyper-parameter analysis: Doubling the component weights causes only slight performance drops, while changing β1 or β2 alone has negligible influence.The authors therefore prefer altering β1 and β2 simultaneously in practice.

V. DISCUSSION

The architecture transfers structural and semantic information from cumbersome segmentation networks to lightweight networks through modular knowledge distillation. Experiments show improved lightweight-network performance, while future work remains for 3D settings and multi-model distillation.

  • Architecture: The architecture uses three distillation modules to guarantee basic transfer, emphasize important neurons, and extract inter-class semantic information.Its design supports encoder-decoder teacher and student networks, including heterogeneous architectures.
  • Extensibility: The modules are replaceable, extensible, and applicable to other semantic segmentation problems requiring distillation.The authors describe adding or replacing modules and transferring relationships between classified regions to general segmentation tasks.
  • Experimental results: 0.678 and 0.679 Dice scores were achieved by MobileNetV2 after distillation from PSPNet and UNet++, exceeding the teachers’ 0.659 and 0.644 scores.These results were reported for kidney tumor segmentation, although the proposed explanation requires future verification.
  • Limitations: 3D extension remains constrained by feature-map computation, storage demands, small objective-to-background ratios, and dataset suitability.The authors note that LiTS and KiTS19 have slice counts that cause the z dimension to disappear during convolution, despite the architecture being extendable to 3D.
  • Limitations: Integrating multiple teachers or students remains challenging because their feature maps can have different sizes while passing through standardized distillation modules.The paper identifies multi-model distillation as future work.
  • Experimental results: The lightweight network achieved performance comparable to state-of-the-art cumbersome networks after structural information was transferred through distillation.The conclusion characterizes the experiments as demonstrating a substantial improvement in the lightweight network.
Loading 2108.09987v1…