Registration-Assisted Prototypical Learning

• The paper introduces a registration-assisted prototypical learning framework that integrates spatial alignment modules with prototype-based segmentation to address domain variability.
• It employs a 3D U-Net backbone and affine registration to normalize anatomical differences, significantly improving Dice scores and reducing inter-institution gaps.
• The approach further incorporates contour-aware losses and vision foundation models to refine segmentation boundaries while overcoming annotation scarcity.

Registration-assisted prototypical learning is a framework that fuses spatial registration modules with prototype-based learning to achieve robust few-shot image segmentation and cross-domain medical image registration. It addresses major challenges in adapting segmentation or registration networks to novel classes or subjects, particularly across varying imaging institutions, protocols, or anatomical variability. At its core, this approach brings explicit spatial alignment and anatomical correspondence into the prototypical learning paradigm, enabling networks to generalize with limited annotated data and across heterogeneous domains (Li et al., 2022, Li et al., 2022, Xu et al., 17 Feb 2025).

1. Prototypical Learning in Few-shot Segmentation

Prototypical learning is an episodic meta-learning strategy where class-wise feature centroids ("prototypes") are extracted from a small labeled support set and used for classification or segmentation of query samples by feature similarity. In 3D medical image segmentation, a support example consists of a volumetric scan $X_s$ and its binary mask $M_s(c)$ for a novel anatomical class $c$. Features are extracted via a 3D backbone $f_\theta$, yielding $F_s$ and $F_q$ as dense feature volumes for support and query inputs.

The segmentation prototype for class $c$ is defined by masked spatial pooling:

$$p_c = \frac{1}{|V|}\sum_{x\,:\,M_s(c)(x)=1} F_s(x), \quad V = \{x \mid M_s(c)(x)=1\},$$

with a background prototype $p_0$ defined analogously for $M_s(c)(x)=0$. The similarity between query features and prototypes, typically via cosine similarity, determines voxel-level class scores:

$$s_c(x) = \frac{F_q(x)\cdot p_c}{\|F_q(x)\|\|p_c\|}, \quad s_0(x) = \frac{F_q(x)\cdot p_0}{\|F_q(x)\|\|p_0\|},$$

and a softmax gives class probabilities per voxel. Both global and local (windowed) prototype pooling may be used; local prototypes provide enhanced spatial sensitivity, particularly for structures with high spatial variability (Li et al., 2022).

The few-shot segmentation task is trained using a Dice or cross-entropy loss over the predicted and ground truth masks:

$$\mathcal{L}_{\mathrm{few}} = 1 - \frac{2\sum_x \hat M_q(c)(x)\,M_q(c)(x)}{\sum_x \hat M_q(c)(x) + \sum_x M_q(c)(x)}.$$

2. Integration of Registration Modules

Direct prototype comparison is confounded by spatial variability (field-of-view, orientation, anatomical differences) between institutions or subjects. Registration-assisted prototypical learning introduces a registration/alignment module to compensate for these global differences.

A dedicated alignment network (e.g., a lightweight conv head with global average pooling and fully connected layers) predicts affine parameters $T_s$, $T_q$ that transform support and query features into an atlas or canonical space. This transformation is supervised via a similarity metric with respect to a reference anatomical atlas $A$, typically using a Dice loss:

$$T^* = \arg\min_T \left[ \mathcal{D}(T(S^*_{\mathrm{base}}), A) + \lambda \mathcal{R}(T) \right],$$

where $S^*_{\mathrm{base}}$ is the predicted base-class segmentation, $\mathcal{D}$ is a Dice-based similarity, and $\mathcal{R}$ regularizes the affine deformation (e.g., penalizing deviation from identity).

Once aligned, features and masks are warped via differentiable resampling. Prototypes are computed and utilized in the atlas space; predictions are mapped back to physical space by inverting the query transformation. This spatial normalization reduces inter-institutional variance and enables generalization to unseen domains (Li et al., 2022, Li et al., 2022).

3. Joint Optimization and Network Architecture

The pipeline is trained end-to-end with a composite objective aggregating segmentation, alignment, and prototype-based losses:

$$\mathcal{L} = \mathcal{L}_{\mathrm{few}} + \alpha\,\mathcal{L}_{\mathrm{seg}} + \beta\,\mathcal{L}_{\mathrm{align}}.$$

Here, $\mathcal{L}_{\mathrm{seg}}$ represents supervised Dice loss on base-class segmentation, and $\mathcal{L}_{\mathrm{align}}$ enforces post-alignment consistency with the reference atlas.

Network architectures are typically based on 3D U-Net backbones (4-level encoder/decoder, channel widths [32, 64, 128, 256], batch normalization, ReLU). Registration heads produce affine parameters for feature-space warping. Notably, 3D architectures capture volumetric context with fewer parameters than slice-wise 2D models ($\sim$5.7M vs. 23.5M parameters) and are critical for robust cross-institution adaptation. Prototype aggregation can be performed across support examples from multiple domains to further enhance invariance (Li et al., 2022, Li et al., 2022).

4. Extensions to Contour-Awareness and Foundation Model Integration

More recent frameworks have evolved registration-assisted prototypical learning by incorporating vision foundation models and advanced anatomical regularization. For example, SAM-assisted registration leverages prompt-driven segmentation masks as auxiliary channels, ensuring anatomical priors are consistently injected throughout training and inference. This is augmented with:

• Prototype Learning: Prototypes $P^k$ for each anatomical region are extracted by masked pooling; prototype contrastive and alignment losses are used to attract voxel features to their class centers and align class centroids across moving/fixed pairs:

$$L_{\mathrm{proto}} = L_{\mathrm{contrast}} + L_{\mathrm{align}}$$

• Contour-aware Loss: Chamfer-based contour loss measures bidirectional distances between predicted and reference mask boundaries, encouraging sharp boundary alignment:

$$L_{\mathrm{contour}} = \frac{1}{|C_m'|}\sum_{i\in C_m'}\min_{j\in C_f} \|i - j\|^2 + \frac{1}{|C_f|}\sum_{j\in C_f}\min_{i\in C_m'} \|j - i\|^2$$

These additional losses demonstrably improve performance in fine-grained, boundary-sensitive registration tasks, particularly with ambiguous boundaries or complex anatomy (Xu et al., 17 Feb 2025).

5. Experimental Protocols and Results

Large-scale multi-institutional datasets are used to evaluate registration-assisted prototypical learning. Protocols split institutions and anatomical classes into base and novel groups, enforcing that novel-class queries are supported only by limited labeled examples (one to five shots) from possibly different institutions.

Key findings from cross-institution male pelvic segmentation (Li et al., 2022, Li et al., 2022):

Method	Avg. Dice (%)	Params (M)	Notes
2D LSNet (baseline)	33.6–44	23.5	Slice-by-slice
3D, no alignment	34.7	5.7	Volumetric, no registration
3D + seg-head only	38.0	5.7	Segmentation supervision
3D + seg-head + align (proposed)	41.3–55	5.7	Full model, p < 0.01 vs. 2D

When query and support are from different institutions, the absence of alignment introduces a $\sim12\%$ Dice drop, while proper alignment reduces the gap to less than 3%. Ablation studies confirm positive, additive contributions from architecture, spatial registration, and mask conditioning modules. Removing affine regularization or prototype branches consistently degradates performance.

Additional studies using anatomical foundation models and contour-aware constraints yield gains of $+4$–$5\%$ Dice over baseline registration methods, with improvements most pronounced in cases with complex geometry or ambiguous anatomical boundaries (Xu et al., 17 Feb 2025).

6. Applications, Limitations, and Prospects

Registration-assisted prototypical learning addresses two central limitations in clinical deployment of deep models: annotation scarcity and cross-domain generalization. Its benefits include explicit anatomical prior injection, reduced inter-institution variability, and parameter efficiency via 3D representations.

Limitations are present: the dependence on affine registration restricts the framework's flexibility versus fully deformable models; availability of suitable multi-class atlases may confine domain applicability. Non-rigid extensions (e.g., B-splines), multi-modal and multi-contrast input streams, and hierarchical or dynamic prototype formulations are actively investigated. Integration of self-supervised cycles between registration and segmentation is a future direction, as is meta-learning atlas priors across diverse institutions (Li et al., 2022, Li et al., 2022, Xu et al., 17 Feb 2025).

A plausible implication is that advances in vision foundation model integration and explicit anatomical constraints will further bridge the domain gap for few-shot and weakly supervised medical image analysis.