A Bayesian Optimization through Sequential Monte Carlo and Statistical Physics-Inspired Techniques
Anton Lebedev Thomas Warford M. Emre Şahin










A. Lebedev et al.



The Hartree Centre, Keckwick Ln, Warrington, UK
 {anton.lebedev, emre.sahin}@stfc.ac.uk University of Manchester,  M13 9PL, Manchester, UK
thomas.warford@student.machester.ac.uk


A. Lebedev et al.






Springer Copyright Notice

Copyright (c) 2023

This work is subject to copyright. All rights are reserved by the
Publisher, whether the whole or part of the material is concerned,
specifically the rights of translation, reprinting, reuse of illustra-
tions, recitation,broadcasting, reproduction on microfilms or in
any other physical way, and transmission or information storage
and retrieval, electronic adaptation, computer software,or by sim-
ilar or dissimilar methodology now known or hereafter developed.

Published in: Lecture Notes in Computer Science 14077 proceedings of the 23rd International Conference Prague, Czech Republic, July 3 - 5, 2023 Proceedings, Part V


A Bayesian Optimization through Sequential Monte Carlo and Statistical Physics-Inspired Techniques
    Anton Lebedev1 Thomas Warford2 M. Emre Şahin1 

    
School of Physics, Nankai University, Tianjin, 300071, China 

==================================================================================================





§ ABSTRACT 

In this paper, we propose an approach for an application of Bayesian optimization using Sequential Monte Carlo (SMC) and concepts from the statistical physics of classical systems. Our method leverages the power of modern machine learning libraries such as NumPyro and JAX, allowing us to perform Bayesian optimization on multiple platforms, including CPUs, GPUs, TPUs, and in parallel. Our approach enables a low entry level for exploration of the methods while maintaining high performance. We present a promising direction for developing more efficient and effective techniques for a wide range of optimization problems in diverse fields.







§ INTRODUCTION


Bayesian optimization of ever-growing models has become increasingly important in recent years and significant effort has been invested in achieving a reasonable runtime-to-solution. Unfortunately, most optimization tasks are implemented and optimized within a specific framework, resulting in a single optimized model. The proliferation of such implementations is difficult, as it requires both domain expertise and knowledge of the specific framework and programming language. 


Probabilistic programming frameworks such as Stan <cit.> and (Num)Pyro <cit.>, provide such support for efficient optimisation methods such as Hamiltonian Monte Carlo (HMC) algorithm which can explore complex high-dimensional probability distributions. These frameworks are powerful tools for statistical analysis and inference.

Stan excels at handling complex hierarchical models with ease, which is often challenging in other probabilistic programming frameworks. Its user-friendly interface makes it accessible to those with little experience in Bayesian inference and statistical modelling, making it a popular choice for the accurate and efficient analysis of complex models. 
It provides domain experts with a performant tool to perform the said task with little to no programming knowledge. It achieves this by defining a "scripting language" for models and translation of these into C++ code.
It suffers, however, from a lack of inherent parallelism and a formulation of its methods in heavily-templated C++. 
NumPyro, built on JAX <cit.>, enables efficient exploration of high-dimensional probabilistic models using different methods, while JAX itself combines the flexibility and ease-of-use of NumPy with the power and speed of hardware accelerators for efficient numerical computing. Additionally, JAX offers compatibility and portability by supporting code execution on a variety of hardware. 


Inspired by HMC and SMC descriptions in <cit.>, we implemented an HMC algorithm in Python using the NumPyro and JAX frameworks and with DeepPPL <cit.> utilized to translate existing Stan models into their Python equivalents.
In this paper we present preliminary findings from our implementation of SMC for Bayesian parameter searches developed with its physical origins intact.



§ METHOD DESCRIPTION

Upon review of the SMC algorithm <cit.> and its HMC kernel, it has become apparent that the approach resembles the cooling process of an ideal gas in a potential field with unknown minima, and that the algorithm's development would benefit from an understanding based on physical systems. To facilitate this understanding, we have undertaken a reformulation of the SMC and HMC algorithms to more accurately reflect the particle ensembles of statistical physics. 
A simple first step was the reintroduction of a temperature into the expressions
for probability, seeing as the probabilities used in these methods are the maximum-entropy
probabilities of an ensemble (collection) of particles at a fixed energy:

    p_i =  e^-E_i/(k_B T)/∑_j=1^N e^-E_j/(k_B T) .

Here k_B is the Boltzmann constant (carrying the dimensions of energy per degree of temperature), T the temperature and E_i = ℋ(q_i,p_i) is the energy of the particle i at position q_i with a momentum p_i given a Hamiltonian:

    ℋ(q,p) = p^2/2m + V(q) .

As is common, the Hamiltonian encodes the dynamics of the system. 
Since we seek to determine the parameters that maximise the log-likelihood of the model the potential
can be defined<cit.> as:

    V(q) := -ln(P(q|X) ) ,

where P(q|X) is the posterior probability density of the model, where X is the vector containing the
observations and q is the (position) vector in parameter space - the parameters of the model.

The formulation of (<ref>)
commonly used in mathematics implies T = 1k_B or an arbitrary definition of k_B := 1.
The former fixes a degree of freedom of the method, whilst the latter allows for a variation, but removes the 
dimensionality of the constant which, in conjunction with T ensures that the argument of the exponential
remains a dimensionless number or quantity. Whilst functionally of no consequence retaining the physical dimensions of the respective quantities allows for sanity checks of formulae during development.

The distribution of the "auxiliary momentum" described in <cit.> is naturally dependent on T,
with a higher temperature T resulting in a broader distribution of the momenta (faster particles) and a wider area of the
initial sample space covered.

The HMC process is rather simple and described in <cit.> in sufficient detail.
In contrast, sequential Monte-Carlo is provided in a less legible form in <cit.>. Hence we privde here the simple
version we turn our attention to:

    
  * Initialise the position samples to a normal distribution on ℝ^n and the momentum samples to a normal distribution
    with temperature T.
    
  * Iterate for a given number of SMC steps:
    
        
  * Propagate each particle in the ensemble using HMC.
        
  * Determine the lowest energy for all particles, subtract it from every energy (renormalisation) and store it for future processing.
        
  * Compute and store the average parameter value.
        
  * Determine the effective size N_𝑒𝑓𝑓𝑒𝑐𝑡𝑖𝑣𝑒 of the ensemble according with <cit.>.
        
  * If resampling is necessary select N_𝑒𝑓𝑓𝑒𝑐𝑡𝑖𝑣𝑒 particles with largest weights (lowest energies) and duplicate
        them according to their probability until the original ensemble size is reached. Then reset the momenta to a thermal distribution and the weights of the particles to 1/N.
    
    
  * Compute the moving average of the stored averages, using the stored energies as weights in accordance with (<ref>).


Here we must note the rescaling of the weights by subtraction of the lowest energy, which is physically motivated by the freedom to choose the origin of the energy scale. Similarly, the weighted moving average in the last step results in a smoothing of the excursions of the mean after a resampling step by weighting means with large associated energies exponentially smaller (the higher the energy the more unlikely a configuration is to occur).



§ NUMERICAL EXPERIMENTS




 §.§ Models





To demonstrate the effectiveness of the implementation we have selected two simple models:


  * A sequence of M independent tosses of two coins - the Coin Toss (CT) model.


  * Item Response Model with Two-Parameter Logistic (IRT 2PL).




  §.§.§ CT model


The CT model assumes complete ignorance of the a-priori coin bias p and its maximum a-posteriori probability estimator  can be determined formally to be:

    P̂_MAP = K/N .

Here K is the number of observed heads and N = 40 is the total number of observations.
This allows us to check the numerical approximation of (<ref>) as well as its 
gradients, ensuring proper functioning of the implementation. The true parameters of the coin bias for each of the two coins are

    p_1 = 1/2 ,  p_2=3/4 .

The potential of the CT model, along with a few sample trajectories of the particles propagating therein are shown in
fig. <ref> prior to the constraining to the support of the model, and in fig. <ref> after.
Here we note that we chose not to invert the momentum in case the new phase-space point (q,p) passes
the Metropolis-Hastings acceptance test, contrary to alg. 1 of <cit.>. As can be seen in the figures, such an inversion
results in a rather slow sampling of the potential in the case of a smooth potential. It is, however, beneficial for a rough potential
and hence likely better in most practical applications.



  §.§.§ IRT 2PL

Item Response Theory (IRT) <cit.> is a statistical model that is widely used in a variety of research fields to analyze item responses in assessments or surveys. The  two-parameter logistic (2PL) model is a specific type of IRT model that assumes each item has two parameters:  the difficulty parameter and the discrimination parameter.
The model assumes that the item responses y_i for i = 1, …, I are Bernoulli distributed with a logit link function. The probability of responding correctly to item i is given by:


    P(y_i=1 | θ, a_i, b_i) = 1/1 + exp(-a_i(θ - b_i))


where θ is the person's latent ability, a_i is the discrimination parameter for item i, and b_i is the difficulty parameter for item i. 



 §.§ Experimental Results

Given the physical interpretation of the SMC iteration and the weighted average derived therefrom we 
expected - a-priori - a rather rapid and smooth convergence towards the true parameter values.
Given a smooth potential, as in the case of the CT model, it is furthermore expected that convergence to
the estimated parameters is faster for lower temperatures T.



  §.§.§ Quality of Estimation




As can be seen in fig. <ref> , all estimates of the bias of a fair coin converge within 5 iterations to
the true value 12 (indicated by the dashed red line). In case of a rather high temperature (more precisely: thermal energy) of T = 10 1/k_B
the apparent limit value deviates noticeably from the true value, given the remarks above this is not unexpected and
will be remedied in a future iteration of the method.

In the case of the biased coin, with p_2 = 3/4 it becomes obvious that the current development stage of the method may suffer from a freeze-in (c.f. T=1 case) resp. a bias in the implementation.
For a comparison reference, non-public, implementation results are marked as '+', showing a similar convergence
behaviour but without the offsets.



  §.§.§ Performance - HMC





Although a reference implementation is available for comparison, its drawbacks are its complexity and limited execution architectures.

Our development utilises JAX and NumPyro, allowing us to run the method for variable models on a variety of architectures. Here we present the performance data obtained for the pure HMC implementation on different architectures. One can observe, in fig. <ref> the behaviour of the run time when running our version of HMC on multiple CPUs, parallelised via MPI. For the CT model run times for up to 2048 particles are dominated by communication and data management overhead. This is a surprisingly small number, given the simplicity of the model and the limited amount of observed data fed into it. Overall the run time growth is sub-linear up until 65538 particles for 1 CPU. The resulting speed-up of using more than one CPU is depicted in fig. <ref>, demonstrating a sub-linear increase even with two CPUs. This observation confirms that the model is too small to scale effectively, at least for fewer than 128k particles (c.f. the IRT model below).






In contrast to the CT model is the IRT model, whose runtimes are displayed in fig. <ref>.
One can immediately see, that the run time is dominated by overhead only below 256 particles per device (here: CPU). It is also important to note, that the run time using 4 CPU cores is larger than when only 3 are utilised. We attribute this to the apparent and unintentional spawning of multiple Python processes per MPI process. Our observations show that each MPI process spawns roughly 2 Python processes. This will be a point of future investigation. This observation explains the apparent super-linear speed-up for the CPUs displayed in fig. <ref>.
One can also observe the order of magnitude reduction in the overall run time in the case of 65538 particles,
when one GPU is used. The sub-linear speed-up from 1 GPU to 2 GPUs can here be explained by the fact, that the
devices were asymmetrically bound to the system (PCIE-x16 vs. PCIE-x8) as well as that two devices of different generations were used: a GTX 1080 Ti and an RTX 3060.

The latter point shows the flexibility of JAX and NumPyro, which allowed us to run the same method on multiple CPU cores and multiple, heterogeneous, GPUs using MPI without having to modify the code!



§ CONCLUSIONS AND FUTURE WORK


In conclusion, we have demonstrated that using the NumPyro and JAX frameworks along with intuition from classical mechanics and statistical physics, it is possible to re-create an SMC Bayesian optimisation process, whilst enriching it with an intuitive understanding of the respective steps.

The selected framework enabled us to create a simple, easy to maintain, implementation of HMC and SMC that 
can be parallelised to multiple computing devices of varying architectures on demand.
Our implementation outperforms a similar implementation in Stan by a factor of ∼ 2 on a single CPU (core)
and scales well to multiple CPUs/GPUs, with larger gains obtained for larger models (or models with more observation
data) and more sampling particles.

In the near future, we plan to extend MPI parallelism from the core HMC stage to the entire SMC iteration, as well
as perform a thorough performance analysis and optimisation of our code, since preliminary checks indicate the existence
of, e.g., unnecessary host-device data transfers.

On the theoretical side we plan to continue the reformulation of SMC in the language of statistical physics and
expect the possibility to include maximum-entropy methods and thermalisation/annealing into the framework, utilising the
long history of MC in physics <cit.>.
Having a formulation of SMC that utilises terminology of (statistical) physics we hope to be able to
extend the approach from the classical onto the quantum domain. This holds the potential of including
existing quantum resources into Bayesian optimisation processes without requiring the user to know
the intricacies of the new architecture.



§ ACKNOWLEDGEMENTS

During PRACE Summer of High-Performance Computing 2022, TW was able to implement parallel HMC using the discussed formulations and included correctness checks based on physical systems and models with known closed-form solutions in the implementation.







 splncs04
 

