# Neural approximation of Procrustes-Wasserstein transport maps

Davide Adamo, Marco Corneli, Manon Vuillien, Emmanuelle Vila

<figure><img src="image_1.png" /></figure>

## To cite this version:

Davide Adamo, Marco Corneli, Manon Vuillien, Emmanuelle Vila. Neural approximation of Procrustes-Wasserstein transport maps. 2026. hal-05483278

HAL Id: hal-05483278

https://hal.science/hal-05483278v1

Preprint submitted on 29 Jan 2026

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_2.png" /></figure>

Distributed under a Creative Commons CC BY 4.0 - Attribution - International License
Neural approximation of Procrustes-Wasserstein transport maps

Davide Adamo<fnref n="1" />,² Marco Corneli<fnref n="1" />,² Manon Vuillien<fnref n="1" /> Emmanuelle Vila³

**Abstract**

Optimal Transport (OT) provides a rigorous mathematical framework for comparing empirical distributions, such as point clouds. However, standard formulations based on the Wasserstein distances are highly sensitive to global rigid transformations in the feature space. Unlike other approaches being computationally prohibitive (e.g. Gromov-Wasserstein) the Procrustes-Wasserstein (PW) distance addresses this limitation in a more scalable way by directly incorporating alignment into the optimization problem. However, current approaches remain confined to discrete settings and lacks generalization to unseen samples. In this work, we extend the PW framework to the continuous domain by formally deriving its dual and semi-dual formulations. Leveraging results in both convex analysis and, more recently, neural optimal transport, we introduce NeuralPW, a framework that jointly learns rigid alignments and continuous transport maps. We empirically demonstrate that NeuralPW accurately recovers geometric transformations, explicitly decoupling global pose alignment from local non-rigid deformation. Finally, we validate our approach on a real-world case study in archaeozoology, demonstrating its ability to robustly match high-density point clouds where discrete primal solvers are computationally intractable.

**1. Introduction**

By lifting the comparison of data from point-wise metrics to a geometric analysis of probability distributions, optimal transport (OT, Villani et al., 2008; Peyré et al., 2019) is attracting more and more interest in modern machine learning. Among the most impacted areas we mention generative modeling (Arjovsky et al., 2017; Gulrajani et al., 2017),

imaging (Solomon et al., 2015; Feydy et al., 2017) and domain adaptation (Courty et al., 2016; Balaji et al., 2020). Central to this success is the Wasserstein distance (Kantorovich, 1942), which leverages the underlying geometry of a common metric space the distributions to be compared live in. While OT has broad utility, the Wasserstein distance has proven central in shape analysis and geometric matching (Eisenberger et al., 2020). In these settings, shapes are naturally represented as empirical measures, such as point clouds, and OT provides a theoretically grounded metric for comparisons, interpolations and deformation analysis. However, in some applications such as shape matching or clustering it is crucial to introduce some forms of invariance with respect to rigid transformations such as rotations or symmetries that, conversely impact the Wasserstein distance (e.g. the distance between a rotated shapes and itself before rotation is not zero).

To address this “alignment gap”, a branch of OT has focused on isometry invariant formulations that disentangle pose from shape. A first approach in this direction is represented by the Gromov-Wasserstein (GW) distance (Mémoli, 2011). GW relaxes the requirement of a shared ambient space by comparing distributions based on their intra-domain relational structures (e.g. pairwise similarity matrices) rather than absolute positions. This flexibility is powerful for matching objects across distinct metric spaces, such as in shape correspondence (Mémoli, 2011; Vayer et al., 2019), word embedding alignment (Alvarez-Melis & Jaakkola, 2018) and graph matching (Vayer et al., 2018; Vincent-Cuaz et al., 2021). Nevertheless, computing standard GW entails solving a non-convex quadratic optimization problem which is notoriously computationally expensive ($O(N^4)$ naively, with $N$ denoting the number of points in the support of the empirical distributions to be compared). While approaches such as sliced (Vayer et al., 2019) and quantized (Chowdhury et al., 2021) versions of GW have improved scalability, GW methods often discard valuable ambient geometric information when distributions do reside in the same metric space.

For the specific but common scenario where distributions lie in the same space ($\mathbb{R}^d$) but differ by an unknown rigid transformation, the **Procrustes-Wasserstein** (PW, Zhang et al., 2017) distance offers a more specialized and efficient alternative. First formalized in the context of unsuper-

<footnote n="1">¹Université Côte d'Azur, UMR 7264 CEPAM, CNRS, Nice, France. ²Université Côte d'Azur, Inria, CNRS, Laboratoire J.A.Dieudonné, Maasai team, Nice, France. ³Université Lumière Lyon II, UMR 5133 Archéorient CNRS, Lyon, France. Correspondence to: Davide Adamo <davide.adamo@univ-cotedazur.fr>.</footnote>
vised word embedding alignment (Zhang et al., 2017; Grave et al., 2019; Alvarez-Melis et al., 2019), the PW distance couples the classical orthogonal Procrustes problem with OT matching. Zhang et al. (2017) formulated a joint optimization over a transport plan $\pi$ and an orthogonal matrix $P \in O(d)$ to simultaneously align the supports and match the masses of probability distributions. Subsequently, Grave et al. (2019) improved stability via convex relaxations over the permutation polytope, while Alvarez-Melis et al. (2019) introduced convexity annealing to eliminate the dependence on ad hoc initialization and generalized the PW framework to broader invariance classes. Recent contributions have further expanded both the application scope and the theoretical understanding of the PW framework. In the domain of cross-lingual learning, Ramírez et al. (2024) demonstrated that many prevailing unsupervised embedding alignment methods (e.g. MUSE and ICP) implicitly optimize a PW-like objective; they proposed a direct solver combining Sinkhorn iterations (Cuturi, 2013) with Procrustes analysis, yielding superior alignment over heuristic baselines. Beyond NLPs, Yang et al. (2022) applied the framework to unsupervised network alignment, introducing HackGAN, which integrates a PW loss within a generative (CycleGAN) architecture to map graph node embeddings across disparate networks. On the theoretical front, Even et al. (2024) established rigorous theoretical bounds for the PW problem, determining the precise sample complexity required to recover the ground-truth alignment in high-dimensional regimes ($d \gg \log N$) using alternating minimization strategies. Most recently, Adamo et al. (2025) formally established the metric properties of the PW distance over a quotient space of rotated/reflected discrete measures and provided an alternating minimization scheme, referred to as ping-pong in Even et al. (2024), for computing PW matching from samples. However, all these existing PW approaches remain confined to the discrete primal setting, meaning that they estimate a static matching matrix (transport plan) between samples and do not yield a continuous transport map $T: \mathbb{R}^d \to \mathbb{R}^d$. Consequently, i) they lack the ability to generalize to unseen samples or to provide a dense deformation field, limiting their insights and ii) their computational cost is soon prohibitive.

**Our contribution.** We extend PW to the continuous framework, introducing a unified method that learns both rigid alignment and non-rigid deformation in a differentiable manner. Our main contributions are: (i) a formal derivation of the dual and semidual formulations of the PW problem and an explicit characterization of the PW transport map (Theorem 3.2); (ii) the introduction of NeuralPW, a framework that leverage neural networks to estimate the PW transport map from the data; (iii) a detailed analysis of the PW problem between centered Gaussian distributions, deriving the explicit closed-form solution for the PW distance and the associated PW map; (iv) we empirically demon-

strate that NeuralPW accurately disentangles global pose differences from local non-rigid geometric variations, overcoming the limitations of Wasserstein neural OT solvers on unaligned data; (v) our approach is validated on a real-world scientific case study from archaeozoology, involving high-resolution 3D point clouds of animal skulls. We show that NeuralPW aligns and matches dense anatomical structures at scales where discrete primal solvers are computationally intractable.

## 2. OT background

**Notation.** In the following we work with probability measures $\mu, \nu \in P_2(\mathbb{R}^d)$. In other words, when saying that $X \sim \mu$ we mean that $X$ follows the probability distribution $\mu$ and is intended to have finite second order moment, e.g. $\mathbb{E}(\|X\|_2^2) < \infty$. With a slight abuse of notation, $L^1(\mu)$ denotes the set of measurable functions $f$ such that $\int f(x)d\mu(x) < \infty$.

### 2.1. Wasserstein distance

Given a pair of probability measures $(\mu, \nu)$, the squared 2-Wasserstein problem is (Kantorovich, 1942)

$$W_2^2(\mu, \nu) := \min_{\pi \in \Pi(\mu, \nu)} \mathbb{E}_{(X,Y) \sim \pi} \|X - Y\|_2^2, \quad (\text{OT})$$

where $\Pi(\mu, \nu)$ is the set of joint distributions (or couplings) $\pi \in P_2(\mathbb{R}^d \times \mathbb{R}^d)$ with marginals $\mu$ and $\nu$. In the above optimization problem, $W_2(\cdot, \cdot)$ defines a distance between $\mu$ and $\nu$ and for an in depth inspection of its properties the reader is referred to (Villani et al., 2009; Santambrogio, 2015; Peyré et al., 2019). Other Wasserstein distances (p-norms in place of the 2 norm in the above equation) as well as entropic optimal transport (Nutz, 2021) are not considered in this work. In case at least one between $\mu$ and $\nu$ has a density with respect to the Lebesgue measure (i.e. $X$ or $Y$ are continuous random variables) there exists a unique transport or Monge map $T: \mathbb{R}^d \to \mathbb{R}^d$ linked with the (unique) optimal $\pi$ from Eq. (OT) by

$$W_2^2(\mu, \nu) = \mathbb{E}_{(X,Y) \sim \pi} \|X - Y\|_2^2 = \mathbb{E}_{X \sim \mu} \|X - T(X)\|_2^2,$$

where $T(X) \sim \nu^1$. Moreover, Eq. (OT) admits the following dual formulation

$$W_2^2(\mu, \nu) = \sup_{(\varphi, \psi) \in \Phi} \{\mathbb{E}_{\mu}[\varphi(X)] + \mathbb{E}_{\nu}[\psi(Y)]\}, \quad (\text{DOT})$$

where $\Phi := \{(\varphi, \psi) \in L^1(\mu) \times L^1(\nu) \mid \forall (x, y) \ \varphi(x) + \psi(y) \le \|x - y\|_2^2\}$ and $\varphi$ and $\psi$ are usually referred to as the Kantorovich potentials. In case the Monge map $T$ exists the above supremum is attained at $(\bar{\varphi}, \bar{\psi})$ and $T$ is related

<footnote n="1">¹This can be written in short-hand notation by means of the push-forward operator $\pi$, i.e. $T_{\pi}\mu = \nu$.</footnote>
to $\bar{\varphi}$ by the following relation

$$T(x) = \nabla (\|x\|_2^2 - 2\bar{\varphi}(x)). \quad (1)$$

The above facts relate to the Brenier's theorem and are proven in (Peyré et al., 2019), Th. 2.1, and (Chewi et al., 2024), Th.1.14. Finally (for our purposes), the dual problem in Eq. (DOT) can be rephrased in a more convenient way, allowing us to make a parallel with convex analysis

$$W_2^2(\mu, \nu) = C_{\mu,\nu} - 2 \inf_{f \in L^1(\mu)} \{\mathbb{E}_{\mu}[f(X)] + \mathbb{E}_{\nu}[f^*(Y)]\}, \quad (2)$$

where $C_{\mu,\nu} := \mathbb{E}_{X \sim \mu} [\|X\|_2^2] + \mathbb{E}_{Y \sim \nu} [\|Y\|_2^2]$, $f(x) := \|x\|_2^2 - 2\varphi(x)$ and $f^*$ is the convex conjugate of $f$ (Appendix A). Note that, by definition of $f$ and from Eq. (1) we have that $T(x) = \nabla \bar{f}(x)$. The Monge map and the semidual formulation of the Wasserstein problem are two key ingredients of this work.

### 2.1.1. LEARNING THE MONGE MAP FROM THE DATA

Given observed realisations $x_1, \dots, x_N$ and $y_1, \dots, y_M$ from $\mu$ and $\nu$, respectively, one might wish to estimate $T$ from the data. This central topic is often treated “separately” in the statistical (e.g. Hütter & Rigollet, 2021; Manole et al., 2024) and deep learning (e.g. Makkuva et al., 2020; Korotin et al., 2022) rooms, although some works seek to make a bridge between the two communities (Divol et al., 2025). In the context of the 2-Wasserstein distance, the common point between most of the above cited works is the empirical counterpart of the semidual formulation of the Wasserstein problem (2), where $\mu$ and $\nu$ are replaced by their empirical cumulative distribution functions. A proper choice of the parametric functional class adopted to estimate the Kantorovich potential was proved to allow one to escape the curse of dimensionality and obtain a rate of convergence at least $O(n^{-1/2})$ (Theorem 3.7, Chewi et al., 2024). In contrast, (Manole et al., 2024) relies on the empirical counterpart of the primal formulation in Eq. (OT) and adopts a 1 nearest neighbor approach to map a new test data point $x$ to the target domain, via the estimated transport plan $\hat{\pi}$. Still based on the semidual formulation in Eq. (2) the works of (Makkuva et al., 2020) and (Korotin et al., 2022) rely on neural networks to learn the Monge map. Makkuva et al. (2020) approximate a reparameterization of the dual potentials $(\varphi, \psi)$ using input convex neural networks (ICNNs, Amos et al., 2017). This approach is discussed in more detail in next sections. More recently, Neural Optimal Transport (NOT, Korotin et al., 2022) rewrites the semidual objective as a saddle-point optimization problem (Corollary 1, Korotin et al., 2022). With the objective of finding a stochastic OT map $T(x, z)$, where $z$ is a noise introduced to create stochasticity, NOT approaches the minimax problem directly parametrizing $T$ and the dual potential $g$ with general neural networks like MLPs. Indeed, sufficiently

large NNs can arbitrarily well approximate stochastic transport maps between distributions (Theorem 1, Korotin et al., 2022).

## 3. The Procrustes-Wasserstein problem

The additional ingredients we now introduce are the orthogonal group $O(d) = \{P \in \mathbb{R}^{d \times d} : P^T P = I_d\}$ and the cost function $C_P(x, y) := \|P x - y\|_2$. The squared 2-Procrustes-Wasserstein cost is defined by

$$PW_2^2(\mu, \nu) := \inf_{\substack{P \in O(d) \\ \pi \in \Pi(\mu, \nu)}} \mathbb{E}_{X,Y \sim \pi} [C_P(X, Y)^2]. \quad (\text{PW})$$

This quantity was recently proven to be an OT distance with respect to a properly defined quotient space, which considers two measures as equivalent if they only differ by a rotation/symmetry of their support (Adamo et al., 2025). The following proposition links $PW_2$ and $W_2$ distances.

**Proposition 3.1.** For any $P \in O(d)$, the Kantorovich transport cost between $P_{\#}\mu$ (the push-forward of $\mu$ by $P$) and $\nu$ is

$$W_2^2(P_{\#}\mu, \nu) = \min_{\pi \in \Pi(\mu, \nu)} \mathbb{E}_{X,Y \sim \pi} [C_P(X, Y)^2],$$

where $\Pi(\mu, \nu)$ is the set of couplings $\pi \in \mathcal{P}(\mathbb{R}^d \times \mathbb{R}^d)$ with marginals $\mu$ and $\nu$.

*Proof.* See Appendix B. □

As an immediate consequence, one has

$$PW_2^2(\mu, \nu) = \inf_{P \in O(d)} W_2^2(P_{\#}\mu, \nu), \quad (3)$$

meaning that solving the PW problem is equivalent to find the optimal orthogonal transformation $P$ that minimizes the standard $W_2^2$ distance between the source $P_{\#}\mu$ and $\nu$. The above formulation is particularly convenient since it immediately allows us to obtain the following dual formulation for (PW)

$$PW_2^2(\mu, \nu) = \inf_{P \in O(d)} \sup_{(\varphi, \psi) \in \Phi_P} \mathcal{F}_{P_{\#}\mu, \nu}(\varphi, \psi), \quad (\text{DPW})$$

where we define

$$\mathcal{F}_{P_{\#}\mu, \nu}(\varphi, \psi) := \mathbb{E}_{U \sim P_{\#}\mu}[\varphi(U)] + \mathbb{E}_{Y \sim \nu}[\psi(Y)],$$

and $\Phi_P := \{(\varphi, \psi) \in L^1(P_{\#}\mu) \times L^1(\nu) \mid \forall (u, y) \in \mathcal{S}(u, y) + \mathcal{S}(y, u) \}$.

### 3.1. Semidual formulation of $PW_2$

We now focus on the dual formulation of $W_2(P_{\#}\mu, \nu)$. The inequality constraint in $\Phi_P$ can be read as

$$\varphi(u) + \psi(y) \le \|u - y\|_2^2$$

$$\varphi(u) + \psi(y) \le \|u\|_2^2 + \|y\|_2^2 - 2\langle u, y \rangle$$

$$f(u) + g(y) \ge \langle u, y \rangle,$$
where we adopted the same reparameterization seen in Eq. (2), namely $(f, g) \in L^1(P_\sharp\mu) \times L^1(\nu)$ and $f(u) = \|u\|_2^2 - 2\varphi$ and $g(y) = \|y\|_2^2 - 2\psi(y)$. Hence, for a fixed $P$, the dual problem now reads

$$W_2^2(P_\sharp\mu, \nu) = C_{\mu,\nu} - 2 \inf_{(f,g) \in \Phi'_P} \mathcal{F}_{P_\sharp\mu,\nu}(f,g)$$

with $\Phi'_P := \{(f,g) \in L^1(P_\sharp\mu) \times L^1(\nu) \mid \forall (u,y) \ f(u) + g(y) \ge \langle u, y \rangle\}$ and $C_{\mu,\nu}$ the same constant introduced previously. We can move one step forward by observing that, for a given, arbitrary $f$ in the above equation, the best we can do to minimize $\mathcal{F}_{P_\sharp\mu,\nu}(f,g)$ is to choose $g$ as small as possible while satisfying the constraint $g(y) \ge \langle u, y \rangle - f(u)$, which is done by setting

$$g(y) = \sup_{u \in \text{Dom}(f)} \{\langle u, y \rangle - f(u)\} =: f^*(y),$$

where $f^*$ is the convex conjugate of $f$. Using the fact that $f^*$ is convex by construction and hence $f^{***} = f^*$ (Appendix A) as well as the symmetry of $\mathcal{F}_{P_\sharp\mu,\nu}(f,g)$ in its arguments, it is not hard to show that the above problem simplifies to the following **semidual** formulation

$$W_2^2(P_\sharp\mu, \nu) = C_{\mu,\nu} - 2 \inf_{f \in L^1(P_\sharp\mu)} \mathcal{F}_{P_\sharp\mu,\nu}(f, f^*), \quad (4)$$

whose solution is given by a closed convex function $\bar{f} \in L^1(P_\sharp\mu)$. As a corollary, it holds that

$$\inf_{f \in L^1(P_\sharp\mu)} \mathcal{F}_{P_\sharp\mu,\nu}(f, f^*) = \inf_{\substack{f \in L^1(P_\sharp\mu) \\ f \text{ convex}}} \mathcal{F}_{P_\sharp\mu,\nu}(f, f^*). \quad (5)$$

Before discussing how to estimate the $\bar{f}$ solving the above minimization problem, we enounce the following

**Theorem 3.2 (Characterization of the PW transport map).**
Let $\mu, \nu \in \mathcal{P}_2(\mathbb{R}^d)$ with finite second order moments and assume $\mu$ is absolutely continuous w.r.t. the Lebesgue measure. Let $\bar{P} \in \mathcal{O}(d)$ be the optimal orthogonal transformation minimizing Eq. (PW). Then the PW transport map $T_{PW} : \mathbb{R}^d \to \mathbb{R}^d$ is given by

$$T_{PW}(x) = \bar{\nabla} \bar{f} \circ \bar{P}(x)$$

where there exists a unique convex function $\bar{f} : \mathbb{R}^d \to \mathbb{R}$, whose gradient is the Kantorovich Monge map between measures $\bar{P}_\sharp\mu$ and $\nu$.

*Proof.* See Appendix C. $\square$

# 4. Learning isometry invariant OT maps

Following the intuition of Makkuva et al. (2020), for any function $g \in L^1(\nu)$, differentiable $\nu$-a.e., it holds that

$$\langle Y, \nabla g(Y) \rangle - f(\nabla g(Y)) \le \langle Y, \nabla f^*(Y) \rangle - f(\nabla f^*(Y)) \\ = f^*(Y) \quad \nu\text{-a.e.}$$

due to the properties of the convex conjugate (Appendix A). In particular

$$\sup_{\substack{g \in L^1(\nu) \\ g \text{ diff. le } \nu\text{-a.e.}}} \mathbb{E}[\langle Y, \nabla g(Y) \rangle - f(\nabla g(Y))] = \mathbb{E}[f^*(Y)],$$

where the expectation is taken w.r.t. $Y \sim \nu$. Thus using that $\sup f = \inf(-f)$ we can re-rewrite Eq. (4) as

$$W_2^2(P_\sharp\mu, \nu) = C_{\mu,\nu} + 2 \sup_{f \in L^1(P_\sharp\mu)} \inf_{\substack{g \in L^1(\nu) \\ g \text{ diff. le } \nu\text{-a.e.}}} \mathcal{V}_{P_\sharp\mu,\nu}(f,g),$$

where $\mathcal{V}_{P_\sharp\mu,\nu}(f,g) := \mathbb{E}_{Y \sim \nu}[f(\nabla g(Y)) - \langle Y, \nabla g(Y) \rangle] - \mathbb{E}_{U \sim P_\sharp\mu}[f(U)]$. Thus the dual PW objective (DPW) reads

$$PW_2^2(\mu, \nu) = \inf_{P \in \mathcal{O}(d)} W_2^2(P_\sharp\mu, \nu). \quad (6)$$

## 4.1. Optimization over $P$, $f$ and $g$.

The dependence of Eq. (6) on $P$ can be made explicit by the following display

$$PW_2^2(\mu, \nu) = C_{\mu,\nu} + 2 \inf_{P \in \mathcal{O}(d)} \sup_{f \in L^1(\mu)} \inf_{g \in L^1(\nu)} \mathcal{V}_{\mu,\nu}(P, f, g), \quad (7)$$

where $g$ is differentiable $\nu$-a.e. and $\mathcal{V}_{\mu,\nu}(P, f, g) := \mathbb{E}_{Y \sim \nu}[f(\nabla g(Y)) - \langle Y, \nabla g(Y) \rangle] - \mathbb{E}_{X \sim \mu}[f(PX)]$.

**Implementation.** We provide estimates for the optimal rigid alignment $(\bar{P})$ and transport map $(\nabla \bar{f})$ starting from the semidual PW formulation of Eq. (7). The heart of our strategy consists into parameterizing $f$ and $g$ by means of fully connected neural networks as Makkuva et al. (2020) did, but with two crucial differences:

**Remark 4.1 (Convexity and neural nets).** As motivated in Makkuva et al. (2020) (Remark 3.4) one does not need ICNNs to approximate $g$. Here we claim that ICNNs are not needed to approximate $f$ anymore, because of Eq. (5): the global minimum $\bar{f}$ is convex and due to the universal approximation theorem we do not need to restrict ourselves to convex architectures.

**Remark 4.2 (Gradient of $g$).** Although one can safely approximate $g$ via neural nets, it is worth noticing that in Eq. (7) we only need $\nabla g$. Henceforth neural nets can be introduced to directly approximate that gradient. This in not an entirely new idea, see for instance Dam et al. (2019) in a generative context.

We acknowledge that, although sounding, Remark 4.1 would need to be stated more formally, in particular due to the non convexity of the objective function and the existence of local stationary points. However we leave it for future works while clearly stating that, in all our experiments, the empirical evidence supported such claim and the estimated $\bar{f}$ systematically happen to be convex without imposing
convexity to the architecture. Hence, we parameterize the potential function f and the gradient potential ∇g using standard MLPs, and treat the entries of the matrix P ∈ ℝ$^{d×d}$ as learnable parameters.

We denote by $\theta_f$ and $\theta_g$ the network parameters of f and $\nabla g$, respectively, and we solve the optimization problem (7) as a min-max-min game. At each iteration, we draw batches $\{x_i\}_{i=1}^R \sim \mu$ and $\{y_j\}_{j=1}^R \sim \nu$. The empirical objective function is defined as

$$ \min_{P} \max_{\theta_f} \max_{\theta_g} \mathcal{J}(P, \theta_f, \theta_g) \qquad (8) $$

where

$$ \mathcal{J}(P, \theta_f, \theta_g) := \frac{1}{R} \sum_{i=1}^{R} f(\nabla g(y_i)) - \langle y_i, \nabla g(y_i) \rangle - f(Px_i). $$

We solved this using Adam (Kingma, 2014) for the potential parameters $\theta_f$ and $\theta_g$ and stochastic gradient descent (SGD) for the alignment matrix P. With respect to the orthogonality constraint of P, we perform a projection step via a SVD after the gradient updates. The full training procedure is summarized in Algorithm 1.

**Algorithm 1 Neural Procrustes-Wasserstein**

**Input:** source $\mu$, target $\nu$, batch size $R$, inner iterations $K_g$, $K_f$, $K_P$, total iterations $T$

Initialize: $P \leftarrow P_0$ (e.g. spectral)

**for** $t \in (1, \dots, T)$ **do**

**for** $k \in (1, \dots, K_g)$ **do**

Sample batches $\{x_i\}_{i=1}^R \sim \mu$, $\{y_j\}_{j=1}^R \sim \nu$

Update $\theta_g$ to minimize $\mathcal{J}(P, \theta_f, \theta_g)$ via Adam

**end for**

**for** $k \in (1, \dots, K_f)$ **do**

Sample batches $\{x_i\}_{i=1}^R \sim \mu$, $\{y_j\}_{j=1}^R \sim \nu$

Update $\theta_f$ to maximize $\mathcal{J}(P, \theta_f, \theta_g)$ via Adam

**end for**

**for** $k \in (1, \dots, K_P)$ **do**

Sample batch $\{x_i\}_{i=1}^R \sim \mu$

Update P to minimize $\mathcal{J}(P, \theta_f, \theta_g)$ via SGD

Project: $P \leftarrow UV^\top$ where $U, S, V^\top \leftarrow \text{SVD}(P)$

**end for**

**end for**

**Remark 4.3 (Initialization).** In discrete PW approaches (see e.g. Alvarez-Melis et al., 2019; Adamo et al., 2025), the optimization of the primal typically starts with an initial guess of the transport plan. In contrast, our continuous formulation, based on the optimization of the semi-dual, requires an initialization of P. The aim is to find a good enough rotation/reflection matrix $P_0 \in \mathcal{O}(d)$ that aligns the source and target measures. Since the joint optimization of $(P, f, g)$ is non-convex, a random initialization of P may lead to suboptimal local minima. To mitigate this, we

adopt a spectral initialization strategy. We effectively treat the (centred) measures $\mu$ and $\nu$ as if they were Gaussian $\mathcal{N}(0, \Sigma_\mu)$ and $\mathcal{N}(0, \Sigma_\nu)$ and, leveraging the closed-form solution derived in Theorem 5.1, we set $P_0 = UV^\top$, where U and V are the eigenvector matrices of $\Sigma_\mu$ and $\Sigma_\nu$ respectively (eigenvectors/eigenvalues are sorted). Geometrically, this operation corresponds to an initial Principal Component Analysis (PCA) alignment, ensuring that the principal axes of the source distribution are pre-aligned with those of the target.

## 5. PW₂ between centered Gaussian measures

In this section, we derive a closed-form expression for the PW distance between two centered Gaussian distributions. As a generalization of the known closed Wasserstein formulation for Gaussian measures (Gelbrich, 1990; Bhatia et al., 2019; Masarotto et al., 2019), we demonstrate that, within the PW framework, the optimal alignment admits an analytical solution based on the spectral decomposition of the covariance matrices. A similar result was recently explored by Toukam (2025). However, we derived it independently and via an alternative proof relying on the nuclear norm formulation of the transport cost and the closed form of the PW map.

Consider two anisotropic centered Gaussian measures $\mu = \mathcal{N}(0, \Sigma_\mu)$ and $\nu = \mathcal{N}(0, \Sigma_\nu)$ on $\mathbb{R}^d$. $\Sigma_\mu, \Sigma_\nu \in \mathbb{S}_{++}^d$ are symmetric positive definite matrices. Recall that the classical squared 2-Wasserstein distance between $\mu$ and $\nu$ has a closed form

$$ W_2^2(\mu, \nu) = \mathrm{tr}(\Sigma_\mu) + \mathrm{tr}(\Sigma_\nu) - 2\mathrm{tr}[(\Sigma_\mu^{1/2} \Sigma_\nu \Sigma_\mu^{1/2})^{1/2}]. $$

The squared 2-Procrustes-Wasserstein distance incorporates an optimization over the orthogonal group $\mathcal{O}(d)$ to align the distributions

$$ PW_2^2(\mu, \nu) = \inf_{P \in \mathcal{O}(d)} W_2^2(\mu, P_\#\nu) $$

where $P_\#\nu = \mathcal{N}(0, P\Sigma_\nu P^\top)$, e.g. the covariance of the rotated target is given by $P\Sigma_\nu P^\top$. Substituting this into $W_2^2(\mu, P_\#\nu)$, we obtain

$$ PW_2^2(\mu, \nu) = \inf_{P \in \mathcal{O}(d)} \{\mathrm{tr}(\Sigma_\mu) + \mathrm{tr}(P\Sigma_\nu P^\top) - 2\mathcal{T}(P)\}, \qquad (9) $$

where $\mathcal{T}(P) := \mathrm{tr}[(\Sigma_\mu^{1/2} P\Sigma_\nu P^\top \Sigma_\mu^{1/2})^{1/2}]$. Observing that the trace is invariant under cyclic permutations and thanks to the orthogonality of of P, we have $\mathrm{tr}(P\Sigma_\nu P^\top) = \mathrm{tr}(\Sigma_\nu)$. Consequently, minimizing the PW distance is equivalent to maximizing the coupling term $\mathcal{T}(P)$. We seek thus for

$$ \bar{P} = \arg \max_{P \in \mathcal{O}(d)} \mathcal{T}(P). $$
Let the spectral decomposition of the covariances be

$$ \Sigma_{\mu} = U \Lambda_{\mu} U^{\top} \quad \text{and} \quad \Sigma_{\nu} = V \Lambda_{\nu} V^{\top}, $$

where $\Lambda_{\mu}$, $\Lambda_{\nu}$ are diagonal matrices containing the sorted eigenvalues (either both in ascending or descending order) of $\Sigma_{\mu}$ and $\Sigma_{\nu}$, and $U, V \in O(d)$ are the respective eigenvector matrices.

**Theorem 5.1.** The squared PW$_2$ distance between centered Gaussian measures $\mu$ and $\nu$ admits the closed form

$$ PW_2^2(\mu, \nu) = tr(\Sigma_{\mu}) + tr(\Sigma_{\nu}) - 2 \sum_{i=1}^{d} \sqrt{\lambda_{\mu}^{i} \lambda_{\nu}^{i}}, \quad (10) $$

where $\{\lambda_{\mu}^i\}_{i=1}^d$ and $\{\lambda_{\nu}^i\}_{i=1}^d$ are the sorted eigenvalues of $\Sigma_{\mu}$ and $\Sigma_{\nu}$, respectively. Moreover, the optimal alignment is $\bar{P} = UV^\top$ and the optimal PW transport map is

$$ T_{PW}(x) = V \Lambda_{\nu}^{1/2} \Lambda_{\mu}^{-1/2} U^{\top} x. $$

*Proof.* See Appendix D. □

Theorem 5.1 has the following geometric intuition: denoting $X \sim \mu$ and $Y \sim \nu$ the source and the target measures, respectively, the PW transport map $T_{PW}$ first rotates (via $U^\top$) the source $X$ to align its principal axis with the canonical ones; then $\Lambda_{\mu}^{-1/2}$ scales to the unit sphere and $\Lambda_{\nu}^{1/2}$ stretch and match to the target's eigenvalues. Finally $V$ optimally align to the target's orientation (the eigenvector basis). A practical behaviour of such effects is reported in Figure 1 and particularly visible (in decoupled actions) in Figure 5 in the Appendix.

# 6. Experiments

In this section we validate NeuralPW on synthetic Gaussian measures and demonstrate its utility for point cloud matching based on the MNIST dataset. MNIST data are centered and scaled to the unit ball. Moreover, we initialize NeuralPW using the spectral strategy described in Remark 4.3.

## 6.1. The Gaussian case

We first consider a controlled setting where the ground truth PW distance is known (Theorem 5.1). We first generate source $X \sim N(0, \Sigma_{\mu})$ and target $Y \sim N(0, \Sigma_{\nu})$ point clouds ($N = M = 5000$) as shown in Figure 1. Then our aim is to match and transport (whenever possible) the source data to the target with three different approaches: (i) plug-in estimation, i.e. the sample covariance matrices (and their spectral decomposition) are injected the closed-forms in Theorem 5.1; (ii) discrete OT, solving the primal PW optimization problem via the alternating minimization solver referred to as *ping-pong* (Adamo et al., 2025; Even

et al., 2024); (iii) our NeuralPW approximation. Figure 1 reports in rows the estimated alignment ($\bar{P}_{\#\mu}$), transport ($\bar{T}_{\#\mu}(\bar{P}_{\#\mu})$) and matching across methods (in columns). All methods correctly estimate the true rotation $P$ (first row). In the second row one sees the transported $X$, not available for primal based methods. The bottom row shows the color gradient transfer from $Y$ to $X$, based on the estimated matching between points. Whereas in primal based methods this can be done via the estimated transport plan, for all other methods (including ground truth), we obtain the discrete matching by assigning each transported source point $\bar{T}_{PW}(x_i)$ to its nearest neighbor in the target $Y$. NeuralPW achieves a squared distance of 0.4116, closely matching the analytical ground truth 0.4268 and recovers a smooth transport map that closely approximate the theoretical one.

## 6.2. MNIST couplings

We now illustrate the advantage of PW over standard Wasserstein distance in a simple case of point clouds clustering. Standard neural OT solvers (Makkuva et al., 2020) are highly sensitive to the global pose of the input distributions. We illustrate this limitation using a subset of the MNIST dataset. We select six images from digit classes 1 and 3, and convert them into 2D point clouds. We consider two scenarios: (i) *aligned*: the digits retain their original orientation; (ii) *rotated*: we apply a random rotation $R \in SO(2)$ to each point cloud individually. For both scenarios, we compute the pairwise distance matrices ($12 \times 12$) using the standard neural OT approximation of $W_2$ (Makkuva et al., 2020) and our NeuralPW framework for PW$_2$. To quantify the discriminative power of the resulting matrices, we perform agglomerative (hierarchical) clustering on the distance matrices and report the Adjusted Rand Index (ARI). The pairwise distance matrices are visualized in Figure 2. In the aligned scenario (top row), $W_2$ struggles to perfectly separate the classes (ARI $\approx$ 0.2). This highlights that even “aligned” datasets contain intrinsic rigid variations that inflate the Wasserstein cost (already noticed in Adamo et al., 2025). In contrast, NeuralPW aligns these variations, achieving perfect separation (ARI = 1.0). This behaviour is especially highlighted in the rotated scenario (bottom row), where $W_2$ collapses completely (ARI = 0.0), as the inter-class shape differences are dominated by the intra-class pose variations. Contrarily, NeuralPW remains invariant to these transformations, recovering the underlying geometric class structure perfectly (ARI = 1.0).

# 7. PW for the study of morphological deformation in archaeozoology

To illustrate both the scalability and the practical utility of NeuralPW when dealing with large 3D point clouds, we present a real-world case study from archaeozoology. Due
<figure><img src="image_3.png" /><figcaption>Figure 1. PW comparison between two Gaussian measures X and Y (left). Top row: the source is aligned to the target. Middle row: transported source. Bottom row: visualized source-target matching (color transfer). Estimated quantities are reported with an overline (\bar{\cdot}).</figcaption></figure>

<figure><img src="image_4.png" /><figcaption>Figure 2. Wasserstein distance matrices ($W_2$) computed using the method of Makkuva et al. (2020) and Procrustes-Wasserstein distance matrices ($PW_2$) computed using NeuralPW. In the first row, MNIST digits 1 and 3 (transformed into point clouds) are directly compared. In the second row, the same digits are compared after applying a random rotation.</figcaption></figure>

to its established theoretical foundations, optimal transport has the potential to be a valuable tool in support of archeo-zoologists seeking to document historical animal biodiversity and quantify the impact of anthropogenic pressures on faunal communities over time. In particular, we used NeuralPW to detect refined delineation of extinct taxa and lineages exhibiting specific adaptations to prolonged environmental fluctuations and human-induced selective pressures. To address similar tasks, bioarcaheologists have relied on established approaches such as Geometric Morphometrics (GMM, Bookstein, 1997) and diffeomorphometry (Durleman et al., 2013). These methods allow one to quantify the morphological variability in both living and fossil organisms and are particularly well suited to the archaeological record. One of their main strengths lies in their adaptability

to a wide range of biological structures (organs, bones and teeth). However, while powerful, they often face two critical limitations: (i) reliance on manual registration or landmark placement to align 3D shapes, introducing human bias; and (ii) the use of simplified morphological templates (averages) to infer evolutionary scenarios. NeuralPW overcomes these limitations by performing automated pairwise comparisons of high-resolution point clouds. By jointly learning the alignment $\bar{P}$ and transport $\nabla \bar{f}$, our framework eliminates the need for pre-processing steps like ICP (Chetverikov et al., 2002) registration. Furthermore, the learned transport map $T_{PW}$ acts as a data-driven diffeomorphism, capturing the dense deformation between specimens without reducing complex anatomies to sparse landmarks. This facilitates a more robust interpretation of species evolution from global pose to local anatomical details.

## 7.1. Case study: Pine Marten skulls

We selected two complete skulls of modern Pine Martens (Martes martes) from the osteological collection of the CEPAM laboratory (UMR 7264, CNRS, Université Côte d'Azur, Nice, France) (Figure 7, in appendix). The specimens, referred to as Skull A (ID: CA025) and Skull B (ID: CA023), correspond to two adult males (Figure 8, in appendix). Standard osteometric measurements (Von Den Drivesch, 1976) reveal distinct size variations: Skull B measures 8.3cm in length (Akrokranion to Prosthion) and 5.2cm in Zygomatic breadth, while Skull A is smaller, measuring 7.7cm and 4.8cm, respectively (see Table 1, in appendix, for a full description). The specimens were digitized using an Artec Spider surface scanner following the protocol in Vuillien (2020), resulting in high-density meshes ($N \approx 500K$ vertices). The use of high-resolution 3D models is strictly necessary to provide a detailed and reliable assessment of
Neural approximation of Procrustes-Wasserstein transport maps

<figure><img src="image_5.png" /><figcaption>**Figure 3.** PW matching between two high-resolution Pine Marten skulls (*N* = 500K). Our neural approach successfully recover the global rotation to align the source (left-most) with the target (middle-left), producing the aligned source *PX* (middle-right). The final transport (right-most) captures shape deformations.</figcaption></figure>

<figure><img src="image_6.png" /><figcaption>**Figure 4.** Runtime comparison (in seconds) between the discrete *ping-pong* solver (Adamo et al., 2025) and our NeuralPW framework. The x-axis represents the point cloud size *N* on logarithmic scale. While the discrete solver suffers from prohibitive computational scaling as *N* increases, NeuralPW maintains a manageable computational footprint, enabling a robust transport of high-density distributions (*N* > 10⁵).</figcaption></figure>

the morphological features of the two martens, allowing the
discussion of intraspecific characteristics.

We perform a PW “matching” between the two high-resolution scans. NeuralPW is trained with a fixed batch size of 500 and for a maximum of 1000 epochs. The potential *f* and gradient potential ∇*g* are optimized with learning rates of 10⁻⁴, while the orthogonal matrix *P* uses a learning rate of 2 · 10⁻³.

**Scalability of PW solvers.** A key advantage of our neural approximation is its ability to scale to datasets where discrete solvers become intractable. Standard discrete OT solvers, including the alternating minimization scheme for PW (Alg. 1, Adamo et al., 2025), also called *ping-pong* in Even et al. (2024), typically scale quadratically with the number of samples (*O*(*N*²)), or cubically if using standard linear programming. In Figure 4, we compare the runtime of NeuralPW against the discrete *ping-pong* solver across vary-

ing downsampled resolutions. While *ping-pong* exhibits
explosive computational cost, becoming prohibitive and
memory-constrained beyond *N* = 10⁴, NeuralPW main-
tains a manageable runtime. Indeed, its complexity depends
primarily on the batch size and network capacity rather than
the total cardinality of the dataset. This allows bioarchaeol-
ogists to leverage the full resolution of modern 3D scanners
without artificial downsampling.

**Discussion.** Focusing on high-resolution point clouds, NeuralPW estimates an optimal alignment $\bar{P}$ and transport $\nabla \bar{f}$, yielding an estimated squared distance of $PW_2^2 = 3.1079$. As shown in Figure 3, the method accurately disentangles pose from shape. The aligned source $\bar{P}X$ matches the orientation of the target $Y$, while the transported source $T_{PW}(X) = \nabla \bar{f}(\bar{P}X)$ successfully morphs the anatomical features of Skull A to match Skull B. In the basioccipital region, for instance, significant differences are observed in the shape of the occipital condyles and the nuchal tubercle. Specifically, Skull A exhibits distinct curvature in the region between the condyles and the tubercle compared to Skull B. This characteristic is specifically visible by the transported source $T_{PW}(X)$ in green (see Figure 9, in appendix). Future work will focus on extending NeuralPW to population-level analysis, and possibly highlighting the connections with PW barycenters.

8. Conclusion

In this work we have extended the PW framework to the con-
tinuous domain, formally introducing the dual and semidual
formulations of the PW problem and derived characteriza-
tion of PW transport maps. We have then approximated
the PW cost and maps, via a neural network-based method
named NeuralPW. We have further derived the closed-form
solution of the PW distance between centered Gaussian
measures and provide experimental setup to highlight the
neural estimations. Finally we have presented a case study
in archaeozoology that enhance future directions of OT ap-
proaches in real domains.

## Acknowledgements

The work has received financial support from the CNRS through the MITI interdisciplinary programs and the Junior Professor Chair (Chaire de Professeur Junior, CPJ) funded by the French National Research Agency (ANR). We would like to thank Arch’AI’Story project (Ministère de l’Enseignement Supérieur et de la Recherche and Univeristy Côte d’Azur) for funding this project.

## Impact Statement

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

## References

Adamo, D., Corneli, M., Vuillien, M., and Vila, E. An in depth look at the procrustes-wasserstein distance: properties and barycenters. *arXiv preprint arXiv:2507.00894*, 2025.

Alvarez-Melis, D. and Jaakkola, T. Gromov-wasserstein alignment of word embedding spaces. In *Proceedings of the 2018 conference on empirical methods in natural language processing*, pp. 1881–1890, 2018.

Alvarez-Melis, D., Jegelka, S., and Jaakkola, T. S. Towards optimal transport with global invariances. In *The 22nd International Conference on Artificial Intelligence and Statistics*, pp. 1870–1879. PMLR, 2019.

Amos, B., Xu, L., and Kolter, J. Z. Input convex neural networks. In *International conference on machine learning*, pp. 146–155. PMLR, 2017.

Arjovsky, M., Chintala, S., and Bottou, L. Wasserstein generative adversarial networks. In *International conference on machine learning*, pp. 214–223. PMLR, 2017.

Balaji, Y., Chellappa, R., and Feizi, S. Robust optimal transport with applications in generative modeling and domain adaptation. *Advances in Neural Information Processing Systems*, 33:12934–12944, 2020.

Bhatia, R., Jain, T., and Lim, Y. On the bures-wasserstein distance between positive definite matrices. *Expositiones mathematicae*, 37(2):165–191, 2019.

Bookstein, F. L. Landmark methods for forms without landmarks: morphometrics of group differences in outline shape. *Medical image analysis*, 1(3):225–243, 1997.

Chetverikov, D., Svirko, D., Stepanov, D., and Krsek, P. The trimmed iterative closest point algorithm. In *2002 International Conference on Pattern Recognition*, volume 3, pp. 545–548. IEEE, 2002.

Chewi, S., Niles-Weed, J., and Rigollet, P. Statistical optimal transport. *arXiv preprint arXiv:2407.18163*, 3, 2024.

Chowdhury, S., Miller, D., and Needham, T. Quantized gromov-wasserstein. In *Joint European Conference on Machine Learning and Knowledge Discovery in Databases*, pp. 811–827. Springer, 2021.

Courty, N., Flamary, R., Tuia, D., and Rakotomamonjy, A. Optimal transport for domain adaptation. *IEEE transactions on pattern analysis and machine intelligence*, 39(9): 1853–1865, 2016.

Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. *Advances in neural information processing systems*, 26, 2013.

Dam, N., Hoang, Q., Le, T., Nguyen, T. D., Bui, H., and Phung, D. Three-player wasserstein gan via amortised duality. In *International Joint Conference on Artificial Intelligence 2019*, pp. 2202–2208. Association for the Advancement of Artificial Intelligence (AAAI), 2019.

Divol, V., Niles-Weed, J., and Pooladian, A.-A. Optimal transport map estimation in general function spaces. *The Annals of Statistics*, 53(3):963–988, 2025.

Durrleman, S., Pennec, X., Trouvé, A., Braga, J., Gerig, G., and Ayache, N. Toward a comprehensive framework for the spatiotemporal statistical analysis of longitudinal shape data. *International journal of computer vision*, 103 (1):22–59, 2013.

Eisenberger, M., Toker, A., Leal-Taixé, L., and Cremers, D. Deep shells: Unsupervised shape correspondence with optimal transport. *Advances in Neural information processing systems*, 33:10491–10502, 2020.

Even, M., Ganassali, L., Maier, J., and Massoulié, L. Aligning embeddings and geometric random graphs: Informational results and computational approaches for the procrustes-wasserstein problem. *Advances in Neural Information Processing Systems*, 37:70730–70764, 2024.

Feydy, J., Charlier, B., Vialard, F.-X., and Peyré, G. Optimal transport for diffeomorphic registration. In *International conference on medical image computing and computer-assisted intervention*, pp. 291–299. Springer, 2017.

Gelbrich, M. On a formula for the 12 wasserstein metric between measures on euclidean and hilbert spaces. *Mathematische Nachrichten*, 147(1):185–203, 1990.
Grave, E., Joulin, A., and Berthet, Q. Unsupervised alignment of embeddings with wasserstein procrustes. In *The 22nd International Conference on Artificial Intelligence and Statistics*, pp. 1880–1890. PMLR, 2019.

Gulrajani, I., Ahmed, F., Arjovsky, M., Dumoulin, V., and Courville, A. C. Improved training of wasserstein gans. *Advances in neural information processing systems*, 30, 2017.

Horn, R. A. and Johnson, C. R. *Topics in matrix analysis*. Cambridge university press, 1994.

Hütter, J.-C. and Rigollet, P. Minimax estimation of smooth optimal transport maps. *The Annals of Statistics*, 49(2): 1166–1194, 2021.

Kantorovich, L. V. On the translocation of masses. In *Dokl. Akad. Nauk. USSR (NS)*, volume 37, pp. 199–201, 1942.

Kingma, D. P. Adam: A method for stochastic optimization. *arXiv preprint arXiv:1412.6980*, 2014.

Korotin, A., Selikhanovych, D., and Burnaev, E. Neural optimal transport. *arXiv preprint arXiv:2201.12220*, 2022.

Makkuva, A., Taghvaei, A., Oh, S., and Lee, J. Optimal transport mapping via input convex neural networks. In *International Conference on Machine Learning*, pp. 6672–6681. PMLR, 2020.

Manole, T., Balakrishnan, S., Niles-Weed, J., and Wasserman, L. Plugin estimation of smooth optimal transport maps. *The Annals of Statistics*, 52(3):966–998, 2024.

Masarotto, V., Panaretos, V. M., and Zemel, Y. Procrustes metrics on covariance operators and optimal transportation of gaussian processes. *Sankhya A*, 81(1):172–213, 2019.

Mémoli, F. Gromov–wasserstein distances and the metric approach to object matching. *Foundations of computational mathematics*, 11(4):417–487, 2011.

Nutz, M. Introduction to entropic optimal transport. *Lecture notes*, Columbia University, 2021.

Peyré, G., Cuturi, M., et al. Computational optimal transport: With applications to data science. *Foundations and Trends® in Machine Learning*, 11(5-6):355–607, 2019.

Ramírez, G., Dangovski, R., Nakov, P., and Soljačić, M. On a novel application of wasserstein-procrustes for unsupervised cross-lingual alignment of embeddings. In *Proceedings of the 17th Workshop on Building and Using Comparable Corpora (BUCC)@ LREC-COLING 2024*, pp. 1–11, 2024.

Santambrogio, F. Optimal transport for applied mathematicians. *Birkhäuser, NY*, 55(58-63):94, 2015.

Solomon, J., De Goes, F., Peyré, G., Cuturi, M., Butscher, A., Nguyen, A., Du, T., and Guibas, L. Convolutional wasserstein distances: Efficient optimal transportation on geometric domains. *ACM Transactions on Graphics (ToG)*, 34(4):1–11, 2015.

Toukam, K. M. Procrustes wasserstein metric: A modified benamou-brenier approach with applications to latent gaussian distributions. *arXiv preprint arXiv:2503.16580*, 2025.

Vayer, T., Chapel, L., Flamary, R., Tavenard, R., and Courty, N. Optimal transport for structured data with application on graphs. *arXiv preprint arXiv:1805.09114*, 2018.

Vayer, T., Flamary, R., Courty, N., Tavenard, R., and Chapel, L. Sliced gromov-wasserstein. *Advances in Neural Information Processing Systems*, 32, 2019.

Villani, C. et al. *Optimal transport: old and new*, volume 338. Springer, 2008.

Villani, C. et al. *Optimal transport: old and new*, volume 338. Springer, 2009.

Vincent-Cuaz, C., Vayer, T., Flamary, R., Corneli, M., and Courty, N. Online graph dictionary learning. In *International conference on machine learning*, pp. 10564–10574. PMLR, 2021.

Von Den Driesch, A. *A guide to the measurement of animal bones from archaeological sites*, volume 1. Peabody museum press, 1976.

Vuillien, M. *Systèmes d’élevage et pastoralisme en Provence et dans les Alpes méridionales durant la Protohistoire: Nouvelles perspectives en archéozoologie*. PhD thesis, Université Côte d’Azur, 2020.

Yang, L., Wang, X., Zhang, J., Yang, J., Xu, Y., Hou, J., Xin, K., and Wang, F.-Y. Hackgan: Harmonious cross-network mapping using cyclegan with wasserstein–procrustes learning for unsupervised network alignment. *IEEE Transactions on Computational Social Systems*, 10 (2):746–759, 2022.

Zhang, M., Liu, Y., Luan, H., and Sun, M. Earth mover’s distance minimization for unsupervised bilingual lexicon induction. In *Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing*, pp. 1934–1945, 2017.
A. Convexity

Consider $f: \mathbb{R}^d \to \mathbb{R} \cup \{+\infty\}$, then $f$ is said to be convex if and only if

$$f(tx + (1-t)y) \leq tf(x) + (1-t)f(y) \quad \forall x, y, \in \mathbb{R}^d, \forall t \in [0, 1].$$

If $f \in C^1(\mathbb{R}^d)$ then $f$ is convex iff $f(y) \geq f(x) + \langle \nabla f(x), y - x \rangle$. Moreover, a function $f \in C^2(\mathbb{R}^d)$ is convex iff $Hf(x) \geq 0$ (i.e. the Hessian is at least positive semi-definite).

**Definition A.1.** Let $f: \mathbb{R}^d \to \mathbb{R} \cup \{+\infty\}$. The convex conjugate of $f$ is the function $f^*: \mathbb{R}^d \to \mathbb{R} \cup \{+\infty\}$ defined by

$$f^*(x^*) := \sup_{x \in \mathbb{R}^d} \{\langle x, x^* \rangle - f(x)\}.$$

Note that for every function $f: \mathbb{R}^d \to \mathbb{R} \cup \{+\infty\}$, $f^*$ is a convex function. Moreover $f^*$ is closed (e.g. $(f^*)^{-1}_{(-\infty, a]}$ is closed in $\mathbb{R}^d$).

From the definition of the convex conjugate we have the Fenchel-Young inequality that reads

$$f(x) + f^*(y) \geq \langle x, y \rangle, \qquad (11)$$

where equality holds if and only if $y \in \nabla f(x)$, e.g. $f(x) + f^*(\nabla f(x)) = \langle x, \nabla f(x) \rangle$. By symmetry

$$f(\nabla f^*(y)) + f^*(y) = \langle \nabla f^*(y), y \rangle \qquad (12)$$

B. The Procrustes-Wasserstein problem

Let $\mu, \nu \in P_2(\mathbb{R}^d)$ and let $X \sim \mu$, $Y \sim \nu$ be a source and target distributions, respectively. Fix the orthogonal group $O(d) = \{P \in \mathbb{R}^{d \times d} : P^T P = I\}$. For $P \in O(d)$ define the cost $c_P(x, y) := \|Px - y\|_2$.

B.1. Equivalent primal formulation

We start with the following

**Proposition B.1.** For any $P \in O(d)$, the Kantorovich transport cost between $P_{\#}\mu$ (the pushforward of $\mu$ by $P$) and $\nu$ is

$$W_2^2(P_{\#}\mu, \nu) = \min_{\pi \in \Pi(\mu, \nu)} \mathbb{E}_{X,Y \sim \pi} [c_P(X, Y)^2],$$

where $\Pi(\mu, \nu)$ is the set of couplings $\pi \in P(\mathbb{R}^d \times \mathbb{R}^d)$ with marginals $\mu$ and $\nu$.

*Proof.* First let us remark that if two random variables $X \sim \mu$ and $U \sim P_{\#}\mu$ are considered, with $\mu$ admitting a density denoted by $p_X(\cdot)$, than by change of variables in multiple integrals we have

$$\mathbb{E}(U) = \mathbb{E}(PX) = \int_{\mathbb{R}^d} Pxp_X(x)dx = \int_{\mathbb{R}^d} up_X(P^T u)|\det P^T|du = \int_{\mathbb{R}^d} up_X(P^T u)du$$

where the last equality comes from the orthogonality of $P$ and shows that

$$p_U(u) = p_X(P^T u) \qquad (13)$$

where $p_U(\cdot)$ is nothing but the density of $P_{\#}\mu$. Now, by definition of Wasserstein (squared) distance it holds that $W_2^2(P_{\#}\mu, \nu) := \min_{\gamma \in \Pi(P_{\#}\mu, \nu)} \mathbb{E}_{U,Y \sim \gamma} [||U - Y||_2^2]$. Moreover

$$\begin{aligned} \mathbb{E}_{U,Y \sim \gamma} [||U - Y||_2^2] &= \int_{\mathbb{R}^d \times \mathbb{R}^d} ||u - y||_2^2 \gamma(u, y) du dy \\ &= \int_{\mathbb{R}^d \times \mathbb{R}^d} ||Px - y||_2^2 \gamma(Px, y) dx dy = \mathbb{E}_{X,Y \sim \pi} ||Px - Y||_2^2, \end{aligned}$$
where change of variables was employed again and we defined $\pi(x, y) := \gamma(Px, y)$. It is easy to check that $\pi$ is a well defined density in $\Gamma(\mu, \nu)$. In particular

$$ \int_{\mathbb{R}^d} \pi(x, y)dy = \int_{\mathbb{R}^d} \gamma(Px, y)dy = p_U(Px) = p_X(x), $$

where the penultimate and last equalities come from the marginal constraints on $\gamma$ and Eq. (13), respectively. Finally, since $\pi = \gamma \odot (P, Id)$ and $P$ is an isomorphism, the correspondence between $\pi$ and $\gamma$ is one to one and the proposition is proven. □

**Note.** Equivalently one may view $\Pi$ as the set of couplings between $P_{\#}\mu$ and $\nu$. Indeed, there is a one-to-one correspondence between couplings $\tilde{\pi}$ of $P_{\#}\mu$ and $\nu$ and couplings $\pi$ of $\mu$ and $\nu$ given by $\pi = (P^{-1} \times \text{Id})_{\#}\tilde{\pi}$, but for practical derivations it is convenient to keep $\pi \in \Pi(\mu, \nu)$ and the integrand $\|Px - y\|_2^2$.

Thence, the squared 2-Procrustes-Wasserstein cost is defined by

$$ PW_2^2(\mu, \nu) := \min_{P \in \mathcal{O}(p)} \min_{\pi \in \Pi(\mu, \nu)} \mathbb{E}_{X,Y \sim \pi} [c_P(X, Y)^2] = \min_{P \in \mathcal{O}(p)} W_2^2(P_{\#}\mu, \nu), $$

telling us that we “just” need to find the orthogonal transformation $P$ that minimizes the usual $W_2^2$ distance between the rotated source $P_{\#}\mu$ and $\nu$.

## B.2. Dual formulation

In order to derive the dual formulation of the Procrustes-Wasserstein problem, we come back to the marginal constraint imposed to the joint distribution $\pi$. We can reformulate these constraint, by saying that $\pi \in \Pi(\mu, \nu)$ if for all bounded measurable functions $\varphi, \psi : \mathbb{R}^d \to \mathbb{R}$

$$ \mathbb{E}_{x,y \sim \pi} [\varphi(x)] = \mathbb{E}_{x \sim \mu} [\varphi(x)], \quad \mathbb{E}_{x,y \sim \pi} [\psi(y)] = \mathbb{E}_{y \sim \nu} [\psi(y)]. $$

Let us now fix $P \in \mathcal{O}(d)$. We introduce potential functions (Lagrange multipliers) $\varphi \in L^1(\mu)$ and $\psi \in L^1(\nu)$. The Lagrangian that enforces marginal constraints if

$$ \begin{aligned} \mathcal{L}(\pi, \varphi, \psi) &= \mathbb{E}_{X,Y \sim \pi} [c_P(X, Y)^2] + \mathbb{E}_{X \sim \mu} [\varphi(X)] - \mathbb{E}_{X,Y \sim \pi} [\varphi(X)] + \mathbb{E}_{Y \sim \nu} [\psi(Y)] - \mathbb{E}_{X,Y \sim \pi} [\psi(Y)] \\ &= \mathbb{E}_{X,Y \sim \pi} [c_P(X, Y)^2 - \varphi(X) - \psi(Y)] + \mathbb{E}_{X \sim \mu} [\varphi(X)] + \mathbb{E}_{Y \sim \nu} [\psi(Y)]. \end{aligned} $$

By the classical Kantorovich duality, the inner infimum over $\pi$ admits the dual formulation

$$ W_2^2(P_{\#}\mu, \nu) = \sup_{(\varphi, \psi) \in \Phi_P} \{\mathbb{E}_{X \sim \mu}[\varphi(X)] + \mathbb{E}_{Y \sim \nu}[\psi(Y)]\}, $$

where

$$ \Phi_P := \{(\varphi, \psi) \in L^1(\mu) \times L^1(\nu) \mid \forall (x, y) \quad \varphi(x) + \psi(y) \le \|Px - y\|_2^2, \quad (\mu, \nu) - \text{a.e.}\}. $$

Putting back the outer infimum over the orthogonal transformations, the Procrustes-Wasserstein dual formulation becomes

$$ PW_2^2(\mu, \nu) = \inf_{P \in \mathcal{O}(p)} \sup_{(\varphi, \psi) \in \Phi_P} \{\mathbb{E}_{X \sim \mu}[\varphi(X)] + \mathbb{E}_{Y \sim \nu}[\psi(Y)]\}. $$

## C. Proof of Theorem 3.2

*Proof.* To prove the Theorem it suffices to explicit the link given by Proposition 3.1 and Eq. (3). Denoting $\bar{P}$ the minimizer of Eq. (3), then the problem to solve is $W_2^2(\bar{P}_{\#}\mu, \nu)$. Since $\mu$ is absolutely continuous with respect to the Lebesgue measure, the push-forward $\bar{P}_{\#}\mu$, under the diffeomorphism $x \mapsto \bar{P}x$, is also absolutely continuous. By Brenier’s Theorem (Peyré et al., 2019), there exists a unique Monge map $T_{\bar{P}_{\#}\mu \to \nu}$ (form $\bar{P}_{\#}\mu$ to $\nu$), given by

$$ T_{\bar{P}_{\#}\mu \to \nu}(u) = \nabla \bar{f}(u), $$

where $\bar{f}$ is a convex function. The PW transport map acting on $x$ is therefore the gradient of $f$ evaluated at $u = \bar{P}x$. □
D. PW₂ between centered Gaussian measures

We report here the detailed proof for Theorem 5.1.

Proof. (Closed PW formulation) We want to maximize the term $\mathcal{T}(P) = \text{tr}[(\Sigma_{\mu}^{1/2} P \Sigma_{\nu} P^{\top} \Sigma_{\mu}^{1/2})^{1/2}]$.

First, notice that for any matrix A, the term tr((A^T)^{1/2}) corresponds to the nuclear norm $||A||_*$ = $\sum_i \lambda^i$, where $\lambda^i$ are the singular values of A. Let us define $A := \Sigma_{\mu}^{1/2} P \Sigma_{\nu}^{1/2}$. Then

$$ \mathcal{T}(P) = \|\Sigma_{\mu}^{1/2} P \Sigma_{\nu}^{1/2}\|_{*}. $$

Now, considering the diagonalization of the covariance matrix $\Sigma_{\mu}$, we have that $\Sigma_{\mu}^{1/2} = U \Lambda_{\mu}^{1/2} U^{\top}$ (respectively for $\nu$). Substituting the spectral decompositions and exploiting the unitary invariance of the nuclear norm, e.g. $||XYZ||_* = ||Y||_*$ for unitary X, Z (X, Z are real orthogonal matrices), we have

$$ \begin{aligned} \mathcal{T}(P) &= \|U \Lambda_{\mu}^{1/2} U^{\top} P V \Lambda_{\nu}^{1/2} V^{\top}\|_{*} \\ &= \|\Lambda_{\mu}^{1/2} (U^{\top} P V) \Lambda_{\nu}^{1/2}\|_{*}. \end{aligned} $$

Let $Q := U^{\top} P V$. Since $U, P, V \in \mathcal{O}(d)$, it follows that $Q \in \mathcal{O}(d)$. The optimization problem transforms into

$$ \max_{Q \in \mathcal{O}(d)} \|\Lambda_{\mu}^{1/2} Q \Lambda_{\nu}^{1/2}\|_{*}. $$

Using the Von Neumann trace inequality (or the singular value inequality for products of matrices), we have the upper bound

$$ \|\Lambda_{\mu}^{1/2} Q \Lambda_{\nu}^{1/2}\|_{*} \le \sum_{i=1}^{d} \sqrt{\lambda_{\mu}^{i}} \sqrt{\lambda_{\nu}^{i}}, $$

where the singular values of $\Lambda_{\mu}^{1/2}$ and $Q\Lambda_{\nu}^{1/2}$ are $\sqrt{\lambda_{\mu}^{i}}$ and $\sqrt{\lambda_{\nu}^{i}}$, respectively.

This upper bound is achieved when $Q$ is the identity matrix, e.g. $\bar{Q} = \mathbf{I}_d$. Recovering the optimal rotation $P$, we have $\mathbf{I}_d = U^{\top} \bar{P} V \Rightarrow \bar{P} = UV^{\top}$. Finally, substituting the matrix value $\mathcal{T}(\bar{P}) = \sum_{i=1}^{d} \sqrt{\lambda_{\mu}^{i} \lambda_{\nu}^{i}}$ into Eq. (9) yields the stated result of Eq. (10).

(Recovering the PW map) The aligned target distribution is drawn from $\bar{P}_{\#\nu}$, where his covariance matrix is

$$ \Sigma_{\bar{P}_{\#\nu}} = \bar{P} \Sigma_{\nu} (\bar{P})^{\top} = (UV^{\top})(V \Lambda_{\nu} V^{\top})(VU^{\top}) = U \Lambda_{\nu} U^{\top}. $$

Note that the source covariance $\Sigma_{\mu}$ and the aligned target covariance $\Sigma_{\bar{P}_{\#\nu}}$ share the same eigenvectors matrix $U$. Therefore, they commute.

We recall that the 2-Wasserstein Monge map between $\mu$ and $\bar{P}_{\#\nu}$ is $T(x) = \Sigma_{\mu}^{-1/2} (\Sigma_{\mu}^{1/2} \Sigma_{\bar{P}_{\#\nu}} \Sigma_{\mu}^{1/2})^{1/2} \Sigma_{\mu}^{-1/2} x$. Since, in our context, $\Sigma_{\mu}$ and $\Sigma_{\bar{P}_{\#\nu}}$ commute, the Monge map simplifies drastically to element-wise scaling of the eigenvalues (Masarotto et al., 2019), e.g.

$$ T(x) = U \Lambda_{\nu}^{1/2} \Lambda_{\mu}^{-1/2} U^{\top} x. $$

Since $T$ maps $X \sim \mu$ to $Y' \sim \bar{P}_{\#\nu}$, and we know that $Y' = \bar{P}Y$ (where $Y \sim \nu$), we can recover the target $Y$ by applying the inverse rigid transform $(\bar{P})^{\top}$, e.g. $Y = (\bar{P})^{\top} Y' = (\bar{P})^{\top} T(X)$. Substituting we get

$$ T_{PW}(x) = (VU^{\top})(U\Lambda_{\nu}^{1/2}\Lambda_{\mu}^{-1/2}U^{\top})x = V\Lambda_{\nu}^{1/2}\Lambda_{\mu}^{-1/2}U^{\top}x. $$

□

Remark D.1. While we prove the statement by exploiting the formulation of $\mathcal{T}(P)$ in terms of the nuclear norm, Toukam (2025) approaches the maximization via the singular value inequality for matrix product, i.e. $\sum_i \lambda_{AB}^i \le \sum_i \lambda_A^i \lambda_B^i$ for couples of matrices A, B (Theorem 3.3.14, Horn & Johnson, 1994). Algebraically, these derivations are equivalent, as both rely on the Von Neumann trace inequality which dictates that the trace is maximized when the principal axes of the spectral decomposition are aligned (yielding to $Q = \mathbf{I}_d$). Nevertheless, we provide here the explicit formulation of the PW Monge map $T_{PW}$, while Toukam (2025) does not detail it. They only observe that, in this context, the Monge map is the one between the optimally rotated source and the target distributions.
<figure><img src="image_7.png" /><figcaption>Figure 5. Geometric interpretation of the PW map between two Gaussian measures $X \sim \mathcal{N}(0, \Sigma_\mu)$ and $Y \sim \mathcal{N}(0, \Sigma_\nu)$.</figcaption></figure>

## E. Experimental settings

We provide here the detailed experimental setup and hyperparameters to ensure the reproducibility of the results reported in Section 6.

**NeuralPW implementation details.** For all experiments, we initialized the alignment matrix $P$ using the spectral initialization strategy described in Remark 4.3. The potentials $f: \mathbb{R}^d \to \mathbb{R}$ and $\nabla g: \mathbb{R}^d \to \mathbb{R}^d$ were parameterized using two MLPs with 2 hidden layers and Softplus activations. We employed the alternating optimization scheme described in Algorithm 1 with the following inner iteration settings: $K_g = 10$, $K_f = 5$ and $K_P = 5$. A crucial detail of our implementation concerns the orthogonality constraint on $P$. While standard projected Gradient Descent requires projecting $P$ onto the manifold $\mathcal{O}(d)$ (via SVD) after every gradient step, we performed the orthogonal projection only once at the end of the outer loop (iteration $t$), rather than after each inner update $k = 1, \dots, K_P$. We empirically observed that delaying the projection allows the optimization trajectory to traverse the ambient Euclidean space more freely before snapping back to the manifold, resulting in faster convergence and more stable minima. The network parameters $\theta_f$ and $\theta_g$ were optimized using Adam with learning rates of $10^{-4}$, while the matrix $P$ was optimized using SGD with a learning rate of $2 \cdot 10^{-3}$.

**Gaussian case.** We generated source and target data by sampling $N = M = 5000$ points from two centered Gaussian measures $\mathcal{N}(0, \Sigma_\mu)$ and $\mathcal{N}(0, \Sigma_\nu)$. NeuralPW was trained for $T = 1000$ total iterations with a batch size of $R = 500$. For the discrete baseline comparison, we utilized the alternating minimization (*ping-pong*) scheme from Adamo et al. (2025, Algorithm 1). We set the convergence tolerance to $10^{-4}$ and the maximum number of alternating steps (e.g. *ping-pong* steps) to 30, which was sufficient for convergence in all trials.

**MNIST case.** We selected the first six images from the MNIST test dataset for classes 1 and 3 (Figure 6, left). Each $28 \times 28$ image was converted into a 2D point cloud by extracting the coordinates of active pixels. To normalize the data, we centered each cloud at its Euclidean barycenter and scaled it to be enclosed within the unit sphere (Figure 6, center). This constitutes our *aligned* scenario. For the *rotated* scenario, we applied a random rotation to each point cloud individually (Figure 6, right). We compared NeuralPW against the standard neural OT solver of Makkuva et al. (2020). For the baseline, we parameterized the potentials $f$ and $g$ using ICNNs with 4 hidden layers. We trained using the minimax game reported in Algorithm 1 of Makkuva et al. (2020) for $T = 1000$ iterations with a batch size of 50. The learning rates were set to $2 \cdot 10^{-3}$ for $f$ and $10^{-4}$ for the dual potential $g$. For NeuralPW, we maintained the same batch size ($R = 50$) and a maximum of training iterations of $T = 1000$, utilizing the architecture and projection strategy described above to estimate the $PW_2$ distance.
<figure><img src="image_8.png" /><figcaption>Figure 6. Left: the 12 MNIST images considered. Center: the associated 2D point clouds used in the *aligned* scenario. Right: the randomly *rotated* point clouds.</figcaption></figure>

F. Additional material for Section 7

<figure><img src="image_9.png" /><figcaption>Figure 7. Picture of a modern Pine Marten (*Martes martes*) © Martre des pins, un animal commun de nos forêts</figcaption></figure>

<figure><img src="image_10.png" /><figcaption>Figure 8. Pictures of Marten skulls used in the case study. Collection of © M. Vuillien, CEPAM osteological collection (UMR 7264, CNRS, France).</figcaption></figure>

<table><caption>Table 1. Inventory for the two modern specimens of *Martes martes* belonging to [Anonymous Laboratory]. All the measurements are in millimeters. Following Von Den Driesch (1976), columns refer to: (1) total length between Akrokranion - Prosthion; (23) Zygomatic breadth; (18) greatest mastoid breadth; (19) greatest breadth of the occipital condyles; (20) greatest breadth of the foramen magnum; (21) height of the foramen magnum.</caption><thead><tr><th>SPECIMEN</th><th>(1)</th><th>(23)</th><th>(18)</th><th>(19)</th><th>(20)</th><th>(21)</th></tr></thead><tbody><tr><td>CA023</td><td>82.53</td><td>51.75</td><td>39.10</td><td>20.78</td><td>11.43</td><td>9.80</td></tr><tr><td>CA025</td><td>77.25</td><td>48.03</td><td>38.87</td><td>18.84</td><td>10.54</td><td>7.30</td></tr></tbody></table>
<figure><img src="image_11.png" /><figcaption>Figure 9. Morphological variation of the basioccipital region between the two Pine Marten specimens. Top: textured 3D model view © M. Vuillien.</figcaption></figure>