H2AL: Hyperbolic Hierarchy-aware Aggregative Learning for Registration-based Few-shot Medical Image Segmentation

AI summary

Overview

• Research area: Computer vision for medical imaging — specifically registration-based few-shot medical image segmentation (RFMIS), combining deformable image registration and 3D anatomical segmentation under limited labels.
• Technical level: Advanced. The paper assumes familiarity with Riemannian geometry (Poincaré ball model, Möbius operations, exponential/logarithmic maps), contrastive learning, multi-task optimization, and 3D medical imaging metrics.
• Scope: The paper proposes a framework that models the coarse-to-fine hierarchy of anatomical structures in hyperbolic space and injects that hierarchy back into Euclidean features, jointly training registration and segmentation in one end-to-end stage, and evaluates it across five experimental settings on brain MRI and cardiac CT.

What This Paper Is About

Registration-based few-shot medical image segmentation works by warping a labeled image onto unlabeled ones so the labels can be propagated as pseudo-labels; if the warping is anatomically wrong, the pseudo-labels are wrong and segmentation suffers, especially for small structures. Existing methods optimize and infer these features in flat Euclidean space, treating anatomical regions as separate and unstructured, which ignores the coarse-to-fine hierarchy that radiologists actually use. H²AL's goal is to capture that hierarchy using hyperbolic space while keeping the semantic richness of Euclidean features, so that both the deformation and the final segmentation become anatomically more consistent and better at discriminating small, ambiguous regions.

Key Contributions

1. The H²AL framework for RFMIS, which combines consistent deformation with hierarchy-aware anatomical representation to improve final segmentation quality in the registration-based few-shot setting.
2. The Hyperbolic Hierarchy-aware Infusion (H2I) module, which jointly exploits hyperbolic hierarchical priors and Euclidean semantic richness to learn anatomically discriminative representations.
3. A Gradient Aggregation (GA) training strategy for one-stage, end-to-end training that aggregates the task-specific gradients from the registration and segmentation decoders to update the shared encoder, promoting collaborative rather than alternating optimization.
4. Extensive experiments on brain and cardiac datasets under five settings showing H²AL outperforms state-of-the-art RFMIS methods in both segmentation and registration, with visualizations verifying that jointly modeling Euclidean semantics and hyperbolic hierarchy improves discrimination of small anatomical structures.

Main Findings

• Segmentation results: On the Brain dataset (T1 MRI), H²AL reaches 84.17 ± 1.09 Dice / 1.70 ± 0.17 HD95 in the Atlas setting, 85.19 ± 1.23 / 1.41 ± 0.14 in 1-shot, and 85.68 ± 1.52 / 1.32 ± 0.15 in 5-shot. On the Cardiac dataset (CT), it reaches 89.24 ± 2.11 / 4.15 ± 1.47 in 1-shot and 90.24 ± 2.01 / 3.56 ± 1.30 in 5-shot. On the Brain dataset it achieves the best performance across all settings; on Cardiac it ranks first in the 1-shot setting and second in the 5-shot setting, while consistently surpassing Bi-JROS.

• Registration results: In the Brain 5-shot registration comparison, H²AL obtains 83.36 ± 1.64 Dice and 0.386 ± 0.011 NCC, the highest of the compared methods, while keeping the percentage of non-diffeomorphic regions (|J_φ| ≤ 0) at 0.36 ± 0.09 — lower than most competing methods, indicating a favorable balance between deformation flexibility and topology preservation.

• Large gains for small structures: The paper reports that for small anatomical structures, H²AL improves Dice scores by 1.47% (registration) and 1.84% (segmentation) over state-of-the-art methods.

• Better agreement on small structures: Confusion matrices for structures under 1% volume on the Brain dataset in the 1-shot setting show higher agreement with ground truth than competing approaches for both registration-derived pseudo-labels and segmentation.

• Robustness under registration failure: In the four most severe registration failure cases (two per dataset), H²AL achieved positive Dice gains in two of four cases and showed smaller drops than BRBS or Bi-JROS in the others.

• Ablation confirms each component adds value: On the Brain dataset in the Atlas setting, the baseline reaches 78.01 registration Dice and 80.05 segmentation Dice (totals). Adding gradient aggregation raises these to 78.58 and 80.73; adding hyperbolic learning to 79.04 and 80.95; adding TSHCL to 80.75 and 82.69; adding GIB to 81.96 and 84.17 with both branches enabled. Small-structure segmentation Dice rises from 78.06 (baseline) to 82.37 with GIB enabled for both tasks.

• Not reported: The paper does not report training or inference wall-clock times, hyperparameter values for λ_smo, λ_sim, λ_TSHCL^reg, λ_dice, λ_TSHCL^seg, or τ, nor does it give parameter counts for the models.

Methodology in Plain English

The framework uses a shared encoder feeding two task-specific decoders — one for registration, one for segmentation — following the common RFMIS architecture, but with three additions.

First, features produced by each decoder pass through a Hyperbolic Hierarchy-aware Infusion (H2I) module. Inside it, a Transformation-guided Supervised Hyperbolic Contrastive Learning (TSHCL) step maps Euclidean features into the Poincaré ball (a model of hyperbolic space) and uses the pseudo-labels from registration to decide which embeddings should be attracted and which repelled. Unlike prior hyperbolic methods, the pull on true-positive pairs grows with their hyperbolic distance to encourage compactness, while the push on negative pairs decays with distance to avoid over-repelling already-distant samples. The registration branch uses the previous-iteration pseudo-label as supervision; the segmentation branch uses the current one.

Second, a Gated Infusion Block (GIB) maps the hyperbolic embeddings back to Euclidean space through the logarithmic map and uses a gated attention mechanism to modulate the Euclidean features via an affine transformation. Two task-specific GIBs are instantiated, one per task, so registration and segmentation can each select the hierarchical cues they need.

Third, instead of the common two-stage training (pretrain the encoder, then alternate between decoders), the authors update each decoder separately and then aggregate the two gradients to update the shared encoder in a single step. This is intended to avoid task dominance and produce features useful for both tasks at once.

The registration loss combines a smoothness term, a similarity term, and the TSHCL term; the segmentation loss combines a multi-class Dice loss with TSHCL. Training uses an NVIDIA A40 GPU with 48 GB memory, a single Adam optimizer with initial learning rate 1×10⁻⁴ and batch size 1, plus 3D augmentations including random rotations and isotropic scaling.

Data: the Brain dataset has 410 3D T1-weighted MRI scans drawn from OASIS, PPMI, ADNI, and ABIDE, with 330 for training and 80 for testing, resampled to 128×128×128. The Cardiac dataset has 125 3D CT scans from MM-WHS, ASOCA, and CAT08, with 100 for training and 25 for testing, resampled to 144×144×128. Metrics are Dice, NCC, and HD95.

Why This Matters

This work argues that the geometry in which anatomical representations live matters, not just the loss functions or supervision signals. By showing that hierarchical structure can be explicitly encoded rather than left implicit, it offers a route to better pseudo-labels and better boundary behavior for small, ambiguous structures — the exact failure mode that limits registration-based few-shot segmentation in clinical use. It is also, per the authors, the first investigation of hyperbolic hierarchy-aware learning in the RFMIS dual-task setting.

Real-world applications:

• Pediatric and rare-disease imaging, where annotated scans are scarce and structures such as small lesions or sub-structures are clinically critical — one contributing affiliation is a children's hospital.
• Cross-site and cross-scanner deployment, since the Cardiac data deliberately spans different cardiac phases and anatomical variations to test inter-subject and inter-site conditions.
• Low-annotation clinical workflows, where a handful of labeled volumes (the 1-shot and 5-shot settings studied here) must be propagated to a large unlabeled archive.
• Image registration pipelines used in atlas construction, longitudinal tracking, and radiotherapy planning, where topology-preserving deformation fields matter for downstream safety.

Industry relevance: Medical imaging software vendors and hospital imaging departments that need to scale segmentation across large unlabeled archives without proportionate annotation budgets; the code is publicly available at https://github.com/JiamingCai469/H2AL.

Future Directions

• Extending beyond brain and cardiac: The evaluations cover two anatomical regions and five settings; whether the hierarchy-aware infusion transfers to other organs, modalities, or 2D imaging is untested here.
• Reducing geometry overhead: Mapping features in and out of the Poincaré ball adds operations; the paper reports no runtime or memory comparison against the two-stage baselines, leaving efficiency claims to be quantified.
• Making the hierarchy explicit: The coarse-to-fine radiologist hierarchy is referenced as being detailed in the supplementary material, but the main paper depends on pseudo-label supervision to shape the embedding space; learning the hierarchy itself, rather than deriving it from pseudo-labels, is an open question.
• Combining with prototype-based few-shot methods: The authors note that their insight may generalize to broader few-shot settings, but no prototype-based experiments are presented.

Target Audience

Researchers and graduate students working on medical image analysis, few-shot learning, or deformable registration; practitioners building clinical segmentation pipelines with constrained annotation budgets; and anyone interested in hyperbolic representation learning beyond the standard graph and knowledge-graph domains. Readers will need a working knowledge of deep learning for 3D medical images and basic differential geometry to follow the methodology fully, though the motivation and results are accessible without it.

Authors’ abstract

Registration-based Few-shot medical image segmentation (RFMIS) aims to generate pseudo-labels for unlabeled images by warping a labeled image through registration. However, existing methods primarily perform pixel-level optimization and inference in Euclidean space, treating anatomical structures as flat and disjoint. This neglect of inherent hierarchies degrades pseudo-label quality and weakens the discrimination of ambiguous regions, limiting the segmentation performance. To overcome this challenge, we propose a Hyperbolic Hierarchy-aware Aggregative Learning framework for RFMIS, termed H2AL, that enhances both deformation plausibility and anatomical discrimination for dual-task learning. Specifically, we introduce a Hyperbolic Hierarchy-aware Infusion (H2I) module, which leverages the hierarchical modeling capability of hyperbolic space to learn precise hierarchy-aware representations via transformation-guided supervised hyperbolic contrastive learning, and injects such hierarchical priors into Euclidean space through a gated infusion block while preserving semantic richness. Furthermore, we propose an end-to-end joint optimization algorithm by gradient aggregation, where the gradients from the registration and segmentation decoders, embedding semantic and hierarchical cues, are aggregated to update the shared encoder to promote collaborative learning across tasks. Extensive experiments on two anatomical regions, with five experimental settings, demonstrate the effectiveness and efficiency of our method in both registration and segmentation. The code is publicly available at https://github.com/JiamingCai469/H2AL.