Minibatch training of neural network ensembles via trajectory sampling
Most iterative neural network training methods use estimates of the loss
function over small random subsets (or minibatches) of the data to update the
parameters, which aid in decoupling the training time from the (often very
large) size of the training datasets. Here, we show that a minibatch approach
can also be used to train neural network ensembles (NNEs) via trajectory
methods in a highly efficient manner. We illustrate this approach by training
NNEs to classify images in the MNIST datasets. This method gives an improvement
to the training times, allowing it to scale as the ratio of the size of the
dataset to that of the average minibatch size which, in the case of MNIST,
gives a computational improvement typically of two orders of magnitude. We
highlight the advantage of using longer trajectories to represent NNEs, both
for improved accuracy in inference and reduced update cost in terms of the
samples needed in minibatch updates.
Jamie F. Mair, Luke Causer, Juan P. Garrahan
# Minibatch training of neural network ensembles via trajectory sampling

###### Abstract

Most iterative neural network training methods use estimates of the loss function over small random subsets (or _minibatches_) of the data to update the parameters, which aid in decoupling the training time from the (often very large) size of the training datasets. Here, we show that a minibatch approach can also be used to train neural network ensembles (NNEs) via trajectory methods in a highly efficient manner. We illustrate this approach by training NNEs to classify images in the MNIST datasets. This method gives an improvement to the training times, allowing it to scale as the ratio of the size of the dataset to that of the average minibatch size which, in the case of MNIST, gives a computational improvement typically of two orders of magnitude. We highlight the advantage of using longer trajectories to represent NNEs, both for improved accuracy in inference and reduced update cost in terms of the samples needed in minibatch updates.

## I Introduction

Traditional machine learning (ML) applications aim to train a single model, usually by adjusting the parameters that define a complex function approximator like a neural network (NN), to perform well on some desired outcome as measured by a proxy loss function. A high performance on an appropriate loss function will entail a high performance on the metrics one cares about, for example accuracy in a classification problem [1]. There is strong empirical evidence from numerical experiments that increasing the scale of single models improves performance, as for example in the timely class of large language models (LLMs) [2].

However, to counteract the seemingly ever-increasing size of LLMs, there has also been significant work towards devising smaller models with similar capabilities in order to reduce the computational cost of training and of inference. A notable recent example is Stanford's Alpaca [3], based on LLaMA [4; 5], which can match GPT3.5 [6] despite being over an order of magnitude smaller. Another possibility is to replace one large model by an _ensemble_ of smaller models which can provide similar or better inferences while also being less costly to train and evaluate [7; 8]. This is the class of problems we focus on here.

Recently, we introduced an approach to train collectively an ensemble of models [9], in particular neural network ensembles (NNEs) where predictions at inference time are aggregated in a committee-like fashion (for classification, the ensemble prediction is the most voted for option, while for scoring, the ensemble prediction is the mean ensemble score). In Ref. [9] we defined an NNE in terms of the trajectory of the model parameters under a simple (discrete in time, diffusive in parameter space) dynamics, and trained it by biasing the trajectory that defines the NNE towards a small time-integrated loss. That is, once training is converged, the NNE corresponds to a discrete trajectory of the model parameters sampled from a distribution of trajectories exponentially "tilted" to have low time-integrated loss. This approach is borrowed from the study of glassy systems [10], where biasing dynamics according to time-integrated observables (e.g. the dynamical activity [11; 12]) is known to access low energy states for the configurations in the trajectory. Such low-loss trajectories can be accessed via importance sampling in trajectory space, such as transition path sampling (TPS) [13] as adapted to stationary dynamics and large deviation problems [14]. The ensuing trained NNE is a collection of NN models correlated by the underlying dynamics of the parameters and with a low value of the total loss due to the tilting.

While Ref. [9] provides a proof of principle of the trajectory sampling approach, it suffers from a significant computational bottleneck: importance sampling is a Monte Carlo scheme on trajectories, where updates are determined according to changes in the (time-integrated) loss evaluated over the _whole training set_, so that each Monte Carlo iteration scales with the size of the training data. For example, when training for the textbook MNIST digit classification problem in Ref. [9], we used only a small amount of the entire training dataset (2048 samples from the available 60000) to make the problem tractable for a comprehensive study. On the contrary, it is well known that ML models generalise poorly with small datasets [1]. This computational limitation makes the method of Ref. [9] impractical for more complex tasks. This has to be contrasted with gradient descent [1], where there is no need to sample faithfully from a distribution, so that the gradient of the loss can be estimated efficiently only on very small subsets of training data, known as _minibatches_, giving rise to stochastic gradient descent (where the noise from the difference between the minibatch estimate and the full loss actually helps convergence to a good local minimum [1]).

In this paper, we resolve the problem above by implementing a minibatch method in the trajectory sampling
used in the training of the NNEs. We build on the approach of Ref. [15] for doing Monte Carlo sampling with small data batches. We show that our new method reduces the training cost by a factor given by the ratio of the average minibatch size (which we determine in an adaptive manner) to the size of the dataset. We illustrate this more efficient method on MNIST classification (using the whole MNIST dataset), showing in this case a computational gain of about two orders of magnitude. Our minibatch approach also allows us to highlight the key features of the trajectory NNE method, showing the advantage of using longer trajectories to represent NNEs both in terms of accuracy and data requirement for training.

The rest of the paper is organised as follows. In Sec. II we describe the theory, reviewing the idea of NNEs as trajectories of a stochastic dynamics, training as tilted trajectory sampling, and the central approach to perform mini-batch trajectory Monte Carlo. In Sec. III we present the adaptive minibatch trajectory sampling method for training NNEs. We illustrate the method with two examples in Sec. IV, an exactly solvable linear perceptron, and the full MNIST digit classification problem. In Sec. V we give our conclusions, and further technical details are provided in the Appendices.

## II Theoretical background

### Neural network ensemble as a trajectory of neural networks

In Ref. [9], we proposed that a NNE could be obtained by evolving the parameters of a NN model under a suitable stochastic dynamics, where the NNE is composed of the sequence of NNs in time. If, at time step \(t\), the NN is defined by \(\mathbf{\theta}_{t}\), this dynamics would give rise to a trajectory \(\mathbf{\theta}_{1}\rightarrow\mathbf{\theta}_{2}\rightarrow\ldots\rightarrow\mathbf{ \theta}_{\tau}\), with the NNE as the set of visited models under the dynamics, \(\mathbf{\Theta}=[\mathbf{\theta}_{1},\mathbf{\theta}_{2},\ldots,\mathbf{\theta}_{\tau}]\). As the aim is to minimise the loss over the ensemble

\[\mathcal{L}(\mathbf{\Theta})=\sum_{t=1}^{\tau}L(\mathbf{\theta}_{t}) \tag{1}\]

where \(L(\mathbf{\theta}_{t})\) is the standard loss for the \(t\)-th model (see below for a specific form of the loss), training is equivalent to finding a suitable dynamics whose typical trajectories are those with low time-aggregated loss, \(\mathcal{L}(\mathbf{\Theta})\).

Once trained, this dynamics is defined in terms of (in general time-dependent) stochastic dynamics \(\mathcal{M}(\tau,\sigma,s)\equiv\{M_{i;\sigma,s}\}_{t=1}^{\tau-1}\), where \(M_{i;\sigma,s}(\mathbf{\theta}^{\prime}|\mathbf{\theta})\) are the transition probabilities at each time step, such that the NNE corresponds to a trajectory generated using dynamics \(\mathcal{M}(\tau,\sigma,s)\). This approach is illustrated in Fig. 1(a): the NNE is a discrete-time trajectory, where each state along the trajectory corresponds to one of the NNs that form the ensemble. Starting from the first model, \(\mathbf{\theta}_{1}\), each subsequent model is sampled according

Figure 1: (a) An NNE as a stochastic trajectory, and a sketch of trajectory sampling. Each state in the trajectory corresponds to one NN model in the NNE. The proposed path sampling update from trajectory \(\mathbf{\Theta}\) to trajectory \(\mathbf{\Theta}^{\prime}\) is via a stochastic bridge in which only one model is modified. (b) The loss function of a model with parameters \(\mathbf{\theta}\). Each layer represents a data element. The loss, \(L(\mathbf{\theta})\), is given by the average of the individual loss for each of the elements of the dataset, \(\{l_{i}(\theta)\}_{iN}\). The minibatch estimate of the loss is instead the average over a random selection \(\{r_{j}\}_{1:b}\) of \(b\) data points in the dataset. (c) Exactly solvable linear perceptron: mean NNE loss \(\langle\mathcal{L}\rangle\) per NNE size \(\tau\), as a function of \(s\) for various \(\tau\). Lines are analytical results. Symbols are the numerical results using the minibatch TPS algorithm for training (with \(20\times 10^{6}\) TPS epochs as a “burn-in” followed by \(20\times 10^{6}\) epochs for convergence of training).
to \(\mathcal{M}(\tau,\sigma,s)\). Three hyperparameters determine the dynamics that produces the NNE: the final time \(\tau\) sets the number of models in the ensemble; \(\sigma\) sets the "stiffness" of the chain (see below for details), that is, how correlated subsequent models are to each other, with small \(\sigma\) corresponding to large stiffness; and \(s\) controls the level of the overall loss, with larger \(s\) corresponding to lower loss. The central idea is that a correlated chain of models generated as a trajectory from dynamics \(\mathcal{M}(\tau,\sigma,s)\) at large \(s\) provides a well trained NNE [9].

### Learning as a trajectory sampling problem

Obtaining a suitable dynamics \(\mathcal{M}(\tau,\sigma,s)\) that produces well-trained NNEs as its typical trajectories is a difficult task. We can however resolve this problem by means of trajectory sampling techniques [9]. Consider as a starting point an untrained dynamics \(\mathcal{M}(\tau,\sigma,0)\) with the same transition probabilities at every time step \(t\)[9]

\[M_{\sigma}(\mathbf{\theta}_{t+1}|\mathbf{\theta}_{t})\propto\exp\left[-\frac{1}{2 \sigma^{2}}(\mathbf{\theta}_{t}-\mathbf{\theta}_{t+1})^{2}\right], \tag{2}\]

with \(\int_{\mathbf{\theta}^{\prime}}M_{\sigma}(\mathbf{\theta}^{\prime}|\mathbf{\theta})=1\). This dynamics corresponds to a discrete-time Gaussian diffusion process that knows nothing about the loss (1). As such, a typical trajectory drawn from it will correspond to a random (and therefore untrained) NNE. As indicated above, the parameter \(\sigma\) sets the variance of the diffusive steps, so that for smaller \(\sigma\) subsequent models are more correlated, while for \(\sigma\to\infty\) all the models of the chain are uncoupled.

The dynamics (2) produces an _ensemble of trajectories_ (and therefore an ensemble of NNEs) with each trajectory having probability

\[P(\mathbf{\Theta};\sigma)=\frac{1}{\mathcal{Z}_{\tau}(\sigma)}p(\theta_{1})\prod_ {t=1}^{\tau-1}\exp\left[-\frac{1}{2\sigma^{2}}(\mathbf{\theta}_{t}-\mathbf{\theta}_{t+ 1})^{2}\right], \tag{3}\]

given by the product of the \(M_{\sigma}\) at each step. Here \(p(\theta_{1})\) is the probability used to draw the first model, and \(\mathcal{Z}_{\tau}(\sigma)\) a normalisation constant (the "partition sum" of the trajectory ensemble). In order to obtain trajectories with low overall loss, what we aim is to define a new trajectory ensemble that is exponentially "tilted" with respect to (3), as is standard in large deviation studies of dynamics (e.g., Ref. [16]), that is [9]

\[P(\mathbf{\Theta};\sigma,s)=\frac{1}{\mathcal{Z}_{\tau}(\sigma,s)}e^{-s\mathcal{L }(\mathbf{\Theta})}P(\mathbf{\Theta};\sigma). \tag{4}\]

For large \(s\), a typical trajectory from this ensemble will correspond to a NNE with low overall loss. The learned dynamics \(\mathcal{M}(\tau,\sigma,s)\) of the previous subsection would be the dynamics that produces trajectories distributed according to the tilted distribution (4).

One way to avoid having to determine the \(\mathcal{M}(\tau,\sigma,s)\) dynamics explicitly is to directly sample trajectories from the tilted distribution (4). In this way, convergence of the training, that is, finding \(\mathcal{M}(\tau,\sigma,s)\), coincides with convergence of the trajectory sampling of (4), as we do here by means of an importance sampling method in trajectory space based on transition path sampling (TPS) [13].

### Monte Carlo in trajectory space

Consider a Monte Carlo scheme for sampling trajectories, specifically a Metropolis-Hastings approach [13]: given a current trajectory \(\mathbf{\Theta}\), the probability to change to a new trajectory \(\mathbf{\Theta}^{\prime}\) is given by

\[p(\mathbf{\Theta}^{\prime}|\mathbf{\Theta})=g(\mathbf{\Theta}^{\prime}|\mathbf{\Theta})A(\mathbf{ \Theta}^{\prime},\mathbf{\Theta}), \tag{5}\]

where the factor \(g(\mathbf{\Theta}^{\prime}|\mathbf{\Theta})\) is the probability to propose the move, and \(A(\mathbf{\Theta}^{\prime},\mathbf{\Theta})\) is that to accept it. For the above to converge to (4) we need to impose that it obeys detailed balance with respect to (4), which implies

\[\frac{A(\mathbf{\Theta}^{\prime},\mathbf{\Theta})}{A(\mathbf{\Theta},\mathbf{\Theta}^{\prime} )}=\frac{P(\mathbf{\Theta}^{\prime};\sigma)g(\mathbf{\Theta}|\mathbf{\Theta}^{\prime})}{P( \mathbf{\Theta};\sigma)g(\mathbf{\Theta}^{\prime}|\mathbf{\Theta})} \tag{6}\]

If the proposed moves obey detailed balance with respect to the original untilted dynamics (3),

\[\frac{g(\mathbf{\Theta}^{\prime}|\mathbf{\Theta})}{g(\mathbf{\Theta}|\mathbf{\Theta}^{\prime} )}=\frac{P(\mathbf{\Theta}^{\prime};\sigma)}{P(\mathbf{\Theta};\sigma)}, \tag{7}\]

then the acceptance ratio reduces to

\[\frac{A(\mathbf{\Theta}^{\prime},\mathbf{\Theta})}{A(\mathbf{\Theta},\mathbf{\Theta}^{\prime} )}=e^{-s\left[\mathcal{L}(\mathbf{\Theta}^{\prime})-\mathcal{L}(\mathbf{\Theta}) \right]} \tag{8}\]

In standard TPS, (7) is realised by proposing trajectories by simply running the original dynamics (3) (via "shooting" or "shifting" moves, see Ref. [13]). This approach, however, carries an exponential cost in the time extent of trajectories, since the loss difference in the exponent of (8) scales linearly with time. This can be mitigated [9] by proposing small changes in a trajectory, see Fig. 1(a): the proposed trajectory is one where only the state at one time is modified; as this has to obey (7), it has to be done as a _Brownian bridge_[17; 18]. That is, as conditioned dynamics starting in the previous state and retuning to the state after the one changed (see Ref. [9] for details).

## III Minibatch path sampling

While the Brownian bridge version of TPS ameliorates the exponential-in-time cost in the trajectory sampling, there is another source of computational slowness coming from the evaluation of the trajectory loss, cf. (8). With the Brownian bridge TPS, Fig. 1(a), only a single model changes between the current and proposed trajectory, and the change in trajectory loss in (8) is therefore
given by that the change of that model's loss. If this change is at time \(t\), this requires the evaluation of \(L(\mathbf{\theta}_{t}^{\prime})\), which at training is the average of the loss under that model for each of the \(N\) training data points,

\[L(\mathbf{\theta}_{t}^{\prime})=\frac{1}{N}\sum_{i=1}^{N}l_{i}(\mathbf{\theta}_{t}^{ \prime}) \tag{9}\]

where \(l_{i}(\mathbf{\theta}_{t})\) is the loss for the inference for data point \(i\). This means that in each Monte Carlo iteration, computing the change in loss inevitably scales with the training set size \(N\) (together with a cost that depends on the size and architecture of the NN being considered). This evaluation can therefore become computationally infeasible for larger datasets. For example, in Ref. [9], we had to reduce the training dataset by almost an order of magnitude to show a proof-of-principle of the method with for the MNIST classification problem.

In contrast to gradient descent, one cannot simply replace the loss over the whole training set for an estimate based on a small subset, or minibatch. For gradient descent, the error that this introduces becomes a source of noise, converting it into stochastic gradient descent (and its adaptive variants [19, 1, 20]). This in turn gives rise to the usual advantages that an exploit/explore strategy brings, in this case to minimise the loss locally descending the gradient versus exploration of the loss landscape. Since Monte Carlo aims to sample from a distribution, Eq. (4) in our case, a straightforward replacement of the loss by a minibatch approximation would lead to failure of the necessary detailed balance condition.

This problem has been considered before in the context of Bayesian inference, where so-called "tall datasets" make Monte Carlo inefficient, see e.g., Refs. [21, 22, 23]. In what follows we build on the approach put forward in Ref. [15] to develop an adaptive minibatch trajectory sampling method.

### Monte Carlo with minibatches

We first describe the scheme of Ref. [15] in the context of the Monte Carlo annealing of a system with degrees of freedom \(\mathbf{\Theta}\) (a NNE in our case) and target distribution (4), and in the next subsection we extend the approach to integrate it with TPS in an adaptive manner.

Let us define the quantity \(\Delta(\mathbf{\theta}^{\prime},\mathbf{\theta})\) as the logarithm of the change in weight under a proposed move,

\[\Delta(\mathbf{\Theta}^{\prime},\mathbf{\Theta})=-s\left[\mathcal{L}(\mathbf{\Theta}^{ \prime})-\mathcal{L}(\mathbf{\Theta})\right], \tag{10}\]

and choose our acceptance function as

\[A(\mathbf{\Theta}^{\prime},\mathbf{\Theta})=(1+e^{\Delta(\mathbf{\Theta}^{\prime},\mathbf{ \Theta})})^{-1}, \tag{11}\]

which satisfies the detailed balance condition (8). Monte Carlo works by generating a proposed move from \(g(\mathbf{\Theta}^{\prime}|\mathbf{\Theta})\), and then accepting the move if

\[A(\mathbf{\Theta}^{\prime},\mathbf{\Theta})>V, \tag{12}\]

where \(V\) is a uniformly distributed random number, \(V\sim\mathcal{U}(0,1)\). As (11) is a logistic (or sigmoid) function whose inverse is also its derivative, we can equivalently write the acceptance test as \(\Delta(\mathbf{\Theta}^{\prime},\mathbf{\Theta})>X_{\text{log}}\), where \(X_{\text{log}}\) is a logistically sampled random variable. As this distribution is symmetric around zero, we can equally write the test (12) as

\[\Delta(\mathbf{\Theta}^{\prime},\mathbf{\Theta})+X_{\text{log}}>0. \tag{13}\]

The loss \(\mathcal{L}(\mathbf{\Theta})\) that enters in (10) is the average of the loss over the entire training data set

\[\mathcal{L}(\mathbf{\Theta})=\frac{1}{N}\sum_{i=1}^{N}\mathcal{L}_{i}(\mathbf{\Theta} )=\frac{1}{N}\sum_{i=1}^{N}\sum_{t=1}^{\tau}l_{i}(\mathbf{\theta}_{t}), \tag{14}\]

where \(\mathcal{L}_{i}(\mathbf{\Theta})\) is the (trajectory) loss for the inference on data point \(i\). Consider now an approximation of the loss difference in terms of a random minibatch of size \(b\)

\[\mathcal{L}(\mathbf{\Theta})\approx\frac{1}{b}\sum_{j=1}^{b}\mathcal{L}_{r(j)}( \mathbf{\Theta}), \tag{15}\]

where \(r(j)\) specifies a random permutation of the indices of the elements in the training dataset, cf. Fig. 1(b). In terms of the above we can define an approximation to (10) [15],

\[\Delta^{*}(\mathbf{\Theta},\mathbf{\Theta}^{\prime})=-\frac{s}{b}\sum_{j=1}^{b}\left[ \mathcal{L}_{r(j)}(\mathbf{\Theta}^{\prime})-\mathcal{L}_{r(j)}(\mathbf{\Theta}) \right]. \tag{16}\]

If one could replace \(\Delta(\mathbf{\Theta},\mathbf{\Theta}^{\prime})\) by \(\Delta^{*}(\mathbf{\Theta},\mathbf{\Theta}^{\prime})\), then there would be a computational gain that would scale as the ratio of the size \(N\) of the training dataset to that of the minibatch \(b\) (as it is necessary to compute \(\Delta(\mathbf{\Theta},\mathbf{\Theta}^{\prime})\) for each Monte Carlo iteration). The way to do so is as follows [15].

Since the elements of the minibatch are chosen in an identical and independent manner, from the central limit theorem and for large enough \(b\), we expect \(\Delta^{*}\) to be normally distributed around \(\Delta\) with some variance \(\rho^{2}(\Delta^{*})\). That is, \(\Delta^{*}=\Delta+X_{\text{norm}}\), where \(X_{\text{norm}}\) is an approximately normal zero-mean random correction of variance \(\rho^{2}(\Delta^{*})\). As a logistically distributed random variable, \(X_{\text{log}}\), is almost normally distributed, we can write \(X_{\text{log}}\) as \(X_{\text{log}}=X_{\text{norm}}+X_{\text{corr}}\), where \(X_{\text{corr}}\) is the (hopefully small) correction to normality. Inserting this decomposition of \(X_{\text{log}}\) into the acceptance test (13), we can replace \(\Delta\) with the minibatch estimate:

\[\Delta^{*}(\mathbf{\Theta}^{\prime},\mathbf{\Theta})+X_{\text{corr}}>0. \tag{17}\]

This new acceptance only depends on the minibatch estimates of the loss and -- if accurate -- will be efficient for \(b\ll N\). The test (17) will asymptotically give the correct acceptance distribution provided that (i) that the fluctuations \(X_{\text{norm}}\) of \(\Delta^{*}\) around \(\Delta\) are normally distributed (which can be checked to adjust the size of \(b\)), and (ii)
that the distribution \(C_{\rm corr}(X;\rho)\) for the random correction \(X_{\rm corr}\) can be numerically computed with low error (see [15] for analysis of the errors).

Thanks to the CLT, condition (i) is relatively easy to satisfy for a large enough minibatch sample size. If the batch size grows beyond the size of the dataset, we do not need this approximation and can use (12). Condition (ii) holds when the sample error on \(\Delta^{*}\) is sufficiently small [15]. In practice one can only compute the distribution of \(C_{\rm corr}(X;\rho)\) accurately enough only for standard deviations of \(X_{\rm norm}\) such that \(\rho\lesssim 1.1\) (see below for our implementation).

As computing \(C_{\rm corr}(X;\rho)\) is numerically expensive (see Appendix A for details) for each empirical \(\rho\approx 1\), we can instead compute a \(\rho^{2}=1\) correction and then add a further "normal correction" with small variance \(1-\rho^{2}\) to make a total normal random variable with fixed variance \(1\). Putting all of these together, we get the minibatch acceptance test that we use

\[\Delta^{*}(\mathbf{\Theta},\mathbf{\Theta}^{\prime})+X_{\rm nc}+X_{\rm corr}>0, \tag{18}\]

where \(X_{\rm nc}\) stands for the normal correction random variable, of zero mean and variance \(1-\rho^{2}\).

### Generalisation to enable training with TPS

The TPS scheme that we use relies on proposing trajectory updates, cf. Fig. 1(a), consisting of _bridging_ moves (for changes in the middle of the trajectory) and _shooting_ moves to get changes in the endpoints, see Ref. [9] for details. In either case, only a single model (i.e. a single time step) is altered, thus reducing the size of the update and improving the acceptance rate.

Once a candidate trajectory is proposed, the minibatch acceptance criterion of the previous subsection is applied. The specific steps are as follows, defining an adaptive minibatch scheme:

* We draw \(m\) random samples from the training set, which are used to calculate an estimate of \(\Delta^{*}(\mathbf{\Theta},\mathbf{\Theta}^{\prime})\) and \(\rho^{2}(\Delta^{*})\), cf. (16). If \(\rho^{2}>1\), \(m\) more samples are drawn without replacement, updating \(\Delta^{*}\) and \(\rho\) accordingly (and terminating if all samples are used). In this way, we form a minibatch of overall size \(b\) such that the sample variance of \(\Delta^{*}\) is strictly less than or equal to \(1\).
* We draw the random correction \(X_{\rm nc}\) from a normal distribution and \(X_{\rm corr}\) from \(C_{\rm corr}(X;\rho)\). With these we use (18) to accept or reject the proposed change to the trajectory. [If the total minibatch size \(b\) equals \(N\), then we use the original test (12) as \(\Delta^{*}\) coincides with \(\Delta\) in that case.] Unlike Ref. [15], exact sampling of (4) is not required to effectively train our NNE, and we do not test the normality assumption of \(\Delta^{*}(\mathbf{\Theta},\mathbf{\Theta}^{\prime})\). This is equivalent to setting their threshold \(\delta\to\infty\) in Ref. [15]. Justification for this simplifying choice is given in Appendix C.

```
1:input Initial trajectory \(\mathbf{\Theta}_{1}\), dataset \(\{x_{1},\ldots,x_{N}\}\), trajectory length \(\tau\), trajectory coupling \(\sigma\), minibatch chunk size \(m\), pre-computed correction \(C_{\rho=1}(X)\) distribution, cut-off hyperparameters \(c_{0}\) and \(c_{1}\) and training epochs \(E\).
2:output Sequence of trajectories \(\{\mathbf{\Theta}_{1},\mathbf{\Theta}_{1},\ldots,\mathbf{\Theta}_{E+1}\}\)
3:for\(k\in[1,2,\ldots,E]\)do
4: Sample proposal \(\mathbf{\Theta}^{\prime}\) using \(g(\mathbf{\Theta}^{\prime}|\mathbf{\Theta}_{k},\tau,\sigma)\) (i.e. shooting or bridging)
5: Sample \(\Delta^{*}(\mathbf{\Theta}^{\prime},\mathbf{\Theta}_{k})\) and \(\rho^{2}\) using \(m\) randomly selected samples, without replacement
6:\(b\gets N\)
7:while\(\rho^{2}>1\)and\(b<N\), and \(\max(\frac{|\Delta^{*}(\mathbf{\Theta},\mathbf{\Theta}^{\prime})|}{\rho}-c_{1},0)<=c_{0}\)do
8: Select \(m\) more randomly selected samples, without replacement and update estimates for \(\Delta^{*}(\mathbf{\Theta}^{\prime},\mathbf{\Theta}_{k})\) and \(\rho^{2}\)
9:\(b\gets b+m\)
10:endwhile
11:if\(b=N\)then
12: Sample random number \(V\sim\mathcal{U}(0,1)\)
13:if\(V<g(\Delta^{*}(\mathbf{\Theta}^{\prime},\mathbf{\Theta}_{k}))\)then
14: Accept with \(\mathbf{\Theta}_{k+1}\leftarrow\mathbf{\Theta}^{\prime}\)
15:else
16: Reject with \(\mathbf{\Theta}_{k+1}\leftarrow\mathbf{\Theta}_{k}\)
17:endif
18:continue
19:endif
20: Sample \(X_{\rm nc}\sim\mathcal{N}(0,1-\rho^{2})\) and \(X_{\rm corr}\sim C_{\rho=1}(X)\)
21:if\(\Delta^{*}(\mathbf{\Theta}^{\prime},\mathbf{\Theta}_{k})+X_{\rm nc}+X_{\rm corr}>0\)then
22: Accept with \(\mathbf{\Theta}_{k+1}\leftarrow\mathbf{\Theta}^{\prime}\)
23:else
24: Reject with \(\mathbf{\Theta}_{k+1}\leftarrow\mathbf{\Theta}_{k}\)
25:endif
26:endfor
```

**Algorithm 1** Minibatch TPS Training

One issue with this simple algorithm is that the minibatch size grows if the sample variance is much larger than \(1\). The biasing parameter, \(s\), scales the sample variance with \(s^{2}\), while taking \(b\) samples only reduces this variance by a factor of \(b\). For fixed \(\mathbf{\Theta}\) and \(\mathbf{\Theta}^{\prime}\), we would expect the minibatch size to change with \(s\) to compensate for the increased sample variance. For \(s\to\infty\), this would revert the minibatch method to the original full-dataset acceptance test. To avoid this, we introduce a _cut-off test_ which halts increasing the minibatch size when \(\Delta^{*}\) is sufficiently far away from the origin and performs an alternative acceptance test. The alternative acceptance test is broken into two stages. Firstly, we approximate the acceptance function to be equal to \(0\) when \(x<-c_{0}\) and \(1\) when \(x>c_{0}\) for some positive constant \(c_{0}\). We choose \(c_{0}\) to be sufficiently high that the approximate acceptance function is unchanged for \(-c_{0}\leq x\leq c_{0}\) without having to renormalise. Secondly, we choose another threshold for which the true mean \(\Delta\) is approximately guaranteed to be within \(\Delta^{*}\pm c_{1}\rho\). We use \(\max(\rho^{-1}|\Delta^{*}(\mathbf{\Theta},\mathbf{\Theta}^{\prime})|-c_{1},0)>c_{0}\) as our cut-off acceptance test. In our experiments, using
\(c_{1}=10\) and \(c_{0}=5\) gave good results. The final combined algorithm is given in Algorithm 1. Setting \(c_{0}\) or \(c_{1}\) to be high will make this cut-off less likely to be triggered, increasing minibatch size and therefore computational cost, however, setting them too close to \(0\) will result in an inaccurate acceptance test which does not obey detailed balance.

## IV Examples of training NN Ensembles via minibatch trajectory sampling

We now apply the method of Sec. III for the training of NNEs in two illustrative problems. The first is that of a linear perceptron, which is simple enough to be solved exactly, allowing us to directly compare our method with the expected results. The second is the more complex, but now standard, problem of MNIST digit classification [24].

### NNE of linear perceptrons

We first test the method with a linear classification problem, also considered in Ref. [9]. This problem can be defined as follows: we generate a set of independent random \(D\) dimensional points, \(\{\mathbf{x}_{i}\}_{i=1}^{N}\) (setting \(x_{N}=1\)), together with a \(D\)-dimensional random weight vector \(\mathbf{w}\) that we use to assign labels \(y_{i}=\mathbf{w}\cdot\mathbf{x}_{i}\) to each of the points \(\mathbf{x}_{i}\). The aim is to train the parameters \(\mathbf{\Theta}\) of an ensemble of \(\tau\) linear perceptrons, where the prediction of the \(t\)-th perceptron for the \(i\)-th data point is \(\mathbf{\theta}_{t}\cdot\mathbf{x}_{i}\). For training, we consider the mean-squared sample loss for each model in the NNE, which for data point \(i\) reads

\[l_{i}^{\rm(MSE)}(\mathbf{\theta}_{t})=\frac{1}{2}(y_{i}-\mathbf{\theta}_{t}\cdot\mathbf{x }_{i})^{2}. \tag{19}\]

The NNE loss over the training dataset, cf. (14), in turn reads

\[\mathcal{L}(\mathbf{\Theta})=\frac{1}{N}\sum_{i=1}^{N}\sum_{t=1}^{\tau}\frac{1}{2 }(y_{i}-\mathbf{\theta}_{t}\cdot\mathbf{x}_{i})^{2} \tag{20}\]

We now implement the adaptive minibatch estimation of the trajectory loss described in Alg. 1 to train this NNE for various \(\tau\) and \(s\). Figure 1(c) demonstrates that the numerics obtained in this way coincide with the analytic results from the exact trajectory distribution (4) [9]. This is an elementary proof-of-principle of the method.

Exact distributions for \(\mathbf{w}\) and \(\mathbf{x}\), along with experimental hyperparameters are provided in Appendix B.1.

### NNE for MNIST digit classification

The second problem we consider is that of an ensemble of models for classification of digits using the standard set of handwritten MNIST images [24], see Fig. 2(a). In this case, each NN in the NNE is a small convolutional neural network (CNN) whose architecture is described in Appendix B.2. For a given image \(X\), one of these CNNs with parameters \(\mathbf{\theta}\) provides the probability \(y(k|X;\mathbf{\theta})\) that the image corresponds to digit \(k\), for \(k=0,\ldots,9\). The appropriate loss function is the mean cross entropy, which for the data point \(i\) and the \(t\)-th model reads

\[l_{i}^{\rm(MNIST)}(\mathbf{\theta}_{t})=-\sum_{k=0}^{9}\delta_{z_{i},k}\log y(k|X_ {i};\mathbf{\theta}_{t}), \tag{21}\]

where \(z_{i}\) is the true classification of \(X_{i}\). The training loss for the NNE then reads

\[\mathcal{L}(\mathbf{\Theta})=-\frac{1}{N}\sum_{i=1}^{N}\sum_{t=1}^{\tau}\sum_{k=0} ^{9}\delta_{z_{i},k}\log y(k|X_{i};\mathbf{\theta}_{t}), \tag{22}\]

We train trajectories for a fixed number of epochs \(E\) (i.e., TPS iterations), which we choose to be large enough for the trajectory loss to appear to converge, which we track in terms of the minibatch loss estimate to avoid further computational cost. One can observe convergence for a small number of sample runs in Figs. 3(a, b). Specifically, we allow for a "burn-in" of the initial \(20\times 10^{6}\) epochs, and subsequently observe the average trajectory loss for the next \(20\times 10^{6}\) epochs. For each set of hyperparameters \((\tau,s)\) we perform six independent trainings starting each training run from a random initial seed trajectory. The time-averaged loss thus obtained is shown in Fig. 2(b) as a function of \(s\) for various NNE sizes \(\tau\). We note the following: (i) for every \(\tau\) the loss per model in the trained NNE decreases with \(s\), as should be the case when converging to (4); (ii) the larger \(\tau\) the lower the loss, indicating that the longer trajectories give rise to more accurate NNEs; (iii) there appears to be a transition from high to low loss with \(s\) which could indicate (dynamical) phase coexistence, as seen in many other trajectory ensemble problems [25].

A similar trend to that of the loss is observed in the accuracy on the generalisation test set. In Fig. 2(c) we show the accuracy of the final ensembles obtained after all training epochs. We use the NNE to collectively make a prediction on each sample by letting each model "vote" for their predicted class, with the class with the most votes getting selected (in the event of a tie, the smaller digit is selected). We plot this ensemble accuracy, averaged over six independent runs, for the same hyperparameters of panel (b).

As a proxy for the computational cost of training, in Fig. 2(d) we show the average batch size per epoch \(\langle b\rangle\) necessary for training to a certain value of the loss per model. Since the size of the training set is \(N=6\times 10^{4}\), the ratio \(N/\langle b\rangle\) gives the computational gain of using the minibatch method. From Figs. 2(b, c), we know that longer trajectories can yield a lower loss at smaller values of \(s\). From Fig. 2(d), we see that longer trajectories (larger NNEs) are also computationally more efficient to
train, requiring a smaller mean minibatch size, \(\langle b\rangle\), than smaller NNEs for the same level of overall loss, showing a computational gain in excess of two orders of magnitude of the minibatch method to the original full loss method of Ref. [9].

## V Conclusions

In this paper we have presented a variant of the minibatch Monte Carlo method of Ref. [15] adapted to the sampling of trajectories that correspond to neural network ensembles [9]. We have shown that this technique can be used to train NNEs via trajectory sampling to give an improvement in computational efficiency up to two orders of magnitude. While we have focused, for concreteness, on supervised learning applications, we note that an adaptive trajectory sampling technique like the one presented here should be also very useful in Monte Carlo based _reinforcement learning_ (RL), where datasets do not have a fixed size. We expect that this method will provide a stable training technique on these RL problems, which have exhibited brittle behaviour when continuously trained on changing objectives [26; 27; 28; 29].

Our results here add to the growing number of recent works studying the training dynamics of NNs from the statistical mechanics point of view, see e.g., Refs. [30; 31; 32; 33; 34]. Most of these consider the training of a single NN in terms of a stochastic dynamics akin to thermal annealing, cf. Ref. [32]. In contrast, our approach based on sampling trajectories of NNs shares more similarities to training by quantum annealing, see for example Refs. [35; 36]. Note that this similarity is not referring to actual unitary dynamics, but to the fact the computation of a trajectory ensemble in (4) is similar to that of a quantum partition sum (in terms of imaginary-time trajectories). Furthermore, the improved computational efficiency provided by the minibatch method we introduced here allowed us to highlight the benefit of larger NNEs (i.e., longer trajectories) capable of accessing lower loss regions of state space using far less data than single NNs or small ensembles.

## Code availability

Our TPS implementation package is available through GitHub, TransitionPathSampling.jl[37], together with the source code to generate the figures and results in the paper [38].

###### Acknowledgements.

 We acknowledge support from EPSRC Grant no. EP/V031201/1 and University of Nottingham grant no. FiF1/3. LC was supported by an EPSRC Doctoral prize from the University of Nottingham. Simulations were performed using the University of Nottingham Augusta HPC cluster, and the Sulis Tier 2 HPC platform hosted by the Scientific Computing Research Technology Platform at the University of Warwick. Sulis is funded by EPSRC Grant EP/T022108/1 and the HPC Midlands+ consortium. We thank the creators and community of the Julia programming language [39], and acknowledge use of the packages CUDA.jl[40; 41], Makie.jl[42] and ForwardDiff.jl[43].

Figure 2: (a) Representative examples of the MNIST digit images used for training. (b) Mean value of the time-averaged loss vs \(s\). Each data point is calculated after \(2\times 10^{7}\) TPS epochs over the following \(2\times 10^{7}\) TPS epochs. (c) Final accuracy on the standard \(10,000\) test images of a single trained NNE (via majority vote), taken at the end of the \(4\times 10^{7}\) epochs. (d) Average batch size per epoch, \(\langle b\rangle\) as a function of the converged mean NNE loss after \(2\times 10^{7}\) TPS epochs. The mean batch size per epoch is obtained from the \(4\times 10^{7}\) TPS training epochs.
## Appendix A Correction Distribution

Seita et al. [15] show that we can calculate the \(X_{\text{corr}}\) distribution numerically. We introduce a parameter \(V\) to specify the range of values to sample the distribution over. We construct two discrete vectors \(X\) and \(Y\) with the elements of \(X\) going linearly from \(-2V\) to \(+2V\) and \(Y\) going from \(-V\) to \(+V\). The vector \(X\) has \(4N+1\) elements and the vector \(Y\) has \(2N+1\) elements.

From here, we define a matrix \(M\) with elements

\[M_{ij}=\Phi_{\sigma}(X_{i}-Y_{j}), \tag{10}\]

where \(\Phi_{\sigma}\) is the cumulative distribution function (CDF) of a normal distribution with variance \(\sigma^{2}\). Additionally, we construct a new vector \(v\) such that

\[v_{i}=S(X_{i}), \tag{11}\]

where \(S\) is the logistic sigmoid function, i.e. the CDF of a logistically distributed random variable. Finally, we define the vector \(u\) to be \(u_{j}=C_{\sigma}(Y_{j})\) which is our target to calculate. This can be calculated using the formula

\[u=(M^{T}M+\lambda I)^{-1}M^{T}v, \tag{12}\]

where \(\lambda\) is regularisation parameter. We followed recommendations from Seita et al [15] and used \(V=10\), \(N=4000\) and \(\lambda=10\) to construct our numerical approximation of \(C_{\sigma}\). We set any negative elements equal to zero and re-normalise the CDF to ensure the area under the curve will equal 1.

Fortunately, the sampling algorithm allows us to calculate the distribution for a single value of \(\sigma\) to save on computation and memory. This distribution can be calculated once and cached for future use. We can alter the acceptance condition to be

\[\Delta^{*}(\mathbf{\Theta},\mathbf{\Theta}^{\prime})+X_{\text{nc}}+X_{\text{corr}}>0, \tag{13}\]

where \(X_{\text{corr}}\) is sampled when \(\sigma=1\) and \(X_{\text{nc}}\sim\mathcal{N}(0,1-\text{Var}[\Delta^{*}])\), requiring that \(\text{Var}[\Delta^{*}]<1\).

### Sampling the correction distribution

The distribution can be efficiently sampled using the CDF: the cumulative sum of the probability distribution function (PDF). The CDF is a monotonically increasing set of \(y\) values from 0 to 1. These values have corresponding \(X\) values in the domain \(-2V\) to \(2V\). In order to sample this distribution we draw a random number, \(u\), uniformly between 0 and 1. We find the \(X\) which corresponds to the intersection of \(u=C_{\sigma}(X)\) by bisection on the discretised points and then linear interpolation between discretised values.

## Appendix B Experiment Configurations

Here, we provide the exact parameters used to generate data provided in the results.

### Linear Perceptron

The linear perceptron model was trained on a simple 1D problem, which was generated via \(\mathbf{y}=m\mathbf{x}+c\), where \(x_{i}\sim\mathcal{U}(0,1)\) and \(m\sim\mathcal{U}(-1,1)\) and \(c\sim\mathcal{U}(-2,0)\). We randomly sampled 256 points to use in the distribution. The minibatch method used batch sizes of 32 and set \(\sigma=0.1\) for the coupling between models in the trajectories. Experiments were run for \(4\times 10^{7}\) epochs. Averages of observables were taken by discarding the first half of the data, to allow for a _burn-in_ time, and using minibatch estimates on the data.

Figure 3: Training curves for MNIST classification NNE. (a) Average ensemble loss per model in the NNE as a function of cumulative data usage \(D\), for \(s=50\) and various NNE sizes \(\tau\). Note that the abscissa is scaled by \(\tau\), so curves are in terms of “per-model” epochs. In terms of \(D/\tau\), for equivalent training time lower losses are reached for larger NNFs. (Loss curves has been down sampled for clarity.) Inset: same in linear \(D\) scale. (b) Same training curves but now plotted in terms of epochs \(E\).
Mnist

The convolutional neural network (CNN) model architecture used for our MNIST experiments was as follows:

1. Input \(28\times 28\) single channel image.
2. Convolution layer with a \(5\times 5\) kernel and 16 output channels.
3. \(2\times 2\) max pooling layer.
4. Convolution layer with a \(3\times 3\) kernel and 8 output channels.
5. \(4\times 4\) max pooling layer.
6. Fully connected dense layer with 10 outputs.
7. Softmax layer to normalise probabilities.

This model would output a normalised probability vector for each input image, specifying the "likelihood" of the image being a certain digit. The model contained 1906, 32-bit, floating point parameters.

For our experiments, we did not anneal the \(s\) parameter, but instead, chose a fixed duration of \(2\times 10^{7}\) epochs to allow for some "burn-in" time. The models were then run for another \(2\times 10^{7}\) epochs, to allow for measuring the loss as an observable. Accuracies were only measured at the end of the \(4\times 10^{7}\) epochs, due to the high computational demand.

We ran all of our experiments using \(\sigma=0.05\) and only changed a random 25% of the parameters of each model on each perturbation. Each perturbation changed only a single model in the trajectory, uniformly randomly.

We ran 6 independent experiments for each set of presented parameters and calculated averages to present the results in Figure 2, along with calculating the error bars using the averages' sample variance.

## Appendix C Normality Investigation

To justify setting the error threshold \(\delta\rightarrow\infty\), we ran a training experiment for two different \(\tau\) at \(s=50\) using a base batch size of 240. These experiments were run for \(20,000\) epochs and samples at intervals of \(5,000\) epochs. Histograms of the individual \(\Delta\) samples across the entire batch are plotted, along with a fitted curve showing the expected normal distribution given the empirical mean and variance of the batch data. This is presented in Figure 4.
