Exploiting Intermediate Reconstructions in Optical Coherence Tomography for Test-Time Adaption of Medical Image Segmentation

This paper proposes IRTTA, a test-time adaptation method that leverages informative intermediate reconstructions from iterative Optical Coherence Tomography algorithms to dynamically adjust a frozen segmentation network's normalization parameters, thereby improving performance and providing uncertainty estimates without altering the reconstruction or downstream models.

The Gist

Imagine you are trying to identify a specific landmark in a foggy city. You have a map (a medical image), but it's blurry and distorted because you took the photo with a cheap camera instead of a professional one.

Usually, doctors and AI models look at the final, clearest version of that photo to make a diagnosis. But this paper suggests a clever trick: don't just look at the final photo; watch the whole process of how the photo gets cleared up.

Here is a simple breakdown of the paper's idea, using everyday analogies:

1. The Problem: The "Cheap Camera" vs. The "Pro Camera"

Medical centers often use expensive, high-quality machines (like a professional DSLR) to get crystal-clear images of the inside of your eye. However, many clinics use cheaper, lower-quality machines (like a smartphone camera).

• The Issue: AI models are trained on the "Pro Camera" images. When they look at "Smartphone Camera" images, they get confused and make mistakes because the pictures look different (noisier, different colors).
• The Current Fix: Scientists use special algorithms to "clean up" the cheap images, turning them into something that looks like the high-quality ones. This is like using a photo-editing app to remove the blur.

2. The Hidden Treasure: The "Reconstruction Trajectory"

Here is the twist: The process of cleaning up the image isn't instant. It happens in steps, like a time-lapse video.

• Step 1: The image is very blurry and noisy.
• Step 50: It's getting clearer, but still a bit fuzzy.
• Step 100: It looks perfect.

Most people throw away Steps 1 through 99 and only use Step 100. The authors of this paper say, "Wait a minute! Those blurry steps in the middle actually hold a lot of useful information!"

3. The Solution: IRTTA (The "Smart Guide")

The authors created a new method called IRTTA. Think of it like this:

Imagine you are an AI trying to find a specific building (a tumor or fluid) in that city.

• The Old Way: You look at the final, perfect photo and try to guess where the building is. If the photo is slightly different from what you learned in school, you might get lost.
• The IRTTA Way: You have a Smart Guide (a small helper AI) who watches the photo being cleaned up step-by-step.
  • At Step 1 (very blurry), the Guide says, "Hey, it's very foggy right now, so don't trust the colors too much."
  • At Step 50 (getting clearer), the Guide says, "Okay, the shapes are forming, but the edges are still soft."
  • At Step 100 (perfect), the Guide says, "Now we can be very confident."

The Guide doesn't change the main AI's brain (the part that knows what a building looks like). Instead, it just tweaks the AI's "glasses" (specifically, how it normalizes the image) to match the current level of clarity.

4. The "No-Label" Superpower

Usually, to teach an AI to adapt to a new camera, you need a teacher to show it the correct answers (labeled data). But in a real clinic, you don't have those answers for every new patient.

• How IRTTA learns without a teacher: It uses a concept called Entropy.
  • Imagine the AI is guessing where the building is. If it's totally confused, its guesses are all over the place (high entropy).
  • If it's confident, its guesses are all the same (low entropy).
  • The system automatically adjusts the "glasses" until the AI is most confident across all the steps of the cleaning process. It's like tuning a radio until the static disappears and the music is clear.

5. The Bonus: A "Confidence Meter"

Because the system looks at the image at every step of the cleaning process, it can tell you how sure it is.

• If the AI sees the building clearly at Step 10, Step 50, and Step 100, it's very confident.
• If the AI sees the building at Step 10, but gets confused at Step 50, and then sees it again at Step 100, the system flags that area as "Uncertain."

This gives doctors a "heat map" showing exactly where the AI is unsure, which is incredibly valuable for safety.

Summary

• The Goal: Make cheap medical images work as well as expensive ones for AI diagnosis.
• The Trick: Instead of ignoring the "blurry middle steps" of image cleaning, use them to help the AI adjust its vision in real-time.
• The Result: Better accuracy in finding diseases (like fluid in the eye) and a built-in "confidence meter" that tells doctors when to double-check the AI's work.

It's like realizing that the journey of cleaning a dirty window is just as important as the clean window itself, and using that journey to help your eyes see better.

Technical Summary

1. Problem Statement

Medical image segmentation models, particularly those based on deep learning, are typically trained on high-fidelity datasets from university hospitals. However, in primary care, low-cost imaging devices (e.g., Optical Coherence Tomography or OCT scanners) are often used. These devices produce images with lower signal-to-noise ratios (SNR) and different intensity distributions, causing a domain shift that degrades segmentation performance.

While advanced iterative reconstruction algorithms (such as diffusion models) are used to enhance image quality from low-cost devices, standard evaluation protocols rely solely on the final reconstructed image. This approach ignores the rich semantic information contained in the intermediate reconstruction steps (the trajectory) generated during the iterative process. Furthermore, existing methods that try to leverage reconstruction often require joint training of the reconstruction and segmentation networks, which is computationally expensive and requires labeled data.

The Core Challenge: How to improve the performance of a pre-trained segmentation network on low-quality, domain-shifted data by utilizing the intermediate steps of an iterative reconstruction process, without access to ground-truth labels at test time and without modifying the reconstruction process itself.

2. Methodology: IRTTA

The authors propose IRTTA (Intermediate Reconstruction for Test-Time Adaption), a framework that adapts a frozen downstream segmentation network using the reconstruction trajectory.

Key Components:

1. Iterative Reconstruction Trajectory:

   • The method utilizes a diffusion-based reconstruction model (specifically GARD for OCT) that generates a sequence of images $x = (x_0, x_1, \dots, x_{S-1})$ over $S$ time steps.
   • $x_0$ is typically noise or a pseudoinverse, and $x_{S-1}$ is the final high-quality reconstruction.
   • Intermediate images $x_i$ are extracted after data-consistency projections and serve as inputs to the segmentation network.

2. Temporal Modulation Network ($g_\Psi$):

   • Instead of retraining the entire segmentation backbone ($f_\theta$), IRTTA keeps the backbone frozen.
   • A lightweight modulation network ($g_\Psi$) is introduced. It takes the current reconstruction time-step $t_i$ (encoded via sinusoidal embeddings) as input.
   • $g_\Psi$ predicts modulation parameters $(\gamma, \beta)$ for the normalization layers (e.g., BatchNorm or LayerNorm) of the frozen backbone.
   • Residual Affine Transformation: The modulation is applied as a residual: $\bar{z} = e^\gamma \odot \bar{x} + \beta$.
   • Zero-Initialization: The final layer of $g_\Psi$ is initialized to zero (in log-space for scaling) to ensure that at the start of adaptation, the network behaves exactly like the original pre-trained model ($\bar{z} = \bar{x}$), preserving initial performance.

3. Test-Time Adaptation (TTA) via Entropy Minimization:

   • Since no ground truth labels are available at test time, the modulation network $\Psi$ is optimized using unsupervised entropy minimization.
   • The objective is to minimize the average prediction entropy across all steps in the trajectory:
     $$L(\Psi) = -\sum_{i=1}^{S} \frac{1}{d} \sum_{p=1}^{d} \sum_{c=1}^{C} \hat{y}_{i,p,c} \log(\hat{y}_{i,p,c})$$
   • This encourages the model to produce confident predictions across the entire reconstruction trajectory.

4. Uncertainty Estimation:

   • The method naturally generates multiple predictions (one for each time step).
   • The final segmentation is the ensemble mean of these predictions.
   • Semantic Uncertainty: The pixel-wise entropy of the mean prediction ($H(\hat{y}_\mu)$) serves as an uncertainty map. High entropy regions indicate ambiguous anatomical structures or potential lesions that were degraded in the input but recovered during reconstruction.

3. Key Contributions

• Novel Modulation Framework: IRTTA is the first method to exploit the full trajectory of iterative reconstructions for test-time adaptation by modulating normalization layers based on the reconstruction time-step.
• Zero-Shot Uncertainty Estimation: It provides semantically meaningful uncertainty maps without retraining the backbone or requiring additional architectural changes. The uncertainty correlates with anatomical ambiguity rather than just image boundaries.
• Superior Adaptation Performance: The method achieves state-of-the-art results in test-time adaptation for medical segmentation, outperforming existing TTA baselines (TENT, CoTTA) and unsupervised domain adaptation (UDA) methods that require source domain data during training.
• Efficiency: It adapts only a small set of parameters (the modulation network) while keeping the heavy segmentation backbone frozen, making it computationally efficient.

4. Experimental Results

The method was evaluated on the RETOUCH benchmark using OCT data from three devices: Cirrus, Topcon, and Spectralis. Spectralis served as the high-quality target domain, while Cirrus and Topcon were the low-quality source domains.

• Quantitative Performance (Dice Similarity Coefficient):

  • Cirrus $\to$ Spectralis: IRTTA achieved a mean Dice score of 0.603, outperforming the baseline GARD (0.553), SCUNet (0.551), and other TTA/UDA methods (e.g., SVDNA at 0.433, SegClr at 0.512).
  • Topcon $\to$ Spectralis: IRTTA achieved 0.444, again leading among TTA methods and surpassing SVDNA (0.448) and SegClr (0.517) in specific metrics, demonstrating strong generalizability.
  • Comparison to Supervised: While IRTTA (0.603) trails the fully supervised upper bound (0.645), it significantly closes the gap with unsupervised methods.

• Uncertainty Quantification:

  • Calibration: IRTTA reduced the Expected Calibration Error (ECE) from 0.013 (GARD) to **0.007**.
  • Precision-Recall: The PRAUC improved from 0.532 to 0.672 on the Cirrus dataset.
  • Qualitative: Uncertainty maps correctly highlighted ambiguous lesions and anatomical regions, whereas baseline methods often highlighted only object boundaries.

• Ablation Studies:

  • Trajectory Usage: Utilizing the full trajectory yielded better results than adapting only on the final image.
  • Hyperparameters: Performance saturated at 100 adaptation steps and 10 reconstruction steps ($S=10$). Excessively large embedding sizes (>64) led to overfitting.

5. Significance and Conclusion

The paper demonstrates that the iterative nature of modern reconstruction algorithms is an underutilized source of semantic information. By treating the reconstruction trajectory as a sequence of domain-shifted inputs, IRTTA enables a frozen segmentation network to adapt dynamically to the evolving image quality.

Significance:

• Clinical Impact: It allows for accurate segmentation on low-cost medical devices without requiring expensive retraining or access to labeled data from the target domain.
• Interpretability: The generated uncertainty maps provide clinicians with valuable insights into which regions of a scan are ambiguous, aiding in diagnostic confidence.
• Generalizability: The approach is agnostic to the specific generative model used for reconstruction and can be extended to other modalities like MRI and CT.

The authors conclude that while a gap remains compared to fully supervised models, IRTTA offers a robust, efficient, and highly effective solution for domain adaptation in medical imaging, with future work planned to explore fusion mechanisms and other imaging modalities.