Test-Time Adaptation in Optical Coherence Tomography Using Trajectory-Aligned Time-Independent Flow

Abstract: Optical coherence tomography (OCT) is essential in ophthalmology, but inconsistent image quality especially in low-cost devices hinders automated analysis. To address this, we introduce a flow-matching-based test-time adaptation method that generates high-quality surrogate images from noisy inputs. Typically, domain gaps between test and training data cause pixel distribution mismatches during the denoising process. We overcome this by matching the test image's histogram to synthetic reference trajectories, successfully aligning the input with expected distributions. Additionally, we remove the network's time conditioning to account for slight deviations in real-world noise distributions. Our approach achieves state-of-the-art performance in segmenting critical biomarkers for two stages of Age-related Macular Degeneration (AMD). Code is available: https://github.com/Veit21/tta-flow.

• The paper introduces TTA-Flow, which applies trajectory-aligned histogram matching with unconditional flow to enhance test-time adaptation in OCT segmentation.
• The method outperforms conventional and diffusion-based TTA techniques, achieving superior DSC scores and markedly reduced FID across multiple datasets.
• Experimental results demonstrate that eliminating time-conditioning and using histogram matching significantly boost robustness and preserve anatomical details.

Trajectory-Aligned Time-Independent Flow for Test-Time Adaptation in Optical Coherence Tomography

Motivation and Problem Formulation

Automated OCT analysis is central to clinical workflows in ophthalmology, but prevalent distribution shifts—especially from hardware variability and lower-cost imaging devices—degrade segmentation efficacy. Classical TTA approaches targeting model weights via entropy minimization or pseudo-labeling have shown inconsistent improvement in medical image segmentation under domain shift. Recent generative paradigms (GANs, flows, diffusion models) attempt to remap test images onto the training manifold during inference, but they typically assume synthetic or ideal noise, rarely reflecting real-world test distributions. Downsampling to accommodate re-noising (as in DDA) compromises fine lesion resolution. Disease-specific cross-modal priors, e.g., CLIP, require extensive disease and modality adaptation.

The paper proposes TTA-Flow: a flow-matching generative model equipped with trajectory-aligned histogram matching and an explicit omission of time-conditioning for robust test-time adaptation of noisy OCT images. The hypothesis is that histogram alignment and unconditional flow modeling circumvent mismatches between real and synthetic noise, facilitating high-fidelity restoration and accurate segmentation of ophthalmic biomarkers.

Methodology

The TTA-Flow framework leverages Flow Matching (FM), parameterizing a deterministic vector field, $\mathbf{v}_\theta$, to transport Gaussian noise to training images along prescribed trajectories governed by an ordinary differential equation. The path between noise and image is enforced via linear interpolation, and regression is performed against the target velocity vector. Crucially, inference omits explicit time/noise level conditioning to accommodate the diverse, unknown noise statistics of real-world test acquisitions, following empirical results from Sun et al. [sun2025noiseconditioningnecessarydenoising].

To bridge the gap between empirical test distributions and theoretical synthetic ones, the model generates reference trajectories and computes average intensity histograms at each trajectory point. Incoming test scans are histogram-matched to the reference trajectory, establishing statistical congruence prior to flow-based restoration. Restoration commences from the matched point, integrating the ODE to recover the denoised image, subsequently passed to downstream segmentation models.

Figure 1: Schematic overview of the TTA-Flow framework. The process begins by (a) generating reference trajectories and (b, c) calculating an average histogram $\bar{H}_s$ for each time step; histogram matching aligns test scans (d), and flow matching reconstructs (e) images for downstream evaluation.

Experimental Design

Evaluations span two datasets: RETOUCH (annotated fluid lesions from three OCT devices) and a paired in-house collection (GA, late AMD, imaged with both Topcon and Spectralis devices). The generative Flow Matching model is trained exclusively on high-SNR Spectralis acquisitions. Segmentation models are likewise trained only on the training domain to ensure no leakage from target domains.

Comparisons include classical TTA (TENT, CoTTA), diffusion-based adaptation (CPDM, DDA), and standard FM. Quantitative metrics are DSC (dice similarity coefficient) for lesion and GA segmentation and FID for perceptual image fidelity.

Figure 2: Comparison to the state-of-the-art on downstream fluid and GA segmentation (Cirrus, Topcon; RETOUCH and in-house GA); TTA-Flow preserves anatomy and reduces noise, outperforming baselines.

Results

TTA-Flow achieves mean DSCs of 58.6 (Cirrus → Spectralis) and 51.0 (Topcon → Spectralis) for fluid segmentation, exceeding diffusion- and flow-based baselines and approaching supervised upper bounds. The method demonstrates resilience across lesion types, outperforming entropy-based TENT on multi-class structures. GA segmentation (Topcon → Spectralis) yields DSCs of 56.1/57.1, and FID is notably reduced from 131.2 (FM) and 85.4 (conditional TTA-Flow) to 32.2 (unconditional TTA-Flow).

Qualitative comparisons confirm anatomical preservation and noise attenuation (Figure 2). Downsampling and data-fidelity terms, as used in DDA and CPDM, are not required; TTA-Flow robustly adapts without hallucinating or omitting lesions.

Ablation studies demonstrate the impact of histogram matching and time-independent flow modeling. Histogram matching elevates mean DSC from 50.8 (FM) to 58.6 (TTA-Flow) on Cirrus. Removing time-conditioning further boosts robustness and perceptual quality, with insensitivity to reference time selection (Figure 3).

Figure 3: Comparison of downstream lesion segmentation performance (DSC) for different reference time points $s_{\text{target}$; unconditional models outperform conditional counterparts consistently.

Implications and Future Directions

The proposed trajectory-aligned, time-independent flow matching for TTA in OCT segmentation constitutes a strong practical advance: histogram preprocessing provides a statistically sound, transparent mechanism for domain adaptation with negligible computational overhead. Unconditional flows demonstrate application-agnostic adaptability, freeing the model from explicit noise-level constraints and obviating complex data-fidelity tuning. The paradigm substantially reduces domain gap in retinal biomarker segmentation from low-cost devices, relevant to rural telemedicine, large-scale screening, and device heterogeneity.

Theoretically, this approach underscores the relevance of empirical distribution matching and the limitations of explicit noise-level modeling in generative restoration. Future work should explore automatic noise-level estimation for adaptive histogram reference selection, integration with Plug-and-Play priors, and multi-modal domain adaptation across imaging modalities.

Conclusion

TTA-Flow introduces a robust, generative test-time adaptation method for OCT that combines trajectory-aligned histogram matching and time-independent flow models. These modifications deliver superior lesion restoration and segmentation, outperforming classical and diffusion-based TTA baselines in challenging, real-world domain shifts. The framework reconciles the empirical noise statistics of test samples with the synthetic trajectory manifold, enabling practical deployment in heterogeneous clinical environments. Adaptive reference estimation and acceleration constitute natural extensions, promising fully automated, high-fidelity image restoration in medical imaging workflows (2606.18876).