MoMA: Momentum Contrastive Learning with Multi-head Attention-based Knowledge Distillation for Histopathology Image Analysis

TL;DR

MoMA addresses data scarcity and privacy constraints in computational pathology by distilling knowledge from a high-quality teacher to a target model via momentum-contrastive learning and multi-head attention. The framework integrates a momentum-updated teacher, a memory-bank contrastive objective, and attention-based reweighting to produce context-aware representations, enabling effective cross-task and cross-domain knowledge transfer. Across three distillation scenarios (same, relevant, irrelevant tasks) and multiple datasets, MoMA consistently outperforms baselines, with ablations confirming the necessity of multi-head attention and the benefits of combining L_CE, L_NCE, and L_KL losses. The approach demonstrates robust generalization to unseen data and provides practical guidelines for task-specific knowledge transfer in histopathology, even when source data cannot be accessed directly.

Abstract

There is no doubt that advanced artificial intelligence models and high quality data are the keys to success in developing computational pathology tools. Although the overall volume of pathology data keeps increasing, a lack of quality data is a common issue when it comes to a specific task due to several reasons including privacy and ethical issues with patient data. In this work, we propose to exploit knowledge distillation, i.e., utilize the existing model to learn a new, target model, to overcome such issues in computational pathology. Specifically, we employ a student-teacher framework to learn a target model from a pre-trained, teacher model without direct access to source data and distill relevant knowledge via momentum contrastive learning with multi-head attention mechanism, which provides consistent and context-aware feature representations. This enables the target model to assimilate informative representations of the teacher model while seamlessly adapting to the unique nuances of the target data. The proposed method is rigorously evaluated across different scenarios where the teacher model was trained on the same, relevant, and irrelevant classification tasks with the target model. Experimental results demonstrate the accuracy and robustness of our approach in transferring knowledge to different domains and tasks, outperforming other related methods. Moreover, the results provide a guideline on the learning strategy for different types of tasks and scenarios in computational pathology. Code is available at: \url{https://github.com/trinhvg/MoMA}.

Figures (15)

• Figure 1: Overview of the MoMA: Attention-Augmented Momentum Contrast Knowledge Distillation framework. A batch of input images is encoded by the student encoder ($f^S$), and the momentum teacher ($f^T$), and each feature representation is re-weighted with regard to other images in the batch as the context. A classifier is added on top of the student encoder. The student model is jointly optimized by contrastive loss and cross-entropy loss.
• Figure 2: Overview of distillation flow across different tasks and datasets. 1) Supervised task is always conducted, 2) Feature distillation is applied if a well-trained teacher model is available, and 3) Vanilla ${L}_{KD}$ is employed if teacher and student models conduct the same task. SSL stands for self-supervised learning.
• Figure 3: Box plots for same task distillation: All the KD models utilize the pre-trained weights from PANDA.
• Figure 4: Box plots for relevant task distillation. All the KD models utilize the pre-trained weights from PANDA.
• Figure 5: Bar plots for irrelevant task distillation. All the KD models utilize the pre-trained weights from ImageNet.