End-to-End Learnable Registration

• End-to-end learnable registration is a unified paradigm that integrates feature extraction, correspondence estimation, and transformation optimization into a differentiable framework.
• It eliminates the need for traditional post-processing, enabling joint training of rigid and deformable alignment modules under supervised and unsupervised paradigms.
• This approach improves registration accuracy, robustness, and computational efficiency across modalities such as medical imaging and 2D/3D matching.

End-to-end learnable registration refers to a class of methods in which all components of the registration pipeline—from data-driven correspondence estimation to transformation parameterization and optimization—are jointly trained by backpropagating a task-specific loss directly from the registered output. By structuring registration as a fully differentiable system, these frameworks discard the need for non-differentiable refinement steps (e.g., external ICP, RANSAC, classical numerical solvers) and enable the network to optimize representations, correspondences, and solvers simultaneously for the ultimate geometric or task-driven objective, such as alignment accuracy or downstream classification. End-to-end learnable registration encompasses both rigid and non-rigid (deformable) transformations, spans medical imaging, 2D/3D matching, functional alignment, and dynamic data, and supports both supervised and unsupervised paradigms. This article surveys foundational methodologies, mathematical formulations, representative architectures, empirical benchmarks, and key theoretical insights in this rapidly evolving domain.

1. Fundamental Principles and Taxonomy

End-to-end learnable registration unifies the learning of representations (e.g., features, segmentations), correspondence estimation (e.g., similarity matrices, assignment rules), weighting or attention mechanisms, and transformation solvers (e.g., rigid, affine, or deformable) within a single computational graph. The common criteria are:

• Full differentiability: Every component—including classical solvers (e.g., SVD, Procrustes, point-to-plane systems), spatial transformers, attention modules, or EM steps—must support gradient backpropagation to enable parameter updates based solely on the alignment loss, ideally from raw data.
• Joint optimization: The pipeline is typically optimized end-to-end using stochastic gradient descent or variants, with loss functions directly reflecting geometric alignment (e.g., $L_2$ discrepancy, Dice, SSIM), or indirectly through downstream tasks (e.g., classification, synthesis).
• Unified training vs. post-refinement: Classical methods often append a non-learned, hand-engineered refinement at inference; end-to-end systems are designed to avoid this, instead recursively mimicking optimization steps or embedding them as "known operators" (Jaganathan et al., 2021).

Taxonomically, systems are divided across:

• Transformation class: Rigid (SE(3)), affine, non-parametric/deformable, diffeomorphic, or piecewise models.
• Data type: Images, surfaces, point clouds, or functional/temporal data.
• Supervision: Supervised (with ground-truth warps or landmarks), weakly supervised (with task losses), unsupervised (image similarity or cycle consistency).

2. Core Architectures: From Known Operators to Neural Solvers

Canonical end-to-end registration architectures fuse three components:

A. Representation learning

• Features (e.g., CNNs/U-Nets for images (Vos et al., 2017, Shan et al., 2017, Han et al., 2018), sparse convolutions/transformers for point clouds (Qiao et al., 2020, Gojcic et al., 2020, Jin et al., 2023, Zhang et al., 2021)), or, for temporal data, 1D-CNNs and elastic warping (Jiang et al., 30 Jan 2025).

B. Correspondence or alignment module

• Soft or hard matching matrices, attention, or correspondence prediction via neural networks or combinatorial Sinkhorn/Hungarian assignments (Zhang et al., 2021, Qiao et al., 2020).
• Recurrent modules (e.g., GRUs or iterative RAFT blocks) to mimic classical iterative refinement for rigid or non-rigid alignment (Jaganathan et al., 2021, Feng et al., 2020).
• Unrolled optical-flow modules (e.g., RAFT (Jaganathan et al., 2021)), pixel or contour-wise flow regression for 2D/3D.

C. Geometric solver / transformation head

• Differentiable closed-form solvers for rigid transforms (Kabsch via SVD (Lu et al., 2019), weighted Procrustes (Zhang et al., 2021)), or linearized point-to-plane solvers as known operators (Jaganathan et al., 2021).
• Non-rigid: spatial transformers with B-spline (Vos et al., 2017), momentum-parametrized stationary velocity fields (Shen et al., 2019), or diffeomorphic flows with scaling-and-squaring (Chen et al., 2024, Jena et al., 2024).
• Adjoint-based or implicit differentiation layers, e.g., optimizers treated as differentiable layers (Jena et al., 2024).

Because all modules are differentiable, the chain rule propagates gradients from final loss through features, matches, and solver back to all trainable parameters.

3. Mathematical Formulations and Losses

The core mathematical challenges are encoding the geometric objective as a differentiable function and ensuring that the argmin or solution step is compatible with end-to-end learning.

• Rigid registration (point clouds or images): The minimization

$$\min_{R \in SO(3), t \in \mathbb{R}^3} \sum_i w_i \| R x_i + t - y_i \|^2$$

is solved by differentiable weighted SVD (Kabsch) (Lu et al., 2019) or—augmented for ambiguity and hard assignment—by soft-to-hard permutation learning (Zhang et al., 2021).

• Deformable/nonparametric: The network predicts a dense displacement field $u(x)$ (free-form or parametrized with splines, momenta, or velocity fields), composing a transformation $\varphi(x) = x + u(x)$ or $\varphi = \exp(v(x))$ (scaling-and-squaring) (Vos et al., 2017, Shen et al., 2019, Jena et al., 2024, Chen et al., 2024). The overall loss combines data (image) term and spatial regularization:

$$\mathcal{L} = \mathcal{D}(f, m \circ \varphi) + R[u]$$

often at multiple scales.

• Joint optimization over representations, correspondences, and transformations: The critical insight is that stochastic gradient descent minimizes the total expected registration error by adapting both feature extractors and geometric solvers, as in (Lawin et al., 2020, Han et al., 2018, Jaganathan et al., 2021).

4. Exemplary Methods and Modalities

Representative System	Domain/Type	Key Innovations/Modules
Deep Iterative 2D/3D Registration (Jaganathan et al., 2021)	2D/3D medical rigid	Unrolled RAFT+PointNet+++PPC
DIRNet (Vos et al., 2017)	2D/3D deformable images	FCN, B-spline STN, unsup. loss
TM-SPR (Chen et al., 2024)	Diffeomorphic image registration	Learnable spatial regularizer
Recurrent MV Alignment (Feng et al., 2020)	Nonrigid 3D surfaces/point clouds	GRU blend-skin, multi-view loss
3D Point Cloud S2H (Zhang et al., 2021)	Rigid point cloud	Soft-to-hard PPM, Hungarian
Q-REG (Jin et al., 2023)	Point cloud, curvature-aware	Single-correspondence solver
AVSM (Shen et al., 2019)	3D medical images	Joint affine + vSVF
DeepICP (Lu et al., 2019)	Rigid 3D point cloud	Weighted keypoint, SVD
End2Reg (Pettinari et al., 15 Dec 2025)	RGB-D, registration+segmentation	Gumbel-Softmax STE mask pass
DeepFRC (Jiang et al., 30 Jan 2025)	Functional data	Elastic warping + classifier
DIO (Jena et al., 2024)	Image registration	Implicit diff. optimizer

These systems collectively cover rigid, affine, non-parametric, diffeomorphic, multimodal, non-rigid, and hybrid registration tasks in 2D/3D/temporal/fMRI/point-based and multi-modal pipelines.

5. Benchmarks, Ablations, and Empirical Gains

Empirical results consistently demonstrate the following traits of end-to-end pipelines, relative to classical or staged approaches:

• Superior accuracy and robustness: End-to-end models achieve state-of-the-art matching recall (e.g., Q-REG (Jin et al., 2023)), registration error (e.g., End2Reg (Pettinari et al., 15 Dec 2025)), and empirical Dice/SSIM for dense alignment (Vos et al., 2017, Chen et al., 2024).
• Improved generalization: Learning registration-aware features (e.g., via implicit optimization, (Jena et al., 2024)) increases robustness to large initial misalignments and domain shift (Chen et al., 2024, Jena et al., 2024).
• Computational efficiency: Non-iterative or learned iterative update blocks produce an order-of-magnitude speedup over CPU-based or staged optimization toolkits (Jaganathan et al., 2021, Vos et al., 2017, Shan et al., 2017).
• Ablative analysis: End-to-end vs. staged learning (e.g., two-step segmentation-then-registration) typically yields 10–30% lower error and reduced outlier rates (Pettinari et al., 15 Dec 2025).
• Plug-and-play extensibility: Architectures such as DIO (Jena et al., 2024) and RLL (Lawin et al., 2020) permit swapping the registration solver at inference without retraining features, or generalization to unseen transformation models.

6. Theoretical Analysis and Interpretability

Several works formalize the advantage of end-to-end training theoretically:

• Gradient propagation and optimization bias: Differentiating through the geometric solver (e.g., SVD, LM, implicit optimizer) enables the feature network to learn invariances induced by the transformation class and regularizer, biasing representations toward globally optimal alignment (Han et al., 2018, Jena et al., 2024).
• Elimination of virtual-point collapse: Rigid registration with soft assignments admits ambiguous solutions, which are resolved by hard assignment mechanisms (e.g., partial permutation matrix in S2H (Zhang et al., 2021)).
• Label-aware or region-specific regularization: End-to-end learning of spatially varying deformation weights $\lambda(p)$ supports interpretability—$\lambda$-maps reveal which anatomical regions are tightly or loosely constrained (Chen et al., 2024).
• Task-driven segmentation and selection: End-to-end joint optimization allows segmentation decoders (e.g., Gumbel-Softmax+STE in End2Reg (Pettinari et al., 15 Dec 2025)) to learn points or regions that are optimal for geometric registration, even surpassing weakly supervised mask priors.

7. Limitations, Open Questions, and Future Prospects

Despite marked empirical and theoretical progress, several open challenges exist:

• Dependence on annotated data: High-accuracy methods often require large paired (or weakly paired) datasets (e.g., 2D/3D with ground-truth or weak association (Jaganathan et al., 2021)).
• Generalization under severe appearance drift: While domain randomization helps, clinical and field deployments may require further adversarial or uncertainty-aware strategies (Jaganathan et al., 2021, Chen et al., 2024).
• Scalability to large point clouds or dense volumes: Efficient approximations (e.g., sparse convolutions (Gojcic et al., 2020), low-res mask heads (Chen et al., 2024)) are required for tractable training.
• Interpretability and model selection: Direct mapping between learned features or regularization maps and anatomical priors remains an active area for explainable AI in registration (Chen et al., 2024).
• Extensions to multi-modal, temporal, and multi-objective settings: End-to-end optimization for synthetic-to-real, real-time, 4D dynamics, or multi-task (e.g., synthesis+registration) workflows is increasingly feasible (Tschuchnig et al., 8 Jul 2025, Yiasemis et al., 2024).

In summary, end-to-end learnable registration constitutes a foundational paradigm shift in registration methodology. By fusing differentiable geometric solvers with adaptive, task-driven feature learning, these systems deliver unprecedented registration speed, accuracy, and adaptivity in both rigid and deformable settings. Tangible advances in interpretability, robustness, and task-specificity continue to spur rapid evolution of the field across modality, supervision, and application boundaries.