CVAE-based Causal Representation Learning from Retinal Fundus Images for Age Related Macular Degeneration(AMD) Prediction

This paper proposes a novel CVAE-based causal representation learning framework that utilizes retinal fundus images to disentangle underlying Age-Related Macular Degeneration mechanisms, thereby enhancing the accuracy and reliability of AMD prediction and diagnosis.

The Gist

🩺 The Big Picture: Teaching AI to Understand the "Why," Not Just the "What"

Imagine you are trying to teach a robot to spot a specific type of tree disease.

• Old AI (The Traditional Approach): You show the robot thousands of photos of sick trees. It learns to say, "Ah, that brown spot means it's sick!" It gets really good at guessing, but it doesn't actually understand why the tree is sick. It just memorizes patterns.
• This New AI (The Causal Approach): This paper teaches the robot to understand the story behind the disease. It learns: "The brown spot is caused by a lack of water, which is caused by a blocked root."

The Goal: The researchers wanted to build an AI that doesn't just predict if a patient has Age-Related Macular Degeneration (AMD) (a disease that causes blindness), but actually understands the biological causes of the disease so it can help doctors figure out the best treatment.

---

🧩 The Problem: The "Black Box"

AMD is a sneaky disease. It damages the retina (the "film" at the back of your eye). There are two main types:

1. Dry AMD: Like a slow leak. Waste builds up under the eye, starving the cells.
2. Wet AMD: Like a pipe bursting. New, weak blood vessels grow and leak fluid, causing rapid damage.

Doctors look at eye photos (fundus images) to diagnose this. But the photos are messy. Sometimes it's hard to tell if a dark spot is just a shadow or a serious bleed. Traditional AI looks at the whole picture and guesses, "That looks like AMD." But if the doctor asks, "Why?" the AI can't answer. It's a Black Box.

---

🛠️ The Solution: The "Magic Decoder Ring"

The authors built a special AI system with two main parts working together. Think of it as a Translator and a Detective.

1. The Translator (The Convolutional VAE)

Imagine the eye photo is written in a complex, messy foreign language.

• The VAE (Variational Autoencoder) is a translator. It looks at the messy photo and compresses it into a short, clean summary code (called a "Latent Representation").
• Instead of keeping the whole photo, it extracts the essence: "Here is a dark spot," "Here is a bright cluster," "Here is a fluid leak."
• It then tries to rebuild the photo from that code. If it can rebuild the photo perfectly, it knows it captured all the important details.

2. The Detective (The Graph Autoencoder)

Now that the AI has the "essence" of the disease, it needs to figure out how the pieces connect.

• The Graph Autoencoder (GAE) acts like a detective drawing a family tree.
• It asks: "Does the 'Dark Spot' cause the 'Fluid Leak'? Or does the 'Fluid Leak' cause the 'Dark Spot'?"
• It builds a Causal Map (a Directed Acyclic Graph). It learns that Factor A leads to Factor B, which leads to AMD.

The Magic Trick: By combining these two, the AI learns to separate the different causes of the disease. It learns that "Drusen" (waste buildup) is one specific variable, and "Hemorrhage" (bleeding) is another. It can even change one variable in its mind to see what happens to the image.

---

🧪 The Experiment: Playing with the "Dials"

To prove their AI was smart, the researchers did a cool test. They took the "code" the AI learned and turned the dials up and down.

• The Test: They took a photo of a healthy eye and told the AI, "Turn up the 'Drusen' dial."
• The Result: The AI generated a new image where bright, yellowish spots (drusen) appeared on the retina.
• The Second Test: They turned up the "Fluid/Hemorrhage" dial.
• The Result: The AI generated an image with dark, cloudy patches (bleeding).

Why this matters: This proves the AI didn't just memorize the picture; it learned the mechanism. It understands that if you increase the "Drusen" variable, the image changes in a specific, medically accurate way.

---

🏆 The Results: Better Diagnosis and Future Treatments

1. It's a Great Doctor:
When they used this "causal code" to predict who has AMD, the AI was incredibly accurate (over 92% accuracy). It was better than standard AI models because it focused on the real causes rather than just random patterns in the pixels.

2. It's a Time Machine for Treatment:
This is the most exciting part. Because the AI understands the causes, doctors could theoretically use it to simulate treatments.

• Imagine: A doctor says, "If we give this patient a drug that stops the bleeding, what will their eye look like in 6 months?"
• The AI can take the patient's current "code," turn down the "bleeding" dial, and generate a picture of what their eye would look like after the treatment works.

🚧 The Limitations (The "But...")

The authors are honest about what they couldn't do yet:

• Resolution: The AI is good at seeing big blobs (bleeds, waste), but it struggles to draw tiny blood vessels perfectly. It's like a painter who is great at landscapes but bad at painting individual leaves.
• Data: They needed a lot of data, and sometimes the "labels" (what the disease actually is) aren't perfect.

💡 The Takeaway

This paper is a bridge between Artificial Intelligence and Medical Science.

• Old AI: "I see a disease."
• New AI (This Paper): "I see a disease, I know why it's there, I can show you what it looks like if we fix the cause, and I can predict if you will get sick."

It's a step toward AI that doesn't just act like a calculator, but acts like a thinking partner for doctors.

Technical Summary

1. Problem Statement

Age-Related Macular Degeneration (AMD) is a leading cause of vision impairment, particularly the "Wet" subtype, which involves rapid retinal damage due to neovascularization. While deep learning (DL) models have improved AMD detection, they primarily rely on correlation rather than causality. Current models often act as "black boxes," failing to:

• Understand the underlying causal mechanisms of AMD (e.g., the progression from drusen to neovascularization).
• Provide reliable explanations for interventions (treatments).
• Distinguish between normal and AMD fundus images with high robustness, especially given the complexity of retinal pathologies and class imbalance in medical datasets.

The research aims to bridge this gap by developing a framework that not only predicts AMD status but also learns latent causal representations to explain why a diagnosis is made and to simulate intervention effects.

2. Methodology

The authors propose a novel framework combining Convolutional Variational Autoencoders (CVAE) with Graph Autoencoders (GAE) to perform explicit latent causal representation learning.

A. Data Source

• Dataset: Retinal Fundus Multi-Disease Image (RFMiD) dataset.
• Preprocessing: Images were resized to $176 \times 176 \times 3$.
• Class Balancing: Due to extreme imbalance (Normal:AMD $\approx$ 1820:100), random under-sampling was applied to create a 200:100 training ratio.

B. Model Architecture

The framework consists of two main components trained simultaneously:

1. Convolutional VAE (CVAE):

   • Encoder: Three convolutional layers (64, 64, 16 filters) followed by dense layers to map input images ($X$) to a latent space ($Z$).
   • Decoder: Symmetrical convolutional layers to reconstruct the input image.
   • Objective: Minimize reconstruction error (MSE) and KL-divergence (ELBO loss) to ensure the latent space captures essential image features.

2. Graph Autoencoder (GAE) for Causal Discovery:

   • Input: The latent vector $Z$ concatenated with the AMD label $Y$ (forming $W = \langle Z, Y \rangle$).
   • Mechanism: Uses a Neural Causal DAG (Directed Acyclic Graph) approach based on Structural Causal Models (SCM).
   • Structure: Implements a Graph Autoencoder where an adjacency matrix $A$ (representing causal edges) is learned as trainable weights between encoder and decoder MLPs.
   • Constraint: Enforces the DAG constraint (acyclicity) using the NOTEARS-style continuous optimization constraint: $tr(e^{A \odot A}) = d$.
   • Objective: Minimize reconstruction error of $W$ while learning the causal structure $A$.

C. Optimization

• Total Loss: $\mathcal{L}_{total} = \mathcal{L}_{VAE} + \mathcal{L}_{GAE}$.
• Training: Simultaneous gradient descent using Adam optimizers.
• Hyperparameters: Latent dimension $|Z|=8$. The model was trained for 850 epochs.
• Causal Graph Extraction: The learned adjacency matrix $A$ was binarized based on weight magnitude thresholds to extract the final causal graph.

D. Downstream Tasks

• Causal Disentanglement: Validated by modifying individual latent variables ($z_i$) and observing changes in reconstructed images (e.g., increasing $z_4$ to simulate drusen).
• AMD Prediction: The extracted latent vectors $Z$ were used as input features for standard ML/DL classifiers (Random Forest, Gradient Boosting, Extra Trees, and a Deep Neural Network) to predict AMD status.

3. Key Contributions

1. Explicit Latent Causal Modeling: Unlike previous works that focus on feature correlation, this study explicitly models the causal DAG of AMD mechanisms directly from fundus images using a CVAE-GAE hybrid.
2. Causal Disentanglement: The model successfully separated distinct AMD factors into specific latent dimensions. Specifically:
   • $z_0$ corresponded to hemorrhage/intraretinal fluids.
   • $z_4$ corresponded to drusen.
   • The model identified a potential causal link between drusen and neovascularization, aligning with emerging clinical literature.
3. Intervention Simulation: The framework allows for "what-if" analysis. By altering latent variables (e.g., simulating the removal of fluid), the model can generate synthetic retinal images showing the potential outcome of treatments, aiding clinical decision-making.
4. Robust Prediction: Demonstrated that using causal latent representations as features yields superior prediction performance compared to raw image inputs in certain metrics, particularly in handling class imbalance.

4. Experimental Results

A. Model Convergence & Reconstruction

• The CVAE and GAE losses converged stably after 850 epochs ($\mathcal{L}_{VAE} \approx 0.0027$, $\mathcal{L}_{GAE} \approx 0.4348$).
• Visual Quality: Reconstructed images accurately captured pathological features like hypo-reflective fluids (hemorrhage) and hyper-reflective clusters (drusen/exudates).

B. Causal Graph Validation

• Structural Hamming Distance (SHD): The extracted causal subgraph had an SHD of 2.0 when compared to the domain-knowledge benchmark graph, indicating high structural accuracy.
• Disentanglement:
  • Modifying $z_4$ (drusen) specifically altered hyper-reflective areas without affecting other regions.
  • Modifying $z_0$ (fluids) specifically altered dark, hypo-reflective regions.
  • Statistical tests (Wilcoxon rank sum) confirmed $z_4$ was significantly correlated with drusen diagnosis ($p < 0.05$).

C. AMD Prediction Performance

Using the latent causal features ($Z$) as input for downstream classifiers:

• Best Performer: The Deep Neural Network (DNN) achieved the highest balanced performance.
  • Accuracy: 92.1%
  • Weighted F1-Score: 91.9%
  • Specificity: 97.0%
  • ROC AUC: 87.10%
• Comparison: While Random Forest achieved a higher ROC AUC (91.72%), it suffered from lower precision/recall balance compared to the DNN. The DNN demonstrated the most robustness against false positives in the imbalanced test set.

5. Significance and Future Work

• Clinical Impact: This research moves beyond simple classification to mechanistic understanding. By identifying causal factors, the model offers a pathway for explainable AI (XAI) in ophthalmology, allowing clinicians to visualize how specific treatments (interventions) might alter retinal pathology.
• Limitations:
  • Resolution: The model struggled to reconstruct fine details like specific blood vessels, likely due to the low resolution of the input fundus images and the inherent limitations of VAEs.
  • Quantitative Metrics: Causal disentanglement was primarily validated qualitatively due to the lack of pixel-level ground truth labels for specific pathologies (e.g., exact drusen boundaries).
• Future Directions: The authors suggest integrating Diffusion Models for higher-fidelity image reconstruction and acquiring fully labeled datasets to enable quantitative disentanglement scoring.

In conclusion, this paper presents a significant step toward Causal AI in medical imaging, demonstrating that latent causal representation learning can simultaneously improve disease prediction accuracy and provide interpretable, intervention-ready insights into AMD pathogenesis.