Test-Time Domain Generalization via Universe Learning: A Multi-Graph Matching Approach for Medical Image Segmentation

TL;DR

The paper introduces Test-Time Domain Generalization via Universe Learning (TTDG-MGM), a multi-graph matching framework that embeds morphological priors of medical images into learnable universe embeddings to achieve robust segmentation under unseen-domain shifts. By enforcing cycle-consistency through a universe-based formulation and enabling unsupervised test-time adaptation, the method learns domain-invariant features and facilitates rapid deployment without labeled target data. Extensive experiments on retinal fundus, polyp, and MRI prostate segmentation demonstrate state-of-the-art performance in both multi-source and single-source DG settings, with ablations confirming the value of universe embeddings, morphology priors, and the multi-graph approach. The approach offers practical impact for medical imaging deployment by improving cross-domain generalization while highlighting downstream considerations such as computational efficiency and handling of small structures.

Abstract

Despite domain generalization (DG) has significantly addressed the performance degradation of pre-trained models caused by domain shifts, it often falls short in real-world deployment. Test-time adaptation (TTA), which adjusts a learned model using unlabeled test data, presents a promising solution. However, most existing TTA methods struggle to deliver strong performance in medical image segmentation, primarily because they overlook the crucial prior knowledge inherent to medical images. To address this challenge, we incorporate morphological information and propose a framework based on multi-graph matching. Specifically, we introduce learnable universe embeddings that integrate morphological priors during multi-source training, along with novel unsupervised test-time paradigms for domain adaptation. This approach guarantees cycle-consistency in multi-matching while enabling the model to more effectively capture the invariant priors of unseen data, significantly mitigating the effects of domain shifts. Extensive experiments demonstrate that our method outperforms other state-of-the-art approaches on two medical image segmentation benchmarks for both multi-source and single-source domain generalization tasks. The source code is available at https://github.com/Yore0/TTDG-MGM.

Figures (9)

• Figure 1: Illustration of multi-graph matching using the universe of nodes. (a) Direct pairwise matching between graphs 1, 2, and 3. (b) Multi-graph matching with the universe of nodes. (c) Equivalence between (a) and (b) (see Lemma 1 in Sec. \ref{['lemma1']} for details).
• Figure 2: Overview of our TTDG framework. During source model training, data from different domains are jointly used to train the Segmentation Network (feature extractor and segmentation head). Feature maps and ground truth masks are utilized to construct graphs $\mathcal{G}_i$ ($i=4$ in this figure legend) and corresponding labels $Y_i$ (Sec. \ref{['sec:graph_generation']}), with universe embeddings learned via back-propagation, incorporating morphological priors (Sec. \ref{['sec:training']}). At test time, multi-graph matching is performed on all target domain images in each batch. Despite style differences, these images share common morphological patterns. Universe embeddings are frozen as prior knowledge to guide the matching, and the segmentation network is fine-tuned via back-propagation for efficient adaptation (Sec. \ref{['sec:testing']}).
• Figure 3: Visualization comparison of segmentation results and Grad-CAM outputs from the final layer of the backbone network for the No Adapt baseline, VPTTA chen2024each, NC-TTT osowiechi2024nc, and our proposed method on retinal fundus images. Additional visual comparisons are provided in the supplementary material.
• Figure 4: The performance of DSC (%) and inference time (s) with different universe sizes $d$ are shown, with the experimental setup identical to that in Table \ref{['tab:tabel_fundus']}. The first curve represents the No Adapt baseline ($d = 0$), while the subsequent curve corresponds to the results obtained with a universe size of $d = 120$.
• Figure A1: Visualization comparison of segmentation results for the No Adapt baseline, DeY-Net wen2024denoising, VPTTA chen2024each, NC-TTT osowiechi2024nc, and our method in retinal fundus segmentation. The five rows from top to bottom display the final segmentation results for tests conducted on Sites A to E. Different colors represent the segmentation instances of different classes identified by the network.

Theorems & Definitions (2)

• Lemma 1: Cycle-consistency, Universe Matching
• proof