Differentiable PAC–Bayes Objectives with Partially Aggregated Neural Networks

Abstract

We make two related contributions motivated by the challenge of training stochastic neural networks, particularly in a PAC–Bayesian setting: (1) we show how averaging over an ensemble of stochastic neural networks enables a new class of partially-aggregated estimators, proving that these lead to unbiased lower-variance output and gradient estimators; (2) we reformulate a PAC–Bayesian bound for signed-output networks to derive in combination with the above a directly optimisable, differentiable objective and a generalisation guarantee, without using a surrogate loss or loosening the bound. We show empirically that this leads to competitive generalisation guarantees and compares favourably to other methods for training such networks. Finally, we note that the above leads to a simpler PAC–Bayesian training scheme for sign-activation networks than previous work.

1. Introduction

The use of stochastic neural networks has become widespread in the PAC–Bayesian and Bayesian deep learning [1] literature as a way to quantify predictive uncertainty and obtain generalisation bounds. PAC–Bayesian theorems generally bound the expected loss of randomised estimators, so it has proven easier to obtain non-vacuous numerical guarantees on generalisation in such networks.

However, we observe that when training these in the PAC–Bayesian setting, the objective used is generally somewhat divorced from the bound on misclassification loss itself, often because non-differentiability leads to difficulties with direct optimisation. For example, Langford and Caruana [2], Zhou et al. [3], and Dziugaite and Roy [4] all initially train non-stochastic networks before using them as the mode of a distribution, with variance chosen, respectively, through a computationally-expensive sensitivity analysis, as a proportion of weight norms, or by optimising an objective with both a surrogate loss function and a different dependence on the Kullback–Leibler (KL) divergence from their bound.

In exploring methods to circumvent this gap, we also note that PAC–Bayesian bounds can often be straightforwardly adapted to aggregates or averages of estimators, leading directly to analytic and differentiable objective functions (for example, [5]). Unfortunately, averages over deep stochastic networks are usually intractable or, if possible, very costly (as found by [6]).

Motivated by these observations, our main contribution is to obtain a compromise by defining new and general “partially-aggregated” Monte Carlo estimators for the average output and gradients of deep stochastic networks (Section 3), with the direct optimisation of PAC–Bayesian bounds in mind. Although our main focus here is on the use of this estimator in a PAC–Bayesian application, we emphasise that the technique applies generally to stochastic networks and thus has links to other variance-reduction techniques for training them, such as the pathwise estimator used in the context of neural networks by [7] amongst many others or Flipout [8]; indeed, it can be used in combination with these techniques. We provide proofs (Section 4) that this application leads to lower variances than a Monte Carlo forward pass and lower variance final-layer gradients than REINFORCE [9].

A further contribution of ours is a first application of this general estimator to non-differentiable “signed-output” networks (with a final output $\in \{-1,+1\}$ and arbitrarily complex other structure, see Section 4). As well as reducing variances as stated above, a small amount of additional structure in combination with partial-aggregation enables us to extend the pathwise estimator to the other layers, which usually requires a fully differentiable network and eases training by reducing the variance of gradient estimates.

We adapt a binary classification bound (Section 5) from Catoni [10] to these networks, yielding straightforward and directly differentiable objectives when used in combination with aggregation. Closing this gap between objectives and bounds leads to improved theoretical properties.

Further, since most of the existing PAC–Bayes bounds for neural networks have a heavy dependency on the distance from initialisation of the obtained solution, we would intuitively expect these lower variances to lead to faster convergence and tighter bounds (from finding low-error solutions nearer to the initialisation). We indeed observe this experimentally, showing that training PAC–Bayesian objectives in combination with partial aggregation leads to competitive experimental generalisation guarantees (Section 6), and improves upon naive Monte Carlo and REINFORCE.

As a useful corollary, this application also leads us to a similar but simpler PAC–Bayesian training method for sign-activation neural networks than Letarte et al. [6], which successfully aggregated networks with all sign activation functions $\in \{+1,-1\}$ and a non-standard tree structure, but incurred an exponential KL divergence penalty and a heavy computational cost (so that in practice they often resorted to a Monte Carlo estimate). Further, the lower variance of our obtained estimator predictions enables us to use the Gibbs estimator directly (where we draw a single sample function for every new example), leading to a modified bound on the misclassification loss which is a factor of two tighter without a significant performance penalty.

We discuss further and outline future work in Section 7.

2. Background

We begin here by setting out our notation and the requisite background.

Generally, we consider parameterised functions, $\{f_{\theta }:\mathcal{X}\to \mathcal{Y}|\theta \in \Theta \subset {\mathbb{R}}^N\}$, in a specific form, choosing $\mathcal{X}\subset {\mathbb{R}}^{d_0}$ and an arbitrary output space $\mathcal{Y}$ which could be for example $\{-1,+1\}$ or $\mathbb{R}$. We wish to find functions minimizing the out-of-sample risk, $R\left(f\right)={\mathbb{E}}_{(x,y)\sim \mathcal{D}}\ell (f\left(\mathit{x}\right),y)$, for some loss function ℓ, for example the 0-1 misclassification loss for classification, ${\ell }_{0-1}(y,y^{\prime })=\mathit{1}\{y\ne y^{\prime }\}$, or the binary linear loss, ${\ell }_{\mathrm{lin}}(y,y^{\prime })=\frac{1}{2}(1-y{y}^{\prime })$, with $\mathcal{Y}=\{+1,-1\}$. These must be chosen based on an i.i.d. sample $S={\left\{({\mathit{x}}_i,y_i)\right\}}_{i=1}^m\sim {\mathcal{D}}^m$ from the data distribution $\mathcal{D}$, using the surrogate of in-sample empirical risk, $R_S\left(f\right)=\frac{1}{m}{\sum }_{i=1}^m\ell (f\left({\mathit{x}}_i\right),y_i)$. We denote the expected and empirical risks under the misclassification and linear losses, respectively $R^{0-1},R^{\mathrm{lin}},R_S^{0-1}$ and $R_S^{\mathrm{lin}}$.

In this paper, we consider learning a distribution (PAC–Bayesian posterior), Q, over the parameters $\theta$. PAC–Bayesian theorems then provide bounds on the expected generalization risk of randomised classifiers, where every prediction is made using a newly sampled function from our posterior, $f_{\theta },\theta \sim Q$.

We also consider averaging the above to obtain Q-aggregated prediction functions,

$$F_Q\left(\mathit{x}\right):={\mathbb{E}}_{\theta \sim Q}{f}_{\theta }\left(\mathit{x}\right).$$

(1)

In the case of a convex loss function, Jensen’s inequality lower bounds the risk of the randomised function by its Q-aggregate: $\ell (F_Q\left(x\right),y)\le {\mathbb{E}}_{f\sim Q}\ell (f\left(x\right),y)$. The equality is achieved by the linear loss, a fact we will exploit to obtain an easier PAC–Bayesian optimisation objective in Section 5.

2.1. Analytic Q-Aggregates for Signed Linear Functions

Q-aggregate predictors are analytically tractable for “signed-output” functions (here the sign function and “signed” functions have outputs $\in \{+1,-1\}$, as the terminology “binary”, used sometimes in the literature, suggests to us too strongly an output $\in \{0,1\}$) of the form $f_w\left(\mathit{x}\right)=sign(\mathit{w}·\mathit{x})$ under a normal distribution, $Q\left(w\right)=N(\mathit{\mu },\mathbb{I})$, as specifically considered in a PAC–Bayesian context for binary classification by [5], obtaining an differentiable objective similar to the SVM. Provided $\mathit{x}\ne \mathit{0}$:

$$F_Q\left(\mathit{x}\right):={\mathbb{E}}_{w\sim N(\mathit{\mu },\mathbb{I})}sign(\mathit{w}·\mathit{x})=erf\left(\frac{\mathit{\mu }·\mathit{x}}{\sqrt{2}\parallel \mathit{x}\parallel }\right).$$

(2)

In Section 4, we will consider aggregating signed output ($f\left(x\right)\in \{+1,-1\}$) functions of a more general form.

2.2. Monte Carlo Estimators for More Complex Q-Aggregates

The framework of Q-aggregates can be extended to less tractable cases (for example, with $f_{\theta }$ a randomised or a “Bayesian” neural network, see, e.g., [1]) through a simple and unbiased Monte Carlo approximation:

$$F_Q\left(\mathit{x}\right)={\mathbb{E}}_{\theta \sim Q}{f}_{\theta }\left(\mathit{x}\right)\approx \frac{1}{T}\sum _{t=1}^T{f}_{{\theta }^t}\left(\mathit{x}\right):={\widehat{F}}_Q\left(\mathit{x}\right).$$

(3)

If we go on to parameterize our posterior Q by $\varphi \in \Phi \subset {\mathbb{R}}^N$ as $Q_{\varphi }$ and wish to obtain gradients without a closed form for $F_{Q_{\varphi }}\left(\mathit{x}\right)={\mathbb{E}}_{\theta \sim Q_{\varphi }}{f}_{\theta }\left(\mathit{x}\right)$, there are two possibilities. One is REINFORCE [9], which requires only a differentiable density function, $q_{\varphi }\left(\theta \right)$ and makes a Monte Carlo approximation to the left hand side of the identity ${\nabla }_{\varphi }{\mathbb{E}}_{\theta \sim q_{\varphi }}{f}_{\theta }\left(\mathit{x}\right)={\mathbb{E}}_{\theta \sim q_{\varphi }}[f_{\theta }\left(\mathit{x}\right){\nabla }_{\varphi }log{q}_{\varphi }\left(\theta \right)]$.

The other is the pathwise estimator, which additionally requires that $f_{\theta }\left(\mathit{x}\right)$ be differentiable w.r.t. $\theta$, and that the probability distribution chosen has a standardization function, $S_{\varphi }$, which removes the $\varphi$ dependence, turning a parameterised $q_{\varphi }$ into a non-parameterised distribution p: for example, $S_{\mu ,\sigma }\left(X\right)=(X-\mu )/\sigma$ to transform a general normal distribution into a standard normal. If this exists, the right hand side of ${\nabla }_{\varphi }{\mathbb{E}}_{\theta \sim q_{\varphi }}{f}_{\theta }\left(\mathit{x}\right)={\mathbb{E}}_{ϵ\sim p}{\nabla }_{\varphi }{f}_{S_{\varphi }^{-1}\left(ϵ\right)}\left(\mathit{x}\right)$ generally yields lower-variance estimates than REINFORCE (see for a modern survey [11]).

The variance introduced by REINFORCE can make it difficult to train neural networks when the pathwise estimator is not available, for example when non-differentiable activation functions are used. Below we find a compromise between the analytically closed form of (2) and the above estimator that enables us to make differentiable certain classes of network and extend the pathwise estimator where otherwise it could not be used. Through this we are able to stably train a new class of network.

2.3. PAC–Bayesian Approach

We use PAC–Bayes in this paper to obtain generalisation guarantees and theoretically-motivated training methods. The primary bound utilised is based on the following theorem, valid for a loss taking values in $[0,1]$:

Theorem 1

([10], Theorem 1.2.6). Given probability measure P on hypothesis space $\mathcal{F}$ and $\alpha >1$, for all Q on $\mathcal{F}$ with probability at least $1-\delta$ over $S\sim {\mathcal{D}}^m$,

$${\mathbb{E}}_{f\sim Q}R\left(f\right)\le \underset{\lambda >1}{inf}{\Phi }_{\lambda /m}^{-1}\left[{\mathbb{E}}_{f\sim Q}{R}_S\left(f\right)+\frac{\alpha }{\lambda }\Delta \right]$$

with ${\Phi }_{\gamma }^{-1}\left(t\right)=\frac{1-exp(-\gamma t)}{1-exp(-\gamma )}$ and $\Delta =KL\left(Q|P\right)-log\delta +2log\left(\frac{log{\alpha }^2\lambda }{log\alpha }\right)$.

This slightly opaque formulation (used previously by [3]) gives essentially identical results when $\mathrm{KL}/m$ is large to the better-known “small-kl” PAC–Bayes bounds originated by Langford and Seeger [12], Seeger et al. [13]. It is chosen because it leads to objectives that are linear in the empirical loss and KL divergence, like

$${\mathbb{E}}_{f\sim Q}{R}_S\left(f\right)+\frac{KL\left(Q|P\right)}{\lambda }.$$

(4)

This objective is minimised by a Gibbs distribution and is closely related to the evidence lower bound (ELBO) usually optimised by Bayesian Neural Networks [1]. Such a connection has been noted throughout the PAC–Bayesian literature; we refer the reader to [14] or [15] for a formalised treatment.

3. The Partial Aggregation Estimator

Here we outline our main contribution: a reformulation of Q-aggregation for neural networks leading to different, lower-variance, Monte Carlo estimators for their outputs and gradients. These estimators apply to networks with a dense final layer, and arbitrary stochastic other structure (for example convolutions, residual layers or a non-feedforward structure). Specifically, they take the form

$$f_{\theta }\left(\mathit{x}\right)=A(\mathit{w}·{\mathit{\eta }}_{{\theta }^{\neg \mathit{w}}}\left(\mathit{x}\right))$$

(5)

with $\theta =vec(\mathit{w},{\theta }^{\neg \mathit{w}})\in \Theta \subset {\mathbb{R}}^D$, $\mathit{w}\in {\mathbb{R}}^d$, and ${\theta }^{\neg \mathit{w}}\in {\Theta }^{\neg \mathit{w}}\subset {\mathbb{R}}^{D-d}$ the parameter set excluding $\mathit{w}$, for the non-final layers of the network. These non-final layers are included in ${\mathit{\eta }}_{{\theta }^{\neg \mathit{w}}}:\mathcal{X}\to {\mathcal{A}}^d\subseteq {\mathbb{R}}^d$ and the final activation is $A:\mathbb{R}\to \mathcal{Y}$. For simplicity we have used a one-dimensional output but we note that the formulation and below derivations trivially extend to a vector-valued function with elementwise activations. We require the distribution over parameters to factorise like $Q\left(\theta \right)=Q^w\left(\mathit{w}\right)Q^{\neg \mathit{w}}\left({\theta }^{\neg \mathit{w}}\right)$, which is consistent with the literature.

We recover a similar functional form to that considered in Section 2.1 by rewriting the function as $A(\mathit{w}·\mathit{a})$ with $\mathit{a}\in {\mathcal{A}}^d$ the randomised hidden-layer activations. The “aggregated” activation function on the final layer, which we define as $I\left(\mathit{a}\right):=\int A(\mathit{w}·\mathit{a})\mathrm{d}{Q}^{\mathit{w}}\left(\mathit{w}\right)$, may then be analytically tractable. For example, with $\mathit{w}\sim N(\mathit{\mu },\mathbb{I})$ and a sign final activation, we recall (2) where $I\left(\mathit{a}\right)=erf\left(\frac{\mathit{\mu }·\mathit{a}}{\sqrt{2}\parallel \mathit{a}\parallel }\right)$.

Using these definitions we can write the Q-aggregate in terms of the conditional distribution on the activations, $\mathit{a}$, which takes the form ${\tilde{Q}}^{\neg \mathit{w}}\left(\mathit{a}\right|\mathit{x}):=\left({\mathit{\eta }}_{(·)}\left(\mathit{x}\right)\right)\circ Q^{\neg \mathit{w}}$, (i.e., the distribution of ${\mathit{\eta }}_{{\theta }^{\neg \mathit{w}}}\left(\mathit{x}\right)|\mathit{x}$, with ${\theta }^{\neg \mathit{w}}\sim Q^{\neg \mathit{w}}$). The Q-aggregate can then be stated as

$$\begin{array}{cc}\hfill F_Q\left(\mathit{x}\right)& :={\mathbb{E}}_{\theta \sim Q}\left[f_{\theta }\left(\mathit{x}\right)\right]\hfill \\ \hfill & ={\int }_{{\theta }^{\neg \mathit{w}}}\left[{\int }_{{\mathbb{R}}^d}A(\mathit{w}·{\mathit{\eta }}_{{\theta }^{\neg \mathit{w}}}\left(\mathit{x}\right))\mathrm{d}{Q}^{\mathit{w}}\left(\mathit{w}\right)\right]\mathrm{d}{Q}^{\neg \mathit{w}}\left({\theta }^{\neg \mathit{w}}\right)\hfill \\ \hfill & ={\int }_{{\theta }^{\neg \mathit{w}}}I\left({\mathit{\eta }}_{{\theta }^{\neg \mathit{w}}}\left(\mathit{x}\right)\right)\mathrm{d}{Q}^{\neg \mathit{w}}\left({\theta }^{\neg \mathit{w}}\right)\hfill \\ \hfill & ={\int }_{{\mathcal{A}}^d}I\left(\mathit{a}\right)\mathrm{d}\{\left({\mathit{\eta }}_{(·)}\left(\mathit{x}\right)\right)\circ Q^{\neg \mathit{w}}\}\left(\mathit{a}\right)\hfill \\ \hfill & =:{\int }_{{\mathcal{A}}^d}I\left(\mathit{a}\right)\mathrm{d}{\tilde{Q}}^{\neg \mathit{w}}\left(\mathit{a}\right|\mathit{x}).\hfill \end{array}$$

In most cases, the final integral cannot be calculated exactly or involves a large summation, so we resort to a Monte Carlo estimate, for each $\mathit{x}$ drawing T samples of the randomised activations, ${\left\{{\mathit{a}}^t\right\}}_{t=1}^T\sim {\tilde{Q}}^{\neg w}\left(\mathit{a}\right|\mathit{x})$ to obtain the “partially-aggregated” estimator

$$F_Q\left(\mathit{x}\right)={\int }_{{\mathcal{A}}^d}I\left(\mathit{a}\right)\mathrm{d}{\tilde{Q}}^{\neg \mathit{w}}\left(\mathit{a}\right|\mathit{x})\approx \frac{1}{T}\sum _{t=1}^{T}I\left({\mathit{a}}^t\right)={\widehat{F}}_Q^{\ast }\left(\mathit{x}\right).$$

(6)

This is quite similar to the original estimator from (3), but in fact the aggregation of the final layer may significantly reduce the variance of the outputs and also make better gradient estimates possible, as we will show below.

3.1. Reduced Variance Estimates

Proposition 1.

Lower variance outputs: For a neural network as defined by Equation (5) and the unbiased Q-aggregation estimators defined by Equations (3) and (6),

$${\mathbb{V}}_Q\left[{\widehat{F}}_Q^{\ast }\left(\mathit{x}\right)\right]\le {\mathbb{V}}_Q\left[{\widehat{F}}_Q\left(\mathit{x}\right)\right].$$

Proof. 

Treating $\mathit{a}$ as a random variable, always conditioned on $\mathit{x}$, we have

$$\begin{array}{cc}\hfill & {\mathbb{V}}_Q\left[{\widehat{F}}_Q\left(\mathit{x}\right)\right]-{\mathbb{V}}_Q\left[{\widehat{F}}_Q^{☆}\left(\mathit{x}\right)\right]={\mathbb{E}}_Q|{\widehat{F}}_Q{\left(\mathit{x}\right)|}^2-{\mathbb{E}}_Q{\left|{\widehat{F}}_Q^{☆}\left(\mathit{x}\right)\right|}^2\hfill \\ \hfill & =\frac{1}{T}{\mathbb{E}}_{\mathit{a}|\mathit{x}}\left[{\mathbb{E}}_{\mathit{w}}{\left|A(\mathit{w}·\mathit{a})\right|}^2-{\left|{\mathbb{E}}_{\mathit{w}}A(\mathit{w}·\mathit{a})\right|}^2\right]\hfill \\ \hfill & =\frac{1}{T}{\mathbb{E}}_{\mathit{a}|\mathit{x}}\left[{\mathbb{V}}_{\mathit{w}}\left[A(\mathit{w}·\mathit{a})\right]\right]\ge 0.\hfill \end{array}$$

 □

From the above we see that the aggregate outputs estimated through partial-aggregation have lower variances. Next, we consider the two unbiased gradient estimators for the distribution over final-layer weights, $\mathit{w}$, arising from partial-aggregation or REINFORCE (as would be used, for example, where the final layer is non-differentiable). Assuming $Q^{\mathit{w}}$ has a density, $q_{\varphi }\left({\theta }^{\mathit{w}}\right)$, parameterised by $\varphi$, these use forward samples of ${\{{\mathit{w}}^t,{\theta }^{\neg \mathit{w},\left(t\right)}\}}_{t=1}^T$ as:

$$\begin{array}{cc}\hfill \widehat{\mathit{G}}\left(\mathit{x}\right)& :=\frac{1}{T}\sum _{t=1}^{T}A({\mathit{w}}^t·{\mathit{\eta }}^t){\nabla }_{\varphi }log{q}_{\varphi }\left({\mathit{w}}^t\right)\hfill \\ \hfill {\widehat{\mathit{G}}}^{\ast }\left(\mathit{x}\right)& :=\frac{1}{T}\sum _{t=1}^T{\nabla }_{\varphi }{I}_{q_{\varphi }}\left({\mathit{\eta }}^t\right).\hfill \end{array}$$

Proposition 2.

Lower variance gradients: Under the conditions of Proposition 1 and the above definitions,

$${\mathrm{Cov}}_Q\left[{\widehat{\mathit{G}}}^{\ast }\left(\mathit{x}\right)\right]⪯{\mathrm{Cov}}_Q\left[\widehat{\mathit{G}}\left(\mathit{x}\right)\right]$$

where $A⪯B\iff B-A$ is positive semi-definite. Thus, for all $\mathit{u}\ne \mathit{0}$, $\mathbb{V}[{\widehat{\mathit{G}}}^{\ast }\left(\mathit{x}\right)·\mathit{u}]\le \mathbb{V}[\widehat{\mathit{G}}\left(\mathit{x}\right)·\mathit{u}]$.

Proof. 

Writing $\mathit{v}:={\nabla }_{\varphi }log{q}_{\varphi }\left(\mathit{w}\right)$ and using the unbiasedness of the estimators,

$$\begin{array}{cc}\hfill & {\mathrm{Cov}}_Q\left[{\widehat{\mathit{G}}}_Q\left(\mathit{x}\right)\right]-{\mathrm{Cov}}_Q\left[{\widehat{\mathit{G}}}_Q^{☆}\left(\mathit{x}\right)\right]\hfill \\ \hfill & ={\mathbb{E}}_Q\left[{\widehat{\mathit{G}}}_Q\left(\mathit{x}\right){\widehat{\mathit{G}}}_Q{\left(\mathit{x}\right)}^T\right]-{\mathbb{E}}_Q\left[{\widehat{\mathit{G}}}_Q^{☆}\left(\mathit{x}\right){\widehat{\mathit{G}}}_Q^{☆}{\left(\mathit{x}\right)}^T\right]\hfill \\ \hfill & =\frac{1}{T}{\mathbb{E}}_{\mathit{a}|\mathit{x}}\left[{\mathbb{E}}_w\left[A{(\mathit{w}·\mathit{a})}^2\mathit{v}{\mathit{v}}^T\right]-{\nabla }_{\varphi }{I}_{q_{\varphi }}\left({\mathit{\eta }}^t\right)\left({\nabla }_{\varphi }{I}_{q_{\varphi }}{\left({\mathit{\eta }}^t\right)}^T\right)\right]\hfill \\ \hfill & =\frac{1}{T}{\mathbb{E}}_{\mathit{a}|\mathit{x}}\left[{\mathrm{Cov}}_{\mathit{w}}[A(\mathit{w}·\mathit{a}){\nabla }_{\varphi }log{q}_{\varphi }\left(\mathit{w}\right)]\right]⪰0\hfill \end{array}$$

where in the final line we have used that ${\nabla }_{\varphi }{I}_{q_{\varphi }}\left({\mathit{\eta }}^t\right)={\nabla }_{\varphi }{\mathbb{E}}_w\left[A(w·{\mathit{\eta }}^t)\right]={\mathbb{E}}_{\mathit{w}}\left[A(\mathit{w}·{\mathit{\eta }}^t)\mathit{v}\right]$. □

3.2. Single Hidden Layer

For clarity (and to introduce notation to be used in Section 4.2) we will briefly consider the case of a neural network with one hidden layer, $f_{\theta }\left(\mathit{x}\right)=A_2({\mathit{w}}_2·{\mathit{A}}_1\left(W_1\mathit{x}\right))$. The randomised parameters are $\theta =vec({\mathit{w}}_2,W_1)$, $W_1\in {\mathbb{R}}^{d_1\times d_0}$, ${\mathit{w}}_2\in {\mathbb{R}}^{d_1}$ and the elementwise activations are ${\mathit{A}}_1:{\mathbb{R}}^{d_1}\to {\mathcal{A}}_1^{d_1}\subseteq {\mathbb{R}}^{d_1}$ and $A_2:\mathbb{R}\to \mathcal{Y}$. We choose the distribution $Q\left(\theta \right)=:Q_2\left(w_2\right)Q_1\left(W_1\right)$ to factorise over the layers. This is identical to the above and sets ${\mathit{\eta }}_{W_1}\left(\mathit{x}\right)={\mathit{A}}_1\left(W_1\mathit{x}\right)$.

Sampling $\mathit{a}$ is straightforward if sampling $W_1$ is. Further, if the final layer aggregate is differentiable, and so is the hidden layer activation ${\mathit{A}}_1$, we may be able to use the lower-variance pathwise gradient estimator for gradients with respect to $Q_1$. We note that this may be possible even if the activation $A_2$ is not differentiable, as in Section 4, where we extend the pathwise estimator where we could not otherwise use it.

Computationally, we may implement the above by analytically finding the distribution on the “pre-activations” $W_1\mathit{x}$ (trivial for the normal distribution) before sampling this and passing through the activation. With the pathwise estimator this is known as the “local reparameterization trick” [16], which can lead to considerable computational savings on parallel minibatches compared to direct hierarchical sampling, $\mathit{a}=A_1\left(W_1\mathit{x}\right)$ with $W_1\sim Q_1$. We will utilise this in all our reparameterizable dense networks, and a variation on it to save computation when using REINFORCE in Section 4.2 and Section 6.

4. Aggregating Signed-Output Networks

Here we consider a first practical application of the aggregation estimator to stochastic neural networks with a final dense sign-activated layer. We have seen above that this partial aggregation leads to better-behaved training objectives and lower-variance gradient estimates across arbitrary other network structure, It may also allow use of pathwise gradients for the other layers, which would not be possible otherwise due to the non-differentiability of the final layer.

Specifically, these networks take the form of Equation (5) with the final layer activation a sign function and weights drawn from a unit variance normal distribution, $Q^w\left(\mathit{w}\right)=\mathcal{N}(\mathit{\mu },\mathbb{I})$. The aggregate $I\left(\mathit{a}\right)$ is given by Equation (2). Normally-distributed weights are chosen because of the simple analytic forms for the aggregate (reminiscent of the tanh activation occasionally used for neural networks) and KL divergence (effectively an $L^2$ regularisation penalty); we note however that closed forms are available for other commonly-used distributions such as the Laplace.

Using Equations (3) and (6) with independent samples ${\left\{({\mathit{w}}^t,{\theta }^{\neg w,\left(t\right)})\right\}}_{t=1}^T\sim Q$ and ${\mathit{\eta }}^t:={\mathit{\eta }}_{{\theta }^{\neg \mathit{w},\left(t\right)}}\left(\mathit{x}\right)$ leads to the two unbiased estimators for the output (henceforth assuming the technical condition ${\mathbb{P}}_{\mathit{\eta }|\mathit{x}}\{\mathit{\eta }=\mathbf{0}\}=0$ that allows aggregation to be well-defined).

$$\begin{array}{cc}\hfill {\widehat{F}}_Q\left(\mathit{x}\right)& :=\frac{1}{T}\sum _{t=1}^{T}sign({\mathit{w}}^t·{\mathit{\eta }}^t)\hfill \end{array}$$

(7)

$$\begin{array}{cc}\hfill {\widehat{F}}_Q^{\ast }\left(\mathit{x}\right)& :=\frac{1}{T}\sum _{t=1}^{T}erf\left(\frac{\mathit{\mu }·{\mathit{\eta }}^t}{\sqrt{2}\parallel {\mathit{\eta }}^t\parallel }\right).\hfill \end{array}$$

(8)

It follows immediately from Propositions 1 and 2 that the latter and the associated gradient estimators have lower variances than the former or the REINFORCE gradient estimates (which we would otherwise have to use due to the non-differentiability of the final layer).

4.1. Lower Variance Estimates of Aggregated Sign-Output Networks

We clarify the situation with the lower variance estimates further below. In particular, we find that the reduction in variance from using the partial-aggregation estimator is controlled by the norm $\parallel \mathit{\mu }\parallel$, so that for small $\parallel \mathit{\mu }\parallel$ (as could be expected early in training) the difference can be large, while as $\parallel \mathit{\mu }\parallel$ grows, the difference in variance is controlled and we could reasonably revert to the Monte Carlo (or Gibbs) estimator. Note also that as $F_Q\left(\mathit{x}\right)\to \pm 1$ (as expected after training), both variances disappear.

We also show that a stricter condition than Proposition 2 holds on the variances of the aggregated gradients here, and thus that the non-aggregated gradients are noisier in all cases than the aggregate.

Proposition 3.

With the definitions given by Equation (7), for all $\mathit{x}\in {\mathbb{R}}^{d_0}$, $T\in \mathbb{N}$, and Q with normally-distributed final layer,

$$0\le {\mathbb{V}}_Q\left[{\widehat{F}}_Q\left(\mathit{x}\right)\right]-{\mathbb{V}}_Q\left[{\widehat{F}}_Q^{\ast }\left(\mathit{x}\right)\right]\le \frac{1}{T}\left(1-{\left|erf\left(\frac{\parallel \mathit{\mu }\parallel }{\sqrt{2}}\right)\right|}^2\right).$$

Proof. 

The left identity follows directly from Proposition 1. We also have

$$\begin{array}{cc}\hfill {\mathbb{V}}_Q\left[{\widehat{F}}_Q\left(\mathit{x}\right)\right]-{\mathbb{V}}_Q\left[{\widehat{F}}_Q^{☆}\left(\mathit{x}\right)\right]& =\frac{1}{T}{\mathbb{E}}_{\mathit{a}|\mathit{x}}\left[{\mathbb{V}}_w[sign(\mathit{w}·\mathit{a})]\right]\hfill \\ \hfill & =\frac{1}{T}\left(1-{\mathbb{E}}_{\mathit{a}|\mathit{x}}{\left|erf\left(\frac{\mathit{\mu }·\mathit{\eta }}{\sqrt{2}\parallel \mathit{\eta }\parallel }\right)\right|}^2\right)\hfill \\ \hfill & \le \frac{1}{T}\left(1-{\left|erf\left(\frac{\parallel \mathit{\mu }\parallel }{\sqrt{2}}\right)\right|}^2\right).\hfill \end{array}$$

 □

Proposition 4.

Under the conditions of Proposition 3,

$$\mathrm{Cov}\left[{\widehat{G}}^{\ast }\left(\mathit{x}\right)\right]⪯\mathrm{Cov}\left[\widehat{G}\left(\mathit{x}\right)\right]+\frac{1-2/\pi }{T}\mathbb{I}.$$

Thus, for all $u$ with $\parallel u\parallel =1$,

$$\mathbb{V}[{\widehat{G}}^{\ast }\left(\mathit{x}\right)·\mathit{u}]\le \mathbb{V}[\widehat{G}\left(\mathit{x}\right)·\mathit{u}]+\frac{1-2/\pi }{T}.$$

Proof. 

It is straightforward to show that

$$\begin{array}{cc}\hfill \widehat{\mathit{G}}\left(\mathit{x}\right)& :=\frac{1}{T}\sum _{t=1}^{T}sign({\mathit{w}}^t·{\mathit{\eta }}^t)(\mathit{\mu }-{\mathit{w}}^t)\hfill \\ \hfill {\widehat{\mathit{G}}}^{\ast }\left(\mathit{x}\right)& :=\frac{1}{T}\sum _{t=1}^T\frac{{\mathit{\eta }}^t}{\parallel {\mathit{\eta }}^t\parallel }\sqrt{\frac{2}{\pi }}exp\left[-\frac{1}{2}{\left(\frac{\mathit{\mu }·{\mathit{\eta }}^t}{\parallel {\mathit{\eta }}^t\parallel }\right)}^2\right]\hfill \\ \hfill \mathrm{Cov}\left[\widehat{\mathit{G}}\left(\mathit{x}\right)\right]& =\frac{1}{T}\left(\mathbb{I}-\mathit{G}{\mathit{G}}^T\right)\hfill \\ \hfill \mathrm{Cov}\left[{\widehat{\mathit{G}}}^{\ast }\left(\mathit{x}\right)\right]& =\frac{1}{T}\left(\mathbb{E}\left[\frac{\mathit{\eta }{\mathit{\eta }}^T}{{\parallel \mathit{\eta }\parallel }^2}\frac{2}{\pi }{e}^{-{\left(\frac{\mathit{\mu }·\mathit{\eta }}{\parallel \mathit{\eta }\parallel }\right)}^2}\right]-\mathit{G}{\mathit{G}}^T\right)\hfill \end{array}$$

so for $\mathit{u}\ne 0$,

$$\begin{array}{cc}\hfill & T{\mathit{u}}^T\left(\mathrm{Cov}\left[\widehat{\mathit{G}}\left(\mathit{x}\right)\right]-\mathrm{Cov}\left[{\widehat{\mathit{G}}}^{\ast }\left(\mathit{x}\right)\right]\right)\mathit{u}\hfill \\ \hfill & ={\parallel \mathit{u}\parallel }^2-\frac{2}{\pi }\mathbb{E}\left[\frac{{|\mathit{u}·\mathit{\eta }|}^2}{{\parallel \mathit{\eta }\parallel }^2}{e}^{-{\left(\frac{\mathit{\mu }·\mathit{\eta }}{\parallel \mathit{\eta }\parallel }\right)}^2}\right]\ge {\parallel \mathit{u}\parallel }^2\left(1-\frac{2}{\pi }\right)>0.\hfill \end{array}$$

Above we have brought $\mathit{u}$ inside the term with an expectation, which is then bounded using Cauchy–Schwarz on $|\mathit{u}·\mathit{\eta }|/\parallel \mathit{\eta }\parallel \le \parallel u\parallel$, and $e^{-\left|t\right|}\le 1$ for all $t\in \mathbb{R}$. □

4.2. All Sign Activations

Here we examine an important special case previously examined by Letarte et al. [6]: a feed-forward network with all sign activations and normal weights. This takes the form

$$f_{\theta }\left(\mathit{x}\right)=sign({\mathit{w}}_L·sign(W_{L-1}\cdots sign\left(W_1\mathit{x}\right)\cdots ))$$

with $\theta :=vec({\mathit{w}}_L,\cdots ,W_1)$ and $W_l:={[{\mathit{w}}_{l,1}\cdots {\mathit{w}}_{l,d_l}]}^T$; $l\in \{1,...,L\}$ are the layer indices. We choose unit-variance normal distributions on the weights, which factorise into $Q_l\left(W_l\right)={\prod }_{i=1}^{d_l}{q}_{l,i}\left({\mathit{w}}_{l,i}\right)$ with $q_{l,i}=\mathcal{N}({\mathit{\mu }}_{l,i},{\mathbb{I}}_{d_{l-1}})$. In the notation of Section 3, ${\mathit{\eta }}_{{\theta }^{\neg w}}\left(\mathit{x}\right)=sign(W_{L-1}\cdots sign\left(W_1\mathit{x}\right)\cdots )$ is the final layer activation, which could easily be sampled by mapping $\mathit{x}$ through the first $L-1$ layers with draws from the weight distribution.

Instead, we go on to make an iterative replacement of the weight distributions on each layer by conditionals on the layer activations to obtain the summation

$$\begin{array}{cc}\hfill & F_Q\left(\mathit{x}\right)=\sum _{{\mathit{a}}_1\in {\{+1,-1\}}^{d_1}}{\tilde{Q}}_1\left({\mathit{a}}_1\right|\mathit{x})\times \cdots \hfill \\ \hfill \cdots & \times \sum _{{\mathit{a}}_{L-1}\in {\{+1,-1\}}^{d_{L-1}}}{\tilde{Q}}_{L-1}\left({\mathit{a}}_{L-1}\right|a_{L-2})erf\left(\frac{{\mathit{\mu }}_L·{\mathit{a}}_{L-1}}{\sqrt{2}\parallel {\mathit{a}}_{L-1}\parallel }\right).\hfill \end{array}$$

(9)

The number of terms is exponential in the depth so we instead hierarchically sample the ${\mathit{a}}_l$. Like local reparameterisation, this leads to a considerable computational saving over sampling a separate weight matrix for every input. The conditionals can be found in closed form: we can factorise individual hidden units ${\tilde{Q}}_l\left({\mathit{a}}_l\right|{\mathit{a}}_{l-1}):={\prod }_{i=1}^{d_l}{\tilde{q}}_{l,i}\left(a_{l,i}\right|{\mathit{a}}_{l-1})$, and find their activation distributions (with ${\mathit{a}}_0:=\mathit{x}$ and z a dummy variable):

$$\begin{array}{cc}\hfill {\tilde{q}}_{l,i}(a_{l,i}=\pm 1|{\mathit{a}}_{l-1})& ={\int }_0^{\infty }\mathcal{N}(z;\pm {\mathit{\mu }}_{l,i}·{\mathit{a}}_{l-1},\parallel {\mathit{a}}_{l-1}{\parallel }^2)\mathrm{d}z\hfill \\ \hfill & =\frac{1}{2}\left[1\pm erf\left(\frac{{\mathit{\mu }}_{l,i}·{\mathit{a}}_{l-1}}{\sqrt{2}\parallel {\mathit{a}}_{l-1}\parallel }\right)\right].\hfill \end{array}$$

A marginalised REINFORCE-style gradient estimator for conditional distributions can then be used; this does not necessarily have better statistical properties but in combination with the above is much more computationally efficient. This idea of “conditional sampling” is inspired by the local reinforce trick. Using samples ${\left\{({\mathit{a}}_1^t\cdots {\mathit{a}}_{L-1}^t)\right\}}_{t=1}^T\sim \tilde{Q}$,

$$\frac{\partial F_Q\left(\mathit{x}\right)}{\partial {\mathit{\mu }}_{l,i}}\approx \frac{1}{T}\sum _{t=1}^{T}erf\left(\frac{{\mathit{\mu }}_L·{\mathit{a}}_{L-1}^t}{\sqrt{2}\parallel {\mathit{a}}_{L-1}^t\parallel }\right)\frac{\partial }{\partial {\mathit{\mu }}_{l,i}}log{\tilde{q}}_{l,i}\left(a_{l,i}^t\right|{\mathit{a}}_{l-1}^t).$$

(10)

This formulation along with Equation (9) resembles the PBGNet model of [6], but derived in a very different way. Indeed both are equivalent in the single-hidden-layer case, but with more layers PBGNet uses an unusual tree-structured network to make the individual activations independent and avoid an exponential computational dependency on the depth in Equation (9). This makes the above summation exactly calculable but is also still not efficient enough in practice, so they resort further to a Monte Carlo approximation: informally, this draws new samples for every layer l based on an average of those from the previous layer, ${\mathit{a}}_l|{\left\{{\mathit{a}}_{l-1}^{\left(t\right)}\right\}}_{t=1}^T\sim \frac{1}{T}{\sum }_{t=1}^T\tilde{Q}\left({\mathit{a}}_l\right|{\mathit{a}}_{l-1}^{\left(t\right)})$.

This is all justified within the tree-structured framework but leads to an exponential KL penalty which—as hinted by Letarte et al. [6] and shown empirically in Section 6—makes PAC–Bayes bound optimisation strongly favour shallower such networks. Our formulation avoids this, is more general—applying to alternative network structures—and we believe it is significantly easier to understand and use in practice.

5. PAC–Bayesian Objectives with Signed-Outputs

We now move to obtain binary classifiers with guarantees for the expected misclassification error, $R^{0-1}$, which we do by optimizing PAC–Bayesian bounds. Such bounds (as in Theorem 1) will usually involve the non-differentiable and non-convex misclassification loss ${\ell }_{0-1}$. However, to train a neural network we need to replace this by a differentiable surrogate, as discussed in the introduction.

Here we adopt a different approach by using our signed-output networks, where since $f\left(\mathit{x}\right)\in \{+1,-1\}$, there is an exact equivalence between the linear and misclassification losses, ${\ell }_{0-1}(f\left(\mathit{x}\right),y)={\ell }_{\mathrm{lin}}(f\left(\mathit{x}\right),y)$, avoiding an extra factor of two from the inequality ${\ell }_{0-1}\le 2{\ell }_{\mathrm{lin}}$.

Although we have only moved the non-differentiability into f, the form of a PAC–Bayesian bound and the linearity of the loss and expectation allow us to go further and aggregate,

$${\mathbb{E}}_{f\sim Q}{\ell }_{0-1}(f\left(\mathit{x}\right),y)={\mathbb{E}}_{f\sim Q}{\ell }_{\mathrm{lin}}(f\left(\mathit{x}\right),y)={\ell }_{\mathrm{lin}}(F_Q\left(\mathit{x}\right),y)$$

(11)

which allows us to use the tools discussed in Section 4 to obtain lower-variance estimates and gradients. Below we prove a small result to show the utility of this:

Proposition 5.

Under the conditions of Proposition 3 and $y\in \{+1,-1\}$,

$$\begin{array}{cc}\hfill {\mathbb{V}}_Q[& {\ell }_{\mathrm{lin}}({\widehat{F}}_Q^{\ast }\left(\mathit{x}\right),y)]\le {\mathbb{V}}_Q\left[{\ell }_{\mathrm{lin}}({\widehat{F}}_Q\left(\mathit{x}\right),y)\right]\hfill \\ \hfill & \le {\mathbb{V}}_{f\sim Q}\left[{\ell }_{0-1}(f\left(\mathit{x}\right),y)\right]=\frac{1}{4}(1-|F_Q\left(\mathit{x}\right){|}^2).\hfill \end{array}$$

Proof. 

$${\mathbb{V}}_Q\left[{\ell }_{\mathrm{lin}}({\widehat{F}}_Q\left(\mathit{x}\right),y)\right]={\mathbb{E}}_Q{\left|\frac{1}{2}(y{F}_Q\left(\mathit{x}\right)-y{\widehat{F}}_Q\left(\mathit{x}\right))\right|}^2=\frac{1}{4}{\mathbb{V}}_Q\left[{\widehat{F}}_Q\left(\mathit{x}\right)\right]$$

and a similar result for ${\widehat{F}}_Q^{\ast }$. $f={\widehat{F}}_Q$ if $T=1$ and ${\ell }_{\mathrm{lin}}(f\left(\mathit{x}\right),y)={\ell }_{0-1}(f\left(\mathit{x}\right),y)$. The result then follows from this and Proposition 3. □

Combining (11) with Theorem 1, we obtain a directly optimizable, differentiable bound on the misclassification loss without introducing the above-mentioned factor of 2.

Theorem 2.

Given P on θ and $\alpha >1$, for all Q on θ and $\lambda >1$ simultaneously with probability at least $1-\delta$ over $S\sim {\mathcal{D}}^m$,

$${\mathbb{E}}_{\theta \sim Q}{R}^{0-1}\left(f_{\theta }\right)\le {\Phi }_{\lambda /m}^{-1}\left[R_S^{\mathrm{lin}}\left(F_Q\right)+\frac{\alpha }{\lambda }\Delta \right]$$

with ${\Phi }_{\gamma }^{-1}\left(t\right)=\frac{1-exp(-\gamma t)}{1-exp(-\gamma )}$, $f_{\theta }:{\mathbb{R}}^d\to \{+1,-1\},\theta \in \Theta$, and $\Delta =KL\left(Q|P\right)-log\delta +2log\left(\frac{log{\alpha }^2\lambda }{log\alpha }\right)$.

Thus, for each $\lambda$, which can be held fixed (“fix-$\mathit{\lambda }$”) or simultaneously optimized throughout training for automatic regularisation tuning (“optim-$\mathit{\lambda }$”), we obtain a gradient descent objective:

$$R_S^{\mathrm{lin}}\left({\widehat{F}}_Q^{\ast }\right)+\frac{\mathrm{KL}\left(Q\right|P)}{\lambda }.$$

(12)

6. Experiments

All experiments (

Table 1. Average (from ten runs) binary-MNIST losses and bounds ($\delta =0.05$) for the best epoch and optimal hyperparameter settings of various algorithms. Hyperparameters and epochs were chosen by bound if available and non-vacuous, otherwise by training linear loss. Bold numbers indicate the best values and standard deviation is reported in italics.

	mlp	pbg	Reinforce		Fix-$\mathit{\lambda }$			Optim-$\mathit{\lambda }$
			sign	relu	sign	sgmd	relu	sign	sgmd	relu
Train Linear	0.78	8.72	26.0	18.6	8.77	7.60	6.35	6.71	6.47	5.41
error, $1\sigma$	0.08	0.08	0.8	1.4	0.04	0.19	0.10	0.11	0.18	0.16
Test 0–1	1.82	5.26	25.4	17.9	8.73	7.88	6.51	6.85	6.84	5.61
error, $1\sigma$	0.16	0.18	1.0	1.5	0.23	0.30	0.19	0.27	0.21	0.20
Bound 0–1	-	40.8	100	100	21.7	18.8	15.5	22.6	19.3	16.0
error, $1\sigma$	-	0.2	0.0	0.0	0.04	0.17	0.04	0.03	0.31	0.05

) run on “binary”-MNIST, dividing MNIST into two classes, of digits 0–4 and 5–9. Neural networks had three hidden layers with 100 units per layer and sign, sigmoid (sgmd) or relu activations, before a single-unit final layer with sign activation. Q was chosen as an isotropic, unit-variance normal distribution with initial means drawn from a truncated normal distribution of variance $0.05$. The data-free prior P was fixed equal to the initial Q, as motivated by Dziugaite and Roy [4] (Section 5 and Appendix B).

The objectives fix-$\mathit{\lambda }$ and optim-$\mathit{\lambda }$ from Section 5 were used for batch-size 256 gradient descent with Adam [17] for 200 epochs. Every five epochs, the bound (for a minimising $\lambda$) was evaluated using the entire training set; the learning rate was then halved if the bound was unimproved from the previous two evaluations. The best hyperparameters were selected using the best bound achieved in these evaluations through a grid search of initial learning rates $\in \{0.1,0.01,0.001\}$, sample sizes $T\in \{1,10,50,100\}$. Once these were selected training was repeated 10 times to obtain the values in Table 1.

$\lambda$ in optim-$\lambda$ was optimised through Theorem 2 on alternate mini-batches with SGD and a fixed learning rate of ${10}^{-4}$ (whilst still using the objective (12) to avoid effectively scaling the learning rate with respect to empirical loss by the varying $\lambda$). After preliminary experiments in fix-$\lambda$, we set $\lambda =m=60,000$, the training set size, as is common in Bayesian deep learning.

We also report the values of three baselines: reinforce, which uses the fix-$\lambda$ objective without partial-aggregation, forcing the use of REINFORCE gradients everywhere; mlp, an unregularised non-stochastic relu neural network with tanh output activation; and the PBGNet model (pbg) from Letarte et al. [6]. For the latter, a misclassification error bound obtained through ${\ell }_{0-1}\le 2{\ell }_{\mathrm{lin}}$ must be used as their test predictions were made through the sign of a prediction function $\in [-1,+1]$, not $\in \{+1,-1\}$. Further, despite significant additional hyperparameter exploration, we were unable to train a three layer network through the PBGNet algorithm directly comparable to our method, likely because of the exponential KL penalty (in their Equation 17) within that framework; to enable comparison, we therefore allowed the number of hidden layers in this scenario to vary $\in \{1,2,3\}$. Other baseline tuning and setup was similar to the above, see the Appendix A for more details.

During evaluation reinforce draws a new set of weights for every test example, equivalent to the evaluation of the other models; but doing so during training, with multiple parallel samples, is prohibitively expensive. Two different approaches to straightforward, not partially-aggregated, gradient estimation for this case suggest themselves, arising from different approximations to the Q-expected loss of the minibatch, $B\subseteq S$ (with data indices $\mathcal{B}$). From the identities

$$\begin{array}{cc}\hfill {\nabla }_{\varphi }{\mathbb{E}}_{\theta \sim q_{\varphi }}{R}_B\left(f_{\theta }\right)& ={\mathbb{E}}_{\theta \sim q_{\varphi }}\frac{1}{\left|B\right|}\sum _{i\in \mathcal{B}}\ell (f_{\theta }\left({\mathit{x}}_i\right),y_i){\nabla }_{\varphi }log{q}_{\varphi }\left(\theta \right)\hfill \\ \hfill & =\frac{1}{\left|B\right|}\sum _{i\in \mathcal{B}}{\mathbb{E}}_{\theta \sim q_{\varphi }}\ell (f_{\theta }\left({\mathit{x}}_i\right),y_i){\nabla }_{\varphi }log{q}_{\varphi }\left(\theta \right)\hfill \end{array}$$

we obtain two slightly different estimators for ${\nabla }_{\varphi }{\mathbb{E}}_{\theta \sim q_{\varphi }}{R}_B\left(f_{\theta }\right)$:

$$\begin{array}{cc}\hfill & \frac{1}{T\left|B\right|}\sum _{t=1}^T\sum _{i\in \mathcal{B}}\ell (f_{{\theta }^{(t,i)}}\left({\mathit{x}}_i\right),y_i){\nabla }_{\varphi }log{q}_{\varphi }\left({\theta }^{(t,i)}\right)\hfill \\ \hfill & \frac{1}{T\left|B\right|}\sum _{i\in \mathcal{B}}\sum _{t=1}^T\ell (f_{{\theta }^t}\left({\mathit{x}}_i\right),y_i){\nabla }_{\varphi }log{q}_{\varphi }\left({\theta }^t\right).\hfill \end{array}$$

The first draws many more samples and has lower variance but is much slower computationally; even aside from the $O\left(\right|B\left|\right)$ increase in computation, there is a slowdown as the optimised BLAS matrix routines cannot be used, and the very large matrices involved may not fit in memory (see for more information [16]).

Therefore, as is standard in the Bayesian Neural Network literature with the pathwise estimator, we use the latter formulation, which has a similar computational complexity to local-reparameterisation and our marginalised REINFORCE estimator (10). We should note though that in preliminary experiments, the alternate estimator did not appear to lead to improved results. This clarifies the advantages of marginalised sampling, which can lead to lower variance with a similar computational cost.

7. Discussion

The experiments demonstrate that partial-aggregation enables training of multi-layer non-differentiable neural networks in a PAC–Bayesian context, which is not possible with REINFORCE gradients and a multiple-hidden-layer PBGNet [6]. These obtained only vacuous bounds, and our misclassification bounds also improve those of a single-hidden-layer PBGNet.

Our experiments raise a couple of questions: firstly, why is it that lower variance estimates empirically lead to tighter bounds? We speculate that the faster convergence of SGD in this case takes us to a more “local” minimum of the objective, closer to our initialisation. Since most existing PAC–Bayes bounds for neural networks have a very strong dependence on this distance from initialisation through the KL term, this leads to tighter bounds. This distance could also be reduced through other methods we consider out-of-scope, such as the data-dependent bounds employed by Dziugaite and Roy [18] and Letarte et al. [6].

A second and harder question is asking why the non-stochastic mlp model obtains a lower overall error. The bound optimisation is empirically quite conservative, but does not necessarily lead to better generalisation; understanding this gap is a key question in the theory of deep learning.

In future work we will develop significant new tools to extend partial-aggregation to multi-class classification, and to improve test prediction bounds for $sign\left({\widehat{F}}_Q\left(\mathit{x}\right)\right)$ with $T>1$, as in PBGNet, which gave slightly improved predictive performance despite the inferior theoretical guarantees.