Graph Relation Distillation for Efficient Biomedical Instance Segmentation

TL;DR

This work tackles the efficiency gap in biomedical instance segmentation by introducing a graph-based knowledge distillation framework. It combines Instance Graph Distillation (IGD) to transfer instance-level features and relations and Affinity Graph Distillation (AGD) to transfer boundary structure, with intra-image and inter-image variants enabled by a memory-bank that captures global relationships across images. The approach yields lightweight student models with less than $1\%$ of the parameters and less than $10\%$ of the inference time of the teacher while maintaining strong segmentation performance across 2D and 3D biomedical datasets. Overall, the graph-based distillation strategy significantly narrows the teacher–student gap and enables practical deployment for resource-constrained biomedical imaging tasks.

Abstract

Instance-aware embeddings predicted by deep neural networks have revolutionized biomedical instance segmentation, but its resource requirements are substantial. Knowledge distillation offers a solution by transferring distilled knowledge from heavy teacher networks to lightweight yet high-performance student networks. However, existing knowledge distillation methods struggle to extract knowledge for distinguishing instances and overlook global relation information. To address these challenges, we propose a graph relation distillation approach for efficient biomedical instance segmentation, which considers three essential types of knowledge: instance-level features, instance relations, and pixel-level boundaries. We introduce two graph distillation schemes deployed at both the intra-image level and the inter-image level: instance graph distillation (IGD) and affinity graph distillation (AGD). IGD constructs a graph representing instance features and relations, transferring these two types of knowledge by enforcing instance graph consistency. AGD constructs an affinity graph representing pixel relations to capture structured knowledge of instance boundaries, transferring boundary-related knowledge by ensuring pixel affinity consistency. Experimental results on a number of biomedical datasets validate the effectiveness of our approach, enabling student models with less than $ 1\%$ parameters and less than $10\%$ inference time while achieving promising performance compared to teacher models.

Figures (6)

• Figure 1: Workflow of our proposed graph relation distillation method for biomedical instance segmentation, which includes two schemes. The instance graph distillation (IGD) scheme constructs instance graphs from embeddings of the teacher-student network pair and enforces the consistency of graphs constructed by the teacher, while the affinity graph distillation (AGD) scheme converts pixel embeddings into pixel affinities that encode structured information of instance boundaries and enforces the student model to generate affinities similar to its teacher model. These two schemes take charge of the knowledge distillation mechanism and are carried out at both intra-image and inter-image levels for global instance relations. The symbol $\odot$ represents dot product operation. The red arrow indicates the loss function.
• Figure 2: Visual comparisons on three 2D datasets. We use networks ResUNet (T1) and MobileNet (S2) as the teacher and student networks, respectively. Over-merge and over-segmentation in the results of the student network are highlighted by red and white boxes, respectively.
• Figure 3: 2D visual comparisons of segmentation results on the CREMI-C and AC3/4 dataset.
• Figure 4: 3D visual comparisons on the CREMI-C and AC3/4 dataset. Red and black arrows indicate over-segmentation and over-merge, respectively.
• Figure 5: A visual example of different embedding maps predicted by student networks distilled with different knowledge distillation methods. We use networks ResUNet (T1) and MobileNet (S2) as the teacher and student networks, respectively.