Learning Latent Variable Models via Jarzynski-adjusted Langevin Algorithm

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.

Contributions. To address the issues mentioned above, in this paper, we build on a numerical technique to sample probability paths, which we term Jarzynski adjusted Langevin algorithm (JALA). 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 on this idea:

• 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 developed in Carbone et al. (2023) and is closely related to sequential Monte Carlo (SMC) samplers (Del Moral et al., 2006) -and can be seen as a weighted-ensemble of the ULA.

• In Section 3, using the JALA as the core component, we propose a numerical method for EM, using JALA, which we term Jarzynski adjusted Langevin algorithm for EM (JALA-EM).

This method uses the JALA for estimating the gradient of the marginal likelihood, which is then used to update the parameters via a gradient-based optimiser. The resulting JALA-EM algorithm is 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 (PŁ) 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 our nonasymptotic result. This is just a first step, as any other gradient-based 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 LVMs 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/jamescuin/jala-em.

We first introduce here the Jarzynski adjusted Langevin algorithm (JALA) which is a sampling method for time-varying probability distributions. The main idea behind the method is to run ULA on these time-varying potentials and correct the resulting bias (due to time-dependence) using the Jarzynski equality. The resulting JALA is a sampler for a sequence of evolving measures.