Learning Latent Variable Models via Jarzynski-adjusted Langevin Algorithm

TL;DR

This paper introduces JALA-EM, a novel algorithm combining nonequilibrium statistical mechanics and sequential Monte Carlo methods to efficiently estimate parameters and perform model selection in latent variable models.

Contribution

It develops a new sampling algorithm based on Jarzynski's equality integrated with Langevin dynamics, enabling recursive marginal likelihood estimation for latent variable models.

Findings

JALA-EM achieves competitive accuracy with existing methods.

The approach enables efficient model selection via recursive marginal likelihood estimation.

The method converges reliably under regularity assumptions.

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.