Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MIRAM: Masked Image Autoencoders Across Multiple Scales with Hybrid-Attention Mechanism for Breast Lesion Risk Prediction

Hung Q. Vo, Pengyu Yuan, Zheng Yin, Kelvin K. Wong, Chika F. Ezeana, Son T. Ly, Stephen T. C. Wong, Hien V. Nguyen

TL;DR

This work tackles the challenge of limited annotated medical data for breast lesion analysis by introducing MIRAM, a self-supervised framework that performs masked image reconstruction across multiple scales. By reconstructing both original and higher-resolution mammogram scales and employing multiple decoders with token-interpolation, MIRAM learns finer-grained, robust representations capable of improving pathology classification and mass margins tasks on CBIS-DDSM. The approach outperforms state-of-the-art SSL methods, with gains in AP and AUC, and demonstrates scalable efficiency through linear self-attention options for the high-resolution decoder. The study further shows that using lesion-level annotated data during pre-training yields the best downstream performance, highlighting the importance of targeted SSL data curation in medical imaging. Overall, MIRAM offers a practical, scalable SSL pathway for enhancing breast lesion risk prediction from mammograms.

Abstract

Self-supervised learning (SSL) has garnered substantial interest within the machine learning and computer vision communities. Two prominent approaches in SSL include contrastive-based learning and self-distillation utilizing cropping augmentation. Lately, masked image modeling (MIM) has emerged as a more potent SSL technique, employing image inpainting as a pretext task. MIM creates a strong inductive bias toward meaningful spatial and semantic understanding. This has opened up new opportunities for SSL to contribute not only to classification tasks but also to more complex applications like object detection and image segmentation. Building upon this progress, our research paper introduces a scalable and practical SSL approach centered around more challenging pretext tasks that facilitate the acquisition of robust features. Specifically, we leverage multi-scale image reconstruction from randomly masked input images as the foundation for feature learning. Our hypothesis posits that reconstructing high-resolution images enables the model to attend to finer spatial details, particularly beneficial for discerning subtle intricacies within medical images. The proposed SSL features help improve classification performance on the Curated Breast Imaging Subset of Digital Database for Screening Mammography (CBIS-DDSM) dataset. In pathology classification, our method demonstrates a 3\% increase in average precision (AP) and a 1\% increase in the area under the receiver operating characteristic curve (AUC) when compared to state-of-the-art (SOTA) algorithms. Moreover, in mass margins classification, our approach achieves a 4\% increase in AP and a 2\% increase in AUC.

MIRAM: Masked Image Autoencoders Across Multiple Scales with Hybrid-Attention Mechanism for Breast Lesion Risk Prediction

TL;DR

This work tackles the challenge of limited annotated medical data for breast lesion analysis by introducing MIRAM, a self-supervised framework that performs masked image reconstruction across multiple scales. By reconstructing both original and higher-resolution mammogram scales and employing multiple decoders with token-interpolation, MIRAM learns finer-grained, robust representations capable of improving pathology classification and mass margins tasks on CBIS-DDSM. The approach outperforms state-of-the-art SSL methods, with gains in AP and AUC, and demonstrates scalable efficiency through linear self-attention options for the high-resolution decoder. The study further shows that using lesion-level annotated data during pre-training yields the best downstream performance, highlighting the importance of targeted SSL data curation in medical imaging. Overall, MIRAM offers a practical, scalable SSL pathway for enhancing breast lesion risk prediction from mammograms.

Abstract

Self-supervised learning (SSL) has garnered substantial interest within the machine learning and computer vision communities. Two prominent approaches in SSL include contrastive-based learning and self-distillation utilizing cropping augmentation. Lately, masked image modeling (MIM) has emerged as a more potent SSL technique, employing image inpainting as a pretext task. MIM creates a strong inductive bias toward meaningful spatial and semantic understanding. This has opened up new opportunities for SSL to contribute not only to classification tasks but also to more complex applications like object detection and image segmentation. Building upon this progress, our research paper introduces a scalable and practical SSL approach centered around more challenging pretext tasks that facilitate the acquisition of robust features. Specifically, we leverage multi-scale image reconstruction from randomly masked input images as the foundation for feature learning. Our hypothesis posits that reconstructing high-resolution images enables the model to attend to finer spatial details, particularly beneficial for discerning subtle intricacies within medical images. The proposed SSL features help improve classification performance on the Curated Breast Imaging Subset of Digital Database for Screening Mammography (CBIS-DDSM) dataset. In pathology classification, our method demonstrates a 3\% increase in average precision (AP) and a 1\% increase in the area under the receiver operating characteristic curve (AUC) when compared to state-of-the-art (SOTA) algorithms. Moreover, in mass margins classification, our approach achieves a 4\% increase in AP and a 2\% increase in AUC.

Paper Structure

This paper contains 19 sections, 7 equations, 3 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Methodology
  3. MIRAM
  4. General Idea
  5. Encoded Tokens Interpolation
  6. Linear self-attention for long-range sequences
  7. SS Pre-training Stage
  8. Datasets
  9. For pre-training stage
  10. For downstream classification tasks
  11. Experimental Results
  12. Experimental Setup and Evaluation Metrics
  13. Architecture Setup
  14. Training Setup
  15. Evaluation Metrics
  16. ...and 4 more sections

Figures (3)

  • Figure 1: Visualization of Learned Features on Pathology Classification Task. Our proposed MIRAM method shows improvements over the original MAE in terms of localizing the lesion's region and capturing finer details.
  • Figure 2: The MIRAM SSL Framework: Two different scales are reconstructed -- one for the original size, while the other is twice the size of both the original height and width.
  • Figure 3: Consistent pathology classification performance of proposed MIRAM-Nystromformer ($224\times{224}$) on CBIS-DDSM official test set over 100 fine-tuning epochs: Cyan curve is for the original MAE, Red curve is for our proposed MAE.