LOCAL MINIMA IN TRAINING OF DEEP NETWORKS
There has been a lot of recent interest in trying to characterize the error surface of deep models. This stems from a long standing question. Given that deep networks are highly nonlinear systems optimized by local gradient methods, why do they not seem to be affected by bad local minima? It is widely believed that training of deep models using gradient methods works so well because the error surface either has no local minima, or if they exist they need to be close in value to the global minimum. It is known that such results hold under very strong assumptions which are not satisfied by real models. In this paper we present examples showing that for such theorem to be true additional assumptions on the data, initialization schemes and/or the model classes have to be made. We look at the particular case of finite size datasets. We demonstrate that in this scenario one can construct counter-examples (datasets or initialization schemes) when the network does become susceptible to bad local minima over the weight space.
1 INTRODUCTION
Deep Learning (LeCun et al., 2015; Schmidhuber, 2015) is a fast growing subfield of machine learning, with many impressive results. One particular criticism often brought up against this family of models is the fact that it relies on non-convex functions which are optimized using local gradient descent methods. This means one has no guarantee that the optimization algorithm will converge to a meaningful minimum or even that it will converge at all. However, this theoretical concern seems to have little bearing in practice.
In Dauphin et al. (2013) a conjecture had been put forward for this based on insights from statistical physics which point to the scale of neural networks as a possible answer. The claim is that the error structure of neural networks might follow the same structure as that of random Gaussian fields which have been recently understood and studied in Fyodorov & Williams (2007); Bray & Dean (2007). The critical points of these functions, as the dimensionality of the problem increases, seem to have a particularly friendly behaviour where local minima align nicely close to the global minimum of the function. Choromanska et al. (2015) provides a study of the conjecture by mapping deep neural models onto spin-glass ones for whom the above structure holds. These work has been extended further (see Section 2 for a review of the topic).
We believe many of these results do not trivially extend to the case of finite size datasets/finite size models. The learning dynamics of the neural network in this particular case can be arbitrarily bad. Our assertions are based on constructions of counterexamples that exploit particular architectures, the full domain of the parameters and particular datasets.
2 LITERATURE REVIEW
One view, that can be dated back to Baldi & Hornik (1989), about why the error surface of neural networks seems well behaved is the one stated in Dauphin et al. (2013). We would refer to this hypothesis as the “no bad local minima” hypothesis. In Baldi & Hornik (1989) it is shown that an MLP with a single linear intermediate layer has no local minima, only saddle points and a global minimum. This intuition is carried further by Saxe et al. (2014; 2013), where deep linear models are studied. While, from a representational perspective, deep linear models are not useful, the hope is
that the learning dynamics of such models can be mathematically understood while still being rich enough to mirror the dynamics of nonlinear networks. The findings of these works are aligned with Baldi & Hornik (1989) and suggest that one has only to go through several saddles to reach a global minimum.
These intuitions are expressed clearly for generic deep networks in Dauphin et al. (2013). The key observation of this work is that intuitions from low-dimensional spaces are usually misleading when moving to high-dimensional spaces. The work makes a connection with deep results obtained in statistical physics. In particular Fyodorov & Williams (2007); Bray & Dean (2007) showed, using the Replica Theory (Parisi, 2007), that random Gaussian error functions have a particular friendly structure. Namely, if one looks at all the critical points of the function and plots error versus the (Morse) index of the critical point (the number of negative eigenvalues of the Hessian), these points align nicely on a monotonically increasing curve. That is, all points with a low index (note that every minimum has this index equal to 0) have roughly the same performance, while critical points of high error implicitly have a large number of negative eigenvalue which means they are saddle points.
These observations align also with the theory of random matrices (Wigner, 1958) which predicts the same behaviour for the eigenvalues of a random matrix as the size of the matrix grows. The claim of Dauphin et al. (2013) is that same structure holds for neural network as well, when they become large enough. Similar claim is put forward in Sagun et al. (2014). The conjecture is very appealing as it provides a strong argument why deep networks end up performing not only well, but also reliably so. Choromanska et al. (2015) provides a study of the conjecture that rests on recasting a neural network as a spin-glass model.To obtain this result several assumptions need to be made, which the authors of the work, at that time, acknowledged that were not realistic in practice. The same line of attack is taken by Kawaguchi (2016).
Goodfellow et al. (2016) argues and provides empirical evidence that while moving from the original initialization of the model along a straight line to the solution (found via gradient descent) the loss seems to be only monotonically decreasing, which speaks towards the apparent convexity of the problem. Soudry & Carmon (2016); Safran & Shamir (2015) also look at the error surface of the neural network, providing theoretical arguments for the error surface becoming well-behaved in the case of over-parametrized models.
A different view, presented in Lin & Tegmark (2016); Shamir (2016), aligned with this work, is that the underlying easiness of optimizing deep networks does not simply rest just in the emerging structures due to high-dimensional spaces, but is rather tightly connected to the intrinsic characteristics of the data these models are run on.
3 THEORETICAL EXAMPLES AND RESULTS
We propose to analyze the error surface of rectified MLPs on finite datasets. The approach we take is a construction one. We build examples of datasets and model initializations that result in bad learning dynamics. In most examples we use ReLU units, as they are the most commonly used activation functions for both classification and regression tasks (e.g. in deep reinforcement learning (Mnih et al., 2015; 2016)). It is worth noting, however, that the phenomena we are demonstrating are not limited in nature to ReLU setup and they manifest themselves also for non-saturating activation functions like sigmoids.
3.1 LOCAL MINIMA IN A RECTIFIER-BASED REGRESSION
In Figure 1 we present 3 examples of local minima for regression using a single layer with 1, 2 and 3 hidden rectifier units on 1-dimensional data. For the sake of simplicity of our presentation we will describe in detail the case with 1 hidden neuron, the other two cases can be treated similarly. In case of one hidden neuron the regression problem becomes
arg min w,b,v,c
L(w, b, v, c) = n∑
i=1
(v · ReLU(wxi + b) + c− yi)2 . (1)
Consider a dataset D1 (see Figure 1 (a)): (x1, y1) = (5, 2), (x2, y2) = (4, 1), (x3, y3) = (3, 0), (x4, y4) = (1,−3), (x5, y5) = (−1, 3).
Proposition 1. For the dataset D1 and L defined in Equation (1) the point v = 1, b = −3, w = 1, c = 0 is a local minimum of L, which is not a global minimum.
Proof. See Appendix B.4.
Remark 1. The point (1,−3, 1, 0) is a minimum, but it is not a “strict” minimum - it is not isolated, but lies on a 1-dimensional manifold at which L ≡ 18 instead. Remark 2. One could ask whether blind spots are the only reasons for bad behaviour of rectifier nets. The answer is actually negative, and as following examples show – they can be completely absent in local optima, at the same time exisiting in a global solution!
Proposition 2. Let us consider a dataset D2 with d = 1, given by points (x1, y1) = (−1, 5), (x2, y2) = (0, 0), (x3, y3) = (1,−1), (x4, y4) = (10,−3), (x5, y5) = (11,−4), (x6, y6) = (12,−5) (Figure 1 (b)). Then, for a rectifier network with m = 2 hidden units and a squared error loss the set of weights w = (−5,−1),b = (1,−8),v = (1,−1), c = −1 is a global minimum (with perfect fit) and the set of weights w = (−3,−1),b = (4 + 13 ,−10),v = (1,−1), c = −3 is a suboptimal local minimum.
Proof. Analogous to the previous one.
Maybe surprisingly, the global solution has a blind spot - all neurons deactivate in x3. Nevertheless, the network still has a 0 training error. This shows that even though blind spots were used previously to construct very bad examples for neural nets, sometimes they are actually needed to fit the dataset.
Proposition 3. Let us consider a dataset D3 with d = 1, given by points (x1, y1) = (−1, 3), (x2, y2) = (0, 0), (x3, y3) = (1,−1), (x4, y4) = (10,−3), (x5, y5) = (11,−4), (x6, y6) = (12,−6) (Figure 1 (c)). Then, for a rectifier network with m = 3 hidden units and a squared error loss the set of weights w = (−1.5,−1.5, 1.5),b = (1, 0,−13− 16 ),v = (1, 1,−1), c = −1 is a better local minimum than the local minimum obtained for w = (−2, 1, 1),b = (3+ 23 ,−10,−11),v = (1,−1,−1), c = −3.
Proof. Completely analogous, using the fact that in each part of the space linear models are either optimal linear regression fits (if there is just one neuron active) or perfect (0 error) fit when two neurons are active and combined.
Note that again that the above construction is not relying on the blind spot phenomenon. The idea behind this example is that if, due to initial conditions, the model partitions the input space in a suboptimal way, it might become impossible to find the optimal partitioning using gradient descent. Let us call (−∞, 6) the region I, and [6,∞) region II. Both solutions in Proposition 3 are constructed in such way that each one has the best fit for the points assigned to any given region, the only difference being the number of hidden units used to describe each of them. In the local optimum two neurons are used to describe region II, while only one describes region I. Symmetrically, the better solution assigns two neurons to region I (which is more complex) and only one to region II.
We believe that the core idea behind this construction can be generalized (in a non-trivial way) to high-dimensional problems. We plan to extend the construction as future work.
Remark 3. In the examples we used only ReLU (and in one case a sigmoid) activation functions, as they are the most common used in practice. The similar examples can be constructed for different activation functions, however the constructions need some modifications and get more technically complicated.
3.2 BAD INITIALIZATION
In this section we prove some general results regarding bad initialization phenomenon.
Proposition 4. There exist an infinite amount of normalized (whitened) datasets, such that for any feed forward rectifier network architecture and an arbitrary ∈ [0, 1), there exists a normal distribution used to initialize the weights and biases initialized to 0 such that with probability at least 1− the gradient based techniques using log loss never achieve 0 training error nor they ever converge (gradient is never zero). Furthermore, this dataset can have a full rank covariance matrix and be linearly separable.
Proof. See Appendix B.1.
Even though the above construction requires control over the means of the normal distributions the weights are drawn from, as one can see in Figure 2, they do not have to be very large in practice. In particular, if one uses an initialization with σ as prescribed by LeCun et al. (1998) or Glorot & Bengio (2010) then the value of µ = 0.24 is sufficient to break the learning, even if we have 10, 000 hidden units in each of 100 hidden layers. Using fixed σ = 0.01 instead fails even with µ = 0.07.
hin+hout
). In
both cases the original papers used µ = 0. Rightmost one: Proposition 5, probability of learning failing with increasing number of layers when the initialization is fully correct.
It is worth noting that even though this observations is about the existence of such dataset, our proof is actually done by construction, meaning that we show a way to build infinite amount of such datasets (as opposed to purely existential proofs). We would like to remark that it was well known that the initialization was important for the behaviour of learning (Glorot & Bengio, 2010; LeCun et al., 1998; Sutskever et al., 2013; Pascanu et al., 2013). Here we are exploiting these ideas in order to better understand the error surface of the model.
If we do not care about the lack of convergence, and we are simply interested in learning failure, we can prove an even stronger proposition, which works for every single dataset:
Proposition 5. For every dataset, every feed forward rectifier network built for it, and every distribution used to initialize both weights and biases such that E[w] = 0,E[b] = 0,Var[w] > 0,Var[b] ≥ 0, the probability that the gradient based training of any loss function will lead to a trivial model (predicting the same label for all datapoints) goes to 1 as the number of hidden layers goes to infinity.
Proof. See Appendix B.2.
We can extend the previous proposition to show that for any regression dataset a rectifier model has at least one local minimum with a large basin of attraction (over the parameter space). Again, we rely on the blind spots of the rectified models. We show that there exists such blind spot that corresponds to a region in parameter space of same dimensionality (codimension 0). The construction relies on the fact that the dataset is finite. As such, it is bounded, and one can compute conditions for the weights of any given layer of the model such that for any datapoint all the units of that layer are deactivated. Furthermore, we show that one can obtain a better solution than the one reached from such a state. The formalization of this result is as follows.
We consider a k-layer deep regression model using m ReLU units ReLU(x) = max(0, x). Our dataset is a collection (xi, yi) ∈ Rd × R, i = 1, . . . , N . We denote hn(xi) = ReLU(Wnhn−1(xi) + bn) where the the ReLU functions are applied component-wise to the vector Wnhn−1(xi) and h0(xi) = xi. We also denote the final output of the model by M(xi) = Wkhk−1 + bk. Solving the regression problem means finding
arg min (Wn)kn=1,(bn) k n=1 L((Wn)kn=1, (bn)kn=1) = N∑ i=1 [M(xi)− yi]2 . (2)
Let us state two simple yet in our opinion useful Lemmata.
Lemma 1 (Constant input). If x1 = . . . = xN , then the solution to regression (2) has a constant outputM≡ y1+...+yNN (the mean of the values in data).
Proof. Obvious from the definitions and the fact, that y1+...+yNN = arg min c
∑N i=1(c− yi)2.
Lemma 2. If there holds W1xi < −b1 for all i-s, then the model M has a constant output. Moreover, applying local optimization does not change the values of W1, b1.
Proof. Straightforward from the definitions.
Combining these two lemmata yields:
Corollary 1. If for any 1 ≤ j ≤ k there holds Wnhn−1 < −bn for all i-s then, after the training, the modelM will output y1+...+yNN .
We will denote M({a1, . . . , aL}) = a1+...+aLL the mean of the numbers a1, . . . , aL. Definition 1. We say that the dataset (xi, yi) is decent if there exists r such that M({yp : xp = xr} 6= M({yp : p = 1, . . . , N}).
Theorem 1. Let θ = ((Wn)kn=1, (bn)kn=1) be any point in the parameter space satisfying Wnhn(xi) < −bn (coordinate-wise) for all i-s. Then
i) θ is a local minimum of the error surface,
ii) if the first layer contains at least 3 neurons and if the dataset (xi, yi) is decent, then θ is not a global minimum.
Proof. See Appendix B.3.
4 EMPIRICAL EXAMPLES
We start our examples with experiments using MNIST dataset which show that bad initialization can lead to significant obstacles in the training process.
4.1 BAD INITIALIZATION ON MNIST
Figure 3 shows the training error of rectified MLP on the MNIST dataset for different seeds and different model sizes. The learning algorithm used is Adam (Kingma & Ba, 2014) and everything except initialization, when specifically stated, follows an accepted protocol (see Appendix A). The results show that models that are not initialized in a good interval do not seem to converge to a good solution of the problem even after 1,000,000 updates. Depth does not seem to be able to resolve the bad initialization of the model. The bottom row experiments are similar to those presented in Zhang et al. (2017), though more limited in their scope. They explore the correlation between the structure in the data and learning, and, at least in appearance, they do not seem to support our working hypothesis that the structure is essential. It is worth noticing though that the initialization is even more important in that setting; destroying the structure makes the model significantly more susceptible to bad initializations than when trained on the data with unpermuted labels (second column of Figure 3, the network requires at least 400 units to be able to achieve 0 training error).
The bad initializations used in these experiments are meant to target the blind spots of the rectifiers. The main idea is that by changing the initialization of the model (the mean of the normal distribution used to sample weights) one can force all hidden units to be deactivated for the most or for all examples in the training set. This prevents said examples from being learned, even though the task might be linearly separable. The construction may seem contrived, but it has important theoretical consequences. It shows that one can not prove well behaved learning for finite sized neural networks applied to finite sized data, without taking into account the initialization or data. We formalize this idea in the Proposition 4, making the observation that the effect can be achieved by either changing the initialization of the model, or the data. In particular, by introducing specific outliers, one can force most of the data examples in the dataset to be in the blind spot of the neural network, despite being whitened.
Details of the experimental setup are given in Appendix A. The results presented in Figure 3 (bottom row), suggest that (from an optimization perspective) the important relationship is not only the one between the inputs and targets, but also between the inputs and the way the model partitions the input space (in here we focus on rectifier models which are, from a mathematical perspective, piece-wise linear functions). To empirically test if this is a viable hypothesis we consider the MNIST dataset, where we scale the inputs by a factor τ . The intuition is not to force the datasets into the blind spot of the model, but rather to concentrate most of the datapoints in very few linear regions (given by
the initialization of the MLP). While these results do not necessarily point towards the model being locked in a bad minimum, they suggest that learning becomes less well behaved (see Fig. 4).
Additional results on a simple Zig-Zag regression task are given Figure 5. The dataset itself is in the left panel, the results are visualized in the right panel. Similarly as in the MNIST case, the experiments suggest that as data becomes more concentrated in the same linear regions (of the freshly initialized model) the learning becomes really hard, even if the model has close to 3000 units.
4.2 THE JELLYFISH - SUBOPTIMAL MODELS IN CLASSIFICATION USING RELU AND SIGMOIDS
To improve our understanding of learning dynamics beyond exploiting blind spots, we look at one of the most theoretically well-studied datasets, the XOR problem. We analyze the dataset using a single hidden layer network (with either ReLU or sigmoid units).
A first observation is that while SGD can solve the task with only 2 hidden units, full batch methods do not always succeed. Replacing gradient descent with more aggressive optimizers like Adam does not seem to help, but rather tends to make it more likely to get stuck in suboptimal solutions (Table 1).
By exploiting observations made in the failure modes observed for the XOR problem, we were able to construct a similar dataset, the Jellyfish, that results in suboptimal learning dynamics. The dataset
is formed of four datapoints, where the positive class is given by [1.0, 0.0], [0.2, 0.6] and the negative one by [0.0, 1.0], [0.6, 0.2]. The datapoints can be seen in the Figure 6.
Compared to the XOR problem it seems the Jellyfish problem poses even more issues, especially for ReLU units, where with 4 hidden units one still only gets 2 out of 3 runs to end with 0 training error (when using GD). One particular observation (see Figure 6) is that in contrast with good solutions, when the model fails on this dataset, its behaviour close to the datapoints is almost linear. We argue hence, that the failure mode might come from having most datapoints concentrated in the same linear region of the model (in ReLU case), hence forcing the model to suboptimally fit these points.
5 DISCUSSION
Previous results (Dauphin et al., 2013; Saxe et al., 2014; Choromanska et al., 2015) provide insightful description of the error surface of deep models under general assumptions divorced from the specifics of the architecture. While such analysis is very valuable not only for building up the intuition but also for the development of the tools for studying neural networks, it only provides one facade of the problem. In this work we move from the generic to the specific. We show that for finite sized models/finite sized datasets one does not have a globally good behaviour of learning regardless of the model size (and even of the ratio of model size to the dataset size).
The overwhelming amount of empirical evidence points towards learning being well behaved in practice. We argue that the way to reconcile these observations is to show that the well-behaved learning dynamics are local and conditioned on the data structure, initialization and perhaps on other architectural choices. One can imagine a continuum ranging from the very specific, where every detail of the setup is important to attain good learning dynamics, to the generic, where learning is globally well behaved regardless of dataset or initialization. We believe that an important step forward in the theoretical study of the neural networks can be made by identifying where this class of models falls on that continuum. In particular, what are the most generic sets of constraints that need to be respected in order to attain the good behaviour. Our results focus on constructing counterexamples which result in a bad learning dynamics. While this does not lead directly to sufficient conditions for well-behaved systems, we hope that by carving out the space of possible conditions we are moving forward towards that goal.
Similar to Lin & Tegmark (2016) we put forward a hypothesis that the learning is only well behaved conditioned on the structure of the data. We point out, that for the purpose of learning, this structure can not be divorced from the particular initialization of the model. We postulate that learning becomes difficult if the data is structured such that there exist regions with a high density of datapoints (that belong to different classes) and the initialization results in models that assign these points to very few linear regions. While constraining the density per region alone might not be sufficient, it can provide a good starting point to understand learning for rectifier models. Another interesting question arising in that regard is what are the consequences on overfitting for enforcing a relatively low density of points per linear regions? Understanding of the structure of the error surface is an extremely challenging problem. We believe that as such, in agreement with a scientific tradition, it should be approached by gradually building up a related knowledge base, both by trying to obtain positive results (possibly under weakened assumptions, as it was done so far) and by studying the obstacles and limitations arising in concrete examples.
ACKNOWLEDGMENTS
We would want to thank Neil Rabinowitz for insightful discussions and the reviewers for their valuable remarks and suggestions.
A EXPERIMENTAL DETAILS
For the experiments depicted in Figure 3 we used MNIST dataset with data divided by 255 and one-hot encoded labels and we:
• ran 100 jobs, with the number of hidden units h sampled from [100, 2000] jointly for all the hidden layers, meaning that the model was 784− h− h− ...− h− 10,
• used Adam as an optimizer, with learning rate of 1e− 4 (and other arguments default), • each job was ran for 1,000,000 updates, • batch size used was 200.
For the experiments depicted in Figure 4 we used MNIST dataset with data divided by 255 and one-hot encoded labels and we:
• ran 400 jobs, with the number of hidden units h sampled from [100, 2000] for MNIST experiments and [100, 3000] for the Zig-Zag problem,
• we used Adam as an optimizer with learning rate randomly sampled to be either 1e− 4 or 1e− 3 (and other arguments default),
• each job was ran for 2,000,000 updates for the Zig-Zag problem and 1,200,000 updates in the MNIST case,
• batch size used was 10 for the Zig-Zag problem (dataset size is 20 points) and 50 for MNIST experiment.
• Initialization of weights and biases was from a Gaussian with 0 mean and standard deviation 1√
fanin .
B PROOFS
B.1 PROOF OF PROPOSITION 4
For simplicity, let us consider a network with only one hidden layer of h units and a binary classification task, implying we have a single output neuron. Analogous result holds for any number of classes and for any arbitrary depth of the network as well. We use the following notation: h(xi) is a vector of activations of hidden units when presented with i-th training sample, M(xi) is the activation of the output neuron given the same sample, W is matrix of hidden weights, b is a vector of biases in the hidden layer and finally v, c are the weights and bias in the output layer. The whole classification becomes M(xi) = vReLU(Wxi + b) + c. Let us now consider a dataset where N − 1 points have all features negative, and a single point, denoted xi∗ (with positive label) which has all the features positive. We can always find such points that their coordinate-wise mean is equal to 0 and their standard deviation equals 1, since we can place all N − 1 points very close to the origin, and the point xi∗ arbitrary far away in the positive part of the input space. Our dataset is therefore normalized (whitened) and it can have a full rank covariance matrix (since the construction does not depend on nothing besides signs of the features). We want to compute P (∀i 6=i∗h(xi) = 0 ∧ h(xi∗) > 0). Since, by construction ∀i 6=i∗xi < 0 and xi∗ > 0, it is also true that if all the weights W are nonnegative (and at least one is positive), then all the activations (after ReLU) of the hidden units will be 0 besides the one positive activation of h(xi∗) which comes directly from the assumption that biases are initialized to 01. Consequently
P (∀i 6=i∗h(xi) = 0 ∧ h(xi∗) > 0) ≥ P (W > 0), and given that weights initializations are independent samples from the same distribution we get
P (W > 0) = dh∏ i=1 ∫ ∞ 0 N (µ, σ2) = (∫ ∞ 0 N (µ, σ2) )dh ,
where µ, σ are parameters of the distribution we use to initialize weights and d is input space dimensionality. All that is left is to show that during any gradient based optimization these weights will not be corrected, which requires one more assumption - that the output weights are positive as well. If this is true, then ∀i 6=i∗M(xi) = 0 (again using that the output bias is zero) andM(xi∗) > 0 and of course P (v > 0) = ( ∫∞ 0 N (µ, σ2))h. Now we have a fully initialized model which maps all the samples to 0, and one positive sample to some positive value. Consequently, given the fact that we use log loss, there holds ∂L/∂vk > 0 and ∂L/∂wkl ≥ 0 for all k, l. Indeed, since these changes are all increasing the probability of good classification of xi∗ and all remaining points are in inactive part of ReLUs thus they cannot contribute to partial derivatives. Therefore, during any gradient based (as well as stochastic and mini-batch based) learning the projection of samples mapped to 0 will not change, and the projection of xi∗ will grow to infinity (so the sigmoid approaches 1). Consequently we constructed an initialization scheme which with probability at least
P (W,v > 0) = P (W > 0)P (v > 0) = (∫ ∞ 0 N (µ, σ2) )dh+h = 1− , (3)
gives the initial conditions of the net, where despite learning with log loss, we always classify all, arbitrary labeled N −1 points to the same label (since they are all mapped to the same output value) and we classify the unique positive sample with the valid label. Furthermore - the optimization never finishes, and there exists a network with a better accuracy, as we can label N −1 points in any manner, including making it linearly separable by labeling according to the first feature only.
In order to generalize to arbitrary number of layers, we would similarly force all the weights to be positive, thus with parametrization of k-layer deep network θ = (Wn)kn=1 we would simply get
P (∀1≤n≤kWn > 0) = k∏
n=1
P (Wn > 0) = (∫ ∞ 0 N (µ, σ2) )dh+(k−2)h2+h = 1− ,
1We can still construct similar proof for biases taken from normal distributions as well.
and finally, if the biases are not zero, but have some some arbitrary values (fixed or sampled) we simply adjust the size of the weights accordingly. Instead of having them bigger than 0 we would compute probability of having them big enough to make the whole first layer to produce 0s for every point besides xi∗ and analogously for the remaining layers.
Furthermore, it is possible to significantly increase the probability of failure if we do not mind the situation in which the learning process is not even starting. Proposition 5 addresses this situation.
B.2 PROOF OF PROPOSITION 5
Proof. Let us notice, that since we are using ReLU activations, the activations on jth hidden layer hj are non-negative. Consequently, if the following layer has only negative weights and non-positive biases, then hj+1 ≡ bj+1 (as all the ReLUs are inactive), so the network has to output exactly the same value for every single point. Furthermore, during the gradient based learning we will never change those weights as gradients going through ReLUs will be equal to zero.
Let us now consider a deep rectifier network with k hidden layers with h neurons each. If only k > 2 we can use the above observation to compute for every j > 1:
P (All neurons inactive in hidden layer j) ≥ P (bj ≤ 0,Wj < 0) = P (bj ≤ 0)P (Wj < 0),
and consequently, due to the assumptions about expected values
P (Learning fails with k hidden layers) ≥ P (∃j>1All neurons inactive in hidden layer j) ≥ P (∃j>1∀l,ibjl ≤ 0 ∧Wji < 0) = 1− P (∀j>1∃lbjl > 0 ∨ ∃iWji ≥ 0)
= 1− k−1∏ j=1 P (∃lbjl > 0 ∨ ∃iWji ≥ 0)
As every layer has the same number h of neurons,the values (∃lbjl > 0 ∨ ∃iWji ≥ 0) are equal for every j. Therefore
1− k−1∏ j=1 P (∃lbjl > 0 ∨ ∃iWji ≥ 0) = 1− P (∃lb1l > 0 ∨ ∃iW1i ≥ 0)k−1
= 1− (1− P (∀ib1i ≤ 0 ∧W1i < 0))k−1
= 1− [1− cdfb(0)hcdfW(0)h 2 ]k−1.
Due to assumptions about distributions of biases and weights we know that
0 < cdfW(0) < 1, 0 ≤ cdfb(0) < 1
so lim k→∞ 1− [1− cdfb(0)hcdfW(0)h 2 ]k−1 = 1.
B.3 PROOF OF THEOREM 1
Claim i) is a direct consequence of Corollary 1. It remains to prove ii). For that it is sufficient to show an example of a set of weighs θ̂ = ((Ŵn)kn=1, (b̂n) k n=1) such that L((Wn)kn=1, (bn)kn=1) > L((Ŵn)kn=1, (b̂n)kn=1). Let r be such that M({yp : xp = xr}) 6= M({yp : p = 1, . . . , N}). Such point exists by assumption that the dataset is decent. LetH be a hyperplane passing through xr such that none of the points xs 6= xr lies onH. Then there exists a vector v such that |vT (xs−xr)| > 2 for all xs 6= xr. Let γ = vTxr. We define W1 in such a way that the first row of W1 is v , the second row is 2v and the third one is v again, and if the first layer has more than 3 neurons, we put all the remaining rows of W1 to be equal zero. We choose the first three biases of b1 to be −γ + 1, −2γ and −γ − 1 respectively. We denote µ = M({yp : xp 6= xr}) and ν = M({yp : xp = xr}). We then choose W2 to be a matrix whose first row is (ν − µ, µ − ν, ν − µ, 0, . . . , 0) and the other
rows are equal to 0. Finally, we choose the bias vector b2 = (µ, 0, . . . , 0)T . If our network has only one layer the output is
(ν −µ)ReLU(vTxp− γ+ 1)− (ν −µ)ReLU(2vTxp− 2γ) + (ν −µ)ReLU(vTxp− γ− 1) +µ.
For every xp = xr this yields (ν − µ) · 1 − 0 + 0 + µ = ν. For any xp 6= xr we either have vTxp − γ < −2 yielding 0− 0 + 0 + µ = µ or vTxp − γ > 2 yielding (ν − µ)(vTxp − γ + 1)− (ν − µ)(2vTxp − 2γ) + (ν − µ)(vTxp − γ − 1) + µ = µ. In case the network has more than 1 hidden layer we set all Wn = I (identity matrix) and bn = 0 for n = 3, . . . , k. If we denote µ̄ = M({yp : p = 1, . . . , N}) (mean of all labels), we get:
L((Ŵn)kn=1, (b̂n)kn=1) = ∑
xp 6=xr
(yi − µ)2 + ∑
xp=xr (yi − ν)2 <∑ xp 6=xr (yi − µ̄)2 + ∑ xp=xr (yi − µ̄)2 = ∑ yi (yi − µ̄)2 = L((Wn)kn=1, (bn)kn=1).
We used the fact that for any finite set A the value M(A) is a strict minimum of f(c) = ∑
a∈A(a− c)2 and the assumption that ν 6= µ̄.
B.4 PROOF OF PROPOSITION 1
There holdsL(1,−3, 1, 0) = 0+0+0+9+9 = 18, andL(−7,−4, 1, 0) = 4+1+0+9+0 = 14, thus (1,−3, 1, 0) cannot be a global minimum. It remains to prove that (1,−3, 1, 0) is a local minimum, i.e. that L(1 + δw,−3 + δb, 1 + δv, δc) ≥ L(1,−3, 1, 0) for |δw|, |δb|, |δv|, |δc| sufficiently small. We need to consider two cases: ReLU activated at x3. In that case
L(1 + δw,−3 + δb, 1 + δv, δc) =
((1 + δv)(3 + 3δw − 3 + δb) + δc)2 + ((1 + δv)(4 + 4δw − 3 + δb) + δc − 1)2 + ((1 + δv)(5 + 5δw − 3 + δb) + δc − 2)2 + (δc + 3)2 + (δc − 3)2.
We introduce new variables x = (δw + 1)(1 + δv) − 1, y = (δb − 3)(1 + δv) + 3, z = δc. The formula becomes
(3x+ y + z) 2 + (4x+ y + z) 2 + (5x+ y + z) 2 + 2z2 + 18 ≥ 18,
which ends the proof in this case. ReLU deactivated at x3. In that case
L(1 + δw,−3 + δb, 1 + δv, δc) = δ2c + ((1 + δv)(4 + 4δw − 3 + δb) + δc − 1) 2 +
((1 + δv)(5 + 5δw − 3 + δb) + δc − 2)2 + (δc + 3)2 + (δc − 3)2 = (4x+ y + z) 2 + (5x+ y + z) 2 + 3z2 + 18 ≥ 18
(we used x = (δw + 1)(1 + δv)− 1, y = (δb − 3)(1 + δv) + 3, z = δc again). Note that due to the assumption that |δw|, |δb|, |δv|, |δc| are sufficiently small the ReLU is always activated at x1, x2 and deactivated at x4, x5.

This paper studies the error surface of deep rectifier networks, giving specific examples for which the error surface has local minima. Several experimental results show that learning can be trapped at apparent local minima by a variety of factors ranging from the nature of the dataset to the nature of the initializations. This paper develops a lot of good intuitions and useful examples of ways that training can go awry.
Even though the examples constructed in this paper are contrived, this does not necessarily remove their theoretical importance. It is very useful to have simple examples where things go wrong. However the broader theoretical framing of the paper appears to be going after a strawman.
“The underlying easiness of optimizing deep networks does not simply rest just in the emerging structures due to high dimensional spaces, but is rather tightly connected to the intrinsic characteristics of the data these models are run on.” I believe this perspective is already contained in several of the works cited as not belonging to this perspective. Choromanska et al., for instance, analyze Gaussian inputs, and so clearly make claims based on characteristics of the data the models are run on. More broadly, the loss function is determined jointly by the dataset and the model parameters, and so no account of the error surface can be separated from dataset properties. It is not clear to me what ‘emerging structures due to high dimensional spaces’ are, or what they could be, that would make them independent of the dataset and initial model parameters. The emerging structure of the error surface is necessarily related to the dataset and model parameters.
Again, a key worry with this paper is that it is aiming at a strawman: replica methods characterize average behavior for infinite systems, so it is not surprising that specific finite sized systems might yield poor optimization landscapes. The paper seems surprised that training can be broken with a bad initialization, but initialization is known to be critical, even for linear networks: saddle points are not innocuous, with bad initializations dramatically slowing learning (e.g. Saxe et al. 2014).
It seems like the proof of proposition 5 may have an error. Suppose cdf_b(0) = 0 and cdf_W(0)=1/2. We have P(learning fails) >= 1 - 1/2^{h^2(k-1)}, meaning that the probability of failure _increases_ as the number of hidden units increases. It seems like it should rather be (ignoring the bias) p(fails) >= 1 - [ 1 - p(w<0)^h^2]^{k-1}. In this case the limit as k-> infinity depends on how h scales with k, so it is no longer necessarily true that “one does not have a globally good behaviour of learning regardless of the model size.”
The paper also appears to insufficiently distinguish between local minima and saddle points. Section 3.1 states it shows training being stuck in a local minimum, but this is based on training with a fixed budget of epochs. It is not possible to tell whether this result reflects a genuine local minimum or a saddle point based on simulation results.
It may also be the case that, while rectifiers suffer from genuine blind spots, sigmoid or soft rectifier nonlinearities may not. On the XOR problem with two hidden nodes, for instance, it was thought that were local minima but in fact there are none (e.g. L. Hamey, “Analysis of the error surface of the XOR network with two hidden nodes,” 1995).
If the desire is simply to show that training does not converge for particular finite problems, much simpler counterexamples can be constructed and would suffice: set all hidden unit weights to zero, for instance.
In the response to prereview questions, the authors write ‘If the “complete characterization” [of the error surface] was indeed universally valid, we would not be able to break the learning with the initialization’ but, as mentioned previously, the basic results for even deep linear networks show that a bad initialization (at or near a saddle point) will break learning. Again, it seems this paper is attacking a straw man along the lines of “nothing can possibly go wrong with neural network training.” No prior theoretical result claims this.
The Figure 2 explanation seems counterintuitive to me. Simply scaling the input, if the weight matrices are initialized with zero biases, will not change the regions over which each ReLU activates. That is, this manipulation does not achieve the goal of concentrating “most of the data points in very few linear regions.” A far more likely explanation is that the much weaker scaling has not been compensated by the learning algorithm, but the algorithm would converge if run longer. The response notes that training has been conducted for an order of magnitude longer than required for the unscaled input to converge, but the scaling on the data is not one but five orders of magnitude—and indeed the training does converge without issue for scaling up to four orders of magnitude. The response notes that Adam should compensate for the scaling factor, but this depends on the details of the Adam implementation—the epsilon factor used to protect against division by zero, for example.
This paper contains many interesting results, but a variety of small technical concerns remain.


We present a variational approximation to the information bottleneck of Tishby et al. (1999). This variational approach allows us to parameterize the information bottleneck model using a neural network and leverage the reparameterization trick for efficient training. We call this method “Deep Variational Information Bottleneck”, or Deep VIB. We show that models trained with the VIB objective outperform those that are trained with other forms of regularization, in terms of generalization performance and robustness to adversarial attack.
1 INTRODUCTION
We adopt an information theoretic view of deep networks. We regard the internal representation of some intermediate layer as a stochastic encoding Z of the input source X , defined by a parametric encoder p(z|x;θ).1 Our goal is to learn an encoding that is maximally informative about our target Y , measured by the mutual information between our encoding and the target I(Z, Y ;θ), where
I(Z, Y ;θ) = ∫ dx dy p(z, y|θ) log p(z, y|θ)
p(z|θ)p(y|θ) . 2 (1)
Given the data processing inequality, and the invariance of the mutual information to reparameterizations, if this was our only objective we could always ensure a maximally informative representation by taking the identity encoding of our data (Z = X), but this is not a useful representation of our data. Instead we would like to find the best representation we can obtain subject to a constraint on its complexity. A natural and useful constraint to apply is on the mutual information between our encoding and the original data, I(X,Z) ≤ Ic, where Ic is the information constraint. This suggests the objective:
max θ
I(Z, Y ;θ) s.t. I(X,Z;θ) ≤ Ic . (2)
Equivalently, with the introduction of a Lagrange multiplier β, we can maximize the objective function RIB(θ) = I(Z, Y ;θ)− βI(Z,X;θ). (3) Here our goal is to learn an encoding Z that is maximally expressive about Y while being maximally compressive about X , where β ≥ 0 controls the tradeoff.3 This approach is known as the information bottleneck (IB), and was first proposed in Tishby et al. (1999). Intuitively, the first term in RIB encourages Z to be predictive of Y ; the second term encourages Z to “forget” X . Essentially it forces Z to act like a minimal sufficient statistic of X for predicting Y .
The IB principle is appealing, since it defines what we mean by a good representation, in terms of the fundamental tradeoff between having a concise representation and one with good predictive power (Tishby & Zaslavsky, 2015a). The main drawback of the IB principle is that computing mutual information is, in general, computationally challenging. There are two notable exceptions: the first
1 In this work, X,Y, Z are random variables, x, y, z and x,y, z are instances of random variables, and F (·;θ) and f(·;θ) are functionals or functions parameterized by θ.
2 Note that in the present discussion, Y is the ground truth label which is independent of our parameters so p(y|θ) = p(y).
3 Note that, in our notation, large β results in a highly compressed representation. In some works, the IB principle is formulated as the minimization of I(Z,X)− βI(Z, Y ), in which case large β corresponds to high mutual information between Z and Y , and hence low compression.
is when X , Y and Z are all discrete, as in Tishby et al. (1999); this can be used to cluster discrete data, such as words. The second case is when X , Y and Z are all jointly Gaussian (Chechik et al., 2005). However, these assumptions both severely constrain the class of learnable models.
In this paper, we propose to use variational inference to construct a lower bound on the IB objective in Equation 3. We call the resulting method VIB (variational information bottleneck). By using the reparameterization trick (Kingma & Welling, 2014), we can use Monte Carlo sampling to get an unbiased estimate of the gradient, and hence we can optimize the objective using stochastic gradient descent. This allows us to use deep neural networks to parameterize our distributions, and thus to handle high-dimensional, continuous data, such as images, avoiding the previous restrictions to the discrete or Gaussian cases.
We also show, by a series of experiments, that stochastic neural networks, fit using our VIB method, are robust to overfitting, since VIB finds a representation Z which ignores as many details of the inputX as possible. In addition, they are more robust to adversarial inputs than deterministic models which are fit using (penalized) maximum likelihood estimation. Intuitively this is because each input image gets mapped to a distribution rather than a unique Z, so it is more difficult to pass small, idiosyncratic perturbations through the latent bottleneck.
2 RELATED WORK
The idea of using information theoretic objectives for deep neural networks was pointed out in Tishby & Zaslavsky (2015b). However, they did not include any experimental results, since their approach for optimizing the IB objective relied on the iterative Blahut Arimoto algorithm, which is infeasible to apply to deep neural networks.
Variational inference is a natural way to approximate the problem. Variational bounds on mutual information have previously been explored in Agakov (2004), though not in conjunction with the information bottleneck objective. Mohamed & Rezende (2015) also explore variational bounds on mutual information, and apply them to deep neural networks, but in the context of reinforcement learning. We recently discovered Chalk et al. (2016), who independently developed the same variational lower bound on the IB objective as us. However, they apply it to sparse coding problems, and use the kernel trick to achieve nonlinear mappings, whereas we apply it to deep neural networks, which are computationally more efficient. In addition, we are able to handle large datasets by using stochastic gradient descent, whereas they use batch variational EM.
In the supervised learning literature, our work is related to the recently proposed confidence penalty (entropy regularization) method of (Pereyra et al., 2016). In this work, they fit a deterministic network by optimizing an objective that combines the usual cross entropy loss with an extra term which penalizes models for having low entropy predictive distributions. In more detail, their cost function has the form
JCP = 1
N N∑ n=1 [H(p(y|yn), p(y|xn))− βH(p(y|xn))] (4)
where H(p, q) = −∑y p(y) log q(y) is the cross entropy, H(p) = H(p, p) is the entropy, p(y|yn) = δyn(y) is a one-hot encoding of the label yn, and N is the number of training examples. (Note that setting β = 0 corresponds to the usual maximum likelihood estimate.) In (Pereyra et al., 2016) they show that CP performs better than the simpler technique of label smoothing, in which we replace the zeros in the one-hot encoding of the labels by > 0, and then renormalize so that the distribution still sums to one. We will compare our VIB method to both the confidence penalty method and label smoothing in Section 4.1.
In the unsupervised learning literature, our work is closely related to the work in Kingma & Welling (2014) on variational autoencoders. In fact, their method is a special case of an unsupervised version of the VIB, but with the β parameter fixed at 1.0, as we explain in Appendix B. The VAE objective, but with different values of β, was also explored in Higgins et al. (2016), but from a different perspective.
The method of Wang et al. (2016b) proposes a latent variable generative model of both x and y; their variational lower bound is closely related to ours, with the following differences. First, we do
not have a likelihood term for x, since we are in the discriminative setting. Second, they fix β = 1, since they do not consider compression.
Finally, the variational fair autoencoder of Louizos et al. (2016) shares with our paper the idea of ignoring parts of the input. However, in their approach, the user must specify which aspects of the input (the so-called “sensitive” parts) to ignore, whereas in our method, we can discover irrelevant parts of the input automatically.
3 METHOD
Following standard practice in the IB literature, we assume that the joint distribution p(X,Y, Z) factors as follows:
p(X,Y, Z) = p(Z|X,Y )p(Y |X)p(X) = p(Z|X)p(Y |X)p(X) (5) i.e., we assume p(Z|X,Y ) = p(Z|X), corresponding to the Markov chain Y ↔ X ↔ Z. This restriction means that our representation Z cannot depend directly on the labels Y . (This opens the door to unsupervised representation learning, which we will discuss in Appendix B.) Besides the structure in the joint data distribution p(X,Y ), the only content at this point is our model for the stochastic encoder p(Z|X), all other distributions are fully determined by these and the Markov chain constraint.
Recall that the IB objective has the form I(Z, Y ) − βI(Z,X). We will examine each of these expressions in turn. Let us start with I(Z, Y ). Writing it out in full, this becomes
I(Z, Y ) = ∫ dy dz p(y, z) log p(y, z)
p(y)p(z) =
∫ dy dz p(y, z) log
p(y|z) p(y) . (6)
where p(y|z) is fully defined by our encoder and Markov Chain as follows: p(y|z) = ∫ dx p(x, y|z) = ∫ dx p(y|x)p(x|z) = ∫ dx
p(y|x)p(z|x)p(x) p(z) . (7)
Since this is intractable in our case, let q(y|z) be a variational approximation to p(y|z). This is our decoder, which we will take to be another neural network with its own set of parameters. Using the fact that the Kullback Leibler divergence is always positive, we have
KL[p(Y |Z), q(Y |Z)] ≥ 0 =⇒ ∫ dy p(y|z) log p(y|z) ≥ ∫ dy p(y|z) log q(y|z) , (8)
and hence
I(Z, Y ) ≥ ∫ dy dz p(y, z) log
q(y|z) p(y)
(9)
= ∫ dy dz p(y, z) log q(y|z)− ∫ dy p(y) log p(y) (10)
= ∫ dy dz p(y, z) log q(y|z) +H(Y ) . (11)
Notice that the entropy of our labels H(Y ) is independent of our optimization procedure and so can be ignored.
Focusing on the first term in Equation 11, we can rewrite p(y, z) as p(y, z) = ∫ dx p(x, y, z) =∫
dx p(x)p(y|x)p(z|x) (leveraging our Markov assumption), which gives us a new lower bound on the first term of our objective:
I(Z, Y ) ≥ ∫ dx dy dz p(x)p(y|x)p(z|x) log q(y|z) . (12)
This only requires samples from both our joint data distribution as well as samples from our stochastic encoder, while it requires we have access to a tractable variational approximation in q(y|z). We now consider the term βI(Z,X):
I(Z,X) = ∫ dz dx p(x, z) log
p(z|x) p(z) =
∫ dz dx p(x, z) log p(z|x)− ∫ dz p(z) log p(z) . (13)
In general, while it is fully defined, computing the marginal distribution of Z, p(z) =∫ dx p(z|x)p(x), might be difficult. So let r(z) be a variational approximation to this marginal.
Since KL[p(Z), r(Z)] ≥ 0 =⇒ ∫ dz p(z) log p(z) ≥ ∫ dz p(z) log r(z), we have the following upper bound:
I(Z,X) ≤ ∫ dx dz p(x)p(z|x) log p(z|x)
r(z) . (14)
Combining both of these bounds we have that I(Z, Y )− βI(Z,X) ≥ ∫ dx dy dz p(x)p(y|x)p(z|x) log q(y|z)
− β ∫ dx dz p(x)p(z|x) log p(z|x)
r(z) = L . (15)
We now discuss how to compute the lower bound L in practice. We can approximate p(x, y) = p(x)p(y|x) using the empirical data distribution p(x, y) = 1N ∑N n=1 δxn(x)δyn(y), and hence we can write
L ≈ 1 N N∑ n=1 [∫ dz p(z|xn) log q(yn|z)− β p(z|xn) log p(z|xn) r(z) ] . (16)
Suppose we use an encoder of the form p(z|x) = N (z|fµe (x), fΣe (x)), where fe is an MLP which outputs both the K-dimensional mean µ of z as well as the K ×K covariance matrix Σ. Then we can use the reparameterization trick (Kingma & Welling, 2014) to write p(z|x)dz = p( )d , where z = f(x, ) is a deterministic function of x and the Gaussian random variable . This formulation has the important advantage that the noise term is independent of the parameters of the model, so it is easy to take gradients.
Assuming our choice of p(z|x) and r(z) allows computation of an analytic Kullback-Leibler divergence, we can put everything together to get the following objective function, which we try to minimize:
JIB = 1
N N∑ n=1 E ∼p( ) [− log q(yn|f(xn, ))] + βKL [p(Z|xn), r(Z)] . (17)
As in Kingma & Welling (2014), this formulation allows us to directly backpropagate through a single sample of our stochastic code and ensure that our gradient is an unbiased estimate of the true expected gradient.4
4 EXPERIMENTAL RESULTS
In this section, we present various experimental results, comparing the behavior of standard deterministic networks to stochastic neural networks trained by optimizing the VIB objective.
4.1 BEHAVIOR ON MNIST
We start with experiments on unmodified MNIST (i.e. no data augmentation). In order to pick a model with some “headroom” to improve, we decided to use the same architecture as in the (Pereyra et al., 2016) paper, namely an MLP with fully connected layers of the form 784 - 1024 - 1024 - 10, and ReLu activations. (Since we are not exploiting spatial information, this correpsonds to the “permutation invariant” version of MNIST.) The performance of this baseline is 1.38% error. (Pereyra et al., 2016) were able to improve this to 1.17% using their regularization technique. We were able to improve this to 1.13% using our technique, as we explain below.
In our method, the stochastic encoder has the form p(z|x) = N (z|fµe (x), fΣe (x)), where fe is an MLP of the form 784 − 1024 − 1024 − 2K, where K is the size of the bottleneck. The first K outputs from fe encode µ, the remaining K outputs encode σ (after a softplus transform).
4 Even if our choice of encoding distribution and variational prior do not admit an analytic KL, we could similarly reparameterize through a sample of the divergence (Kingma & Welling, 2014; Blundell et al., 2015).
The decoder is a simple logistic regression model of the form q(y|z) = S(y|fd(z)), where S(a) = [exp(ac)/ ∑C c′=1 exp(ac′)] is the softmax function, and fd(z) = Wz + b maps the K dimensional latent code to the logits of the C = 10 classes. (In later sections, we consider more complex decoders, but here we wanted to show the benefits of VIB in a simple setting.)
Finally, we treat r(z) as a fixed K-dimensional spherical Gaussian, r(z) = N (z|0, I). We compare our method to the baseline MLP. We calso consider the following deterministic limit of our model, when β = 0. In this case, we obtain the following objective function:
JIB0 = − 1
N N∑ n=1 Ez∼N (fµe (xn),fΣe (xn)) [logS(yn|fd(z)] (18)
When β → 0, we observe the VIB optimization process tends to make fΣe (x) → 0, so the network becomes nearly deterministic. In our experiments we also train an explicitly deterministic model that has the same form as the stochastic model, except that we just use z = fµe (x) as the hidden encoding, and drop the Gaussian layer.
4.1.1 HIGHER DIMENSIONAL EMBEDDING
To demonstrate that our VIB method can achieve competitive classification results, we compared against a deterministic MLP trained with various forms of regularization. We use a K = 256 dimensional bottleneck and a diagonal Gaussian for p(z|x). The networks were trained using TensorFlow for 200 epochs using the Adam optimizer (Kingma & Ba, 2015) with a learning rate of 0.0001. Full hyperparameter details can be found in Appendix A.
The results are shown in Table 1. We see that we can slightly outperform other forms of regularization that have been proposed in the literature while using the same network for each. Of course, the performance varies depending on β. These results are not state of the art, nor is our main focus of our work to suggest that VIB is the best regularization method by itself, which would require much more experimentation. However, using the same architecture for each experiment and comparing to VIB as the only source of regularization suggests VIB works as a decent regularizer in and of itself. Figure 1(a) plots the train and test error vs β, averaged over 5 trials (with error bars) for the case where we use a single Monte Carlo sample of z when predicting, and also for the case where we average over 12 posterior samples (i.e., we use p(y|x) = 1S ∑S s=1 q(y|zs) for zs ∼ p(z|x), where S = 12). In our own investigations, a dozen samples seemed to be sufficient to capture any additional benefit the stochastic evaluations had to offer in this experiment5.
We see several interesting properties in Figure 1(a). First, we notice that the error rate shoots up once β rises above the critical value of β ∼ 10−2. This corresponds to a setting where the mutual information between X and Z is less than log2(10) bits, so the model can no longer represent the fact that there are 10 different classes. Second, we notice that, for small values of β, the test error
5 A dozen samples wasn’t chosen for any particular reason, except the old addage that a dozen samples are sufficient, as mirrored in David MacKay’s book (MacKay, 2003). They proved sufficient in this case.
is higher than the training error, which indicates that we are overfitting. This is because the network learns to be more deterministic, forcing σ ≈ 0, thus reducing the benefits of regularization. Third, we notice that for intermediate values of β, Monte Carlo averaging helps. Interestingly, the region with the best performance roughly corresponds to where the added benefit from stochastic averaging goes away, suggesting an avenue by which one could try to optimize β using purely statistics on the training set without a validation set. We have not extensively studied this possibility yet.
In Figure 1(c), we plot the IB curve, i.e., we plot I(Z, Y ) vs I(Z,X) as we vary β. As we allow more information from the input through to the bottleneck (by lowering β), we increase the mutual information between our embedding and the label on the training set, but not necessarily on the test set, as is evident from the plot.
In Figure 1(d) we plot the second term in our objective, the upper bound on the mutual information between the images X and our stochastic encoding Z, which in our case is simply the relative entropy between our encoding and the fixed isotropic unit Gaussian prior. Notice that the y-axis is a logarithmic one. This demonstrates that our best results (when β is between 10−3 and 10−2) occur where the mutual information between the stochastic encoding and the images is on the order of 10 to 100 bits.
4.1.2 TWO DIMENSIONAL EMBEDDING
To better understand the behavior of our method, we refit our model to MNIST using a K = 2 dimensional bottleneck, but using a full covariance Gaussian. (The neural net predicts the mean and the Cholesky decomposition of the covariance matrix.) Figure 1(b) shows that, not surprisingly, the classification performance is worse (note the different scaled axes), but the overall trends are the
same as in the K = 256 dimensional case. The IB curve (not shown) also has a similar shape to before, except now the gap between training and testing is even larger.
Figure 2 provides a visualization of what the network is doing. We plot the posteriors p(z|x) as a 2d Gaussian ellipse (representing the 95% confidence region) for 1000 images from the test set. Colors correspond to the true class labels. In the background of each plot is the entropy of the variational classifier q(y|z) evaluated at that point.
We see several interesting properties. First, as β increases (so we pass less information through), the embedding covariances increase in relation to the distance between samples, and the classes start to overlap. Second, once β passes a critical value, the encoding “collapses”, and essentially all the class information is lost. Third, there is a fair amount of uncertainty in the class preditions (q(y|z)) in the areas between the class embeddings. Fourth, for intermediate values of β (say 10−1 in Figure 2(b)), predictive performance is still good, even though there is a lot of uncertainty about where any individual image will map to in comparison to other images in the same class. This means it would be difficult for an outside agent to infer which particular instance the model is representing, a property which we will explore more in the following sections.
4.2 BEHAVIOR ON ADVERSARIAL EXAMPLES
Szegedy et al. (2013) was the first work to show that deep neural networks (and other kinds of classifiers) can be easily “fooled” into making mistakes by changing their inputs by imperceptibly small amounts. In this section, we will show how training with the VIB objective makes models significantly more robust to such adversarial examples.
4.2.1 TYPES OF ADVERSARIES
Since the initial work by Szegedy et al. (2013) and Goodfellow et al. (2014), many different adversaries have been proposed. Most attacks fall into three broad categories: optimization-based attacks (Szegedy et al., 2013; Carlini & Wagner, 2016; Moosavi-Dezfooli et al., 2016; Papernot et al., 2015; Robinson & Graham, 2015; Sabour et al., 2016), which directly run an optimizer such as L-BFGS or ADAM (Kingma & Ba, 2015) on image pixels to find a minimal perturbation that changes the model’s classification; single-step gradient-based attacks (Goodfellow et al., 2014; Kurakin et al., 2016; Huang et al., 2015), which choose a gradient direction of the image pixels at some loss and then take a single step in that direction; and iterative gradient-based attacks (Kurakin et al., 2016),
which take multiple small steps along the gradient direction of the image pixels at some loss, recomputing the gradient direction after each step.6
Many adversaries can be formalized as either untargeted or targeted variants. An untargeted adversary can be defined as A(X,M) → X ′, where A(.) is the adversarial function, X is the input image, X ′ is the adversarial example, and M is the target model. A is considered successful if M(X) 6= M(X ′). Recently, Moosavi-Dezfooli et al. (2016) showed how to create a “universal” adversarial perturbation δ that can be added to any image X in order to make M(X + δ) 6= M(X) for a particular target model.
A targeted adversary can be defined as A(X,M, l)→ X ′, where l is an additional target label, and A is only considered successful if M(X ′) = l.7 Targeted attacks usually require larger magnitude perturbations, since the adversary cannot just “nudge” the input across the nearest decision boundary, but instead must force it into a desired decision region.
In this work, we focus on the Fast Gradient Sign (FGS) method proposed in Goodfellow et al. (2014) and the L2 optimization method proposed in Carlini & Wagner (2016). FGS is a standard baseline attack that takes a single step in the gradient direction to generate the adversarial example. As originally described, FGS generates untargeted adversarial examples. On MNIST, Goodfellow et al. (2014) reported that FGS could generate adversarial examples that fooled a maxout network approximately 90% of the time with = 0.25, where is the magnitude of the perturbation at each pixel. The L2 optimization method has been shown to generate adversarial examples with smaller perturbations than any other method published to date, which were capable of fooling the target network 100% of the time. We consider both targeted attacks and untargeted attacks for the L2 optimization method.8
4.2.2 ADVERSARIAL ROBUSTNESS
There are multiple definitions of adversarial robustness in the literature. The most basic, which we shall use, is accuracy on adversarially perturbed versions of the test set, called adversarial examples.
It is also important to have a measure of the magnitude of the adversarial perturbation. Since adversaries are defined relative to human perception, the ideal measure would explicitly correspond to how easily a human observer would notice the perturbation. In lieu of such a measure, it is common to compute the size of the perturbation using L0, L1, L2, and L∞ norms (Szegedy et al., 2013; Goodfellow et al., 2014; Carlini & Wagner, 2016; Sabour et al., 2016). In particular, the L0 norm measures the number of perturbed pixels, the L2 norm measures the Euclidean distance between X and X ′, and the L∞ norm measures the largest single change to any pixel.
4.2.3 EXPERIMENTAL SETUP
We used the same model architectures as in Section 4.1, using a K = 256 bottleneck. The architectures included a deterministic (base) model trained by MLE; a deterministic model trained with dropout (the dropout rate was chosen on the validation set); and a stochastic model trained with VIB for various values of β.
For the VIB models, we use 12 posterior samples of Z to compute the class label distribution p(y|x). This helps ensure that the adversaries can get a consistent gradient when constructing the perturbation, and that they can get a consistent evaluation when checking if the perturbation was successful
6 There are also other adversaries that don’t fall as cleanly into those categories, such as “fooling images” from Nguyen et al. (2014), which remove the human perceptual constraint, generating regular geometric patterns or noise patterns that networks confidently classify as natural images; and the idea of generating adversaries by stochastic search for images near the decision boundary of multiple networks from Baluja et al. (2015).
7 Sabour et al. (2016) proposes a variant of the targeted attack, A(XS ,M,XT , k)→ X ′S , where XS is the source image, XT is a target image, and k is a target layer in the model M . A produces X ′S by minimizing the difference in activations of M at layer k between XT and X ′S . The end result of this attack for a classification network is still that M(X ′S) yields a target label implicitly specified by XT in a successful attack.
8 Carlini & Wagner (2016) shared their code with us, which allowed us to perform the attack with exactly the same parameters they used for their paper, including the maximum number of iterations and maximum C value (see their paper for details).
(i.e., it reduces the chance that the adversary “gets lucky” in its perturbation due to an untypical sample). We also ran the VIB models in “mean mode”, where the σs are forced to be 0. This had no noticeable impact on the results, so all reported results are for stochastic evaluation with 12 samples.
4.2.4 MNIST RESULTS AND DISCUSSION
We selected the first 10 zeros in the MNIST test set, and use the L2 optimization adversary of Carlini & Wagner (2016) to try to perturb those zeros into ones.9 Some sample results are shown in Figure 3. We see that the deterministic models are easily fooled by making small perturbations, but for the VIB models with reasonably large β, the adversary often fails to find an attack (indicated by the green borders) within the permitted number of iterations. Furthermore, when an attack is succesful, it needs to be much larger for the VIB models. To quantify this, Figure 4 plots the magnitude of the perturbation (relative to that of the deterministic and dropout models) needed for a successful attack as a function of β. As β increases, the L0 norm of the perturbation decreases, but both L2 and L∞ norms increase, indicating that the adversary is being forced to put larger modifications into fewer pixels while searching for an adversarial perturbation.
Figure 5 plots the accuracy on FGS adversarial examples of the first 1000 images from the MNIST test set as a function of β. Each point in the plot corresponds to 3 separate executions of three different models trained with the same value of β. All models tested achieve over 98.4% accuracy on the unperturbed MNIST test set, so there is no appreciable measurement distortion due to underlying model accuracy.
Figure 6 plots the accuracy on L2 optimization adversarial examples of the first 1000 images from the MNIST test set as a function of β. The same sets of three models per β were tested three times, as with the FGS adversarial examples.
We generated both untargeted and targeted adversarial examples for Figure 6. For targeting, we generate a random target label different from the source label in order to avoid biasing the results with unevenly explored source/target pairs. We see that for a reasonably broad range of β values, the VIB models have significantly better accuracy on the adversarial examples than the deterministic models, which have an accuracy of 0% (the L2 optimization attack is very effective on traditional model architectures).
Figure 6 also reveals a surprising level of adversarial robustness even when β → 0. This can be explained by the theoretical framework of Fawzi et al. (2016). Their work proves that quadratic classifiers (e.g., xTAx, symmetric A) have a greater capacity for adversarial robustness than linear classifiers. As we show in Appendix C, our Gaussian/softmax encoder/decoder is approximately quadratic for all β <∞.
4.2.5 IMAGENET RESULTS AND DISCUSSION
VIB improved classification accuracy and adversarial robustness for toy datasets like MNIST. We now investigate if VIB offers similar advantages for ImageNet, a more challenging natural image classification. Recall that ImageNet has approximately 1M images spanning 1K classes. We preprocess images such that they are 299x299 pixels.
Architecture
We make use of publicly available, pretrained checkpoints10 of Inception Resnet V2 (Szegedy et al., 2016) on ImageNet (Deng et al., 2009). The checkpoint obtains 80.4% classification accuracy on the ImageNet validation set. Using the checkpoint, we transformed the original training set by applying the pretrained network to each image and extracting the representation at the penultimate layer. This new image representation has 1536 dimensions. The higher layers of the network continue to classify this representation with 80.4% accuracy; conditioned on this extraction the classification
9 We chose this pair of labels since intuitively zeros and ones are the digits that are least similar in terms of human perception, so if the adversary can change a zero into a one without much human-noticeable perturbation, it is unlikely that the model has learned a representation similar to what humans learn.
10Available at the Tensorflow Models repository in the Slim directory: https://github.com/ tensorflow/models/tree/master/slim
model is simply logistic regression. To further speed training, we whitened the 1536 dimensional representation.
Under this transformation, the experiment regime is identical to the permutation invariant MNIST task. We therefore used a similar model architecture. Inputs are passed through two fully connected layers, each with 1024 units. Next, data is fed to a stochastic encoding layer; this layer is characterized by a spherical Gaussian with 1024 learned means and standard deviations. The output of the stochastic layer is fed to the variational classifier–itself a logistic regression, for simplicity. All other hyperparameters and training choices are identical to those used in MNIST, more details in Appendix A.
Classification
We see the same favorable VIB classification performance in ImageNet as in MNIST. By varying β, the estimated mutual information between encoding and image (I(Z,X)) varies as well. At large values of β accuracy suffers, but at intermediate values we obtain improved performance over both a deterministic baseline and a β = 0 regime. In all cases our accuracy is somewhat lower than the original 80.4% accuracy. This may be a consequence of inadequate training time or suboptimal hyperparameters.
Overall the best accuracy we achieved was using β = 0.01. Under this setting we saw an accuracy of 80.12%–nearly the same as the state-of-the-art unmodified network– but with substantially smaller information footprint, only I(X,Z) ∼ 45 bits. This is a surprisingly small amount of information; β = 0 implies over 10,000 bits yet only reaches an accuracy of 78.87%. The deterministic baseline, which was the same network, but without the VIB loss and a 1024 fully connected linear layer instead of the stochastic embedding similarly only achieved 78.75% accuracy. We stress that regressions from the achievable 80.4% are likely due to suboptimal hyperparameters settings or inadequate training.
Considering a continuum of β and a deterministic baseline, the best classification accuracy was achieved with a β = 0.01 ∈ (0, 1). In other words, VIB offered accuracy benefit yet using a mere ∼ 45 bits of information from each image.
Adversarial Robustness
We next show that the VIB-trained network improves resistance to adversarial attack.
We focus on the Carlini targeted L2 attack (see Section 4.2.1). We show results for the VIB-trained network and a deterministic baseline (both on top of precomputed features), as well as for the original pretrained Inception ResNet V2 network itself. The VIB network is more robust to the targeted L2 optimization attack in both magnitude of perturbation and frequency of successful attack.
Figure 7 shows some example images which were all misclassified as “soccer balls” by the deterministic models; by contrast, with the VIB model, only 17 out of 30 of the attacks succeeded in being mislabeled as the target label.11 We find that the VIB model can resist about 43.3% of the attacks, but the deterministic models always fail (i.e., always misclassify into the targeted label).
Figure 8 shows the absolute pixel differences between the perturbed and unperturbed images for the examples in Figure 7. We see that the VIB network requires much larger perturbations in order to fool the classifier, as quantified in Table 2.
5 FUTURE DIRECTIONS
There are many possible directions for future work, including: putting the VIB objective at multiple or every layer of a network; testing on real images; using richer parametric marginal approximations, rather than assuming r(z) = N (0, I); exploring the connections to differential privacy (see e.g., Wang et al. (2016a); Cuff & Yu (2016)); and investigating open universe classification problems (see e.g., Bendale & Boult (2015)). In addition, we would like to explore applications to sequence prediction, where X denotes the past of the sequence and Y the future, while Z is the current representation of the network. This form of the information bottleneck is known as predictive information (Bialek et al., 2001; Palmer et al., 2015).
A HYPERPARAMETERS AND ARCHITECTURE DETAILS FOR EXPERIMENTS
All of the networks for this paper were trained using TensorFlow (Abadi et al., 2016). All weights were initialized using the default TensorFlow Xavier initialization scheme (Glorot & Bengio, 2010) using the averaging fan scaling factor on uniform noise. All biases were initialized to zero. The Adam optimizer (Kingma & Ba, 2015) was used with initial learning rate of 10−4, (β1 = 0.5, β2 = 0.999) and exponential decay, decaying the learning rate by a factor of 0.97 every 2 epochs. The networks were all trained for 200 epochs total. For the MNIST experiments, a batch size of 100 was used, and the full 60,000 training and validation set was used for training, and the 10,000 test images for test results. The input images were scaled to have values between -1 and 1 before fed to the network.
All runs maintained an exponential weighted average of the parameters during the training run; these averaged parameters were used at test time. This is in the style of Polyak averaging Polyak & Juditsky (1992), with a decay constant of 0.999. Our estimate of mutual informations were measured in bits. For the VIB experiments in all sections, no other form of regularization was used.
For the 256 dimensional gaussian embeddings of Section 4.1.1, a linear layer of size 512 was used to create the 256 mean values and standard deviations for the embedding. The standard deviations were made to be positive by a softplus transformation with a bias of -5.0 to have them initially be small.
σ = log (1 + exp(x− 5.0)) (19)
For the 1024 dimensional Imagenet embeddings of Section 4.2.5, a sigma bias of 0.57 was used to keep the initial standard deviations near 1 originally, and a batch size of 200 was used.
For the 2 dimensional gaussian embeddings of Section 4.1.2, a linear layer was used with 2+4 = 6 outputs, the first two of which were used for the means, and the other 4 were reshaped to a 2 × 2 matrix, the center was transformed according to a softplus with a bias of -5.0, and the off diagonal components were multiplied by 10−2, while the upper triangular element was dropped to form the Cholesky decomposition of the covariance matrix.
B CONNECTION TO VARIATIONAL AUTOENCODERS
We can also consider unsupervised versions of the information bottleneck objective. Consider the objective: max I(Z,X)− βI(Z, i), (20) similar to the information theoretic objective for clustering introduced in Slonim et al. (2005).
Here the aim is to take our dataX and maximize the mutual information contained in some encoding Z, while restricting how much information we allow our representation to contain about the identity of each data element in our sample (i). We will form a bound much like we did in the main text. For the first term, we form a variational decoder q(x|z) and take a bound:
I(Z,X) = ∫ dx dz p(x, z) log
p(x|z) p(x)
(21)
= H(x) + ∫ dz p(x) ∫ dx p(x|z) log p(x|z) (22)
≥ ∫ dz p(x) ∫ dx p(x|z) log q(x|z) (23)
= ∫ dx p(x) ∫ dz p(x|z) log q(x|z). (24)
Here we have dropped the entropy in our data H(X) because it is out of our control and we have used the nonnegativity of the Kullbach-Leibler divergence to replace our intractable p(x|z) with a variational decoder q(x|z).
Turning our attention to the second term, note that: p(z|i) = ∫ dx p(z|x)p(x|i) = ∫ dx p(z|x)δ(x− xi) = p(z|xi), (25)
and that we will take p(i) = 1N .
So that we can bound our second term from above I(Z, i) = ∑ i ∫ dz p(z|i)p(i) log p(z|i) p(z) (26)
= 1
N ∑ i ∫ dz p(z|xi) log p(z|xi) p(z)
(27)
≤ 1 N ∑ i ∫ dz p(z|xi) log p(z|xi) r(z) , (28)
Where we have replaced the intractable marginal p(z) with a variational marginal r(z).
Putting these two bounds together we have that our unsupervised information bottleneck objective takes the form
I(Z,X)− βI(Z, i) ≤ ∫ dx p(x) ∫ dz p(z|x) log q(x|z)− β 1
N ∑ i KL[p(Z|xi), r(Z)]. (29)
And this takes the form of a variational autoencoder (Kingma & Welling, 2014), except with the second KL divergence term having an arbitrary weight β.
It is interesting that while this objective takes the same mathematical form as that of a Variational Autoencoder, the interpretation of the objective is very different. In the VAE, the model starts life as a generative model with a defined prior p(z) and stochastic decoder p(x|z) as part of the model, and the encoder q(z|x) is created to serve as a variational approximation to the true posterior p(z|x) = p(x|z)p(z)/p(x). In the VIB approach, the model is originally just the stochastic encoder p(z|x), and the decoder q(x|z) is the variational approximation to the true p(x|z) = p(z|x)p(x)/p(z) and r(z) is the variational approximation to the marginal p(z) = ∫ dx p(x)p(z|x). This difference in interpretation makes natural suggestions for novel directions for improvement.
This precise setup, albeit with a different motivation was recently explored in Higgins et al. (2016), where they demonstrated that by changing the weight of the variational autoencoders regularization term, there were able to achieve latent representations that were more capable when it came ot zeroshot learning and understanding ”objectness”. In that work, they motivated their choice to change the relative weightings of the terms in the objective by appealing to notions in neuroscience. Here we demonstrate that appealing to the information bottleneck objective gives a principled motivation and could open the door to better understanding the optimal choice of β and more tools for accessing the importance and tradeoff of both terms.
Beyond the connection to existing variational autoencoder techniques, we note that the unsupervised information bottleneck objective suggests new directions to explore, including targetting the exact marginal p(z) in the regularization term, as well as the opportunity to explore tighter bounds on the first I(Z,X) term that may not require explicit variational reconstruction.
C QUADRATIC BOUNDS FOR STOCHASTIC LOGISTIC REGRESSION DECODER
Consider the special case when the bottleneck Z is a multivariate Normal, i.e., z|x ∼ N(µx,Σx) where Σx is a K × K positive definite matrix. The parameters µx,Σx can be constructed from a deep neural network, e.g.,
µx = γ1:K(x)
chol(Σx) = diag(log(1 + exp(γK+1:2K))) + subtril(γ2K+1:K(K+3)/2),
where γ(x) ∈ RK(K+3)/2 is the network output of input x.
Suppose that the prediction is a categorical distribution computed as S(Wz) where W is a C × K weight matrix and logS(x) = x − lse(x) is the log-soft-max function with lse(x) = log ∑K k=1 exp(xk) being the log-sum-exp function.
This setup (which is identical to our experiments) induces a classifier which is bounded by a quadratic function, which is interesting because the theoretical framework Fawzi et al. (2016) proves that quadratic classifiers have greater capacity for adversarial robustness than linear functions.
We now derive an approximate bound using second order Taylor series expansion (TSE). The bound can be made proper via Browne & McNicholas (2015). However, using the TSE is sufficient to sketch the derivation.
Jensen’s inequality implies that the negative log-likelihood soft-max is upper bounded by:
− log E [S(WZ)|µx,Σx] ≤ −E [logS(WZ)|µx,Σx] = −Wµx + E [lse(WZ)|µx,Σx] = −Wµx + E [lse(Z)|Wµx,WΣx] .
The second order Taylor series expansion (TSE) of lse is given by, lse(x+ δ) ≈ lse(x) + δT S(x) + 12δT [ diag(S(x))− S(x)S(x)T ] δ.
Taking the expectation of the TSE at the mean yields,
EN(0,WΣxW T )[ lse(Wµx + δ)] ≈ lse(Wµx) + EN(0,WΣxW T )[δT ]S(Wµx)+ + 12 EN(0,WΣxW T )[δ T [ diag(S(Wµx))− S(Wµx)S(Wµx)T ] δ]
= lse(Wµx) + 1 2 tr(WΣxW T [ diag(S(Wµx))− S(Wµx)S(Wµx)T ] )
= lse(Wµx) + 1 2 tr(WΣxW T diag(S(Wµx)))− 12S(Wµx)T WΣxW T S(Wµx)
= lse(Wµx) + 1 2 √ S(Wµx) T WΣxW T √ S(Wµx)− 12S(Wµx)T WΣxW T S(Wµx)
The second-moment was calculated by noting,
E[XT BX] = E tr(XXT B) = tr(E[XXT ]B) = tr(ΣB).
Putting this altogether, we conclude, E [S(WZ)|µx,Σx] ' S(Wµx) exp ( − 12 √ S(Wµx) T WΣxW T √ S(Wµx) + 12S(Wµx)T WΣxW T S(Wµx) ) .
As indicated, rather than approximate the lse via TSE, we can make a sharp, quadratic upper bound via Browne & McNicholas (2015). However this merely changes the S(Wµx) scaling in the exponential; the result is still log-quadratic.

Thank you for an interesting read. I personally like the information bottleneck principle and am very happy to see its application to deep neural networks. To my knowledge, this is the first paper that applies IB to train deep networks (the original papers only presented the concept), but see below for the note of independent work claim.
The derivation of the variational lowerbound is very clear, even for those who are not very familiar with variational inference. Also the explanation of the IB principle is clear. Experimental results seem to be very promising.
I found the presentation for the model a bit confusing. In variational inference/information maximisation, p usually denotes the model and q represents the "inference engine". This means the choice of inference method is independent to the modelling procedure. However the presented VIB assumed p(x, y) as the **underlying data distribution** (and approximated by the empirical distribution), thus here the model is actually q(y|z)p(z|x). Then the authors presented p(y|x) as the **predictive distribution** in page 8, paragraph 2 of section 4.2.3. Predictive in what sense? I guess you meant p(y|x) = \int q(y|z) p(z|x) dz in this case, but this makes the two definitions contradict to each other!
The authors have made an interesting connection to variational auto-encoder and the warm-up training (by tuning beta). However, even when the loss function formula is the same to the variational lowerbound used in VAE (in this case beta = 1), the underlying model is different! For example, r(z) in VIB is the variational approximation to p(z) (which means r(z) is not a component in the model), while in VAE it is the prior distribution which is actually defined in the modelling procedure. Similaly p(z|x) in VIB is included in the model, while in VAE that is the approximate posterior and can be independently chosen (e.g. you can use p(x|z) as a deep NN but p(z|x) as a deep NN or a Gaussian process).
In summary, I think the presentation for the modelling procedure is unclear. I hope these point would be made clearer in revision since the current presentation makes me uncomfortable as a Bayesian person. In the VAE part, it's better to clearly mention the difference between VIB and VAE, and provide some intuitions if the VIB interpretation is preferred.
Typos:
Eq. 9-11: did you mean q(y|z) instead of q(z|y)?
Fig 2 "as beta becomes smaller": did you mean "larger"?
**claim for independent work**
The authors claimed that the manuscript presented an independent work to Chalk et al. 2016 which is online since May 2016. It seems to me that nowadays deep learning research is very competitve that many people publish the same idea at the same time. So I would trust this claim and commend the authors' honesty, but in case this is not true, I would not recommend the manuscript for acceptance.