Analytically Tractable Inference in Deep Neural Networks
Abstract
Since its inception, deep learning has been overwhelmingly reliant on backpropagation and gradient-based optimization algorithms in order to learn weight and bias parameter values.
The
Tractable Approximate Gaussian Inference (TAGI) algorithm was shown to be a viable and scalable alternative to backpropagation for shallow fully-connected neural networks.
In this paper, we are demonstrating how TAGI matches or exceeds the performance of backpropagation, for training classic deep neural network architectures.
Although TAGI’s computational efﬁciency is still below that of deterministic approaches relying on backpropagation, it outperforms them on classiﬁcation tasks and matches their performance for information maximizing generative adversarial networks while using smaller architectures trained with fewer epochs.
INTRODUCTION
Since its inception, deep learning has been overwhelmingly reliant on backpropagation and gradient-based optimization algorithms in order to learn weight and bias parameter values.
The Tractable Approximate Gaussian Inference (TAGI) algorithm
[Goulet et al., 2020] was shown to be a viable and scalable alternative to backpropagation.
Nevertheless, the method’s performance was so far only demonstrated for shallow fully-connected neural networks, because its formulation is not compatible with existing libraries, e.g.,
TensorFlow or Pytorch.
In this paper, we demonstrate how
TAGI matches or exceeds the performance of backpropagation, for classic deep neural network architectures.
Sectionﬁrst presents the fundamentals of TAGI on which sectionbuilds in order to adapt the approach to convolutional neural networks (CNN) [LeCun and Bengio, 1995] and generative adversarial networks (GAN)
[Goodfellow et al., 2014].
In parallel, this same section presents how common tools such as hidden state normalization [Goodfellow et al., 2016], i.e. batch normalization [Ioffe and Szegedy, 2015] and layer normalization [Ba et al., 2016], and observation noise decay [Neelakantan et al., 2015] can be integrated into the TAGI’s framework.
Finally, section 4 compares the performance of TAGI with the state-of-the-art (SOTA) performance for image classiﬁcation and image generation benchmarks.
TAGI
TAGI relies upon two key concepts [Goulet et al., 2020].
First, the joint distribution between the observations and a neural network’s parameters is approximated by a multivariate Gaussian distribution.
Learning the parameters, i.e., the mean vector and covariance matrix deﬁning the joint probability density function (PDF) requires two simplifying assumptions; (1) the product of Gaussian random variables is also Gaussian with its moments, expected values, covariance, and cross covariances calculated analytically using moment generating functions, and (2) the non-linear activations such as rectiﬁed linear unit (ReLU), Tanh, and logistic sigmoid are locally linearized at the expected value of the input hidden unit, thus the output expected value, covariance, and cross-covariance for the activation unit can be computed analytically.
The second key concept employed by TAGI has two folds; (1) it requires simplifying the covariance matrices for hidden layers and parameters to a diagonal structure, and (2) it performs the Gaussian conditional inference using a recursive layer-wise procedure which allows for an analytically tractable inference that scales linearly with respect to the number of weight parameters.
Despite the simplifying assumptions mentioned above, TAGI was shown to match the performance of feedforward neural networks (FNN) trained with backpropagation for regression and classiﬁcation benchmarks.
In the following section, we will show how to apply TAGI to deep architecture formulations in order to demonstrate its superiority on benchmark problems.
FORMULATIONS
The fundamental mathematical operations, i.e., additions, multiplications and non-linear activations, in deep neural networks are no different from those in FNN.
Therefore, the mathematical formulation developed for deﬁning the joint
PDF for observations and parameters in FNN can be readily employed in deeper architectures.
This section covers the speciﬁcities related to CNNs and GANs.
3.1 CNN
The main building blocks of CNNs are the convolution and pooling layers whereas the convolution operation is performed using ﬁxed-size ﬁlters whose weights are learned from data.
As the convolution operation
[He and Sun, 2015] involves the same operations as a FNN, i.e., multiplication, addition and non-linear transformation, TAGI’s mathematical developments can be directly applied to perform the operations in the convolution layer.
Small changes occur in the pooling operations which reduce the output size following the convolutional layers.
Additional changes also take place in two key tools that are employed with CNNs; hidden states’ normalization, and observation noise decay.
3.1.1 Pooling layer
Existing methods such as the stochastic max pooling introduced by Peters and Welling [2018] can be directly be applied with TAGI.
In addition to this method, we adapt an average pooling approach to the context of TAGI, where, as its name indicates, the output is the average of the activation units in the pooling kernel.
For a K -elements pooling kernel, the output Gaussian random variable is deﬁned as
[Equation]
3.1.2 Layer and batch normalization
In the context of TAGI, layer normalization [Ba et al., 2016] can be done using the following principles; For a layer of
Gaussian activation units A =
[ A 1 A 2 · · · A A ] (cid:124) , the normalization is deﬁned as
[Equation]
where µ A and σ A are the mean and standard deviation of a
Gaussian mixture reduction [Runnalls, 2007] made from all activation units, each having a probability 1 / A
so that
[Equation]
The resulting expected value, covariance, and crosscovariance for the normalized vector of activation units ˜
A are given by
[Equation]
By using the short-hand notation deﬁned by Goulet et al.
[2020] for hidden units at a layer j + 1 , i.e. Z + ≡ Z ( j +1) , and for the weights, bias, an hidden units at a layer j , i.e. [Equation]
[Equation] Z [Equation] j ) } , we can use the terms from equation 4 in the TAGI’s formulation to update the covariance between hidden states at successive layers
[Equation]
where J is the diagonal Jacobian matrix of the activation function evaluated at µ Z , i.e., J = diag (cid:0) ∇ z σ ( µ Z ) (cid:1) .
The covariance between the weight and bias parameters at a layer j and the hidden states at a layer j + 1 are
[Equation]
The developments of Equations 5 and 6 are provided in
Appendix A. These updated expected value, covariance, and cross-covariance terms can be directly employed in the
TAGI’s layer-wise recursive inference procedure.
The procedure for batch normalization [Ioffe and Szegedy, 2015] is analogous to the layer normalization except that the normalization is made for each hidden units, over a batch of
B observations.
The Gaussian activation units contributing to the normalization are A j =
[ A 1 j A 2 j · · · A B j ] (cid:124) , so that the number of activation units A in Equations 3 and 4 is replaced by the number of observations per batch B .
3.1.3 Observation noise decay
In the original TAGI formulation for FNN, the observation errors’ standard deviation parameters σ V were considered as constants during the training.
Throughout empirical experimentation with CNNs, we noticed that we could improve the performance when using a decay equation
[Equation]
where e is the current epoch and η ∈ (0 , 1] is the decay factor hyperparameter that needs to be learned outside of the TAGI analytic inference procedure.
This approach is similar to what is done in standard deep neural networks trained with backpropagation where noise is added to the gradient [Neelakantan et al., 2015] with a decay schedule.
In the case of gradient-based learning, this noise consists in discrete samples added to the gradient itself whereas for
TAGI, it consists in additional variance on the output layer so that the update during the inference step will put more weight on the prior rather than on the likelihood, with this effect diminishing with time.
3.2 GENERATIVE ADVERSARIAL NETWORKS
GANs build on a generative approach to create realizations from the input covariates distribution f ( x ) .
In practice, the input x typically takes the form of images.
The GAN architecture consists in a discriminator (Dnet) and a generator (Gnet) network as depicted in Figure 1.
The generator network is used to create synthetic data and the discriminator is a binary classiﬁcation network that allows distinguishing synthetic (i.e. fake) and real data.
The generator takes random noise as input.
Figure 1 presents the schematic representation of a GAN where the magenta arrows depict the dependencies in the network and correspond to the ﬂow of information during the forward propagation of uncertainty through the network.
The blue arrows running backward represent the ﬂow of information during the inference procedure.
In Figure 1a, the Dnet takes as input either a fake image generated from an input noise sample that is passed to the Gnet, or a real image; once passed through the Dnet, that image is either classiﬁed as fake or real in the output layer z
( O ) .
The parameters of the Dnet are then updated in the inference step by relying on the true label for that image.
During this inference step, the parameters of the Gnet are
[Figure]
a GAN.
The magenta arrows correspond to the ﬂow of information during the forward propagation of uncertainty through the network and the blue arrows represent the ﬂow of information during the inference procedure.
noise x c x d
Gnet
Fake images
Dnet
Pnet z ( O ) P fake
Qnet z ( O ) Q x d , x c Real images real (a) Discriminator noise x c x d
Gnet
Fake images
Dnet
Pnet z ( O ) P real
Qnet z ( O ) Q x d , x c (b) Generator not updated.
For the generator presented in Figure 1b, an input noise sample is passed to the Gnet in order to return a fake image that is then passed through the Dnet.
We then update the Gnet’s parameters by forcing the observation for the output layer of the Dnet, z
( O ) , to be classiﬁed as a real image.
Note that although the inference procedure passes through the Dnet, only the parameters from the Gnet are updated for the generator.
(InfoGAN)
[Chen et al., 2016] allow controlling the latent space by disentangling interpretable data representations such as classes, rotation, and width.
As a GAN, an infoGAN consists in two networks: discriminator and generator.
With infoGAN, the discriminator network has two purposes: (1) distinguishing the fake and real data, and (2) learn the interpretable representations.
Figure 2a shows the architecture of the discriminator where the Pnet is employed to classify the real and fake data just like in normal GANs, while the
Qnet allows predicting the latent variables associated with an input image.
Both the Pnet and Qnet are outputs from the shared Dnet taking as input either a real or fake image that has been created by the Gnet.
During inference, the parameters of the Dnet, Pnet and Qnet are updated by relying on the true label for that image as well as the latent variables in the case of fake images.
The inputs of the generator networks consist in an input noise sample, a categorical variable vector ( x d ), and continuous variable vector ( x c ).
For example, in the case of the MNIST dataset [LeCun et al., 1998], the discrete variables represent the digits from 0 to 9 and the continuous variable can be the width and rotation of each digit.
In practice, x d and x c are randomly generated during training and can be speciﬁcally queried once the training is completed.
Note that again only the Gnet’s parameters are updated in Figure 2b, even if during inference, the information has to ﬂow backward through the Pnet, Qnet and
Dnet.
For image datasets, GANs architectures commonly employ both CNNs and FNNs which involve the operations such as multiplication, addition and non-linear transformation.
Therefore, TAGI’s mathematical developments can be directly applied to perform the operations in GAN and infoGAN architectures without having to rely on any gradient whatsoever, or to develop any new equations other than those presented in Section 3.1 for CNNs.
BENCHMARKS
In this section, we compare the performance of TAGI for deep neural networks with existing deterministic or
Bayesian neural networks trained with backpropagation.
In the case of TAGI, the hyperparameters are kept the same for all benchmarks with a batch size B = 16 , a maximal number of epochs E = 50 , a noise discount factor η = 0 . 975 , and all weights are initialized using the He’s approach [He et al., 2015].
These hyperparameters are kept constant across all benchmarks to demonstrate the robustness of the approach towards them.
The only hyperparameter speciﬁc to each experiment is the initial observation noise standard deviation σ 0 V .
Additionally, for the MNIST [LeCun et al., 1998] and
CIFAR10 [Krizhevsky et al., 2009] datasets, the input data was ﬁrst normalized between 0 and 1 and then the mean of the normalized data was subtracted.
For the celebA dataset [Liu et al., 2015], we preprocess the input data by subtracting the mean and dividing by the standard deviation of the input data.
4.1 MNIST
The ﬁrst benchmark is the MNIST digit classiﬁcation task.
Table 1 compares the classiﬁcation accuracy with the results from Wan et al.
[2013] where both approaches use the same
CNN architecture with 2 convolutional layers (32-64) and a fully connected layer with 150 hidden units.
For this experiment, the initial observation noise standard deviation σ 0 V = 1 .
Results show that TAGI can already reach a classiﬁcation error < 2% after the ﬁrst epoch and after E =epochs, the performance is slightly better than the results obtained with backpropagation trained with over 20 times as many epochs.
In a second experiment on the same dataset, we employ an infoGAN (see §3.2) in order to test the performance of TAGI on unsupervised learning tasks.
The hyperparameters speciﬁc to this experiment are the two observation noise initial standard deviations for the output of the Pnet and Qnet that [Table] BP [Wan et al
[Table] [Table] are taken as σ [Table]
= σ 0 ,QV = 3 .
As a reference, we employ the input latent space and infoGAN architecture described in
Chen et al.
The input latent space includes 62 Gaussian noise variables, one categorical variable x d containing ten classes, and two continuous variables for representing the width and rotation of the digit.
In the case of TAGI, we employ the same input latent space with a modiﬁcation to the architecture.
Speciﬁcally, for the Gnet, we use half as many convolutional ﬁlters and half as many hidden units in the output’s fully connected layer.
For the Qnet, we use a fully connected layer containing 300 hidden units instead of 128.
Overall, our architecture contains ≈ 2.1M parameters versus 13.3M for the network used as reference [Chen et al., 2016], and we train over 50 epochs whereas they train with 100 epochs.
The complete architecture’s details are provided in Appendix B.2.1.
The results obtained using TAGI are presented in Figures 3a–c while the results obtained with backpropagation are shown in Figures 3d–f. Note that the results for the MNIST trained with backpropagation were obtained using the code from Mohit [2019].
Overall, we denote an advantage for the network trained with backpropagation as TAGI can only disentangle 9 out of 10 digits.
Moreover, for TAGI, the two continuous variables forming the latent space do not have an effect as clear as in the case of backpropagation.
The reason for this sub-par performance is that for TAGI, the output layer uncertainty is homoscedastic whereas, in the case of Chen et al.
[2016], it is heteroscesastic so that they optimize both the expected values and variances for the two continuous variables on the output layer.
With its current formulation TAGI can only handle homoscedastic output uncertainty which jeopardizes its capacity at disentangling the latent space with the unsupervised setup where we learn both the categorical and continuous variables at once.
4.2 CIFAR-
The second benchmark is the CIFAR-10 image classiﬁcation task.
Experiments are conducted for a simple network with three convolutional layers, as well as for the Resnet architecture [He et al., 2016].
For both cases, we also experimented with the usage of layer and batch normalization.
[Figure]
The details regarding the 3-layers convolutional network are presented in Appendix B.1.2.
Note that we only apply the data augmentation techniques including random cropping and horizontal ﬂipping [Osawa et al., 2019] to the Resnet architecture for the comparison purpose with the neural networks trained with backpropagation.
In the case of TAGI, the initial observation noise standard deviation was set to the same value as for the MNIST experiment, i.e., σ 0 V = 1 .
Table 2 presents the classiﬁcation error rate for TAGI-trained networks as well as the results from several papers using deterministic as well as Bayesian methods for the same network architectures.
Note that in the case of Resnet network, the results presented are the average from three runs.
First, for the 3-layer convolutional architecture, the results show that TAGI matches the SOTA performance while using a lower number of epochs for the experiments with the batch normalization and without normalization.
TAGI outperforms the SOTA performance for the experiment with the layer normalization.
Second, for the Resnet18 architecture, TAGI outperforms deterministic and Bayesian approach with a third of the epochs for training.
Note that the performance of TAGI still matches the performance of MCdropout and VOGN even if we employ half as many convolutional ﬁlters in the Resnet18 architecture, i.e., Resnet18half.
Table 3 presents the predictive uncertainty calibration metrics [Guo et al., 2017] for the Resnet18 and Resnet18half architecture.
The results show that for the Resnet18, TAGI performs poorly on the NLL and ECE criteria, yet, for the same criteria, it performs on par with other methods for the
Resnet18half.
One of reason behind this limitation is that
TAGI is currently limited to using homoscedastic observation noise parameters σ V , which is not the case for existing methods relying on backpropagation.
Despite this limitation, none of the backpropagation-based methods match the
AUROC criterion obtained with TAGI.
4.3 CELEBA
In the last experiment, we applied infoGAN to the CelebA dataset for which we downsampled the original images from 3 × 224 × 224 to 3 × 32 × 32 .
The hyperparameters speciﬁc to this experiment are the two observation noise initial standard deviations for the output of the Pnet and Qnet that are taken as σ 0 ,PV = 3 and σ 0 ,QV = 8 .
As a reference, we use the input latent space and the infoGAN architecture described in
Chen et al.
The input latent space contains a Gaussian noise vector including 128 variables along with ten categorical variables x d each having ten classes.
The speciﬁcations of the architecture used with TAGI is presented in
Appendix B.2.2.
Overall the TAGI-trained architecture contains approximately 0.7M parameters in comparison with 4.1M for the architecture which serves here as reference.
Results for the infoGan trained with backpropagation were obtained using the code from Mohit [2019].
Figure 4 compares a sample of the training images with those generated with both the architecture from Chen et al.
[2016] trained with backpropagation over 100 epochs, and with an architecture containing six times fewer parameters trained with TAGI over 50 epochs.
These results show that despite using a smaller network, TAGI performs at par with the larger network trained with backpropagation.
Figure 5 presents the variations in images for the latent variable controlling the hair colour.
Again, TAGI performs at par with the larger network trained with backpropagation.
Other latent variables such as the gender, azimuth, hairstyle, and contrasts are presented in Appendix D [Table] [Table]
DISCUSSION
For the infoGAN experiments, the results have shown that
TAGI can reach a performance at par with a network trained with backpropagation while using signiﬁcantly smaller architectures and a lower number of epochs.
We opted for this approach because TAGI’s implementation is not yet optimized for computational efﬁciency so the current computing time is approximately one order of magnitude higher than when training the same deterministic neural network using backpropagation.
Throughout the experiments, TAGI employs far fewer hyperparameters than traditional neural networks trained with backpropagation.
Moreover, TAGI’s robustness toward these hyperparameter values allows keeping most of them constant across various tasks.
The main hyperparameters that still need to be adapted with the size of the network employed are the initial output noise standard deviations, σ 0 V .
In the future, these hyperparameters could be passed as hidden states that are learned online like other weights and bias, while simultaneously adding the capacity to model heteroscedastic output noise variance.
In addition, we can improve the classiﬁcation performance with the same architecture by using smaller batch sizes, which, however, decreases the potential for parallelization.
The same applies, for larger batch sizes where the losses in classiﬁcation performance is compensated by an increase in computation parallelization.
A systematic study of the [Figure] [Figure] effect of batch sizes remains to be explored.
CONCLUSION
This paper demonstrates the applicability of TAGI to deep neural networks.
The results obtained from different tasks conﬁrm that TAGI is able to perform analytical inference for the parameters of deep neural network.
Although the computational efﬁciency is still below that of standard approaches relying on backpropagation, TAGI outperforms these approaches on classiﬁcation tasks and match the performance for infoGANs while using smaller architectures trained with fewer epochs.
Future code developments oriented toward efﬁciency should reduce that gap.
Author Contributions
Both authors contributed equally.
Acknowledgements
The ﬁrst author was ﬁnancially supported by research grants from Hydro-Quebec, and the Natural Sciences and Engineering Research Council of Canada (NSERC).
APPENDIX A CROSS-COVARIANCES FOR LAYER NORMALIZATION
Assuming that Z , W , and B are independent on each other, [Figure]
[Equation]
[Figure] [Table] [Figure] [Table]
[Equation]
[Figure]
[Table]
[Equation]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Equation]
=
[Equation]
[Equation] 8) where A = J ( Z − µ Z ) + σ ( Z ) with σ ( . ) being the activation function.
[Equation]
B MODEL ARCHITECTURES
This appendix contains the speciﬁcations for each model architecture in the experiment section.
D refers to a layer depth; W refers to a layer width; H refers to the layer height in case of convolutional or pooling layers; K refers to the kernel size; P refers to the convolutional kernel padding; S refers to the convolution stride; σ refers to the activation function type; ReLU refers to rectiﬁed linear unit; lReLU refers to leaky rectiﬁed linear unit.
B.1 CLASSIFICATION B.1.1 MNIST B.1.2 CIFAR
Layer
D × W × H K × K P S σ
Input 1 ×
×- - - -
Convolutional 32 × 27 × 27 4 ×1 1 ReLU
Pooling 32 × 13 × 13 3 ×0Convolutional 64 × 9 ×5 ×0 1 ReLU
Pooling 64 × 4 ×3 ×0Fully connected 150 × 1 ×- - - ReLU
Output 11 × 1 ×- - - -
Layer
D × W × H K × K P S σ Input 3 × 32 ×- - - -
Convolutional 32 × 32 × 32 5 ×2 1 ReLU
Pooling 32 × 16 × 16 3 ×1Convolutional 32 × 16 × 16 5 ×2 1 ReLU
Average pooling 32 × 8 ×3 ×1Convolutional 64 × 8 ×5 ×2 1 ReLU
Average pooling 64 × 4 ×3 ×1Fully connected 64 × 1 ×- - -
ReLU
Output 11 × 1 ×- - - -
B.2 INFOGAN B.2.1 MNIST
Layer
D × W × H K × K P S σ
{+Input 1 × 
+}
×- - - -
Convolutional 32 × 28 × 28 3 ×1 1 lReLU
Batch normalization 32 × 28 × 28 - - -
Average pooling 32 × 14 × 14 3 ×1Convolutional
[Table]
[Table] [Table]
[Table] [Table] [Table]
[Table] [Table]
C SUPPLEMENTARY RESULTS FOR CLASSIFICATION
[Table] [Table]
[Figure]