Multi-OCT-SelfNet: Integrating Self-Supervised Learning with Multi-Source Data Fusion for Enhanced Multi-Class Retinal Disease Classification

TL;DR

A self-supervised framework based on large language models, SwinV2, is developed to gain a deeper understanding of multi-modal dataset representations, enhancing the model's ability to extrapolate to new data for the detection of eye diseases using optical coherence tomography (OCT) images.

Abstract

In the medical domain, acquiring large datasets poses significant challenges due to privacy concerns. Nonetheless, the development of a robust deep-learning model for retinal disease diagnosis necessitates a substantial dataset for training. The capacity to generalize effectively on smaller datasets remains a persistent challenge. The scarcity of data presents a significant barrier to the practical implementation of scalable medical AI solutions. To address this issue, we've combined a wide range of data sources to improve performance and generalization to new data by giving it a deeper understanding of the data representation from multi-modal datasets and developed a self-supervised framework based on large language models (LLMs), SwinV2 to gain a deeper understanding of multi-modal dataset representations, enhancing the model's ability to extrapolate to new data for the detection of eye diseases using optical coherence tomography (OCT) images. We adopt a two-phase training methodology, self-supervised pre-training, and fine-tuning on a downstream supervised classifier. An ablation study conducted across three datasets employing various encoder backbones, without data fusion, with low data availability setting, and without self-supervised pre-training scenarios, highlights the robustness of our method. Our findings demonstrate consistent performance across these diverse conditions, showcasing superior generalization capabilities compared to the baseline model, ResNet-50.

Figures (10)

• Figure 1: A graphical depiction of our methodology, consolidating fundamental concepts and procedural steps.
• Figure 2: Overview of the Framework: (a) In the initial pre-training phase, the framework utilizes masked image autoencoder as a self-supervised task to learn representations from unlabeled images. In this process, a random subset of image patches is masked and fed into the auto-encoder to reconstruct it. (b) In this phase, the pre-trained encoder from the first phase is employed along with a linear classifier for the classification task. The learned weights from the pre-training phase are transferred to the fine-tuning phase.
• Figure 3: Illustration of the data combination process for the Self-Supervised Pre-training phase. Training and validation sets from three distinct Optical Coherence Tomography (OCT) datasets are merged to form a unified training and validation set, enhancing diversity and richness in the model's representation learning.
• Figure 4: Distribution of Retinal Disease Samples Across Three Datasets: Grouped-bar diagrams show sample counts in training, validation, and test sets for each retinal disease category in datasets DS1, DS2, and DS3. The Donut charts display the overall percentage distribution per dataset.
• Figure 5: The progression results of the Multi-OCT-SelfNet-SwinV2 model on sample validation images across various epochs illustrate its learning process in reconstructing input images. From left to right, (a) is the corresponding ground truth image, (b) is the masked image input, and (c)-(e) is the reconstructed images in different epochs.