HypCBC: Domain-Invariant Hyperbolic Cross-Branch Consistency for Generalizable Medical Image Analysis

TL;DR

This work tackles the challenge of generalization under distribution shifts in medical imaging by leveraging hyperbolic geometry. It introduces HypCBC, a backbone-agnostic, two-branch framework that maps Euclidean features into a high-dimensional hyperbolic space (128D) while enforcing domain-invariant information through a compact 2D bottleneck via cross-branch consistency. By freezing a Euclidean backbone and attaching lightweight hyperbolic projections, the method achieves strong in-distribution accuracy and improved domain generalization across dermatology, histopathology, and retinal imaging benchmarks, with significant gains over Euclidean baselines. The results, including ablations on latent dimension and consistency, demonstrate the practical viability and robustness of hyperbolic embeddings for generalizable medical AI, offering a scalable path toward safer clinical deployment.

Abstract

Robust generalization beyond training distributions remains a critical challenge for deep neural networks. This is especially pronounced in medical image analysis, where data is often scarce and covariate shifts arise from different hardware devices, imaging protocols, and heterogeneous patient populations. These factors collectively hinder reliable performance and slow down clinical adoption. Despite recent progress, existing learning paradigms primarily rely on the Euclidean manifold, whose flat geometry fails to capture the complex, hierarchical structures present in clinical data. In this work, we exploit the advantages of hyperbolic manifolds to model complex data characteristics. We present the first comprehensive validation of hyperbolic representation learning for medical image analysis and demonstrate statistically significant gains across eleven in-distribution datasets and three ViT models. We further propose an unsupervised, domain-invariant hyperbolic cross-branch consistency constraint. Extensive experiments confirm that our proposed method promotes domain-invariant features and outperforms state-of-the-art Euclidean methods by an average of $+2.1\%$ AUC on three domain generalization benchmarks: Fitzpatrick17k, Camelyon17-WILDS, and a cross-dataset setup for retinal imaging. These datasets span different imaging modalities, data sizes, and label granularities, confirming generalization capabilities across substantially different conditions. The code is available at https://github.com/francescodisalvo05/hyperbolic-cross-branch-consistency .

Figures (8)

• Figure 1: Given a frozen Euclidean feature extractor $\Phi$ that outputs $\mathbf{f}$, ERM applies a Euclidean linear probe over a 128D projection $\mathbf{h}_{128\mathrm{D}}$. HypERM additionally uses a fixed exponential map $\exp^c_{128}$ to classify over hyperbolic embeddings $\mathbf{z}_{128\mathrm{D}}$. Our method, HypCBC, introduces a second projection $\mathbf{h}_{2\mathrm{D}}$ followed by $\exp^c_{2}$ to yield $\mathbf{z}_{2\mathrm{D}}$. The logits of this low-dimensional branch are used as targets in the KL loss, promoting domain-agnostic information transfer into the high-dimensional branch.
• Figure 2: Given an input image-label pair $(\mathbf{x},y) \in \mathcal{X} \times {\mathcal{Y}}$, we extract an image embedding $\mathbf{f}=\Phi(\mathbf{x})\in\mathbb{R}^{n}$, where $n$ depends on the chosen backbone. This is projected via two heads into Euclidean embeddings $\mathbf{h}_1\in\mathbb{R}^{128}$ and $\mathbf{h}_2\in\mathbb{R}^{2}$. Each is mapped into its respective Poincaré ball $\mathbb{D}_c^d$ by $\exp^c_d$, yielding hyperbolic embeddings $\mathbf{z}_1$ and $\mathbf{z}_2$. Both branches incur cross-entropy losses on their Multiclass Logistic Regression (MLR) logits $\hat{\mathbf{y}}_1$ and $\hat{\mathbf{y}}_2$. In addition, $\hat{\mathbf{y}}_2$ also supervises $\hat{\mathbf{y}}_1$ via a KL-based consistency term. The high-dimensional branch captures fine-grained, class-specific features for inference, and the low-dimensional branch enforces an information bottleneck that promotes domain-invariant representations.
• Figure 3: Overview of dataset domains. Top-Left: Camelyon17-WILDS, showing representative "tumor" patches from each contributing hospital (H0--H4). Bottom-Left: Retina, showing representative fundus images ($y=4$) from each dataset domain (APTOS 2019, DeepDR, IDRiD, Messidor-2). Colored frames indicate the train (pink), validation (orange), and test (purple) subsets. Right: Fitzpatrick17k, showing representative "malignant" skin-lesion images from each of the three skin-tone groups (I--II, III--IV, V--VI).
• Figure 4: Figure \ref{['fig:domain-accuracy-a']} shows the domain AUC ($\downarrow$) vs. latent dimension $d$: lower is better (more domain-invariant), while Figure \ref{['fig:domain-accuracy-b']} plots the label AUC ($\uparrow$) vs. $d$: higher is better (more discriminative). The curves are shown for Euclidean (ERM) and hyperbolic embeddings (HypERM). Results are reported on in-distribution splits where all domains appear in train/val/test, yielding $3, 5,$ and $4$ domains for Fitzpatrick17k, Camelyon17-WILDS, and Retina, respectively.
• Figure 5: Improvement in AUC of two-branch consistency regularization over single-branch baseline ($\Delta\mathrm{AUC}$) for Euclidean and hyperbolic manifolds, while varying bottleneck dimension $d_2$. These are termed CBC and HypCBC, respectively. From left to right: (1) Fitzpatrick17k leave-one-domain-out folds (I-II, III-IV, V-VI), differentiated with bar hatches, (2) Camelyon17-WILDS (test), and (3) Retina (test). While Euclidean regularization gains vary by dataset, hyperbolic regularization provides positive $\Delta\mathrm{AUC}$ across every task and reasonable (i.e., $< 128$) bottleneck size.