VAMAE: Vessel-Aware Masked Autoencoders for OCT Angiography

1. Problem Statement

Optical Coherence Tomography Angiography (OCTA) is a non-invasive imaging modality crucial for diagnosing retinal diseases (e.g., diabetic retinopathy, glaucoma). However, learning robust representations from OCTA data presents unique challenges:

• Sparse Structures: Unlike natural images, OCTA scans consist of sparse, filamentary vessel structures (occupying only 10–15% of the image) against a uniform background.
• Topological Constraints: Clinical diagnosis relies on global vascular topology (branching patterns, connectivity, and the Foveal Avascular Zone boundaries) rather than just local texture.
• Failure of Standard SSL: Existing self-supervised learning (SSL) methods, particularly Masked Autoencoders (MAE), are designed for dense natural images. They use uniform random masking, which in OCTA leads to a trivial pretext task where the model reconstructs uninformative background pixels rather than critical vessel structures. This results in fragmented predictions and poor preservation of vascular connectivity.

2. Methodology: VAMAE Framework

The authors propose VAMAE (Vessel-Aware Masked Autoencoders), a geometry-informed self-supervised pretraining framework designed specifically for OCTA. The methodology consists of three core components:

A. Multi-Target Preprocessing Pipeline

Instead of relying solely on raw intensity, the pipeline decomposes each OCTA image into three complementary representations to provide rich supervision signals:

1. Intensity ($I$): The raw OCTA image.
2. Vesselness ($V$): Generated using multi-scale Frangi filtering to enhance tubular structures across different vessel calibers.
3. Skeleton ($S$): Extracted via morphological thinning to provide a 1-pixel-wide topological representation of vessel connectivity and branching.

B. Vessel-Aware Masking Strategy

To address the sparsity issue, VAMAE replaces uniform random masking with an adaptive, anatomy-guided strategy:

• Patch-wise Priors: For each image patch, the system calculates vessel density (from the vesselness map) and skeletal presence (from the skeleton map).
• Hybrid Importance Score: A score $w_i$ is computed for each patch:
  $$w_i = \alpha \cdot (0.5d_i + 0.5s_i) + (1 - \alpha) \cdot \epsilon_i$$
  Where $d_i$ is density, $s_i$ is skeletal presence, $\epsilon_i$ is random noise, and $\alpha=0.6$ is a weighting factor.
• Adaptive Selection: The top 75% of patches with the highest scores are selected for masking. This ensures the model is forced to reconstruct vessel-rich regions rather than empty background, while the stochastic term ($\epsilon_i$) prevents overfitting.

C. Architecture and Pretraining Objective

• Model: An asymmetric Vision Transformer (ViT) encoder-decoder. The encoder processes only visible patches; the decoder reconstructs all patches.
• Multi-Target Reconstruction: The decoder utilizes three task-specific MLP heads to simultaneously predict:
  1. Reconstructed Intensity ($\hat{I}$)
  2. Reconstructed Vesselness ($\hat{V}$)
  3. Reconstructed Skeleton ($\hat{S}$)
• Loss Function: A weighted sum of losses computed only on masked patches:
  $$L_{total} = 0.3 L_{intensity} + 0.5 L_{vesselness} + 0.2 L_{skeleton}$$
  The weights prioritize vesselness (structural prior) while balancing appearance and topology.

3. Key Contributions

1. VAMAE Framework: The first vessel-aware masked autoencoder specifically designed for self-supervised pretraining on OCTA data.
2. Geometry-Informed Masking: Introduction of a masking strategy that prioritizes vessel-rich regions using Frangi vesselness and skeleton-derived priors, solving the "trivial background" problem of standard MAE.
3. Multi-Target Learning: A novel objective that jointly reconstructs intensity, vesselness, and skeletons, enabling the model to capture appearance, local structure, and global topology simultaneously.
4. State-of-the-Art Performance: Demonstration that geometry-aware SSL can outperform fully supervised baselines and existing medical SSL methods, particularly in low-data regimes.

4. Experimental Results

The method was evaluated on the OCTA-500 benchmark for large vessel, vein, and Foveal Avascular Zone (FAZ) segmentation.

• Label Efficiency:
  • VAMAE achieved 82.4% Dice score with 100% labeled data, outperforming standard Random MAE (76.8%) and other medical SSL baselines (MedMAE, BioVessel-Net).
  • Crucially, with only 50% labeled data, VAMAE achieved 78.4% Dice, matching or exceeding the performance of other methods trained on 100% of the data. This effectively reduces annotation costs by 50%.
• Cross-Task Generalization: VAMAE showed consistent improvements across all tasks:
  • Large Vessels: +5.6% over Random MAE.
  • FAZ: +2.8% over Random MAE.
  • Veins: +5.1% over Random MAE (showing strong transfer to low-contrast structures).
• Ablation Studies:
  • Masking: Switching from random to vessel-aware masking alone improved performance by +4.8%.
  • Multi-Target: Adding vesselness and skeleton targets provided an additional +0.8% gain. Removing the vesselness target caused the largest performance drop (-3.2%), confirming its importance.
• Robustness: The model maintained significant gains across healthy controls and mild-to-moderate diabetic retinopathy cases, though performance degraded slightly in severe cases with disrupted vascular architecture.

5. Significance and Conclusion

• Domain Adaptation: The paper demonstrates that treating medical images as generic natural images is suboptimal. Incorporating domain-specific geometric priors (vesselness, topology) into the SSL pretext task is essential for learning transferable representations in OCTA.
• Clinical Impact: By achieving high accuracy with limited annotations, VAMAE addresses the bottleneck of expensive expert labeling in medical imaging, making advanced vessel segmentation more accessible.
• Theoretical Insight: The authors provide an information-theoretic justification, showing that vessel-aware masking increases the information density per masked patch, forcing the encoder to learn discriminative vascular features rather than memorizing background noise.

In summary, VAMAE successfully bridges the gap between general self-supervised learning and the specific geometric constraints of retinal vasculature, setting a new benchmark for OCTA analysis.