KD-OCT: Efficient Knowledge Distillation for Clinical-Grade Retinal OCT Classification

Summary

This study introduces KD-OCT, a novel knowledge distillation framework designed to compress computationally intensive deep learning models for retinal OCT classification into efficient, deployable models. By distilling a powerful ConvNeXtV2-Large teacher into a lightweight EfficientNet-B2 student, KD-OCT achieves near-teacher diagnostic performance for classifying normal, drusen, and CNV cases while significantly reducing model size and inference time. This enables real-time, high-accuracy AMD screening suitable for edge deployment.

Medical Relevance

This research is crucial for making advanced AI diagnostics accessible and practical in ophthalmology, enabling early, efficient, and widespread screening for vision-threatening conditions like Age-related Macular Degeneration (AMD) and Choroidal Neovascularization (CNV) directly at the point of care.

AI Health Application

The AI application is a deep learning-based diagnostic assistant that classifies retinal OCT images to detect and screen for age-related macular degeneration (AMD) and choroidal neovascularization (CNV). The specific innovation is making these models computationally efficient for real-time and 'edge deployment' in clinical settings, thereby improving accessibility and speed of diagnosis for vision-threatening conditions.

Key Points

  • Addresses the computational challenges of deploying state-of-the-art deep learning models (e.g., ConvNeXtV2-Large) for retinal OCT classification in clinical settings.
  • Proposes KD-OCT, a novel knowledge distillation framework, to compress a high-performance teacher model into a lightweight student model.
  • The teacher model is a ConvNeXtV2-Large, enhanced with advanced augmentations, stochastic weight averaging (SWA), and focal loss for robust training.
  • The student model is a more efficient EfficientNet-B2, chosen for its balance of efficiency and capability.
  • KD-OCT employs real-time distillation using a combined loss function that balances soft knowledge transfer from the teacher with hard ground-truth supervision.
  • Evaluated on the Noor Eye Hospital (NEH) dataset using patient-level cross-validation, ensuring robust generalization.
  • Achieves near-teacher performance with substantial reductions in model size and inference time, outperforming comparable multi-scale or feature-fusion OCT classifiers in efficiency-accuracy balance.
  • The compressed student model surpasses most existing frameworks, facilitating its deployment on edge devices for AMD and CNV screening.

Methodology

The study employs a knowledge distillation approach, termed KD-OCT. A powerful ConvNeXtV2-Large model, pre-trained with advanced augmentations, stochastic weight averaging, and focal loss, serves as the teacher. A lightweight EfficientNet-B2 model is trained as the student. Distillation occurs in real-time, utilizing a combined loss function that balances soft targets (probabilities) from the teacher with hard labels (ground truth) for supervision. The framework was evaluated using patient-level cross-validation on the Noor Eye Hospital (NEH) dataset for classifying normal, drusen, and CNV cases.

Key Findings

KD-OCT successfully compressed a high-performance ConvNeXtV2-Large teacher into an EfficientNet-B2 student, achieving diagnostic performance nearly equivalent to the teacher model. This compression resulted in substantial reductions in model size and inference time, making the model significantly more efficient. The KD-OCT student model demonstrated superior efficiency-accuracy balance compared to other multi-scale or feature-fusion OCT classifiers and exceeded most existing frameworks in overall performance despite its compact size.

Clinical Impact

The development of KD-OCT enables the deployment of high-accuracy, real-time retinal OCT classification models directly on edge devices in clinical settings. This facilitates efficient and widespread screening for AMD and CNV, potentially improving early detection rates, streamlining patient management, and increasing access to advanced diagnostic capabilities, particularly in resource-constrained environments.

Limitations

The abstract does not explicitly state limitations of the KD-OCT method itself, but it highlights the computational demands of state-of-the-art models as the problem being addressed. Specific limitations regarding the generalizability to diverse populations or different OCT devices are not discussed within the abstract.

Future Directions

While not explicitly stated as future directions, the paper's emphasis on enabling 'edge deployment for AMD screening' implies continued efforts towards broader clinical implementation, real-world validation, and potentially extending the framework to classify a wider range of retinal pathologies or integrate into a complete clinical workflow.

Medical Domains

Ophthalmology Retinal Imaging Macular Degeneration Diagnostic Imaging

Keywords

Knowledge Distillation Retinal OCT Age-related Macular Degeneration (AMD) Choroidal Neovascularization (CNV) Deep Learning EfficientNet ConvNeXtV2 Edge Computing Clinical AI

Abstract

Age-related macular degeneration (AMD) and choroidal neovascularization (CNV)-related conditions are leading causes of vision loss worldwide, with optical coherence tomography (OCT) serving as a cornerstone for early detection and management. However, deploying state-of-the-art deep learning models like ConvNeXtV2-Large in clinical settings is hindered by their computational demands. Therefore, it is desirable to develop efficient models that maintain high diagnostic performance while enabling real-time deployment. In this study, a novel knowledge distillation framework, termed KD-OCT, is proposed to compress a high-performance ConvNeXtV2-Large teacher model, enhanced with advanced augmentations, stochastic weight averaging, and focal loss, into a lightweight EfficientNet-B2 student for classifying normal, drusen, and CNV cases. KD-OCT employs real-time distillation with a combined loss balancing soft teacher knowledge transfer and hard ground-truth supervision. The effectiveness of the proposed method is evaluated on the Noor Eye Hospital (NEH) dataset using patient-level cross-validation. Experimental results demonstrate that KD-OCT outperforms comparable multi-scale or feature-fusion OCT classifiers in efficiency- accuracy balance, achieving near-teacher performance with substantial reductions in model size and inference time. Despite the compression, the student model exceeds most existing frameworks, facilitating edge deployment for AMD screening. Code is available at https://github.com/erfan-nourbakhsh/KD- OCT.

Comments

7 pages, 5 figures