AdamFlow: Adam-based Wasserstein Gradient Flows for Surface Registration in Medical Imaging

Abstract: Surface registration plays an important role for anatomical shape analysis in medical imaging. Existing surface registration methods often face a trade-off between efficiency and robustness. Local point matching methods are computationally efficient, but vulnerable to noise and initialisation. Methods designed for global point set alignment tend to incur a high computational cost. To address the challenge, here we present a fast surface registration method, which formulates surface meshes as probability measures and surface registration as a distributional optimisation problem. The discrepancy between two meshes is measured using an efficient sliced Wasserstein distance with log-linear computational complexity. We propose a novel optimisation method, AdamFlow, which generalises the well-known Adam optimisation method from the Euclidean space to the probability space for minimising the sliced Wasserstein distance. We theoretically analyse the asymptotic convergence of AdamFlow and empirically demonstrate its superior performance in both affine and non-rigid surface registration across various anatomical structures.

• The paper introduces a novel AdamFlow algorithm that extends the Adam optimizer to the Wasserstein space for robust surface registration in medical imaging.
• The paper leverages a hybrid optimization strategy combining sliced Wasserstein distance and Chamfer metrics, achieving state-of-the-art accuracy on liver, pancreas, and left ventricle datasets.
• The paper establishes convergence guarantees and computational efficiency with O(N log N) complexity, enabling scalable registration of high-resolution surface meshes.

AdamFlow: Adam-based Wasserstein Gradient Flows for Surface Registration in Medical Imaging

Problem Formulation and Theoretical Framework

The paper introduces a distributional optimisation framework for surface registration in medical imaging by modelling 3D surface meshes as probability measures and framing registration as the minimisation of a discrepancy functional between these probability measures. The core contribution is the utilisation of the sliced Wasserstein distance (SWD) as a discrepancy metric and the development of the AdamFlow algorithm, which generalises Adam—an adaptive, momentum-based optimiser—from the Euclidean to the space of probability measures.

Instead of traditional closest-point-based registration algorithms (e.g., ICP), which are sensitive to initialisation, susceptible to local optima, and computationally inefficient for high-resolution meshes, AdamFlow operates in the Wasserstein space by leveraging functional gradients, particularly the Wasserstein gradient. For non-rigid registration, the approach introduces a composite objective combining SWD for coarse, global alignment with the Chamfer distance for fine, local refinement, under Laplacian regularisation to preserve mesh quality.

Figure 1: Illustration of affine and non-rigid surface registration as distributional optimisation between source and target surfaces, performed by integrating a Wasserstein gradient flow to optimise an objective functional.

AdamFlow is rigorously formulated as a PDE (over the space of probability measures and coupled momenta) whose marginal on the measure space asymptotically converges to the set of critical points of the registration objective, even for non-convex functionals. The authors formally prove this convergence property (under the condition $4\alpha-\beta<3$ for the AdamFlow moment hyperparameters).

Algorithmic Contributions

Central to AdamFlow is the extension of the standard Adam update rule—reliant on adaptive first- and second-moment estimates of the gradient—to operate on functional gradients in the Wasserstein space rather than Euclidean gradients. Under a particle-based approximation, AdamFlow updates particle locations (representing mesh points) and moment estimates, integrating simultaneously affine and non-rigid transformations depending on the optimisation stage.

The registration of surface meshes, both affine and non-rigid, is computationally efficient due to the O($N\log N$) complexity of SWD, where $N$ is the number of mesh vertices. The coarse-to-fine scheme allows for robust, large-scale deformations to be captured during the SWD phase and precise fitting during the Chamfer phase. The particle-based implementation (using Euler integration) is parallelisable, further improving scalability.

Empirical Results

Affine Registration

AdamFlow demonstrates clear, consistent improvements in accuracy and efficiency compared to classical registration techniques such as ICP and momentum-based Wasserstein gradient flows (WGF, HBF, Nesterov). On liver, pancreas, and left ventricle datasets, the AdamFlow method achieves superior average symmetric surface distance (ASSD) and 90th percentile Hausdorff distance (HD90) with statistical significance ($p<0.05$), while requiring only a marginal computational overhead compared to simpler optimisers.

Figure 2: Comparative results for affine surface registration, showing that AdamFlow achieves decreased registration errors relative to alternatives.

Figure 3: Qualitative comparisons for affine registration of various organs; AdamFlow consistently provides more accurate surface alignment than ICP.

Non-Rigid Registration

The hybrid objective for non-rigid registration yields substantial gains over both SWD-only and Chamfer-only objectives, with AdamFlow again delivering the most rapid and stable convergence across the organ datasets. The analysis includes a sensitivity study on regularisation weights and the number of SWD projections, indicating robust hyperparameter regimes. Coarse SWD alignment enables the subsequent Chamfer phase to avoid poor local minima associated with classical nearest-neighbor matching.

Figure 4: Qualitative coarse-to-fine non-rigid registration via AdamFlow, with robust global and accurate local correspondence.

Strong numerical results demonstrate the method’s practical significance:

Organ	Objective	Method	ASSD (mm) ↓	HD90 (mm) ↓
Liver	Hybrid	AdamFlow	0.720	1.264
Pancreas	Hybrid	AdamFlow	0.415	0.742
Left Ventricle	Hybrid	AdamFlow	0.402	0.658

AdamFlow achieves state-of-the-art non-rigid registration accuracy and runtime across all tested datasets.

Computational Complexity and Efficiency

A key strength is the runtime profile: SWD (with just 4 Monte Carlo projections) enables pairwise distances between large meshes to be computed in millisecond-scale time, dominating ICP and Chamfer in both efficiency and scaling. In full registration, surface alignment for large-scale meshes (up to 50,000 vertices) is achievable within seconds on commodity hardware.

Figure 5: Measured runtime for discrepancy metrics as a function of point cloud size; SWD with modest projection count achieves lowest compute times.

Hyperparameter Analysis

Empirical studies examine the effects of the number of SWD projections, mesh Laplacian regularisation weights, and AdamFlow’s learning rate and momentum coefficients. The main conclusions are:

• Increasing the number of projections improves alignment accuracy up to a modest value ($L=8$), beyond which marginal gains diminish relative to cost.
• Regularisation is essential for mesh quality (preserving smoothness and avoiding topological defects), but once above a threshold, further increases have little effect.
• AdamFlow’s convergence and registration accuracy are stable to typical variations in learning rate and momentum parameters, confirming robustness.

  Figure 6: Effect of the number of Monte Carlo projections $L$ on AdamFlow registration accuracy; modest $L$ suffices for strong results.

  

  Figure 7: Influence of mesh Laplacian regularisation weight on accuracy and mesh quality.

  

  Figure 8: Qualitative effect of different regularisation weights on non-rigid surface registration, demonstrating the necessity of Laplacian terms for plausible deformations.

  

Figure 9: Impact of learning rate on AdamFlow’s affine registration; optimal $\eta$ yields stable, accelerated convergence.

Theoretical and Practical Implications

Theoretically, AdamFlow substantiates the extension of adaptive, moment-based optimisation to the Wasserstein space, providing proof of convergence to critical points for arbitrary (potentially non-convex) functionals—a property not previously established for practical surface registration settings. This bridges the gap between theory-driven optimal transport frameworks and high-performance geometric optimisation as demanded by medical imaging applications.

Practically, the proposed algorithm drastically reduces sensitivity to initialisation and enables robust, scalable mesh registration for population-level or high-resolution anatomical studies, including application domains such as statistical shape modelling, inter/intra-subject variability analysis, and longitudinal studies.

Potential future directions include incorporating richer mesh features (normals, curvature), designing diffeomorphic or topology-aware regularisation terms to prevent self-intersection, and further characterisation of the global convergence rate. The methodology also lays groundwork for extending the flow-based optimisation approach to other structured data domains.

Conclusion

The paper presents a rigorous and empirically validated framework, AdamFlow, that marries adaptive, momentum-based optimisation with distributional geometry for efficient, accurate surface registration in medical imaging. By generalising Adam to the Wasserstein space and coupling SWD and Chamfer objectives, AdamFlow achieves superior performance compared to traditional and existing Wasserstein flow methods. The formal convergence analysis and systematic empirical evaluation collectively establish a new standard approach for mesh-based anatomical correspondence problems.

AdamFlow’s open-source implementation further enables direct adoption for future work in medical image analysis, shape modelling, and geometric learning.