PaperlayerView exact revision on arXiv

Source-linked AI summary

Compression of Vision Foundation Models for Retinal OCT via Growing-Neural-Gas Codebooks and Graph-Guided Masked Image Modeling

Ali Yavari, Tilman Schmoll, Rainer A. Leitgeb, Wolfgang Drexler, Thomas Schlegl

arXiv:2610.33340v1cs.CV

TL;DR

Retinal vision foundation models are computationally demanding, and existing masked-image-modeling approaches do not transfer their semantic knowledge into lightweight backbones. The paper distills a frozen teacher into a graph-structured discrete codebook for student pretraining; across classification and segmentation benchmarks, the smaller student performs close to or sometimes above its teacher, depending on the teacher, task, metric, and capacity.

  • Problem

    Large, GPU-intensive foundation models hinder resource-constrained retinal imaging, while existing MIM approaches do not transfer their semantic knowledge into lightweight backbones and one-hot targets treat all incorrect codes equally.

    Source: §1 · pp. 2–3

  • Method

    A topology-aware distillation framework uses a frozen teacher’s discrete codebook as the reconstruction target for masked image modeling, learning the codebook with Growing Neural Gas.

    Source: §1 · p. 3

  • Results

    Across three classification and two segmentation benchmarks, the smaller student performed close to, and in some settings above, its teacher, with differences depending on teacher, task, metric, and capacity.

    Source: §1 · p. 3, §4 · p. 22

  • Takeaways & Limitations

    Three benchmark datasets were external to student pretraining and acquired at independent centers on two other devices, supporting cross-site and cross-device evaluation of the compact student.

    Source: §4 · pp. 22, 24

    Evaluation used retrospective public benchmarks and a single modality, OCT B-scans, and does not replace prospective clinical validation.

    Source: §4 · p. 24

Abstract

from arXiv · show

Vision foundation models for retinal imaging are large and computationally demanding, limiting point-of-care deployment. We ask how much of such a model's downstream utility survives compression into a smaller network. We propose a topology-aware discretization framework that distills a continuous foundation model into a discrete codebook, then uses it as the reconstruction target for masked image modeling (MIM). It combines gradient-free codebook learning via Growing Neural Gas, which yields an adjacency graph over the codes, with a graph-distance soft cross-entropy loss that replaces one-hot MIM targets with topologically graded supervision. With student pretraining on 383 public Duke iAMD optical coherence tomography (OCT) volumes, the pipeline distills an 86M-parameter teacher into a 22M-parameter student with at least 3.6x fewer FLOPs. The student performs close to, and in some settings above, its teacher across three classification and two segmentation benchmarks, with the magnitude and direction of the difference depending on the teacher, task, metric, and student capacity. On the external OLIVES benchmark, it surpasses its MIRAGE teacher on every metric (98.06% vs. 95.22% balanced accuracy). It works with two Vision Transformer (ViT-Base) teachers, retinal (MIRAGE) and natural-image (DINOv3); with the out-of-domain teacher, the student outperforms it on retinal-layer segmentation, indicating that the added value of retina-specific pretraining is task-dependent: largest for segmentation, while a general-domain teacher is competitive for classification. These results yield compact retinal models potentially suitable for deployment, transferring across OCT datasets, devices, and clinical tasks, including three datasets not used in the student's pretraining.

1. Introduction

Retinal OCT foundation models offer transferable representations, but their computational demands and the scarcity of large annotated or pooled datasets constrain practical use. The paper proposes topology-aware discrete distillation and reports that a much smaller student retains performance close to, and sometimes above, its teacher across five benchmarks.

1. Introduction

  • Large foundation models impose computational and infrastructure requirements that obstruct real-time use on portable devices or in resource-constrained clinical settings.Pretraining broadly generalizable models also requires large, heterogeneous datasets pooled across institutions and devices, which many centers cannot readily acquire and curate.Source: §1 · p. 2
  • Existing discrete-token MIM lacks relational structure among codes, while prior compact-ViT distillation work did not study graph-structured discrete targets on small, unlabeled in-domain medical cohorts.The paper identifies this as a gap between large natural-image corpora and lightweight medical backbones.Source: §1 · p. 3
  • The framework uses Growing Neural Gas to learn a discrete codebook and adjacency graph, then applies graph-distance soft cross-entropy to give masked prediction topologically graded supervision.It distills a frozen teacher into a compact student and is instantiated with retinal MIRAGE and natural-image DINOv3 teachers using Duke iAMD OCT data.Source: §1 · p. 3
  • The 22M-parameter student distilled from an 86M-parameter teacher performs close to, and on individual tasks above, its teacher across three classification and two segmentation benchmarks.The size reduction is roughly fourfold; the magnitude and direction of performance differences depend on teacher, task, metric, and student capacity.Source: §1 · p. 3
  • The same pipeline works with retinal and natural-image teachers, and the findings support compact retinal models as a route toward deployment without collecting the large multi-source corpora used for original pretraining.Performance differences across tasks and teachers remain variable.Source: §1 · p. 4

2. Materials and methods

The methods describe distilling frozen vision-foundation-model features into a graph-structured codebook and pretraining a compact student with graph-guided masked image modeling on unlabeled Duke iAMD OCT. They also detail task-specific downstream implementations and evaluation across classification and segmentation benchmarks, including cross-dataset tests, repeated-seed classification testing, and patient-level segmentation aggregation.

2.1. Framework overview

  • The two-stage framework learns a discrete tokenizer from a frozen teacher, uses its tokens and code graph as masked-image-modeling targets for a compact student, then evaluates the student on classification and segmentation.The tokenizer is used to generate training targets; the student backbone is retained for downstream tasks.Source: §2.1 · p. 4

2.2. Pretraining dataset

  • Self-supervised pretraining used 383 Duke iAMD volumes; because it used no labels, Duke evaluation was label-free transductive rather than label-leaking.One volume was excluded for an annotation inconsistency; expert retinal-boundary annotations defined downstream segmentation targets.Source: §2.2 · pp. 4, 6

2.3. Benchmark datasets

  • Five benchmarks across four public OCT datasets cover three classification and two segmentation tasks, with Kermany, OLIVES, and AROI reserved for cross-dataset transfer evaluation.Duke iAMD is evaluated transductively; classification uses the central B-scan for volumetric datasets, while segmentation uses annotated B-scans and patient-level splits.Source: §2.3 · p. 6

2.4. GNG-based image tokenizer

  • Growing Neural Gas (GNG) learns a codebook and adjacency graph alongside encoder training by updating codes from detached patch latents, without gradient-based codebook learning.GNG moves the winning code and its neighbors, links jointly winning codes, and inserts new codes in high-error regions until the 512-code budget is reached.Source: §2.4 · pp. 8–10

2.5. Masked image modeling with GNG graph guidance

  • The student’s masked-patch targets soften the one-hot code labels by GNG graph distance, penalizing confusion with nearby codes less while retaining a class-token aggregation branch.Soft-target support is limited to codes within three graph hops; a separate aggregation objective uses intermediate patch tokens to supervise the class-token representation.Source: §2.5 · pp. 10–12

2.6. Implementation details

  • Implementation used separate MIRAGE and DINOv3 teacher tokenizers and ViT-S and ViT-B students, with graph-soft-target parameters τ = 1.0 and R = 3 fixed across models and datasets.Tokenizer training used 512 codes in 256 dimensions; student MIM pretraining ran for 180 epochs on the same 383-volume corpus.Source: §2.6 · pp. 12–13
  • Segmentation used frozen ViT encoders with ConvNeXt decoders trained by cross-entropy plus Dice on patient-stratified Duke iAMD and AROI splits (269/38/76 and 17/2/4 patients).Duke iAMD segmentation used four classes; AROI used eight classes, including retinal layers and lesions.Source: §2.6 · p. 13
  • Classification used frozen-encoder linear probes with cross-entropy and 0.1 label smoothing; Kermany used its author-provided split, while Duke iAMD and OLIVES used stratified patient-level splits.Duke iAMD split 267/39/77 subjects; OLIVES split 60/9/18 patients (1,134/176/284 volumes).Source: §2.6 · pp. 13–14

2.7. Evaluation metrics

  • Classification used weighted one-vs-rest AUROC, weighted average precision, and balanced accuracy across five probe seeds; graph ablation used paired tests with Benjamini–Hochberg correction.The ablation significance tests were conditional on each encoder’s single pretraining run and did not capture pretraining stochasticity.Source: §2.7 · p. 14
  • Segmentation metrics were foreground-class Dice, IoU, and HD95, scored per B-scan then averaged per patient; background was excluded, and HD95 is reported in in-plane pixels.Patient-level means and standard deviations were reported, with layer and lesion scores averaged separately where applicable.Source: §2.7 · p. 14

3. Results

Across OCT classification and segmentation, compressed GNG-ViT students often matched or exceeded their teachers, with outcomes varying by teacher, task, and metric. Ablations distinguish tokenizer-only features from stage-2 pretraining, assess graph-guided targets, and quantify inference costs.

3.1. Ocular disease classification

  • GNG-ViT-Small-MIRAGE exceeded MIRAGE in 7 of 9 classification comparisons, including OLIVES BAcc 95.22 → 98.06; Base-MIRAGE improved on MIRAGE in all 9 comparisons.Together, MIRAGE-distilled students improved on their teacher in 16 of 18 comparisons overall.Source: §3.1 · p. 15, Table 2 · p. 16

3.2. Segmentation of retinal layers and lesions

  • On AROI retinal layers, GNG-ViT-Small-DINOv3 raised Dice from 83.76 to 89.03 and IoU from 75.83 to 82.58; on Duke iAMD, Dice rose from 85.34 to 88.34.DINOv3 HD95 fell from 8.02 ± 6.17 to 1.95 ± 0.63 on AROI layers and from 19.16 ± 9.07 to 8.40 ± 5.01 on Duke iAMD.Source: §3.2 · p. 15, Table 3 · p. 18
  • GNG-ViT-Small-MIRAGE stayed within 0.49 Dice points on AROI layers and 0.64 on Duke iAMD layers, while MIRAGE-distilled students surpassed MIRAGE on AROI lesion Dice and IoU.Small-MIRAGE lesion Dice and IoU were 48.75 and 40.28, versus 46.29 and 38.66 for MIRAGE; Base-MIRAGE achieved the best lesion Dice and IoU, 51.13 and 42.16.Source: §3.2 · p. 16, Table 3 · p. 18

3.3. Ablation study

  • Tokenizer-only features ranked differently by task: DINOv3-derived features led on Kermany classification, while MIRAGE-derived features led on every AROI segmentation metric.Both tokenizer encoders lagged the corresponding full pipelines, suggesting most reported downstream gains arose from stage-2 MIM pretraining.Source: §3.3 · p. 18, Table 4 · p. 20, Table 5 · p. 20
  • Graph-guided targets significantly improved 10 of 18 classification comparisons and worsened 2; Base-MIRAGE Kermany BAcc rose from 73.90 to 89.98 (+16.1 points).The gain remained significant after BH-FDR correction, while Small-model AUROC and AP on OLIVES dipped significantly versus OneHot-Small.Source: §3.3 · p. 19, Table 6 · p. 21
  • For segmentation, graph targets were associated with modest overlap gains and lower boundary error; Small-MIRAGE AROI-layer HD95 fell from 3.61 to 1.79 versus OneHot.AROI lesion metrics showed high patient-level variability and no consistent ordering between variants.Source: §3.3 · p. 19, Table 7 · p. 21
  • The selected codebook used 512 codes and 256-dimensional vectors, while soft-target settings were τ = 1.0 and R = 3.At τ = 1, R = 2 and R = 3 yielded nearly identical Kermany BAcc, 84.25 and 84.30, respectively.Source: §3.3 · p. 20

3.4. Inference cost

  • ViT-Small students required 6.37 encoder GFLOPs per image, 3.6× fewer than MIRAGE-Base and 16.9× fewer than DINOv3-Base; Base students were not compressed.Classification latency fell to 1.66 ms, while segmentation cost fell less because the shared decoder dominates.Source: §3.4 · p. 20, Table 8 · p. 22

4. Discussion and conclusion

The framework retains much of teacher performance in substantially smaller retinal OCT students, while the benefit of retinal-specific pretraining and graph-guided targets varies by task and student capacity. Its evaluation spans external datasets and devices, but remains bounded by retrospective public benchmarks and untested methodological alternatives.

4. Discussion and conclusion

  • The 22M-parameter MIRAGE student stayed within one Dice point of its teacher on retinal layers and reached 98.06 balanced accuracy versus 95.22 for its teacher on OLIVES.Across three classification and two segmentation benchmarks, student performance was close to or sometimes above the teacher, with differences depending on teacher, task, metric, and capacity.Source: §4 · p. 22
  • Graph-distance soft targets improved several outcomes, particularly for Base students, and lowered retinal-layer boundary error, but effects varied by task and capacity.Some Small-model settings showed neutral or adverse effects; GNG supplies the adjacency graph used to define the graph-distance targets.Source: §4 · pp. 22–24
  • DINOv3 teachers led on Kermany and Duke iAMD classification, while MIRAGE led on OLIVES balanced accuracy and all segmentation benchmarks.DINOv3-distilled students substantially outperformed their teacher on retinal-layer segmentation; the DINOv3 tokenizer-only encoder reached 84.15 Dice on AROI layers, while its distilled student reached 89.03 versus 90.30 for the MIRAGE teacher.Source: §4 · pp. 22–23
  • The ViT-Small student uses 22M parameters versus 86M for its teacher and at least 3.6× fewer encoder FLOPs, but desktop-GPU measurements do not establish deployment readiness.Student pretraining used a single public cohort of 383 volumes; the authors note this may lower the barrier for smaller centers seeking a task-specific backbone.Source: §4 · p. 24
  • Kermany, OLIVES, and AROI were external to student pretraining and came from independent centers on two other devices, but Kermany informed hyperparameter selection.OLIVES and AROI were the only benchmarks untouched by both pretraining and model selection.Source: §4 · p. 24
  • Evaluation is limited to retrospective public OCT B-scan benchmarks, without prospective validation or demographic-stratified analysis.The study also did not compare conventional VQ tokenizers in the full pipeline or alternative graph topologies, and teacher-agnosticism was tested only with two ViT-Base teachers at D=768.Source: §4 · pp. 24–25
  • The authors propose testing broader multi-device, multi-disease, and multi-institution OCT data, larger teachers and code budgets, and links between codes and retinal anatomy or pathology.They also suggest the approach could extend beyond OCT B-scans as a recipe for compressing medical vision foundation models.Source: §4 · p. 25

CRediT authorship contribution statement

The authors divide conceptualization, analysis, investigation, methodology, software, visualization, supervision, project administration, and writing contributions across the team.

CRediT authorship contribution statement

  • Ali Yavari contributed conceptualization, formal analysis, investigation, methodology, software, visualization, and writing.He is credited with the original draft and review and editing.Source: CRediT authorship contribution statement · p. 25
  • Tilman Schmoll contributed project administration, supervision, and writing; Rainer A. Leitgeb contributed formal analysis, methodology, supervision, visualization, and writing.Both are credited with writing review and editing.Source: CRediT authorship contribution statement · p. 25
  • Wolfgang Drexler contributed supervision and writing, while Thomas Schlegl contributed conceptualization, formal analysis, methodology, visualization, and writing.Both are credited with writing review and editing.Source: CRediT authorship contribution statement · p. 25

Declaration of competing interest

The competing-interest declaration reports Carl Zeiss Meditec affiliations and consulting relationships for three authors.

Declaration of competing interest

  • Tilman Schmoll is a Carl Zeiss Meditec employee, and Wolfgang Drexler and Rainer A. Leitgeb are company consultants who received financial compensation.The other authors declare no conflicts of interest.Source: Declaration of competing interest · p. 25

Data availability

The study used publicly available pretraining and benchmark datasets, with no in-house or private data.

Data availability

  • Duke iAMD, Kermany, OLIVES, and AROI are publicly available datasets, and the study used no in-house or private data.Duke iAMD supplied pretraining data and one downstream benchmark; Kermany, OLIVES, and AROI supplied the other downstream benchmarks.Source: Data availability · p. 26

Funding

Carl Zeiss Meditec AG supported the work but had no role in its conduct or publication decisions.

  • Carl Zeiss Meditec AG supported the work.Source: Funding · p. 26
  • Tilman Schmoll is employed by the company, and Wolfgang Drexler and Rainer A. Leitgeb serve as its consultants.Their contributions are detailed in the CRediT authorship contribution statement.Source: Funding · p. 26
  • The company had no involvement in study design, data collection, analysis or interpretation, manuscript preparation, or the decision to submit.Source: Funding · p. 26

Related papers

Loading 2610.33340v1…