Question:
Regarding the description of the DWGF (Diffusion-regularized Wasserstein Gradient Flow) algorithm, which of the following statements are correct:
A. The diffusion regularization term introduced in the algorithm is designed to amplify the difference between the distribution of the latent variable and the prior distribution along the entire diffusion path.
B. This algorithm is a training-free method that directly utilizes the pre-trained latent diffusion model (such as Stable Diffusion) as a prior.
C. The algorithm employs an adaptive optimizer to discretize the Wasserstein gradient flow, aiming to accelerate convergence and handle ill-conditioned optimization problems.
D. From the test results in the text, it can be seen that compared to RLSD and PSLD, the output is more clear and relatively smooth/average, lacking detailed variations.

Answer: BCD

2.2
Wasserstein gradient of R[µ(z0|y)]
To fully leverage the information provided by the generative model, we first construct
the prior regularization term R(µ(z0|y)) based on the pretrained latent diffusion model.
Analogous to Wang et al. (2023); Luo et al. (2023), we consider a weighted KL divergence
along the diffusion process:
R(µ(z0|y)) := Dw,[0,T ]
KL
(µ(z0|y)∥pθ−(z0)) =
Z T
0
w(s)DKL(µ(zs|y)∥pθ−(zs))ds,
(9)
where w(s) : [0, T] →R+ is a time-dependent weighting term and the densities µ(zs|y) :=
R
Z N(zs; αsz0, σ2
sI)µ(z0|y)dz0 are pushforwards of µ through the forward transition kernel
of the diffusion model. We remark that the weighted KL divergence in (9) is a well-defined
functional and admits favorable properties (cf. Section A.4) as summarised below.
Theorem 2.1. The weighted KL divergence Dw,[0,T ]
KL
(µ(z0|y)∥pθ−(z0)) (9) is i) nonnegative,
ii) convex in the first component µ(z0|y), and iii) is minimized if and only if the standard
KL divergence DKL(µ(z0|y)∥pθ−(z0)) is minimized.
Accordingly, we obtain the gradient of the first variation of R (cf. Section A.2):
∇z0
δR[µ]
δµ
= Es,ϵ

˜w(s)

∇zs log
Z
Z
p(zs|z0)µ(z0,t|y)dz0 −∇zs log pθ−(zs)
 ∂zs
∂z0

,
(10)
where the expectation is taken over s ∼U(0, T), ϵ ∼N(0, I), we set ˜w(s) := Tw(s), and
p(zs|z0) = N(zs; αsz0, σ2
sI) is the forward kernel of the diffusion model. To approximate
the integral in (10), we use a particle-based approach to form a Monte-Carlo approximation
using {z(i)
0 }i ∼µ(z0|y) (Wang et al., 2023; Lim and Johansen, 2024).
2.3
Algorithmic Considerations
Final Gradient Flow. We can simply combine the gradients derived in (7) and (10) to
obtain the final gradient flow in (5). As we have an intractable term, we will simulate N
identical copies of this ODE, where the integral in (10) is approximated using these particles
as mentioned. We mention some practical aspects of this implementation below.
Deterministic Encoding
During our experiments, we observe that the diagonal log
covariance matrix of the approximate Gaussian posterior only contains small values (≈−17),
thus we may view the encoding process as deterministic and take Eϕ−: X →Z as a map to
the mean of the Gaussian. In this case, (8) reduces to (1/ρ2)[E˜rϕ−(z0|x0)[Dϕ−(z0)] −x0] ≈
(1/ρ2)[Dϕ−(Eϕ(x0))−x0], whose vector Jacobian product with the decoder gradient resembles
the data consistency term ∇zt||E[z0|zt] −E(D(E[z0|zt]))||2 introduced in Rout et al. (2023);
Song et al. (2024).
Adaptive Optimizer
To accelerate convergence on the challenging optimization landscapes
common in imaging inverse problems, we employ a non-standard discretization for the flow
in (5). Instead of a simple Euler step, we treat the drift terms as gradients and apply the
Adam optimizer (Kingma and Ba, 2014). We discuss this choice in Section C.
3
3
Experiments
We evaluate our method DWGF on the FFHQ dataset (Karras et al., 2019) downsampled
to a resolution of 512 × 512. We compare with Posterior Sampling with Latent Diffusion
(PSLD) (Rout et al., 2023), one of the established baselines for solving inverse problems with
a latent diffusion prior and Repulsive Latent Score Distillation (RLSD) (Zilberstein et al.,
2025), which is a recent method based on similar ideas from gradient flows. We adopt same
the experimental settings as in Zilberstein et al. (2025) and report the best results therein.
In this work, we only evaluate our method on box inpainting and super-resolution. We report
standard metrics: Peak Signal-to-Noise Ratio (PSNR), Learned Perceptual Image Patch
Similarity (LPIPS) (Zhang et al., 2018), and FrÃľchet Inception Distance (FID) (Heusel
et al., 2017).
...
As shown in the quantitative results (Table 1), our algorithm achieves comparable performance
in terms of PSNR and LPIPS but suffers from poor FID. Examining the qualtitative results
(Figures 1a and 1b), we conjecture that this is caused by the mode-seeking behavior of the KL
divergence, leading to blurry reconstructions. We point out that additional regularizations
such as entropic regularizations (Wang et al., 2024b) or repulsive potentials (Corso et al.,
2024; Zilberstein et al., 2025; Hu et al., 2025) can be incorporated into our functional in (3)
to tackle mode collapse.

A Gradient Flow Approach to Solving Inverse
Problems with Latent Diffusion Models
Tim Y. J. Wang
Department of Mathematics
Imperial College London
tw1320@ic.ac.uk
O. Deniz Akyildiz
Department of Mathematics
Imperial College London
deniz.akyildiz@imperial.ac.uk
Abstract
Solving ill-posed inverse problems requires powerful and flexible priors. We
propose leveraging pretrained latent diffusion models for this task through
a new training-free approach, termed Diffusion-regularized Wasserstein
Gradient Flow (DWGF). Specifically, we formulate the posterior sampling
problem as a regularized Wasserstein gradient flow of the Kullback-Leibler
divergence in the latent space. We demonstrate the performance of our
method on standard benchmarks using StableDiffusion (Rombach et al.,
2022) as the prior.
1
Introduction
Inverse problems are ubiquitous in science and engineering.
They involve finding the
underlying true signal x0 from the corrupted observation y. In this work, we are primarily
concerned with inverse problems of the form:
y = A(x0) + ϵ,
ϵ ∼N(0, σ2
yI),
(1)
where A : X →Y is a known forward corruption operator and ϵ is the additive Gaussian noise.
Through the lens of Bayesian inference (Stuart, 2010), the solution to the problem in (1) can
be elegantly viewed as sampling from the posterior p(x0|y) ∝p(y|x0)pdata(x0) by placing
a prior on x0. While traditional hand-crafted priors are often chosen for mathematical
convenience, they struggle with high-dimensional, ill-posed inverse problems.
In recent years, alternative priors based on diffusion models (Sohl-Dickstein et al., 2015; Ho
et al., 2020; Song et al., 2021) have become the state-of-the-art for solving imaging inverse
problems. A common strategy is to modify the sampling process of a pixel-space diffusion
model to guide the sampler towards the posterior p(x0|y) through various techniques, for
instance, proximal (Zhu et al., 2023; Wu et al., 2024), gradient-based (Chung et al., 2023;
Song et al., 2023a; Boys et al., 2024), variational inference (Mardani et al., 2024; Zilberstein
et al., 2025), or sequential Monte-Carlo methods (Cardoso et al., 2024; Dou and Song, 2024;
Chen et al., 2025).
However, adapting these techniques to more computationally efficient latent-space diffusion
models (Rombach et al., 2022) is not straightforward, as the model defines a prior p(z0)
on a latent variable instead of the ground truth signal x0. Approaches such as Rout et al.
(2023); Song et al. (2024) adapt diffusion posterior sampling (Chung et al., 2023) to the
latent space with data consistency regularization, while Zilberstein et al. (2025) builds on
RED-Diff (Mardani et al., 2024) and simulate a particle system for an augmented distribution
q(x0, z0|y) defined on the product of pixel and latent space.
In this work, we depart from adapting existing pixel-space methods. Instead, we formulate
the posterior sampling problem from first principles as a Wasserstein gradient flow of a KL
2nd Workshop on Frontiers in Probabilistic Inference: Sampling Meets Learning, 39th Conference
on Neural Information Processing Systems (NeurIPS 2025).
arXiv:2509.19276v1  [stat.ML]  23 Sep 2025
divergence functional directly in the latent space. Using the diffusion prior, we develop a
regularized gradient flow, which we term Diffusion-regularized Wasserstein Gradient Flow
(DWGF). We derive a system of ordinary differential equations (ODEs) that approximates
this flow, providing a principled method for solving inverse problems with latent diffusion
priors.
Notations.
We use X = Rdx to denote the ambient (pixel) space, Y ⊆Rdx the observation
space, and Z = Rdz the latent space. We denote the space of probability measures on Rd
with finite q-th moments as Pq(Rd); in this work, we focus on the cases where q = 2.
2
A Wasserstein Gradient Flow Approach
Assume that we are given a pretrained latent diffusion model, which consists of a decoder
pϕ−(x0|z0), encoder rϕ−(z0|x0), and a diffusion generative model pθ−(z0) in the latent space.
Given the inverse problem defined in (1), we would like to leverage this model to obtain an
approximate posterior distribution, denoted here
qµ(x0|y) =
Z
pϕ−(x0|z0)µ(z0|y)dz0.
(2)
In order to obtain this distribution, one needs to approximate the unknown distribution
µ(z0|y). We thus consider the following regularized optimization problem over P2(Z):
µ⋆(z0|y) ∈arg min
µ∈P2(Z)
F[µ(z0|y)] + γR[µ(z0|y)],
(3)
where F[µ] = DKL (qµ(x0|y)∥p(x0|y)) and R : P2(Z) →R+ can be taken as any regularization
that makes the minimizers of the gradient flow well-defined (Crucinio et al., 2024), and γ > 0
controls the strength of regularization. In order to solve the problem in (3), we adopt a
gradient descent approach in the space of probability measures . The gradient flow of any
functional L[µ] on P2(Z) starting at some ¯µ0 is given by (Figalli and Glaudo, 2021):
∂µt
∂t = ∇z0 ·

µt∇z0
δL[µ]
δµ

,
µ0 = ¯µ0
(4)
The PDE above is the continuity equation corresponds to the following ODE (Ambrosio
et al., 2005, Chapter 8):
dz0,t
dt
= −∇z0
δL[µ]
δµ
= −

∇z0
δF[µ]
δµ
+ γ∇z0
δR[µ]
δµ

,
z0,0 ∼¯µ0,
(5)
where we have set L[µ] := F[µ] + γR[µ]. The solution to (5) defines a flow map Tt(z0,0) = zt
such that µt = (Tt)#¯µ0. To simulate (5), we need to compute the functional derivative
δL[µ]/δµ and its gradient. In what follows, we will construct the regularization and derive
the first variations.
2.1
Wasserstein gradient of DKL (qµ(x0|y)∥p(x0|y))
Wasserstein-2 (W2) gradient of the first term F[µ] = DKL(qµ(x0|y)∥p(x0|y)) can be derived
similarly to Wang et al. (2023). In particular, the first variation is given by (see Section A.1
for the full derivation)
δF[µ]
δµ
= Ep(x0|z0)[log qµ(x0|y) −log p(x0|y)].
(6)
Then to obtain the first term of the drift in the ODE (5), we compute
∇z0
δF[µ]
δµ
:= Eϵ∼N(0,I)

∇x0(log q(gϕ−(ϵ, z0)|y) −log p(gϕ−(ϵ, z0)|y))∂Dϕ−(z0)
∂z0

,
(7)
where we have used the chain rule and the reparameterization trick gϕ−(ϵ, z0) :=
Dϕ−(z0) + ρϵ for the Gaussian decoder. As noted by Wang et al. (2024a), the first term
∇x0 log q(gϕ−(ϵ, z0)|y) is zero due to reparameterization (cf. Section A.3).
2
We note that the only intractable term in (7) is the posterior score ∇x0 log p(x0|y), which
admits the decomposition ∇x0 log p(x0|y) = ∇x0 log p(y|x0) + ∇x0 log p(x0). While the first
term ∇x0 log p(y|x0) is tractable for being the gradient of the Gaussian likelihood (1), the
data score ∇x0 log p(x0) requires approximation. Assuming regularity conditions that permit
interchanging the gradient and integral, we can express the prior score as an expectation,
which we approximate using the encoder of the pretrained latent diffusion model:
∇x0 log p(x0) =
Z
Z
[∇x0 log p(x0|z0)]p(z0|x0)dz0 ≈
Z
Z
[∇x0 log p(x0|z0)]˜rϕ−(z0|x0)dz0,
(8)
where ˜rϕ−(z0|x0) is the approximate posterior distribution given by a Variational Autoencoder
(VAE) (Kingma and Welling, 2013), which is part of the pretrained latent diffusion model.
2.2
Wasserstein gradient of R[µ(z0|y)]
To fully leverage the information provided by the generative model, we first construct
the prior regularization term R(µ(z0|y)) based on the pretrained latent diffusion model.
Analogous to Wang et al. (2023); Luo et al. (2023), we consider a weighted KL divergence
along the diffusion process:
R(µ(z0|y)) := Dw,[0,T ]
KL
(µ(z0|y)∥pθ−(z0)) =
Z T
0
w(s)DKL(µ(zs|y)∥pθ−(zs))ds,
(9)
where w(s) : [0, T] →R+ is a time-dependent weighting term and the densities µ(zs|y) :=
R
Z N(zs; αsz0, σ2
sI)µ(z0|y)dz0 are pushforwards of µ through the forward transition kernel
of the diffusion model. We remark that the weighted KL divergence in (9) is a well-defined
functional and admits favorable properties (cf. Section A.4) as summarised below.
Theorem 2.1. The weighted KL divergence Dw,[0,T ]
KL
(µ(z0|y)∥pθ−(z0)) (9) is i) nonnegative,
ii) convex in the first component µ(z0|y), and iii) is minimized if and only if the standard
KL divergence DKL(µ(z0|y)∥pθ−(z0)) is minimized.
Accordingly, we obtain the gradient of the first variation of R (cf. Section A.2):
∇z0
δR[µ]
δµ
= Es,ϵ

˜w(s)

∇zs log
Z
Z
p(zs|z0)µ(z0,t|y)dz0 −∇zs log pθ−(zs)
 ∂zs
∂z0

,
(10)
where the expectation is taken over s ∼U(0, T), ϵ ∼N(0, I), we set ˜w(s) := Tw(s), and
p(zs|z0) = N(zs; αsz0, σ2
sI) is the forward kernel of the diffusion model. To approximate
the integral in (10), we use a particle-based approach to form a Monte-Carlo approximation
using {z(i)
0 }i ∼µ(z0|y) (Wang et al., 2023; Lim and Johansen, 2024).
2.3
Algorithmic Considerations
Final Gradient Flow. We can simply combine the gradients derived in (7) and (10) to
obtain the final gradient flow in (5). As we have an intractable term, we will simulate N
identical copies of this ODE, where the integral in (10) is approximated using these particles
as mentioned. We mention some practical aspects of this implementation below.
Deterministic Encoding
During our experiments, we observe that the diagonal log
covariance matrix of the approximate Gaussian posterior only contains small values (≈−17),
thus we may view the encoding process as deterministic and take Eϕ−: X →Z as a map to
the mean of the Gaussian. In this case, (8) reduces to (1/ρ2)[E˜rϕ−(z0|x0)[Dϕ−(z0)] −x0] ≈
(1/ρ2)[Dϕ−(Eϕ(x0))−x0], whose vector Jacobian product with the decoder gradient resembles
the data consistency term ∇zt||E[z0|zt] −E(D(E[z0|zt]))||2 introduced in Rout et al. (2023);
Song et al. (2024).
Adaptive Optimizer
To accelerate convergence on the challenging optimization landscapes
common in imaging inverse problems, we employ a non-standard discretization for the flow
in (5). Instead of a simple Euler step, we treat the drift terms as gradients and apply the
Adam optimizer (Kingma and Ba, 2014). We discuss this choice in Section C.
3
3
Experiments
We evaluate our method DWGF on the FFHQ dataset (Karras et al., 2019) downsampled
to a resolution of 512 × 512. We compare with Posterior Sampling with Latent Diffusion
(PSLD) (Rout et al., 2023), one of the established baselines for solving inverse problems with
a latent diffusion prior and Repulsive Latent Score Distillation (RLSD) (Zilberstein et al.,
2025), which is a recent method based on similar ideas from gradient flows. We adopt same
the experimental settings as in Zilberstein et al. (2025) and report the best results therein.
In this work, we only evaluate our method on box inpainting and super-resolution. We report
standard metrics: Peak Signal-to-Noise Ratio (PSNR), Learned Perceptual Image Patch
Similarity (LPIPS) (Zhang et al., 2018), and FrÃľchet Inception Distance (FID) (Heusel
et al., 2017).
Inpainting (Box)
SR (×8)
Method
FID ↓PSNR ↑
LPIPS ↓
FID ↓PSNR ↑LPIPS ↓
PSLD (Rout et al., 2023)
57.70
22.72
0.082
81.31
24.82
0.314
RLSD (Zilberstein et al., 2025)
29.18
24.98
0.079
65.42
28.39
0.286
Ours
118.05
27.56
0.184
101.94
32.71
0.193
Table 1: Results for large box inpainting and 8× super-resolution, all with noise σy = 0.001
on the FFHQ-512 validation dataset.
True
Degraded
Particle 1
Particle 2
Particle 3
Particle 4
(a) Diversity of the particles x0 produced by DWGF on large
box inpainting.
SRx8
True
Degraded
DWGF
PSLD
RLSD
Inpaint
(b) Qualitative comparison of the DWGF, PSLD, and RLSD
on inpainting and super-resolution (8×) tasks.
Figure 1: Qualitative results on the FFHQ dataset downsampled to 512 × 512 resolution.
As shown in the quantitative results (Table 1), our algorithm achieves comparable performance
in terms of PSNR and LPIPS but suffers from poor FID. Examining the qualtitative results
(Figures 1a and 1b), we conjecture that this is caused by the mode-seeking behavior of the KL
divergence, leading to blurry reconstructions. We point out that additional regularizations
such as entropic regularizations (Wang et al., 2024b) or repulsive potentials (Corso et al.,
2024; Zilberstein et al., 2025; Hu et al., 2025) can be incorporated into our functional in (3)
to tackle mode collapse.
4
Conclusion
We proposed DWGF, a novel approach of solving inverse problems using latent diffusion
models as priors. We note that our approach can be extended to the conditional setting where
we jointly optimize the prompt embedding as in Spagnoletti et al. (2025) using Euclidean-
Wasserstein gradient flows (Kuntz et al., 2023). Additionally, it would be interesting to
explore control variates (Wang et al., 2024a) for variance reduction.
4
Acknowledgements
We thank the anonymous reviewers for their constructive comments. TW is supported by
the Roth Scholarship from the Department of Mathematics, Imperial College London.
References
Akyildiz, O. D., Girolami, M., Stuart, A. M., and Vadeboncoeur, A. (2025). Efficient prior
calibration from indirect data. SIAM Journal on Scientific Computing, 47(4):C932–C958.
Ambrosio, L., Gigli, N., and Savaré, G. (2005). Gradient flows: in metric spaces and in the
space of probability measures. Springer.
Boys, B., Girolami, M., Pidstrigach, J., Reich, S., Mosca, A., and Akyildiz, O. D. (2024).
Tweedie moment projected diffusions for inverse problems. Transactions on Machine
Learning Research. Featured Certification.
Cardoso, G., el idrissi, Y. J., Corff, S. L., and Moulines, E. (2024). Monte carlo guided de-
noising diffusion models for bayesian linear inverse problems. In The Twelfth International
Conference on Learning Representations.
Chen, H., Ren, Y., Min, M. R., Ying, L., and Izzo, Z. (2025). Solving inverse problems via
diffusion-based priors: An approximation-free ensemble sampling approach. arXiv preprint
arXiv:2506.03979.
Chung, H., Kim, J., Mccann, M. T., Klasky, M. L., and Ye, J. C. (2023). Diffusion posterior
sampling for general noisy inverse problems. In The Eleventh International Conference on
Learning Representations.
Corso, G., Xu, Y., Bortoli, V. D., Barzilay, R., and Jaakkola, T. S. (2024). Particle guidance:
non-i.i.d. diverse sampling with diffusion models. In The Twelfth International Conference
on Learning Representations.
Crucinio, F. R., De Bortoli, V., Doucet, A., and Johansen, A. M. (2024). Solving a class
of fredholm integral equations of the first kind via wasserstein gradient flows. Stochastic
Processes and their Applications, 173:104374.
Dou, Z. and Song, Y. (2024).
Diffusion posterior sampling for linear inverse problem
solving: A filtering perspective. In The Twelfth International Conference on Learning
Representations.
Engel, E., Dreizler, R. M., and Dreizler, R. M. (2011). Density Functional Theory: An
Advanced Course. Theoretical and Mathematical Physics. Springer Nature.
Figalli, A. and Glaudo, F. (2021). An invitation to optimal transport, Wasserstein distances,
and gradient flows. EMS textbooks in mathematics. European Mathematical Society.
Geng, Z., Deng, M., Bai, X., Kolter, J. Z., and He, K. (2025). Mean flows for one-step
generative modeling. arXiv preprint arXiv:2505.13447.
Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A. (2012). A kernel
two-sample test. Journal of Machine Learning Research, (25):723–773.
Hagemann, P., Hertrich, J., AltekrÃĳger, F., Beinert, R., Chemseddine, J., and Steidl, G.
(2024). Posterior sampling based on gradient flows of the MMD with negative distance
kernel. In The Twelfth International Conference on Learning Representations.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. (2017). Gans
trained by a two time-scale update rule converge to a local nash equilibrium. In Proceedings
of the 31st International Conference on Neural Information Processing Systems.
Ho, J., Jain, A., and Abbeel, P. (2020). Denoising diffusion probabilistic models. In Advances
in Neural Information Processing Systems.
5
Hu, Y., Mei, K., Sahraee-Ardakan, M., Kamilov, U. S., Milanfar, P., and Delbracio, M. (2025).
Kernel density steering: Inference-time scaling via mode seeking for image restoration.
arXiv preprint arXiv:2507.05604.
Karras, T., Laine, S., and Aila, T. (2019). A style-based generator architecture for generative
adversarial networks. In 2019 IEEE/CVF Conference on Computer Vision and Pattern
Recognition (CVPR), pages 4396–4405.
Kingma, D. P. and Ba, J. (2014). Adam: A method for stochastic optimization. International
Conference on Learning Representations.
Kingma, D. P. and Welling, M. (2013).
Auto-encoding variational bayes.
CoRR,
abs/1312.6114.
Kuntz, J., Lim, J. N., and Johansen, A. M. (2023). Particle algorithms for maximum
likelihood training of latent variable models. In Proceedings of The 26th International
Conference on Artificial Intelligence and Statistics, volume 206 of Proceedings of Machine
Learning Research.
Lim, J. N. and Johansen, A. M. (2024). Particle semi-implicit variational inference. In The
Thirty-eighth Annual Conference on Neural Information Processing Systems.
Lim, J. N., Kuntz, J., Power, S., and Johansen, A. M. (2024). Momentum particle maximum
likelihood. In Forty-First International Conference on Machine Learning.
Luo, W., Hu, T., Zhang, S., Sun, J., Li, Z., and Zhang, Z. (2023). Diff-instruct: A universal
approach for transferring knowledge from pre-trained diffusion models. In Thirty-seventh
Conference on Neural Information Processing Systems.
Mardani, M., Song, J., Kautz, J., and Vahdat, A. (2024). A variational perspective on
solving inverse problems with diffusion models. In The Twelfth International Conference
on Learning Representations.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. (2022). High-resolution
image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference
on Computer Vision and Pattern Recognition (CVPR).
Rout, L., Raoof, N., Daras, G., Caramanis, C., Dimakis, A., and Shakkottai, S. (2023).
Solving linear inverse problems provably via posterior sampling with latent diffusion
models. In Thirty-seventh Conference on Neural Information Processing Systems.
Santambrogio, F. (2016). Euclidean, Metric, and Wasserstein gradient flows: an overview.
arXiv preprint arXiv:1609.03890.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. (2015). Deep un-
supervised learning using nonequilibrium thermodynamics. In Proceedings of the 32nd
International Conference on Machine Learning, volume 37 of Proceedings of Machine
Learning Research.
Song, B., Kwon, S. M., Zhang, Z., Hu, X., Qu, Q., and Shen, L. (2024). Solving inverse prob-
lems with latent diffusion models via hard data consistency. In The Twelfth International
Conference on Learning Representations.
Song, J., Vahdat, A., Mardani, M., and Kautz, J. (2023a). Pseudoinverse-guided diffusion
models for inverse problems. In International Conference on Learning Representations.
Song, Y., Dhariwal, P., Chen, M., and Sutskever, I. (2023b). Consistency models. Interna-
tional Conference on Machine Learning.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. (2021).
Score-based generative modeling through stochastic differential equations. In International
Conference on Learning Representations.
6
Spagnoletti, A., Prost, J., Almansa, A., Papadakis, N., and Pereyra, M. (2025). LATINO-
PRO: LAtent consisTency INverse sOlver with PRompt optimization. arXiv preprint
arXiv:2503.12615.
Stuart, A. M. (2010). Inverse problems: a bayesian perspective. Acta numerica, 19:451–559.
Vadeboncoeur, A., Girolami, M., and Stuart, A. M. (2025). Efficient deconvolution in
populational inverse problems. arXiv preprint arXiv:2505.19841.
Villani, C. (2003). Topics in optimal transportation. Graduate studies in mathematics ; v.
58. American Mathematical Society.
Wang, P., Fan, Z., Xu, D., Wang, D., Mohan, S., Iandola, F., Ranjan, R., Li, Y., Liu, Q.,
Wang, Z., and Chandra, V. (2024a). SteinDreamer: Variance reduction for text-to-3D
score distillation via stein identity. arXiv preprint arXiv:2401.00604.
Wang, P., Xu, D., Fan, Z., Wang, D., Mohan, S., Iandola, F., Ranjan, R., Li, Y., Liu,
Q., Wang, Z., and Chandra, V. (2024b). Taming mode collapse in score distillation for
text-to-3D generation. In 2024 IEEE/CVF Conference on Computer Vision and Pattern
Recognition (CVPR), Seattle, WA, USA.
Wang, Z., Lu, C., Wang, Y., Bao, F., Li, C., Su, H., and Zhu, J. (2023). Prolificdreamer:
High-fidelity and diverse text-to-3d generation with variational score distillation.
In
Thirty-seventh Conference on Neural Information Processing Systems.
Wu, H., He, L., Zhang, M., Chen, D., Luo, K., Luo, M., Zhou, J.-Z., Chen, H., and Lv, J.
(2024). Diffusion posterior proximal sampling for image restoration. In ACM Multimedia
2024.
Xie, S., Xiao, Z., Kingma, D. P., Hou, T., Wu, Y. N., Murphy, K. P., Salimans, T., Poole, B.,
and Gao, R. (2024). EM distillation for one-step diffusion models. In The Thirty-eighth
Annual Conference on Neural Information Processing Systems.
Yin, T., Gharbi, M., Park, T., Zhang, R., Shechtman, E., Durand, F., and Freeman, W. T.
(2024a). Improved distribution matching distillation for fast image synthesis. In The
Thirty-eighth Annual Conference on Neural Information Processing Systems.
Yin, T., Gharbi, M., Park, T., Zhang, R., Shechtman, E., Durand, F., and Freeman, W. T.
(2024b). Improved distribution matching distillation for fast image synthesis. arXiv
preprint arXiv:2405.14867.
Yin, T., Gharbi, M., Zhang, R., Shechtman, E., Durand, F., Freeman, W. T., and Park, T.
(2024c). One-step diffusion with distribution matching distillation. In 2024 IEEE/CVF
Conference on Computer Vision and Pattern Recognition (CVPR).
Zhang, R., Isola, P., Efros, A. A., Shechtman, E., and Wang, O. (2018). The unreasonable
effectiveness of deep features as a perceptual metric. arXiv preprint arXiv:1801.03924.
Zhu, Y., Zhang, K., Liang, J., Cao, J., Wen, B., Timofte, R., and Van Gool, L. (2023).
Denoising Diffusion Models for Plug-and-Play Image Restoration. In 2023 IEEE/CVF
Conference on Computer Vision and Pattern Recognition Workshops (CVPRW).
Zilberstein, N., Mardani, M., and Segarra, S. (2025). Repulsive latent score distillation
for solving inverse problems. In The Thirteenth International Conference on Learning
Representations.
7
A
Derivations
Notations. We denote the space of probability measures on Rd with finite q-th moments
as Pq(Rd). We equip the space P2(Rd) with the usual Wasserstein-2 scalar product (Figalli
and Glaudo, 2021) and denote its tangent space as TP2(Rd). For simplicity, we assume
all distributions considered in this work admit differentiable densities with respect to the
Lebesgue measure.
A.1
Derivation of the first variation of F
We first state a standard result on the first variation of the KL divergence (cf. Villani (2003);
Santambrogio (2016); Figalli and Glaudo (2021)):
Lemma A.1. For p, q ∈P2(Rd), the first variation of the KL divergence is given by:
δDKL(q∥p)
δq
[x] = log q(x) −log p(x) + 1,
∀x ∈Rd
(11)
Proof. Using the equivalent definition of first variation, we can write for m ∈TP2(Rd) and
t > 0:
DKL(q + tm∥p) = DKL(q∥p) + t

m, δDKL(q∥p)
δq

+ o(t),
(12)
where the inner product is defined ⟨m, f⟩:=
R
Rd f(z)m(z)dz for all f, m ∈TP2(Rd). Using
Taylor expansion of (z+t) log(z+t) = z log z+t(log z +1)+o(t), we can write DKL(q+tm∥p)
as:
DKL(q + tm∥p) =
Z
(q(x) + tm(x))[log(q(x) + tm(x)) −log p(x)]dx
(13)
=
Z
q(x) log q(x)dx −
Z
q(x) log p(x)dx
(14)
+ t
Z
(log q(x) −log p(x) + 1)m(x)dx + o(t),
(15)
which shows the desired result by matching the terms.
Lemma A.2. For F : P2(Z) →R+ with F[µ] :=
R
w(z)µ(z)dz (cf. Section 2) for any fixed
w(z) : Rd →R, we have:
δF[µ]
δµ
[z] = w(z)
(16)
Proof. Similar to the proof above, for any m ∈TP2(Z) and t ∈R, we have:
F[µ + tm] =
Z
w(z)(µ(z) + tm(z))dz
(17)
=
Z
w(z)µ(z)dz + t
Z
w(z)m(z)dz.
(18)
From the definition of the functional derivative, this implies the desired result.
Proposition A.3. For qµ(x0|y) =
R
p(x0|z0)µ(z0|y)dz0 and p(x0|y) both in P2(Rd) we have:
δDKL(qµ(x0|y)∥p(x0|y))
δµ
= Ep(x0|z0)[log qµ(x0|y) −log p(x0|y)]
(19)
Proof. Using the chain rule of functional derivatives (Engel et al., 2011, Appendix A.3) and
the two lemmas above, we obtain:
δDKL(qµ(x0|y)∥p(x0|y))
δµ
=
Z δDKL(qµ(x0|y)∥p(x0|y))
δq
[x] · δqµ(x0|y)
δµ
[z0]dx
(20)
=
Z
[log qµ(x0|y) −log p(x0|y)]p(x0|z0)dx,
(21)
where we have set w(z0) = p(x0|z0) in Theorem A.2 as the Gaussian decoder distribution.
8
We now derive the gradient of the first variation.
Proof of (7). By Theorem A.3 above, we see that
δDKL(qµ(x0|y)∥p(x0|y))
δµ
= Ep(x0|z0)[log qµ(x0|y) −log p(x0|y)].
(22)
Assume standard regularity assumptions, we can exchange the expectation and the gradient.
An application of the reparameterization trick gϕ−(ϵ, z0) := Dϕ−(z0) + ρϵ yields:
∇z0
δF[µ]
δµ
= ∇z0Ep(z0|x0)[log qµ(x0|y) −log p(x0|y)]
(23)
= Eϵ∼N(0,I)[∇z0(log qµ(gϕ−(ϵ, z0,t)|y) −log p(gϕ−(ϵ, z0,t)|y))]
(24)
= Eϵ∼N(0,I)

∇x0(log qµ(gϕ−(ϵ, z0,t)|y) −log p(gϕ−(ϵ, z0,t)|y))∂Dϕ−(z0,t)
∂z0,t

(25)
where we have used the chain rule in the last step.
A.2
Derivation of the first variation of R
Similar to Theorem A.1 above, we can derive the first variation of the weighted KL divergence
(9). Recall that the weighted KL divergence for p, q ∈P2(Rd) is given by:
Dw,[0,T ]
KL
(q∥p) =
Z
[0,T ]
w(s)DKL(qs∥ps)ds,
(26)
where both qs and ps are defined as pushforwards through a Markov kernel: qs(xs) :=
R
p(xs|x0)q(x0)dx0, where we define p(xs|x0) := N(xs; αsx0, σ2
sI).
Proposition A.4. For p, q ∈P2(Rd), the first variation of the weighted KL divergence is
given by:
δDw,[0,T ]
KL
(q∥p)
δq
[x] =
Z
[0,T ]
w(s)Ep(xs|x0)[log qs −log ps]]ds,
∀x ∈Rd
(27)
Proof. Applying the chain rule for functional derivatives again, we obtain:
δDw,[0,T ]
KL
(q∥p)
δq
[x] =
Z δDw,[0,T ]
KL
(q∥p)
δDKL(qs∥ps) [s]δDKL(qs∥ps)
δq
[z]ds
(28)
=
Z
w(s)Ep(xs|x0)[log qs −log ps]ds
(29)
where we have used Theorem A.2 to obtain the first variation of the first term and the proof
of Theorem A.3 for the second term.
We can now derive the gradient of the first variation in (10).
Proof of (10). Substituting the definition for p = pθ−(z0), q = µ(z0|y) with p(zs|z0) =
N(zs; αsz0, σ2
sI) into Theorem A.4 and taking the gradient yields:
∇z0
δR[µ]
δµ
= ∇z0
Z
[0,T ]
w(s)Ep(zs|z0) [(log µ(zs|y) −log pθ−(zs))] ds
(30)
=
Z
[0,T ]
w(s)∇z0Eϵ∼N (0,I) [(log µ(gs(ϵ, z0)|y) −log pθ−(gs(ϵ, z0)))] ds
(31)
=
Z
[0,T ]
w(s)Eϵ∼N(0,I) [(∇z0 log µ(gs(ϵ, z0)|y) −∇z0 log pθ−(gs(ϵ, z0)))] ds
(32)
=
Z
[0,T ]
w(s)Eϵ∼N(0,I)

(∇zs log µ(gs(ϵ, z0)|y) −∇zs log pθ−(gs(ϵ, z0))) ∂zs
∂z0

,
(33)
9
where the second line follows from the reparameterization trick gs(ϵ, z0) = αsz0 + σsϵ, the
third line follows from an exchange of the gradient and the integral, and we used the chain
rule in the final line.
The marginal density µ(gs(ϵ, z0)|y) = µ(zs|y) can be approximated using an integral
R
Z p(zs|z0)µ(z0,t|y)dz0, whence (10) follows.
A.3
Decoder Reparameterization
We now show that the gradient ∇z0Ep(x0|z0)[log qµ(x0|y)] is in fact zero with a reparametriz-
able Gaussian decoder p(x0|z0) (Wang et al., 2024a). Assuming standard regularity conditions,
we have:
∇z0Ep(x0|z0)[log qµ(x0|y)] = Eϵ∼N (0,I)

∇z0 log N(x0; Dϕ−(z0), ρ2I)

(34)
= Eϵ∼N (0,I)

∇z0 −
1
2ρ2 ∥Dϕ−(z0) + ρϵ −Dϕ−(z0)∥2

= 0.
(35)
A.4
Proof to Theorem 2.1
We restate the theorem here for convenience:
Theorem A.5. The weighted KL divergence Dw,[0,T ]
KL
(µ(z0|y)∥pθ−(z0)) (9) is i) nonnegative,
ii) convex in the first component µ(z0|y), and iii) is minimized if and only if the standard
KL divergence DKL(µ(z0|y)∥pθ−(z0)) is minimized.
Proof. We note that the first two properties follows from those of the KL divergence. For the
third property, we note that the weighted KL is zero if and only if DKL(µ(zt|y)∥pθ−(zt)) = 0
for every t, since the weights w(t) are nonnegative. Now using the argument in Wang et al.
(2023), we see that µ(zt|y) = pθ−(zt) if and only if their characteristic functions are equal
φµ(zt|y)(s) = φpθ−(zt)(s). But we have:
φµ(zt|y)(s) = φµ(z0|y)(αts) · φN (0,I)(σts)
(36)
φpθ−(zt)(s) = φpθ−(z0)(αts) · φN (0,I)(σts),
(37)
since we have from the forward process zt = αtz0 + σtϵ for ϵ ∼N(0, I). The third property
thus follows from the positivity of the KL divergence.
B
Algorithm
We now give the pseudocode for simulating the gradient flow in Algorithm 1 below; we use
the notation zs,k to denote the particles diffused to the time s ∈[0, T] and at step k ∈N
of the Euler discretization of the ODE (5). Note that instead of approximating the drift
corresponding to the diffusion regularization in (10) by sampling time s uniformly from
[0, T], we follow Mardani et al. (2024) and adopt a deterministic schedule for sampling s
to be s ∈{T, T −1, · · · , 0}. In our experiments, we take w(s) = cσ2
s/αs with a constant
c ∈(0, 1) similarly to Mardani et al. (2024) .
C
Experimental Details
Implementation
In practice, we use the Adam optimizer (Kingma and Ba, 2014) with
a learning rate of lr = 1.0 and default hyperparameters (β1, β2) = (0.9, 0.999) to solve the
ODE as in Algorithm 1; this is done by viewing u(z(i)
0,k) + γv(z(i)
0,k) as the gradient of a loss
function to be optimized. This approach effectively introduces momentum and an adaptive
diagonal preconditioning at each step, which helps mitigate ill-conditioning and flat minima.
While unconventional for numerical solutions to ODEs, this heuristic is supported by recent
work that formally incorporates momentum (Lim et al., 2024) and preconditioning (Lim and
Johansen, 2024) into the simulation of gradient flows.
We set the balancing coefficient γ = 0.15 and the data consistency weight to λ = 0.1ρ2
(cf. Section 2.3), where we choose the decoder standard deviation to be ρ = 10−3. We
10
Algorithm 1 Diffusion-regularized Wasserstein Gradient Flow (DWGF)
1: Inputs: Observation y ∈Y, differentiable forward operator A : X →Y, initial particles
{z(i)
0,0}N
i=1, diffusion weights w(t), regularization strength γ, data consistency weight λ,
pretrained latent diffusion model with score sθ−(·, ·) : Z × R+ →Z and its VAE with
encoder-decoder pair (Eϕ−, Dϕ−)
2: Set counter k ←0
3: for s ∈{T, T −1, · · · , 0} do
4:
Sample ϵ(i) ∼N(0, I) and decode x(i)
0
←gϕ−(ϵ, z(i)
0,k) for i ∈[N]
5:
Data likelihood ∇x(i)
0 ℓ(x(i)
0 ) ←∇x(i)
0 (−
1
2σ2y ∥y −A(x(i)
0 )∥2
2)
6:
u(z(i)
0,k) ←

−λ
ρ2 [Dϕ−(Eϕ−(x(i)
0 ))] −x(i)
0 ] −∇x(i)
0 ℓ(x(i)
0 )
 ∂Dϕ−(z(i)
0,k)
∂z(i)
0,k
(38)
7:
Sample ν(i) ∼N(0, I) for all i ∈[N]
8:
Compute z(i)
s,k ←αsz(i)
0,k + σsν(i)
9:
v(z(i)
0,k) ←w(s)

∇z(i)
s,k log 1
N
N
X
j=1
N(z(i)
s,k; αsz(j)
0,k, σ2
sI) −sθ−(z(i)
s,k, s)

∂z(i)
s,k
∂z(i)
0,k
(39)
10:
Simulate z(i)
0,k+1 ←OptimizerStep(z(i)
0,k, u(z(i)
0,k) + γv(z(i)
0,k)) for all i ∈[N]
11:
k ←k + 1
12: end for
13: return Particles decoded {Dϕ−(z(i)
0,k)}N
i=1
choose Stable Diffusion v2.1 (Rombach et al., 2022) as the base model and set the number
of sampling steps T = 999, which are the same as Zilberstein et al. (2025).
Evaluation
We follow the experimental setups in Zilberstein et al. (2025) and evaluate our
methods on the first 100 images from the validation set of FFHQ (Karras et al., 2019) using
4 particles. Due to limited computational budget, we do not reproduce the experiments in
Zilberstein et al. (2025) but choose to report their best results for each task therein, namely
the non-repulsive version of RLSD.
D
Further Discussion and Related Works
Further Discussions
Despite potentials of DWGF in inverse problems, it has some
significant limitations. Firstly, our method does not yet match the performance of state-of-
the-art in terms of FID (Section 3), which we attribute to the lack of proper regularization
and insufficient hyperparameter tuning. Secondly, analogous to other variational approaches
(Mardani et al., 2024; Zilberstein et al., 2025), DWGF requires a large number of sampling
steps. A promising research direction is thus to integrate our framework with recent advances
in few-step and consistency models(Song et al., 2023b; Luo et al., 2023; Yin et al., 2024a;
Geng et al., 2025), which are emerging contenders to traditional iterative models. Finally,
our method uses a particle cloud to approximate the intractable integral in solving the
Wasserstein gradient flow, which may lead to considerable memory requirements.
Related Works
Gradient flow approaches have been adopted in the context of general
inverse problems in recent approaches (Akyildiz et al., 2025; Vadeboncoeur et al., 2025). In
imaging, Hagemann et al. (2024) uses a gradient flow of the Maximum Mean Discrepancy
(MMD) (Gretton et al., 2012), which is parametrized as a sequence of pushforward maps.
Zilberstein et al. (2025) uses an interacting particle system approach with repulsive potential,
but they compute the marginal score of the latent distribution ∇zt log q(z(i)
t |y) as −(z(i)
t
−
11
αtz(i)
0 )/σ2
t = −ϵ(i)/σt, which only holds when q(z0|y) is Gaussian, hence their approach
reduces to a Bures-Wasserstein gradient flow, similar to RED-Diff (Mardani et al., 2024).
Our work is most related to Wang et al. (2023), which treat the parameters of an 3D
representation MLP as particles to be optimized via Wasserstein gradient flow. However,
their method operates in a conditional setting and involves training another network to
approximate the score. Also related to our approach are works on score-distillation, which
involve the training of a generator pθ to approximate the output of a diffusion model (Yin
et al., 2024c,b; Luo et al., 2023; Xie et al., 2024).
12


Question:
Based on the understanding of the entire paper, among the following statements regarding the approximate calculation of Wasserstein distance matrices and their applications, the following are correct:
A. The matrix completion method (W₂-MC) is inferior to the Nyström method (W₂-Nystrom) only under specific conditions.
B. W₂-Nyström has been proven to be a better choice due to its higher accuracy, stronger theoretical guarantees, and lower computational complexity.
C. This study found that the biconcentric matrix B of the Wasserstein distance matrix has the same characteristics as the original distance matrix D. Therefore, the feature decomposition of D can be directly performed to simplify the MDS calculation, and based on this, a fast classification pipeline was designed.
D. W₂-Nyström employs a full-column sampling method, which retains more global structural information compared to scattered sampling. Moreover, it is approximately a single-step analytical process, thereby reducing the convergence problems and instability that may arise from iterative optimization.

Answer: B

Here we describe two algorithms for estimating the Wasserstein distance matrix to compute em-
beddings of images. Algorithm 1, W2-MC, randomly computes a given number of entries of the
distance matrix, and then uses a matrix completion algorithm to estimate the full squared distance
matrix. We leave the completion algorithm as a generic plug-in for existing algorithms, and consider
one designed for distance matrices in our experiments. We sample from the strict upper triangle,
include the diagonal, and reflect the entries across the diagonal to obtain the observed distance
matrix. We apply the same procedure after completion to ensure that the resulting matrix retains
the properties of a distance matrix.
Algorithm 2, W2-Nystr¨om, uses the Nystr¨om method to complete the distance matrix.
In
particular, we compute only a set of columns of the distance matrix Cobs = D(:, I), and estimate
Dest = CobsUC⊤
obs where U = Cobs(I, :) = D(I, I). As in W2-MC, we also enforce the distance
matrix structure by enforcing the 0 diagonal.
Note that for large N, the computational cost of both algorithms is dominated by the cost of
computing the W2 distances, and the cost of the matrix completion subroutine MC is often negligible
by comparison. Thus, algorithms that limit sampling complexity of the matrix are more desirable.
For this reason, we do not truncate the rank of the Nystr¨om approximation to D in W2-Nystr¨om,
which is usually done to save computation time.
...
4.1
Fixed sample budget
We first do a fixed sample budget comparison of W2-MC and W2-Nystr¨om for completing the
Wasserstein distance matrix of 2, 000 points from the OrganCMNIST dataset from the MedMNIST
benchmark [21], comprising coronal plane CT scans of 11 different organs downsampled to size 28×
28. For uniform sampling, we subsample the matrix at different rates γ ∈{3%, 5%, 10%, 20%, 25%}.
We then match the number of observed entries between uniform sampling and the Nystr¨om method
by solving the equation c(c−1)
2
+ c(n −c) = γ · n(n−1)
2
for c (the number of columns selected by
Nystr¨om). This equation equates the number of off-diagonal entries used by both methods. For
MC in W2-MC, we use the factorized Gram matrix formulation from [14] based on the relationship
between the Gram matrix and the square distance matrix:
Find
P ∈Rn×q,
s.t.
(PP ⊤)ii + (PP ⊤)jj −2(PP ⊤)ij = Dij,
P ⊤1 = 0,
(2)
where q denotes the estimate for the rank of the underlying points. The centering constraint can
be removed by reparameterization. Defining A : Rn×n →R|Ω| as
A(X)α = Xii + Xjj −2Xij,
α = (i, j) ∈Ω,
the augmented Lagrangian is
L(Q; Λ) = 1
2

A(QQ⊤) −b + Λ

2
2 ,
(3)
with b is the vector of observed distances and Λ ∈R|Ω|.
The algorithm starts with an initial
estimate of Q (randomly initialized in our experiments) and then alternates between two steps:
updating Q via Barzilai–Borwein descent, and updating the Lagrange multipliers. For the dataset
in consideration, we set the rank estimate to be 200.
Table 1 reports the results averaged over 10 random trials. We observe that for fixed sample
budget, the Nystr¨om method consistently yields better accuracy compared to matrix completion.
While both Nystr¨om and typical matrix completion algorithms scale similarly up to logarithmic
factors, the Nystr¨om method is a single-step procedure, whereas the cost of the matrix completion
algorithm depends on the target accuracy, with the number of iterations required to reach a given
error often dominating the runtime. In our experiments, we run W2-MC for a sufficient number
of iterations, as we consistently observe stable convergence across trials. In practice, choosing a
suitable stopping criterion for matrix completion problems could also be challenging. Therefore, this
fixed sample budget experiment leads us to conclude that computing full columns of W2 distance
measurements may be preferred for approximating Wasserstein distance matrix. Studying how far
this generalizes is the subject of future work.

Recovering Wasserstein Distance Matrices from Few Measurements
Muhammad Rana1, Abiy Tasissa3, HanQin Cai4, Yakov Gavriyelov1, and Keaton Hamm1,2
1Department of Mathematics, University of Texas at Arlington
2Division of Data Science, University of Texas at Arlington
3Department of Mathematics, Tufts University, Medford, MA 02155, USA.
4Department of Statistics and Data Science and Department of Computer Science,
University of Central Florida
Abstract
This paper proposes two algorithms for estimating square Wasserstein distance matrices from
a small number of entries. These matrices are used to compute manifold learning embeddings
like multidimensional scaling (MDS) or Isomap, but contrary to Euclidean distance matrices, are
extremely costly to compute. We analyze matrix completion from upper triangular samples and
Nystr¨om completion in which O(d log(d)) columns of the distance matrices are computed where
d is the desired embedding dimension, prove stability of MDS under Nystr¨om completion, and
show that it can outperform matrix completion for a fixed budget of sample distances. Finally,
we show that classification of the OrganCMNIST dataset from the MedMNIST benchmark is
stable on data embedded from the Nystr¨om estimation of the distance matrix even when only
10% of the columns are computed.
1
Introduction
The bare representation of data in many applications is high-dimensional [1]. In today’s era of
big data, this high dimensionality poses significant challenges and necessitates specialized methods.
Manifold learning refers to a class of techniques that compress data representations into lower-
dimensional forms, making them more suitable for analysis and computation. A widely used linear
technique is principal component analysis (PCA). However, linear approaches can be restrictive
when the data lie on or near a nonlinear manifold. To address this, nonlinear manifold learning
techniques aim to obtain compact representations while preserving local geometric structures [2].
We note that, in the standard way they are used, all of the aforementioned methods typically treat
data as vectors in a Euclidean space.
While many applications yield vector data, there are a number of domains that produce data
that may be best modeled as distributions or measures. For example, gene expression and natural
language processing data can be represented as distributions over gene networks or dictionaries,
respectively, and image data can be considered as discrete distributions over a 2-dimensional pixel
grid. Recent works have used this notion to improve understanding of data geometry and per-
formance on tasks such as classification or clustering [3, 4, 5]. One of the advantages of treating
images, for example, as measures, is that spatial relationships are not obscured as they are under
vectorization. Moreover, the natural notion of a metric structure on the space of measures is that
of the Wasserstein metric from optimal transport (OT) [6].
1
arXiv:2509.19250v1  [stat.ML]  23 Sep 2025
Several works have explored manifold learning using probability measures and the Wasserstein
metric, e.g., [5, 7]. This approach has shown improvement in task performance and data modeling,
but it comes with the significant drawback of slow computation time. In particular, computing
one Wasserstein distance exactly for discrete measures supported on m points costs Ω(m3 log(m))
flops. To form neighbor embeddings for N data measures requires computing the square distance
matrix Dij = W2(µi, µj)2, carrying total complexity O(N2m3 log(m)).
Mitigating this high computational cost can be approached in two ways: 1) speeding up a single
Wasserstein distance computation, and 2) computing fewer entries of D to form an approximation
of it. Many methods from computational OT explore the first approach, e.g., through entropic
regularization, multiscale methods, or Nystr¨om approximations to the cost matrix in the discrete
formulation of the Kantorovich problem (see [6] for a survey). In contrast, relatively few works have
explored the problem of estimating an entire Wasserstein distance matrix directly from samples.
The focus of this paper is to propose algorithms for estimating square Wasserstein distance
matrices for use in neighbor embeddings by computing as few entries as possible. Our approach
is based on exploiting the approximate low-rank structure of Wasserstein squared distances: if the
matrix D is approximately rank r, we propose two algorithms based on matrix completion and
the Nystr¨om method which allow us to estimate D from O(r log(r)) fully observed columns and
maintain stability of the MDS embedding.
2
Technical background
2.1
Optimal transport and MDS
The Kantorovich formulation of the quadratic Wasserstein metric on discrete probability measures
µ = Pn
i=1 αiδxi and ν = Pm
i=1 βiδyi is
W2(µ, ν) :=
min
P∈Γ(µ,ν) ⟨C, P⟩
(1)
where Γ(µ, ν) = {P ∈Rm×n
+
: P1m = α, P ⊤1n = β} are couplings with marginals µ and ν,
and Cij = ∥xi −yi∥2 [6].
This problem computes an optimal matching of mass between two
discrete measures. Given measures {µi}N
i=1 ⊂W2(Rn), the Wassmap algorithm [5, 7] computes
Dij = W2(µi, µj)2, a double centering of it B = −1
2HDH where H = I −1
N 11⊤, with truncated
SVD Bd = VdΣdV ⊤
d . The embedding vectors in Rd are then rows of VdΣd.
2.2
Matrix completion
Given a low-rank matrix with missing entries, the problem of estimating the unknown entries
is known as the matrix completion problem. This problem has been studied extensively due to
its broad applications, most notably in recommendation systems [8]. A naive formulation seeks
to minimize the rank of the matrix subject to the constraints imposed by the observed entries;
however, rank minimization is computationally hard in general [9]. A common convex relaxation is
to minimize the nuclear norm, which serves as a surrogate for rank. In this convex setting, seminal
theoretical results show that, under random sampling, an N × N rank-r matrix can be exactly
recovered from O(νrN log N) observed entries, where ν is an incoherence parameter capturing
the conditioning of the underlying matrix [10].
Given the computational challenges of convex
methods, non-convex algorithms have also been explored, and similar recovery guarantees have
been established in this setting [11]. In this paper, we focus on the completion of distance matrices.
In the case of Euclidean distance matrices, it can be shown that the squared distance matrix
2
has rank at most d + 2, where d is the embedding dimension of the underlying points [12]. This
structural property enables the use of low-rank optimization techniques for Euclidean distance
matrix completion [13, 14].
2.3
Nystr¨om method
Consider a symmetric matrix in block form D =
 A
B
B⊤
C

, where A ∈Rm×m and B ∈Rm×n. We
are interested in the regime m ≪n, where the aim is to estimate the block C using only the smaller
blocks A and B. Suppose that rank(A) = rank(D). Then, one can show that the columns of the
submatrix
B
C

lie in the span of the columns of
 A
B⊤

. Specifically, there exists an X ∈Rm×n
such that B = AX and C = B⊤X. If A is invertible, then X = A−1B, which yields the exact
relation C = B⊤A−1B. In the more general case where A is not invertible, we may instead use
the Moore–Penrose pseudoinverse to obtain X = A†B leading to C = B⊤A†B. This procedure is
known as the Nystr¨om method [15].
While it is not a completion algorithm per se, the Nystr¨om method can be used to complete a
distance matrix. In particular, if only the columns containing A and B⊤are computed, then the
entries of C can be recovered, or in the general case approximated, by B⊤A†B.
Notations For a distance matrix D, owing to symmetry and zero diagonal entries, we only consider
samples from the strictly upper-triangular part.
We denote by Ω⊂{(i, j) : i ∈[N], j > i}
the set of sampled index pairs, with cardinality |Ω| = m. Given a matrix D and an index set
I ⊂[N] := {1, . . . , N}, we write D(:, I) for the submatrix of D consisting of the columns indexed
by I. Analogously, D(I, :) denotes the submatrix formed by the rows indexed by I. PΩ(D)ij = Dij
if (i, j) ∈Ωand is 0 otherwise. If rank(D) = r and D = VrΣrV ⊤
r , then the incoherence ν of D is
p
n/r maxi ∥|Vr(i, :)∥.
3
Algorithms
Here we describe two algorithms for estimating the Wasserstein distance matrix to compute em-
beddings of images. Algorithm 1, W2-MC, randomly computes a given number of entries of the
distance matrix, and then uses a matrix completion algorithm to estimate the full squared distance
matrix. We leave the completion algorithm as a generic plug-in for existing algorithms, and consider
one designed for distance matrices in our experiments. We sample from the strict upper triangle,
include the diagonal, and reflect the entries across the diagonal to obtain the observed distance
matrix. We apply the same procedure after completion to ensure that the resulting matrix retains
the properties of a distance matrix.
Algorithm 2, W2-Nystr¨om, uses the Nystr¨om method to complete the distance matrix.
In
particular, we compute only a set of columns of the distance matrix Cobs = D(:, I), and estimate
Dest = CobsUC⊤
obs where U = Cobs(I, :) = D(I, I). As in W2-MC, we also enforce the distance
matrix structure by enforcing the 0 diagonal.
Note that for large N, the computational cost of both algorithms is dominated by the cost of
computing the W2 distances, and the cost of the matrix completion subroutine MC is often negligible
by comparison. Thus, algorithms that limit sampling complexity of the matrix are more desirable.
For this reason, we do not truncate the rank of the Nystr¨om approximation to D in W2-Nystr¨om,
which is usually done to save computation time.
3
Algorithm 1: W2-MC
1 Input: Data measures {µi}N
i=1, sample size m, matrix completion subroutine MC
2 Randomly sample m entries Ω
3 Compute observed distance matrix Dobs = PΩ(D), where Dij = W2(µi, µj)2
4 diag(Dobs) = 0
▷enforce 0 diagonal
5 Dobs = Dobs + D⊤
obs
▷reflect across diagonal
6 Dest = MC(Dobs)
7 Dest = 1
2(Dest + D⊤
est)
▷enforce symmetry
8 Output: Estimated distance matrix Dest
Algorithm 2: W2-Nystr¨om
1 Input: Data measures{µi}N
i=1, number of columns c
2 Randomly sample c columns I ⊂[N]
3 Compute columns of distance matrix Cobs =
(
W2(µi, µj)2,
i ∈[N], j ∈I
0,
otherwise.
4 Dest = CobsUC⊤
obs, where U = Cobs(I, :)
5 diag(Dest) = 0
▷enforce 0 diagonal
6 Output: Estimated distance matrix Dest
The Nystr¨om method is well-studied, and using some existing bounds, we are able to prove a
stability result for MDS embeddings on Wasserstein distance matrices. We assume that the data
measures in W2(Rn) are isometric to a set of Euclidean vectors {yi}N
i=1 ⊂Rd. Examples of such sets
have been studied and include common image articulations like translations, scalings, and shears
of a fixed base measure [16, 7]. We assume noise (or error) in the computation of the W2 distances,
and consider how the MDS embedding {zi}N
i=1 ⊂Rd coming from Dest of W2-Nystr¨om compares
to {yi}. MDS embeddings are centered and equivalent up to orthogonal transformation, so we
measure error by a Procrustes distance, defined as minQ∈O(d) ∥Z −QY ∥2, where O(d) is the set
of orthogonal d × d matrices. Note that Z and Y are matrices whose columns are the embedded
vectors.
Theorem 3.1. Suppose {µi}N
i=1 ⊂W2(Rn), and Dij = W2(µi, µj)2 has incoherence ν.
Sup-
pose there exist {yi}n
i=1 ⊂Rd such that Dij = ∥yi −yj∥2. Let λ > 0, δ ∈(0, 1), and compute
m ≳λν2(d + 2) max{log(d + 2), log(1/δ)} columns uniformly randomly in W2-Nystr¨om with er-
ror Cobsi,j = W2(µi, µj)2 + Eij such that ∥E∥2 ≤σd+2(D)/3.
If ∥Y †∥2∥E∥1/2
2
≤1/
√
2, then
Dest = CobsUC⊤
obs yields an MDS embedding {zi}N
i=1 ⊂Rd such that
min
Q∈O(d) ∥Z −QY ∥2 ≤(1 +
√
2)∥Y †∥2∥E∥2 ×

5
N
(d + 2)(1 −δ) + 3
s
N
(d + 2)(1 −δ) + 2
!
with probability at least 1 −max{re−λ(δ+(1−δ) log(1−δ)), δ}.
Proof. A simple modification of the proof of [17, Theorem 3.3] implies the upper bound with
(1 +
√
2)∥Y †∥2∥D −Dest∥2. Then [18, Theorem 3] implies that the columns selected would yield
4
an exact Nystr¨om decomposition of D without the presence of noise. Thus, the hypotheses of
[19, Remark 4.7] are satisfied, and yield that ∥D −Dest∥2 ≤5∥Vd+2(I, :)†∥2
2 + 3∥Vd+2(I, :)†∥2 + 2
where I is the set of columns selected by Algorithm 2 and Vd+2 are the first d + 2 left singular
vectors of D. Finally, [20, Lemma 3.4] shows that with the given probability, ∥Vd+2(I, :)†∥2 ≤
p
N/((d + 2)(1 −δ)) for all δ ∈(0, 1).
A simple corollary is that if we compute O((d + 2) log(d + 2)) columns of D without error, then
the MDS embedding is exact up to orthogonal transformation. Exact MDS embeddings of W2 data
were studied in [7].
4
Numerical Experiments
4.1
Fixed sample budget
We first do a fixed sample budget comparison of W2-MC and W2-Nystr¨om for completing the
Wasserstein distance matrix of 2, 000 points from the OrganCMNIST dataset from the MedMNIST
benchmark [21], comprising coronal plane CT scans of 11 different organs downsampled to size 28×
28. For uniform sampling, we subsample the matrix at different rates γ ∈{3%, 5%, 10%, 20%, 25%}.
We then match the number of observed entries between uniform sampling and the Nystr¨om method
by solving the equation c(c−1)
2
+ c(n −c) = γ · n(n−1)
2
for c (the number of columns selected by
Nystr¨om). This equation equates the number of off-diagonal entries used by both methods. For
MC in W2-MC, we use the factorized Gram matrix formulation from [14] based on the relationship
between the Gram matrix and the square distance matrix:
Find
P ∈Rn×q,
s.t.
(PP ⊤)ii + (PP ⊤)jj −2(PP ⊤)ij = Dij,
P ⊤1 = 0,
(2)
where q denotes the estimate for the rank of the underlying points. The centering constraint can
be removed by reparameterization. Defining A : Rn×n →R|Ω| as
A(X)α = Xii + Xjj −2Xij,
α = (i, j) ∈Ω,
the augmented Lagrangian is
L(Q; Λ) = 1
2

A(QQ⊤) −b + Λ

2
2 ,
(3)
with b is the vector of observed distances and Λ ∈R|Ω|.
The algorithm starts with an initial
estimate of Q (randomly initialized in our experiments) and then alternates between two steps:
updating Q via Barzilai–Borwein descent, and updating the Lagrange multipliers. For the dataset
in consideration, we set the rank estimate to be 200.
Table 1 reports the results averaged over 10 random trials. We observe that for fixed sample
budget, the Nystr¨om method consistently yields better accuracy compared to matrix completion.
While both Nystr¨om and typical matrix completion algorithms scale similarly up to logarithmic
factors, the Nystr¨om method is a single-step procedure, whereas the cost of the matrix completion
algorithm depends on the target accuracy, with the number of iterations required to reach a given
error often dominating the runtime. In our experiments, we run W2-MC for a sufficient number
of iterations, as we consistently observe stable convergence across trials. In practice, choosing a
5
Rate / # of cols.
W2-Nystr¨om
W2-MC
25% / 268
1.58 × 10−2
5.20 × 10−2
±1.96 × 10−3
±1.91 × 10−2
20% / 211
1.85 × 10−2
3.42 × 10−2
±2.2 × 10−3
±8.26 × 10−5
10% / 103
2.85 × 10−2
5.65 × 10−2
±2.60 × 10−3
±6.01 × 10−3
5% / 51
4.82 × 10−2
1.01 × 10−1
±8.59 × 10−3
±2.50 × 10−2
3% / 30
7.21 × 10−2
1.32 × 10−1
±1.20 × 10−2
±9.77 × 10−3
Table 1: Relative error (mean ± standard deviation) of Nystr¨om and matrix completion under
different sampling rates of a 2000 × 2000 OrganCMNIST W2 distance matrix. For the Nystr¨om
method, the numbers of columns associated with the sample rates are also reported.
suitable stopping criterion for matrix completion problems could also be challenging. Therefore, this
fixed sample budget experiment leads us to conclude that computing full columns of W2 distance
measurements may be preferred for approximating Wasserstein distance matrix. Studying how far
this generalizes is the subject of future work.
4.2
Classification stability
We next consider the stability of classification performance under the Nystr¨om approximation as
the number of sampled columns varies. Specifically, we compute 5%, 10%, 20%, . . . , 100% of the
columns of the 2000 × 2000 OrganCMNIST distance matrix using W2-Nystr¨om. The resulting
approximation Dest is embedded via MDS, followed by a 90 −10 train-test split. For the MDS
embedding, we set the dimension to 120, chosen such that the retained singular values account
for 97% of the total sum of singular values. Classification is then performed on the embeddings
using linear discriminant analysis (LDA), 1-nearest neighbor (KNN1), kernel SVM with Gaussian
kernel (SVM RBF), and random forest. We note that the dataset has 11 underlying classes, and we
evaluate the quality of the classification using classification accuracy, defined as the proportion of
correctly identified data points. Figure 1 shows classification accuracy as a function of the number
of columns sampled.
The main conclusion we draw from Figure 1 is that the classification performance is quite stable
after computing about 10–20% of the columns of the distance matrix. This result and Theorem 3.1
suggest that the Nystr¨om Wassmap algorithm (MDS on the output of Algorithm 2) is stable after
sampling O((d + 2) log(d + 2)) columns of D.
5
Conclusion
We showed that the Nystr¨om method yields a good approximation of W2 distance matrices, leading
to stable MDS embeddings. On the OrganCMNIST dataset, we showed that classification accuracy
is stable even when only sampling a small fraction of columns, and for a fixed sampling budget, the
Nystr¨om method outperforms a distance matrix completion approach. Unlike classical Euclidean
distances, where the tight link between Gram and distance matrices grounds the theory of multidi-
mensional scaling, the structural properties of W2 distance matrices are far less understood. This
6
t
Figure 1: Classification accuracy (mean ± 1 standard deviation of 100 trials) vs. number of columns
computed in Algorithm 2 for 2000 OrganCMNIST data points.
gap makes it especially important to develop principled methods that reduce the cost of computing
W2 distances.
In future work, we will explore theoretical conditions under which W2 distance
matrices admit approximate low-rank structure or fast spectral decay, as well as explore hybrid
strategies that combine Nystr¨om-selected columns with partial entry sampling similar following
ideas similar to [22].
Acknowledgements
Abiy Tasissa acknowledges partial support from the National Science Foundation through grant
DMS-2208392. HanQin Cai acknowledges partial support from the National Science Foundation
through grant DMS-2304489.
Muhammad Rana and Keaton Hamm were partially supported by a Research Enhancement
Program grant from the College of Science at the University of Texas at Arlington and by the
Army Research Office under Grant Number W911NF-23-1-0213. The views and conclusions con-
tained in this document are those of the authors and should not be interpreted as representing the
official policies, either expressed or implied, of the Army Research Office or the U.S. Government.
The U.S. Government is authorized to reproduce and distribute reprints for Government purposes
notwithstanding any copyright notation herein.
References
[1] Nanning Zheng and Jianru Xue, Statistical learning and pattern analysis for image and video
processing, Springer Science & Business Media, 2009.
7
[2] Marina Meil˘a and Hanyu Zhang, “Manifold learning: What, how, and why,” Annual Review
of Statistics and Its Application, vol. 11, no. 1, pp. 393–417, 2024.
[3] James Mathews, Maryam Pouryahya, Caroline Moosm¨uller, Ioannis G. Kevrekidis, Joseph O.
Deasy, and Allen Tannenbaum, “Molecular phenotyping using networks, diffusion, and topol-
ogy: soft-tissue sarcoma,” Scientific Reports, vol. 9, 2019, Article number: 13982.
[4] Matt Kusner, Yu Sun, Nicholas Kolkin, and Kilian Weinberger,
“From word embeddings
to document distances,” in International conference on machine learning. PMLR, 2015, pp.
957–966.
[5] Wei Wang, John A Ozolek, Dejan Slepˇcev, Ann B Lee, Cheng Chen, and Gustavo K Rohde,
“An optimal transportation approach for nuclear structure-based pathology,” IEEE transac-
tions on medical imaging, vol. 30, no. 3, pp. 621–631, 2010.
[6] Gabriel Peyr´e and Marco Cuturi, “Computational optimal transport: with applications to
data science,” Foundations and Trends in Machine Learning, vol. 11, no. 5-6, pp. 355–607,
2019.
[7] Keaton Hamm, Nick Henscheid, and Shujie Kang, “Wassmap: Wasserstein isometric mapping
for image manifold learning,” SIAM Journal on Mathematics of Data Science, vol. 5, no. 2,
pp. 475–501, 2023.
[8] Andy Ramlatchan, Mengyun Yang, Quan Liu, Min Li, Jianxin Wang, and Yaohang Li, “A
survey of matrix completion methods for recommendation systems,” Big Data Mining and
Analytics, vol. 1, no. 4, pp. 308–323, 2018.
[9] Raghu Meka, Prateek Jain, Constantine Caramanis, and Inderjit S Dhillon, “Rank minimiza-
tion via online learning,”
in Proceedings of the 25th International Conference on Machine
learning, 2008, pp. 656–663.
[10] Emmanuel J Cand`es, Justin Romberg, and Terence Tao, “Robust uncertainty principles: Exact
signal reconstruction from highly incomplete frequency information,” IEEE Transactions on
information theory, vol. 52, no. 2, pp. 489–509, 2006.
[11] Yuejie Chi, Yue M Lu, and Yuxin Chen, “Nonconvex optimization meets low-rank matrix
factorization: An overview,” IEEE Transactions on Signal Processing, vol. 67, no. 20, pp.
5239–5269, 2019.
[12] Leo Liberti, Carlile Lavor, Nelson Maculan, and Antonio Mucherino,
“Euclidean distance
geometry and applications,” SIAM review, vol. 56, no. 1, pp. 3–69, 2014.
[13] Nilson JM Moreira, Leonardo T Duarte, Carlile Lavor, and Cristiano Torezzan, “A novel low-
rank matrix completion approach to estimate missing entries in euclidean distance matrix,”
Computational and Applied Mathematics, vol. 37, no. 4, pp. 4989–4999, 2018.
[14] Abiy Tasissa and Rongjie Lai, “Exact reconstruction of euclidean distance geometry problem
using low-rank matrix completion,” IEEE Transactions on Information Theory, vol. 65, no.
5, pp. 3124–3144, 2018.
[15] Christopher Williams and Matthias Seeger, “Using the nystr¨om method to speed up kernel
machines,” Advances in neural information processing systems, vol. 13, 2000.
8
[16] Varun Khurana, Harish Kannan, Alexander Cloninger, and Caroline Moosm¨uller, “Supervised
learning of sheared distributions using linearized optimal transport,” Sampling Theory, Signal
Processing, and Data Analysis, vol. 21, no. 1, 2023.
[17] Alexander Cloninger, Keaton Hamm, Varun Khurana, and Caroline Moosm¨uller, “Linearized
wasserstein dimensionality reduction with approximation guarantees,” Applied and Computa-
tional Harmonic Analysis, vol. 74, pp. 101718, 2025.
[18] Ameet Talwalkar and Afshin Rostamizadeh, “Matrix coherence and the nystrom method,”
arXiv preprint arXiv:1408.2044, 2014.
[19] Keaton Hamm and Longxiu Huang, “Perturbations of cur decompositions,” SIAM Journal
on Matrix Analysis and Applications, vol. 42, no. 1, pp. 351–375, 2021.
[20] Joel A Tropp, “Improved analysis of the subsampled randomized hadamard transform,” Ad-
vances in Adaptive Data Analysis, vol. 3, no. 01n02, pp. 115–126, 2011.
[21] Jiancheng Yang, Rui Shi, Donglai Wei, Zequan Liu, Lin Zhao, Bilian Ke, Hanspeter Pfister,
and Bingbing Ni, “Medmnist v2-a large-scale lightweight benchmark for 2d and 3d biomedical
image classification,” Scientific Data, vol. 10, no. 1, pp. 41, 2023.
[22] HanQin Cai, Longxiu Huang, Pengyu Li, and Deanna Needell, “Matrix completion with cross-
concentrated sampling: Bridging uniform sampling and cur sampling,” IEEE Transactions on
Pattern Analysis and Machine Intelligence, vol. 45, no. 8, pp. 10100–10113, 2023.
9