Learning Latent Variable Models via Jarzynski-adjusted Langevin Algorithm
Authors: James Cuin
Davide Carbone
O. Deniz Akyildiz
Venue: NeurIPS 2025 poster ()
Keywords: Latent variable models
Parameter estimation
Sequential Monte Carlo
Langevin algorithms
Primary area: probabilistic_methods
TL;DR: We propose a novel estimation method for latent variable models.

Abstract: We utilise a sampler originating from nonequilibrium statistical mechanics, termed here Jarzynski-adjusted Langevin algorithm (JALA), to build statistical estimation methods in latent variable models. We achieve this by leveraging Jarzynski’s equality and developing algorithms based on a weighted version of the unadjusted Langevin algorithm (ULA) with recursively updated weights. Adapting this for latent variable models, we develop a sequential Monte Carlo (SMC) method that provides the maximum marginal likelihood estimate of the parameters, termed JALA-EM. Under suitable regularity assumptions on the marginal likelihood, we provide a nonasymptotic analysis of the JALA-EM scheme implemented with stochastic gradient descent and show that it provably converges to the maximum marginal likelihood estimate. We demonstrate the performance of JALA-EM on a variety of latent variable models and show that it performs comparably to existing methods in terms of accuracy and computational efficiency. Importantly, the ability to recursively estimate marginal likelihoods—an uncommon feature among scalable methods—makes our approach particularly suited for model selection, which we validate through dedicated experiments.

Learning Latent Variable Models via Jarzynski-adjusted Langevin Algorithm
James Cuin Department of Mathematics Imperial College London
Davide Carbone
Laboratoire de Physique de l'Ecole Normale Supérieure,
Université PSL, CNRS,
London, UK
Sorbonne Université, Université de Paris
jamie. cuinemimpeenmmmnall...... acc..alaaa...................... ac. uk
Paris, France
Rawidd.carbonee@gee. ens .
O. Deniz Akyildiz Department of Mathematics Imperial College London
London, UK
deniz denizzakkyildibeeimmeriaaa.. ac.
Abstract
We utilise a sampler originating from nonequilibrium statistical mechanics, termed here Jarzynski-adjusted Langevin algorithm (JALA), to build statistical estimation methods in latent variable models. We achieve this by leveraging Jarzynski's equality and developing algorithms based on weighted version of the unadjusted Langevin algorithm (ULA) with recursively updated weights. Adapting this for latent variable models, we develop sequential Monte Carlo (SMC) method that provides the maximum marginal likelihood estimate of the parameters, termed JALA-EM. Under suitable regularity assumptions on the marginal likelihood, we provide nonasymptotic analysis of the JALA-EM scheme implemented with stochastic gradient descent and show that provably converges to the maximum marginal likelihood estimate. We demonstrate the performance of JALA-EM on variety of latent variable models and show that performs comparably to existing methods in terms of accuracy and computational efficiency. Importantly, the ability to recursively estimate marginal likelihoods- an uncommon feature among scalable methods -aaies approach partieularly suited for model selection, which we validate through dedicated experiments.
Introduction
Real world data often contains latent structures (Whiteley al., 2025) that can be described well with latent variable model (LVM). As result, LVMs have become ubiquitous in modern statistical research, from natural language processing (Wang et al., 2023) to bioinformatics (Shen et al., 2009), due to their flexibility in capturing complex and hidden processes underlying observed data. It is therefore of significant interest to fir LVMs, i.e., learn its parameters from data.
A particularly prominent estimation paradigm for LVMs is that of maximum marginal likelihood estimation (MMLE) (Dempster et al., 1977). Given some fixed observed data y. the task of MMLE is to compute the maximum likelihood estimate of the parameters, denoted 0, by maximising the marginal likelihood Po (v). This intractable quantity is defined through marginalising the joint likelihood of the observed data y and the latent variable æ, the latent variable æ. More precisely, let Rd be the observed data, z E &d be the latent variables, and 0 E Rd be the parameters of interest. Given
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

some fixed observed data y, we define the joint likelihood function (z y) R x Rd - R. Our main task in this paper to develop methods for finding the maximisers of the marginal likelihood po(v)= Rdd -> IR. i.e., we aim to solve
Or E argmax log Po (v), where po(u)= - -fe.. po(a, y)dx. DER20
(I)
Jd=
Indeed, the integral in (1) is often analytically intractable or numerically expensive, making direct maximisation of the marginal likelihood infeasible in many relevant applications.
Historically, the gold standard for solving the problem in (1) is the Expectation- Expectation-Maximiiiee Maximisation (EM) algorithm (Dempster et al., 1977), which is two-step iterative procedure that alternates between estimating an expectation Q( (00,0x=1) = _ (srly) Plog y)] (E-step) and updating the parameter via 0% arg max Q 8)-1) (M-step). This algorithm requires the implementation of two general procedures, integration for the E-step and maximisation for the M-step. which are generally intractable for complex statistical models. This has resulted in many numerical strategies for implementing these steps efficiently, which gave rise to a large number of variants and approximations of the FM algorithm. Farly examples include simulation based approaches for the E-step for models where it is possible to sample from the posterior pa (zly) for given @ (Wei and Tanner, 1990), termed Monte Carlo EM. Similarly, the M-step is often approximated using numerical optimisation techniques as exact maximisation is intractable (Liu and Rubin, 1994; Meng and Rubin, 1993; Lange, 1995). In general, however, these schemes are still impossible to implement as the posterior is rarely amenable to exact sampling, which resulted in the use of Markov chain Monte Carlo (MCMC) algorithms for the E-step. There has been sigmificant body of work in this direction, see, e.g.. Atchadé et al. (2017); Caffo et al. (2005); Delyon et al. (1999); Fort and Moulines (2003); De Bortoli et al. (2021); 1); Gruffaz et al. (2024). Most notably, De Bortoli et al. (2021) used unadjusted Langevin chains for the E-step and gradient descent schemes for the M-step, which is most related to our work. This approach may result in long running times for the Markov chain approximating the Etstep, which also complicates the analysis. To overcome such limitations, significant body of work developed an interacting particle systems approach to solve the MMLE problem (Kuntz et al., 2023; Caprio et al., 2025; Sharrock et al., 2024; Akyildiz et al., 2025) including extensions for nondifferentiabl models (Encinar et al., 2025) and accelerated schemes (Lim et al., 2024; Oliva and Akyildiz, 2024) with applications generative modelling (Wang et al., 2025; Oliva et al., 2025; Marks al., 2025) and inverse problems (Glyn-Davies et al., 2025). These works use particle system instead of a MCMC chain for the E-step, which is computationally efficient. However, most of these methods run the particles in the space of latent variables independently (except Sharrock et al. (2024))- which could also be improved with further interaction. Moreover, these methods do not allow an easy computation of the model likelihood po (y) for a given 0.
There are alternative, closer to our approach, sequential Monte Carlo (SMC)-based methods. Johansen et al. (2008) considers an SMC algorithm on an extended target measure that concentrates on the MMLE solution. Crucinio (2025) provides a general SMC method that bears similarities to ours, but is designed for a specific parameter update mechanism (rather than general optimisers). In the context of particle filtering, Hamiltonian Monte Carlo approaches have been explored, as in Septier and Peters (2015).
To address the issues mentioned above, in this paper, we build on numerical technique to sample probability paths, which we term Jarzynski adjusted Langevin algorithm (JALAA). JALA is a Langevin Monte Carlo (LMC) method to sample from time-varying probability measures, corrected using a SMC rather than Metropolis steps. The key idea is to run biased dynamics, specifically that of the unadjusted Langevin algorithm (ULA) with no Metropolis correction, and subsequently correct for the bias in sampling via an exponentially weighted factor referred to as a Jarzynski factor. Building this idea:
Contributions.
In Section 2, we formulate the Jarzynski adjusted Langevin algorithm (JALA) for sampling from time-varying sequence of distributions. The algorithm is based on the sampler de- veloped in Carbone et al. (2023) and is closely related to sequential Monte Carlo (SMC) samplers (Del Moral et al., 2006) - and can be seen weighted-ensemble of the ULA.
In Section 3, using the JALA as the core component, we propose numerical method for EM, using JALA, which we term Jarzynski adjusted Langevin algorithm for EM (IALA-EM). This method uses the JALA for estimating the gradient of the marginal likelihood, which

---

is then used to update the parameters via gradient- optimiser. The resulting JALA- EM algorithm a sequential Monte Carlo method that provides the maximum marginal likelihood estimate of the parameters.
In Section 4, we provide a convergence analysis, under log-concavity and Polyak- Łojasiewicz (Pt.) conditions. In particular, we provide a nonasymptotic analysis of the JALA-EM method which is implemented via stochastic gradient descent (SGD). We utilise the convergence analysis of SGD algorithms to prove nonasymptotic result. This is just first step, as any other gradient-basee optimiser can be used in the JALA-EM algorithm and their theoretical properties can be used to prove convergence of the JALA-EM algorithm.
Finally, in Section 5, we demonstrate the performance of JALA-EM on a variety of L.Ms and provide empirical evidence that it successfully estimates model parameters for various regression models and that it can also be used in model selection unlike the methods above.
Computational Cost and Code. Experiments were run on a personal computer and a Google Colab T4 GPU. The code can be found in https::/ github....... com/ jateecuuinnaaaaaaaeee
2 Jarzynski-adjusted Langevin algorithm
We first introduce here the Jarzynski adjusted Langevin algorithm (JALA) which is a sampling method for time- vare-ng probability distributions. The main idea behind the method is to run ULA these time-varying potentials and correct the resulting bias (due to time-dependence) using the Jarzynski equality. The resulting JALA is sampler for a sequence of evolving measures.
Consider a sequence of target distributions (TT))20 on Rd and assume that we wish to use a ULA-based strategy to sample from these distributions. As opposed to the case with static target measure, the problem is significantly harder, as a naive application of the ULA, i.c., iterating Xa-1 = X2 VVV 2 + 22h where Uz(2z) = log TT. (1)), will not yield samples from the target distribution TT at time k. In fact, even introducing Metropolis step would not correct the bias (even asymptotically). as the distributions (T*))20 are evolving. One possible solution is the idea of employing Jarzynski-based correction for sampling from time-evolving measures which was studied by Carbone et al. (2023). As opposed to the setting therein, we first introduce the method for time-varying potentials (Ux) k20, and then will apply it to the
of LVMs in Section 3.
Proposition 1 (Adapted from Proposition of Carbone et al. (2023)). Assume that Zy =
e duu da < oc for all k 2 O. Let the sequences (Xx)k20 and (Az) be given by the eration rule
Xk+1 = Xx (X,))(X)) + 22h EK+1,
Xg ~ To:
A441= = Av ax+i (X Xx) + an (XR, Xx+1), where (Ex BeEt are i.i.d ((2,,, Ia), and we define
Ag = O,
ax( x.) Un(2t) + (x, x1) [il(z1) + (i)((a))
(3)
Then, for all le E No and a test function
yp RRd. + RR we have
[p(( (E]
Em [+(X)]
Z = ZulE (EAL)
E [e-A]]
where the expectations the right-hand side, of both equations, are over the law of the joint process (Xg,As), as defined through (2).
See Appendix B.1 for a proof. Eq. (4) provides powerful tool for estimating expectations with respect to the time-evolving measures Tg. In particular, it allows us to compute the expectation of any test function & with respect to the target distribution and using the Jarzynski factor A to correct for the bias introduced by the unadjusted Langevin dynamics. In practice, the expectation in (4) cannot be computed, hence a particle algorithm can be used, which makes the method an instance of an SMC method. This result then can used for estimating parameters of an LVM or fitting energy-based models (EBMs) as observed in Carbone et al. (2023), as time-varying potentials (U=):20 2:20 can indexed by some parameter 0g and the sampler can be used to compute the gradient of the log- marginal likelihood.
by running the joint process (XR, A&)
3

---

Algorithm JALA-EM
I: Inputs: Observed data y E IR4; potential energy U: number of particles N N; number of iterations K E N; step-size h > 0; gradient-based optimiser OPT; ESS threshold C >> 0; and set of walkers {Xx] ke(n) sampled from pe, (xly).
2: Initialise Aj = 0 for all i [A]]
3: for k = 0. 0...., K- do
4: 1 = / 2j=1 Σj= exi,
for foraliie [N].
I/ compute weights II estimate the gradient I| update the parameter
5: 9k Ei=1 w VOU (0x.
6: 0k+1 = (P 9k)
7: for i 1,..., N do
Xi X - hV (0p, xi) + v2h
A A au+1 (X X + as (x Xx+1),
where EL- ~ -N Id,) and ar is defined as in (2) with U2 = U(k.
8: end for
9: Resample w.r.t wh+ and set A/41 = 0 if ESS -1 < C.
end for Outputs: 0x. (w M (x
11:
The definition of Aş update in relation to SMC can be interpreted particular case of Eq. (12) in Del Moral et al. (2006), or later on Sec. 2.4 in Heng et al. (2020), where the backward transition kernel L, L,,, x,-1 is chosen to be M, (X,, X,-1), given that M, X,) is the forward transition kernel (i.e. ULA kernel in the case of (2)). Exemplary illustrations of this situation in the context of MCMC design can be found in Nilmeier al. (201 and in Fig. 2 of Schönle et al. (2025). A theoretical molivation for this particular choice is that in the limit h > 0 one recovers A&+1 Ax + (XU,(Xx) + Oh3/2), meaning that the continuous time limit for the evolution of the log-weights coherent with Jarzynski work up to O(h1/2), i.e. the discretization error of Euler-Maruyama ULA, cfr. Carbone et al. (2024). Moreover, for sufficiently small h the ULA kernel becomes reversible and "I-step" ergodic (cfr. Mattingly al. (2002)), ensuring that for a finite sample size the estimator (4) is not affected by variance and bias issues due to poor exploration. In practice, different strategies for tuning h are available, see for instance Kim al. (2025).
Remark 1.
3 Jarzynski--djusted Langevin algorithm for MMLE
We now outline Jarzynski adjusted Langevin algorithm for EM (JALA- (ALA-EMM, in which the key idea is to approximate the nonequilibrium corrections required by Jarzynski's identity through a population of particles that evolve under ULA, whilst simultaneously performing parameter updates.
Let y be the observed data and po (z. y) be joint distribution of an LVM. For the ease of notation, let U(f log (* where U will be referred as joint potential. Our aim is to compute the minimisers of (0) appealing scheme would be to implement gradient based optimiser, using V, log pv)).. However, as we have discussed, this gradient cannot be computed in closed form either. Nevertheless, for sufficiently regular models (see, e.g., Douc et al. (2014, Appendix D)), one can write this gradient as
log Pa()) wr.t. as is fixed. Given that this cannot be done analytically, an
VoV(0)= - V (o)
2)) Eno(a y) (VU x))
(5)
where 200.z) Po y) and (y( is the posterior distribution of latent variables. This identity (which is known as the Fisher's identity) is behind many recent algorithms to implement
MMLE procedures, see, e.g., De Bortoli et al. (2021) for an implementation of EM based on this identity (using ULA chains to sample from pa (=
Instead of an MCMC-based method, we aim using the JALA to approximate the expectation in (5). This can indeed done straightforwardly using Proposition 1, by replacing Ue(xr) with U U(2,, .) and letting (0x (220 evolve with some discrete-time protocol (c.g.. SGD). As noted in the last section, this

---

computation can only be performed using particle scheme and then setting p(x)) := (oU x)). This is the main idea behind JALA-EM. To develop our method, we instantiate the method using N particles, resulting in the system:
X X - NV, (0s, X%) + v 2h
(6)
1 A ax+1 (X++1; Xi.) + aa (X,,
(7)
where ag is defined as in (3) with U2 == U(Ok, -). At iteration k. the gradient can be estimated by computing the normalised weights w A/ مj=1 Aii, and then computing the weighted average
N Σ
9k
w VoU (i).
(8)
This can then be used for generic first-order optimiser. We outline the method in Algorithm 1. A few remarks are in order for our method.
We note that the method (6)-7) is SMC method, where Ag the unnormalised log-weights. It is thus vulnerable to standard problems of SMC samplers, such as weight degeneracy. For this to not happen, the sequence (0k)w20 needs to 'vary slowly' which can be achieved with small step-sizes or adaptive approaches, see, e.g., Zhou et al. (2016). To keep track of the degeneracy, we use the effective sample size (ESS) as ESS; takes values in between 1 and N. Initially, since A = 0, we have ESS, = N, but it decreases with k. as generally observed in SMC. When this happens, a resampling operation triggered (see Appendix
Remark 2.
( (Elvira et al., 2022). This quantity
C.I for details).
4 Nonasymptotic bounds for JALA-EM
In this section, we establish nonasymptotic error rates for our algorithm, where the optimiser OPT is assumed to be first-order method, specifically, stochastic gradient descent (SGD):
OPT. (Ok; 9k) 02 Ykflk:
where (7w)220 is sequence of step-sizes and (gx) are the estimates of VV == V Po() for (0k)
Remark 3. Our framework would allow for deriving nonasymptotic bounds using any optimiser (notably adaptive optimisers), soon the analysis for biased gradients (e.g. Surendran al. (2024)) is available. This is in contrast to similar methods like Kuntz et al. (2023); Akyildiz et al. (2025) where theory strictly requires gradient descent type update for the parameters, albeit in practice adaptive optimisers are used in experiments. See Appendix A.1 for more thorough discussion of this point.
4.1 Bounding the MSE of the stochastic gradient
In order to establish nonasymptotic rates for JALA-EM, the first challenge is to notice that our stochastic gradient estimate given by (8) is biased, as opposed to the usual SGD setting. Therefore, we need to first provide an analysis our estimator w.r.t. the true gradient, which requires the following assumptions.
A1. We assume that (i) the sequence (Ak) has finite exponential moments, i.e. supk20 2le2Ae) <
oo, (ii) there exists measurable set A C Rd, independent of k. such that: Leb(A) > 0 (positive Lebesgue measure) and sup, sup U(0,x)) x) S M for some < co, (iii) the gradient satisfies || sup(0.z)eRdo xRd= V (0, (0xz)l...
Three remarks are in order regarding A1. First, the finite exponential moments are about the weights of the SMC method, specifically, we require weights assumption. Secondly, A (ii) immediately implies Z₁ = Ze. 2 SA U(( dx SA dr L A) > O, and just correspond to ask po, (xly) to not concentrate mass zero measure set. Thirdly, we require the gradient to be bounded, as it acts as test function in our importance sampling (IS)-type estimator. This assumption can be dropped at the expense significantly complicating the analysis, see, c.g., Agapiou al. (2017). We thus keep this assumption for simplicity.
bounded, which is usually standard
5

---

Assume that we run the system (2) with U+(+) = U( -). At iteration k, for any 0% Rde, we have that
E [V Xx)e (x),
(O2( Epo, (xlp) [VUU(Ok,
IE [eeAA]
where in the last equality we have used Proposition with sp(x)) VoU (0g, x) and so the expectations on the right-hand side are over the law of the joint process, as defined in (2). In practice, however, we run a particle system given (6)-(7) and estimate this gradient as (8). Our estimator an IS-type estimator and
prove the following proposition to bound the mean squared error (MSE).
Proposition 2. Under A A for any sample size, N 2 1, we have that
E [OL) snl2]] 4C IV for k 2 O, where C= sup E
(10)
N
(EAL)) 21 [Eeee) sup& Cj <0.
self-contained proof can be found in Appendix B.2.
Remark 4. We note that Proposition 2 is for our scheme without resampling. The full scheme with resampling requires nontrivial analysis which we leave for future work. In light of this, one has to choose N = ((1fe) to obtain that MSE S With resampling, the picture could be different and the number of particles can be adapted to keep the MSE stable, see, egg. Elvira et al. (2016).
4.2 Nonasymptotic convergence
Now, in light of Proposition 2, Algorithm optimisation problem.
can be understood to be solving the unconstrained
0. := arg (())((() arg max log pe (v BER4
(11)
PER
where we have access to biased and noisy gradient estimators. Below, analyse this scheme in two distinct assumptions on the negative marginal log-likelihood V, first under the PL condition and then strong convexity condition, utilising the results from (Demidovich et al., 2023).
For our first result, we impose the following smoothness condition, together with the PL condition on the marginal likelihood. For brevity we use the notation V(0-) V..
A2. Assume the function V is differentiable and L-smooth, that is, for all (0,0') E Rd x Rd.
(0) -VV(0') s LIl0 1111
(12)
and also that V is bounded from below by V. E R. A3. (Polyak-tajasieviie) There exists constant, μ >> O, such that for all 0 e Rdd
(VV(2)) 2 (0) (V)()))))))))) V.).
(13)
Theorem 1. Assume that OPT(Ou, gk), Algorithm 1, refers to the SGD scheme outlined in (9). Under Al-A3, provided that fixed step-size is chosen such that 0 < 2/- S min (1/4L, 2/p} we have, for every ls E
E (Vx) v.] (1 (1-
TH
8LY V
V NN
80
2
nN
where 6, = ((0o) - V, and C = sup E (( /(B [(B
(eAu1) 1) <
xc.
proof can be found in Appendix B.3. Note that this provides a result on the convergence of marginal likelihood since V@))= = - log Py()).
Remark 5. To expand this result and write in a more intuitive form, we see that Theorem ! essentially implies that
E [V(0k)
- v.] < 60 (1-2)/2)² + +1/2))
This turn provides us a guideline on how to scale parameters, in particular, the number of iterations k and the number of samples N. To obtain E [V (0x) V.] < O(e) with 0 < & 1, must

---

choose = = and k Ol Notably, due to the apperance of N in the second term, we do not need take & - for our method to converge, and rather should take N large. This is in contrast to theoretical results obtained in Akyildiz et al. (2024, 2025) and Caprio et al. (2025) for alternative diffusion-based MMLLE methods, where for convergence, one needs to pick N large and 7 very small. See Appendix A.2 for more discussion.
We next impose stronger assumption on the negative log marginal likelihood to strengthen this result.
A4. The following holds: (VV - VV 0"; 2 (n/2)]10 - 0'12
Remark 6. Note that there are practical examples where A4 almost satisfied. For example, if U(0, x) is strongly convex in (), then A4 can be shown to be satisfied via Prekopa-Leinde inequality (see, e.g., Saumard and Wellner (2014, Theorem 3.8)). For example, can be shown that U(0, x) is strictly convex in (0,x)for the Bayesian logistic regression example, see Kuntz et al. (2023, Prop. 1).
Theorem 2. Under Al, A2, and A4, fix step-size such that 0 < min (1/4L, 2/ then have
E [110k - 1-] S (1 7/4 &
V
8IV 22N
80 80(2/14)
2
N2N
where 6n = V(0o) V, and = sup, E [(d4)) /(B [edu]) xc.
The proof of this theorem follows from the fact that the strong convexity (A4) implies ||0 0-1 s (2)))(V(()- V.), hence the bound under PL condition in Theorem can be applied directly.
Experimental results
5.1 Bayesian logistic regression Wisconsin cancer data
To benchmark JALA-EM's performance, relative to that of the particle gradient descent (PGD) and stochastic optimization via unadjusted Langevin algorithm (SOUL. algorithms, introduced in Kuntz et al. (2023) and De Bortoli et al. (2021) respectively, we consider the extensively studied Bayesian logistic regression task using the Wisconsin Breast Cancer dataset, as described in De Bortoli et al. (2021). Specifically, we employ a model with a Bernoulli likelihood for the binary classification of dy 683 malignant benign samples, and the regression weights w E dd are assigned an isotropic Gaussian prior p(w O o -N(w I (6) As is the case in Kuntz et al. (2023), fix the prior variance to og = 5.0, and estimate the unique maximiser of the marginal likelihood.
To ensure robust comparison, algorithm-specife step-sizes were tuned via 3-fold cross- cross- validation set, D Dtrain,vals where the Log Pointwise Predictive Density
on the respective (LPPD) was utilised the evaluation metrie, since this signifies predictive performance of the fitted
1.0
0.45
0.5
KDB
0.30
0.0
0.1
2
-2
-1
2
Time
PGD
SOUL
JALA EM
200 lter.
Figure 1: Bayesian logistic regression: Parameter estimates for PGD, SOUL., and JALA-EM, for N 100 and tuned step-sizes, with a iteration milestone 200 indicated (left). KDE of the second coordinate of the posterior approximation for the model weights (right).
7

---

model. Subsequently, final algorithm runs were constrained to 5.0 second wall-clock duration to assess practical computational efficiency. Comprehensive experimental details can be found in Appendix C.2.
The parameter estimates and posterior approximations, illustrated in Figures 1, 5 and for particle counts N E {100, 50, 10). highlight some of JALA-ENS distinct characteristics. Across the particle counts, JALA-EM demonstrates rapid convergence for the parameter 0, consistently matching or even outperforming PGD and SOUL in terms of wall-clock time, while reaching the same estimated value. Although across algorithms all posterior approximations, for the representative regression weight W2, are comparable at N = 100, the posterior fidelity of PGD and SOUL noticeably degrades as N decreases. In contrast, JALA-ER maintains high-quality approximations, even with fewer particles, which we attribute to its effective particle management due to reweighting. Despite JALA-EN requiring an additional step-size to be tuned compared to PGD, it provides a compelling combination of fast parameter estimation and accurate posterior approximation in this setting.
5.2 Error model selection Bayesian regression
To demonstrate JALA-EM'S unique and dual capability concurrently estimate both parameters and the marginal likelihood, we consider a Bayesian model selection problem. Specifically, we aim to differentiate between two competing, nested Bayesian linear regression models, in which one the true generating process G of the synthetic dataset D. with model preference quantified via Bayes Factors derived from these marginal likelihood estimates. Specifically, these models aim to capture the relationship V/ = X,w + E; through latent weight vector E Rd=, where E; represents the observation error. The first model, .MG. assumes i.i.d. Gaussian observation errors, E; ~ NO d²). while the second, .M,, assumes Student-t distributed errors, ~ Student-t(O,, o² v), providing robustness against outliers. For both models, an isotropic Gaussian prior is placed on the regression weights w. p(w] a) NO0 a la,). Notably, for numerical stability and to ensure positivity, we use JALA-EM to estimate the log-transformed model parameters, Φi = logo² Φ1 = log a, and for Mr, Ø3
log
Indeed, to estimate the log marginal likelihood log 2Mm for model M, we have, by Proposition 1, that after K steps
log 2,K = log ZM, + log
Σe
Eeki
log N 1og
(14)
where Zm.o is the marginal likelihood evaluated the initial estimates, M, Furthermore, we adopt the following version of (14) which includes resampling as usual in SMC contexts, cfr. Del Moral al. (2006):
log mM log MM +
log
exp
Ain)
Ko =1
(15)
N
m=
where ky is the step at which the j-th resampling event occurs, with j = 1,..., J, whilst A is the log-incremental weight of particle at step m, and Ih 2m= A,, represents the cumulative log-weight of particle between two resampling steps.
Notably, the computation of the marginal likelihood is model dependent, as for Mg this is analytically tractable, whereas for Mr generally intractable, and thus we estimate it using Importance Sampling, with proposal distribution derived from Gaussian approximation to the posterior of w using initial parameter guesses. Comprehensive experimental details and corresponding derivations can be found in Appendix C.4.
with features drawn from (0, la_). True parameters are set to a = 1.0, o² = 1.0, and μ* = 4 (for G = Mr), with w, drawn from (0,0, a,- II,))))))) ). This core experimental trial was repeated 100 times, with JALA- EM configured with N = 50 and K = 250. The results, exemplified by Figure 7, highlight rapid convergence to sensible parameter estimates and estimated marginal likelihoods clearly differentiating between the true and misspecified models. Notably, when G -Ma, log 2M quickly stabilises close to its true analytical value, and precisely matches the analytical marginal likelihood derived
Our numerical experiments focus on generating with d, 500 and dę
8

---

G Me
G- Mp
2125
-210
-2000
211
-225-
2200
250
2225
2070
2011
2275
-2700-
2000
-2725
Analytic log MM--
-2730
.00
150
200
250
150
200
200
Iteration.
Iteration.
JALA-EM (Ma)
JALA-EM (M2)
Analytic log, ZMoA
Figure 2: Error model selection: Marginal likelihood estimates for JALA-EM fitting Me and Mr. where the true underlying model is the former (left) and the latter (right). Here, dy = 1500, dg = 8, and N 50. Parameter estimates initialised at values perturbed away from the true values.
from the iterated parameter estimates, corroborating Proposition 1. Additionally, when G = Mr. JALA-EM fitting Mr is more effective recovering the true parameter values, and importantly its log M,,,,, stabilises at higher value. In fact, in Figure 2, where dy 1500, JALA- LL'E''s ability to effectively leverage larger datasets for more confident model selection is highlighted, consistent with Bayesian principles. Lastly, we note that, across the 100 repetitions, model selection based on the higher estimated marginal likelihood correctly identified the true underlying model class in 100% of trials when .M.. and 99% of trials when MMr, indicating JALA- AMM's robustness and utility for model selection.
5.3 Model order selection Bayesian regression
related, but distinet problem is that of model order selection in the context of competing, nested Bayesian polynomial regression models. In this case, the candidate models share the same probabilis- tie structure for both their weights and observation errors, and instead differ in the model specific features utilised, which are derived from the polynomial basis expansion up to and including order denoted by up (xx())) for the p-th model My. To be clear, these models aim capture the relationship y: (I,)) Wp sú, through a latent vector wp E RR+1 where e; (000 o²).
Following the comprehensive setup described in Appendix C.5, we compare the performance of JALA-EM to that of competitive baseline, which notably leverages Ordinary Least Squares (OLS) for weight estimation followed by Bayesian Information Criterion (BIC) for model selection, over a range of true model orders P. [2, 8). where repeat the core experimental trial 100 times for each order. As seen Table 1, we observe JALA-EM to match or outperform the aforementioned baseline across the values of p. and dy considered, where outperformance is notably more pronounced for higher model orders. We attribute this to JALA-EM leveraging an approximation of an integral over the latent parameters, which yields more robust trade-off between model complexity and fit than the baseline's is reliance on a point estimate and BIC approximation, particularly in the non-asymptotic settings considered.
5.4 Bayesian neural network
As motivated Kuntz et al. (2023), Bayesian neural network (BNN) models represent a more complex task due to the typically multimodal nature of their posterior distributions. In this setting, we first consider a BNN for the binary classification task of distinguishing between the MNIST handwritten digits 4 and 9, before extending this to multi-class setting distinguishing between the digits 2, 4, 7, and 9, for a variety of BNN model capacities. Full experimental details are provided in Appendix C.3.
Following the setup of Kuntz al. (2023), we observe short transient phases in the evolution of parameter cstimates-alle albeit converging to different local maxima-a well in evaluation metrics

---

12.5
10.3
Test Eroor
LPPD
0.2
7,5
0.4
3.0
-5
2.5
0.2
20
0.0
200 Iteration,
500
100
203
300
100
Iteration, &
PGD
SOUL
JALA-EM
Tost Error a
LPPD/
Figure 3: Bayesian neural network: LPPD and Test Error, for PGD, SOUL, and JALA-EM particle clouds, with N 100 and global step-size of 0.1, as in Kuntz et al. (2023) (left). Parameter estimates, (a, (3) (log(aw) log(a.)). where due and o are the variances of the zero-mean isotropic Gaussian priors applied to the weights W and (right).
for PGD, SOUT,, and JALAA-E,, as illustrated Figure Notably, the predictive performance of JALA-EM is comparable to that of SOUT. across all experiments, which is to be expected due to the correlations that exist between particles due to the resampling mechanism present in JALA-EM. Although PGD converges faster in terms of wall-clock time, it does not support marginal likelihood estimation, which is key advantage our method. Overall, this example highlights the reasonable scalability of our proposal to more complex and high-dimensional inference problems.
6 Conclusions
We provided an SMC-based EM algorithm that exploits unadjusted Langevin dynamics together with SMC corrections implement MMLE procedure. Our method is general and can be interfaced with different optimisation methods for efficiency. One distinct feature of our method is the ability compute normalising constants on-the-fly. hence the ability to select models during training runs. We demonstrate the performance of our method on variety of regression examples, as well as Bayesian neural network example.
Limitations and future work. Besides SOUL and PGD, other benchmarks (such as using Metropolis- Hastings for the E-step gradient-free optimisation for the M-step) would be interesting consider. Additionally, our work can use large number computational tricks from particle filtering literature, since our method is closely related. Our method SMC-based, hence there computational limitations due to weight computations. Without a careful implementation, one can run into weight degeneracy issues (2008), One can resort to alternative strategies, such as tempering (Crucinio, 2025) or weight clipping (Martino et al., 2018) to improve the performance.
common in SMC, especially in high-dimensions, see, e.g., Bickel al.
Broader impact. Our method provides a general purpose training scheme for fitting LVMs which are ubiquitous in science and engineering. Our methodology help faster and more reliable model testing in a variety of scenarios, accelerating the process of model development and selection for a range number of applications. However, standard training algorithm, it could be also used to train models for harmful activities, but this risk is not any larger than standard training method.
Acknowledgements
J. is supported by EPSRC through the Modern Statistics and Statistical Machine Learning (StatML) CDT programme, grant no. EP/S023151/1. D.C. worked under the auspices of Italian National Group of Mathematical Physics (GNFM) for the support. O. D. A. is grateful to Department of Statistics, London School of Economics and Political Science (LSH) for the hospitality during the preparation of this work. We thank anonymous referees for their insightful feedback which improved our work.
INdAM. D.C. expresses their gratitude to Marylou Gabrié
10

---

References
Sergios Agapiou, Omiros Papaspiliopoulos, Daniel Sanz-Alonso, and Andrew M Stuart. Importance sampling: intrinsic dimension and computational cost. Statistical Science, pages 405-431, 2017.
O Deniz Akyildiz, Michela Ottobre, and lain Souttar. A multiscale perspective on maximum marginal likelihood estimation. arXiv preprint 0tti:224060448877 2024.
O Deniz Akyildiz, Francesca Romana Crucinio, Mark Girolami, Tim Johnston, and Sotirios Sabanis. Interacting particle Langevin algorithm for maximum marginal likelihood estimation. ESAIM: Probability and Statistics, 29:243-280, 2025.
Ömer Deniz Akyildiz and Joaquín Míguez. Convergence rates for optimised adaptive importance samplers. Statistics and Compuring, 31:1-17, 2021.
Yves F Atchadé, Gersende Fort, and Eric Moulines. On perturbed proximal gradient algorithms. Journal of Machine Learning Research, 18110):1-3332011 2017.
Peter Bickel, Bo Li, and Thomas Bengtsson. Sharp failure rates for the bootstrap particle filter in high dimensions. In Pushing the limits contemporary statistics: Contributions in honor of Jayanta K. Ghosh, volume 3. pages 318-330. Institute of Mathematical Statistics, 2008.
Brian S Caffo, Wolfgang Jank, and Galin L Jones. Ascent-based Monte Carlo expectation- maximization. Journal of the Royal Statistical Society Series B: Statistical Methodology, 67 (2):235-253) 2005.
Rocco Caprio, Juan Kuntz, Samuel Power, and Adam M Johansen. Error bounds for particle gradient descent, and extensions of the log-Soboles and Talagrand inequalities. Journal of Machine Learning Research, 26(103):1-38, 2025.
Davide Carbone. Jarzynski reweighting and sampling dynamics for training energy-based models: theoretical analysis of different transition kernels. Bollettino dell' Uni'oe Matematica Italiana, 2025.
Davide Carbone, Mengjian Hua, Simon Coste, and Eric Vanden-Eijnden. Efficient training of energy- based models using jarzynski equality. In Advances in Neural Information Processing Systems, volume 36, pages 52583-52614, 2023.
Davide Carbone, Mengjian Hua, Simon Coste, and Eric Vanden-Eijnden. Generative models as out-of-equilibrium particle systems: training of energy-based models using non-equilibrium thermodynamics. In International Conference on Nonlinear Dynamics and Applications, pages 287-311. Springer, 2024.
Francesca R Crucinio. A mirror descent approach to maximum likelihood estimation in latent variable models. arXiv preprint arXiv:2501.1589 2025.
Valentin De Bortoli, Alain Durmus, Marcelo Pereyra, and Ana F Vidal. Efficient stochastic optimi- sation by unadjusted Langevin Monte Carlo: Application to maximum marginal likelihood and empirical Bayesian estimation. Statistics and Computing. 31:1-18, 2021.
Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. Sequential Monte Carlo samplers. Journal of the Royal Statistical Society Series B: Statistical Methodology, 68(3):41 1--36, 2006.
Bernard Delyon, Mare Lavielle, and Eric Moulines. Convergence of a stochastic approximation version of the EM algorithm. Annals of statistics, pages 128, 1999.
Yury Demidovich, Grigory Malinovsky, Igor Sokolov, and Peter Richtárik. A guide through the zoo of biased SGD. Advances in Neural Information Processing Systems, 36:23158 -21111 2023.
Arthur P Dempster, Nan M Laird, and Donald B Rubin. Maximum likelihood from incomplete data via the EM algorithm. Journal of the royal statistical society: series B (methodolooical), 39(1): 1-22, 1977.
11

---

Randal Douc and Olivier Cappé. Comparison of resampling schemes for particle filtering. In ISPA 2005. Proceedings of the 4th International Symposium on Image and Signal Processing and Analysis, 2005., pages 64-69. leee, 2005.
Randal Douc, Eric Moulines, and David Stoffer. Nonlinear time series: Theory, methods and applications with R examples. CRC press, 2014.
Víctor Elvira, Joaquín Míguez, and Petar M Djurié. Adapting the number of particles in sequential monte carlo methods through an online scheme for convergence assessment. IEEE' Transactions on Signal Processing, 65(7):1781-1733 2016.
Víctor Elvira, Luca Martino, and Christian Robert. Rethinking the effective sample size. Interna- tional Statistical Review, 90(3):525-550, 2022.
Paula Cordero Encinar, Francesca R Crucinio, and O Deniz Akyildiz. Proximal interacting particle Langevin algorithms. Uncertainty in Artificial Intelligence (UAI), 2025.
Gersende Fort and Eric Moulines. Convergence of the Monte Carlo expectation maximization for curved exponential families. The Annals of Statistics, (1:)::220-12599 2003.
Alex Glyn-Davies, Connor Duffin, leva Kazlauskaite, Mark Girolami, and Ô Deniz Akyildiz. Statisti- cal finite elements via interacting particle Langevin dynamics. SIAM/ASA Journal on Uncertainty Quantification, (000))1200-2227, 2025.
Samuel Gruffaz, Kyurae Kim, Alain Durmus, and Jacob Gardner. Stochastic approximation with biased MCMC for expectation maximization. In International Conference on Artificial Intelligence and Statistics, pages 2332-2340. PMLR, 2024.
Jeremy Heng, Adrian N Bishop, George Deligiannidis, and Arnaud Doucet. Controlled sequential Monte Carlo. The Annals of Statisticy, 48(5):2904-2929, 2020.
Adam M Johansen, Araud Doucet, and Manuel Davy. Particle methods for maximum likelihood estimation in latent variable models. Statistics and Computing, 18:47-57, 2008.
Kyurae Kim, Zuheng Xu, Jacob R Gardner, and Trevor Camphell. Tuning sequential Monte Carlo sam- plers via greedy incremental divergence minimization. In Forty-second International Conference on Machine Learning. 2025.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proceedings of the 3rd International Conference on Learning Reprexentations (ICIR), San Diego, USA, 2015.
Juan Kuntz, Jen Ning Lim, and Adam M Johansen. Particle algorithms for maximum likelihood training of latent variable models. In International Conference on Artificial Intelligence and Statistics, pages 5134-5180. PMLR, 2023.
Kenneth Lange. A gradient algorithm locally equivalent to the EM algorithm. Journal of the Royal Statistical Society: Series B (Methodologiiaal) 57(2):425 437, 1995.
Tiancheng Li, Miodrag Bolic, and Petar M Djuric. Resampling methods for particle filtering: classification, implementation, and strategies. IEEE Signal processing magazine, 32(3):70-86, 2015.
Jen Ning Lim, Juan Kuntz, Samuel Power, and Adam Michael Johansen. Momentum particle maximum likelihood. In International Conference on Machine Learning, pages 29816-29871. PMLR, 2024.
Chuanhai Liu and Donald B Rubin. The ECME algorithm: a simple extension of EM and ECM with faster monotone convergence. Biometrika, 644, 1994.
Joanna Marks, Tim W. J. Wang, and Deniz Akyildiz. Learning latent energy-based models via interacting particle Langevin dynamics. arXiv preprint arXiv:2510.1233 123111 2025.
Luca Martino, Victor Elvira, J Míguez, A Ardés-Rodriguez, and PM Djurić. A comparison of clipping strategies for importance sampling. In 2018 IEEE Statistical Signal Processing Workshop (SSP), pages 558-562. IEEE, 2018.
12

---

Jonathan C Mattingly, Andrew M Stuart, and Desmond Higham. Ergodicity for SDEs and approx- imations: locally Lipschitz vector fields and degenerate noise. Stochastie processes and their applications, (01(2)1185-23))) 2002.
Xiao-Li Meng and Donald B Rubin. Maximum likelihood estimation via the ECM algorithm: A general framework. Biometrika, 80(2)267-277 1993.
Jerome P Nilmeier, Gavin Crooks, David DL Minh, and John D Chodera. Nonequilibrium candidate Monte Carlo is an efficient tool for equilibrium simulation. Proceedings of the National Academy of Sciences, 108(45):E1009-E100
2011.
Paul Felix Valsecchi Oliva and O Deniz Akyildiz. Kinetic Interacting Particle Langevin Monte Carlo. arXiv preprint arXiv:2407.05790 2024.
Paul Felix Valsecchi Oliva, Deniz Akyildiz, and Andrew Dunean. Uniform-in-time convergence bounds for persistent contrastive divergence algorithms. arXiv preprint arXiv:2510.01944 2025.
Sébastien Roch. Mathematical Methods in Data Science: Bridging Theory and Applications with Python. Cambridge Mathematical Textbooks. Cambridge University Press, 2025.
Adrien Saumard and Jon A Wellner. Log-concavity and strong log-concavity; a review. Statistics surveys, 8:45, 2014.
Christoph Schönle, Marylou Gabrié, Tony Lelièvre, and Gabriel Stoltz. Sampling metastable systems using collective variables and Erryynski-Crook paths. Journal of Computational Physics, page 113806, 2025.
Francois Septier and Gareth Peters. Langevin and Hamiltonian based sequential MCMC for efficient Bayesian filtering in high-dimensional spaces. IEEE Journal of selected topics in signal processing, 10(2):312-3211 2015.
Louis Sharrock, Daniel Dodd, and Christopher Nemeth. Tuning-free maximum likelihood training of latent variable models via coin betting. In International Conference on Artificial Intelligence and Statistics, pages 1810-1818. PMLR, 2024.
Ronglai Shen, Adam B Olshen, and Mare Ladanyi. Integrative clustering of multiple genomic data types using a joint latent variable model with application breast and lung cancer subtype analysis. Bioinformatics, 25(22)22906-9) 2912, 2009.
Sobihan Surendran, Adeline Fermanian, Antoine Godichon-Baggiont and Sylvain Le Corff. Non- asymptotic analysis of biased adaptive stochastic approximation. Advances Neural Information Processing Systems, 37:12897-12943, 2024.
Tim YJ Wang, Juan Kuntz, and O Deniz Akyildiz. Training latent diffusion models with interacting particle algorithms. arXiv preprini arXiv:2505. /2412, 2025.
Xinyi Wang, Wanrong Zhu, Michael Saxon, Mark Steyvers, and William Yang Wang. Large language models are latent variable models: Explaining and finding good demonstrations for in-context learning. Advances in Neural Information Processing Systems, 36: 16611561411500 15638, 2023.
Greg CG Wei and Martin A Tanner. A Monte Carlo implementation of the EM algorithm and the poor man's data augmentation algorithms. Journal of the American statistical Association, 85 (411):699-704, 1990.
Nick Whiteley, Annie Gray, and Patrick Rubin-Delanchy Statistical exploration of the manifold hypothesis. Journal of the Royal Statistical Society: Series B, 2025. ISSN 0035-9246.
W H Wolberg and O L Mangasarian. Multisurface method of pattern separation for medical diagnosis applied to breast cytology. Proceedings of the National Academy of Sciences. 97323):913399196 1990.
Yuling Yao, Aki Vehtari, and Andrew Gelman. Stacking for non-mixing Bayesian computations: The curse and blessing of multimodal posteriors. Journal of Machine Learning Research, 23(79):1-45, 2022.
13

---

Yan Zhou, Adam M Johansen, and John AD Aston. Toward automatic model comparison: an adaptive sequential Monte Carlo approach. Journal of Computational and Graphical Statisties, 25 (3):701-726, 2016.
A Discussion and comparisons
In this section, we put our method in more general context by comparing it to the recently proposed body of interacting particle methods for the MMLE problem eited throughout the paper. We will compare in this section, in particular, particle gradient descent (PGD) of Kuntz et al. (2023) and interacting particle Langevin algorithm (IPLA) of Akyildiz et al. (2025), as well as the slow-fast Langevin algorithm (SFLA) method discussed in Akyildiz et al. (2024).
A.1 Algorithmic comparison
Our method is closely related to PGD and interacting particle Langevin algorithm (IPL.A) which rely on particle systems to update the parameter iterates. To be specific, given some initialisation (to, - xx) PGD (Kuntz al., 2023) consists of the following discrete dynamical system:
0k+1
02
(OU( X%),
(16)
X = X - VV x() + 22
(17)
where Z ~ NO(0. are independent standard Gaussian random variables. The IPLA method of Akyildiz et al. (2025) is very similar, with a scaled noise in 0 iterates:
N
0x+1 = 0x
Y
Σ
VOU
X.) + V
Ex+1.
(18)
N
1 x (,, Xi) + 22
(19)
where Ex+1 ~ (O, la.) ) is another independent standard Gaussian random variable. In the above equations, one see parameters updates (16) and (18) play similar role to our optimiser steps in (9). However, this setting, particles are equally weighted (1/N), and the latent variables are updated through ULA step in (17) and (19) similar to our case. One crucial aspect for both methods that the step-size y needs to be same for both and X updates as these algorithms are derived from an Euler discretisation of continuous-timg time dynamics. This is significant limitation in practice as the scales of 9 and X can be very different, and thus, for example, the step-size needs to be chosen very small ensure stability of the X updates, which in turn leads very slow updates for f. related aspect here is that while, for example, Kuntz et al. (2023) uses an adaptive optimiser for 0-updates in challenging experiments, theoretical results for PGD (Caprio et al., 2025) and IPLA both tied to the use of plain gradient updates as in (16) and (18). In contrast, our method allows for different step-sizes for 0 and X updates, as well as the use of adaptive optimisers. This is because our method is not derived from a discretisation of a continuous-ting dynamics, but rather from direct optimisation of the MMLE objective.
To address some of the limitations above, the continuous-time timee dynamics behind PGD and IPLA have been extended in Akyildiz et al. (2024); Oliva et al. (2025) to allow for different time-scales for 0 and X updates. For example, the resulting slow-fast Langevin algorithm (SFLA) method of Akyildiz et al. (2024) consists of the following updates:
Oz-1 02 7VoU- Xs) +
V Ek+1
27
(20)
X = X2 2 V.U Xx) + (27
2 7
(21)
where 3 > 0 an inverse temperature parameter, > 0 scale parameter, and y > is step-size. Here, there is a single latent variable X₁ which is updated with a step-size 7/c that can be much larger than the step-size 7 used for updates (see Oliva et al. (2025) for related particle-based approach). However, this comes at the cost of introducing two additional hyperparameters 3 and which can be difficult to tune in practice. Moreover, as for PGD and IPI.A, theoretical results for SFLA are tied to the use of plain gradient updates as in (20), and do not allow for adaptive optimisers.
14

---

A.2
Theoretical comparison
Both PGD and IPLA come with nonasymptotic bounds; see the analysis in Caprio al. (2025) and Akyildiz et al. (2025). Under strong convexity and gradient smoothness assumptions on U, both methods provide theoretical guarantees of the following form:
0| 0- 0-119/1/²= & O(72/2 + N +e t=>kk),
(22)
where l 0 is strong convexity constant, ? the step-size, 0, is the unique minimiser of the MMLE objective. One can see that, in addition to the fact that these methods have to work with the same step-size y for both 0 and X updates, the above bound also requires y to be very small to control the bias term O Specifically, for the error to vanish, one needs take large N and small On the other hand, the convergence rate is governed by the exponential term (e (,))), which decays faster for larger y,, resulting trade-off between accuracy and convergence speed.
is the number of particles, and
In contrast, result in Theorem provides error bound:
B| 0.1 1/1/////2 s ((1 (/12///2)8/2 - N-1/2).
(23)
In the above display, it suffices to take N large for error to be small. The step-size be chosen independently of N and only affects the convergence speed, thus decoupling the accuracy and convergence speed trade-off.
Our results can also be extended to adaptive optimisers, see, e.g. Surendran et (2024) whereas theoretical results for PGD, IPLA, and SFLA are tied to the use of plain gradient updates.
A.3 General comparison of the methods
We provide the following table summarising the main differences between the methods discussed above.
Feature
PGD
SOUL
SFLA
JALA-EM
Marginal likelihood estimation
X
Different 0 and x step-sizes
/
Thcoretical guarantees with fixed
V
X
Adaptation of theory for adaptive optimizers
X
B Theoretical results
B.1 Proof of Proposition
Proof. This proof follows mutatis mutandis from the proof of Proposition 1 in Carbone et al. (2023), we include here for completeness.
First, note that under the ULA update given in (2), the transition probability density, that is, the probability density of moving from y at time k. is given by
(a v) (4Th) d/2 exp (-1) 1 - (lv-=+NVU(=)))
4h
We also observe that:
k
A
Σ 1X - ax X)-1)), e N.
(24)
Now, the need for this transition probability density becomes apparent when attempting to calculate expectations over the law of the joint process, as we specifically have, for scalar test function yp
E [x(())eA = =/ ((zn)e ...,...,,yyy drk-
(25)
15

---

where p(zo..... 2b) denotes the joint probability density function of the path (Xo. ..b......... ..) at & E N.
and is defined
k
p(To....,55)) (xa)
II 9=0
Bg 1g-1)>
II
3 12g
In light of this, it remains to express Ax in terms of the transition probability density, B. Indeed, we have from (2) that
A
aa (Xq-
(Xxxxxx N.
(26)
Thus, by substituting (3) into (26), have that
k
exp (Ax)
II
exp (x (x::x:xx-11)) Xg-1)),
q=1
II q=1
k
exp U Ug-1(Xx-1) + {(2x - X VUg-1(X + h 1 V
UUg(X)) ] (X X²) vU,(X,) ())
exp (Un(Xo) - Un U2 (Xx))
II exp ブタ (X, Xq-1) (X 1
1
/ 1x-
- (X--1 X²) (X)(X,) 1/ [U,(X,))))
3g(X,r X
exp (Un(X.) Un (Xx)) II
B, (X
where in the last equality have used the fact that
1 (
(4th) 1 exp
Xq-1)
4h
(27)
1 Xy)
(4th)- - cxp
4h |Xg +
1- (Xx
2 (X, Xx-1) VUn-1(X-11 + 11 |U 1 (Xg-
=exp
2 (Xg-1 X²) vU,(X,) //2 [(,x(X,))))
(28)
as we expanded the squares at numerator and denominator and we have notably leveraged the fact that dn IX,- 1²-112 1 [g-1
X, 2= 0.
16

---

Now, recalling that Z - f UU(()) dr and bringing everything together into 25, we obtain
/ Bdlk
E [xo(Xx)e
(((x)e plao. ................................ mk) dag ... .....rr dux,
& Spe plzx) exp (Uo(zo) Un(aa)) II By (Eg. xg-1) po(zo) doo drks
A
Zo S (((x)) exp (zu)()))) II (Z T- dæo dz
g=1
/ (((x)) exp (-Un(zE)) drk.
Zo
where in the last equality we have have iteratively integrated over 2o, up to and including T where we note that, for (1 E and x--1 RRd
Jas (ag Fig 1 1)d
Thus, we (Xxx) obtain
E [e
- / exp (-Ux(Ex)) dxy
1
Zk
Zo Jz
Zo
This completes the proof, since [x(Xx) Ax /B [edk] = Ex, [x() and as generalization of this result to vector-valued is straightforward.
See Carbone (2025) for a further insight on different choices of the transition kernel 3(x, y).
B.2 Proof of Proposition 2
Now, we derive an MSE bound on the error introduced through utilising the sample based approxima- tion
E
VoU (Ok; Xx)cAA]
Eee Voli (oli X 10
E [e:A:]
ΣN e
where, to be clear, the expectations on the left-hand side are over the law of the joint process (N As), as defined through (2). Throughout, we assume that ((x(") A > are i.i.d copies of (Xx, Ak),
or equivalently that no resampling is performed up to time k.
Specifically, we wish bound the quantity
EN Voli (0x, ((.,)(1)) e
E VoU (k., Xn)eA
E FR) E E
WF. W
E
E
Au
Σ ((0) (1)
where we have let W4 = As and F6 (oU Xx) for notational brevity, whilst sample based
approximations, W N N e and WF; = EN (oU x(1) 4( are indicated by an overset bar.
(1)
17

---

We follow the proof of Theorem in Akyildiz and Míguez (2021). First, we note that the following inequality holds,
E
VgU(ON: Xx)e
E VOU(Oks Xxx Je
F
WFR
E.
[eA
Σ cA4
B[W+]
WE
FB WR K WFR B[WA]]
[W W
W &
WF&) WF (W EW&DII
[E[WE]]
W&
Wxx
[[W
WF
W
WE E[W
W M
FR WFR
1
WF
W
E[W R.
W
WFR &
W &
E[A
11 W
[B[
WW
[B WF
W [[
[[[W]]
11
[[WA]
where the first inequality we have repeatedly applied the triangle inequality, and in the fourth equality simply divided the denominator through each term. Furthermore, in the second inequality
we have made use of the fact that Σ> VaU ((() Je < VV Σm" 4 recalling that we assume || V < co, for any 0 E Rº and X E Rd.
Now, utilising the fact that (a - 6)2 < 2(a² + b²), and taking the expectation of both sides, we obtain
W Fi
WF k WA
27
FR WF
W. B[W&]]
27
S
+
+1
W.
B[WA]
E[Wa]
2E
[|E[W&FR]
1F
E [W W
(E)
())
F
2 I13.
(E[W2) (E(WAI)))
22
(B)
W 12
F
E
N
(E)
N
(B[W&])²²
B[W2]
[[[2]
N
(T](W)))
N (B()
E[W31
N
(E
giving the desired result, where in the first equality we have leveraged the fact that the numerators, the right-hand side, are perfect Monte Carlo estimates, and so this equality holds through a standard variance decomposition. We also note that in the second inequality we have again used the fact that |/ so that
F = (OU X²)) is bounded by | ||
|W
| 1 W
W 11 <
E[W2)
|W < Fo |
F& E B[I F < EWlll-
To conclude, we exploit assumption A1
F
EW21
4 F1
E
<
sup
(29)
N
[
N
(E
18

---

B.3 Proof of Theorem 1
We specifically leverage Theorem 4 from Demidovich et al. (2023), ensuring that the respective assumptions are satisfied.
AS. There exist constants, A,,,C,b,c O such that the gradient estimator, for every 0 E Rd satisfies
EVg(), (1 [0]1 (V) C.
(30)
(31)
where (0.). Theorem 3. Let Assumptions A2, A3, and A5 hold. Choose a step size, 7, such that
pb
0<7 min
(32)
[[AA + BB) pb)
then have, every &s E N,
LO7 c
E (0&e) V*] $ 50 +
+
(33)
2pb
where 60 = V(6o) V*
In light of Assumption A5 lacking an intuitive explanation, we turn Assumption Demidovich et al. (2023):
A6. For all f E Rd. there exists A M 0 such that
[llg(0) vV A
(34)
We remark that A6 notably a stronger assumption than that of A5, and through Theorem 13 of Demidovich et al. (2023) we have that Assumption A6 implies Assumption A5 with A = 0, B = 2. C' 2A b 2. and 2 A²
Critically, if the assumptions of Proposition 2 hold, we have that
4C
VoU
E
[VV)V(0)) 9(ww))l A
A
(35)
where C = sups (edu)²] / (TE [) and so Assumption of Demidovich et al. (2023) is satisfied for this A.
Thus, if Assumptions A2 and A3 hold, as well as those of Proposition 2, we have that, for a step size
pub
y
0 < min
min
T((A + WB) pbb
T.(0 + 2μx)
2
1 2
=min
4L 11
for every & E N,
LC* +
((((Bk) V* s (1 -xbb) +
2ub
ub¹
11
2
(1 ? 2
2
7/ k +
8Ly 2U
V o
(1
C +
uN
PN
where 6, = V(0o) - V* To conclude, note that V(0+) V* log pg- (v) - log Po: (v).
19

---

C Experimental details
C.1 Implementation details
In this section, we outline the details of implementation for our method.
The resampling step in Algorithm proceeds follows. Given that it is triggered step k,, first resample the walkers X using the normalized weights w as probability of picking the i-th particle and we reset A.T O. Several routines for resampling have been developed in the context of particle filterin,, cfr. Douc and Cappé (2005); Li et al. (2015). In our experiments we decided adopt systematic resampling, since not only is it computationally efficient, with order (NV), but also the variance of the number of resampled particles is reduced, compared to that of say stratified resampling (Li et al., 2015). As described in Carbone et al. (2023), an extra stage is then necessary to track Zy through the resampling step, that is adopting
Z Zk,
for all k 2 k, until the next resampling step.
C.2 Bayesian logistic regression Wisconsin cancer data
First, we examine JALA-ENO performance, in the context of training a Bayesian logistic regression model, relative to that of the Particle Gradient Descent (PGD) and Stochastic Optimisation via Unadjusted Langevin (SOUL) algorithms, as introduced in Kuntz et al. (2023) and De Bortoli et al. (2021) respectively. As is the case in the former paper, utilise the setup described in the latter paper, which we describe in detail below.
Setup. The empirical Bayesian logistic regression problem is formulated through considering the Wisconsin Breast Cancer dataset Wolberg and Mangasarian (1990), which can freely accessed at
https://archive ics. uci ott.aa//////ttaeeeeerraeeeeeaerr (original))
In particular, this dataset contains 699 samples, obtained between January 1989 and November 1991, where each sample consists of dz = 9 latent variables extracted from digitised images of fine needle aspirates (FNA) of breast masses, such the Clump Thickness and Marginal Adhesion, as well as the corresponding malignant or benign diagnosis. Since missing values exist within this dataset, we opted discard the corresponding samples, resulting in a dataset of dy = 683 samples.
Regarding further data pre-processing, the latent variables, or features, were collated into a matrix X Xdd and were then standardised column-wise to have zero mean and unit standard deviation. To be clear, if Hy and sj are the mean and standard deviation of the j-th feature column respectively, the normalised features, X E Rd, are given by X.,, (X /z)/sj- Furthermore, note that the response is mapped such that y E {0, 10,,, where y/ represents benign case and = 1 malignant one. Lastly, note that the pre-processed dataset, = (X, y). was then split into a training-validationn set, Derain, vol and a hold-out test set Dreat, via an 80-20 stratified split, so that class proportions are maintained across said split.
Model. As mentioned above, we adopt empirical Bayesian logistic regression framework, where, for a single sample (z,, y:), in which x; E Rdd and ys E {0, 1}. the corresponding likelihood follows Bernoulli distribution.
p(, w) a (I.Tw) (1 (α, w))
where, be clear, w E dd are the regression weights, and (z) = (1+e is the standard logistic function. Following Kuntz et al. (2023), we assign an isotropic Gaussian prior to the regression weights, that
p(w o3) N (w|0 (
where La, the d, dimeinnnial unit vector, and of is the prior variance. Notably, the parameter 0 E I. that we wish to estimate. Thus, it follows that the unnormalised negative log pog-posterior, or
20

---

the joint potential, for single particle, indexed by m, conditional on the observed data (in the training-validateo set)
[errasn matl
Uw 0
(log (1+ E (()))))) Mi (xx] (eTTw(m))
0lde
200
m=l
where we refer to Appendix D.1.1 for full details. Indeed, as outlined in Proposition 1 of Kuntz et al. (2023), the marginal likelihood has a unique maximiser.
Approach. To facilitate a robust comparison of JALA-EM with PGD and SOUL, we implement a systematic and reproducible procedure for model fitting and hyperparameter selection. A critical aspect is the choice of step-sizes, as these significantly influence the dynamics and performance of each algorithm. Notably, JALA-EM and SOUL permit distinct step-sizes for their particle (Langevin) updates and their parameter (0) updates. This flexibility is not present in PGD, where a single step-size controls both processes. The practical importance of tuning step-sizes is underscored by the toy hierarchical model in Kuntz et al. (2023), where mean-field analysis yields different optimal step-sizes for PGD and its variants presented within the paper.
Therefore, rather than using a fixed global step-size for all algorithms, or all dataset splits, we perform K-fold cross-validation on Perain,vo to select the step-sizes. Specifically, this tuming is conducted over a predefined grid for the particle update step-sizes and, where applicable (i.e. for SOUL and JALA-EM), for the 8-update step-sizes. The evaluation metric utilised for this hyperparameter tuning, within the cross- validation folds, is the Log Pointwise Predictive Density (LPPD), which assesses the model's s average predictive accuracy on unseen data.
Let P = { (k) === be the ensemble N particles (regression weight vectors) obtained from an algorithm run, to fit the statistical model described above. For each data point (z,, ys) E D Dyal. we first compute the predictive probability of the true observed label, yis averaged over the N particles,
Pavg(yi (Mi 20i, P)
Σ
(a (8)) (1-
(1
(x w()))) 1-v=
The LPPD for Drat is then the average of the natural logarithm of these pointwise predictive probabilities, across the respective dataset,
Doal
1
LPPD(D
log (162,112 P),
Deett
where higher LPPD value signifies superior predictive performance of the fitted statistical model, since the mean KI. divergence between our classifier and the optimal classifier will be smaller (Kuntz al., 2023). Notably this quantity is equivalent the negative of the predictive Negative Log-Likelihood (NLL), which we instead choose to minimise in our implementation. Having selected step-sizes, for the respective algorithms, from our grid, we then re-run the algorithms on the entirety of Dirran,eat, which are then ready to be evaluated on Dtest-
Implementation Details. As is the case in both Akyildiz et al. (2025) and Kuntz et al. (2023), we set the prior variance to of = 5.0, and also initialise the scalar parameter (to estimate) at 6g = 0 for all algorithms. The N particles were initialised by drawing samples independently from the Gaussian prior distribution, p0,] af), which now notably simplifies to (0, 5la.). Regarding the step-size tuning procedure, we opt for K = 3, motivated by the computational cost of the SOUL. algorithm, scaling significantly as N increases. We note that carly stopping during the tuning occurs the 10 over 10 consecutive evaluations,
validation LPPD does not improve by at least ELPPD that are each spaced 10 iterations apart, as we argue that in such cases the stationary phase has been reached. Furthermore, the maximum number of iterations for tuning, during each fold, was set to K = 500. To be clear, the case of SOUL, this refers to the number of outer steps, whereas the number of inner steps is determined by N. Also, note that for JALA-EM, choose C 1/1.05 and utilise systematic resampling in cases in which this threshold is breached, as recommended in Carbone al. (2023).
To form our grid of step-sizes, relevant to the particle updates, we compute baseline step-size, which we denote heuter. Specifically, heater is computed to provide a theoretically grounded scale reflecting
21

---

the curvature of the energy landscape, and to this end we compute the worst-case upper bound on the Hessian, Hoonnd to obtain single, globally relevant Lipschitz constant, L, for determining hederr- Indeed, this is constant matrix that is larger, in the positive semi-defiitte sense, than any Hessian encountered, and thus its largest eigenvalue provides an upper bound on valid across the entire parameter space. In fact, for the Bayesian logistic regression model, this bound takes the form
H Honnd train, pal Xtrain
where the largest eigenvalue estimated numerically through the power iteration method Roch (2025), giving global estimate of L. Lastly, to incorporate small margin of safety, we take heater = 0.99/ L. The grid for the particle update step-size was then constructed as 10 linearly spaced values in the range [0.2 heuler, 2.0 x heater], whereas we limit the grid of 0-update step-sizes be simply (0.05, 0.1, 0.15), again due to the significant time takes to tune SOUL as N becomes large.
Having selected step-size, each algorithm was step-sizes, for which trajectories were limited to a wall-clock duration 5.0 seconds, with an iteration milestone of 200 steps, that we indicate, if reached, in the respective visualisations. To be clear, we did not run the algorithms here for the same number of total iterations, but instead the same wall-clock duration.
the entire of Dirain, using said
Results. Here, we present the comparative performance JALA-EN against PGD and SOUL for the Bayesian logistic regression task described above, for the cases that N 100, N = 50, and N 10, which is explicitly visualised within Figures 4, 5, and respectively. In particular, we focus on the convergence of the estimates of the global parameter 0, the Kernel Density Estimates (KDEs) for the representative regression weight, W2.
0.45
0.5
KDE
0.0
0.1
2
2
Time
SOUL
JALA-EM
200 Iter.
Figure 4: Bayesian logistic regression: Parameter estimates for PGD, SOUL, and JALA-EM, for
N 100 and tuned step-sizes, with a iteration milestone 200 indicated (left). KDE of the second coordinate of the posterior approximation for the model weights (right).
1.0
0.15
0.5
KDB
0.30
0.15
0.1
2
3
-2
0
2
3
Time
PGD
SOUL
JALA EM
200 lter.
Figure 5: Bayesian logistic regression: Parameter estimates for PGD, SOUL,, and JALA-EM, for N 50 and tuned step-sizes, with iteration milestone of 200 indicated (left). KDE of the second coordinate of the posterior approximation for the model weights (right).
22

---

0.60
0.8
0.45
0.5
KDE
0.30
0.15
0.0
0.00
Time (5)
JALA-EM
200
Figure 6: Bayesian logistic regression: Parameter estimates for PGD, SOUL, and JALA-EM, for
N 10 and tuned step-sizes, with iteration milestone of 200 indicated (left). KDE of the second coordinate of the posterior approximation for the model weights (right).
Across the particle counts, JALA-EM demonstrates rapid convergence for the parameter consistently matching or even outperforming PGD and SOUL in terms of wall-clock time, while reaching the same estimated value. Although across algorithms all posterior approximations, for the representative regression weight W2, are comparable at N = 100, the posterior fidelity of PGD and SOUL noticeably contrast, JALAA-EM maintains high-quality approximations, even with
degrades as N decreases. fewer particles, which we attribute to its effective particle management due to reweighting. Although JALA-EM requires an additional step-size be tuned compared PGD, provides a compelling combination of fast parameter estimation and accurate posterior approximation in this setting.
C.3 Bayesian neural network
To investigate the performance of JALA-EM in a more challenging setting, we opt to apply to the task of training Bayesian neural network (BNN). This setting is generally more challenging due to the higher dimensionality of the parameter space and the commonly complex, multimodal nature of the posterior distribution over network weights. Again, we compare JALA-EM with PGD and SOUL.
Setup. To begin, we focus on the setup the binary classification problem, since the multi-class setting is natural extension. To be clear, using the MNIST handwritten digit dataset, the task is to distinguish between images of the digits 4 and 9, which we remap to classes 0 and 1 respectively. Initially, we filter the dataset for these two digits, and then in the interest of computational speed, as done in Kuntz et al. (2023), we draw a subsample of 1000 images, fa E TR. D2 x8=781 and their corresponding labels {0, 1}. ensuring class balance maintained through stratified sampling. This 1000-sample dataset, Daub. forms the basis for further processing. Then, image features are standardised by subtracting the pixel-wise mean and dividing by the pixel-wise standard deviation. with statistics computed over sabb Subsequently, the processed sub is split into final training set, Dusin. and a hold-out test set, Diea. using an 80-20 stratified split to preserve class proportions.
In the multi-class setting, we instead aim to distinguish between images of the digits 2, 4, 7, and which remap to the classes O, 1, 2, and 3 respectively. Again, the dataset is initially filtered for these digits, from which we then draw a subsample of 2500 images through stratified sampling, forming aan here. Subsequently, Dirain and Dies are formed in the exact same manner described in the binary setting above.
Model. We employ two-layer fully connected Bayesian neural network, with notable simplification in the architecture that the bias parameters are set to zero both layers. This model structure, including the absence of biases and the empirical Bayes treatment of prior variance hyperparameters (see below), aligns with Kuntz et al. (2023) and Yao et al. (2022).
Specifically, the first layer maps the input feature vector f. E R to a hidden representation using weights Wi E x Dr and a hyperbolic tangent activation function, with Dj, hidden units. The second layer then maps the hidden representation to the output logits E RD=, using weights W2 R DD.,,, where D.
for binary classification and D, in the multi-class setting. With this in mind, the logits for the e-th class of the i-th sample are then zíe (W₂ tanh(Wi
23

---

while the likelihood for single data point, (f,,4), is p(4,.fs, W1, W2) = (roftmax(2+))::1 Note this is equivalent to using a categorical cross-entropy loss where the log-likelihood term is (log s=f)m
Now, isotropic Gaussian priors are then assigned the weights of both layers, that p(Wila) N(W1 e²a LD, >D)) and (8)))8) (W222 ID.x... The parameters E R and _ E I control the logarithm of the standard deviation of these priors and are themselves estimated from the data as part of empirical Bayes procedure, collectively denoted @ (a, 3). Let Dw, D= and Dw: = D,D, be the total number of parameters in W, and W2, respectively. Then, the joint probability density of the weights and the training labels, Luain
{(f,, ()) given 0, is: p441/ W2)
p( W2, fuiin (0) (W)
II
The unnormalised negative log-posterior, relevant for the Langevin- based algorithms, is then UW1, W2, 0| [uain) = - log PW1. W2, Lumin()) which implicitly includes the log-prior terms.
Approach. The primary goal is to compare the efficacy of JALA-EM against PGD and SOUL in training this BNN. Both the LPPD and classification Test Error on Diest are used to evaluate the average predictive performance. Indeed, given an ensemble of N particles ((W/)) wr)))) Tk== probability for the true label of a test sample f; is
)
Payg (4ils, ((W )
softmax(W2 tanh(W 1:))
(a))
k=
The LPPD for Drest is then:
LPPD(Duua) = [Dusel (f..))D...
Σ
log(max(10- 30 Parg(liffa ((W(h) w )==)))
where probabilities are clipped here the interest of numerical stability.
Implementation Details. First, we describe the details for the binary classification setting. Regarding
initialisation, the initial parameter values are ag 0.0 and 8. 0.0. For all algorithms, N 100 particles are initialised by drawing weights from their respective priors p(Wi and p(W2 (8)). which simplifies to (0. 4) given Op = 3g
Furthermore, each algorithm runs for K K Commonn_tees jitess = 500 iterations. In contrast to the Bayesian logistic regression experiment (see Appendix C.2), this experiment utilises a global, fixed step-size of 0.1, for all algorithms, rather than tuning them via K-fold cross-validatin validation, as this is an example where computational (or expertise) limitations prohibit comprehensive fine-tuning. As such, all algorithms, including JALA-EM, are implemented in JAX. Note that for JALA-EM, EEM, we choose C = 1//1005 also note that, the interest of numerical stability some extra clipping occurs for significantly large norms, and the gradients used for updates are normalised as in Kuntz et al. (2023). The most notable change in the multi-class setting is that we instead utilise N = 50 particles, for all algorithms, in the interest of computational speed.
was the case before. We
Results. In the binary classification scenario, we observe short transient phases in the evolution of parameter estimates- albeit converging to different local maxima as well as in evaluation metrics for PGD, SOUL, and -EM, as illustrated in Figure 3. Notably, the predictive performance of JALA-EM is comparable to that of SOUL, which is to be expected due to the correlations that exist between particles due to the resampling mechanism present in JALA-EM. Although PGD converges faster in terms of wall-clock time, it does not support marginal likelihood estimation, which is key advantage of our method. In the multi-class scenario we obtain average final misclassification error rates (and standard deviations) of 4.74% (0.34%), 9.24% (0.93%). and 9. 9.20% (1.27%) for PGD, JALA-EM, and SOUL respectively, across 10 repeats. Indeed, one can vary the number of hidden neurons, and to this end we repeated the experiment for D₁ 512 hidden neurons, resulting in BNN of 403,456 latent variables, where note we had BNNs of 31, 1,440 and 31,520 latent variables previously, for the binary and multi-class scenarios respectively. In this higher-dimension average final misclassification error rates (and standard deviations) obtained were 4.98% (0.42%) 7.20% (0.81%), and 8.02% (1.02%) for PGD, JALA-EM, and SOUL respectively, across 10 repeats.
case, the
Overall, this example highlights the reasonable scalability of our proposal to more complex and high-dimensional inference problems.
24

---

C.4 Error model selection Bayesiin regression
Having investigated JALA-EM's efficacy in parameter estimation within both straightforward and more complex task, we now examine the algorithm's 's capability to estimate marginal likelihoods. In fact, this is dual capability since parameters of interest are still estimated. particular, we consider the distinct, yet related challenge of Bayesian model selection.
Setup. We consider model selection problem for a dataset, D = ((1i, u:) generated by an underlying model, G. where E 4R and y/s E IR, which we collate into X E 2 and E Rd, respectively. Specifically, we focus on Bayesian linear regression models for D, that aim to capture the relationship y, X,w E; through a latent weight vector Rd where E, represents the observation error. Indeed, JALA-EN can not only be leveraged to estimate the parameters of a given model, but also can be used to estimate the marginal likelihood, facilitating model selection through the use of, say, a Bayes Factor.
Models. In this experiment, we consider pair of competing, nested Bayesian linear regression models, where for both models, an isotropic Gaussian prior placed on these weights, p(wla) (00 (2.)) where 0 is the precision, however, the models importantly differ in their assumptions about E,.
The first model we consider, namely the Gaussian regression model, MM. assumes independent and identically distributed (i.i.d)) Gaussian errors, E; -N o²). leading to the likelihood,
dg
PlulX, W. ,.. = II N Nin o²) = (N(ll o² Ta,,)
i=1
where, for numerical stability, and to ensure the positivity of both o² and a, we choose to work with the logarithmic parametrisations of the parameters of interest. To be clear, estimate Ag (Φ1, $2). with Φ1 log x² and Φ = log a.
The second model we consider is the Student-t regression model, .M,, which assumes i.i.d. Student-t distributed errors, E; ~ Student-r(O,, o²,v), v), leading to greater robustness to outliers, and results in the likelihood
dy
T((U +((u+1)/2) (1+ (y: Xzw)2 -2++1//2
plulX, w. o², M, Mr) = II l(w/2)/TTO'
DO²
where the additional parameter, v, is the degrees freedom. Here, an exponential prior is placed on v, that is p(r|>>) Ave with the rate parameter, A = 0.1 fixed throughout all experiments. Again, for numerical stability, and to ensure positivity, we work with logarithmic parametrisations, and so estimate 0r = (Ø), d22, Φ3). with Φ3 log p, while Φ1 and Φ2 are as before. Notably, as + J0, the Student-t distribution approaches the Gaussian distribution, rendering MG nested within Mr. To ensure that MM provides distinct heavy-tailed modelling and to avoid identifiability issues with Mc. particularly when G = .M., we constrain v E [0.2, 5.0], achieved in our implementation by clipping Φ3 E [lo2( log(5.0)).
Approach. For the synthetic dataset, D. generated by the regression model G E (Ma, we employ JALA-EM to iteratively refine our initial estimates for 0g and 0r, while concurrently estimating the corresponding model's marginal likelihood. Indeed, after K iterations, by Proposition 1, we have that our estimate for the log marginal likelihood is given by
N
log 2M, = log ZM,D + log
Σ
Axx
log N =
where ZM,O o is the marginal likelihood evaluated at the initial estimates, 0, We adopt the following version of (14) which includes resampling as usual in SMC contexts, cfr. Del Moral et al. (2006):
log 2M.K log ZM.
Σ log
1 N
exp
kj
A ho
(36)
j=1
where ky the step at which the j-th resampling event occurs, with j 1,..., whilst A is the log- incremental weight of particle at step m, and Ely Σm=k--+t A represents the cumulative log-weight of particle i between two resampling steps.
25

---

Notably, the method for computing ZM is model dependent, and for Ma this is analytically tractable,
log MM. = log N(y)0, + ta txxx = =- (d,, log((2) + log(det(Zo)) + " '),
where Σg = ofLd, + ab xxx (see Appendix D.2.1). To be clear, of and ag form the initial parameter estimate, 8G,0.
For My. however, ZMr,O is generally intractable, and so we estimate it using Importance Sampling (IS), drawing S samples, {wa)==1- from a proposal distribution q(w). Specifically, we choose q(w)= (V (4g. Eg), where By and Mg are the posterior mean and covariance, respectively. of w under Gaussian likelihood model, p[w] X. y, o 0g) o0)p(w/)o)) using the initial of and og from 0r.0- The IS estimate of the initial marginal likelihood is thus
1
8
ply) X, wsy oi vo, MT)p(uu ((w))
2Mr.0 .0 R3
S
where also makes up 0r.o-
It is then natural to perform model selection using the logarithm of the approximate Bayes Factor, given by log( = log MM log My where a positive value indicates, conditioned on the estimated @ and 0r, that the data provides greater evidence for MM as the true data generating process for D.
Implementation Details. In our numerical experiments, synthetic datasets were generated, com- prising of dy = 500 samples and d, = 8 features, where to be clear, the feature vectors 2;, for each sample, were drawn i.idd. from a standard multivariate Gaussian, NO, IA,.) ). Indeed, two distinct sce- narios for the errors were drawn from E; ~ A000, o²)). whereas in the Student-t case, g = Mr, observation errors were drawn from & Student-t(O, o2 v.). In both cases, the true weight precision a. 1.0, as is the true eиror variance a,, = 1.0, while the true degrees of freedom is set to v, = 4.0. For each of the data-gent generatinn scenarios, we note that the complete experimental trial, that is the data generation. fitting of MM and Mr via JALA-EN EM, and subsequent model selection, was repeated 100 times, for which we report the proportion of trials in which the correct model was recovered. To be clear, to generate single D, single true weight vector, W., is drawn from N(0, o. TT.) ) = _N, I.,) Indeed, we utilise different random seeds for each experimental trial to ensure variability in both the synthetic dataset generation and in the stochastic elements of JALA-EM.
generating process, G, are considered, where in the first
MM. and observation
Regarding the configuration of JALA-EM, we initialise model parameters perturbed from their true values, so that M, = (log o 1, log o. + +,....) )) = 1,. ). where we set log Un log v, + log(4.) (..)) + 1 when G MMr, and log 1/1 log(5.0) when G .MG. corresponding to the upper limit of our constraint for v. The algorithm is run for K = 250 iterations, using N = 50 particles, with Langevin dynamic step-size of h 5 x 10 while the parameter optimisation learning rate is 17 = 5 x 10-3, where OPT is in fact Adam (Kingma and Ba, 2015), with 3 = 0.9, to demonstrate optimisers other than SGD can be leveraged. In the case of Mc. particles are drawn directly from the analytically available posterior, [[w|X, y, (G.o). which is notably a Gaussian. For MMr, however, particles are generated by N independent short MCMC runs, where each chain targets the posterior [w]XX, y, 0r.o), and is initialised from the prior p(w]oo)) Specifically, each MCMC run generates a particle by evolving chain for 200 steps, using Unadjusted Langevin Algorithm (ULA) updates, so that its final state is taken as the particle sample. The ULA step size at step 4, denoted e,, is adapted dynamically based on the -o-norm of the gradient of log-target density, 9t = [ log logp(wel-) |. In fact, we initialise 6o = 1 X 10 and let ct+1 0.9e; if g; 1000d, and e; > 10 Conversely, if ge < 10dg and e < 0.1 then we let f++1 = 1.05c,, and otherwise E, is left unchanged. To ensure we are consistent with Proposition 1, we fix the resampling threshold to C = 0, so that particle resampling is essentially disabled, however note that one can also track Z through the resampling step, outlined in Section 5. Lastly, note that during the IS, S 5000 samples drawn.
We begin by commenting on the behaviour of JALA-EM for single representative (at least for the 100 repeats) synthetic datasets, each comprising of dy = 500 samples, as illustrated in Figure 7. Indeed, not only is JALA-EM'Ss 's ability to converge quickly to sensible parameter estimates illustrated, but so is its capability to differentiate between the true and misspecified model.
Results.
26

---

In the case where the true underlying model is Gaussian, that = MMa, JALA-EM fitting Me quickly estimated the MI-II (see Appendix D.2.2) values (Ø), (2). and what is more, the estimated marginal likelihood quickly stabilises close to its true analytical value. To be clear, by this we mean the analytical value of the marginal likelihood corresponding to the true parameters, 0., and we denote this by log M. An important, yet subtly different quantity, is the analytical value of the marginal likelihood corresponding to the parameters estimated at the iteration step k, which we denote by log M As outlined above, this is a tractable quantity for Mc. and we observe this to correspond exactly with the estimates generated by JALA- ALL-EM in this case, importantly corroborating Proposition 1.
Conversely, when the true underlying model the Student-t, that more effective at recovering the true parameter values, and importantly its log My stabilises at higher value. In fact, in both scenarios the estimated marginal likelihood trajectories differ significantly, suggesting potential for JALA EM to perform robust Bayesian model selection.
MMr JALA-EM fitting Mr
Notably, when applying the previously described procedure with dy 1500, instead of d, - 500, samples, as illustrated in Figure 8, we observe broadly consistent behaviour, albeit with more decisive model discrimination. In fact, with more data the parameter estimates for the correctly specified models maintain their accuracy, while the separation between estimated marginal likelihoods becomes more distinct. This increased distinction is consistent with Bayesian principles, where the increased amount of data provides stronger evidence, and indicates JALA-EM's ability to effectively leverage larger datasets for more confident model selection.
Ma
1.00
Thee
1.0
Tone True
MM MI--II
Mr ML-II
0.8
-740
0.50
0.6-
770
10g
0.00
0.2-
-0.50
800
Amalytic log Zaias
-0.75
0.0
Analytic
200
90
150
200
150
0 AM7
2.0
1.0
Trx 02
Analyic log Zata
-940
15-
0.6
1.0
-800
0.0
-0.4
150
Itstatinn
Deration,
JALA-EM (Ma)
JALA-EM (M4)
Figure Error model selection: Parameter estimates (left & middle) and marginal likelihood estimates (right) for JALA-EM fitting Mg and Mr, where the true underlying model is the former (top) and the latter (bottom). Here, d, 500, d_ 8, and N 50. Parameter estimates are initialised at values perturbed away from the true values.
27

---

6- Ma
1.0
Tree bi
1.
2125
Mg
MMg
-200
2115
2200
0.0
-2228
2270
-1.0
Aualytic log Zaiah
0.0
-2325
Analyti Zuo
200
50
150
200
Mr
1.0
1.0
0.8
mm
-2025
22600
100
0.2
0.0
2725
-0.2
-02
2790
-044
0.4
Analyme Z
200
JALA EM (Ma)
JALA-EM (MM)
Figure 8: Error model selection: Parameter estimates (left & middle) and marginal likelihood estimates (right) for JALA-EM fitting Mg and Mr, where the true underlying model is the former (top) and the latter (bottom). Here, dy 1500, d, 8, and 50. Parameter estimates are initialised at values perturbed away from the true values.
Regarding the aforementioned 100 repetitions of the core experimental trial, with model selection being performed by selecting the model with the higher log 2 we found that the underlying model class correctly identified 100% of the time when = =Ma. and 99% of the time when MMr
C.5 Model order selection Bayesian regression
To further evaluate JALA-EM's performance, in a more complex model selection context, we consider the challenge of Bayesian model selection in polynomial regression setting. Once again, this highlights JALA-EM's 's dual capability to estimate both parameters of interest and marginal likelihoods.
Setup. In a subtly distinct manner, we now consider a model selection problem for a dataset,
D [(x,, 2- generated by an underlying model, G. where both x; R and y: R. Specifically, we focus on a family of Bayesian linear regression models that aim to capture the relationship :: up (a;) T w, + &; through latent weight vector Wp Rd, where Sp (xx) E IRd is model specific feature vector and s; represents the observation error. To be clear, we then collate up (æ:) and y into X, E IRdy xd, and y E IRd, in a similar manner as seen in Appendix C.4.
Models. In particular, we consider a larger set of competing, nested Bayesian linear regression models, (Mp) Notably, these models share the same probabilistic structure for both their weights and observation errors, differing instead in the model specific features utilised.
28

---

Once again, for all models, an isotropic Gaussian prior, p(wla) = aN is placed on the weights, as are i..d Gaussian errors assumed. Indeed, this leads the likelihood,
d
plu|Xp, up, M = II (Villpp(r)) L=1
"p = =N(yl pppp o² Id,,)
where we again choose estimate 0 (di 022, with Φ1 log 0² and Φ2 log a, for numerical stability and guarantees of positivity.
Indeed, the p-th model in the candidate set, Mp. corresponds to a polynomial regression model of order p, that assumes a model specific feature vector (pp (2:) derived from the polynomial basis expansion up and including order p. so that 9p (2;) [1, ",,T,,,,,,,7.. As consequence, we have that d, p + 1, and so Pp (x:)) E R Wp RR++1, and Xp (ddx(++1)
Approach. For the synthetic dataset, D. generated by the true regression model G .M.. we utilise JALA-EM to fit each of the P candidate models, with the aim of estimating p.- For each candidate, JALA-ED EM iteratively refines the initial estimates for 0, while concurrently estimating the corresponding model's marginal likelihood. Indeed, after K' iterations, by Proposition 1, we have that our estimate, of model Mp. for the log marginal likelihood is given by
N
log M,,.KK log Z,,,0 + log
Σ Ax Eetm)
)
-log log N = plul( O Mp.,
where ZM, is the marginal likelihood evaluated at the initial estimates, 0M for model Mp. Notably, in this case, the method for computing ZM,.u is identical for all candidate models, and since we have Gaussian errors in all cases, it analytically tractable for all orders.
log Z,, 0 -logNN o² Id, + ta xxx = --2(0 log(2r) log(det(Zio) + 20020-1 ')),
where Σo = Li + a X (see Appendix D.2.1). To be clear, of and an form the initial parameter estimate, 0,,,..--
Model selection is thus naturally performed by simply selecting the model order that maximises this quantity, namely ş = arg max log M,,.
pE[P)
As competitive baseline, we compare our method to an approach that leverages Ordinary Least Squares (OLS) for weight estimation, followed by selecting the model with the lowest Bayesian Information Criterion (BIC). To be clear, for each candidate model, we first compute the maximum likelihood estimate (MLE) for the weight vector, MLE by minimising the residual sum of squares (RSS). Indeed, under the assumption of Gaussian ertors, this reduces to standard OLS, whose solution given by
P.LLE = (X xT5Tx,))110
where Xp has full rank, so that the Gram matrix is invertible and pLLLLE established the desired minimiser of the RSS. The MLE for the observation variance, o is computed in a similar manner,
22
1ly X
M.MLE
dy
from which the BIC for Mp is then calculated,
BIC, = (p + 2) log dy - 2 2oggpul (.,.LE, o7
&p,MLE, MMML))M)))
where we subsequently select the model order that minimises the BIC, that we select the model order Prase arg min BIC,-
pE[P)
Implementation Details. the numerical experiments, synthetic datasets were generated, compris- ing of dh, samples. In this case, the values of the single feature from which the model specific feature vectors are derived, were drawn i.i.d from a Uniform distribution, 4[-2.5, 2.5). Indeed, the observation errors were drawn from ~ (0 of). where the true error variance is o3 7.5.
29

---

whereas the true weight precision is a, 1.0 so that the true weight vector, w., is drawn from NO, a, Iap)) / (0, 1p+1). Notably, each experimental trial has an associated true polynomial order, P.,, so that y4 = Pp. (x;) w. EG, where we have Sp, (x,) = [ T........... J.
Regarding the configuration of JALA-EM, we again initialise model parameters perturbed from their true values, so that 8,,,.U (log o- + 1, log a, + 1) (1, The algorithm is run for K 200 iterations, using N 50 particles, with Langevin dynamic step-size of h x 10-6, while the parameter optimisation learning rate is 5 x 10-3, where OPT is again Adam, with 31 = 0.9. For all candidate models, which share Gaussian likelihood structure, particles are drawn directly from the analytically available posterior p( (,,,Xp, y, 0 (,,..).. Once again, to ensure we are consistent with Proposition 1, we fix the resampling threshold to 0.
Results. To evaluate the performance of JALA-EM in robust and systematic manner, for each of dy {100, 250, 500} separately, we iterated through a range of true orders, p. [2, 8], and ran 100 repeats of the core experimental trial, with different random seeds for each trial. To be clear, the task for both JAL.A-EM and the aforementioned baseline to select the true model order from a candidate set of polynomial degrees, p E [1, 10). In order to reward close estimates, for cases in which the estimated order is incorrect, we choose utilise the Mean Absolute Error (MAE), over Accuracy, for evaluation. The results, across the true orders, can be found in Table 1.
Indeed, we observe JALA-EM match or outperform the baseline for all model orders considered, achieving notably lower average MAE in cases where the underlying polynomial is of higher-order. In such cases, the true model is more complex, possessing higher-dimente lighe--dimensional latent variable space, in which the Bayesian treatment of JALA-EM appears more effective at balancing model fit and complexity. In contrast, the baseline's 's approach, which notably relies on a point estimate, vi p.M... and the BIC approximation, is less reliable, and in fact exhibits propensity to underfit, due to the marginal improvement in oog-likelhhodd being outweighed by increases in the complexity penalty for our non-asymptotic regimes.
Table 1: Average MAE values model order estimates, for true orders E [2, 8). obtained over
100 experimental trials, for JALA-EM and the OLS & BIC baseline. Here, d, {100, 250, 500) and N = 50, whilst the candidate model orders are E [1, 10). Parameter estimates are initialised at values perturbed away from their true values.
JALA-EM (MAE)
OLS & BIC Baseline (MAE)
True Order
dy
100
dy
250
dy
500
dy 100
dy 250
dy 500
2
0.67
0.45
0.35
0.64
0.46
0.35
3
0.73
0.41
0.25
0.90
0.42
0.34
0.73
0.49
0.30
0.80
0.60
0.30
5
0.51
0.40
0.27
0.64
0.42
0.35
0.35
0.31
0.18
0.49
0.32
0.23
7
0.23
0.15
0.09
0.44
0.28
0.15
8
0.28
0.19
0.13
0.37
0.30
0.20
D Further theoretical results for experiments
D.1 Bayesian logistic regression Wisconsin cancer data
D.I.1 Derivation of the negative log-posterior
Here, we detail the derivation of the unnormalised, negative log- postpoorrrror, for single particle (k)) conditional on the observed data Duain, = {(21, y)) (Puin.al and the parameter 6, which we denote (k)( D train, val).
Indeed, the posterior of the weights and parameter 0, given the data Dinin val. follows from Bayes' theorem,
p(w, 0 ( val) x p(D (DDrrrin wl
30

---

For fixed 8, as is the case when evaluating the potential for specific particle w(*), and by treating the prior p(0) as uniform, the unnormalised posterior for wk) is
p(w (k) (Puain, val. 0 o oc ((riiii, val (x) (b) }p(p)) (k) 10,
The potential function (ww(k) Diaiiai val) is defined as the negative logarithm of this quantity, albeit up to an additive constant e E
U (w rrain, val) log p
(rrrin vall (a) log plog)(()))))))))))))))()) 0, +c,
for which we derive the forms of each term below. To begin, note that the negative log- aoglikelihood for Drain, val is
imn sall
log pl (rrain, vall (A) =
log p(ulkz w())
D Den vall
log(1 + ezi wlth)
you w())
Now, recall that the prior for the weights (kk given A and og is chosen isotropic Gaussian,
((w 0g) Nu 1d, ld.))
Thus have that
1 (k) 0. 0-10,11111 11.11
10,
(2TTa))d.. exp
20
and obtain
- logg (00 (0,00)) = log
20 (w(() 1a, 11.1
(2a0g)dal12 exp
de
log( (2TTO)
20 (A)
1d, 11
2
de
log (TTO +
1
| (k)
1a,
2
20
Since the first term does not depend either (k) or 0, can be considered an additive constant for U, and is thus absorbed into the constant c. Indeed, the relevant part of this negative log-prior is
20% ||w (k) - 0- 12.
Lastly, combining the relevant terms, ignoring the constant e, we arrive at
[Pimi real
(z)(t), 0 ( vat) Σ
[oo((1 + e²d (eT)(t)) - y12, w(]]
+
1
(20)
0-1d.
11,
203
as desired.
D.1.2 Derivation of the Hessian upper bound
Here, we outline the derivation of Hround. the of the single-particle energy function Uww, ((Drnin.val) with respect to the regression weights w. To be clear, the purpose of Hoond is to estimate the global Lipschitz constant I. of the gradient V..U. This constant I., taken as largest eigenvalue Hhound, subsequently informs the calculation of baseline step-size healer used in constructing grids for tuning the particle update step-size.
upper bound for the Hessian matrix
31

---

Recall, as seen above, single-partion energy function is given by:
DDmin wel
(w, 0 D vrain.val) =
log 1 + ex'l yr
w
0 1 de
20
Negative Log. Likelihood (NLL.)
Negutive Log-Prior (NLP)
We are then interested in the Hessian matrix UW D(mmin), which we decompose as 2 (NLL) + V (NLP).
V
Since NLP(w, 0) V,, (. 0) of (w 0 la,): and so this means that the Hessian of the NLP term, with respect to w, is thus constant
20g 100 1d, 11 we observe the gradient, with respect to w, be
V (NLP)
Id. o
The NLL term, on the other hand, gives a Hessian of the form
[Dminoal
V (NLL) = Σ (x w) - oxx w))x,,
which can notably expressed in matrix form as X ((w) uuminvatt where Xurin [Dimin,val x dz matrix, and (ww) is (Drriinvall [Dirainvall diagonal matrix, with its i-th diagonal ele- ment being (w) (z w)(1 o or (z w)). Crucially, (z) is bounded, and specifically
maximised at O, where (0) 0.5, yielding maximum value of 0.5 x (1 - 0.5) 0.25.
Therefore, for all ny and w,
w1 oxx w)) S 0.25, 0 [
implying that each (diagonal) element dw((w) < 0.25, so that the matrix D(w) A 00.25/ / the Loewner order).
/minsull (in
As result, an upper bound for the NLL Hessian is given by
V Xirain.val ((w) Xirain.val 0.25 train |eanoul Xtrain,vl 0.25 vaan X Xaaingal.
Combining the Hessian for the NLP term with the upper bound for the NLL Hessian, we notably obtain a weight- ceight-independent upper bound for the total Hessian (w,
(0, D|Puainvat),
1 o
V
3 X- X Xirain,val + Id.:
4 trainvval
and define the constant matrix Hhound this quantity. To be clear, have
1 XT
Hbound
iniii S Xurain.vll
Idz
(37)
4
06
The largest eigenvalue of this positive semi-definite matrix Hbound provides the global Lipschitz constant L used to determine healer-
D.2 Error model selection Bayesian regression
D.2.1 Derivation the covariance matrix for Gaussian marginal likelihood
For Mg. we have the Bayesian linear regression model = Xx + E, with w ~ N( op la.),
~ (O a3 la,) and w and e independent.
Now, the mean of y is Ely] = E[Xu + = XE(w] + [[e] 0, while the covariance matrix
En Cov(y) is given by
Σo E - (V - Elul 1 Elyy J]
32

---

Expanding yv gives
yy = (Xw+ (XX () = Xww> T + Xwee + ew X Tx + =E",
which gives, upon taking the expectation, and using El 1] El E(ee due to independence and zero means,
Eg wrxw ] + Elss" ], f | + El J.
Since Flww ]] Cov(w) E(w F(w]] we find
0g t ld, and Elee Cov(e) + E(s] -IT
Eg = ogla, + a xxx,,
where 2g is the covariance of under the model, p(y)X,,111111 oū. 0o) = NN/l0, Co), as desired.
D.2.2 MI-II estimation for Ma
Type-II Maximum Likelihood (ML-II estimation determines hyperparameters by maximising the marginal likelihood. Specifically, we achieve this by minimising the negative log-marginal likelihood with respect to the logarithmic parameterisations, Φ1 log o² and D2 log
The log marginal likelihood for Mg
L( )2) logp(y)(Xe = c aa= e
dy 2
1 log(det(<>))
1 - -リ -
log(2T)
2
where E = eth Id, +e d2 XXT.
The ML-II estimates for (Ø1, are found by solving arg ming, .o -( 22) This minimisation performed numerically, typically using gradient-basee optimisation algorithms such as L-BFGS-B,
which we choose to utilise.
33

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We claim that our method performs comparably to the state the art, which demonstrate in experimental section. Furthermore, we also demonstrate the ability to estimate marginal likelihoods which is a special feature of our method.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Motivations of the assumptions are discussed throughout the main and Limitation paragraph is included in the Conclusions. We discuss both theoretical and practical limitations.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
34

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We provide a proof of every theoretical result presented the paper. For the sake of brevity, complete proof' are in Appendix A.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: The experimental setups are detailed in Appendix and in the Experimental Result section with all the hyperparameters choices.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
35

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We provide the code via .iip file, which will also be made publically available.
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurlPS code and data submission guidelines (https://nips. cc/ Seneddee//ooSSummisseee
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/ /nips.cc.. wttpcc/pccppppccceeeeeoooccpcpce gtt////dde/////// (CodeSubmissinn wwwissssiioooooiieeeee (eionnolccy)) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All the details the experiments are in the Experimental Results section and in Appendix B, with detailed description of the setup and hyperparameters, well as in the released code.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We provide in Appendix B results about statistical significance of the experi- ments, i.e. error bar on multiple runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
36

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The experiments were run on personal computer and Google Colab. We specify it in the Introduction.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't i make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics btttp::/neuuriie http://uurrpppccppuuubtt cccpublic/Etleeee eshiidieuell
Answer: [Yes]
Justification: We followed the NeurIPS code of ethics and confirm that paper is in full agreement with the code.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer:
[Yes]
We discuss the broader impact of our work the Conclusions in a specific paragraph.
Justification:
Guidelines:
37

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymentt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: They coincide with the authors. Pre-existent benchmark models that reimplemented are credited in the Related Works and the Experimental Results sections.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
38

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The released is code not an asset, but just due for experiment reproducibility.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
39

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: httppp://neurrppe for what should or should not be described.
40


Reviews:
[{"review_id": "BbKG4RI8rO", "reviewer": "Reviewer_Ygyo", "summary": "This paper contributes to the literature of parameter inference in latent variable models using their marginal likelihood function. The key challenge in such parameter estimation problems is the inability to analytically obtain the marginal likelihood function by integrating out the latent variables. Correspondingly, much effort has been devoted to finding efficient ways to estimate the gradients of the marginal likelihood function so that parameter estimation can be done. This paper proposes a novel methodology for doing so using a combination of unadjusted Langevin algorithm steps to drive a set of particles, and, crucially, compute a set of weights, that are used to estimate the gradient together with a bias correction and allow its use in conjunction with standard optimization algorithms. The authors provide some non-asymptotic results on the performance of method, giving some quantitative bounds on how quickly the sequence of iterates approaches the true value for a given sample size. Some - now standard for this literature - simple experiments are considered to demonstrate the efficacy of the methodology. Of interest in Bayesian statistics is estimation of the marginal likelihood function itself, and the proposed method is able to do that as well.", "questions": "Is Carbone's correction the only one that can be used naturally with ULA like this? Are there any other possible bias correction schemes for ULA? Does it rely on a particular discretization scheme?\n\nGiven that these are particle-based estimates of marginal likelihoods, what are their variance properties like? \n\nI am very interested in unpacking Remark 1 a bit more here... how does the backward kernel connect to bias correction? Can you say a little more about this point please.", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: the algorithm is certainly elegant and demonstrates how ULA samples can be directly bias-corrected. It is easy to implement (in fact, I have personally tried to do so by adjusting some existing logistic regression code and succeeded on my part with a little work). The paper provides a strong and practical answer to a question which is very natural for this literature. The writing is very clear, and the intuition behind the method is supported by theoretical results.\n\nWeaknesses: The Bayesian neural network example is not very convincing and feels overly simplified. It would have been nice to see the method pushed a bit further in terms of what it can do in practice. For example, to do it on a non-binary classification task. And to report the misclassification rate for the task. Also in the main paper, the authors should be more clear in describing what exactly the latent variables are in the logistic regression model (this is said in the Appendix, but to a novice reader this may not be obvious).", "was_revised": true, "final_justification": "All of my questions have been satisfactorily clarified. Thank you to the authors!", "rebuttal": "> Weaknesses: The Bayesian neural network example is not very convincing and feels overly simplified. It would have been nice to see the method pushed a bit further in terms of what it can do in practice. For example, to do it on a non-binary classification task. And to report the misclassification rate for the task.\n\nWe agree that a non-binary classification task would be interesting and have thus extended the MNIST binary classification task described to a multi-class task, where specifically we classify 2500 subsampled images corresponding to the digits 2, 4, 7, and 9 (rather than just 4 and 9), for a variety of BNN model capacities. To be clear, the original BNN has 40 hidden neurons, and thus for the 4-way classification task has 31,520 latent weights. For $N=50$, $K=500$, and again a common global step-size, we obtain average final misclassification error rates (and standard deviations) of 4.74% (0.34%), 9.24% (0.93%), and 9.20% (1.27%) for PGD, JALA-EM, and SOUL respectively. As alluded to, one can vary the number of hidden neurons, and to this end we repeated the experiment for 512 hidden neurons, resulting in a BNN of 403,456 latent weights, and average final misclassification error rates (and standard deviations) of 4.98% (0.42%), 7.20% (0.81%), and 8.02% (1.02%) for PGD, JALA-EM, and SOUL respectively.\n\nAdditionally, we have implemented another Bayesian model selection example, specifically for polynomial regression. Rather than comparing two models, $\\mathcal{M}\\_{G}$ and $\\mathcal{M}\\_{T}$, we consider a series of nested models, $\\mathcal{M}\\_{1}, \\dots,\\mathcal{M}\\_{10}$, where $\\mathcal{M}\\_{p}$ denotes a polynomial of order $p$. We then use estimated marginal likelihoods to select the one best supported by the evidence. \n\nTo robustly evaluate performance we iterate through a range of true orders $p\\_{true} \\in [2,8]$, and run $100$ trials with different random seeds for each order. For each trial, we randomly sample $\\mathbf{w}\\_{true} \\in \\mathbb{R}^{p\\_{true}+1}$, which is then multiplied by the respective expanded polynomial basis matrix $X$, for $M=500$ datapoints (that we randomly sample in $[-2.5, 2.5]$), and is then subsequently perturbed by random Gaussian noise, determined by $\\sigma=7.5$. For each generated (noisy) dataset, the task for the algorithm is to select the best model from a candidate set of polynomial degrees $p \\in [1, 10]$. \n\nFor our algorithm, we set $K = 200$ and $N=50$, selecting the model with the highest final estimated marginal likelohood. Note that the parameters estimated here are $\\phi\\_{1} = \\log(\\sigma^{2})$ and $\\phi\\_{2} = \\log(\\alpha)$, where $\\alpha$ is the precision of a zero-mean Gaussian prior placed on the regression weights, and we initialise both values at 0. As a baseline, we compare our method to an approach in which the model weights are estimated via OLS and the selected model has the lowest Bayesian Information Criterion (BIC).\n\nPerformance is evaluated using the Mean Absolute Error (MAE) to provide a nuanced measure, rewarding close estimates when not perfectly accurate, which is not interpretable from just the Accuracy alone, for example. The MAE across the 100 repeats is summarised below:\n\n|True Order|JALA-EM (MAE)|OLS & BIC (MAE)|\n|-|:-:|:-:|\n|2|0.41|0.35|\n|3|0.30|0.34|\n|4|0.33|0.30|\n|5|0.35|0.35|\n|6|0.24|0.23|\n|7|0.12|0.15|\n|8|0.16|0.20|\n\n​The code for both of these new examples will be made available in the future in the public repository.\n\n\n\n>Also in the main paper, the authors should be more clear in describing what exactly the latent variables are in the logistic regression model (this is said in the Appendix, but to a novice reader this may not be obvious).\n> \n\nMany thanks for this suggestion - we have now clarified that the regression weights, $w \\in \\mathbb{R}^{d\\_{x}=9}$, are the latent variables for this example. On reflection, although implicit through our Appendix, we agree it is better to be explicit about this in the main text.  \n\n\n> Is Carbone's correction the only one that can be used naturally with ULA like this? Are there any other possible bias correction schemes for ULA? Does it rely on a particular discretization scheme?\n> \n\nAnother example of such a correction is MALA. However, a key limitation is that particles evolve independently, without access to global information that could be used for resampling. As a result, when the energy landscape undergoes a phase transition—for example, from unimodal to bimodal—MALA may yield biased results if the proportion of particles in each mode does not match the correct distribution.\n\nThe correction we propose is more general and can be applied to other Markov kernels, as illustrated in [1]. It does not depend on a specific discretization scheme. Instead, one simply needs to adapt the definition of the weights to the particular scheme chosen for the discretized SDE.\n\n> Given that these are particle-based estimates of marginal likelihoods, what are their variance properties like?\n\nWe provide a discussion of the error bound in (10), which is generally nontrivial in the SMC setting. Regarding the variance, it is true that in general there are limited theoretical guarantees—a well-known challenge in the Sequential Monte Carlo literature. As we discussed in the text, variance can be controlled empirically through appropriate resampling strategies. While this control is typically based on the empirical Effective Sample Size (ESS) rather than the population ESS, this is precisely the framework in which practitioners operate in real applications. \n\n> I am very interested in unpacking Remark 1 a bit more here... how does the backward kernel connect to bias correction? Can you say a little more about this point please.\n> \n\nAs discussed in Del Moral’s work, it is known that the variance of the importance weights is minimized when the backward kernel coincides with the true inverse of the forward kernel. However, in practice, such an inverse is rarely available in closed form. This has motivated a line of research aimed at learning or approximating the backward kernel to reduce variance, as explored for instance in [2, 3]. In our setting, we chose the backward kernel to match the forward dynamics, which is a standard and computationally efficient choice. However, we stress that this does not correspond to the inverse kernel, since the forward dynamics (e.g., ULA) is not symmetric and does not satisfy detailed balance. Consequently, the resulting importance weights are essential to correct for this asymmetry and ensure consistency. This point is central in recent contributions such as [1] we mentioned above, where it is shown that properly accounting for the lack of reversibility via the weights can significantly improve accuracy across a wide class of non-symmetric kernels. The correction applies independently of the choice of forward discretization scheme, as long as the associated weights are correctly computed. This flexibility is a strength of the method, allowing it to accommodate a variety of sampling kernels.\n\n[1] Carbone, D. (2025). Jarzynski Reweighting and Sampling Dynamics for Training Energy-Based Models: Theoretical Analysis of Different Transition Kernels. arXiv preprint arXiv:2506.07843.\n\n[2] Doucet, A., Moulines, E., & Thin, A. (2023). Differentiable samplers for deep latent variable models. Philosophical Transactions of the Royal Society A, 381(2247).\n\n[3] Kim, K., Xu, Z., Gardner, J. R., & Campbell, T. (2025) Tuning Sequential Monte Carlo Samplers via Greedy Incremental Divergence Minimization. In Forty-second International Conference on Machine Learning (ICML)."}, {"review_id": "PtIdF4fLYO", "reviewer": "Reviewer_un9b", "summary": "A special case of Del Moral et al's SMC framework applied to Type 2 MLE optimization and similar problems, where the forward kernel is ULA, the backward kernel is the standard choice for deterministic involutions, and the sequence of distribution is given by the top level parameters optimization path.", "questions": "An in any Z estimator, we need to know the true value of Z_0, e.g. this is shown in (4). In traditional SMC Samplers, we choose pi_0 so that it can have a known normalization constant (e.g. the prior, or a variational distribution). However in the context of the paper, how is Z_0 estimated (especially in the high dimensional case, where simple IS would fail)? This is an important thing to clarify as the paper highlights \"JALA-EM’s unique and dual capability to concurrently estimate both parameters and the marginal likelihood\".", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "A strength of the manuscript is the non-asymptotic convergence analysis. As the authors correctly point out, the fact that the gradient estimation is biased makes the analysis less routine than typical stochastic gradient theoretical analysis.\n\nThe paper is well written and explores interesting ideas. I enjoyed reading it. I hope you find these comments below helpful...\n\nAs my above summary description emphasizes, the paper is firmly within the Del Moral 2006 / Neal 2001 framework of SMC Samplers/AIS with backward kernels, something that the authors do explain in remark 1. As a result I feel introducing a new terminology for the methodology \"Jarzynski-adjusted Langevin algorithm\" is just going to add more noise in the literature. In fact, that combination of unadjusted moves with AIS/SMC Sampler keeps being re-invented and renamed. One example of re-invention in the ML literature is \"Markov Chain Monte Carlo and Variational Inference: Bridging the Gap\" by Salimans, Kingma, Welling 2015 (another great example of sketchy citation practice: that 2015 paper cites other works of Neal but not a word of the AIS paper it is re-inventing...). \n\nJust to be clear, I find that combination of AIS/SMC with unadjusted kernels applied to Type 2 MLE interesting, and it becomes more novel when it is combined with the optimization of top level parameters. But I strongly feel the new name is a disservice to the community! An easy change, but one I would feel strongly about.\n\nHere is a specific example why it would be better to rely more on the original literature: let us look at (4) in Proposition 1. The estimator for the normalization constant shown there corresponds to the AIS case, i.e. no resampling. Yet the authors do use resampling later on (as they should!), but in 5.2 they appear to still use the AIS estimator. If instead of relying on a re-invention of the wheel (here, citing Carbone et al 2024), they used the original, i.e. (15) in Del Moral 2006, they would obtain a much improved estimator.\n\nAnother point that I think should be changed is to present the algorithm as taking OPT to be a black box algorithm. Here the problem is that if OPT goes too far in the parameter space, the distributions pi_k and pi_{k+1} will be very far from each other, which in turn leads to weight degeneracy.   The correct formalism would be to take OPT to be a trust-region optimizer. This is related to the Adaptive SMC Samplers literature where an annealing parameter is tuned adaptively (see e.g. Zhou Johansen Aston 2015). Here the fact that a single step of SGD creates an implicit notion of \"not going too far in one step\". \n\n\"as exact maximisation is intractable\" -> that's a bit imprecise... often while the overall objective is non-convex, the inner M-step can still be convex. Often when an optimizer within an M step is used, it is because the optimizer does not have close form solution. \n\nIn practice, adaptive step sizes are important for large scale deployment of SMC samplers, here in two ways. First, in terms of controlling the distance between the successive distribution (this is related to my earlier comment on trust region methods), it is typical that a non-constant step size is needed to avoid degeneracy (see again Zhou 2015). Second, more specifically to this manuscript, there is the step size in the state space x. For the latter, a good reference in the context of unadjusted algorithms with SMC/AIS sampler is Kim, Xu, Gardner, Campbell, 2025 and references therein. \n\nThe experiments are interesting but jump straight from toy models to a black box NN example. I wish there would have been something in between, e.g. a spatial model or a scientific application. \n\nOne obvious numerical comparison that I would have found interesting is to compare the unadjusted SMC with the adjusted SMC one (i.e., the latter being the traditional choice where the backward kernel is the time reversal), while keeping as much as possible the same (except for the step size which would need to be different for the two).", "was_revised": true, "final_justification": "Thanks for the nice discussion! No major adjustment needed in my score.", "rebuttal": "We thank the reviewer un9b for their thoughtful and informative review.\n\n**Name of the method/algorithm**: We would like to address first the concern of the reviewer saying (we do not copy paste it all due to the length limit):\n\n> I strongly feel the new name is a disservice to the community! An easy change, but one I would feel strongly about.\n\nWe fully agree with the reviewer's message about not adding noise in terms of inventing new algorithm names. We also agree with the reviewer's assessment here that our method is a typical SMC algorithm with unadjusted Langevin kernels/proposals. In this respect, one possible 'name' for our method could be \"SMC with ULA proposals\". This would however become quickly untenable in naming also our parameter estimation method.\n\nOur perspective adopting our acronym is to take a dual perspective on the method. While we fully agree the reviewer's perspective on seeing the method as \"SMC with ULA proposals\", we take the MCMC perspective here on seeing it as an \"adjustment\" of ULA - much like Metropolis adjustment. This adjustment takes the form of running identical chains and correcting them using weights, based on Jarzynski identity. Given the roots of the ideas in the statistical physics literature, and also recognising Jarzynski's contributions, we named this special adjusted ULA variant as JALA. Our intention is definitely not to add the noise - but to come up with a specific name to this variant - which has been not specifically named in the literature.\n\nFinally, to explain our perspective, we would like to give an example from MCMC literature. For example, Metropolis-Hastings is a general method. However, a special instance of this method (MH with Langevin proposals) is termed MALA in the literature. We took here a similar approach, we hope that the reviewer sees this as a recognition of this special subclass of methods. We have been also very careful in Section 2 to summarize the relations to Del Moral (2006). We are open to discussion in the rebuttal phase on this point. We will definitely also add the citation [1] the reviewer requested in another comment as this seems very relevant.\n\n> [...] let us look at (4) in Proposition 1. The estimator for the normalization constant shown there corresponds to the AIS case, i.e. no resampling. Yet the authors do use resampling later on (as they should!), but in 5.2 they appear to still use the AIS estimator. If instead of relying on a re-invention of the wheel (here, citing Carbone et al 2024), they used the original, i.e. (15) in Del Moral 2006, they would obtain a much improved estimator.\n> \n\nApologies for any misunderstanding here. We will indeed fix the formula in our Section 5.2, as we believe this is already identical to original formula as pointed out by the reviewer in the implementation. That is, after resampling, we reset the weights as in Del Moral (2006).\n\n> Another point that I think should be changed is to present the algorithm as taking OPT to be a black box algorithm. Here the problem is that if OPT goes too far in the parameter space, the distributions pi_k and pi_{k+1} will be very far from each other, which in turn leads to weight degeneracy. The correct formalism would be to take OPT to be a trust-region optimizer. This is related to the Adaptive SMC Samplers literature where an annealing parameter is tuned adaptively (see e.g. Zhou Johansen Aston 2015). Here the fact that a single step of SGD creates an implicit notion of \"not going too far in one step\".\n\nThank you and we fully agree with this excellent observation as we have seen this in practice. Our optimisers need to take 'small enough' steps for distributions to stay close and remain easy to sample. We will write a remark on this point on the choice of OPT and its impact on sampling performance and will definitely mention the necessity of adaptive methods and will cite the relevant references in our remark.\n\n> \"as exact maximisation is intractable\" -> that's a bit imprecise... \n\nWe agree with this comment. We have reworded this part, now reads:\n\n*as exact maximisation may not be analytically available.*\n\n> *Need for adaptive step sizes*: In practice, adaptive step sizes are important for large scale deployment of SMC samplers, here in two ways. First, in terms of controlling the distance between the successive distribution (this is related to my earlier comment on trust region methods), it is typical that a non-constant step size is needed to avoid degeneracy (see again Zhou 2015). Second, more specifically to this manuscript, there is the step size in the state space x. For the latter, a good reference in the context of unadjusted algorithms with SMC/AIS sampler is Kim, Xu, Gardner, Campbell, 2025 and references therein.\n> \n\nWe agree with this comment, and as we noted above, we will add a discussion on the need of adaptive step-sizes in our setting.\n\n> The experiments are interesting but jump straight from toy models to a black box NN example. I wish there would have been something in between, e.g. a spatial model or a scientific application.\n\nThank you for this feedback - to bridge this gap we have implemented a new experiment, namely that of Bayesian model selection in a polynomial regression setting. The idea is that this serves as an intermediate step in complexity between the initial logistic regression example and the BNN example. Rather than comparing two models, $\\mathcal{M}\\_{G}$ and $\\mathcal{M}\\_{T}$, we consider a series of nested models, $\\mathcal{M}\\_{1}, \\dots,\\mathcal{M}\\_{10}$, where $\\mathcal{M}\\_{p}$ denotes a polynomial of order $p$. We then use estimated marginal likelihoods to select the one best supported by the evidence. \n\nTo robustly evaluate performance we iterate through a range of true orders $p\\_{true} \\in [2,8]$, and run $100$ trials with different random seeds for each order. For each trial, we randomly sample $\\boldsymbol{w}\\_{true} \\in \\mathbb{R}^{p\\_{true}+1}$, which is then multiplied by the respective expanded polynomial basis matrix $X$, for $M=500$ datapoints (that we randomly sample in $[-2.5, 2.5]$), and is then subsequently perturbed by random Gaussian noise, determined by $\\sigma=7.5$. For each generated (noisy) dataset, the task for the algorithm is to select the best model from a candidate set of polynomial degrees $p \\in [1, 10]$. \n\nFor our algorithm, we set $K = 200$ and $N=50$, selecting the model with the highest final estimated marginal likelohood. Note that the parameters estimated here are $\\phi\\_{1} = \\log(\\sigma^{2})$ and $\\phi\\_{2} = \\log(\\alpha)$, where $\\alpha$ is the precision of a zero-mean Gaussian prior placed on the regression weights, and we initialise both values at 0. As a strong baseline, we compare our method to an approach in which the model weights estimated via OLS and the selected model has the lowest Bayesian Information Criterion (BIC).\n\nPerformance is evaluated using the Mean Absolute Error (MAE) to provide a nuanced measure, rewarding close estimates when not perfectly accurate. The MAE across the 100 repeats is summarised below:\n\n|True Order|JALA-EM (MAE)|BIC & OLS (MAE)|\n|-|:-:|:-:|\n|2|0.41|0.35|\n|3|0.30|0.34|\n|4|0.33|0.30|\n|5|0.35|0.35|\n|6|0.24|0.23|\n|7|0.12|0.15|\n|8|0.16|0.20|\n\n​The code for this new example will be made available in the future in the public repository.\n\n> One obvious numerical comparison that I would have found interesting is to compare the unadjusted SMC with the adjusted SMC one (i.e., the latter being the traditional choice where the backward kernel is the time reversal), while keeping as much as possible the same (except for the step size which would need to be different for the two).\n\nWe agree that this is indeed a very interesting direction. There are multiple ways to implement these methods, e.g. a few options are summarised in [1]. While we could not provide this comparison now, we think that (with small enough step-sizes) these methods would perform parameter estimation similarly. We reference an interesting work that supports this view where some alternative approaches are characterised and assessed in the context of MCMC methods [2].\n\n> An in any Z estimator, we need to know the true value of Z_0, e.g. this is shown in (4). In traditional SMC Samplers, we choose pi_0 so that it can have a known normalization constant (e.g. the prior, or a variational distribution). However in the context of the paper, how is Z_0 estimated (especially in the high dimensional case, where simple IS would fail)? This is an important thing to clarify as the paper highlights \"JALA-EM’s unique and dual capability to concurrently estimate both parameters and the marginal likelihood\".\n> \n\nIndeed, the estimation of $Z\\_0$ is a crucial aspect in any Z-estimator framework. In our setting, where $\\pi\\_0$ is not explicitly normalized, we therefore used importance sampling to estimate $Z\\_0$. However we agree that for higher-dimensional models, this presents a genuine challenge where simple importance sampling becomes unreliable. One option to solve this issue is to apply a dedicated SMC/AIS procedure to bridge the initial posterior distribution with a tractable reference distribution and estimate its normalising constant. This can easily be integrated then into our approach and be updated.\n\n[1] Kim, K., Xu, Z., Gardner, J. R., & Campbell, T. Tuning Sequential Monte Carlo Samplers via Greedy Incremental Divergence Minimization. In Forty-second International Conference on Machine Learning (ICML) 2025.\n\n[2] Schönle, Christoph, et al. \"Sampling metastable systems using collective variables and Jarzynski–Crooks paths.\" Journal of Computational Physics, 2025."}, {"review_id": "jBSXHyecZv", "reviewer": "Reviewer_j23q", "summary": "This paper makes use of recent advances in the application of the Jarzynski relation and sequential monte carlo to estimate expectations under target distributions with densities known up to normalization. Here, it is used to perform an estimation of an expectation necessary in an expectation-maximization algorithm for latent variable models, e.g. to perform maximum marginal likelihood estimation. \n\nThe authors make use of the Jarzynski relation to sample under a potential $U_k$ as specified in a previous work [1] by keeping track of importance weights necessary to adapt a local unadjusted langevin dynamics (ULA). The target expectation here is one emerging in the optimization of latent variable models, i.e. to find the maximizers of the marginal likelihood $p_\\theta (y)$ of the joint model $p_\\theta(x,y)$. \n\nThey test their proposed method on bayesian logistic regression and bayesian neural network tasks and show that it is robust and comparable to existing approaches while offering the ability to simultaneously estimate a number of quantities.\n\n\n\n\n\n\n\n\n\n[1] Carbone et al, \"Efficient Training of Energy-Based Models Using Jarzynski Equality\", 2023.", "questions": "- In the notation, I may be missing something. Can the authors clarify why the dynamics in 6 gives access to the expectation (through the use of the weights) of the *conditional density p(x|y)*? Or are you just using Bayes rule and not saying it somewhere? I ask because $U(\\theta, x)$ is defined as $- \\log p_\\theta(x,y)$ and not $-\\log p_\\theta(x|y)$. \n\n- Can the authors more clearly delineate why their approach gives access to certain quantities that the alternative approaches do not? This is mentioned throughout, but it is not clear to me why all the other methods are not amenable to any sort of online likelihood estimation / partition function estimation. Is this truly not available in other approaches? This would help me understand the contribution here a bit better. It does not need be some unique miracle to be a valuable contribution, I just want to clarify that fulfilling a variety of these desiderata in one method is not amenable to the alternative approaches.\n\n- Could you provide a few more experimental realizations? This is really lacking in terms of differentiating with other approaches. \n\n- Could you comment on the sensitivity of the proposed method to the resampling procedure and the frequency of resampling?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "*Strengths*\n\n- The paper does a good job introducing the method, with mostly clear writing, and provides a nice theoretical analysis of the implications of how the proposed algorithm must be implemented (e.g. how to control biases, etc.)\n- Numerical experiments validate the robustness of the approach, and show that it is competitive with existing approaches, while also providing potentially more estimable quantities simultaneously (I want the authors to comment more about this distinguishing factor). \n\n\n*Weaknesses*\n\n- Methodologically, the jarzynski formulation was already completely typified as it is used here in [1]. I guess the authors see the key insight as reapplying it in a new domain, i.e. to estimate a different expectation on which to perform gradient descent. This is a bit lacking.\n- One could argue that applying it in a new domain (which definitely requires new insights!) is sufficient to warrant support of the work, but it is hard to tell of the value of the method given the limited numerical experiments as it pertains to existing approaches. The only comparisons done to other methods is on the Bayesian neural network example, and the authors write that the performance of JALA-EM is comparable to that of SOUL, and slower than PGD. The authors do note that their method provides greater access to certain quantities of interest, like an online estimate of the partition function, but they should make clearer why this is unique to their approach (perhaps with a table of characteristics of the various methods). \n\n\n\n\n[1] Carbone et al, \"Efficient Training of Energy-Based Models Using Jarzynski Equality\", 2023.", "was_revised": true, "final_justification": "The clarification of new experimentation and the novelty of the asymptotic analysis is useful. I recommend an accept.", "rebuttal": "We thank the reviewer for their thoughtful and informative review.\n\n> Weaknesses: Methodologically, the jarzynski formulation was already completely typified as it is used here in [1]. I guess the authors see the key insight as reapplying it in a new domain, i.e. to estimate a different expectation on which to perform gradient descent. This is a bit lacking.\n\nWe thank the reviewer for pointing this out. We agree that the formulation we employ builds upon existing work, notably [1]. Our methodological contributions are two-fold:\n\n- Adapting the methodology in [1] to a statistical problem (MMLE) as pointed out by the reviewer - expanding its scope to a whole class of statistical models.\n- Providing a convergence analysis for this scheme, connecting the literature on biased gradient descent and IS-type estimates of gradients of marginal likelihoods.\n\nThis second aspect is also crucial to our contribution as the theoretical analysis of $\\theta$ iterates was left open in [1]. We will emphasise this a bit more in our updated manuscript.\n\n> Weaknesses: One could argue that applying it in a new domain (which definitely requires new insights!) is sufficient to warrant support of the work, but it is hard to tell of the value of the method given the limited numerical experiments as it pertains to existing approaches. The only comparisons done to other methods is on the Bayesian neural network example, and the authors write that the performance of JALA-EM is comparable to that of SOUL, and slower than PGD. The authors do note that their method provides greater access to certain quantities of interest, like an online estimate of the partition function, but they should make clearer why this is unique to their approach (perhaps with a table of characteristics of the various methods).\n>\n\nWe thank the reviewer for these suggestions. First, we'd like to note further experiments are conducted (see our response below). To answer the last request, we have prepared a table to summarize when/how JALA-EM provides unique properties.\n\n|Feature|PGD|SOUL|JALA-EM|\n|-|:-:|:-:|:-:|\n|Marginal likelihood estimation|✖|✖|✔|\n|Different $\\theta$ and $x$ step-sizes| ✖| ✔| ✔|\n|Theoretical guarantees with fixed $\\gamma$| ✔| ✖|✔|\nEasy adaptation of theory to alternative optimizers|✖| ✖|✔|\n\nWe will add this table to our paper. To expand the last point, we would like to emphasize another nontrivial aspect of our approach, that is, *easy adaptation of theory to alternative optimizers*. In practice, methods like PGD and SOUL are implemented using *adaptive optimizers*, e.g., see [2]. However, their theory can only account for vanilla gradient descent (for SOUL, theory only works for a particular (and non-trivial) decay of step-sizes). In our case, due to the SMC-based approach to gradient estimation, we can apply theory for *any optimizer* provided that guarantees for biased gradients exist. We can easily also do this with adaptive optimizers, adapting results from, e.g., [3]. We will expand these points in our updated manuscript.\n\n> In the notation, I may be missing something. Can the authors clarify why the dynamics in 6 gives access to the expectation (through the use of the weights) of the conditional density p(x|y)? Or are you just using Bayes rule and not saying it somewhere? I ask because $U(\\theta, x)$ is defined as $-\\log p_{\\theta}(x, y)$ and not $-\\log p_{\\theta}(x \\vert y)$.\n>\n\nWe are indeed implicitly using the standard fact from sampling theory that the posterior $p_\\theta(x \\mid y)$ is proportional to the joint density $p_\\theta(x, y)$, that is $-\\log p_\\theta(x \\mid y) = -\\log p_\\theta(x, y) + \\log p_\\theta(y)$, where the marginal $\\log p_\\theta(y)$ is independent of $x$ and thus irrelevant for the dynamics or the importance weights. Therefore, using the joint log-density $-\\log p_\\theta(x, y)$ as the energy function in the ULA dynamics is equivalent to targeting the unnormalized posterior. This is a standard approach in Langevin-based sampling, where only the gradient of the log-density is required up to an additive constant. We will make this point explicit in the revised version to avoid confusion.\n\n> Can the authors more clearly delineate why their approach gives access to certain quantities that the alternative approaches do not? This is mentioned throughout, but it is not clear to me why all the other methods are not amenable to any sort of online likelihood estimation / partition function estimation. Is this truly not available in other approaches? This would help me understand the contribution here a bit better.\n\nIn general, Markov chain Monte Carlo (MCMC) algorithms cannot readily provide estimates of marginal likelihoods (normalising constants). This is the main difficulty behind the approaches like SOUL to provide estimates of these quantities. There are, however, ways to provide estimates of the normalizing constant using MCMC output, see, e.g. [4]. However, these classical approaches are post-processing techniques -- means that unlike JALA-EM, they would require the whole chain for estimation -- rather than providing estimates online.\n\nPGD [2] approach simulates an SDE to minimise an energy functional and similarly they do not provide a way to estimate normalising constants. In fact, it is not clear how one can use their samples to obtain this estimate, unlike the MCMC case above.\n\n> Could you provide a few more experimental realizations? This is really lacking in terms of differentiating with other approaches.\n\nThank you for the suggestion - in an effort to alleviate your concern, we have extended the BNN example to a multi-class setting, where specifically we extended the MNIST binary classification task described to a multi-class task, using 2500 subsampled images corresponding to the digits 2, 4, 7, and 9, for a variety of BNN model capacities. The original BNN has 40 hidden neurons, and thus 31,520 latent weights. For $N=50$, $K=500$, and again a common global step-size, we obtain average final misclassification error rates (and stds) of 4.74% (0.34%), 9.24% (0.93%), and 9.20% (1.27%) for PGD, JALA-EM, and SOUL respectively. Indeed, one can vary the number of hidden neurons, and so we repeated the experiment for 512 hidden neurons, resulting in a BNN of 403,456 latent weights, and average final misclassification error rates (and stds) of 4.98% (0.42%), 7.20% (0.81%), and 8.02% (1.02%) for PGD, JALA-EM, and SOUL respectively.\n\nFor an additional Bayesian model selection example, we implemented a polynomial regression example. Rather than comparing two models, $\\mathcal{M}\\_{G}$ and $\\mathcal{M}\\_{T}$, we consider a series of nested models, $\\mathcal{M}\\_{1}, \\dots,\\mathcal{M}\\_{10}$, where $\\mathcal{M}\\_{p}$ denotes a polynomial of order $p$. We then use estimated marginal likelihoods to select the one best supported by the evidence. \n\nTo robustly evaluate performance we iterate through a range of true orders $p\\_{true} \\in [2,8]$, and run $100$ trials with different random seeds for each order. For each trial, we randomly sample $\\boldsymbol{w}\\_{true} \\in \\mathbb{R}^{p\\_{true}+1}$, which is then multiplied by the respective expanded polynomial basis matrix $X$, for $M=500$ datapoints, and then perturbed by random Gaussian noise ($\\sigma=7.5$). For each noisy dataset, the task is to select the best model from a candidate set of polynomial degrees $p \\in [1, 10]$. \n\nWe set $K = 200$ and $N=50$, selecting the model with the highest final estimated marginal likelihood. The parameters estimated here are $\\phi\\_{1} = \\log(\\sigma^{2})$ and $\\phi\\_{2} = \\log(\\alpha)$, where $\\alpha$ is the precision of a zero-mean Gaussian prior placed on the weights. As a baseline, we compare our method to an approach in which the model weights estimated via OLS and the selected model has the lowest Bayesian Information Criterion (BIC).\n\nPerformance is evaluated using the Mean Absolute Error (MAE) to provide a nuanced measure, rewarding close estimates when not perfectly accurate, which is not interpretable from just the Accuracy, for example. The MAE across the 100 repeats is summarised below:\n\n|True Order|JALA-EM (MAE)|OLS & BIC (MAE)|\n|-|:-:|:-:|\n|2|0.41|0.35|\n|3|0.30|0.34|\n|4|0.33|0.30|\n|5|0.35|0.35|\n|6|0.24|0.23|\n|7|0.12|0.15|\n|8|0.16|0.20|\n\nCode for both of these new examples will be made available in the future in the public repository.\n\n> Could you comment on the sensitivity of the proposed method to the resampling procedure and the frequency of resampling?\n>\n\nRegarding the resampling procedure, we did not explore alternatives beyond the standard systematic resampling method already used in [1], which is known to offer a good trade-off between computational efficiency and variance control. As discussed in that work, systematic resampling tends to be stable across a range of settings, and we believe that further refinements in this direction would yield marginal benefits compared to other factors, such as controlling the variance of the importance weights, which plays a more critical role in the performance of the method.\n\nConcerning the resampling frequency, we adopted a conservative strategy by setting the threshold parameter $C = 1/1.05$ to ensure a high empirical ESS at each step. While this choice helps maintain particle diversity, we agree that a more systematic study of the sensitivity of the algorithm with respect to C could be a valuable follow-up. \n\nWe are keen to discuss if there are any further questions.\n\n[1] Carbone, Davide, et al. \"Efficient training of energy-based models using jarzynski equality.\" (2023), *NeurIPS*.\n\n[2] Kuntz, Juan, et al. \"Particle algorithms for maximum likelihood training of latent variable models.\" (2023), *AISTATS*\n\n[3] Surendran, Sobihan, et al. \"Non-asymptotic analysis of biased adaptive stochastic approximation.\" (2024), *NeurIPS*\n\n[4] Chib, S., & Jeliazkov, I. (2001). Marginal likelihood from the Metropolis–Hastings output. Journal of the American statistical association, 96(453), 270-281."}, {"review_id": "C1HnSiTRMn", "reviewer": "Reviewer_QYFn", "summary": "The authors use jarzynksi-adjusted langevin sampling from nonequilibrium statistical mechanics to develop a weight-version of sequential monte carlo methods.", "questions": "How does the algorithm compare to algorithms such as Hamiltonian sampling for sampling speed? What is the difference in performance between using the Jarzynski sampling for time-varying probability distributions and sequential Hamiltonian sampling as in Septier and Peters, IEEE, 2015?", "limitations": "Well addressed in the paper.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths - The authors describe many ways in which the algorithm may be used in neural networks.  It is flexible and metrics are well established to describe how sampling distribution converges to marginal likelihood.  I thank the authors for their clear derivations and writing. I think this paper will be of interest to those who use Bayesian neural networks.  \n\nWeaknesses - The sampling speed is not clear to be. Generally, Kl Divergence over time is used to compare speed. As authors mentioned, it is computationally expensive due to the weights, so it is not clearly written when this method would be better than gradient descent that makes long jumps such as levy flights.", "was_revised": true, "final_justification": "The authors clearly addressed my concerns.  I appreciate the comparison to OLS and BIC in the rebuttal.  While it is still unclear how this algorithm performs in comparison to momentum-based approaches, I believe that is outside the scope of this manuscript and it, in and of itself, is a solid foundation for work in the future.", "rebuttal": "We thank the reviewer for their thoughtful and informative review.\n\n> Weaknesses: The sampling speed is not clear to be. Generally, Kl Divergence over time is used to compare speed. As authors mentioned, it is computationally expensive due to the weights, so it is not clearly written when this method would be better than gradient descent that makes long jumps such as levy flights.\n\nWe appreciate the reviewer’s observation. It is true that, in our method, computing and applying the importance weights introduces a computational overhead compared to existing algorithms, such as PGD and SOUL. However, we emphasize that the final iteration complexity of our method is the same as these two methods If we denote the cost of a single evaluation of the gradient of $U$ by $C_{\\nabla U}$, then we have that the  computational (iteration) cost of PGD, SOUL, and JALA is $\\mathcal{O}(N C_{\\nabla U})$. We appreciate this obsfucates practical speeds of the respective methods, and considerations related to parallelisation. As indicated by Figure 1 in our paper, our method is practically slower than PGD, due to multiple potential and gradient evaluations, and less parallelizable compared to PGD, due to the synchronization bottleneck introuced by the resampling step.\n\nIn any case, we note that our method enables recursive estimation of the marginal likelihood, hence its applicability to model selection, which recall the PGD method cannot facilitate.\n\nWith respect to gradient descent, we would like to note that our method is an instance of SGD when OPT is chosen as SGD. JALA-EM can make long jumps, depending on the variance of the gradient estimator. However, we remark that this can be beneficial or harmful depending on the model under consideration.\n\n> How does the algorithm compare to algorithms such as Hamiltonian sampling for sampling speed\n\nWe thank the reviewer for the suggestion. In this work, we did not include a benchmark against momentum-based samplers such as Hamiltonian Monte Carlo or underdamped Langevin dynamics, since our method is built on top of ULA-type dynamics, which are overdamped and do not involve auxiliary momentum variables. The theoretical framework can, in principle, be extended to underdamped schemes, and we agree that such an extension would be a promising direction for future research.\n\n\n> What is the difference in performance between using the Jarzynski sampling for time-varying probability distributions and sequential Hamiltonian sampling as in Septier and Peters, IEEE, 2015?\n>\n\nThank you for suggesting this reference. [Septier & Peters, IEEE, 2015] consider a filtering setting (as opposed to our setting where we are interested in parameter estimation in a static model). Therefore, it is not directly possible to provide a comparison between two approaches. That being said, the Hamiltonian proposals in their work can also be adapted to our setting. While they focused on Metropolis-adjusted proposals, our focus was on the Jarzynski-based correction framework to unadjusted overdamped dynamics, and in particular its use in the (static) latent variable estimation context. We will cite this work in our literature review.\n\n[Septier & Peters, IEEE, 2015] Septier, F., & Peters, G. W. (2015). Langevin and Hamiltonian based sequential MCMC for efficient Bayesian filtering in high-dimensional spaces. IEEE Journal of selected topics in signal processing, 10(2), 312-327.\n\n> Paper Formatting Concerns: As a minor edit, the references are distracting being a different color and seem like they take up a lot of space.\n>\n\nWe thank the reviewer for the observation. In the updated version, we will switch to another color to improve readability to reduce visual clutter.\n\nWe are happy to discuss or clarify if the reviewer has any more questions."}, {"review_id": "GbsxqStpyI", "reviewer": "Reviewer_iWiS", "summary": "In this paper, the authors introduce the Jarzynski-adjusted Langevin algorithm (JALA) as a sampling procedure for sequences of evolving probability measures. JALA can be viewed as a type of Expectation–Maximisation algorithm to perform maximum marginal likelihood estimation (MMLE) in latent variable models (LVMs).  Building on nonequilibrium statistical mechanics, JALA runs biased Unadjusted Langevin Algorithm (ULA) dynamics on a time‐varying potential and applies an exponential Jarzynski weight as a bias correction to the estimates of expectations under the target measures. This sampler is then embedded in an EM‐style routine—termed JALA-EM—in which a particle population evolves under ULA, weights are updated via Jarzynski factors, and parameter updates are performed using a generic first‐order optimiser, e.g. SGD, although other optimisers could be used.\n\nThe authors provide a nonasymptotic convergence analysis of JALA-EM under log‐concavity and Polyak–Łojasiewicz conditions, showing that with a suitable choice of step‐size and particle count the method converges to the MMLE with an explicit rate of the form\n$\\E\\bigl[V(\\theta_k)-V^\\star\\bigr] \\le \\delta_0 (1-\\tfrac{\\gamma\\mu}{2})^k + O\\bigl(\\tfrac{\\gamma}{N}\\bigr) + O\\bigl(\\tfrac{1}{N}\\bigr),$\nwhere $V(\\theta)=-\\log p_\\theta(y)$ and $N$ is the number of particles.\n\nEmpirically, JALA-EM is evaluated on (i) Bayesian logistic regression on the Wisconsin breast cancer dataset, demonstrating faster, robust convergence of both parameter estimates and posterior approximations compared to particle gradient descent (PGD) and SOUL; (ii) Bayesian linear regression model selection, where JALA-EM yields accurate marginal‐likelihood estimates on‐the‐fly and correctly identifies the data‐generating model; and (iii) a small Bayesian neural network on MNIST binary digit classification.", "questions": "1) Have the authors tested JALA-EM against SOUL, PGD or other related methods on higher-dimensional examples? And if so, how does JALA-EM compare?\n2) The authors do not comment on the sensitivity of tuning parameters such as $gamma$ and $h$. The theoretical results do not seem to provide a clear indication of how these parameters should be tuned in practice. What advice should be given to practitioners regarding step size selection so that JALA-EM can be easily used out-of-the-box? \n3)", "limitations": "Yes", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Quality\n\nThe paper’s technical quality is high, underpinned by a rigorous nonasymptotic convergence analysis of the JALA–EM algorithm under clearly defined, and reasonable, assumptions such as smoothness and the Polyak–Łojasiewicz condition. This theoretical treatment, uncommon in hybrid sequential Monte Carlo–EM work, provides practitioners with explicit error bounds and convergence rates, lending strong credibility to the method. Empirically, the authors validate their claims through comprehensive benchmarks—including Bayesian logistic regression, linear regression model selection, and a small neural network on MNIST—comparing against state-of-the-art baselines (PGD and SOUL) with well-tuned hyperparameters and wall-clock timing. A small weakness is that the analysis hinges on strong assumptions (e.g., bounded gradients and finite exponential moments of Jarzynski weights) that may not hold in high-dimensional or strongly multimodal settings, potentially limiting the method’s practical applicability. However, the authors do note that some assumptions could be relaxed, albeit leading to a more complex theoretical analysis. Moreover, the experimental scale remains quite modest—datasets of dimension up to nine or small Bayesian neural networks—so it remains unclear how JALA-EM will perform on large, complex latent-variable models. It could also be reasonably argued that the empirical study could be improved if the authors compared JALA-EM against other competing EM-type methods, i.e. not just SOUL and PGD.\n\nClarity\n\nOverall, the paper is well structured: the algorithmic description (Algorithm 1), theoretical propositions, and implementation details are presented in discrete, logical sections that make the paper clear to follow. Notation is consistent with standard conventions in both SMC and Langevin literature, which makes the paper easy to read for those familiar with this literature. That said, the work assumes significant prior familiarity with nonequilibrium thermodynamics (the Jarzynski equality) and advanced SMC formalisms, which could present a steep learning curve for those not already versed in both domains. For example, in equation (2), it isn't immediately clear why $A_k$ is needed (unless the reader is familiar with the Jarzynski method, which is not well-known in the sampling literature), whereas the ULA process on $X_k$ is clear as this is the latent variable to be estimated. \n\nSignificance\n\nA central strength of the proposed approach is its ability to recursively estimate marginal likelihoods “on the fly” during parameter learning, something most scalable EM-type and particle-gradient methods do not offer. This feature enables principled model selection within the same algorithmic run that fits parameters, representing a meaningful advance for practitioners who require both posterior inference and evidence computation. Furthermore, by introducing tools from nonequilibrium statistical mechanics into mainstream latent-variable inference, the paper bridges disciplines which should open new avenues of research between these two areas. However, the necessity of maintaining and updating Jarzynski weights for each particle introduces extra computational and memory overhead compared to vanilla ULA or particle-gradient schemes, potentially negating runtime benefits in some applications. The reliance on resampling to control weight degeneracy also raises concerns, which are well-known in the SMC literature. Selecting effective ESS thresholds or weight-clipping strategies may lead to non-trivial tuning and could introduce additional variance in practice.\n\nOriginality\n\nMethodologically, the use of Jarzynski–adjusted Langevin dynamics within an SMC ensemble to drive an EM algorithm for marginal‐likelihood estimation is genuinely novel and clearly extends prior SMC-EM and Jarzynski-based sampling work. The theoretical contribution—providing the first nonasymptotic convergence guarantees for an SMC–EM algorithm leveraging nonequilibrium corrections—is particularly noteworthy. Nonetheless, the innovation builds fairly directly on recent work by Carbone et al. (2024a,b) and existing SMC-EM variants, so some readers may view the advance as incremental rather than revolutionary. Additionally, the introduction of new hyperparameters (e.g., the Jarzynski step-size and ESS resampling threshold) adds further complexity to the tuning process, which could make the algorithm more challenging for practitioners to successfully implement in practice.", "was_revised": true, "final_justification": "There is a clear consensus among the reviewers to accept the paper. I'm happy with the quality of the paper and its appropriateness for the NeurIPS conference. The authors also addressed the concerns/questions that I raised.", "rebuttal": "We are thankful to the reviewer for their thoughtful and informative review.\n\n> it remains unclear how JALA-EM will perform on large, complex latent-variable models.\n\nTo address the point about complex LVMs (as also asked by other reviewers), we have extended our example on classification with BNNs beyond the binary class setting, where a clear example is one in which we consider MNIST samples corresponding to more than two distinct digits. For example, we classify 2500 subsampled images corresponding to the digits 2, 4, 7, and 9 (rather than just 4 and 9), for a variety of BNN model capacities. To be clear, the original BNN has 40 hidden neurons, and thus for the 4-way classification task has 31,520 latent variables. For $N=50$, $K=500$, and again a common global step-size, we obtain average final misclassification error rates (and standard deviations) of 4.74% (0.34%), 9.24% (0.93%), and 9.20% (1.27%) for PGD, JALA-EM, and SOUL respectively. As alluded to, one can vary the number of hidden neurons, and to this end we repeated the experiment for 512 hidden neurons, resulting in a BNN of 403,456 latent variables, and average final misclassification error rates (and standard deviations) of 4.98% (0.42%), 7.20% (0.81%), and 8.02% (1.02%) for PGD, JALA-EM, and SOUL respectively.\n\n> It could also be reasonably argued that the empirical study could be improved if the authors compared JALA-EM against other competing EM-type methods, i.e. not just SOUL and PGD.\n\nMany thanks for this comment. We specifically focus on cases where E and M steps both require numerical algorithms (e.g. sampling for E-step and optimisation for M-step). We believe SOUL and PGD are the canonical choices in this setting and best performing benchmarks. This is because these methods leverage Langevin-type sampling ideas to implement the sampling and gradient-based optimisers for M-step. Other benchmarks (e.g. using Metropolis-Hastings for E step or gradient-free optimisation for M-step) would be interesting to consider but would underperform given the complexity of the problems we consider here. We will mention these approaches in our work.\n\n> The work assumes significant prior familiarity with nonequilibrium thermodynamics (the Jarzynski equality) and advanced SMC formalisms, which could present a steep learning curve for those not already versed in both domains. For example, in equation (2), it isn't immediately clear why $A_{k}$ is needed (unless the reader is familiar with the Jarzynski method, which is not well-known in the sampling literature), whereas the ULA process on $X_{k}$ is clear as this is the latent variable to be estimated.\n>\n\nMany thanks for this comment. We agree the method as written in Prop. 1 deviates from the standard way of presenting it in other sequential Monte Carlo papers (where $A_k$ (log-weights) are not directly mentioned, but weights are). We will provide the following remark in our paper to make a connection to more standard SMC literature:\n\n**Remark.** *We note that, in connection to more standard SMC literature, the sequence $(A_k)_{k\\geq 0}$ denote *log-weights*. In other words, let us denote weights associated with our particles $\\mathsf{w}_k^i$ for $i = 1,\\ldots, N$. In our context, we denote $A_k = \\log \\mathsf{w}_k$ and $\\mathsf{w}_k = e^{A_k}$. Hence it can be seen that the formula in Prop. 1 and in Algorithm 1 implements a weighted average.*\n\n> Have the authors tested JALA-EM against SOUL, PGD or other related methods on higher-dimensional examples? And if so, how does JALA-EM compare?\n>\n\nThank you for mentioning this, as the performance of JALA-EM in higher-dimensional settings is important. First, note that the original BNN model used for the binary classification of 4 and 9 MNIST digits, in the paper, has $(28 \\times 28 \\times 40) + (40 \\times 2) = 31,440$ latent weights. In the multi-class setting, this is expanded to $(28 \\times 28 \\times 40) +(40 \\times 4) = 31,520$ latent weights, as mentioned above, whereas the larger model has $(28 \\times 28 \\times 512) +(512 \\times 4) = 403,456$ latent weights, providing high-dimensional problem(s) as requested. As outlined above, JALA-EM is comparable to SOUL which is to be expected due to the correlations that exist between particles due to the resampling mechanism present in JALA-EM, and again note that PGD does not support marginal likelihood estimation.\n\n\n> The authors do not comment on the sensitivity of tuning parameters such as $\\gamma$ and $h$. The theoretical results do not seem to provide a clear indication of how these parameters should be tuned in practice. What advice should be given to practitioners regarding step size selection so that JALA-EM can be easily used out-of-the-box?\n>\n\nThank you for highlighting this critical point, since this is also a common challenge for the other particle algorithms considered. To this end, as done in Appendix B.2 of our paper (see the **Approach.** section), we suggest utilising a systematic and reproducible procedure for model fitting and hyperparameter selection. In particular, we perform K-fold cross-validation on $\\mathcal{D}\\_{train, val}$ to select step-sizes. In our case, we have utilised a pre-defined grid for the particle step-sizes, and where applicable (i.e. for JALA-EM and SOUL) for the $\\theta$-update step-sizes also. To form the grid of step-sizes, relevant to the particle updates, we compute a baseline step-size, which we denote $h\\_{euler}$, and the idea here is to provide a theoretically grounded scale reflecting the curvature of the energy landscape. More details pertaining to these grids can notably be found in both Appendix B.2. and Appendix C.1.2. The evaluation metric utilised for this hyperparameter tuning can of course be chosen at the discretion of the practitioner, however we opted for the Log Pointwise Predictive Density (LPPD), which assesses the model’s average predictive accuracy on unseen data. To be clear, the step-size combination yielding the highest average LPPD across the folds is chosen. \n\nWe hope that this cross-validation scheme serves as an initial systematic way in which step-sizes can be chosen, ensuring our method can be used *out-of-the-box*. We agree that further refinement of this process would be a promising direction for future research.\n\nWe are happy to discuss any of these points further if the reviewer has any more concerns."}]

Author rebuttal:

Meta-review: # Summary 

This paper introduces a novel method, referred to as the Jarzynski-Adjusted Langevin Algorithm, for maximum marginal likelihood estimation. The core challenge is estimating the gradient of the marginal likelihood which is an integral over the latent variables and is therefore intractable. JALA addresses this by running a SMC algorithm with an appropriate sequence of distributions and transitions given by the Unadjusted Langevin Algorithm applied to this sequence. 

This approach yields a population of weighted particles, whose gradient estimates are used in a first-order optimization loop. The authors provide non-asymptotic convergence guarantees under standard conditions (log-concavity, Polyak–Łojasiewicz), showing explicit rates depending on particle count and step size.

Empirical results on Bayesian logistic regression, linear regression model selection, and a small Bayesian neural network show that JALA-EM offers competitive performance.

# Recommendation 

All reviewers find the paper worth publication. and therefore I recommend acceptance.
I strongly encourage the authors to incorporate the reviewers’ feedback and points raised during the rebuttal discussion when preparing the final version of the paper.
Decision: Accept (poster)