Learning few-step posterior samplers by unfolding and distillation of diffusion models
Authors: Charlesquin Kemajou Mbakam
Marcelo Pereyra
Jonathan Spence
Venue: Accepted by TMLR ()

Abstract: Diffusion models (DMs) have emerged as powerful image priors in Bayesian computational imaging. Two primary strategies have been proposed for leveraging DMs in this context: Plug-and-Play methods, which are zero-shot and highly flexible but rely on approximations; and specialized conditional DMs, which achieve higher accuracy and faster inference for specific tasks through supervised training. In this work, we introduce a novel framework that integrates deep unfolding and model distillation to transform a DM image prior into a few-step conditional model for posterior sampling. A central innovation of our approach is the unfolding of a Markov chain Monte Carlo (MCMC) algorithm—specifically, the recently proposed LATINO Langevin sampler (Spangnoletti et al., 2025)—representing the first known instance of deep unfolding applied to a Monte Carlo sampling scheme. We demonstrate our proposed unfolded and distilled samplers through extensive experiments and comparisons with the state of the art, where they achieve excellent accuracy and computational efficiency, while retaining the flexibility to adapt to variations in the forward model at inference time.

Published in Transactions on Machine Learning Research (11/2025)
Learning few-step posterior samplers by unfolding and distil- lation of diffusion models
Charlesquin Kemajou Mbakam
emk:2000@ho.ccc
School of Mathemationl and Computer Sciences 8 Marwell Institute for Mathematical Sciences
Heriot- Watt University
Jonathan Spence
JSpence: 5eeernccceedd....oooo uk
School of Mothematics &
Maswell Institute for Mathematicol Sciences University of Edinburyh
Marcelo Pereyra
M.P Mererra Www..
School of Mathematical and Computer Sciences & Maxwell Institute for Mathematical Sciences
Heriol- Wall University
Reviewed on OpenReview: //perrvveeeeeeveee net/ forum? id= DGC
Abstract
Diffusion models (DMs) have emerged as powerful image priors in Bayesian computational imaging. Two primary strategics have been proposed for leveraging DMs in this context: Plug-and-Play methods, which are zero-shot and highly flexible but rely on approximations; and specialized conditional DMs, which achieve higher accuracy and faster inference for specific lasks through supervised training. In this work, we introduce a novel framework that integrates deep unfolding and model distillation to transform DM image prior into few-step conditional model for posterior sampling. A central innovation of our approach is the unfolding of a Markov chain Monte Carlo (MCMC) algorithm specifically, the recently proposed LATINO Langevin sampler (Spagnoletti et al. 2025 representing the first known instance of deep unfolding applied to a Monte Carlo sampling scheme. We demonstrate our proposed unfolded and distilled samplers through extensive experiments and comparisons with the state of the art, where they achieve excellent accuracy and computational efficiency, while retaining the flexibility to adapt to variations in the forward model at inference time. The code is available at
www.ooommnhaaaaleeeaeeeee - UD2M
1
Introduction
We seek to infer an unknown image of interest x² E Rdd from a noisy measurement y = related to x² by a statistical observation model of the form
y ~ P(Ax*),
(1)
where A is a linear measurement operator describing deterministic instrumental aspects of the data acquisition
process and P is statistical model describing measurement noise. This includes, for instance, the canonical linear Gaussian observation model Ar* +n, where n is additive Gaussian noise with covariance o Such models arise frequently in problems related to image deblurring, inpainting, super-resolution. compressive sensing, and tomographic reconstruction (scc, c.g., (Kaipio & Somersalo, 2010 Ongic et al., (0200). Within this context, we are particularly interested in imaging problems where recovering x² from y severely ill-conditioned or ill-posed, leading to significant uncertainty about the solution (Kaipio & Somersalo 2010

---

Published in Transactions on Machine Learning Research (11/2025)
Adopting Bayesian approach, we leverage prior knowledge available in order to regularize the problem and deliver meaningful inferences that are well-posed. This is achieved by treating x' as a realization of random variable Xg and using Bayes' theorem to obtain the posterior distribution of xoly = y, with density given for all zy E Rd by (Robert 2007)
plylzo) p(roly) Jne plyl To )p( (To (o))do
where p(y( denotes the likelihood function associated to Equation and p(ro) the marginal density of x
The choice of p(zo), so-called prior, is crucial and greatly impacts the obtained posterior p(zoly), especially in problems that are ill-conditioned or ill-posed. Bayesian imaging methods traditionally used analytical priors designed to promote solutions with expected structural properties (e.g., smoothness, sparsity, piecewise regularity). However, modern Bayesian imaging methods rely predominantly on machine learning in order to harness vast amounts of prior knowledge available in the form of training data and deliver unprecedented accuracy (Mukherjee et al. 2023). ). In particular, state-of-the-arl Bayesian imaging techniques often use deep generative models as image priors, notably denoising diffusion models (DMs). For an introduction to DMs, we refer the reader to the seminal DM papers Song & Ermon (2019) 2020); Song et al. and to the recent. survey on DMs as image priors for solving inverse problems (Daaas et al. 2024).
The literature on DM-based Bayesian imaging methodology has two main strands. On the one hand, Plug & Play (PnP) methods soek to use pretrained DM in zero-shot manner, in combination with likelihood function p(y|=n)) specified analytically during inference. PnP approaches are flexible by construction and generalize robustly to new measurement models by exploiting knowledge of the likelihood function p(yl explicitly. Notable examples of PnP DM techniques include, e.g., DPS (Chung et al. 2022), DDRM (Kawar ct al. 2022). DifPPIR (Zhu ct al., 2023). IIGMM (Song et al. 2023a). and mid-point (Moufad et al. 2025).
The other strand relies on fine-tuning an unconditional DM representing the prior, in order to construct. y-conditional DM for posterior sampling. This involves modifying DM to take the measurement y input and, through supervised training, specializing it for posterior sampling of xoly= (instead of marginal sampling of xp)- Specializing a DM for a specific imaging task enables the development of Bayesian imaging methods that surpass PnP approaches in both accuracy and computational efficiency, assuming the task is fully defined during training. However, in the absence of additional training, such task-specific DMs often exhibit limited generalization even slightly altered measurement conditions, as they are not inherently designed to exploit the likelihood p(y]zz)) during inference. For more details and comparisons with PnP strategies, see, c.g., I'SB (Liu et al. 2023a), InDI (Delbracio & Milanfar 2023 CDDB (Chung ct al. 2023
Moreover, the most advanced DM-based imaging methods currently available wheether zero-shot or trained for conditional sampling- leverage model distillation to enhance computational efficiency and improve performance. Notably, distillation techniques such as consistency models (Song ct al. 2023c: Kim et al. 2024) and flow matching (Liu et al.] 2022bb Lipman el al., 2023) have significantly reduced the number of required neural
function evaluations (NFEs) per posterior sample from over 10 to fewer than 10, while concurrently improving
sampling quality and accuracy (see (Spagnolett et al. (2025) Garber & Tirer 22025 Zhao et al. (2025)
To bridge the methodological gap between these two strands, this paper introduces a novel framework for constructing conditional DMs which offers the computational efficiency and accuracy of conditional DMs obtained via finetuning and distillation, with the flexibility afforded by explicit likelihood functions in inference, thus bringing together the advantages of finetuning and PnP strategies. This is achieved by generalizing deep unfolding (Monga et al. 2021), a paradigm wherein iterative algorithms with a fixed number of steps are mapped onto modular deep neural network architectures. These architectures explicitly incorporate the likelihood function p(y]=n), for example through its gradient proximal operator, in conjunction with a denoiser that is subsequently fine-tuned during training. A main novelty of this work is to present the first application of deep unfolding of Markov chain Monte Carlo (MCMC) sampling scheme, in contrast to prior work on deep unfolding that focused on optimization--asse algorithms (which, unlike MCMC schemes, are not inherently suitable for posterior sampling). Namely, we propose to unfold the sraate-f--the-rrs LATINO zero-shot imaging Langevin method (Spagnoletti et al., 2025), , and train the resulting unfolded LATINO networks by using a supervised consistency trajectory models objective combining distortion, perceptual and adversarial terms (Kim et al. 2024). This leads to distilled conditional DMs that are fast. (low NFFs),
2

---

Published in Transactions on Machine Learning Research (11/2025)
Deblurring
Inpainting
SR (x4)
JPEG (QF=10)
Measurement
UD (Ours)
Figure 1: Qualitative comparison of the proposed Unfolded and Distilled Diffusion Model (UD*M) for posterior sampling on the ImageNet 256 dataset. Tasks: Gaussian Deblurring, random inpainting (70%), super-rcsolution (4x), and restoration of JPEG compression artifacts (QF=10).
generate high-quality samples, are GPU-memory efficient during inference (due the absence of automatic differentiation).. We refer to the proposed method as Unfolded and Distilled Diffusion Model (UD²M). Crucially, UD3Ms support joint training to specific measurement models at inference time (i.e., specific instances of A and P). For illustration, Figure shows some examples of posterior samples obtained with the proposed unfolded and distilled diffusion models (UD*Ms) applied to Gaussian deblurring, random inpainting, super-resolution by factor 4, and JPEG artifact restoration for images from ImageNet dataset.
family of likelihoods, followed by instantaneous specialization
In summary, we present the following two main methodological contributions that underpin the proposed CD2M method:
Deep unfolding of Monte Carlo sampling: In Section 3 we introduce novel framework that unfolds the LATINO Langevin Markov chain Monte Carlo sampler into a trainable neural architecture, representing the first application of deep unfolding to Monte Carlo sampling scheme.
Conditional model distillation for efficient posterior sampling: In Section 3.2
we employ modern
generative model distillation techniques to fine-tune the unfolded LATINO network into a few-step conditional sampler that achicves high accuracy and flexibility while enabling fast inference.
A summary of background material can be found Section 2 and numerical comparisons our method to state-of-the-at models from the literature are presented in Section
1.1 Notation
We denote by p(æly) the posterior distribution of the random variable x given y = Similarly, p(y]<) denoles the likelihood. We use the shortened notation log p(y(z) = (x). For a function f fd - R, the operator prox,f = arg min f() + 1 ||. 21/3/2 denotes the proximal operator of f. The expected value of a random variable f(x) denoted [x)); when the distribution to be integrated clear from context, we occasionally drop the subscript and write Ex)(
3

---

Published in Transactions on Machine Learning Research (11/2025)
2 Background & related works
Diffusion models (DMs) are generative models that draw samples from a distribution of interest p(xo) by iteratively reversing the following "noising" process
3r x + Bidwer
dx.
(2)
2
which is designed to transport Xy to a standard normal random variable Xy ~ (O, 1) as T co, and where 3, represents a noise schedule and w, is a Brownian motion (Ho et al.] 2020). The reverse process, which allows generating samples of Xo by transporting XT ~ N(O, I) back to Xo, given by (Song et al. 2021):
dx. [1- B: xxx= Be x, po)))(x)) dt + VBddw. 2.
where Wa is again Brownian motion, in the reverse direction. The target distribution p(xo) encoded in the so-called score function I2 + 2z, log (oxe), which is approximated using Tweedie's identity (Efron 2011] by a deep neural network denoiser Go(: () that is trained by weighted denoising score matching (DSM), see Song et al. (021).
Denoising diffusion probabilistic models (DDPM) (Ho et al., 2020 and denoising diffusion implicit models (DDIM) (Song et al. ) are the two main met met hodolooss to implement. Equation 3 in time-discrete setting, where the process is initialized with Xy ~ NO, I) for some large finite T and solved itcratively with decreasing progressively to ( = O. Accurate estimation of the scores Vz, log underpinning Equation 3 has been crucial to the success of DDPM and DDIM. This has been achieved using large training datasetss specialized network architectures, and significant high-peeformance computing resources (Ho et al. 2020 Song ct al. 2021). In addition to advancing generative modcling, DMs have become central to modern strategies for solving inverse problems in computational imaging (see, e.g,, the recent survey paper (Daras el al. 2024
Distillation of diffusion models Conventional DM approaches produce high quality samples by pro- gressively transporting Xy to Xg. However, they have high computational cost, as generating each sample typically requires performing between 100 and 1000 evaluations of Gg. This drawback can be addressed by distilling Go into a new model Go requiring far fewer steps (e.g., 4-8 steps); see, e.g., Song et al. (2023c); Liu et al. (2023b). Remarkably, modern distillation techniques combining DSM with adversarial training are able further improve the quality of the samples as they dramatically reduce the number of sampling steps, while retaining the same network architecture as the original DM. This is the case, for example, of the so-called consistency trajectory models (Kim et al. 2023,, which we use in our proposed method in Section 4.
Consistency models Consistency models (CMs) are derived from probability flow formulation of Equation given by (Song al. 2023b
Br 2 r1
3t & Va, log dt.
dx. =
(4)
2
This ODE is equivalent to Equation 3 in the sense that both have the same marginals pe(-) of xt- CMs
leverage this property by learning so-called consistency function Go that maps any point de on trajectory
(atheey of Equation 4 backwards to Σyp for some small given 17 > 0. Given Go (I,, + Zay CMs are able to sample the r.v. x,, (xT, T) with xy I) in single step, see Kim et al. (023)) for
details. Two-step CMs achieve superior quality by re-noising X,, Go T)) following Equation for some
intermediate time s (n. T) followed by Go(x,, s) to bring back x, close to Xg. Multi-step CMs apply this stra strategy recursively, in 4 to 8 steps, achieving top performance while retaining computational efficiency (Kim) lt al.) 2024). Recent examples of few-step Bayesian image restoration methods based on CMs include, for instance, the fino-tuned method CoSIGN (Zhao et al. (025)), the zero-shot method CM4IR (Garber & Tirer 2025j, and the zero-shot method LATINO (Spagnoletti et al. 2025) which uses a so-called Stable Diffusion latent CM with semantic conditioning via text prompting.
4

---

Published in Transactions on Machine Learning Research (11/2025)
LoRA fine-tuning of diffusion models Low-Rank Adaptation (LoRA) Huu et al. 2022 is technique designed to fine-tune large attention-based architectures such as DMs in manner that achieves comparable performance to full fine-tuning at a fraction of the computational cost. This is achieved by adding a trainable low-rank correction Ag to the modcl's original attention weights, which remain frozen during fine-tuning, and exploiting structural properties of allention layers in order to optimize the weights 0 + Ag at a reduced cost. LoRA fine-tuning has been widely applicd to DMs. For example, in the context of DM distillation, CM-LORA strategies distill a pre-trained DM G.(t, into a CM Ga(- ()) = Go+ (t,() by adjusting Ae with CM objective combining DSM and adversarial training (Luo et al. 2023
Deep unfolding (also known as deep unrolling) is a deep learning paradigm that transforms iterative optimization algorithms with fixed number of iterations into deep neural network architectures, whereby the algorithm's 's steps become trainable layers in a network that is trained end-to-end (Monga et al. 2021). This approach, also known as deep unrolling, provides a powerful template for designing interpretable network architectures that incorporate y and the operator A in a manner that modular and explicit (e.g., via proximal operator layers), together with trainable elements (e.g. U-Net modules) that can be recognized as data-driven regularization terms. Because A is provided explicitly during training, unfolded networks can be easily trained to handle a range of operators (e.g., motion blurs, as represented by a database) and then instantiated for a specific operator during inference. To date, deep unfolding has primarily been implemented by unrolling optimization algorithms. A main innovation of this paper is the unfolding of a stochastic Langevin sampling algorithm instcad (specifically, the LATINO algorithm of Spagnoletti et al. (0025) Also, to the best of our knowledge, this the first instance of deep unfolding in the context of DMs.
Langevin PnP methods for image restoration. DMs can serve as highly informative PnP priors within Langevin Markov chain Monte Carlo posterior samplers. Notably, [Kemjoou Mbakam et al. incorporates DM within PnP unadjusted Langevin algorithm (ULA) (Laumont et al.) 2022) to estimate the posterior mean E(zly). Similarly. [Cocurooux et al. (2024b) embeds a DM within a Split-Gibbs sampler (Vono et al. 2019), which is equivalent to a noisy UL.A PPereyra et. al.; 2023). With regards to other stragies not based
on DMs, we note [oourdoox et al. (2024a)
and Melidonis et al. 2024) which explore the use of normalizing 20222) proposes a general theoretical framework for incorporating
flows as PnP priors, while Holden et al.
variational autoencoder (VAE) and generative adversarial network (GAN) priors.
LAtent consisTency INverse sOlver (LATINO) is a stale-of-the-art zero-shot posterior sampling technique derived from embedding CM prior Go(:,t) within an overdamped Langevin diffusion process of the form dxo., = V log ds + V log p(x..))d + dw., which converges to the posterior of interest p(roly) as s - 00, see Spagnoletti et al. (2025). In its simplest form, LATINO to sample from p( ly) is based on the following recursion for all
k N
x 1k+1) prox (x0 (k)
implicit gradient step, By : Io H> - pogl(y))
x (k+1)- Vāra (k+1)
v(1 oig Jest1, sample X1, (-1 by using forward SDE 2
x
(k+1) (o(((ttt)) (k+1) tt),
transport Xx,
(k+1) with ODE 4 (k+1) back to x
(5)
where 8,13 are step-sizes related to the time-discretization of the Langevin process, 6k41 N I), and the proximal step prOx (5) = arg min Gn (2o) + 23
1 | =113/2 with Jy 2n + - -log p(y((r) is equivalent to an implicit Euler step on Jy: see Spagnoletti et al. (2025) and Appendix B for details. For linear Gaussian observation models of the form y = Ar -n with noise covariance I²I, prox (x) = (6A A + - (5AA y+
o²x), which can be computed efficiently by using a fast linear solver or specialised scheme. Alternatively, for challenging linear operators and non-linear problems, we use an appropriate optimizer to compute an approximate solution of prox (s) = arg min mo (x0) + 28 | - sI s12/2, which is essentially a regularized least squares problem. It is worth mentioning that a more general form of the LATINO recursion generalizes Equation by performing the noise and de-noise step on the latent space of a (deterministic) auto-encoder pair (E, D). This allows embedding CMs trained on the latent space of (E, D), notably modern stable diffusion
Thee steps of the recursion in Spagnoletti al. (2225 are presented a different order, first sampling using then the CM step, and the proximal step last. We find that starting with the proximal step performs better when unrolling iterations.
5

---

Published in Transactions on Machine Learning Research (11/2025)
CMs Yin et al. 2024p. In the present work, we consider CMs trained to denoise directly on pixel space, with & and DD both set to the identity in Spagnoletti et al. 20255
Integrating the CM Go within a Langevin sampler allows LATINO to use the likelihood function zo + p(y(2o) directly and exactly- through its proximal operator without the need for approximation This represents a significant advantage over prevalent alternative DM-based posterior sampling approaches, which attempt to embed the likelihood into the DM itself process that is typically intractable and reliant on likelihood approximations because of time inhomogeneity. As result, LATINO outperforms competing zero-shot DMs in sampling accuracy, with computational cost comparable to Seate-of-the-ase multi-step CMs (i.e., in the order of 4-8 iterations). The proposed UD²M framework unfolds LATINO within a y-conditional distilled DM, which is trained end-to-end in a supervised manner for posterior sampling.
3 Proposed framework
3.1 UD*Ms:
Diffusion model unfolding and distillation for multi-step posterior sampling
We are now ready to present our proposed methodology for lcarning a posterior sampler for (Xuly = y), through unfolding and distillation of a DM that has been trained to sample x, unconditionally (i.e., from the prior p(ro)). We assume the availability of DM Gp(-.t) pre-trained to approximate Go(z,, t) z [[x,|xx = xạ) with x Va + vI- āee, following Equation 2 From this DM, we seek to derive a distilled model La that, in single step, samples approximately from the conditional distribution p(xuly, xz). This distilled model La is then embedded within a multi-step sampling scheme that iteratively generates x; ~ p(x,|>0 = x =) exactly through the noising process of Equation 2 followed by approximate sampling of x p(xoly, (k) by using La Lo(y, .x. ,)). The scheme starts at L = T by sampling from the marginal Xy ~ NN I), and the timestep t is reduced across iterations to reach f = 0 in a few steps (e.g., N = 3 steps). This scheme is summarized Algorithm [ Importantly, unlike zero-shot methods, here we consider that the likelihood function p(y) associated to xx)) belongs a pre-determined (parametrized) class of likelihood functions of interest that is fixed during the training of Lg. By training over wide range of sampled likelihoods from the determined class, retain flexibility to specify the exact parameters of the likelihood during inference time. Specializing Lg in this manner will lcad to significant accuracy and computational cost gains for likelihood models belonging to the training distribution when compared to zero-shot strategies that are fully likelihood agnostic during training. Note that Le samples the conditional random variable (xoly, xt), we do not soek to sample (xoly) directly.
(k)
Our UD2M approach to construct La combines distillation with deep-unfolding, whereby we obtain a trainable architecture from unfolding an iterative algorithm into sequential scheme (Monga et al. 2021). As mentioned previously, a main novelty of this work is to unfold a Langevin sampler to construct a generative model, which we train end-to-end. More precisely, given t E [0, and x, Iz, we implement La zi, t) by unfolding K iterations of the LATINO Langevin sampling algorithm of Equation 5 modified to target the conditional distribution p(xoly, xt). The precise construction of Lp is illustrated in Figure 2 Within each LATINO module, the prior represented by a distilled DM Go+ derived from Gg by LORA fine-tuning where As denotes LORA adaptation weights which we will train end-to-end for conditional sampling (with 0 frozen), and where the likelihood p(y, 2, (ma)
(o)p)))))( associated with the posterior p(xo y, 2.2) is involved via its proximal operator proxsgg.7, (s) arg min 9y,z, (zo) + 28 |1 >113/2 where Jy. log p(y, xa (ro).
As mentioned previously, for linear Gaussian observation models of the form y = Ax + with noise variance
o², this operator is given for all x E Rd by
prox&go.a
(x) (0)=
AT 02
Va 2
2 8
(6)
2The Langevin diffusion time-homogencous process and the resulting itcrates x (k) are asymptotically ergodic, converging to ô-neighbourhood of p(rolyy as + oc. The iterates are not traveling backwards in time through inhomogeneous prncess and therefore there no need to compute the likelihood y w.r.t. a naisy version of xo, as is the case conditional DMs.
6

---

Published in Transactions on Machine Learning Research (11/2025)
2/
Distil. DM step Go+ (1, ts
Distil. DM step Go+ (-, to
Distil. DM step Go+ (,, t6)
se-ttep forward SDE (2)
fs-step forward SDE (2)
15-slep forward SDE (2)
Lo 24. t)
(0)
proX
proXagy.n
L
((2)
5 11,00
(K--))
proX
Zt
Figure 2: Diagram of the proposed conditional sampling architecture, La(y, Fe, derived by deep unfolding K LATINO iterations (Spagnoletti et al. 2025). The prior is introduced via pre-trained unconditional DM Go with LoRA adaptation Ao, while the observation model, mcasurement and noisy state 24 are involved via
gy.z, = - -logp(y. de (z))- The first module is initialized by an estimator Jp such as RAM (Terris et al. 2025 or by setting (0( v) = Aty. y. The unfolded network is finetuned and distilled to sample from (xoly.
where of = - -T,)/a, and the matrix Σ is defined as follows
A
A
I Ti
I
(7)
For other modcls, we solve prox (s) arg min go 9y.z, (20) + 26 2200 s113/2 with gy.r. log p(y, F2(
approximately by using Adam (Kingma & Ba 2014 with vn as warm-starting.
Because of the low-rank structure of Ag, and because we used tied weights across LATINO modules, the number of trainable parameters usually of order of 0.1% of the number contained in 6, which makes adversarial training significantly more stable. Operating with reduced number of parameters also greatly simplifies lcarning and storing adaptations of Go tailored for specific classes of problems of interest, and allows us to easily instantiale problem-specific variants of Lg during inference. Moreover, Equation S can be computed cheaply when singular value decomposition of A is available. Otherwise, if computing an exact solution efficiently is not possible, we use a specialised sub-iteration, a randomizod lincar solver, or some steps of convex optimization scheme (Chambolle & Pock 2016).
Lastly, to keep the number of LATINO modules K small without compromising accuracy, we warm-start the first LATINO module by using auxiliary neural estimator t. x,) input module, paramctrized by set of weights t. For example, in experiments use the recently proposed foundational model RAM (Terris et al. 2025), with stacked operators (A, to take (y,x.) as input. Alternatively, one can also initialize with A Aly. Accordingly, the set of trainable parameters of La is 0 = (Do, v) if fy is finetuned for additional performance; otherwise we have & = As if v0 is frozen or the first LATINO module is initialized directly with A y. The proposed architecture summarized Figure 2 below.
Before discussing our strategy for training Lo, is worth noting that one could in principle design La to target the posterior p(xply) ly) directly, by using unfolding and distillation but bypassing fully the multi-step scheme. However, one-step samplers
notoriously harder to train, and the resulting samples are often of inferior quality [Luhman & Luhman 2021 Mong et al. 2023 Song et al. 2023c Salimans & Ho 2022). Also, sampling p(xoly) directly with LATINO would require unfolding a larger number of iterations compared to sampling p(xoly/ y, 2x), hence offsetting the computational benefits of targeting p(xoly) directly. This is due to the fact that the additional information provided by x, = 2, leads to a posterior distribution that much better conditioned and thus easier sample by Langevin sampling. By using multi-step sampling, we obtain excellent sample quality with reduced number of NFEs (e-g., we K and N = steps).
7

---

Published in Transactions on Machine Learning Research (11/2025)
Algorithm 1 Conditional Diflusion Sampling
Require: Observation y. Time-grid 0 = to < ti ... < <tN = T.
1: Sample ZiN ~N)
Initialize reversed diffusion
2: for n. N,. NN...,1 do
3: Set Lo f- Ltan (Ztn,: 3) using Lo (see Figure 2
Unfolded sample targeting po Tt,, y)
Sample 2in-1 ~ Pin-i (Ito-1 80, I2n))
Reverse DDIM step
end for return Fto
3.2
Training objective
To train our network to sample the conditional random variable (Xoly, Xt); we take inspiration from Kim
iet al. (2024): Bendel et al. (2023) and combine reconstruction and adversarial losses, which has been shown to greatly improve for sample generation quality as opposed using a reconstruction loss alone. More precisely,
adapt the objective of Kim et al. (2024) and use the following objective based on an additional adversarial loss with joinly-optimized weights
arg min Sont
( 1h) Ladv(Do, 0) + we, Ce, (Do, v) + wPsCps(Aen) v),
(8)
arg max
(b) Lad U, 6) was
(9)
where Wh, wps. Was > 0 are regularization hyper- parameters and the remaining loss terms are defined as follows, where we let t be uniformly distributed on []. First, (A,, 0) is the f3-regression loss commonly encountered in DSM on the pixel space,
Leg (Do, v) N ||xo Lso. (x,, y)1
(10)
Second, (ps(Ae,0) promotes perceptual quality via the Learned Perceptual Image Patch Similarity (LPIPS) loss (Zhang et al. 2018)
Lrs(An, v/) = X LPIPS(x0, (xi,
(11)
Third, the adversarial loss, Ladv, is based on Goodfellow et al. 2014
Ladv (Ao, 0, Exo.y log (xxD)(Xo))))) + Ex,,,, log (1 (Lom (),,
(12)
where IR + R is the sigmoid function. The discrimination D& is CNN adapted from Radford al. (2016); Bendel et al. (2023), which we train to maximize the discriminator loss ( To regularize the discriminator, we follow Mescheder et al. (2018) and introduce gradient-sparsity regularizer designed to promote Lipschitz regularity of Dø, defined by
CGs(Ø) = Ex.x. V Do (ux [ux + u)LL),
(13)
where u ~ (((, is independent of xo and y. Note that our presentation of the training objective assumes RAM initialization, the trainable parameters of Lo 9 == (Ao, v). If th is frozen or if the first LATINO module is initialized directly with A'y and Iz, then
the same training objective reduced to D As.
As mentioned previously, using Equation is closcly related to training LS..y as consistency trajectory model (Kim et al. 2024] to sample from xx)) Similarly to CMs, our distilled model is trained for sampling, not for denoising. However, CMs are deterministic whereas L is a stochastic map, and we do not sook enforce consistency of paths as result. Also, we use a discriminator D3 that conditions on y only, whereas training CM for (xply, x,) by following Kim et al. (0224) would require using a discriminator that conditions both y and 24- This would also require larger and complex discriminator architecture
8

---

Published in Transactions on Machine Learning Research (11/2025)
with time embedding, which we do not consider here, we find that lighter discriminator that conditions y only suffices for correctly training the proposed UD sampler.
Before concluding this section, we present two alternative interpretations of the proposed deep unfolded architecture. Both are based on deep unfolding of discretizations of the same overdamped Langevin process that underlies LATINO, specifically the PnP-UT.A scheme introduced by Laumont et al. (2022). Tn the first interpretation, our architecture is viewed as a deep unfolding of PnP-ULA using a pre-trained SNORE denoiser [Renaud et al. 2024) adapted via LoRA. In the second, it is interpreted as deep unfolding of the DM-within-ULi ULA scheme proposed by Mbakam et al. (222)). In both cases, the likelihood is incorporated through its proximal operator (Durmus et al.] 20118), rather than through the more conventional gradient-basod approach. Among these interpretations, we find the deep unfolding of LATINO to be the most meaningful. LATINO is itself a distilled few-step scheme that incorporates the likelihood through its proximal operator and generalizes naturally to imago-text latent. DMs such as Stable Diffusion (Yin et al.) 2024). which are socing increasing adoption. Developing UD2M samplers by finetuning latent DMs is a natural extension of this work and main direction for future research.
4 Experiments
4.1 Experimental setup.
Preliminaries. To demonstrate the effectiveness of the proposed method, we now report a series numerical experiments and comparisons with competing approaches from the state of the art In our experiments, we consider three forms of image deblurring (Gaussian, uniform and motion deblurring), two forms of image inpainting (box and uniformly randomly missing pixels), image super-resolution (SR) by a factor 4, as well as the removal of artifacts from aggressive JPEG compression (as an example of a non-linear restoration problem). Moreover, we assess the performance of our method qualitatively and quantitatively by computing the following quantitative metrics: the Peak Signal-to-Noise Ratio (PSNR) as distortion metric (Wang et al. 2004), the Learned Perceptual Image Patch Similarity (LPIPS) for perceptual quality (Zhang et al.] (2018). and the Frechet Inception Distance (FID) with inception-v3 embeddings for sampling accuracy (Heusel el al. 2017
Datasets and experimental conditions. We use the following two public datasets in our experiments ImageNet [Russakovsky et al. 2015) and LSUN Bedroom (Yu et al. 2015) which have been used extensively in prior work related to image restoration with DMs and CMs in particular. For both datasets, we crop images to a size of 256 x 256 pixels and normalize their pixel amplitude to the range of [o, 1]. For our experiments with the ImageNet dataset, we used one million images from the training set for model training, whereas for the LSUN Bedroom datasct, we used 1.2 million images for training. For computing performance metrics, we use a test set of 1500 images from ImageNet and a Lest set of 300 images from LSUN bedroom. Regarding the parameters of each experiment, for ImageNet restoration tasks we follow Chung et al. (2023), with the exception of motion deblurring where we follow Zhu et al. (2023). For LSUN bedroom restoration tasks, for random inpainting, Gaussian deblurring and SR (x4), we use the setup as Garber & Tirer (2025) | with a noise-level a= = 0.025; we follow Zhao et al. 2025 for box inpainting. For more details about the considered experiments, please see Appendix C.3
Comparisons with the state of the art We compare our proposed method with an extensive selection of zero-shot and finc-tuned diffusion models from the state of the art For the experiments with ImageNet, we compare with the diffuuson bridges I'SB (Liu et al. 2023a) and CDDB (Chung et al. 2023). which are trained specifically for posterior sampling for ImageNet and for particular tasks. In addition, we compare with the two zero-shot methods DifIPIR (Zhu et al. 2023) and DDRM (Kawar et al. 2022). which rely on the same pretrained DM for ImageNet as our method, but differ in how they incorporate the data fidelity term.
For the experiments with LSUN bedroom, we use the same methods mentioned previously and two additional ones that rely on a CM trained and distilled for this datasct: the zero-shot method CM4IR (Garber & Tirer 2025). and CoSIGN (Zhao et al. (2025), a learning- based method which upgrades an unconditional CM into conditional CM for posterior sampling by learning an operator-specffic ControlNet. guidance that is
9

---

Published in Transactions on Machine Learning Research (11/2025)
attached onto the CM Zhao et al more computationally efficient than our unfolding approach, but lacks flexibility and leads to less accurate results as it does not use an explicit likelihood function.
2025). Our experiments show that CoSIGN's ControlNet guidance is
Implementation of the proposed method We implement our method as follows. For our deep unfolded architecture, we use K = 3 LATINO modules with a pre-trained diffusion model d we use the U-NET architecture of Ho et al. (2020) with time-embeddings and attention layers trained on a linear noise schedule 31 10-1 up B1000 0.02. With regards to training, we use the training objective detailed in Section 3.2 and LoRA adaptation with rank 5, while keeping original weights 0 frozen. Moreover, for each experiment, with the exception of JPEG artifact removal, we implement and train two versions of our model. In one version, the first LATINO module is initialized directly with a weighted average of the inputs A'y and Ty. In the other, the first LATINO module is initialized by a RAM module (Terris et al. 2025) with stacked operators (I, A) that take He and y as input, and which we train end-to-end together with he LoRA adaptation Ag by using the objective of Section 3.2 We do not consider RAM initialization for the JPEG artifact removal task because RAM is not suitable for non-linear image restoration problems. For more details about the model's architecture, training parameters, GPU and time consumption, please see Appendix C
Extension to non-linear inverse problems Similarly to Spagnoletti et al. 2025): to apply our method to JPEG artifact restoration and other non-linear image restoration problems we solve proxsg, (s) = arg min gy( + # 2o- 2/ with yy. logp(y, xy 20) approximately by using an optimizer -we use Adam KKngmma & Ba 2014) with s as warm-starting. We do not use RAM initialization.
4.2 Results
Table D and Table 2 below summarize the performance metrics for our proposed method, with (w/) and without (wo) RAM initialization, for the considered ImageNet and LSUN Bedroom experiments respectively. For comparison, we also report the performance metrics for the alternative approaches from the state of the art. The results for the ImageNet motion deblurring task are reported separately in Table 0 as there are no publicly available finetuned models for posterior sampling for that lask for comparison. As mentioned previously, we do not use RAM initialization for JPEG artifact removal because RAM does not support non-lincar problems, and we do not use RAM initialization for box inpainting because it is a lightweight. model with receptive field that is not sufficiently large to meaningfully inpaint large regions.
We observe from Table ! and Table 2 that the proposed approach performs very strongly relative to the state of the art across all tasks and all metries, especially when RAM initialization is used. Notably, our approach very clearly outperforms the alternative methods from the state-of-the-art in terms of FID and LPIPS scores, indicating excellent perceptual and sampling quality. For example, for the SR (x4) task, our method with RAM initialization achieves a FID score of 11.9 for ImageNet and 19.48 for LSUN Bedrooms, largely outperforming the second best methods (CDDB and CoSIGN respectively) which achieve scores of 19.88 and 40.48. In addition, the proposed approach is highly computationally efficient, as it only requires in the order of 10 NFEs per posterior sample, key computational advantage relative to non-distilled methods. This illustrates the benefit of training with the CM objective of Equation X that induces distillation. Having said that, out method is not as efficient as CoSIGN, which does not have the cost of evaluating K = 3 unfolded L.ATINO modules per step of the multi-step scheme.
For illustration and qualitative evaluation, Figure B Figure 4 and Figure 5 depict examples of posterior samples for the ImageNet SR (x4), Gaussian deblurring and JPEG restoration tasks, respectively. We observe that our proposed approach (w/ RAM initialization) delivers sharp posterior samples that recover significant amount of the fine detail in the ground truth images, and without noticcable noise color bias. With regards to the LSUN experiments, Figure O and Figure 7 present results for the SR (x4) and box inpainting tasks, respectively. Again, we observe that our proposed approach (w/ RAM initialization) delivers sharp posterior samples with fine dotail and no noticeable noise or color bias, in agreement with the
strong performance metrics reported in Table 2
For ImageNet, we use the un- on-(class) (claclllsoondtiionn condittoonl model http::///tthhh..co///pap////aatt and for LSUN bedroom the model tnttp:///githn catiub..coo/peeeeeeeee wwwp.///eaeeeepottct.ttaefeeee diffuuisooo
10

---

Published in Transactions on Machine Learning Research (11/2025)
Table 1: Quantitative results for the deblurring and super-resolution tasks. PSNR [dB] (t), LPIPS (L) and
FID (L) results on ImageNet dataset with noise level 0.01. Comparisons with ISB and CDDB are taken from
Chung et al. (2023).
Deblurring
SR
JPEG
Gaussian
Uniform
x4
QF=10
Methods
NFEs
PSNR
LPIPS
FID
PSNR
LPIPS
FID
PSNR
LPIPS
FID
PSNR
LPIPS
FID
Ours RAM)
38.77
4.61
35.57
11.14
24.42
0.15
20.69
27.52
0.18
35.16
Ours (w/ RAM)
35.97
3.30
36.96
0.01
2.69
26.70
11.9
CDDB
37.02
5.01
31.26
0.19
23.15
26.41
0.2
19.88
26.34
0.26
19.48
12SB
1000
36.01
5.8
30.75
23.01
24.13
20.35
DiEPIR
100
28.10
0.13
21.53
31.44
0.10
20.20
20.39
70.45
DDRM
36.73
0.07
4.34
29.21
0.21
19.97
26.05
0.27
46.49
26.33
0.33
47.02
Table 2: Deblurring, inpainting and super-resolution. PSNR (dB) (t), LPIPS (4) and FID (U) results on LSUN bedroom dataset with noise level 0.025. Comparisons with PSB and, CDDB and CoSIGN on the Box inpainting and SR (x4) problems are taken from Zhao et al. 22025 The remaining comparisons are obtainod
from Garber & Tirer 2025
Inpainting
Deblurring Gaussian
SR
Box
Random
x
Methods
NFFs
PSNR
LPIPS
FTD
PSNR
LPIPS
FID
PSNR
LPIPS
FID
PSNR
LPTPS
FID
Ours (wo RAM)
6"
22.22
0.11
8.45
22.32
0.19
43.3
26.73
0.26
88.97
24.51
0.19
32.71
Ours (w/ RAM)
12
27.75
0.06
8.61
29.67
0.08
13.23
25.48
0.12
19.48
COSIGN
2
22.61
0.14
38.64
23.22
0.37
19.74
0.342
26.13
0.22
40.84
CDDB
1000
23.74
0.13
45.20
27.31
0.24
54.20
FSB
1000
23.21
0.28
55,10
27.23
0.21
53.40
DiAPIR
100
13.42
0.34
66.79
23.8
0.36
97.69
27.48
0.32
64.81
23.83
0.45
101.92
DDRM
20
18.90
0.22
51.50
19.16
0.55
28.94
0.22
25.09
0.37
CM4I
21.98
0.24
42.83
25.28
0.33
159.43
29.00
0.21
53.19
26.14
0.30
124.22
"We have implemented our method with K = and 3 (so NFEs) for all experiments, except. box inpainting where we have used
5 instead 3 (so NFEs), as this is more challenging task.
Ours (wo RAM) Ours (w/ RAM) Figure 3: Comparison of posterior samples for the task SR (x4) with noise level 0.01 on ImageNet 256.
Ground Truth
Degraded
DifIPIR
DDRM
CDDB
4.3
Ablation Studies
We now proceed to examine and quantify the role of different key elements of our proposed method. This is achieved by appropriately modifying deactivating these elements and quantifying changes in performance.
11

---

Published in Transactions on Machine Learning Research (11/2025)
Ground Truth
Degraded
DifPPIR
DDRM
CDDB
Ours (wo RAM) Ours (w/ RAM)
Figure Comparison of posterior samples for the task Gaussian deblurring with noise level o 0.01 on ImageNet 256.
Ground Truth
Degraded
CDDB
Ours (wo RAM)
Figure 5: Comparison of posterior samples for the task JPEG artifact removal (QF=10) with noise level
= 0.01 on ImageNet 256.
4.3.1 Choice of the number of sampling steps N
First we assess the effect of modifying the number of sampling steps N in the multi-step scheme. Note that N is specified during test time, so modifying does not requires retraining the model. Table 3 a) shows the performance metrics for N (1, 3, 27) for the LSUN SR (x4) model with K LATINO modules and RAM initialization (recall that in our previous SR (x4) experiments We observe in Table 3() that the performance metrics do not change significantly with variations of N, indicating that the proposed approach is robust to different choices of N. Low values of N lead to some small amount estimation bias towards the mean, and hence to a small increase PNSR at the expense of LPIPS and FID performance, whereas larger values of N lead to mild decrease overall performance due accumulation of errors (see Kim al. (2023) for details) well as an increase in computational cost from additional NFEs. As mentioned previously, in experiments used N 3, which found to reliably
have used K = and N 3).
12

---

Published in Transactions on Machine Learning Research (11/2025)
Ground Truth
Degraded
CM4IR
CoSIGN
Ours (wo RAAM Ours (w/ RAM)
Figure 6: Comparison of posterior samples for the task SR (x4) with noise level o = 0.025 on LSUN Bedroom.
Ground Truth
Degraded
DDRM
CDDB
CoSIGN
Ours (wo RAM)
Figure 7: Comparison of posterior samples for the task box inpainting with noise level o 0.025 on LSUN
Bedroom.
provide good performance with low computational cost across all the considered tasks (with the exception of box inpainting, when N = as the task is challenging and we don't use RAM initialization).
4.3.2 Effect of unfolding
We now assess the benefit of unfolding multiple LATINO modules per sampling step, as opposed to using a single LATINO module, with the same or possibly a larger number of sampling steps. Table 3(b) reports the performance metrics when retraining the network with a single LATINO module (K = 1) and RAM initialization. By comparing Table 3 b) with Table 3a), we observe a noticeable deterioration performance sampling quality as measured by FID, even if the reduction the number of LATINO modules per step is compensated by additional sampling steps. We conclude that there is clear benefit in distilling conditional DM architectures from deep unfolding, which is consistent with the literature on deep unfolding of optimization
Table Reconstruction metrics for varying number of conditional diffuusion steps (N) and unfolded iterations (K).
N
PSNR
LPIPS
FID
NFEs $
N
PSNR
LPIPS
FID
NFEs
1
25.93
0.11
14.07
4
1
25.60
0.11
17.95
2
3
25.88
0.10
13.15
12
3
25.61
0.11
16.59
6
25.54
0.10
15.10
36
24.72
0.12
20.42
18
27
25.02
0.11
17.47
108
22.82
0.12
21.93
54
(a) K=3
(b) K-1
13

---

Published in Transactions on Machine Learning Research (11/2025)
algorithms MMnga et al, 2021 p. It is also worth mentioning that modifying K without retraining the DM leads to severe visual artifacts (results not reported here), so it is not possible to modify K during inference.
4.3.3 The choice of the rank of the LoRA adaptation A
We now assess the effect of the LoRA fine-tuning rank parameter. As explained previously, rather than fine-tuning the full model, we use a LoRA approach whereby we train a low-rank correction Ag that is added to the original model weights 8, which remain frozen. The number of trainable parameters in Ag is directly related to the choice of the rank of Ag. Setting the rank too low deteriorates performance because it constrains the model too much, whereas setting the rank too high increases the computational cost of fine-tuning the model and can lead to some overfitting. In order to illustrate this, we train again our model for ImageNet SR (x4) by setting the rank to and to 20 (in our previous results we have used rank 5). The number of trainable parameters in Ag is 165888 when the rank of Ag 2, 414720 when the rank is 5, and 1658880 when it is 20; this represents 0.03%, 0.08% and 0.3% of the total number of parameters in 0, respectively. The results for this study are summarized in Table 4 We observe that the model's performance, measured by PSNR and LPIPS, is robust to large variations in the rank of Ag. However, the rank of Ae has noticeable effect on FID performance, which drops significantly if it is set too low, indicating clear deterioration in sampling quality (e.g., in Table the FID for rank 2 is almost twice that of rank 5).
Table 4: Ablation study evaluating the impact of varying LoRA ranks on the ImageNet dataset for the super-resolution task.
rank(Ae) 2
5
20
PSNR
26.05
26.70
26.32
LPIPS
0.11
0.08
0.09
FID
19.75
11.9
13.03
|
165888
414720
1658880
4.3.4 Generalization w.r.t. to the noise variance o²
Models obtained by deep unfolding often generalize better to different levels of measurement noise, by comparison other end-to-end training rategies that are not aware of the value of the noise variance o² during test time. To demonstrate this capability, we follow (Garber Tirer 2025) and evaluate our LSUN bedroom SR (x4) model trained for 0.025 (with N 3 and K 3) in a more challenging situation (a 0.05). The results are summarized in Table 5 For comparison, also report the results obtained with the CoSTGN and CM4IR models applied to the same problem (results from (Garber & Tirer 2025 Table 2)), where CoSIGN is also finetuned for a = 0.025 and where CM4IR, as zero-shot method, is provided the correct value of a = 0.05 during test time. We observe that our proposed method is able to use the value of a = 0.05 effectively during test time, even if training was performed using o = 0.025, similarly to zero-shot methods like CM4IR. Conversoly, CoSIGN struggles to generalize the higher noise level and exhibits poor PSNR and LPIPS metrics as result. Again, by comparison to CM4IR, we achieve significantly better perceptual quality as reflected by lower LPIPS metric. For completeness, Figure S depicts examples of obscrvations and samples for our method when a = 0.0125, a = 0.025 (in-distribution) and = 0.05.
Table 5: Reconstruction statistics for our model applied to the SR (x4) task on LSUN bedroom, with varying noise level in the observation space.
Ours
CoSIGN
CM4IR
Noise lovel o
PSNR
LPIPS
PSNR
LPIPS
PSNR
LPIPS
0.025 (in-distribution)
25.12
0.178
26.10 20.35
0.205
26.14
0.295
0.05
24.11
0.214
0.509
25.60
0.320
14

---

Published in Transactions on Machine Learning Research (11/2025)
0.125
0.025
0.05
Cround Truth
Observation
Sample
Obscrvation
Sample
Obscrvation
Sample
Figure 8: Comparison of posterior samples for the task SR (x4) on LSUN Bedroom, for three levels of noise,
Table 6: Reconstruction metrics computed the LSUN Bedroom validation set using forward operators for the model trained T.SUN Bedroom SR (x4).
CoSIGN
Ours
Task Operator
PSNR
LPIPS
PSNR t
LPIPS +
SR. (x8) dut-o--distribution
20.11
0.55
20.40
0.32
SR (44 in-distribution
26.13
0.22
25.40
0.12
4.3.5 Generalization w.r.t. to more challenging degradations
We now explore the capacity of our proposed approach to deal with more challenging forward models not seen during training. For this exploratory analysis, we apply our LSUN Bedroom SR (x4) model to a more challenging SR x8 task (we keep 0.025, as training). Again, we report comparisons with the LSUN Bedroom SR x4 CoSIGN model. Although both models trained for SR (x4), our method receives the correct SR (x8) forward operator during inference. Conversely, CoSIGN does not have access to the SR (x8) forward operator, it relies on pre-trained ControlNet guidance. Through this exploratory experiment, we seek to show that models obtained by deep unfolding are more robust to challenging degradations (or forward operators) not encountered during training, by comparison to other end-to-end training strategies. A summary of these results is presented in Table G and an example in Figure 9 We observe that, while both methods struggle with this difficult restoration task, our method clearly outperforms CoSIGN which is unable use the forward operator during test time and exhibits strong reconstruction artifacts as a result. For completeness, Figure 9 also presents the result obtained with the CM4IR zero-shot method, which also outperforms CoSIGN by using the forward operator A, but suffers from some smoothing because of the approximations involved in zero-shot inference.
NUM
Ground Truth
Degraded
CoSIGN
CM4IR
Ours (w/ RAM)
Figure 9: Comparison of posterior samples for the task SR (x8), with noise level 0.025, LSUN Bedroom with models trained for LSUN Bedroom SR (x4), to illustrate the models' capacity to generalize to more challenging tasks.
15

---

Published in Transactions on Machine Learning Research (11/2025)
4.3.6 Task specific versus universal image restoration models
In our previous experiments, we considered that A is unique and known a priori, or it is not unique but belongs to a known class of operators specific for the considered task (e.g., motion blurs). In principle, one could also consider training a single common model for a much wider range of tasks. In particular, it would be interesting to train "universal" Bayesian image restoration model, considered in [Terris] et al. (2025) for MMSE estimation. Unfortunately, developing such general models requires using much more sophisticated architecture than the one we consider here. For example, Terris et al. (2025) combines unfolding with several additional key elements not considered here (o.g., multiscale decomposition of the operators, Krylov updates). Without these substantial improvements, our proposed approach suffers from a drop in performance when we try to generalize too widely. To illustrate this phenomenon, we trained a single posterior sampling model for the ImageNet dataset, simultaneously on Gaussian, uniform and motion deblurring, random inpainting and super-resolution tasks. The corresponding results are summarized in Table 7 where we also report the results obtained with our task specific models for comparison. We observe that the task specific models outperform the model trained jointly on all the considered tasks, and in all the considered performance metrics (the difference is most noticcable for FID performance). The extension of the UD3Ms framework to support the development of general purpose posterior samplers that can tackle a wide range of image restoration tasks and a wide range of datasets, without retraining or compromising performance, is a main perspective for future work.
Table 7: Qualitative results of generalized models evaluated on the ImageNet dataset. include the FID score, average PSNR, and average LPIPS.
The reported metrics
Gaussian
Deb
Motion Deb
Uniform Deb
Inp 70%
SR (x4)
Spec
Univ
DiffPIR
Spec
Univ
Spec
Univ
DiffPIR
Spec
Univ
Spec
Univ
DifPPIR
PSNR
35.97
34.80
28.1
33.40
31.10
36.96
31.35
31.44
31.15
30.16
26.70
25.74
20.39
LPIPS
0.01
0.03
0.13
0.04
0.04
0.01
0.08
0.1
0.03
0.06
0.08
0.12
0.36
FID
3.30
4.94
21.53
4.08
6.56
2.69
19.6
20.2
7.87
13.21
11.9
21.34
70.45
5
Conclusion
5.1
Concluding remarks and perspectives for future work
Diffussio models are powerful image priors for Bayesian computational imaging (Daras et al. 2024h. The literature offers two main strategics to construct Bayesian imaging techniques based on diffusion models: Plug & Play methods that are zero-shot and hence highly flexible but reliant on approximations, and y-conditional diffusion models that achieve superior accuracy and speed through supervised specialization for particular task. The framework introduced in this paper, unfolded and distilled diffusion models (UD?Ms), leverages deep unfolding and distillation in order to upgrade diffusion model representing the image prior into a few-step consistency model for posterior sampling. UD2M samplers offer excellent accuracy and high computational efficiency, while retaining the flexibility required to generalize to variations in the forward operator and noise distribution during inference. A main novelty of UD2M is to unfold a Markov chain Monte Carlo algorithm, namely the recently proposed LATINO Langevin sampler (Spagnoletti et al. 2025), the first. example of deep unfolding of Monte Carlo scheme.
With regards to limitations of our proposed framework and perspectives for future work. We are cognizant that diffusion models can reproduce biases stemming from unbalanced or biased training data and distribution shift, diminishing the reliability of our framework. Future work should develop appropriate guardrails so that DM-based posterior samplers can be deployed safely and reliably. Also, we hope and anticipate that these issues will be partially mitigated by progress and democratization of deep generative modeling technology and the advancement of self-supervised deep generative models. Moreover, while our samplers are more computationally efficient than some alternatives, they still rely on very large models and therefore consume significant amount of energy and compute resources. The development of more frugal architectures and
16

---

Published in Transactions on Machine Learning Research (11/2025)
sampling strategies is an important perspective for future work. Furthermore, mentioned previously, it would be interesting to study architectures and training strategies to develop UD²M samplers that can be applied to a wide range of image restoration tasks and datasets without the need for additional training, or in a few-shot manner, as is the case of the foundational image reconstruction model of Terris et al. (2025). In addition, as mentioned previously, many state-of-the-arl DMs leverage the latent space representation of an autoencoder. For example, Stable Diffusion models (Yin al. 2024) that allow semantic conditioning through text prompting. Extending UD²M samplers to these DMs is a promising direction for future work.
Lastly, in parallel with this work, the concurrent works Wang et al. [2025] Elata al. 2025 also seek to bridge supervised and zero-shot strategies for Bayesian imaging with diffusion models. Wang et al. (2025) relies on unfolding & conventional proximal gradient optimization scheme and uses an 82 loss in order to obtain a conditional diffusion model, without distillation. Elata et al. (2025) proposes a novel approach to incorporate the forward operator A and the mcasurement y directly within Transformer or UNet architectures commonly used for diffusion models. Contrasting and combining UD2M with these alternative strategics is important perspective for future work.
5.1.1 Broader impact statement
Computational imaging research generally poses some ethical dilemmas, espocially when involves significant methodological innovation. Although the innovations studied in the project do not raise any direct concerns, are aware that the proposed Bayesian imaging framework could be adapted and transferred towards military applications by a hostile actor and for violating citizen privacy, for example. In addition, our framework relies on deep generative modeling technology that currently suffers from biases, which could lead to leading to unfair or discriminatory outcomes. We hope that these critical issues will be resolved by progress and democratization of generative modeling technology.
Morcover, the innovations presented in this paper fall within the "limited risk" level of the EU AI Act. We have followed the Act's recommendations for this level; i.e., focus the full transparency of the training data and the methods used. This also links in with our open research philosophy available open-source implementations for research codes, with clear documentation and demonstrations. We see this as an essential requirement for conducting research in a democratic knowledge society,
will release publicly
Acknowledgments
We are grateful to Andres Almansa, Abdul-Lateef Haji-Ali, Alessio Spagnoletti and Konstantinos Zygalakis for useful discussion. This work was supported by UKRI Engincering and Physical Sciences Research Council (EPSRC) (EP/V006134/1, FPP7553481/1) We acknowledge the use of the HWC high-performance computing facility (DMOG) and associated support services in the completion of this work.
References
Matthew Bendel, Rizwan Ahmad, and Philip Schniter. A regularized conditional GAN for posterior sampling in image recovery problems. Advances in neural information processing systems, (6:68673-68666 2023.
Antonin Chambolle and Thomas Pock. An introduction to continuous optimization for imaging. Acta Numerica, 25:161 319, 2016.
Mostafa Cherif, Tobías I Liaudat, Jonathan Kern, Christophe Kervazo, and Jérôme Bobin. Uncertainty quantification for fast reconstruction methods using augmented equivariant bootstrap: Application to radio interferometry. Xixiv preprint 0rXiv:2410.233 28.23778 2024.
Hyungii Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffussio posterior sampling for general noisy inverse problems. In The Eleventh International Conference on Learning Represenlations, 2022.
Hyungjin Chung, Jeongsol Kim, and Jong Chul Ye. Direct diffusion bridge using data consistency for inverse problems. In Thirty Conference on Neural Information Processing Systems, 2023. URL https://openreviev.c http://eeeeeee.eeee..o//o/t//t/r
17

---

Published in Transactions on Machine Learning Research (11/2025)
Florentin Cocurdoux, Nicolas Dobigeon, and Pierre Chainais. Normalizing flow sampling with Langevin dy- namics in the latent space. Machine Leurning, 13((11:830118834) 2024a. doi: 10. 10.1007/s109944444-000222-3 URL http:///...rrg11000000/000000000
Florentin Cocurdoux, Nicolas Dobigeon, and Pierre Chainais. Plug-and-play split Gibbs sampler: Embodding deep generative priors Bayesian inference. IEEE Transactions on Image Processing, 2024b.
Giannis Daras, Hyungjin Chung, Chich-Hsin Lai, Yuki Mitsufuji, Jong Chul Ye, Peyman Milanfar, Alexan- dros G. Dimakis, and Mauricio Delbracio. A survey diffusion models for inverse problems, 2024. URL
https://arxivv org h/t///b//2/0 /4100 241000000800003333
Mauricio Delbracio and Peyman Milanfar. Inversion by direct iteration: An alternative to denoising diffusion for image restoration. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL Featured Certification, Outstanding Certification.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat GANs on image synthesis. Advances in neural information processing systems, 34:8780-877 8794, 2021.
Alain Durmus, Eric Moulines, and Marcelo Pereyra. Efficient Bayesian computation by proximal Markov chain Monte Carlo: When Langevin meets Moreau. SIAM Journal Imaging Sciences, 11(():473-506 2018.
Bradley Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association, 106 1496)11602-111,, 2011.
Noam Elata, Hyungjin Chung, Jong Chul Ye, Tomer Michaeli, and Michael Elad. Invfussion: Bridging supervised and zero-shot diffusion for inverse problems, 2025. URL httpp::/aaxxx...or///////22 01689
Tomer Garber and Tom Tirer. Zero-shot image restoration using few-step guidance of consistency models (and beyond). aixiv preprint (rXiv:2412.20530 2025.
Ian Goodfellow, Jean Pouget-Abadie, Mchdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in Newral Informalion processing systems, 27, 2014.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by two time-scale update rule converge to local Nash equilibrium. Advances in neural information processing systems, 10, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information processing systems, 33:6840-6851, 2020.
Matthew Holden, Marcelo Pereyra, and Konstantinos C. Zygalakis. Bayesian imaging with data-driven priors encoded by neural networks. SIAM Journal on Imaging Sciences, (5(2):892-924, 2022. doi:
10.1133//1MM140631.. URL https:://doi org.l0...111..11 11w1111711000 1137/2111106613
Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. ICLR, 1(2):3, 2022.
Jari Kaipio and E Somersalo. Statistical and computational inverse problems. Applied Mathematical Sciences. Springer, New York, NY, October 2010.
Bahjat Kawar, Michael Flad, Stefano Ermon, and Jiaming Song. Denoising diffusson restoration models. Advances in Neural Information Processing Systems, 35:23593-2338 236066 2022.
Charlesquin Kemajou Mbakam, Marcclo Percyra, and unadjusted Langevin method for empirical Bayesian estimation in semi-blin image deblurring problems. STAM Journal on Imaging Sciences. to appear.
Giovannelli. stochastic optimisation
18

---

Published in Transactions on Machine Learning Research (11/2025)
Dongjun Kim, Chich-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning probability flow ode trajectory of diffusion. arXiv preprint arXiv:2310. 02279, 2023.
Dongjun Kim, Chich-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning probability flow ODE trajectory of diffusion. In 12th International Conference on Learning Representations, 2024. URL https://openreview nett//orruumddddym]88eee
Diederik P Kingma and Jimmy Ba.
Adam:
method for stochastic optimization. arv preprinl
arXiv:1412.00 6980, 2014.
Rémi Laumont, Valentin De Bortoli, Andrés Almansa, Julie Delon, Alain Durmus, and Marcelo Pereyra. Bayesian imaging using Plug & Play priors: When Langevin meets Tweedie. SIAM Journal on Imaging Sciences, 15(2):701-737 2022.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference Learning Representations, 2023. URL Po
https://openreviiet
G.H. Liu, A. Vahdat, D.A. Huang, E. Theodorou, W. Nie, and A. Anandkumar. 12SB: Image-to-inage Schrödinger bridge. In 40th International Conference on Machine Learning (ICML), 2023a.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Represenlations, 2023b. URL /
ped//forumiddxxjjttin
Eric Luhman and Troy Luhman. Knowledge distillation iterative generative models for improved sampling speed, 2021.
Simian Luo, Yiqin Tan, Suraj Patil, Daniel Gu, Patrick von Platen, Apolinário Passos, Longbo Huang, Jian Li, and Hang Zhao. Lem-lora: A universal stable-diffusion acceleration module, 2023. URL https:
//arxiv.org/abs/////
Pierluigi Maponi. The solution of linear systems by using the Sherman Morrison formula. Linear algebra and its applications, 420(2-3):276 2944 2007.
S. Martin, A. Gagneux, P. Hagemann, and G. Steidl. PnP-Flow: Plug-and-play image restoration with flow matching. In International Conference for Learning Representations, 2025 (to appear).
Charlesquin Kemajou Mbakam, Jean-Francois Giovannelli, and Marcelo Pereyra. Empirical Bayesian image restoration by Langevin sampling with denoising diffusion implicit prior, 2024. URL https:
S Mclidonis, M Holden, Y Altmann, M Pereyra, and K C Zygalakis. Empirical Bayesian imaging with large-scale push-forward generative priors. IEEE Signal Process. Lell., 31:631 635, 2024.
Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 14297 14306, 2023. doi:
Lars Mescheder, Andreas Geiger, and Sebastian Nowozin. Which training methods for GANs do actually converge? In International conference on machine learning, pp. 3481-3490. PMI.R, 2018.
Vishal Monga, Yuelong Li, and Yonina C. Eldar. Algorithm unrolling: Interpretable, efficient deep learning for signal and image processing. IEEE Signal Processing Magazine, 38:18-44, 2021. ISSN 15580792. doi: 10. 10.1109/MS9 20109/MMPPP0 20......00300690000
19

---

Published in Transactions on Machine Learning Research (11/2025)
Badr Moufad, Yazid Janati, Lisa Bedin, Alain Oliviero Durmus, randal douc, Eric Moulines, and Jimmy Olsson. Variational diffusion posterior sampling with midpoint guidance. In The Thirleenth International Conference on Learning Representations, 2025. URL https://openrevierr met:fforummii---Ett http:////paaaaare.e..eeeeeooataa
Subhadip Mukherjee, Andreas Hauptmann, Ozan Öktem, Marcelo Pereyra, and Carola-Bibiane Schönlieb. Learned reconstruction methods with convergence guarantees: A survey of concepts and applications. IEEE
Signal Processing Magazine, 40(1):164-182, 2023. doi:
Gregory Ongie, Ajil Jalal, Christopher A. Metzler, Richard G. Baraniuk, Alexandros G. Dimakis, and Rebecca Willett. Deep learning techniques for inverse problems in imaging. IEEE Journal on Selected Areas in Information Theory, (11):99-56 -6, 2020. doi: 10110///AATT.O 2020.2991566..
Marcelo Pereyra. Proximal Markov Chain Monte Carlo algorithms. Statistics and Computing, 26(4):745- 2016.
Marcelo Pereyra, Luis A. Vargas- Mieles, and Konstantinos C. Zygalakis. The split Gibbs sampler revisited: Improvements to algorithmic structure and augmented target distribution. SIAM Journal on Imaging Sciences, 16(4):2040 2071, 2023. doi: 10.1137/22M1506122. httpss////// 1ttp:/00.111/2222113 001002222000
Alcc Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with doep convolutional generative adversarial networks. In Proceedings of the International Conference on Learning Representations, 2016.
Marien Renaud, Jean Prost, Arthur Leclaire, and Nicolas Papadakis. Plug-and-play image restoration with stochastic denoising regularization. In Proceedings of the 41st International Conference on Machine Learning, ICML 24, JMLR.org, 2024.
Christian P. Robert. The Bayesion choice: From decision-theoretic foundations to computational implemen- lation, volume Second edition. Springer, 2007.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115(():211122)) 252, 2015. doi: 10.1007 (01007/s1100 1////11111----016---1-111
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In The Tenth International Conference on Learning Representations, 2022.
C. Schuler, M. Hirsch, S. Harmeling, and B. Schölkopf. Learning to deblur. IEEE Transactions on Pattern Analysis and Machine Intelligence, 38(7):1439-1444 July 2016. doi:
Michael Soloveitchik, Tzvi Diskin, Efrat Morin, and Ami Wiesel. Conditional Frechet inception distance. arXiv preprint arXiv:2103. 11521, 2021.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising Diffusion Implicit Models. In International Conference on Learning Representations, 2021. URL /atpps:/openrevvee: wtt.://rww.....e.oooooo/a/oaaaaa
Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-guided diffusion models for inverse problems. In International Conference on Learning Representations, 2023a. URL https
//app.ssmantiisccholla. wwwpsschhhaarrrroo///oeoo argg/oorpssDD:: 2592987115
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Newral Information Processing Systems, 2019. URL
Org/CorpusID: 196470877
Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems, 33:12438-12445 2020.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishck Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations.
20

---

Published in Transactions on Machine Learning Research (11/2025)
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. 2023b.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 32211-322522 PMLR, 23-29 Jul 2023c.
A. Spagnoletti, J. Prost, A. Amansa, N. Papadakis, and M. Pereyra. LATINO-PRO: LAtent consisTency INverse sOlver with PRompt Optimization. arXiv preprint (arXiv:2503.138) 12615), 2025.
Julian Tachella and Marcelo Pereyra. Equivariant bootstrapping for uncertainty quantification in imaging inverse problems.. arXiv preprinl arXiv:2310. 11838, 2023.
lightweight foundation model for computational imaging. arXiv preprinl ar arXim:2003.0800 2025. Matthieu Terris, Samuel Hurault, Maxime Song, and Julian Tachella. Reconstruct anything model: A
Maxime Vono, Nicolas Dobigeon, and Pierre Chainais. Split-and-auurnented Gibbs sampler- Application to large-scale inference problems. IEEE Transactions on Signal Processing, 0766) 1661, 2019.
Yuanhao Wang, Shirin Shoushtari, and Ulugbek S. Kamilov. Diff-unfolding: model-based score learning framework for inverse problems, 2025. URL
Zhou Wang, Alan C Bovik, Hamid R. Sheikh, and Eero P Simoncelli. Image quality assessment: From error visibility to structural similarity. IEEE transactions on image processing, 13(4):600-612, 2004.
Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, and William T Freeman. Improved distribution matching distillation for fast image synthesis. In NewrIPS, 2024.
Fishor Yu, Yinda Zhang, Shuran Song, Ari Seff, and Jianxiong Xiao. LSUN: Construction of a large-scale image dataset using deep learning with hurans the loop. arXiv preprinl Xrxiv:1506.00 03365, 2015.
Kai Zhang, Yawei Li, Wangmeng Zuo, Lei Zhang, Lue Van Gool, and Radu Timofte. Plug-and-play image restoration with deep denoiser prior. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44 (10):6360-637) 6376, 2021.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on Computer Vision and Pottern Recognition, pp. 586-595, 2018.
Jiankun Zhao, Bowen Song, and Liyue Shen. CoSIGN: Few-step guidance of consistency model solve general inverse problems. In Aleš Leonardis, Elisa Ricci, Stefan Roth, Olga Russakovsky, Torsten Sattler, and Gül Varol (eds.), Computer Vision ECCV 2024, pp. 108-126. Springer Nature Switzerland, 2025.
Ningning Zhao, Qi Wei, Adrian Basarab, Nicolas Dobigeon, Denis Kouamé, and dean-Yves Tourneret. Fast. Single Image Super- Resolution Using a New Analytical Solution for lg -2222 Problems. IEEE Transactions Image Processing, 25(8):3683-360) 3697, 2016.
Yuanzhi Zhu, Kai Zhang, Jingyun Liang, Jiezhang Cao, Bihan Wen, Radu Timofte, and Lue Van Gool. Denoising Diffusion Models for Plug-and-Play Image Restoration. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pottern Recognition, pp. 1219 1229, 2023.
21

---

Published in Transactions on Machine Learning Research (11/2025)
A Solving the proximal step in Figure 2
This section outlined the approach for solving the proximal step in Figure tailored to each image restoration task in our study.
The degradation model for deblurring is generally express as Ax + where y is the observation, n is zero mean Gaussian noise with variance o². The forward operator A represents circular block matrix, and can be transformed as A = FIFFF where F and FF denote the fast Fourier transform (FFT) and its inverse, respectively, while AA diagonal matrix containing the Fourier coefficients of the blurring kernel. A fast solution of the proximal step Figure 2 is given by
Deblurring
81,k = Ft
(P
Fx Va
F
L.ke
Σズ
I + o 8
+
1
(14)
Inpainting For inpainting task, the degradation process formulated as = A o + n, where A denotes the mask operator and represents the elementwise multiplication operator. By using the well known Sherman-Morrison formula. (Maponi 2007), we obtain the following fast solution for the proximal step in Figure
Ag Lt.k
Tt
8r,k Kk
I
- Ri
o
(15)
Kk
v vara
d0? and denotes the elementwise multiplication operator. 6+o7
where Ką
ruseriresolution In super-resolution task, we formulate the problem by assuming that the low-resolution image is a blurred, downsampled, and noisy version of the high-resolution image. Mathematically, the model is given by y = S.Ax + n, where S, represents the downsampler operator with factor s and represents the circular blur operator. According to Zhao et al. (2016); Zhang et al. (2021) a fast solution for the proximal step in Figure is given by
to o, 1 [ (EA
EAr) ts
81, = Fi
(16)
lik
(> to +
where r r= EAS y Kb Ezx + Vana, Mit -, and o, represents distinct block processing operator with element-wise multiplication. The opcrator ts denotes distinct block downsampler. For more details, see Zhao et. al. (2016); Zhang et al. 2022))
Extension to non-linear tasks: JPEG For the non-linear JPEG restoration task, the observation model is given by y A(x) n, where A denotes the non-linear compression operator and is the zero- mean Gaussian noise with variance o². In this case, computing the exact solution of the proximal step prox (s) arg min (xp) + 28
1 1|zy s1l2/2 with 9y.x, = - log p(v, x1(o)) is not feasible. To address
this challenge, we approximate the solution by using Adam optimizer (Kingma & Ba 2014 with s as warm-starting and we do not use RAM initialization.
B Relation to Langevin Dynamics
In this appendix, present a derivation of the LATINO kernel equation 5 small-timestep discretisation with a score-based prior of the regularized Langevin SDE:
dxo,s V log ds + V log )ds dw..
which is ergodic with respect the posterior distribution p(xly). An alternative derivation, based on large timestep approximation of the Langevin diffusion with consistency model prior can be found in Spagnoletti
22

---

Published in Transactions on Machine Learning Research (11/2025)
al.
consider an auxiliary variable z such that, conditioned 2, we have ~ 6), The marginal process t0,s ~ (z|x = Rl 20,s]) then follows the diffusion
2025 Following Pereyra et al. (2023).
dzo,s Xu.. Zu 2
do Flogp(zo,))de dws:
28
(17)
X0. = E xo . 20.al,
where express the augmented likelihood V llgzp(yiz, 6) in terms of Xo,s and Zo,s through Fisher's Identity as in Pereyra et al. (2023). The short-time continuous dynamics of 2o,4 for h & 11 can be estimated through the proximal splitting scheme
1)) h zo,n E00x(1- (1- xo.n + V2he
Zo,h Prox- log
~~ (0,1).
Sctting h and approxinating the proximal operator by minimum mcan-square denoiser Dş for the convolution of p(z) with O)) additive noise as in Laumont et al. (2022), we oblain the stochastic noise-and-denoise step
Zu,A = D8 Ko,k + V28c).
Ergodic convergence guarantoes for plug-and-plan discrctisation the Langevin diffusion can be found in Laumont et al. 2022). In a related work, Martin et al. (2025 (to appear) consider the use of an identical noise-and-denoise within stochastic gradient descent algorithms with provable convergence guarantees to a maximum-a p-osterrooi estimate. Setting 8 ((1 digs //2r, and multiplying by factor of vou, within the denoiser we have
20,b D (/2-0,))/(( ag Xo, VI- a.g
Applying this kernel iteratively, approximating the MMSE denoiser ((1-a)))22 by the score-based prior Ge(:,ts), we obtain the updates
(k+1) E[xly, ((),
x
= =Proxsg, (z)
(++1) Gool ai, X (k+1)
+ /1 - to,
where relate the conditional mean of x given y and z to the proximal operator of fy since p(x|z) defines conjugate Gaussian prior for p(x|y, 2). This recovers the LATINO algorithm equation under the mapping /ai,x( + VT- at = x (k+1) and (k+1) x (k+1). The kernel equation can thus understood as a
proximal split-Gibbs timestep for the Langevin diffusion with score-based plug-and-play prior. closely related approach was employed in Mbakam al. (2024) based explicit gradient discretisation of the
marginal diffusion Equation 17 and using small number of DDIM iterations to formulate stochastic
plug-and-play denoising prior. Following Pereyra et al. (2023 Laumont et al. (2022 Martin et al. (2025 (to
appeary, the parameter 6 (equivalently, (s) can be tuned to provide asymptotic convergence to the posterior distribution. However, as observed empírically in Appendix D.I. when applied as zero-shot method the MCMC chain arising from the LATINO updates requires a significant number of NFFs to reach stationarity. Following the successful unfolding of optimization procodures for imaging illustrated Monga et al. (2021). ) aim in this work to generate accurate posterior samples unfolding small number K LATINO iterations likelihood-awate neural network. By designing network in this manner, obtain modular and interpretable model with natural embedding of the forward observation model.
In Spagnoletti et al. (2025), the same algorithm is derived for large lime approximations R 1. This derivation relies on the use pre-distilled consistency model prior to construct a kernel which contracts to the same invariant measure as equation 17 key difference for the proposed approach Scction 3 is that we start with a diffusion model designed for small time-steps. Several iterations of LATINO are then distilled to construct conditional consistency model for determined class of image restoration problems.
23

---

Published in Transactions on Machine Learning Research (11/2025)
C
Implementation details
C.1 Models
In this work, we used two distinct pre-trained models for our experiments. The first model is derived from ADM (Dhariwal & Nichol 2021). originally pre-trained on the ImageNet dataset at the resolution of x56 256. The second model is based on a pre-trained model from DDPM (Ho et al. 2020), which was pre-trained on LSUN dataset at the resolution of 256x256. The architectures of these models differ significantly and a detailed comparison their architecture, parameter counts and computational requirements is provided in Table
Regarding LoRA configuration, implemented exclusively on the attention layers of the pre-trained model to improved its adaptability for specific image restoration tasks while reducing the computational overhead.
Table Models summarize: |0] denotes the total number of frozen parameters, Aa indicates the total LoRA parameters and |æs represents the total RAM parameters. Model 1 refers to the ADM pre-trained on the ImageNet dataset, while Model denotes the DDPM pre-trained on the LSUN dataset.
# #Alln
Models
rank( An)
[0]
A
layers
Model
16
552814086
0
Model 2
5
113673219
0
Model 1 w/ LoRA
16
552814086
414720
Model 2 w/ LoRA
5
5
113673219
122880
RAM
0
35618953
0
C.2 Training procedure
We adopted a training procedure that is carefully designed to ensure robust performance across tasks such deblurring, inpainting and super-resoutton. The training objective integrates reconstruction and adversarial losses to enhance sample generation quality. For the reconstruction loss, we used a consistency lerm to enforce data fidelity and a LPIPS term to enhance perceptual quality. Conversely, the adversarial term is designed to promote the model to generate realistic sample. Additionally, the training loss incorporates gradient penalty term to stabilize training by enforcing Lipschitz continuity on the discriminator. To facilitate reproducibility, comprehensive details regarding batch size (bs), learning rate (Ir), optimizer and specific weights for loss term are summarized in Table g
Table 9: Model 1 refers to the ADM pre- pre-trained on the ImageNet dataset, while Model 2 denotes the DDPM pre-trained on the LSUN dataset.
lr
bs
optimizer
weight decay
Weg
WGP
WGS
Model
1e-1
2
AdamW
0.01
1
0.1
0.01
Model 2 R.M.
1e-4 4
AdamW
0.01
1
0.1
0.01
le-5
AdamW
0.01
C.3 Setting experiments
Box inpainting We consider fixed and centered square mask of size 128 128 pixels, K and 0.025.
Random inpainting In this experiment, a random binary mask applied with missing pixel rates of 80% and 70%. We use K iterations for algorithm unrolling and noise level of 0.025.
24

---

Published in Transactions on Machine Learning Research (11/2025)
resortie:: For SR experiments with scaling factor m, images from the LSUN dataset are downsampled by convolving with bicubic filter of size m² x m². In contrast, for the ImageNet dataset, images are first convolved with a 3 x 3 Gaussian kerne and bandwidth of 10, followed by a downsampling operation with factor m. Unless otherwise stated, we consider downsampling factor m = 4. The number of ilerations algorithm unrolling is set K 3, and experiments are conducted with noise levels a E (0.01,0.025),
Deblurring For this experiment, consider types blur.
Gaussian blur: On Imagenet experiments, consider kernel sizes of 5 and 25 25 pixels and bandwidth of 10. For the LSUN bedroom dataset,
consider the anisotropic Gaussian blur used in Garber & Tirer 2025 : Kawar et al. (2022 of size 9 x 9 and with bandwidth (20, 1).
Motion blur:
We generate random motion blur kernels with kernel sizes of 5 5 using the approach
described in Schuler et al. [2016]
Uniform blur with kernel sizes of 9 x 9.
We conduct experiments under noise levels {0.01, 0.025, 0.05) and using iteeaaionn for algorithm
unrolling.
JPEG restoration For this experiment, we applied JPEG compression with quality factor of QF=10 and introduced Gaussian noise with variance 0.01
C.4 Memory and time consumption
Table 10 presents comparative analysis of memory consumption and computational time for the models evaluated this work. The results demonstrate that fine-tuning pre-trained models using LoRA technique reduces GPU memory usage and inference time. The table also includes the memory usage and inference time for RAM.
For the JPEG restoration task, we obtained an inference time of 4.01 seconds. This increased computational time is because the proposed method requires a sub-iterative method to solve the proximal step in Figure 2
Table 10: Models summarize: denotes the number of unfolding steps. Model refers the ADM pre-trained on the TeageNet dataset, while Model denotes the DDPM pre-trained on the LSUN dataset.
Models
K
GPU (GB)
Time (s)
Resolution
Model 1
16.49
256
Model 2
3.39
256
Model 1 w/ LoRA
0.013
0.89
256
Model 2 w/ LoRA
3
0.004
0.53
256
RAM
1.06
0.31
256
D Additional results
In this appendix, we provide additional qualitative results between UD2M applicd to the problems in Section 4 Qualitative comparisons to similar methods for deblurring and super-resolution on ImageNet. may be found in Figures and L3 Qualitative reconstructions obtained via the universal UD²M scheme discussed in
Scction 4.3.6 can be found in Figure 14
4The Gaussian kernel as low-pass pass filter reduce allasing before downsampling. 5 An implementation of this operator can found Deepinverse
25

---

Published in Transactions on Machine Learning Research (11/2025)
Additional qualitative comparisons of UD2M (with RAM) the LSUN bedroom validation set can be found in Figures 15 and T6 for the gaussian deblurring and random inpainting lasks, each with in-distribution noise level 0 = 0.025. To illustrate qualitatively the extensibility of the unfolded architecture to out-of-distribution noise, additional comparisons are presented in Figure 17 using noise level o = 0.05 which was not used in training UD For the comparison methods on LSUN bedroom, we use the reconst ructed images provided in Garber & Tirer (2025). To illustrate the generative capabilities of the UD2M prior, additional reconstructions for box inpainting on LSUN bedroom are provided in Figure [8
D.1 Ablation Study MNIST
We consider further ablation studies for inverse problems MNIST digits, with 60,000 training and 10,000 test images. For ease of computation, we add a 2 x 2 zero-padding to resize each image to shape 32 x 32. The small size of these images allows for efficieut training of several models for extensive ablation studies and for many posterior samples to be generated, allowing for accurate exploration of the learned posterior data. We test UD on a challenging inverse problem on this data, to benchmark the convergence of our method for a small number of NFEs.
Model For our deep unfolded UDRM architecture, we use score-based denoiser with the samne architecture as Dhariwal & Nichol (2021). The score-based denoiser is downsized to take as input grayscale images of size 32x32. We pre-train the score model to generate images from the 32 x 32 MNIST datasot using 220,000 training ilerations with a batch-size of 64. The final model contains 16,761,409 parameters.
Through adversarial distillation, we fine-tune the score-model into a conditional consistency model, generating conditional samples of zo given 2; and an observation y by unfolding K itcrations of the LATINO kernel. To compare models with different mumber of unfolded iterations, we train 4 separate models with K = 1, and K = 8, respectively. For each model, we use the same approach as above and freeze the weights in the pre-trained score and apply a learnable low-rank correction only to weights within the attention layers. We apply LoRA with rank r = 5, resulting in 65,280 trainable weights representing 0.39% of the original model.
Experiment We consider the SR (x4) problem, using additive Gaussian noise with variance of size o² = 0.05². Due to the small size of MNIST digits, down-sampling by a factor of 4 represents a significant. degradation in the quality of the image (see 10 By considering a highly ill-conditioned problem, we aim to highlight the efficioncy of UD²M to converge to an accurate posterior distribution with only a small mumber of NFEs. This allows an effective demonstration of the enhanced convergence rate of the unfolded LATINO iterates over K = iterations, which in contrast to the large K >> 1 number of iterations typically required for MCMC sampling (Pereyra 2016)
Comparison To benchmark against a statc-of-the art zero-shot MCMC sampler, we compare our method to the long-term dynamics of the LATINO kernel equation 5 Using the same pre-trained denoiser as discussed above, we run K Rs 50, 000 iterations of the LATINO algorithm. To evaluate the performance gain by embedding LATINO within distilled diffusson loop through UD*M, we do not apply fine-tuning or distillation to the kernel. Instead, interpreting LATINO a plug-and-play discretisation of the Langevin diffusion, tune the step-size by hand using the theoretical bounds in Laumont et al. 2022
Evaluation Metrics For the evaluation, we consider the PSNR and LPIPS metrics used above. The FID metric implicitly assumes that the Inception-v3 network accurately maps the prior distribution to a Gaussian in the encoding space. Due to limitations with the Inception-v3 network mapping MNIST images to Gaussian distribution, we replace this network with VAE encoder EMNIST pre-trained explicitly in Holden 2022 to encode the MNIST datasct into a 12-dimensional Gaussian distribution. The Frechet distance is then computed between VAE-encoded MNIST digits and VAE-encoded digits reconstructed through our model. Following Bendel et al. (2023), we compute posterior approximation quality using the conditional Frechet inception distance (CFID) Soloveitchik et al. (2021). For approximations & of the ground truth x, the CFID computes the conditional Frechet distance between the embeddings P(EMNIST (x) EMNIST (y)) and FaatnnsrraFFFNNsss(s(( (MNIST(y)) ( averaged over many observations y-
ct al.
((EMNIST(()
26

---

Published in Transactions on Machine Learning Research (11/2025)
2
16
2
16
K
K
14.8
15.5
15.5
15.7
15.9
0.18
0.28
0.22
0.14
0.11
2
16.3
17.1
17.3
17.5
17.5
2
0.07
0.10
0.07
0.05
0.05
4
18.1
18.2
18.2
18.2
18.2
4
0.05
0.04
0.04
0.04
0.04
8
18.5
18.6
18.6
18.6
18.6
0.04 8
0.04
0.04
0.04
0.04
(a)
PSNR
(b) LPIPS
N
N
1
2
4
8
16
2
4
8
16
K
K
1
4.87
3.82
3.44
2.86
2.49
1
6.91
5.16
4.83
4.08
3.53
2
2.81
1.86
1.16
0.78
0.73
2
3.65
2.25
1.47
1.01
0.94
4
0.31
0.24
0.24
0.23
0.22
0.47
0.38
0.37
0.36
0.35
8
0.15
0.15
0.15
0.15
0.15
0.22
0.22
0.22
0.22
0.22
(c) FID
(d)
CFID
Table 11: Comparison of reconstruction metrics for UD models trained on the MNIST dataset for range of unfolded steps K and DDIM steps N. For each observation y, the PSNR computed between the average of 16 samples from the estimated posterior distribution. Since the remaining metrics quantify sample quality, they are computed between single posterior sample and the ground truth.
Results We train instances of the unfolded LATINO archilecture, with K == 1,2,4,, respectively. For each model, we run the unfolded UD²M sampler for N = 1,, 8, 16 iterations for total of NK NFEs. The evaluation metrics for each instance are shown in Table I1 For the PSNR, we compare the ground truth to an average 16 UD samples to approximate the MMSE estimator for the problem. The remaining metrics use a single posterior sample. For each metric, we notice fast. convergence within K iterations. For this ill-conditioned problem, we find small improvement by additionally using K unfolded steps. Qualitative results shown for sample observation in Figure 10 For each K, N pair, we plot. a single UD2M sample, along with the mean and standard deviation estimated from 16 UD samples. The standard deviation plots are amplified by factor of 3x to facilitate visual comparison and align with approximately 99% credible intervals. For each shown standard deviation, we indicate in the upper left corner the correlation to the corresponding absolute error. We observe sharper images and greater sample diversity for K 2 4. Using K = 8, we observe around 10% increase correlation between the posterior variance and the absolute error. each instance, deviations of individual samples from the sample mean illustrate the ability of our method to avoid mode collapse in the learned posterior distribution, providing wide exploration of viable samples. Reference reconstructions are shown for the long-time average of the zero-shot LATINO algorithm over 50,000 NFEs. The convergence of LATINO to an invariant measure is summarized in Figure 11 (left) For our unrolled model, we observe similar qualitative performance to LATINO with K 2 4 iterations with 16KN & 50, 000 NFEs to compute the mean and standard deviation, along with better correlation of the standard deviation to the true error. In all metrics and qualitative comparison, we observe a greater benefit from increasing K opposed N. This highlights the effectiveness of unfolding small number of LATINO iterations to distill the conditional diffusion model to require only small number of iterations.
To benchmark the uncertainty quantification properties of UD*M from accurate posterior coverage, plot empirical coverage probabilities in Figure 11 (right). To focus on important search directions, we compute the empirical coverage on the latent space of the encoder EMNIST- The coverage is computed for K = = 4, 8 unfolded iterations, with N scaled to a fixed number of 16 NFEs. For this calculation, following Tacholaa
& Pereyra 20023): Cheril al. (2024). we improve sample diversity directions corresponding to small eigenvalues of the super-resolution operator by bootstrapping our reconstructed samples through equivariant transformations of the reconstructions
Cherif et al. (2024). We notice that increasing K has positive impact on the empirical coverage probabilities, with accurate results for K 4. This highlights the usefulness of UD2M as tool for accurate uncertainty quantification.
27

---

Published in Transactions on Machine Learning Research (11/2025)
N =
N =4
N
=
=
=
N =
K=2 K=1
0
0
0
0
==1
0
0
0
0
0
G
0
0
K K=2
0
0
0
0
K K=4
0
K=8 K=4
K= 8
0
0
0
0
Samples
Sample Mean
N = 1
N = 2
N 4
N 8
0.73
0.70
0.70
0.72
K ==1
LATINO Baseline (K
50, 000, N 1):
0.70
0.77
0.70
0.68
O Sample
O
K K=2 2
Sample Mean
0.72
0.66
0.71
0.63
K = 4
0.64
0.80
0.84
0.83
0.78
K ==8 8
Sample St. Dev
Sample Standard Deviation
Figure 10: Example reconstructions for the MNIST SR (x4) problem for 21,2,4,8). Top: The ground truth image x and observation y. Middle: Samples (left) and mean estimate (right) from the learned UD²M postcrior distribution with K unfolded LATINO iterations using N DDIM steps. Bottom: The corresponding posterior standard devialion, computed using 16 conditionally independent samples given The correlation of the standard deviation to the absolute error is indicated in the upper left of each standard deviation plot. A reference mean and standard deviation of the zero-shot LATINO kernel after 50,000 iterations are shown in the bottom right.
28

---

Published in Transactions on Machine Learning Research (11/2025)
K 8, N 2
PSNR Sample (max 17.8))
K 4, N 4
PSNR Mean (max 18.3)
K 2, N 8
LPIPS (min = 0.04)
K N 16
0.4
18
Empitical Coverage
0.3
16
PSNR
LPIPS
0.2
0.5
14
0.1
12
100
10
10
103
10
10
0
0.2
0.4
0.6
0.8
Iteration
Percentile
Figure 11: Left: Empirical convergence of the zero-shot LATINO algorithm against number of itcrations. The PSNR of single sample and of the Markov Chain are shown. Additionally, the LPIPS of single posterior sample is shown. Results averaged over 32 independent test images. Right: Empirical coverage probabilities for UD2M for K = 1, 4, 8 and N scaled to fix total of 16 NFEs. The dashed line represents perfect posterior coverage.
29

---

Published in Transactions on Machine Learning Research (11/2025)
Degraded
DifPIR
DDRM
CDDB
Ours (wo RAM) Ours (w/ RAM)
GT
Figure 12: Examples of posterior samples for the task Gaussian deblurring on ImageNet with noise level
= 0.05 and Gaussian kernel bandwidth of 10.
30

---

Published in Transactions on Machine Learning Research (11/2025)
Degraded
DifPIRR
DDRM
CDDB
Ours (wo RAM) Ours (w/ RAM)
GT
Figure 13: Examples of posterior samples for the task SR(x4) with noise level = 0.05 on ImageNet 256
31

---

Published in Transactions on Machine Learning Research (11/2025)
GT
Observation
Gaussian Deblurring
UD M
Observalion
4x resuueerrrsolttoon
UD
Observation
Inpainting (70%)
UD
Figure 14: Examples of posterior samples generated by a "universal" UD sampler on the following tasks on
ImageNet 256: Gaussian Deblurring with a kernel bandwidth of 10, random inpainting (70%), super-resolution
(x4). Noise level 0.01.
32

---

Published in Transactions on Machine Learning Research (11/2025)
Degraded
DifIPIR
DDRM
CM4IR
CoSIGN
Ours (w/ RAM)
GT
Figure 15: Qualitative results from the LSUN bedroom validation set applied to anisotropic Gaussian deblurring with bandwidth (20, 1) and additive Gaussian noise with standard deviation 0.025.
Degraded
DiIPI
DDRM
CM4IR
CoSIGN
Ours (w/ RAM)
GT
Figure 16: Qualitative results from the LSUN bedroom validation set applied to random inpainting with 80%
masked pixels and additive Gaussian noise with standard deviation 0.025.
Degraded
DifIPIR
CM4IR
CoSIGN
Ours (w/ RAM)
GT
Figure 17: Qualitative results from the LSUN bedroom validation set applied to anisotropic Gaussian deblurring with bandwidth (20, 1) (top) and random inpainting with 80% masked pixels (bottom). Each with additive Gaussian noise with standard deviation = 0.05, which is out the training distribution for our model which was trained for 0.025.
33

---

Published in Transactions on Machine Learning Research (11/2025)
Figure 18: Examples of posterior samples for the task box inpainting with noise level a 0.05 on LSUN Bedroom. From top to bottom: Ground truth, degraded observation, UD²M reconstruction (without RAM).
34


Reviews:
[{"review_id": "1L2eeg3P7K", "reviewer": "Reviewer_G1t8", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes UD2M, a few-step posterior sampler for image restoration that \n1. unrolls a small number of LATINO Langevin iterations into a trainable block,\n2. uses an explicit proximal data-consistency step to encode the likelihood, and \n3. plugs in a distilled diffusion prior with lightweight LoRA adapters. \n\nChaining a handful of such blocks across decreasing diffusion noise levels yields fast, good-quality reconstructions. Experiments on deblurring, inpainting, ×4 super-resolution, and JPEG (ImageNet/LSUN) show competitive or improved FID/LPIPS/PSNR at roughly order-10 NFEs, with ablations on LoRA rank and step budgets.\n\nStrengths:\n1. Clear separation of physics (prox) and prior (distilled DM); compelling few-step efficiency.\n2. Broad empirical study across tasks/datasets; strong results for the reported NFE budgets.\n3. Practical design (LoRA, closed-form prox for linear-Gaussian cases) that keeps training and inference lightweight.\n\nWeakness:\n1. Requires task-specific adaptation; prox for non-linear operators relies on approximate solves.\n2. The paper empirically studies the effect of unrolling, showed good results, but provides no non-asymptotic guarantee for the learned $K$-step unrolled map vs. full LATINO iterations.\n3. The trade-off between inner $K$ and outer $N$ under a fixed budget is not isolated.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Empirical evidence is strong. Evidences provided across deblurring, inpainting, ×4 super-resolution, and JPEG on ImageNet/LSUN the method reports competitive/improved FID/LPIPS/PSNR with ablations (LoRA rank, step budgets) and noise-level generalization.\n\nHowever, the theoretical support is limited, with no non-asymptotic bounds for the unrolled sampler vs. full LATINO, and only heuristic treatment of inexact proximal steps.", "audience_yn": "Yes", "audience_explanation": "TMLR readers working on diffusion models, computational imaging, and deep unrolling would find this relevant because it combines an explicit likelihood prox with a LoRA-adapted few-step diffusion prior for fast posterior sampling.", "requested_changes": "1. Provide a error bound between the learned $K$-step unrolled map vs. full LATINO iterations under mild assumptions. If relevant theorems exist, include or discuss them. Or at least analyze and discuss the finite-step error and its effect on the task.\n2. Provide a $(K,N)$ grid under fixed NFE showing the best allocation.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "2TsZbirZ0I", "reviewer": "Reviewer_xFEp", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents UD$^2$M, a framework for enabling pre-trained Diffusion Models (DMs) to perform posterior sampling, in particular in the case of inverse problems in computational imaging. The paper builds upon LATINA, which distills a DM to a Consistency Model (CM), and uses it as a prior in combination with a proximal algorithm with respect to the likelihood of the conditioning variable $y$. The authors add many techniques, such as Algorithm Unfolding of the LATINA Langevin sampling steps, LoRa to make these steps fine-tunable, and adversarial training with a discriminator.\n\nStrengths:\n1. The paper is generally well written, and most parts are clear\n2. There are interesting relations with other parallel works in the literature\n3. There is an extensive evaluation and ablation studies\n\nWeaknesses:\n1. There are quite a lot of modules and techniques, which I feel don't get presented in a very clear order. For example, what is $L_\\vartheta$ concretely? What are its inputs? Why are parameters $\\phi$ introduced before the discriminator $D_\\phi$ in section 3.2?\n2. Although text passages are fairly easy to follow, the math notation is quite hard to understand at times. I would advise a notation (and definitions) section.\n3. I feel like the line between the related works and the proposed method is blurry in the writing. For someone not familiar with each specific subject, it is difficult to follow what is proposed and what already exists. I advise a paragraph (preferably with a list) explaining the contributions of the paper, and what distinguishes it from the related works.\n\nQuestions:\n1. I found the sentence \"Importantly, unlike zero-shot methods, here we consider that the likelihood function $p(y|x_0)$ associated to $p(x_0|y, x_t)$ belongs a class of likelihood functions of interest that is fixed during the training of $L_\\vartheta$, while retaining the flexibility to specify the exact parameters of the likelihood during inference time\" to be a bit confusing. Isn't it fixed also for other methods?\n2. To clarify, is the idea of unfolding LATINO to have a fixed number of specialized layers, fine-tuned with LoRa, instead of having equal \"recurrent\" steps?", "claims_and_evidence_yn": "Yes", "claims_explanation": "The authors demonstrated empirically that UD$^2$M has high performance with flexibility and a low number of NFE, making it efficient.", "audience_yn": "Yes", "audience_explanation": "The paper advances the field of Diffusion Models, which is interesting and a hot topic in the current state of the Machine Learning field.", "requested_changes": "Critical changes are those explained in Weaknesses 1-3. In general, I advise making the paper more accessible to readers not familiar with every single technique or their usefulness.", "broader_impact_concerns": "The Broader Impact Statement provided in the paper addresses fully the possible concerns.", "additional_comments": "Typos:\n- Spangnoletti -> Spagnoletti\n- Introduction:\n\t- Within this context, we are particularly \"interest\" -> interested\n\nThere might be other typos after that point, so I advise a full paper check.", "extra_scores": {}, "extra_text": {}}, {"review_id": "Fa7Ln5GnAj", "reviewer": "Reviewer_qiWz", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Summary Of Contributions:\n\nThe paper has as an overarching goal the task of learning few-step samplers that draw from the posterior for imaging problems that are severely ill-conditioned or ill-posed.\nThe paper proposed the Unfolded and Distilled Diffusion Model method, which unfolds K iterations of the LATINO Langevin posterior sampler inside a conditional diffusion pipeline and then distils to a few-step consistency model. The proposed approach uses likelihoods via proximal operators and supports joint training over a family of likelihoods. Experiments on ImageNet‑256 and LSUN Bedroom across Gaussian deblurring, inpainting, superresolution, and restoration of JPEG compression artifacts report strong FID/LPIPS.\n\n\nStrengths And Weaknesses:\n\nStrengths:\n\n- Clear and well-written\n- The Unfolded and Distilled Diffusion Model method is well presented. As far as my knowledge, the authors claim of this being the first deep unfolding of an MCMC sampler within diffusion‑model imaging is correct. \n- Exact likelihood use via a proximal step inside Langevin avoids common likelihood approximations in DM‑based posterior sampling.\n- Demonstrate decent empirical evidence across ImageNet/LSUN tasks for multiple metrics, ablations useful for key hyperparameters (choice of rank for LoRA, number of sampling steps, etc). Sufficient details given for reproducibility. \n\n\nWeaknesses:\n\n- My main issue is no results explicitly assess posterior validity. The paper does a good job across various metrics typically used to assess image quality and diversity for image tasks; however, if want to present this as posterior sampling, I would like to see something more concrete. For instance, to verify details in the training objective don't lead to an overconfident posterior.\n- No comparison to LATINO\n- Proximal step solving. Reliance on sub‑iterative solvers for some tasks (e.g., non‑linear cases); Appendix A is helpful and many problems may admit fast proximal solutions, but when sub‑iterations are required inference slows.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Yes. Overall, the main contributions are novel and the method's performance is reasonably supported by decent experimentation.", "audience_yn": "Yes", "audience_explanation": "Yes. In particular, this paper should appeal to readers interested in Bayesian computational imaging.", "requested_changes": "Requested Changes:\n- My main requested change is regarding posterior validity. Either include some additional posterior validity checks (some options might be: simulation-based calibration, assess empirical coverage of credible sets, posterior predictive log score, run on toy example and compare with golden standard method) or some softening of claims surrounding posterior sampling.\n\nOther optional changes:\n- LATINO, a direct baseline would clarify how much the proposed approach improves on LATINO.\n\nSpelling and grammar:\n- “we are particularly interest in imaging problems” ->  interested\n- “for solving for inverse problems” -> for solving inverse problems\n- “out method is not as efficient” -> our\n- “LSUN bederoom dataset” -> bedroom\n- “Ground Thruth” -> truth", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)