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

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 .

💡 Research Summary

The paper tackles the persistent problem of domain shift in medical image analysis, where limited data and variations in scanners, protocols, and patient populations often cause deep learning models to fail when deployed in new clinical settings. While most recent domain‑generalization approaches operate on Euclidean feature spaces, the authors argue that the flat geometry of Euclidean space cannot adequately capture the hierarchical and multi‑scale relationships inherent in medical images. To address this, they introduce HypCBC (Domain‑Invariant Hyperbolic Cross‑Branch Consistency), a framework that combines a frozen Euclidean backbone (Vision Transformer variants) with lightweight hyperbolic projection layers and a novel unsupervised consistency regularizer.

Methodology

1. Feature Extraction – A pre‑trained ViT backbone is frozen; its output f∈ℝⁿ is fed into two linear heads producing Euclidean embeddings h₁∈ℝ¹²⁸ and h₂∈ℝ².
2. Hyperbolic Mapping – Both embeddings are projected onto a Poincaré ball Dᶜ using the exponential map exp_c, yielding hyperbolic vectors z₁ and z₂. The curvature parameter c controls the negative curvature of the manifold.
3. Branch Design – The high‑dimensional branch (z₁) is intended to retain fine‑grained, class‑specific information, while the low‑dimensional branch (z₂) acts as an “information bottleneck” that naturally suppresses domain‑specific cues because a 2‑D hyperbolic space cannot encode many discriminative directions.
4. Classification Heads – Each hyperbolic vector passes through a Multiclass Logistic Regression (MLR) head that computes class logits from geodesic distances on the manifold. Both branches are trained with standard cross‑entropy loss.
5. Cross‑Branch Consistency (CBC) – The logits from the low‑dimensional branch are temperature‑scaled and used as a soft target for the high‑dimensional branch via a KL‑divergence term: L_KL = T²·KL(σ(ŷ₂;T)‖σ(ŷ₁;T)). This unsupervised regularizer forces the high‑dimensional representation to align with the domain‑invariant signal encoded in the bottleneck branch, without requiring domain labels or adversarial training.
6. Overall Objective – L = CE(ŷ₁, y) + CE(ŷ₂, y) + λ·L_KL, where λ balances the consistency term.

Experiments

• In‑Distribution (ID) Classification: Evaluated on 11 diverse medical datasets (various modalities, sample sizes 10²–10⁵, label granularities 2–11). Hyperbolic embeddings consistently outperformed Euclidean baselines (p < 0.05).
• Domain Generalization (DG): Benchmarked on Fitzpatrick17k (dermatology), Camelyon17‑WILDS (histopathology), and a cross‑dataset retinal imaging suite. HypCBC achieved an average AUC gain of +2.1 % over state‑of‑the‑art Euclidean DG methods.
• Ablation Studies:
  • Varying the dimensionality of a single‑branch model showed that higher dimensions improve ID accuracy but increase domain variance.
  • Comparing low‑dimensional (2‑D) versus high‑dimensional (128‑D) regularization demonstrated that the 2‑D bottleneck yields robustness gains that cannot be explained solely by additional parameters.
  • Curvature c and bottleneck size were systematically varied, confirming that a modest negative curvature and a very low‑dimensional bottleneck are optimal for suppressing domain‑specific features.

Key Insights

• Hyperbolic space’s exponential volume growth provides ample capacity for fine‑grained distinctions while a tiny 2‑D hyperbolic branch cannot over‑fit to domain artifacts, effectively acting as a domain‑invariant teacher.
• The CBC loss operates at the logit level, allowing cross‑dimensional knowledge transfer without the instability often observed in hyperbolic end‑to‑end training.
• Freezing the backbone and only learning lightweight projections sidesteps the numerical instability reported in prior hyperbolic networks, making the approach computationally efficient and easy to integrate with existing foundation models.

Contributions

1. First comprehensive validation of hyperbolic representation learning for medical imaging, showing statistically significant improvements across a broad set of ID tasks.
2. Introduction of a novel two‑branch hyperbolic training scheme with unsupervised cross‑branch consistency that markedly boosts domain generalization without requiring domain labels or extra data.
3. Detailed analysis of how manifold geometry and bottleneck dimensionality affect the trade‑off between label discrimination and domain invariance.

Future Directions
The authors suggest extending HypCBC to fully end‑to‑end hyperbolic backbones, exploring multi‑domain adversarial extensions, and applying the framework to 3‑D modalities (CT, MRI) or large‑scale federated learning scenarios.

Overall, HypCBC presents a compelling, low‑overhead strategy to harness hyperbolic geometry for robust, generalizable medical image analysis, bridging the gap between theoretical advantages of non‑Euclidean spaces and practical deployment constraints in clinical AI.