On the Information Bottleneck Theory of Deep Learning (Saxe et al.) — part 1/6

Cleaned-text code points [0, 14848). Bibliography omitted here; source.md includes it.

Published as a conference paper at ICLR 2018
ON THE INFORMATION BOTTLENECK
THEORY OF DEEP LEARNING
Andrew M. Saxe, Yamini Bansal, Joel Dapello, Madhu Advani
Harvard University
{asaxe,madvani}@fas.harvard.edu,{ybansal,dapello}@g.harvard.edu
Artemy Kolchinsky, Brendan D. Tracey
Santa Fe Institute
{artemyk,tracey.brendan}@gmail.com
David D. Cox
Harvard University
MIT-IBM Watson AI Lab
davidcox@fas.harvard.edu
david.d.cox@ibm.com
ABSTRACT
The practical successes of deep neural networks have not been matched by theoret-
ical progress that satisfyingly explains their behavior. In this work, we study the
information bottleneck (IB) theory of deep learning, which makes three speciﬁc
claims: ﬁrst, that deep networks undergo two distinct phases consisting of an
initial ﬁtting phase and a subsequent compression phase; second, that the compres-
sion phase is causally related to the excellent generalization performance of deep
networks; and third, that the compression phase occurs due to the diffusion-like
behavior of stochastic gradient descent. Here we show that none of these claims
hold true in the general case. Through a combination of analytical results and
simulation, we demonstrate that the information plane trajectory is predominantly
a function of the neural nonlinearity employed: double-sided saturating nonlineari-
ties like tanh yield a compression phase as neural activations enter the saturation
regime, but linear activation functions and single-sided saturating nonlinearities
like the widely used ReLU in fact do not. Moreover, we ﬁnd that there is no evident
causal connection between compression and generalization: networks that do not
compress are still capable of generalization, and vice versa. Next, we show that
the compression phase, when it exists, does not arise from stochasticity in training
by demonstrating that we can replicate the IB ﬁndings using full batch gradient
descent rather than stochastic gradient descent. Finally, we show that when an
input domain consists of a subset of task-relevant and task-irrelevant information,
hidden representations do compress the task-irrelevant information, although the
overall information about the input may monotonically increase with training time,
and that this compression happens concurrently with the ﬁtting process rather than
during a subsequent compression period.
1 I NTRODUCTION
Deep neural networks (Schmidhuber, 2015; LeCun et al., 2015) are the tool of choice for real-world
tasks ranging from visual object recognition (Krizhevsky et al., 2012), to unsupervised learning
(Goodfellow et al., 2014; Lotter et al., 2016) and reinforcement learning (Silver et al., 2016). These
practical successes have spawned many attempts to explain the performance of deep learning systems
(Kadmon & Sompolinsky, 2016), mostly in terms of the properties and dynamics of the optimization
problem in the space of weights (Saxe et al., 2014; Choromanska et al., 2015; Advani & Saxe, 2017),
or the classes of functions that can be efﬁciently represented by deep networks (Montufar et al.,
2014; Poggio et al., 2017). This paper analyzes a recent inventive proposal to study the dynamics of
learning through the lens of information theory (Tishby & Zaslavsky, 2015; Shwartz-Ziv & Tishby,
2017). In this view, deep learning is a question of representation learning: each layer of a deep neural
network can be seen as a set of summary statistics which contain some but not all of the information
present in the input, while retaining as much information about the target output as possible. The
1

Published as a conference paper at ICLR 2018
amount of information in a hidden layer regarding the input and output can then be measured over the
course of learning, yielding a picture of the optimization process in the information plane. Crucially,
this method holds the promise to serve as a general analysis that can be used to compare different
architectures, using the common currency of mutual information. Moreover, the elegant information
bottleneck (IB) theory provides a fundamental bound on the amount of input compression and target
output information that any representation can achieve (Tishby et al., 1999). The IB bound thus
serves as a method-agnostic ideal to which different architectures and algorithms may be compared.
A preliminary empirical exploration of these ideas in deep neural networks has yielded striking
ﬁndings (Shwartz-Ziv & Tishby, 2017). Most saliently, trajectories in the information plane appear to
consist of two distinct phases: an initial “ﬁtting” phase where mutual information between the hidden
layers and both the input and output increases, and a subsequent “compression” phase where mutual
information between the hidden layers and the input decreases. It has been hypothesized that this
compression phase is responsible for the excellent generalization performance of deep networks, and
further, that this compression phase occurs due to the random diffusion-like behavior of stochastic
gradient descent.
Here we study these phenomena using a combination of analytical methods and simulation. In Section
2, we show that the compression observed by Shwartz-Ziv & Tishby (2017) arises primarily due to
the double-saturating tanh activation function used. Using simple models, we elucidate the effect of
neural nonlinearity on the compression phase. Importantly, we demonstrate that the ReLU activation
function, often the nonlinearity of choice in practice, does not exhibit a compression phase. We
discuss how this compression via nonlinearity is related to the assumption of binning or noise in the
hidden layer representation. To better understand the dynamics of learning in the information plane,
in Section 3 we study deep linear networks in a tractable setting where the mutual information can be
calculated exactly. We ﬁnd that deep linear networks do not compress over the course of training for
the setting we examine. Further, we show a dissociation between generalization and compression.
In Section 4, we investigate whether stochasticity in the training process causes compression in the
information plane. We train networks with full batch gradient descent, and compare the results to those
obtained with stochastic gradient descent. We ﬁnd comparable compression in both cases, indicating
that the stochasticity of SGD is not a primary factor in the observed compression phase. Moreover,
we show that the two phases of SGD occur even in networks that do not compress, demonstrating
that the phases are not causally related to compression. These results may seem difﬁcult to reconcile
with the intuition that compression can be necessary to attain good performance: if some input
channels primarily convey noise, good generalization requires excluding them. Therefore, in Section
5 we study a situation with explicitly task-relevant and task-irrelevant input dimensions. We show
that the hidden-layer mutual information with the task-irrelevant subspace does indeed drop during
training, though the overall information with the input increases. However, instead of a secondary
compression phase, this task-irrelevant information is compressed at the same time that the task-
relevant information is boosted. Our results highlight the importance of noise assumptions in applying
information theoretic analyses to deep learning systems, and put in doubt the generality of the IB
theory of deep learning as an explanation of generalization performance in deep architectures.
2 C OMPRESSION AND NEURAL NONLINEARITIES
The starting point for our analysis is the observation that changing the activation function can
markedly change the trajectory of a network in the information plane. In Figure 1A, we show
our replication of the result reported by Shwartz-Ziv & Tishby (2017) for networks with the tanh
nonlinearity.1 This replication was performed with the code supplied by the authors of Shwartz-Ziv
& Tishby (2017), and closely follows the experimental setup described therein. Brieﬂy, a neural
network with 7 fully connected hidden layers of width 12-10-7-5-4-3-2 is trained with stochastic
gradient descent to produce a binary classiﬁcation from a 12-dimensional input. In our replication we
used 256 randomly selected samples per batch. The mutual information of the network layers with
respect to the input and output variables is calculated by binning the neuron’stanh output activations
into 30 equal intervals between -1 and 1. Discretized values for each neuron in each layer are then
used to directly calculate the joint distributions, over the 4096 equally likely input patterns and true
output labels. In line with prior work (Shwartz-Ziv & Tishby, 2017), the dynamics in Fig. 1 show a
1Code for our results is available at https://github.com/artemyk/ibsgd/tree/iclr2018
2

Published as a conference paper at ICLR 2018
Figure 1: Information plane dynamics and neural nonlinearities. (A) Replication of Shwartz-Ziv &
Tishby (2017) for a network with tanh nonlinearities (except for the ﬁnal classiﬁcation layer which
contains two sigmoidal neurons). The x-axis plots information between each layer and the input,
while the y-axis plots information between each layer and the output. The color scale indicates
training time in epochs. Each of the six layers produces a curve in the information plane with the
input layer at far right, output layer at the far left. Different layers at the same epoch are connected
by ﬁne lines. (B) Information plane dynamics with ReLU nonlinearities (except for the ﬁnal layer
of 2 sigmoidal neurons). Here no compression phase is visible in the ReLU layers. For learning
curves of both networks, see Appendix A. (C) Information plane dynamics for a tanh network of
size 784− 1024− 20− 20− 20− 10 trained on MNIST, estimated using the non-parametric kernel
density mutual information estimator of Kolchinsky & Tracey (2017); Kolchinsky et al. (2017),
no compression is observed except in the ﬁnal classiﬁcation layer with sigmoidal neurons. See
Appendix B for the KDE MI method applied to the original Tishby dataset; additional results using
a second popular nonparametric k-NN-based method (Kraskov et al., 2004); and results for other
neural nonlinearities.
3

Published as a conference paper at ICLR 2018
transition between an initial ﬁtting phase, during which information about the input increases, and a
subsequent compression phase, during which information about the input decreases.
We then modiﬁed the code to train deep networks using rectiﬁed linear activation functions (f(x) =
max(0,x )). While the activities of tanh networks are bounded in the range [−1, 1], ReLU networks
have potentially unbounded positive activities. To calculate mutual information, we ﬁrst trained the
ReLU networks, next identiﬁed their largest activity value over the course of training, and ﬁnally
chose 100 evenly spaced bins between the minimum and maximum activity values to discretize
the hidden layer activity. The resulting information plane dynamics are shown in Fig. 1B. The
mutual information with the input monotonically increases in all ReLU layers, with no apparent
compression phase. To see whether our results were an artifact of the small network size, toy dataset,
or simple binning-based mutual information estimator we employed, we also trained larger networks
on the MNIST dataset and computed mutual information using a state-of-the-art nonparametric
kernel density estimator which assumes hidden activity is distributed as a mixture of Gaussians
(see Appendix B for details). Fig. C-D show that, again, tanh networks compressed but ReLU
networks did not. Appendix B shows that similar results also obtain with the popular nonparametric
k-nearest-neighbor estimator of Kraskov et al. (2004), and for other neural nonlinearities. Thus, the
choice of nonlinearity substantively affects the dynamics in the information plane.
To understand the impact of neural nonlinearity on the mutual information dynamics, we develop
a minimal model that exhibits this phenomenon. In particular, consider the simple three neuron
network shown in Fig. 2A. We assume a scalar Gaussian input distributionX∼N (0, 1), which is
fed through the scalar ﬁrst layer weightw1, and passed through a neural nonlinearityf(·), yielding
the hidden unit activityh =f(w1X). To calculate the mutual information with the input, this hidden
unit activity is then binned yielding the new discrete variableT = bin(h) (for instance, into 30 evenly
spaced bins from -1 to 1 for the tanh nonlinearity). This binning process is depicted in Fig. 2B. In
this simple setting, the mutual informationI(T ;X) between the binned hidden layer activityT and
the inputX can be calculated exactly. In particular,
I(T ;X) = H(T )−H(T|X) (1)
= H(T ) (2)
= −
N∑
i=1
pi logpi (3)
whereH(·) denotes entropy, and we have used the fact thatH(T|X) = 0 sinceT is a deterministic
function ofX. Here the probabilitiespi =P (h≥bi andh<b i+1) are simply the probability that
an inputX produces a hidden unit activity that lands in bini, deﬁned by lower and upper bin limits
bi andbi+1 respectively. This probability can be calculated exactly for monotonic nonlinearitiesf(·)
using the cumulative density ofX,
pi =P (X≥f −1(bi)/w1 andX <f −1(bi+1)/w1), (4)
wheref −1(·) is the inverse function off(·).
As shown in Fig. 2C-D, as a function of the weight w1, mutual information with the input ﬁrst
increases and then decreases for thetanh nonlinearity, but always increases for the ReLU nonlinearity.
Intuitively, for small weightsw1≈ 0, neural activities lie near zero on the approximately linear part
of the tanh function. Therefore f(w1X)≈ w1X, yielding a rescaled Gaussian with information
that grows with the size of the weights. However for very large weightsw1→∞ , the tanh hidden
unit nearly always saturates, yielding a discrete variable that concentrates in just two bins. This
is more or less a coin ﬂip, containing mutual information with the input of approximately 1 bit.
Hence the distribution ofT collapses to a much lower entropy distribution, yielding compression for
large weight values. With the ReLU nonlinearity, half of the inputs are negative and land in the bin
containing a hidden activity of zero. The other half are Gaussian distributed, and thus have entropy
that increases with the size of the weight.
Hence double-saturating nonlinearities can lead to compression of information about the input, as
hidden units enter their saturation regime, due to the binning procedure used to calculate mutual
information. The crux of the issue is that the actual I(h;X) is inﬁnite, unless the network itself
adds noise to the hidden layers. In particular, without added noise, the transformation fromX to the
continuous hidden activityh is deterministic and the mutual informationI(h;X) would generally be
4