Question:
Based on the results obtained from the experiment, the DWGF algorithm performs significantly well in terms of PSNR and LPIPS, but has a poor FID. Regarding this phenomenon, the following analysis is correct:
A. The core reason for this phenomenon is that when minimizing the forward KL divergence, the optimization process drives the approximate distribution q to shrink and closely fit a certain local peak of the true distribution p, while ignoring the behavior tendency of other peaks and low-probability regions.
B. DWGF performed poorly in the FID metric because the model collapse resulted in a narrow distribution of generated samples, failing to cover all variations of the real data, thus leading to insufficient differences between the overall distribution and the real distribution.
C. Unlike DWGF, RLSD introduces the repulsive potential, actively forcing particles to maintain a distance from each other and preventing them from all collapsing onto the same mode.
D. PLSD is a compromise between DWGF and RLSD, and its effect is more balanced.

Answer: C

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:
This paper proposes a novel free-training method called Diffusion-regularized Wasserstein Gradient Flow(DWGF) for solving inverse problems. Regarding the description of DWGF in the paper, the correct statements are:
A. The core idea is to minimize the objective function by simulating the evolution process that follows the fastest descent in the metric space based on a probability distribution.
B. Use an adaptive optimization algorithm to combine the two gradients and update the particle positions.
C. Starting from the first principles, the posterior sampling is formulated as a Wasserstein gradient flow in the latent space, with the goal of minimizing the KL divergence.
D. Applicable to tasks that require high quality and extremely high reconstruction accuracy, but do not demand high diversity in generated samples.

Answer: ABC

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.

...

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 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.

...

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)

...

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
...

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.

...

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.

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