CoRe: Joint Optimization with Contrastive Learning for Medical Image Registration

1. Problem Definition

Medical image registration aims to establish dense anatomical correspondences between images acquired at different times, from different subjects, or via different modalities. The process faces two primary challenges:

• Appearance Variability: Differences in acquisition protocols, scanner hardware, and modalities cause intensity variations, making direct pixel-wise comparisons unreliable.
• Nonlinear Deformations: Anatomical structures undergo complex, non-rigid deformations due to longitudinal changes, inter-subject variability, pathology, or surgery.

Existing methods often struggle to balance robustness to appearance changes with the ability to model complex deformations. Traditional methods rely on hand-crafted features or mutual information (sensitive to noise), while deep learning approaches often require expensive ground-truth annotations or separate pre-training stages that do not fully align feature learning with the downstream registration task.

2. Methodology: The CoRe Framework

The authors propose CoRe (Contrastive learning for medical image Registration), a novel hybrid framework that integrates self-supervised contrastive learning directly into the deformable image registration pipeline via joint optimization.

Core Architecture

The framework consists of three main components:

1. Feature Extractor ($G$): A convolutional neural network (CNN) that extracts feature representations from fixed ($I_f$) and moving ($I_m$) images.
2. Projection Head ($P$): A convolutional block that projects features into an embedding space.
3. Differentiable Optimization Module ($H$): A module that infers the displacement field ($u$) by computing feature correlations over a discretized search window, minimizing a combined objective of smoothness and feature dissimilarity.

Joint Optimization Strategy

Unlike previous approaches that pre-train feature extractors separately or freeze them during registration, CoRe optimizes the feature extractor $G$ simultaneously using two objectives:

1. Registration Loss ($L_{reg}$):

   • Utilizes a self-training scheme with pseudo-labels.
   • The pipeline generates displacement fields, which are refined via iterative instance optimization (forward-backward consistency and double warping) to create pseudo-ground truth.
   • The loss minimizes the Mean Squared Error (MSE) between the predicted displacement field and the refined pseudo-labels.

2. Equivariance Contrastive Loss ($L_c$):

   • Goal: To enforce that feature embeddings for the same anatomical location remain consistent despite geometric deformations.
   • Mechanism: The model applies random affine geometric transformations ($T_g$) to the input images. It enforces equivariance by ensuring that the feature representation of a transformed image $G(T_g(I))$ is consistent with the transformed feature representation of the original image $T_g(G(I))$.
   • Implementation: Uses an InfoNCE loss to maximize the similarity between positive pairs (features from the same anatomical location under different geometric views) and minimize similarity with negative pairs.
   • Key Distinction: The paper specifically focuses on geometric equivariance rather than intensity invariance, arguing that for mono-modal CT data, geometric robustness is more critical than intensity robustness.

Total Loss Function:
$$L = L_{reg} + \alpha \cdot L_c$$
Where $\alpha$ balances the registration and contrastive objectives.

3. Key Contributions

• Unified Joint Optimization: CoRe is the first to integrate self-supervised contrastive learning directly into the registration pipeline, eliminating the need for separate pre-training stages or paired multimodal ground-truth data.
• Equivariance Constraints: The introduction of geometric equivariance constraints ensures that learned features are robust to tissue deformations, a critical requirement for accurate registration.
• Synergy Demonstration: The authors demonstrate that simultaneous optimization of contrastive and registration objectives yields superior results compared to training with either objective alone, proving that representation learning and registration tasks are synergistic.
• State-of-the-Art Performance: The method achieves top-tier performance on both intra-patient (thoracic) and inter-patient (abdominal) datasets without requiring manual segmentation labels for training.

4. Experimental Results

The method was evaluated on two datasets:

• AbdomenCT: Inter-patient registration (30 scans, 13 anatomical structures).
• RadChestCT: Intra-patient longitudinal registration (371 pairs, 22 anatomical structures).

Quantitative Findings:

• Dice Similarity Coefficient (DSC): CoRe achieved the highest DSC scores on both datasets.
  • AbdomenCT: 52.1% (vs. 51.1% for RegCyc and 49.2% for SAMConvex).
  • RadChestCT: 89.4% (vs. 88.5% for DEEDs and 88.1% for SAMConvex).
• Statistical Significance: The improvements were confirmed via Wilcoxon signed-rank tests ($p < 0.0001$).
• Smoothness: The standard deviation of the log Jacobian determinant (SDLogJ), a measure of deformation field plausibility, remained low and comparable to other hybrid methods, indicating physically plausible transformations.
• Ablation Studies:
  • Joint optimization outperformed separate training of registration and contrastive losses.
  • Geometric Equivariance vs. Intensity Invariance: Interestingly, adding non-linear intensity augmentations (common in contrastive learning) did not improve performance on the mono-modal CT datasets and sometimes degraded it. The authors hypothesize that standard CT intensities limit the utility of intensity augmentations, whereas geometric equivariance is the dominant factor for robustness.

5. Significance and Conclusion

The CoRe framework represents a significant advancement in medical image registration by bridging the gap between self-supervised representation learning and classical optimization-based registration.

• Efficiency: It removes the dependency on costly, manually aligned training data or separate pre-training phases.
• Robustness: By explicitly training the feature extractor to be equivariant to geometric deformations, the model learns anatomical features that are inherently stable under the complex tissue shifts found in clinical scenarios.
• Generalizability: The approach is effective across both challenging inter-patient (abdominal) and intra-patient (thoracic) settings, outperforming conventional, pure deep learning, and existing hybrid methods.

The study concludes that integrating contrastive learning directly into the registration objective creates a powerful synergy, enabling the extraction of dense, semantically meaningful features that drive more accurate and robust image alignment.