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

The paper addresses Registration-based Few-shot Medical Image Segmentation (RFMIS), a paradigm that propagates labels via image warping and leverages anatomical correspondence to preserve geometric consistency. In RFMIS, unsupervised registration on large unlabeled data produces pseudo-labeled images for supervised segmentation. The key challenge in this task lies in the quality of pseudo-labels, since inaccurate or anatomically inconsistent labels can degrade the segmentation performance, particularly for small anatomical structures.

The authors identify a critical limitation of existing RFMIS methods: most RFMIS methods optimize and infer at the pixel level in Euclidean output embedding spaces, treating anatomical regions as flat and disjoint classes. However, medical anatomical categories exhibit an intrinsic coarse-to-fine hierarchy defined by radiologists. This neglect of inherent hierarchies "causes confused small or ambiguous anatomical structures, inducing deformation fields that fit local intensities but violate global hierarchical anatomical consistency, leading to predicted segmentation results with weakened anatomical discrimination."

The paper notes that while recent methods in segmentation, registration, and image-text alignment leverage hyperbolic space owing to its exponentially expanding geometry and natural ability to encode hierarchical structures, existing hyperbolic methods primarily focus on hyperbolic embedding, while neglecting the rich semantic information embedded in Euclidean space, resulting in incomplete modeling of anatomical semantics. Furthermore, no prior methods have investigated its potential in RFMIS, which necessitates the joint optimization of registration and segmentation and is more challenging than a single-task setting.

This motivates the central research question: how to effectively leverage both rich semantic information in Euclidean space and anatomical hierarchies in hyperbolic space within a complex dual-task framework.

The authors propose H2AL (Hyperbolic Hierarchy-aware Aggregative Learning), a framework that enhances both deformation plausibility and anatomical discrimination for dual-task learning. The framework consists of:

The H2I module "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." It consists of two key components:

Transformation-guided Supervised Hyperbolic Contrastive Learning (TSHCL) strategy: This maps Euclidean embeddings into hyperbolic space using the Poincaré ball model, Leveraging the natural ability of hyperbolic geometry to model tree-like structures. Unlike previous methods that solely rely on the hyperbolic geometry for implicit representation learning, this approach incorporates multi-label supervision to explicitly guide the formation of a structured embedding space in the Poincaré ball. It computes pairwise similarity using geodesic distance and designs distance-aware contrastive weights: a pull weight for positive pairs that increases with distance to promote compactness, and a push weight for negative pairs that decays with distance to avoid excessive repulsion between unrelated samples. This ensures stronger attraction between dissimilar but true positive samples, and softer repulsion between negatives that are already distant.

Gated Infusion Block (GIB): This component first projects the hierarchy-aware hyperbolic embeddings into Euclidean space via logarithmic mapping, then the projected embeddings are passed through a gated attention mechanism to generate hierarchy-informed attention maps, which are subsequently used to modulate the Euclidean features via an affine transformation. The design takes advantage of the rich semantic representations from the Euclidean space while incorporating the hierarchical priors embedded in hyperbolic geometry. Two task-specific GIBs are instantiated for registration and segmentation, respectively.

The paper notes that existing RFMIS methods employ a two-stage training strategy that the shared encoder is pretrained and then followed by alternating optimization of task-specific decoders, resulting in time-consuming and suboptimal task interaction modeling. To address this, the authors propose a one-stage, end-to-end training strategy, termed learning with Gradient Aggregation (GA). Task-specific decoders are updated separately during backpropagation, then we aggregate them to update the shared encoder in a unified manner using the sum of gradients from both tasks. This approach facilitates synergistic representation learning, encouraging the encoder to learn features that are beneficial to both registration and segmentation tasks and allows the model to better capture inter-task dependencies and avoid task dominance.

For registration, the total loss combines a smoothness loss (to enforce deformation smoothness), a similarity loss (measuring similarity between warped image and target), and the TSHCL loss. For segmentation, the loss combines a multi-class Dice coefficient loss with the TSHCL loss.

The main contributions summarized in the paper are:

• We propose H2AL for RFMIS, which enhances both consistent deformation and anatomical representation to increase final segmentation performance.

• We introduce an H2I module that jointly exploits the hyperbolic hierarchical priors and the Euclidean semantic richness, learning precise and anatomically discriminative representations.

• "We devise a GA-based joint optimization strategy for one-stage end-to-end training, which aggregates task-specific gradients infused with semantic and hierarchical cues to collaboratively optimize the shared encoder across dual tasks."

• Extensive experiments on brain and cardiac datasets show that H2AL outperforms state-of-the-art RFMIS methods in both segmentation and registration.

The method is evaluated on two anatomical regions with five experimental settings:

• Brain dataset: 410 3D T1-weighted MRI scans from OASIS, PPMI, ADNI, and ABIDE, resampled to 128×128×128, with 330 volumes for training and 80 for testing. Evaluated under Atlas-based, 1-shot, and 5-shot settings.

• Cardiac dataset: 125 3D CT scans from MM-WHS, ASOCA, and CAT08, resampled to 144×144×128, with 100 images for training and 25 for testing. Evaluated under 1-shot and 5-shot settings.

Comparison methods include fully supervised methods (UNet, MASSL, Co-manifold), atlas-based methods (DeepAtlas, DataAug, TBIOneShot), and few-shot methods (BRBS, PC-Reg-RT, Bi-JROS), along with registration-only methods (SyN, VoxelMorph). Metrics include Dice, NCC (Normalized Cross-Correlation), and HD95 (95th-percentile Hausdorff distance).

Few-shot methods outperform both fully supervised and atlas-based methods, demonstrating their effectiveness in low-label scenarios. H2AL achieves the best performance across all settings on the Brain dataset, significantly outperforming leading few-shot methods such as BRBS and Bi-JROS. On the Cardiac dataset, it ranks first in the 1-shot setting and second in the 5-shot setting, consistently surpassing Bi-JROS. Notably, for small anatomical structures, it improves Dice scores by 1.47% and 1.84% for registration and segmentation than SOTA methods respectively.

H2AL further leverages hierarchical anatomical information by first aligning coarse parent structures and then refining finer child structures, achieving the highest Dice and NCC scores. Importantly, these gains are not obtained at the expense of topology preservation. H2AL maintains a low percentage of non-diffeomorphic regions, outperforming most competing methods, demonstrating a favorable balance between deformation flexibility and topology preservation.

The paper presents confusion matrices showing our method shows higher agreement with the ground truth than competing approaches, indicating more accurate segmentation and improved registration performance on small structures. This is attributed to strong registration capability enables high-quality pseudo-label generation and joint modeling of hierarchical structures facilitates effective integration of structural priors.

In failure case analysis, despite substantial registration errors, "H2AL consistently achieves superior segmentation performance and even yields positive Dice gains in two of four cases. This robustness stems from the hierarchical modeling strategy, which first aligns parent structures to establish stable global anchors and then refines child structures, mitigating error propagation caused by misalignment."

The ablation study on the Brain dataset under the Atlas setting shows:

• Effect of GA: Both registration and segmentation show consistent performance improvements... suggesting that the two objectives provide complementary supervisory signals rather than competing ones.

• Effect of H2I: "Incorporating TSHCL and GIB into H2I brings consistent gains in both registration and segmentation, especially for small-scale structures... H2I can effectively bridge the complementary strengths of Euclidean and hyperbolic spaces by preserving the semantic richness of Euclidean representations while injecting the hierarchy-aware priors encoded in hyperbolic geometry."

"Euclidean embeddings exhibit compact clusters but weak global structure, whereas hyperbolic embeddings capture clear hierarchies at the cost of compactness. In contrast, our method integrates both spaces via H2I, achieving both compact intra-class clustering and well-structured inter-class separation."

Gradient aggregation accelerates convergence and improves the final performance, as separately induced task updates may lead to inconsistent optimization dynamics and unstable training. The GA strategy encourages more coordinated dual-task learning, alleviates optimization inconsistency, and enhances training stability.

The paper concludes that H2AL enhances both deformation plausibility and anatomical discrimination for dual-task learning through the H2I module and GA-based training. The authors state: More broadly, we believe that jointly modeling Euclidean semantics, hyperbolic hierarchy, and cross-task gradient interaction provides a promising and extensible direction for future medical image analysis.

Here are specific improvements to AI systems based on the H2AL paper, along with what the improved system can do.

1. Integrate hyperbolic hierarchy-aware feature infusion into segmentation networks.

• Improvement: Replace purely Euclidean feature maps in the decoder (or shared encoder) with a dual-space mechanism: first project features into a Poincaré ball, apply supervised hyperbolic contrastive learning with geodesic distances, then project back and modulate the original Euclidean features via a gated attention block (logarithmic mapping + affine transformation).

• What it can do: The AI system simultaneously preserves fine-grained local semantic details (from Euclidean space) and enforces coarse-to-fine anatomical hierarchies (from hyperbolic space). This leads to significantly better segmentation of small or ambiguous anatomical structures, higher Dice scores on low-contrast boundaries, and more anatomically consistent predictions without requiring extra labeled data.

2. Use transformation-guided supervised hyperbolic contrastive learning for medical image representation.

• Improvement: In multi-task or pre-training pipelines, construct positive/negative pairs based on known anatomical correspondences (e.g., registered regions), and compute contrastive loss using geodesic distance in the Poincaré ball with distance-aware pull/push weights — stronger attraction for distant true positives, softer repulsion for already-distant negatives.

• What it can do: The improved system learns an embedding space where both intra-class compactness and inter-class hierarchical separation are explicitly optimized. It becomes more robust to large anatomical variability and can better transfer to downstream tasks such as few-shot segmentation, where only a handful of labeled examples are available.

3. Implement gradient aggregation for joint optimization of multi-task systems.

• Improvement: For any network with a shared encoder and multiple task-specific decoders (e.g., registration + segmentation, detection + segmentation), update each decoder separately via backpropagation, then aggregate the task-specific gradients (sum or weighted sum) to update the shared encoder in a single step, instead of alternating between tasks or using a simple weighted total loss.

• What it can do: The improved system trains end-to-end in one stage, converges faster, avoids task dominance, and learns encoder features that are consistently beneficial to all tasks. This reduces optimization instability and eliminates the need for pretraining plus alternating finetuning, saving time while improving final multi-task performance.

4. Inject hierarchy-aware priors into deformable registration to improve coarse-to-fine alignment.

• Improvement: During registration, use the hyperbolic-infused features to first establish correspondences of large parent structures (e.g., whole brain regions), then refine child structures (e.g., subcortical nuclei), by incorporating hierarchy-aware contrastive supervision into the registration loss (alongside smoothness and similarity terms).

• What it can do: The improved system produces deformation fields that preserve global anatomical topology while accurately aligning small structures. It achieves higher Dice and NCC scores, maintains a low percentage of non-diffeomorphic regions, and generates higher-quality pseudo-labels for subsequent segmentation — especially useful in few-shot or atlas-based settings.

5. Combine hyperbolic hierarchy modeling with Euclidean semantic richness in one unified framework.

• Improvement: Instead of choosing between hyperbolic embedding methods and standard Euclidean learning, use a gated infusion block that dynamically selects how much hierarchical information to inject into Euclidean features, with task-specific instantiations for registration and segmentation.

• What it can do: The improved system achieves both compact intra-class clustering and well-structured inter-class separation, leading to more discriminative representations. It outperforms methods that rely solely on either Euclidean or hyperbolic spaces, and it generalizes to new domains (e.g., from brain to cardiac MRI/CT) with minimal retraining.

6. Leverage hierarchical anatomical modeling for robustness to registration failures.

• Improvement: In few-shot medical image segmentation pipelines that rely on registration-based label propagation, add a hierarchy-aware component that aligns parent structures as stable global anchors before refining child structures.

• What it can do: Even when registration is partially misaligned, the improved system mitigates error propagation and still produces accurate segmentations (positive Dice gains even under substantial registration errors). This makes the AI system more reliable for clinical use where poor initial alignment is common.

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.