NeurEPDiff: Neural Operators to Predict Geodesics in Deformation Spaces

TL;DR

NeurEPDiff introduces a neural operator that efficiently predicts geodesic deformations in high-dimensional spaces, significantly reducing computational costs and enabling resolution-invariant performance for image registration tasks.

Contribution

The paper develops NeurEPDiff, a neural operator that learns the nonlinear mapping of velocity fields for rapid geodesic prediction, a novel approach in diffeomorphic registration.

Findings

NeurEPDiff achieves high registration accuracy on 2D and 3D datasets.

It significantly reduces computation time compared to traditional methods.

The model generalizes well across different image resolutions.

Abstract

This paper presents NeurEPDiff, a novel network to fast predict the geodesics in deformation spaces generated by a well known Euler-Poincar\'e differential equation (EPDiff). To achieve this, we develop a neural operator that for the first time learns the evolving trajectory of geodesic deformations parameterized in the tangent space of diffeomorphisms(a.k.a velocity fields). In contrast to previous methods that purely fit the training images, our proposed NeurEPDiff learns a nonlinear mapping function between the time-dependent velocity fields. A composition of integral operators and smooth activation functions is formulated in each layer of NeurEPDiff to effectively approximate such mappings. The fact that NeurEPDiff is able to rapidly provide the numerical solution of EPDiff (given any initial condition) results in a significantly reduced computational cost of geodesic shooting of diffeomorphisms in a high-dimensional image space. Additionally, the properties of discretiztion/resolution-invariant of NeurEPDiff make its performance generalizable to multiple image resolutions after being trained offline. We demonstrate the effectiveness of NeurEPDiff in registering two image datasets: 2D synthetic data and 3D brain resonance imaging (MRI). The registration accuracy and computational efficiency are compared with the state-of-the-art diffeomophic registration algorithms with geodesic shooting.