Deeper Connections between Neural Networks and Gaussian Processes Speed-up Active Learning

TL;DR

This paper introduces a method to improve active learning for neural networks by approximating Bayesian neural networks with Gaussian processes, enabling efficient uncertainty updates and better out-of-sample predictions, demonstrated on real-world large-scale data.

Contribution

The paper proposes approximating Bayesian neural networks with Gaussian processes to enhance active learning efficiency and accuracy without retraining.

Findings

Outperforms state-of-the-art active learning methods

Efficient uncertainty estimation without retraining

Effective on large-scale chemical and physical data

Abstract

Active learning methods for neural networks are usually based on greedy criteria which ultimately give a single new design point for the evaluation. Such an approach requires either some heuristics to sample a batch of design points at one active learning iteration, or retraining the neural network after adding each data point, which is computationally inefficient. Moreover, uncertainty estimates for neural networks sometimes are overconfident for the points lying far from the training sample. In this work we propose to approximate Bayesian neural networks (BNN) by Gaussian processes, which allows us to update the uncertainty estimates of predictions efficiently without retraining the neural network, while avoiding overconfident uncertainty prediction for out-of-sample points. In a series of experiments on real-world data including large-scale problems of chemical and physical modeling, we show superiority of the proposed approach over the state-of-the-art methods.

Tables

Table 1. Table 1: Summary of the datasets used in experiments with UCI data.

Dataset name	# of samples	# of attributes	Feature to predict
BlogFeedback Buza, (2014)	60021	281	Number of comments
SGEMM GPU Nugteren and Codreanu, (2015)	241600	18	Median calculation time
YearPredictionMSD Bertin-Mahieux et al., (2011)	515345	90	Year
Relative location of CT slices Graf et al., (2011)	53500	386	Relative location
Online News Popularity Fernandes et al., (2015)	39797	61	Number of shares
KEGG Network Shannon et al., (2003)	53414	24	Clustering coefficient

Full text

1 Introduction

In the modern applications of machine learning, especially in engineering design, physics, chemistry, molecular and materials modeling, the datasets are usually of limited size due to the expensive cost of computations. On the other side, such applications typically allow the calculation of additional data at points chosen by the experimentalist. Thus, there is the need to get the highest possible approximation quality with as few function evaluations as possible. Active learning methods help to achieve this goal by trying to select the best candidates for further target function computation using the already existing data and machine learning models based on these data.

Modern active learning methods are usually based either on ensemble-based methods Settles, (2012) or on probabilistic models such as Gaussian process regression Sacks et al., (1989); Burnaev and Panov, (2015). However, the models from these classes often don’t give state-of-the-art results in the downstream tasks. For example, Gaussian process-based models have very high computational complexity, which in many cases prohibits their usage even for moderate-sized datasets. Moreover, GP training is usually done for stationary covariance functions, which is a very limiting assumption in many real-world scenarios. Thus, active learning methods applicable to the broader classes of models are needed.

In this work, we aim to develop efficient active learning strategies for neural network models. Unlike Gaussian processes, neural networks are known for relatively good scalability with the dataset size and are very flexible in terms of dependencies they can model. However, active learning for neural networks is currently not very well developed, especially for tasks other than classification. Recently some uncertainty estimation and active learning approaches were introduced for Bayesian neural networks Gal et al., (2017); Hafner et al., (2018). However, still one of the main complications is in designing active learning approaches which allow for sampling points without retraining the neural network too often. Another problem is that neural networks due to their parametric structure sometimes give overconfident predictions in the areas of design space lying far from the training sample. Recently, Matthews et al., (2018) proved that deep neural networks with random weights converge to Gaussian processes in the infinite layer width limit. Our approach aims to approximate trained Bayesian neural networks and show that obtained uncertainty estimates enable to improve active learning performance significantly.

We summarize the main contributions of the paper as follows:

We propose to compute the approximation of trained Bayesian neural network with Gaussian process, which allows using the Gaussian process machinery to calculate uncertainty estimates for neural networks. We propose active learning strategy based on obtained uncertainty estimates which significantly speed-up the active learning with NNs by improving the quality of selected samples and decreasing the number of times the NN is retrained. 2. 2.

The proposed framework shows significant improvement over competing approaches in the wide range of real-world problems, including the cutting edge applications in chemoinformatics and physics of oil recovery.

In the next section, we discuss the problem statement and the proposed approach in detail.

2 Active learning with Bayesian neural networks and beyond

2.1 Problem statement

We consider a regression problem with an unknown function $f(\mathbf{x})$ defined on a subset of Euclidean space $\mathcal{X}\subset\mathbb{R}^{d}$, where an approximation function $\hat{f}(\mathbf{x})$ should be constructed based on noisy observations {EQA}[c] y = f(x) + ϵ with $\epsilon$ being some random noise.

We focus on active learning scenario which allows to iteratively enrich training set by computing the target function in design points specified by the experimenter. More specifically, we assume that the initial training set $D_{\rm{init}}=\{\mathbf{x}_{i},y_{i}=f(\mathbf{x}_{i})+\epsilon_{i}\}_{i=1}^{N}$ with precomputed function values is given. On top of it, we are given another set of points called “pool” $\mathcal{P}=\{\mathbf{x}_{j}\in\mathcal{X}\}_{j=1}^{N_{\rm{pool}}}$, which represents unlabelled data. Each point $\mathbf{x}\in\mathcal{P}$ can be annotated by computing $y=f(\mathbf{x})+\epsilon$ so that the pair $\{\mathbf{x},y\}$ is added to the training set.

The standard active learning approaches rely on the greedy point selection {EQA}[c]

xnew = argmaxx∈P A(x — ^f, D),

where $A(\mathbf{x}\,|\,\hat{f},D)$ is so-called acquisition function, which is usually constructed based on the current training set $D$ and corresponding approximation function $\hat{f}$.

The most popular choice of acquisition function is the variance $\hat{\sigma}^{2}(\mathbf{x}\,|\,\hat{f},D)$ of the prediction $\hat{f}(\mathbf{x})$, which can be easily estimated in some machine learning models such as Random forest Mentch and Hooker, (2016) or Gaussian process regression Rasmussen, (2004). However, for many types of models (e.g., neural networks) the computation of prediction uncertainty becomes a nontrivial problem.

2.2 Uncertainty estimation and active learning with Bayesian neural networks

In this work, we will stick to the Bayesian approach which treats neural networks as probabilistic models $p(y\,|\,\mathbf{x},\mathbf{w})$. Vector of neural network weights $\mathbf{w}$ is assumed to be a random variable with some prior distribution $p(\mathbf{w})$. The likelihood $p(y\,|\,\mathbf{x},\mathbf{w})$ determines the distribution of network output at a point $\mathbf{x}$ given specific values of parameters $\mathbf{w}$. There is vast literature on training Bayesian networks (see Graves, (2011) and Paisley et al., (2012) among many others), which mostly targets the so-called variational approximation of the intractable posterior distribution $p(\mathbf{w}\,|\,D)$ by some easily computable distribution $q(\mathbf{w})$.

The approximate posterior predictive distribution reads as: {EQA}[c] q(y  — x) = ∫p(y  — x, w)   q(w)   d w.

The simple way to generate random values from this distribution is to use the Monte-Carlo approach, which allows estimating the mean: {EQA}[c] E_q(y  — x)   y ≈1T ∑_t = 1^T ^f(x, w_t),

where the weight values $\mathbf{w}_{t}$ are i.i.d. random variables from distribution $q(\mathbf{w})$. Similarly, one can use Monte-Carlo to estimate the approximate posterior variance $\hat{\sigma}^{2}(\mathbf{x}\,|\,\hat{f})$ of the prediction $y$ at a point $\mathbf{x}$ and use it as an acquisition function: {EQA}[c] A(x — ^f, D) = ^σ^2(x — ^f).

We note that the considered acquisition function formally doesn’t depend on the dataset $D$ except for the fact that $D$ was used for training the neural network $\hat{f}$.

Let us note that the general greedy active learning approach (2.1) by design gives one candidate point per active learning iteration. If one tries to use (2.1) to obtain several samples with the same acquisition function, it usually results in obtaining several nearby points from the same region of design space, see Figure 1. Such behaviour is typically undesirable as nearby points are likely to have very similar information about the target function. Moreover, neural network uncertainty predictions are sometimes overconfident in out-of-sample regions of design space.

There are several approaches to overcome these issues each having its drawbacks:

One may retrain the model after each point addition which may result in a significant change of the acquisition function and lead to the selection of a more diverse set of points. However, such an approach is usually very computationally expensive, especially for neural network-based models. 2. 2.

One may try to add distance-based heuristic, which explicitly prohibits sampling points which are very close to each other and increase values of acquisition function for points positioned far from the training sample. Such an approach may give satisfactory results in some cases, however usually requires fine-tuning towards particular application (like the selection of specific distance function or choice the parameter value which determines whether two points are near or not), while its performance may degrade in high dimensional problems. 3. 3.

One may treat specially normalized vector of acquisition function values at points from the pool as a probability distribution and sample the desired number of points based on their probabilities (the higher acquisition function value, the point is more likely to be selected). This approach usually improves over greedy baseline procedure. However, it still gives many nearby points.

In the next section, we propose the approach to deal with problems above by considering the Gaussian process approximation of the neural network.

2.3 Gaussian process approximation of Bayesian neural network

Effectively, the random function {EQA}[c] ^f(x, w) = E_p(y  — x, w)   y

is the stochastic process indexed by $\mathbf{x}$. The covariance function of the process $\hat{f}(\mathbf{x},\mathbf{w})$ is given by {EQA}[c] k(x, x’) = E_q(w) (^f(x, w) - m(x)) (^f(x’, w) - m(x’)),

where $m(\mathbf{x})=\mathbb{E}_{q(\mathbf{w})}\hat{f}(\mathbf{x},\mathbf{w})$.

As was shown in Matthews et al., (2018); Lee et al., (2017) neural networks with random weights converge to Gaussian processes in the infinite layer width limit. However, one is not limited to asymptotic properties of purely random networks as Bayesian neural networks trained on real-world data exhibit near Gaussian behaviour, see the example on Figure 2.

We aim to make the Gaussian process approximation $\hat{g}(\mathbf{x}\,|\,\hat{f})$ of the stochastic process $\hat{f}(\mathbf{x},\mathbf{w})$ and compute its posterior variance $\hat{\sigma}^{2}(\mathbf{x}\,|\,\hat{f},X)$ given the set of anchor points $X=\{\mathbf{x}_{i}\}_{i=1}^{N}$. Typically, $X$ is a subset of the training sample. Given $X$, Monte-Carlo estimates $\hat{k}(\mathbf{x}^{\prime},\mathbf{x}^{\prime\prime})$ of the covariance function $k(\mathbf{x}^{\prime},\mathbf{x}^{\prime\prime})$ for every pair of points $\mathbf{x}^{\prime},\mathbf{x}^{\prime\prime}\in X\cup\mathbf{x}$ allow computing {EQA}[c]

^σ^2(x — ^f, X) = ^k(x, x) - ^k^T(x) ^K^-1 ^k(x),

where \hat{K}=\bigl{[}\hat{k}(\mathbf{x}_{i},\mathbf{x}_{j})\bigr{]}_{i,j=1}^{N} and \hat{\mathbf{k}}(\mathbf{x})=\bigl{(}\hat{k}(\mathbf{x}_{1},\mathbf{x}),\dots,\hat{k}(\mathbf{x}_{N},\mathbf{x})\bigr{)}^{\mathrm{T}}.

We note that only the trained neural network $\hat{f}(\mathbf{x},\mathbf{w})$ and the ability to sample from the distribution $q(\mathbf{w})$ are needed to compute $\hat{\sigma}^{2}(\mathbf{x}\,|\,\hat{f},X)$.

2.4 Active learning strategies

The benefits of the Gaussian process approximation and the usage of the formula (2) are not evident as one might directly estimate the variance of neural network prediction $\hat{f}(\mathbf{x},\mathbf{w})$ at any point $\mathbf{x}$ by sampling from $q(\mathbf{w})$ and use it as acquisition function. However, the approximate posterior variance $\hat{\sigma}^{2}(\mathbf{x}\,|\,\hat{f},X)$ of Gaussian process $\hat{g}(\mathbf{x}\,|\,\hat{f})$ has an important property that is has large values for points $\mathbf{x}$ lying far from the points from the training set $X$. Thus, out-of-sample points are likely to be selected by the active learning procedure.

Moreover, the function $\hat{\sigma}^{2}(\mathbf{x}\,|\,\hat{f},X)$ depends solely on covariance function values for points from a set $X$ (and not on the output function values). Such property allows updating uncertainty predictions by just adding sample points to the set $X$. More specifically, if we decide to sample some point $\mathbf{x}^{\prime}$, then the updated posterior variance $\hat{\sigma}^{2}(\mathbf{x}\,|\,\hat{f},X^{\prime})$ for $X^{\prime}=X\cup\mathbf{x}^{\prime}$ can be easily computed: {EQA}[c] ^σ^2(x — ^f, X’) = ^σ^2(x — ^f, X) - ^k2(x, x’  — ^f, X)^σ2(x’  — ^f, X),

where $\hat{k}(\mathbf{x},\mathbf{x}^{\prime}\,|\,\hat{f},X)=\hat{k}(\mathbf{x},\mathbf{x}^{\prime})-\hat{\mathbf{k}}^{\mathrm{T}}(\mathbf{x})\hat{K}^{-1}\hat{\mathbf{k}}(\mathbf{x}^{\prime})$ is the posterior covariance function of the process $\hat{g}(\mathbf{x}\,|\,\hat{f})$ given $X$.

Importantly, $\hat{\sigma}^{2}(\mathbf{x}\,|\,\hat{f},X^{\prime})$ for points $\mathbf{x}$ in some vicinity of $\mathbf{x}^{\prime}$ will have low values, which guarantees that further sampled points will not lie too close to $\mathbf{x}^{\prime}$ and other points from the training set $X$. The resulting NNGP active learning procedure is depicted in Figure 3.

3 Experiments

Our experimental study is focused on real-world data to ensure that the algorithms are indeed useful in real-world scenarios. We compare the following algorithms:

• •

pure random sampling;

• •

sampling based on the variance of NN stochastic output from $\hat{f}$, which we refer to as MCDUE (see Gal et al., (2017); Tsymbalov et al., (2018));

• •

the proposed GP-based approaches (NNGP).

The more detailed descriptions of these methods can be found in Appendix A.

Following Gal and Ghahramani, (2016) we use the Bayesian NN with the Bernoulli distribution on the weights, which is equivalent to using the dropout on the inference stage.

3.1 Airline delays dataset and NCP comparison

We start the experiments from comparing the proposed approach to active learning with the one based on uncertainty estimates obtained from a Bayesian neural network with Noise Contrastive Prior (NCP), see Hafner et al., (2018). Following this paper, we use the airline delays dataset (see Hensman et al., (2013)) and NN consisting of two layers with 50 neurons each, leaky ReLU activation function, and trained with respect to NCP-based loss function.

We took a random subset of 50,000 data samples from the data available on the January – April of 2008 as a training set and we chose 100,000 random data samples from May of 2008 as a test set. We used the following variables as input features PlaneAge, Distance, CRSDepTime, AirTime, CRSArrTime, DayOfWeek, DayofMonth, Month and ArrDelay + DepDelay as a target.

The results for the test set are shown in Figure 4. The proposed NNGP approach demonstrates a comparable error with respect to the previous results and outperforms (on average) other methods in the continuous active learning scenario.

3.2 Experiments on UCI datasets

We conducted a series of experiments with active learning performed on the data from the UCI ML repository Dua and Taniskidou, (2017). All the datasets represent real-world regression problems with $15+$ dimensions and $30000+$ samples, see Table 1 for details.

For every experiment, data are shuffled and split in the following proportions: 10% for the training set $D_{train}$, 5% for the test set $D_{test}$, 5% for the validation set $D_{val}$ needed for early-stopping and 80% for the pool $\mathcal{P}$.

We used a simple neural network with three hidden layers of sizes 256, 128 and 128. The details on the neural network training procedure can be found in Appendix B. We performed $16$ active learning iterations with 200 points picked at each iteration.

To compare the performance of the algorithms across the different datasets and different choices of training samples, we will use so-called Dolan-More curves. Let $q^{p}_{a}$ be an error measure of the $a$-th algorithm on the $\mathcal{P}$-th problem. Then, determining the performance ratio $r^{p}_{a}=\frac{q^{p}_{a}}{\min_{x}(q^{p}_{x})}$, we can define the Dolan-More curve as a function of the performance ratio factor $\tau$: {EQA}[c] ρ_a(τ) = #(p:rpa≤τ)np,

where $n_{p}$ is the total number of evaluations for the problem $p$. Thus, $\rho_{a}(\tau)$ defines the fraction of problems in which the $a$-th algorithm has the error not more than $\tau$ times bigger than the best competitor in the chosen performance metric. Note that $\rho_{a}(1)$ is the ratio of problems on which the $a$-th algorithm performance was the best, while in general, the higher curve means the better performance of the algorithm.

The Dolan-More curves for the errors of approximation for considered problems after the 16th iteration of the active learning procedure are presented in Figure 5. We see that the NNGP procedure is superior in terms of RMSE compared to MCDUE and random sampling.

3.3 SchNet training

To demonstrate the power of our approach, we conducted a series of numerical experiments with the state-of-the-art neural network architecture in the field of chemoinformatics “SchNet” Schütt et al., (2017). This network takes information about an organic molecule as an input, and, after special preprocessing and complicated training procedure, outputs some properties of the molecule (like energy). Despite its complex structure, SchNet contains fully connected layers, so it is possible to use a dropout in between them.

We tested our approach on the problem of predicting the internal energy of the molecule at 0K from the QM9 data set Ramakrishnan et al., (2014). We used a Tensorflow implementation of a SchNet with the same architecture as in original paper except for an increased size of hidden layers (from 64 and 32 units to 256 and 128 units, respectively) and dropout layer placed in between of them and turned on during an inference only.

In our experiment, we separate the whole dataset of 133 885 molecules into the initial set of 10 000 molecules, the test set of 5 000 molecules, and the rest of the data allocated as the pool. On each active learning iteration, we perform 100 000 training epochs and then calculate the uncertainty estimates using either MCDUE or NNGP approach. We then select 2 500 molecules with the highest uncertainty from the pool, add them to the training set and perform another active learning iteration.

The results are shown in Figure 6. The NNGP approach demonstrates the most steady decrease in error, with the 25% accuracy increase in RMSE. Such improvement is very significant in terms of the time savings for the computationally expensive quantum-mechanical calculations. For example, to reach the RMSE of 2 kcal/mol starting from the SchNet trained on 10 000 molecules, one need to additionally sample 15 000 molecules in case of random sampling or just 7 500 molecules using the NNGP uncertainty estimation procedure.

3.4 Hydraulic simulator

In the oil industry, to determine the optimal parameters and control the drilling process, engineers carry out hydraulic calculations of the well’s circulation system, that usually are based on either empirical formulas or semi-analytical solutions of the hydrodynamic equations. However, such semi-analytical solutions are known just for few individual cases, while in the other ones, only very crude approximations are usually available (see Podryabinkin et al., (2013) for details). As a result, such calculations have relatively low precision. On the other hand, the full-scale numerical solution of the hydrodynamic equations describing the flow of drilling fluids can provide a sufficient level of accuracy, but it requires significant computational resources and subsequently is very costly and time-consuming. The possible solution to this problem is the use of a surrogate model.

We used a surrogate model for the fluid flow in the well-bore while drilling. The oracle, in this case, is a numerical solver for the hydrodynamic equations Podryabinkin et al., (2013), which, given six main parameters of the drilling process as an input, outputs the unit-less hydraulic resistance coefficient that characterizes the drop in the pressure.

In this experiment, we used a two-layer neural network with $50$ neurons per each layer and LeakyReLU activation function. Initial training and pool sets had $50$ and $20000$ points respectively. We completed $10$ active learning iterations, adding $50$ points per each iteration. The results are shown in Figure 7. Clearly, NNGP is superior in terms of RMSE and MAE, while maximum error for NNGP is 10 times lower.

4 Related work

4.1 Active learning

Active learning Settles, (2012) (also known as adaptive design of experiments Forrester et al., (2008) in statistics and engineering design) is a framework which allows the additional data points to be annotated by computing target function value and then added to the training set. The particular points to sample are usually chosen as the ones that maximize so-called acquisition function. The most popular approaches to construct acquisition function are committee-based methods Seung et al., (1992), also known as query-by-committee, where an ensemble of models is trained on the same or various parts of data, and Bayesian models, such as Gaussian Processes Rasmussen, (2004), where the uncertainty estimates can be directly inferred from the model Sacks et al., (1989); Burnaev and Panov, (2015). However, the majority of existing approaches are computationally expensive and sometimes even intractable for large sample sizes and input dimensions.

Active learning for neural networks is not very well developed as neural networks usually excel in applications with large datasets already available. Ensembles of neural networks (see Li et al., (2018) for a detailed review) often boil down to an independent training of several models, which works well in some applications Beluch et al., (2018), but is computationally expensive for the large-scale applications. Bayesian neural networks provide uncertainty estimates, which can be efficiently used for active learning procedures Gal et al., (2017); Hafner et al., (2018). The most popular uncertainty estimation approach for BNNs, MC-Dropout, is based on classical dropout first proposed as a technique for a neural network regularization, which was recently interpreted as a method for approximate inference in Bayesian neural networks and shown to provide unbiased Monte-Carlo estimates predictive variance Gal and Ghahramani, (2016). However, it was shown Beluch et al., (2018) that uncertainty estimates based on MC-Dropout are usually less efficient for active learning then those based on ensembles. Pop and Fulop, (2018) suggests that it is partially due to the mode collapse effect, which leads to overconfident out of sample predictions of uncertainty.

Connections between neural networks and Gaussian processes recently gain significant attention, see Matthews et al., (2018); Lee et al., (2017) which study random, untrained NNs and show that such networks can be approximated by Gaussian processes in infinite network width limit. Another direction is incorporation of the GP-like elements into the NN structure, see Sun et al., (2018); Garnelo et al., (2018) for some recent contributions. However, we are not aware of any contributions making practical use of GP behaviour for standard neural network architectures.

5 Summary and discussion

We have proposed a novel dropout-based method for the uncertainty estimation for deep neural networks, which uses the approximation of neural network by Gaussian process. Experiments on different architectures and real-world problems show that the proposed estimate allows to achieve state-of-the-art results in context of active learning. Importantly, the proposed approach works for any neural network architecture involving dropout (as well as other Bayesian neural networks), so it can be applied to very wide range of networks and problems without a need to change neural network architecture.

It is of interest whether the proposed approach can be efficient for other applications areas, such as image classification. We also plan to study the applicability of modern methods for GP speed-up in order to improve the scalability of proposed approach.

Acknowledgements

The research was supported by the Skoltech NGP Program No. 2016-7/NGP (a Skoltech-MIT joint project).

Appendix A Details on active learning approaches

Here we provide descriptions for the uncertainty estimation and active learning algorithm used in experiments: MCDUE described in Gal and Ghahramani, (2016); Tsymbalov et al., (2018) and NNGP / M-step NNGP proposed in this paper.

Appendix B NN training details

This section provides details on NN training for UCI dataset experiments.

Neural network had three hidden layers with sizes 256, 128, 128, respectively. Learning rate started at $10^{-3}$, its decay was set to 0.97 and changed every 50000 epochs. Minimal learning rate was set to $10^{-5}$. We reset learning rate for each active learning algorithm in hope of beating the local minima problem. Training dropout rate was set to 0.1. $L_{2}$ regularization was set to $10^{-4}$. Batch size set to 200.

We developed rather complex active learning procedure in order to balance between the stochastic nature of NN training and early stopping.

The algorithm we followed for each experiment on a fixed dataset to initialize the training is as follows:

Initialize NN. Shuffle and split the data. Train for a mandatory number of epochs $epochs_{mandatory}=10000$. Set $warnings=0$, $E^{previous}_{val}=10^{10}$. 2. 2.

For every epoch $current\_epoch$ in $1,\ldots,epochs_{max}=10^{6}$:

1. (a)

Train NN. 2. (b)

If $current\_epoch$ % $step_{ES\_check}=100$:

1. i.

Get RMSE error $E_{val}$ on validation set. 2. ii.

If current $E_{val}$ exceeds the $E^{previous}_{val}$ by $ES_{window}=1\%$: then $warnings:=warnings+1$. Else set $warnings:=0,E^{previous}_{val}:=E_{val}$. 3. iii.

If $warnings>warnings_{max}=3$: break from training procedure.