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

arXiv: https://arxiv.org/abs/2604.02290 · HTML: https://mlsanity.com/paper/2604.02290/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 25 of its 64 references that are in the map.
- First posted on arXiv 2 Apr 2026.

## Builds on (most cited first, 25 of 25)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [A Point Set Generation Network for 3D Object Reconstruction from a Single Image](https://mlsanity.com/paper/1612.00603/) (2016)
- [On the Convergence of Adam and Beyond](https://mlsanity.com/paper/1904.09237/) (2019)
- [Accelerating 3D Deep Learning with PyTorch3D](https://mlsanity.com/paper/2007.08501/) (2020)
- [On the Variance of the Adaptive Learning Rate and Beyond](https://mlsanity.com/paper/1908.03265/) (2019)
- [TotalSegmentator: robust segmentation of 104 anatomical structures in CT images](https://mlsanity.com/paper/2208.05868/) (2022)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Deep Closest Point: Learning Representations for Point Cloud Registration](https://mlsanity.com/paper/1905.03304/) (2019)
- [PointNetLK: Robust & Efficient Point Cloud Registration using PointNet](https://mlsanity.com/paper/1903.05711/) (2019)
- [AbdomenCT-1K: Is Abdominal Organ Segmentation A Solved Problem?](https://mlsanity.com/paper/2010.14808/) (2020)
- [Generalized Sliced Wasserstein Distances](https://mlsanity.com/paper/1902.00434/) (2019)
- [DeepICP: An End-to-End Deep Neural Network for 3D Point Cloud Registration](https://mlsanity.com/paper/1905.04153/) (2019)
- [ImageCAS: A Large-Scale Dataset and Benchmark for Coronary Artery Segmentation based on Computed Tomography Angiography Images](https://mlsanity.com/paper/2211.01607/) (2022)
- [Convergence and Dynamical Behavior of the ADAM Algorithm for Non-Convex Stochastic Optimization](https://mlsanity.com/paper/1810.02263/) (2018)
- [Point-set Distances for Learning Representations of 3D Point Clouds](https://mlsanity.com/paper/2102.04014/) (2021)
- [Topology-Preserving Shape Reconstruction and Registration via Neural Diffeomorphic Flow](https://mlsanity.com/paper/2203.08652/) (2022)
- [DeepCSR: A 3D Deep Learning Approach for Cortical Surface Reconstruction](https://mlsanity.com/paper/2010.11423/) (2020)
- [The varifold representation of non-oriented shapes for diffeomorphic registration](https://mlsanity.com/paper/1304.6108/) (2013)
- [Adam-family Methods for Nonsmooth Optimization with Convergence Guarantees](https://mlsanity.com/paper/2305.03938/) (2023)
- [Integrating Efficient Optimal Transport and Functional Maps For Unsupervised Shape Correspondence Learning](https://mlsanity.com/paper/2403.01781/) (2024)
- [S3M: Scalable Statistical Shape Modeling through Unsupervised Correspondences](https://mlsanity.com/paper/2304.07515/) (2023)
- [Accelerating optimization over the space of probability measures](https://mlsanity.com/paper/2310.04006/) (2023)
- [Diffeomorphic Mesh Deformation via Efficient Optimal Transport for Cortical Surface Reconstruction](https://mlsanity.com/paper/2305.17555/) (2023)
- [Long-Time Asymptotics of the Sliced-Wasserstein Flow](https://mlsanity.com/paper/2405.06313/) (2024)

## Abstract (arXiv)

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.