Distributed training of foundation models for ophthalmic diagnosis

Abstract

Vision impairment affects nearly 2.2 billion people globally, and nearly half of these cases could be prevented with early diagnosis and intervention-underscoring the urgent need for reliable and scalable detection methods for conditions like diabetic retinopathy and age-related macular degeneration. Here we propose a distributed deep learning framework that integrates self-supervised and domain-adaptive federated learning to enhance the detection of eye diseases from optical coherence tomography images. We employed a self-supervised, mask-based pre-training strategy to develop a robust foundation encoder. This encoder was trained on seven optical coherence tomography datasets, and we compared its performance under local, centralized, and federated learning settings. Our results show that self-supervised methods-both centralized and federated-improved the area under the curve by at least 10% compared to local models. Additionally, incorporating domain adaptation into the federated learning framework further boosted performance and generalization across different populations and imaging conditions. This approach supports collaborative model development without data sharing, providing a scalable, privacy-preserving solution for effective retinal disease screening and diagnosis in diverse clinical settings.

Fig. 1. Overview of the four phases of our framework.

a Local learning phase in which a baseline model is trained in a particular dataset and evaluated over test set(s). b Centralized learning approach, comprising pre-training, fine-tuning, and evaluation. c Federated learning (FDL) approach where the pre-training phase is conducted via FDL. d domain adaptation (DAD)-FDL pipeline, where the DAD configuration is distributed before pre-training.

Fig. 2. DL pipelines of the local learning and pre-training.

a Local learning pipeline, in which a pair of images and their label are input to the model after undergoing four transformations: rotation, color jittering, Gaussian blur, and Sobel filter. b Pre-training phase, during which the input image is masked and given to the reconstruction network to train the encoder over time.

Fig. 3. Reconstructed images from the MIM network, where 50% of the images are masked.

a₁–f₁ Columns of original image samples from DS1 to DS6, respectively. a₂–f₁ Column of masked images. a₃–f₃ Columns of reconstructed images.

Fig. 4. Main stages of pre-training and fine-tuning via federated learning (FDL) at the University of North Carolina at Charlotte (UNCC) and the University of Illinois Chicago (UIC).

a Server shares the initial model parameter and the configuration to all nodes. b Each FDL node pre-trains its model and sends the model’s weights to the server. c Finally, the server aggregates the weights and reciprocates them to each client, and they start the fine-tuning step.

Fig. 5. Samples from DS1 to DS7 and their unsupervised noise-transformed version, where all the transformed images have similar pixel value intensity, are shown in rows d to i in alphabetical order.

a₁–g₁ Columns of original images from DS1 to DS7, respectively. a₂–g₂ Columns of transformed images. h Target image used to transform other images based on its pixel value intensity.

Fig. 6. Macro AUC-ROC plot of four models over DS7.

Local, centralized, DAD-FDL-1, and DAD-FDL-5, with DAD-FDL-5 outperforming other methods on DS7.

Fig. 7. Eigen Grad-CAM inference from the norm layer of the final Swin Transformer BlockV2.

Models’ inference samples (choroidal neovasculature (CNV), diabetic macular edema (DME), diabetic retinopathy (DR), Drusen, normal and age-related macular degeneration (AMD)) from DS1 to DS7 datasets, respectively. Checkmarks and crosses indicate correct and incorrect predictions.