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

Qiang Ma; Qingjie Meng; Xin Hu; Yicheng Wu; Wenjia Bai

cs.CVApr 2, 2026

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

Authors: Qiang Ma, Qingjie Meng, Xin Hu, Yicheng Wu, Wenjia Bai

Organizations: Department of Brain Sciences, Imperial College London, London, UK · Department of Computing, Imperial College London, London, UK · School of Computer Science, University of Birmingham, Birmingham, UK · Department of Computer Science, Columbia University, New York, NY, USA ·

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.

Figures & tables

Objective	$F[\mu]$	$\nabla_{W}F[\mu]$
Potential energy	$\int_{\mathbb{R}^{d}}V(x)\mathrm{d}\mu(x)$	$\nabla V(x)$
Internal energy	$\int_{\mathbb{R}^{d}}U(\mu(x))\mathrm{d}x$	$\nabla U^{\prime}(\mu(x))$
Negative entropy	$\int_{\mathbb{R}^{d}}\mu(x)\log\mu(x)\mathrm{d}x$	$\nabla\log\mu(x)$
KL divergence	$KL(\mu\|\|\mu_{*})$	$\nabla\log\frac{\mu(x)}{\mu_{*}(x)}$
2-Wasserstein distance	$\frac{1}{2}W_{2}^{2}(\mu,\mu_{*})$	$x-\mathcal{T}(x)$

Table 1: Commonly used objective functionals and their Wasserstein gradients. For 2-Wasserstein distance,

\mathcal{T}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}

denotes the optimal transport map from

\mu

\mu_{*}

such that the pushforward measure

\mathcal{T}\sharp\mu=\mu_{*}

Figure 1: Illustration of affine surface registration of the left ventricle of the heart and non-rigid surface registration of the liver. The surface registration is formulated as a distributional optimisation problem between a source surface ( $t=0$ ) and a target surface. The registration is performed by integrating a Wasserstein gradient flow for $t\in[0,T]$ to optimise an objective functional.

Dataset	Organ	#Meshes	#Vertices
AbdomenCT-1K	Liver	360	$34,553\pm 6,229$
AbdomenCT-1K	Pancreas	360	$6,194\pm 1,889$
ImageCAS	Left ventricle	500	$10,420\pm 1,551$

Table 2: Dataset information for surface registration. The number of meshes and the average numbers of vertices for the meshes of each anatomical structure are reported.

Panel of Figure 2 — Figure 2: Comparative results and convergence curves for affine and non-rigid surface registration. The registration errors are measured by ASSD (mm) and HD90 (mm). Registration errors decrease with increasing numbers of optimisation iterations.

	Affine	Non-rigid
Organ	$K_{\mathrm{affine}}$	$K_{\mathrm{sw}}$	$K_{\mathrm{cham}}$
Liver	$1,500$	$500$	$200$
Pancreas	$3,000$	$1,200$	$200$
Left ventricle	$1,500$	$100$	$100$

Table 3: The integration steps used for affine and non-rigid registration of different anatomical surfaces.

	Affine	Non-rigid
Method	$\eta_{\mathrm{affine}}$	$\eta_{\mathrm{sw}}$	$\eta_{\mathrm{cham}}$
ICP	$10^{-6}$	–	–
WGF	$10^{-5}$	$0.5$	$0.1$
HBF	$10^{-5}$	$0.5$	$0.1$
Nesterov	$10^{-7}$	$0.005$	$0.005$
AdamFlow	$10^{-2}$	$0.5$	$0.1$

Table 4: The learning rates used in different optimisation methods for affine and non-rigid surface registration.

	ASSD (mm) $\downarrow$			HD90 (mm) $\downarrow$
Method	Liver	Pancreas	Left ventricle	Liver	Pancreas	Left ventricle
ICP	$8.751\pm 3.118$	$6.902\pm 3.453$	$2.955\pm 0.956$	$23.83\pm 13.48$	$20.90\pm 15.25$	$6.870\pm 3.382$
WGF	$7.701\pm 1.953$	$5.780\pm 1.724$	$2.401\pm 0.650$	$16.82\pm 4.715$	$13.77\pm 4.951$	$5.016\pm 1.764$
HBF	$7.670\pm 1.947$	$5.757\pm 1.715$	$2.365\pm 0.644$	$16.76\pm 4.710$	$13.71\pm 4.926$	$4.953\pm 1.758$
Nesterov	$7.555\pm 1.926$	$5.511\pm 1.649$	$2.172\pm 0.631$	$16.55\pm 4.691$	$13.07\pm 4.662$	$4.602\pm 1.733$
AdamFlow	$~\mathbf{7.452}\pm\mathbf{1.880}^{*}$	$~\mathbf{5.381}\pm\mathbf{1.451}^{*}$	$~\mathbf{1.932}\pm\mathbf{0.529}^{*}$	$~\mathbf{16.28}\pm\mathbf{4.597}^{*}$	$~\mathbf{12.75}\pm\mathbf{4.289}^{*}$	$~\mathbf{4.059}\pm\mathbf{1.354}^{*}$

Table 5: Comparison of AdamFlow to different optimisation methods for affine surface registration. The registration errors are measured by ASSD (mm) and HD90 (mm). The best results are bolded. The symbol * indicates that AdamFlow achieves significantly better performance than all baseline methods (paired

t

-test,

p<0.05

		ASSD (mm) $\downarrow$			HD90 (mm) $\downarrow$
Objective	Method	Liver	Pancreas	Left ventricle	Liver	Pancreas	Left ventricle
SWD	WGF	$1.528\pm 0.394$	$1.441\pm 0.487$	$0.838\pm 0.172$	$3.369\pm 1.550$	$3.148\pm 1.427$	$1.680\pm 0.471$
Chamfer	WGF	$0.827\pm 0.264$	$0.548\pm 0.211$	$0.666\pm 0.324$	$1.537\pm 1.362$	$1.184\pm 0.959$	$1.600\pm 1.102$
Hybrid	WGF	$0.755\pm 0.192$	$0.419\pm 0.144$	$0.408\pm 0.052$	$1.346\pm 1.101$	$0.763\pm 0.596$	$0.670\pm 0.157$
Hybrid	HBF	$0.749\pm 0.189$	$0.416\pm 0.141$	$0.407\pm 0.051$	$1.337\pm 1.104$	$0.747\pm 0.558$	$0.668\pm 0.152$
Hybrid	Nesterov	$0.923\pm 0.243$	$0.659\pm 0.874$	$0.553\pm 0.091$	$1.848\pm 1.151$	$1.674\pm 4.632$	$0.965\pm 0.280$

Table 6: Comparison of different objective functionals and optimisation methods for non-rigid surface registration. The registration errors are measured by ASSD (mm) and HD90 (mm). The best results are bolded. The symbol * indicates that AdamFlow achieves significantly better performance than all baseline methods (paired

t

-test,

p<0.05

Figure 3: Qualitative comparisons for affine surface registration of the liver, pancreas, and left ventricle from a source mesh to a target mesh. AdamFlow shows consistently better affine registration accuracy than the traditional ICP algorithm.

Figure 4: Qualitative comparisons for non-rigid surface registration of the liver, pancreas, and left ventricle from a source mesh to a target mesh. AdamFlow is used to minimises the SWD for global surface alignment (coarse), and then optimises the Chamfer distance for local mesh refinement (fine), enabling coarse-to-fine surface registration.

Figure 5: Runtime (millisecond) for computing the discrepancy metrics across different numbers of points.

	Method	Liver	Pancreas	Left ventricle
Affine	WGF	$2.47\pm 0.33$	$2.72\pm 0.21$	$1.54\pm 0.06$
	HBF	$2.53\pm 0.34$	$2.82\pm 0.21$	$1.60\pm 0.05$
	Nesterov	$2.53\pm 0.34$	$2.83\pm 0.20$	$1.60\pm 0.06$
	AdamFlow	$2.65\pm 0.33$	$3.07\pm 0.21$	$1.72\pm 0.05$
Non-rigid	WGF	$2.52\pm 0.43$	$1.40\pm 0.15$	$0.41\pm 0.04$
	HBF	$2.53\pm 0.43$	$1.43\pm 0.15$	$0.41\pm 0.04$

Table 7: Runtime (second) of different optimisation methods for affine and non-rigid surface registration of three organs.

Figure 6: Effect of different numbers of Monte Carlo projections

L

on affine registration of pancreas surfaces with sliced Wasserstein discrepancy using AdamFlow.

Figure 7: Effect of mesh Laplacian regularisation weights

\lambda_{\mathrm{lap}}

on non-rigid registration of liver surfaces via AdamFlow.

Figure 8: Qualitative comparison for non-rigid registration of liver surfaces with different mesh Laplacian regularisation weights $\lambda_{\mathrm{lap}}$ .

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 9 — Figure 9: Effect of hyperparameters on affine registration of left ventricle surfaces using AdamFlow.

Medical Image Registration

Medical image registration aligns scans across time, patients, or imaging modalities to track anatomy and support diagnosis, treatment, and image-guided procedures.

See all 96 papers

Polaffini: A feature-based approach for robust affine and polyaffine image registration
Feb 19, 2026Antoine Legouhy, Cosimo Campo, Ross Callaghan +2Medical Image Registration Deformable Image Registration
Hawkes Institute & Department of Computer Science, University College London, London, UK · Institut Pasteur, Universit´e Paris Cit´e, Unit´e de Neuroanatomie Appliqu´ee et Th´eorique, France · AINOSTICS ltd., Manchester, UK
GLIDE-Reg: Global-to-Local Deformable Registration Using Co-Optimized Foundation and Handcrafted Features
Feb 27, 2026Yunzheng Zhu, Aichi Chien, Kimaya kulkarni +5Medical Image Registration Deformable Image Registration
Medical & Imaging Informatics, Department of Radiological Sciences, David Geffen School of Medicine at UCLA, USA · Magnetic Resonance Research Labs, Department of Radiological Sciences, David Geffen School of Medicine at UCLA, USA · Department of Radiation Oncology, David Geffen School of Medicine at UCLA, Los Angeles, USA
ACorD: Adaptive Correspondence Discovery for Medical Image Registration
Oct 19, 2026Arghavan Rezvani, Kun Han, Xiaohui XieMedical Image Registration Deformable Image Registration
University of California, Irvine

Feb 19, 2026J/K move · Enter open · S save

Polaffini: A feature-based approach for robust affine and polyaffine image registration

Antoine Legouhy, Cosimo Campo, Ross Callaghan, Hojjat Azadbakht+1

Hawkes Institute & Department of Computer Science, University College London, London, UK · Institut Pasteur, Universit´e Paris Cit´e, Unit´e de Neuroanatomie Appliqu´ee et Th´eorique, France · AINOSTICS ltd., Manchester, UK

Medical Image Registration Deformable Image Registration

Open paper AskPDF arXiv

In this work we present Polaffini, a robust and versatile framework for anatomically grounded registration. Medical image registration is dominated by intensity-based registration methods that rely on surrogate measures of alignment quality. In contrast, feature-based approaches that operate by identifying explicit anatomical correspondences, while more desirable in theory, have largely fallen out of favor due to the challenges of reliably extracting features. However, such challenges are now significantly overcome thanks to recent advances in deep learning, which provide pre-trained segmentation models capable of instantly delivering reliable, fine-grained anatomical delineations. We aim to demonstrate that these advances can be leveraged to create new anatomically-grounded image registration algorithms. To this end, we propose Polaffini, which obtains, from these segmented regions, anatomically grounded feature points with 1-to-1 correspondence in a particularly simple way: extracting their centroids. These enable efficient global and local affine matching via closed-form solutions. Those are used to produce an overall transformation ranging from affine to polyaffine with tunable smoothness. Polyaffine transformations can have many more degrees of freedom than affine ones allowing for finer alignment, and their embedding in the log-Euclidean framework ensures diffeomorphic properties. Polaffini has applications both for standalone registration and as pre-alignment for subsequent non-linear registration, and we evaluate it against popular intensity-based registration techniques. Results demonstrate that Polaffini outperforms competing methods in terms of structural alignment and provides improved initialisation for downstream non-linear registration. Polaffini is fast, robust, and accurate, making it particularly well-suited for integration into medical image processing pipelines.

Abstract

Medical Image Registration

Explore similar work

Polaffini: A feature-based approach for robust affine and polyaffine image registration

GLIDE-Reg: Global-to-Local Deformable Registration Using Co-Optimized Foundation and Handcrafted Features

ACorD: Adaptive Correspondence Discovery for Medical Image Registration