Information Bottleneck: Exact Analysis of (Quantized) Neural Networks

1 INTRODUCTION . Improving our theoretical understanding of why over-parameterized deep neural networks generalize well is arguably one of main problems in current machine learning research ( Poggio et al. , 2020 ) . Tishby & Zaslavsky ( 2015 ) suggested to analyze deep neural networks based on their Information Bottleneck ( IB ) concept , which is built on measurements of mutual information ( MI ) between the activations of hidden layers and the input and target ( Tishby et al. , 1999 ) , for an overview see Geiger ( 2020 ) . Shwartz-Ziv & Tishby ( 2017 ) empirically studied the IB principle applied to neural networks and made several qualitative observations about the training process ; especially , they observed a fitting phase and a compression phase . The latter information-theoretic compression is conjectured to be a reason for good generalization performance and has widely been considered in the literature ( Abrol & Tanner , 2020 ; Balda et al. , 2018 ; 2019 ; Chelombiev et al. , 2019 ; Cheng et al. , 2019 ; Darlow & Storkey , 2020 ; Elad et al. , 2019 ; Fang et al. , 2018 ; Gabrié et al. , 2019 ; Goldfeld et al. , 2019 ; Jónsson et al. , 2020 ; Kirsch et al. , 2020 ; Tang Nguyen & Choi , 2019 ; Noshad et al. , 2019 ; Schiemer & Ye , 2020 ; Shwartz-Ziv & Alemi , 2020 ; Wickstrøm et al. , 2019 ; Yu et al. , 2020 ) . The work and conclusions by Shwartz-Ziv & Tishby ( 2017 ) received a lot of critique , with the generality of their claims being doubted ; especially Saxe et al . ( 2018 ) argued that the results by Shwartz-Ziv & Tishby do not generalize to networks using a different activation function . Their critique was again refuted by the original authors with counter-claims about incorrect estimation of the MI , highlighting an issue with the approximation of MI in both studies . Our goal is to verify the claims by Shwartz-Ziv & Tishby and the critique by Saxe et al . in a setting where the MI can be computed exactly . These studies consider neural networks as theoretical entities working with infinite precision , which makes computation of the information theoretic quantities problematic ( for a detailed discussion we refer to Geiger , 2020 , see also Section 3 ) . Assuming continuous input distributions , a deterministic network using any of the standard activation functions ( e.g. , RELU , TANH ) can be shown to have infinite MI ( Amjad & Geiger , 2019 ) . If an empirical input distribution defined by a data set D is considered ( as it is the case in many of the previous studies ) , then randomly-initialized deterministic neural networks with invertible activation functions will most likely result in trivial measurements of MI in the sense that the MI is finite but always maximal , that is , equal to log |D| ( Goldfeld et al. , 2019 ; Amjad & Geiger , 2019 ) . In order to obtain non-trivial measurements of MI , real-valued activations are usually discretized by binning the values , throwing away information in the process . The resulting estimated MI can be shown to be highly dependent on this binning , we refer to Geiger ( 2020 ) for a detailed discussion . Instead of approximating the MI in this fashion , we take advantage of the fact that modern computers – and thus neural networks – are discrete in the sense that a floating point value can typically take at most 232 different values . Because 32-bit precision networks may still be too precise to observe compression ( i.e. , information loss ) , we apply quantization to the neural network system to an extent that we can compute informative quantities , that is , we amplify the effect of the information loss due to the discrete computations in the neural network . One may argue that we just moved the place where the discretization is applied . This is true , but leads to a fundamental difference : previous studies applying the discretization post-hoc rely on the in general false assumption that the binned MI approximates the continuous MI well – and thus introduce measurement errors , which may occlude certain phenomena and/or lead to artifactual observations . In contrast , our computations reflect the true information flow in a network during training . Our study confirms that estimation of MI by binning may lead to strong artifacts in IB analyses and shows that : • Both fitting and compression phases occur in the output SOFTMAX layer . • For the hidden layers , the fitting phase occurs for both TANH and RELU activations . • When using TANH in the hidden layers , compression is only observed in the last hidden layer . • When using RELU , we did not observe compression in the hidden layers . • Even when applying low precision quantization , more complex networks with many neurons in each layer are observed to be too expressive to exhibit compression , as no information is lost . • Our setting excludes that the MI approximation is the reason for these different IB dynamics . The next section introduces the IB concept with a focus on its application to neural networks including the critique and controversy as well as related work . Section 3 discusses issues relating to the estimation of MI , and the idea behind our contribution . Section 4 presents our experiments , results and discussion before we conclude in Section 5 . 2 THE INFORMATION BOTTLENECK . Preliminaries . Given a continuous random variable ( r.v . ) X with density function p ( x ) and support X , the continuous entropy H ( X ) of X is a measure of the uncertainty associated with X and is given by H ( X ) = − ∫ X p ( x ) log p ( x ) dx . Given two r.v.s X and Y with density functions p ( x ) and q ( y ) and supports X and Y , the mutual information I ( X ; Y ) of X and Y is a measure of the mutual “ knowledge ” between the two variables . The symmetric I ( X ; Y ) is given by I ( X ; Y ) = ∫ Y ∫ X p ( x , y ) log p ( x , y ) p ( x ) p ( y ) dxdy . In many cases it is impossible to compute the continuous entropy and MI for continuous r.v.s exactly , due to limited samples or computational limits , or because it may not be finite ( Geiger , 2020 ) . Instead , we often estimate the quantities by their discrete counterparts . When X is a discrete r.v. , we consider the Shannon entropy H ( X ) = − ∑ P ( x ) logP ( x ) . Correspondingly , the mutual information I ( X ; Y ) of two discrete r.v.s X , Y , is given by I ( X ; Y ) = ∑ x , y P ( x , y ) log P ( x , y ) P ( x ) P ( y ) . We have the following useful identity for both the continuous and discrete MI : I ( X ; Y ) = H ( X ) − H ( X|Y ) , ( 1 ) where H ( X|Y ) is the conditional entropy of X given Y . IB Definition . The IB method was proposed by Tishby et al . ( 1999 ) . It is an information theoretic framework for extracting relevant components of an input r.v . X with respect to an output r.v . Y . These relevant components are found by “ squeezing ” the information from X through a bottleneck , in the form of an r.v . T . In other words , T is a compression of X . The idea generalizes rate distortion theory , in which we wish to compress X , obtaining T , such that I ( X ; T ) is maximized subject to a constraint on the expected distortion d ( x , t ) wrt . the joint distribution p ( x , t ) ( Tishby et al. , 1999 ) . In the IB framework , the distortion measure d is replaced by the negative loss in MI between T and the output Y , I ( T ; Y ) . Both IB and rate distortion are lossy compression schemes.1 The data processing inequality ( DPI ) I ( Y ; X ) ≥ I ( Y ; T ) holds for the IB ; that is , the bottleneck r.v . can not contain more information about the label than the input . One drawback of the information bottleneck method is the dependence on the joint distribution , p ( x , y ) , which is generally not known . Shamir et al . ( 2010 ) addressed this issue and showed that the MI , the main ingredient in the method , can be estimated reliably with fewer samples than required for estimating the true joint distribution . As common in the IB literature , whenever we discuss the MI computed on a finite data set D , we assume that p ( x , y ) corresponds to the empirical distribution defined by D , which is true for the experiments in Section 4.1 . In practice , the assumption has to be relaxed to the data being drawn i.i.d . from p ( x , y ) . However , any uncertainty resulting from the finite sample estimation in the latter case is not considered in our discussions . IB In Deep Learning . Tishby & Zaslavsky ( 2015 ) applied the IB concept to neural networks . They view the layers of a deep neural network ( DNN ) as consecutive compressions of the input . They consider the Markov chain Y → X → T1 → T2 → ... → TL = Ŷ , where Ti denotes the i ’ th hidden layer of the L-layer network and TL = Ŷ denotes the output of the network . Again , the bottleneck must satisfy the DPI : I ( Y ; X ) ≥ I ( Y ; T1 ) ≥ I ( Y ; T2 ) ≥ ... ≥ I ( Y ; Ŷ ) , ( 2 ) I ( X ; X ) ≥ I ( X ; T1 ) ≥ I ( X ; T2 ) ≥ ... ≥ I ( X ; Ŷ ) . ( 3 ) Estimating the MI of continuous variables is difficult ( Alemi et al. , 2017 ) , as evident from the many different methods proposed ( Kraskov et al. , 2004 ; Kolchinsky & Tracey , 2017 ; Noshad et al. , 2019 ) . In the discrete case , I ( X ; T ) and I ( T ; Y ) can be computed as I ( X ; T ) = H ( T ) − H ( T |X ) = H ( T ) , ( 4 ) I ( T ; Y ) = I ( Y ; T ) = H ( T ) − H ( T |Y ) , ( 5 ) following from ( 1 ) and using in ( 4 ) the assumption that T is a deterministic function of X . However , for deterministic neural networks the continuous entropies may not be finite ( Goldfeld et al. , 2019 ; Saxe et al. , 2018 ; Amjad & Geiger , 2019 ) . Shwartz-Ziv & Tishby ( 2017 ) estimate the MI via ( 4 ) and ( 5 ) by discretizing T and then computing the discrete entropy . They trained a network ( shown in Figure 1a ) on a balanced synthetic data set consisting of 12-bit binary inputs and binary labels . The network was trained for a fixed number of epochs while training and test error were observed . For every epoch and every layer T , The discretization is done by binning of T : Given upper and lower bounds bu , bl , and m ∈ N , we let B : R→ [ m ] denote the binning operation , that maps x ∈ [ bl , bu ] to the index of the corresponding bin from the set of m uniformly distributed bins in [ bl , bu ] . Overloading the notation , we apply B directly to a vector in Rd in order to obtain the resulting vector in [ m ] d of bin indices . Using discretized T ′ = B ( T ) , I ( X ; T ′ ) and I ( T ′ ; Y ) are then computed directly by ( 4 ) and ( 5 ) , using estimates of P ( T ′ ) , P ( Y ) and P ( T ′|Y ) over all samples D of X . Shwartz-Ziv & Tishby used the TANH activation function for the hidden layers , with bl = −1 , bu = 1 ( bl = 0 for the output SOFTMAX layer ) and m = 30 bins . The estimated I ( X ; T ) and I ( T ; Y ) are plotted in the information plane , providing a visual representation of the information flow in 1For IB , finding the optimal representation T can be formulated as the minimization of the Lagrangian I ( X ; T ) − βI ( T ; Y ) subject to the Markov chain Y → X → T and β ∈ R+ ( Tishby & Zaslavsky , 2015 ) . the network during training ( see example in Figure 1b ) . Based on the obtained results , ShwartzZiv & Tishby ( 2017 ) make several observations ; one notable observation is the occurrence of two phases : an empirical risk minimization phase and a compression phase . The first phase , also referred to as fitting phase , is characterized by increasing I ( T ; Y ) related to a decreasing loss . The subsequent compression phase is characterized by decreasing I ( X ; T ) , and it has been argued that this compression leads to better generalization . Critique and Controversy . The work by Tishby & Zaslavsky ( 2015 ) and Shwartz-Ziv & Tishby ( 2017 ) has jump-started an increasing interest in the IB method for deep learning , with several papers investigating or extending their contributions , see the review by Geiger ( 2020 ) . However , as mentioned , their work also received criticism . In particular , the compression phase as a general phenomenon has been called into question . Saxe et al . ( 2018 ) published a paper refuting several of the claims made by Shwartz-Ziv & Tishby ( 2017 ) . They based their criticism on a replication of the experiment done by Shwartz-Ziv & Tishby ( 2017 ) where they replaced the bounded TANH activation function with the unbounded RELU activation function . When discretizing , they used the maximum activation observed across all epochs for the upper binning bound bu . The article claims that the two phases observed by Shwartz-Ziv & Tishby occurred because the activations computed using TANH saturate close to the boundaries −1 and 1 . The claim is supported by experiments using RELU activations and m = 100 bins , in which the two phases are not observed . The critique paper by Saxe et al . was published at ICLR 2018 , but started a discussion already during the review process the previous year , when Shwartz-Ziv & Tishby defended their paper in the online discussion forum OpenReview.net2 , posting a response titled “ Data falsifying the claims of this ICLR submission all together ” ( Saxe et al. , 2017 ) . The response specifically states that “ The authors don ’ t know how to estimate mutual information correctly ” , referring to Saxe et al . ( 2018 ) , and goes on to provide an example with a network using RELU activations , which does indeed exhibit the two phases . In response , Saxe et al . performed further experiments using different estimators for MI : a state-of-the-art non-parametric KDE approach ( Kolchinsky & Tracey , 2017 ) and a k-NN based estimator ( Kraskov et al. , 2004 ) . The authors still did not observe the two phases claimed by Shwartz-Ziv & Tishby . Following the discussion on OpenReview.net , several other papers have also commented on the controversy surrounding the information bottleneck . Noshad et al . ( 2019 ) presented a new MI estimator EDGE , based on dependency graphs , and tested it on the specific counter example using 2https : //openreview.net/ RELU activations as suggested by Saxe et al . ( 2018 ) , and they observed the two phases . Table I in the review by Geiger ( 2020 ) provides a nice overview over empirical IB studies and if the compression phase was observed ( Darlow & Storkey , 2020 ; Jónsson et al. , 2020 ; Kirsch et al. , 2020 ; Noshad et al. , 2019 ; Raj et al. , 2020 ; Shwartz-Ziv & Tishby , 2017 ) or not ( Abrol & Tanner , 2020 ; Balda et al. , 2018 ; 2019 ; Tang Nguyen & Choi , 2019 ; Shwartz-Ziv & Alemi , 2020 ; Yu et al. , 2020 ) or the results were mixed ( Chelombiev et al. , 2019 ; Cheng et al. , 2019 ; Elad et al. , 2019 ; Fang et al. , 2018 ; Gabrié et al. , 2019 ; Goldfeld et al. , 2019 ; Saxe et al. , 2018 ; Schiemer & Ye , 2020 ; Wickstrøm et al. , 2019 ) . In conclusion , an important part of the controversy surrounding the IB hinges on the estimation of the information-theoretic quantities – this issue has to be solved before researching the information flow . Related Work . The effect of estimating MI by binning has been investigated before , we again refer to Geiger ( 2020 ) for a good overview and discussion . Shwartz-Ziv & Alemi ( 2020 ) consider infinite ensembles of infinitely-wide networks , which renders MI computation feasible , but do not observe a compression phase . Chelombiev et al . ( 2019 ) applies adaptive binning , which , while less prone to issues caused by having the “ wrong ” number of bins , is still an estimation and thus also suffers from the same problems . Goldfeld et al . ( 2019 ) explores IB analysis by use of stochastic neural networks which allow them to show that the compression phase in these noisy networks occurs due to clustering of the hidden representations . While theoretically interesting , the stochastic neural networks are still qualitatively different from deterministic ones and thus not directly applicable in practice . Raj et al . ( 2020 ) conducted an IB analysis of binary networks , where both the activations and weights can only be ±1 , which allows for exact computation of MI . The binary networks are significantly different from the networks used in the original studies , whereas applying the IB analysis to quantized versions of networks from these studies allows for a more direct comparison . At a first glance , the work by Raj et al . ( 2020 ) could be viewed as taking our approach to the extreme . However , the computations and training dynamics of the binary networks are qualitatively different from our study and the original IB work . For example , the binary networks require modifications for gradient estimation ( e.g. , Raj et al . consider straight-through-estimator , approximate sign , and swish sign ) .

Summary: **Update after authors' response:** The authors have managed to clarify some issues and make a couple of small improvements to the manuscript. I am tempted to raise my score to a 7, but to me personally the paper does not quite pass the threshold for an 8 (which is the next possible rating on the conference scale). I am in favor of accepting the paper and will argue so during the reviewers' discussion.  **Summary**  The paper investigates the information-plane analysis of deep neural network training dynamics. The original claim was that neural networks go through distinct fitting and compression phases in SGD training that are well characterized by the information plane. These findings have later been disputed (particularly due to the way mutual information is estimated) which has lead to tens of papers investigating the phenomenon with mixed conclusions. The main idea of this paper is to eliminate the estimation problem by training neural networks with quantized activations, where the discrete mutual information can be computed exactly (in the limit of infinite samples). The original claims, controversy, and follow-up works are introduced and discussed in great detail. Original experiments are repeated with a very similar protocol to facilitate comparability of results. Additional experiments on MNIST are conducted, and ablations are performed. Results show the two distinct phases for some activation functions and some layers, but not for other activation functions.  **Main contributions**  1) Thorough discussion of the history of the Information-plane analysis, the main claims and previous findings and the follow-up papers it has spawned. Significance: The background and related work is well researched and well presented, which is very helpful for readers who have not followed this line of research closely. The only downside is that there is a fairly recent review (which is cited in the paper) which somewhat limits significance.  2) Quantized-activation training to avoid estimation errors when computing mutual information terms. Significance: the idea is very sensible in principle. Unfortunately, as some of the ablations show, the qualitative results can depend strongly on the quantization bit-width chosen. This is unfortunate since the main goal was to eliminate the influence of binning in naive estimation of the mutual information. Nonetheless, the results are exact and reliable for training quantized-activation neural networks - what is not clear is how to choose the quantization. 3) Reproduction of original experiments of Schwartz-Ziv & Tishby and Saxe et al. By sticking as closely as possible to the original training protocols, the new results are as comparable as possible (which does not mean that it is guaranteed that quantized-activation training dynamics are similar to “non-quantized” dynamics analyzed via binning - but empirically this seems to hold to a large degree). The paper also shows interesting ablations and additional experiments on MNIST. Significance: The most important experiments to run are included in the paper and important ablations are shown. The significance of the results could be improved by showing larger scale experiments on different kinds of networks (CNNs, ResNets, Transformers, …).