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

Trinh Thi Le Vuong,Jin Tae Kwak

2023-08-31

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{<a class="link-external link-https" href="https://github.com/trinhvg/MoMA" rel="external noopener nofollow">this https URL</a>}.

Image and Video Processing,Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The problem addressed in this paper is the insufficient generalization ability of models in pathological image analysis due to limitations in data quality and availability. Despite the increasing volume of pathological data, the quality data for specific tasks is still limited, possibly due to privacy and ethical concerns with patient data. The paper proposes a method called MoMA (Momentum Contrastive Learning with Multi-head Attention-based Knowledge Distillation) to overcome these issues using knowledge distillation techniques. Specifically, MoMA adopts a student-teacher framework to learn the target model from a pre-trained teacher model without direct access to the source data, and performs knowledge distillation through momentum contrastive learning and multi-head attention mechanism to provide consistent and context-aware feature representations. This approach allows the target model to absorb the beneficial representations from the teacher model while adapting to the unique characteristics of the target data. Experimental results demonstrate the ability of MoMA to transfer knowledge to different domains and tasks in various scenarios, outperforming other relevant methods, and providing guidance for learning strategies in different types of tasks and scenarios in computational pathology.