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

AI summary

Overview

Research area: Medical image analysis, domain generalization, and geometric deep learning (hyperbolic representation learning).

Technical level: Intermediate. The core idea is conceptually simple (project frozen features into a curved space and add a small second branch), but understanding why it works requires some familiarity with hyperbolic geometry, representation learning, and distribution shift.

Scope: The paper validates hyperbolic embeddings for medical image classification and proposes HypCBC, an unsupervised two-branch training scheme that improves generalization to unseen clinical domains.

What This Paper Is About

Medical AI models are trained on data from specific scanners, hospitals, and patient populations, but they must work reliably on data from entirely different sources. Most methods represent images in flat (Euclidean) space, which the authors argue mismatches the inherently hierarchical structure of clinical data. The paper's goal is to show that embedding medical image features on a hyperbolic manifold, plus a lightweight consistency constraint between two branches, produces classifiers that hold up better under domain shift without needing domain labels or target-domain data.

Key Contributions

1. First comprehensive validation of hyperbolic representation learning for medical imaging. The authors evaluate Euclidean versus hyperbolic classifiers across eleven in-distribution medical datasets spanning nine imaging modalities, sample sizes from 10² to 10⁵, and 2 to 11 label classes, using three ViT backbones (ViT, DeiT3, DINOv2). Hyperbolic embeddings win with statistical significance (Wilcoxon signed-rank, p < 0.05).

2. A novel unsupervised, domain-invariant hyperbolic cross-branch consistency (HypCBC) constraint. A two-branch architecture uses a 2D hyperbolic bottleneck branch to regularize a 128D hyperbolic main branch via a KL-divergence loss on class logits. Unlike domain-adversarial methods, it requires no domain labels and no gradient reversal.

3. Strong domain generalization results across three benchmarks. HypCBC outperforms established Euclidean methods (ERM, IRM, GroupDRO, VREx, DANN, CDANN, CORAL, MMD, plus augmentations like RandAugment, AugMix, Med-C) on Fitzpatrick17k, Camelyon17-WILDS, and a cross-dataset retinal benchmark, by an average of +2.1% AUC.

4. Ablations explaining why the design works. Studies on bottleneck dimension and manifold geometry show that a low-dimensional branch (especially 2D) enforces domain invariance while retaining label discriminability, and that hyperbolic regularization yields consistent gains where Euclidean regularization does not.

Main Findings

• Hyperbolic beats Euclidean on in-distribution classification. Across all models and datasets, hyperbolic classifiers improve accuracy significantly (p < 0.05). With DINOv2, hyperbolic heads win on 10 of 11 tasks, with the largest gains on OCT (+2.70%), Path (+1.48%), and Tissue (+1.36%). With DeiT3, hyperbolic wins 10 of 11; with ViT, 8 of 11.

• The low-dimensional branch controls the invariance/discrimination trade-off. Domain-classification AUC (lower is better) is lowest at a 2D bottleneck for both manifolds and rises with dimension, confirming that compact bottlenecks suppress domain information. Label AUC remains high even at 2D (up to 99% on Camelyon17) and keeps improving with dimension.

• Hyperbolic embeddings keep benefiting from higher dimensions; Euclidean ones plateau. On Fitzpatrick17k and Camelyon17-WILDS, Euclidean label AUC stops improving past a certain dimension, while hyperbolic embeddings gain further, reaching roughly a 3% AUC gap at d = 128.

• HypCBC adds gains on top of hyperbolic ERM. The single-branch hyperbolic model (HypERM) improves over Euclidean ERM by an average of +1.42% AUC; HypCBC adds a further statistically significant +0.66%.

• Bigger domain shift produces bigger hyperbolic gains. Fitzpatrick17k, which has the most extreme shift (skin-tone groups I–II, III–IV, V–VI), shows the largest improvements. HypCBC reaches up to a 2% AUC boost on the III–IV fold, versus only 0.5% for the Euclidean consistency variant.

• Euclidean cross-branch consistency is unreliable; hyperbolic consistency is not. On Camelyon17-WILDS, Euclidean CBC produces negative AUC changes for all bottleneck sizes, while HypCBC yields positive gains across datasets and reasonable bottleneck sizes. Performance collapses only when the bottleneck is as large as the main branch (d₂ = 128).

• The method is insensitive to its two hyperparameters. Sweeping the KL weight λ ∈ {0.1, 0.2, 0.5, 1.0} and temperature T ∈ {1.0, 3.0, 5.0, 10.0} changes results only marginally (standard deviation between 0.1 and 0.7), indicating no need for careful tuning.

• HypCBC competes without domain labels. On Camelyon17-WILDS it matches the top Euclidean method (VREx) even though VREx uses domain labels during training and HypCBC does not. On the Retina test split, HypCBC beats the best Euclidean augmentation method (AugMix) by 0.68%.

Methodology in Plain English

The approach has three layers of design.

First, the authors avoid training a hyperbolic network end-to-end, which prior work shows can be numerically unstable and slow to converge on large datasets. Instead, they freeze a pre-trained Euclidean vision transformer (DINOv2, DeiT3, or ViT) and treat it purely as a feature extractor.

Second, they project those frozen features into the Poincaré ball, a standard model of hyperbolic space with constant negative curvature. The projection uses the exponential map, which is a smooth, invertible transformation that places Euclidean vectors onto the curved manifold. Because hyperbolic space expands exponentially away from the origin, it can represent hierarchical and fine-grained distinctions with more room than flat space allows. Classification on top of these embeddings uses multiclass logistic regression based on geodesic (curved shortest-path) distances rather than ordinary dot products.

Third, they add a second branch. The main branch projects features to 128 dimensions and is what gets used at test time, capturing fine-grained, class-specific detail. A second branch projects to just 2 dimensions, creating a deliberate information bottleneck. Both branches are trained with cross-entropy against the true labels. Additionally, the 2D branch's softmax logits (softened with temperature T = 3) supervise the 128D branch through a KL-divergence term weighted by λ = 0.2. The intuition is that the 2D branch cannot retain enough capacity to encode which domain an image came from, but it can encode coarse class structure — so forcing the main branch to agree with it pushes domain-agnostic information upward while the main branch retains its own domain-sensitive detail.

The overall loss is simply the sum of the two cross-entropy losses plus the weighted KL term. The paper draws an analogy to domain-adversarial training, but replaces the adversarial game and domain labels with a dimensionality bottleneck.

Experiments use curvature c = 1.0, a feature-clipping radius of 1.0 for numerical stability, AdamW with a cosine-annealing schedule (initial learning rate 1e-4, batch size 64), early stopping after 10 epochs without improvement, and five seeds per configuration.

For domain generalization, the benchmarks are: Fitzpatrick17k (dermatology, leave-one-domain-out over three grouped skin-tone domains), Camelyon17-WILDS (histopathology, 422,394 images from five hospitals, standard train/val/test hospital splits), and a Retina benchmark built by training on APTOS 2019 and DeepDR and testing on Messidor-2, with IDRiD as validation. Results are reported as AUC averaged over five seeds.

Why This Matters

Impact on research. The paper makes a case that the geometry of the representation space is an underused lever for robustness in medical imaging, separate from augmentation or invariant-risk objectives. It also demonstrates that hyperbolic learning can be applied to frozen foundation-model features rather than requiring unstable end-to-end training, which lowers the barrier for others to adopt it. The finding that a 2D bottleneck both enforces invariance and enables direct visualization of embeddings suggests a practical diagnostic tool for studying what a model has actually learned.

Real-world applications:

• Dermatology across skin tones. A model trained mostly on lighter skin tones can degrade badly on darker skin. Better invariance to this shift directly affects equity of care.
• Histopathology across hospitals. Tissue-staining protocols, slide scanners, and lab procedures vary between hospitals; a model deployed across a hospital network must not degrade at the sites it was not trained on.
• Retinal screening across cameras. Diabetic retinopathy screening uses fundus cameras with different sensors, grading protocols, and patient cohorts; robustness to these differences is needed for population-scale screening programs.
• General multi-vendor clinical deployment. The same problem recurs across CT, X-ray, OCT, and microscopy — any setting where equipment and protocols differ between training and deployment sites.

Industry relevance. Medical AI vendors face regulatory and commercial pressure to demonstrate performance that is stable across sites rather than tuned to one institution. A method that requires no domain labels, no access to target-domain data, and only a lightweight head on top of an existing frozen backbone is attractive for practical deployment: it can be bolted onto existing feature extractors, requires no re-training of the foundation model, and has hyperparameters that barely matter. Cross-branch consistency also offers a cheaper alternative to collecting domain annotations, which are often unavailable in clinical settings.

Future Directions

• Learnable or adaptive curvature. The current implementation fixes the curvature parameter c at 1.0. The authors note that a static curvature may be suboptimal across datasets and manifold dimensions, and suggest learnable curvature or data- and manifold-adaptive schedules. Note that performance does collapse when the bottleneck grows to 128 dimensions, so the relationship between dimensionality and curvature is an open question.

• Unfreezing the backbone. The method currently freezes the Euclidean feature extractor and only trains lightweight hyperbolic heads. Whether end-to-end or partially fine-tuned hyperbolic training can be made stable enough to add further gains — without the numerical instability reported in prior work — remains unresolved.

• Extending beyond classification. All experiments are image classification or binary prediction tasks. Whether the same cross-branch consistency idea transfers to segmentation, detection, or regression tasks with spatial outputs is untested.

• Broader backbone and modality coverage. The evaluation is limited to ViT-family backbones and specific benchmark datasets. Testing convolution-based and state-space architectures, and settings such as multi-modal imaging or multi-label problems (the Chest dataset was excluded for its multi-label setup), would clarify how general the approach is.

Target Audience

This paper is most useful to machine learning researchers working on domain generalization, robustness, or geometric deep learning, and to medical imaging researchers and engineers building classifiers intended for multi-site or multi-vendor clinical deployment. Practitioners who already use frozen foundation-model features and want a low-cost robustness improvement will find the method directly applicable. Readers without prior exposure to hyperbolic geometry should expect to spend time on Section 3.1, but the architectural idea is implementable without deep mathematical background.

Authors’ 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 .