Estimating Information Flow in DNNs

1 INTRODUCTION . Recent work by Shwartz-Ziv & Tishby ( 2017 ) uses the Information Bottleneck framework ( Tishby et al. , 1999 ; Tishby & Zaslavsky , 2015 ) to study the dynamics of DNN learning . The framework considers the mutual information pair ( I ( X ; T ` ) , I ( Y ; T ` ) ) between the input X or the label Y and the network ’ s hidden layers T ` . Plotting the evolution of these quantities during training , Shwartz-Ziv & Tishby ( 2017 ) made two interesting observations : ( 1 ) while I ( Y ; T ` ) remains mostly constant as the layer index ` increases , I ( X ; T ` ) decreases , suggesting that layers gradually shed irrelevant information about X ; and ( 2 ) after an initial fitting phase , there is a long compression phase during which I ( X ; T ` ) slowly decreases . It was suggested that this compression is responsible for the generalization performance of DNNs . A follow-up paper ( Saxe et al. , 2018 ) contends that compression is not inherent to DNN training , claiming double-sided saturating nonlinearities yield compression while single-sided/non-saturating ones do not necessarily compress . Shwartz-Ziv & Tishby ( 2017 ) and Saxe et al . ( 2018 ) present many plots of ( I ( X ; T ` ) , I ( Y ; T ` ) ) evolution across training epochs . These plots , however , are inadvertently misleading : they show a dynamically changing I ( X ; T ` ) when the true mutual information is provably either infinite or a constant independent of the DNN ’ s parameters ( see ( Amjad & Geiger , 2018 ) for a discussion of further degeneracies related to to the Information Bottleneck framework ) . Recall that the mutual information I ( X ; T ` ) is a functional of the joint distribution of ( X , T ` ) ∼ PX , T ` = PXPT ` |X , and that , in standard DNNs , T ` is a deterministic function of X . Hence , if PX is continuous , then so is T ` , and thus I ( X ; T ` ) = ∞ ( cf . ( Polyanskiy & Wu , 2012-2017 , Theorem 2.4 ) ) . If PX is discrete ( e.g. , when the features are discrete or if X adheres to an empirical distribution over the dataset ) , then the mutual information is a finite constant that does not depend on the parameters of the DNN . Specifically , for deterministic DNNs , the mapping from a discrete X to T ` is injective for strictly monotone nonlinearities such as tanh or sigmoid , except for a measure-zero set of weights . In other words , deterministic DNNs can encode all information about a discrete X in arbitrarily fine variations of T ` , causing no loss of information and implying I ( X ; T ` ) = H ( X ) , even if deeper layers ` have fewer neurons . The compression observed in Shwartz-Ziv & Tishby ( 2017 ) and Saxe et al . ( 2018 ) therefore can not be due to changes in mutual information . This discrepancy between theory and experiments originates from a theoretically unjustified discretization of neuron values in their approximation of I ( X ; T ` ) . To clarify , the quantity computed and plotted in these works is I ( X ; Bin ( T ` ) ) , where Bin is a per-neuron discretization of each hidden activity of T ` into a user-selected number of bins . This I ( X ; Bin ( T ` ) ) is highly sensitive to the selection of bin size ( as illustrated in Fig . 1 ) and does not track I ( X ; T ` ) for any choice of bin size.1 Nonetheless , compression results based on I ( X ; Bin ( T ` ) ) are observed by Shwartz-Ziv & Tishby ( 2017 ) and Saxe et al . ( 2018 ) in many interesting cases . To understand this curious phenomenon we first develop a rigorous framework for tracking the flow of information in DNNs . In particular , to ensure I ( X ; T ` ) is meaningful for studying the learned representations , we need to make the map X 7→ T ` a stochastic parameterized channel whose parameters are the DNN ’ s weights and biases . We identify several desirable criteria that such a stochastic DNN framework should fulfill for it to provide meaningful insights into commonly used practical systems . ( 1 ) The stochasticity should be intrinsic to the operation of the DNN , so that the characteristics of mutual information measures are related to the learned internal representations , and not to an arbitrary user-defined parameter . ( 2 ) The stochasticity should relate the mutual information to the deterministic binned version I ( X ; Bin ( T ` ) ) , since this is the object whose compression was observed ; this requires the injected noise to be isotropic over the domain of T ` analogously to the per-neuron binning operation . And most importantly , ( 3 ) the network trained under this stochastic model should be closely related to those trained in practice . We propose a stochastic DNN framework in which independent and identically distributed ( i.i.d . ) Gaussian noise is added to the output of each of the DNN ’ s neurons . This makes the map from X to T ` stochastic , ensures the data processing inequality ( DPI ) is satisfied , and makes I ( X ; T ` ) reflect the true operating conditions of the DNN , following Point ( 1 ) . Since the noise is centered and isotropic , Point ( 2 ) holds . As for Point ( 3 ) , Section 2 experimentally shows the DNN ’ s learned representations and performance are not meaningfully affected by the addition of noise , for variances β2 not too large . Furthermore , randomness during training has long been used to improve neural network performance , e.g. , to escape poor local optima ( Hinton et al. , 1984 ) , improve generalization performance ( Srivastava et al. , 2014 ) , encourage learning of disentangled representations ( Achille & Soatto , 2018 ) , and ensure gradient flow with hard-saturating nonlinearities ( Gulcehre et al. , 2016 ) . Under the stochastic model , I ( X ; T ` ) has no exact analytic expression and is impossible to approximate numerically . In Section 3 we therefore propose a sampling technique that decomposes the estimation of I ( X ; T ` ) into several instances of a simpler differential entropy estimation problem : estimating h ( S + Z ) given n samples of the d-dimensional random vector S and knowing the distribution of Z ∼ N ( 0 , β2Id ) . We analyze this problem theoretically and show that any differential entropy estimator over the noisy DNN requires at least exponentially many samples in the dimension d. Leveraging the explicit modeling of S + Z , we then propose a new estimator that converges 1Another approach taken in Saxe et al . ( 2018 ) considers I ( X ; T ` + Z ) ( instead of I ( X ; Bin ( T ` ) ) ) , where Z is an independent Gaussian with a user-defined variance . This approach has two issues : ( i ) the values as a function of ` may violate the data processing inequality , and ( ii ) they do not reflect the operation of the actual DNN , which was trained without noise . We focus on I ( X ; Bin ( T ` ) ) because it was commonly used in Shwartz-Ziv & Tishby ( 2017 ) and Saxe et al . ( 2018 ) , and since both methods have a similar effect of blurring T ` . as O ( ( log n ) d/4/ √ n ) , which significantly outperforms the convergence rate of general-purpose differential entropy estimators when applied to the noisy DNN framework . We find that I ( X ; T ` ) exhibits compression in many cases during training of small DNN classifiers . To explain compression in an insightful yet rigorous manner , Section 4 relates I ( X ; T ` ) to the well-understood notion of data transmission over additive white Gaussian noise ( AWGN ) channels . Namely , I ( X ; T ` ) is the aggregate information transmitted over the channel PT ` |X with input X drawn from a constellation defined by the data samples and the noisy DNN parameters . As training progresses , the representations of inputs from the same class tend to cluster together and become increasingly indistinguishable at the channel ’ s output , thereby decreasing I ( X ; T ` ) . Furthermore , these clusters tighten as one moves into deeper layers , providing evidence that the DNN ’ s layered structure progressively improves the representation of X to increase its relevance for Y . Finally , we examine clustering in deterministic DNNs . We identify methods for measuring clustering that are valid for both noisy and deterministic DNNs , and show that clusters of inputs in learned representations typically form in both cases . We complete the circle back to I ( X ; Bin ( T ` ) ) by clarifying why this binned mutual information measures clustering . This explains what previous works were actually observing : not compression of mutual information , but increased clustering by hidden representations . The geometric clustering of hidden representations is thus the fundamental phenomenon of interest , and we aim to test its connection to generalization performance , theoretically and experimentally , in future work . 2 PRELIMINARY DEFINITIONS Tℓ−1 σ ( W Noisy DNNs : For integers k ≤ ` , let [ k : ` ] , { i ∈ Z ∣∣k ≤ i ≤ ` } and use [ ` ] when k = 1 . Consider a noisy DNN with L+ 1 layers { T ` } ` ∈ [ 0 : L ] , with input T0 = X and output TL . The ` th hidden layer , ` ∈ [ L − 1 ] , is described by T ` = f ` ( T ` −1 ) + Z ` , where f ` : Rd ` −1 → Rd ` is a deterministic function of the previous layer and Z ` ∼ N ( 0 , β2Id ` ) ; no noise is injected to the output , i.e. , TL = fL ( TL−1 ) . We set S ` , f ` ( T ` −1 ) and use ϕ for the probability density function ( PDF ) of Z ` . The functions { f ` } ` ∈ [ L ] can represent any type of layer ( fully connected , convolutional , max-pooling , etc. ) . Fig . 2 shows a neuron in the ` th layer of a noisy DNN . To explore the relation between noisy and deterministic DNNs under conditions representative of current machine learning practices , we trained four-layer convolutional neural networks ( CNNs ) on MNIST ( LeCun et al. , 1999 ) . The CNNs used different levels of internal noise , including no noise , and one used dropout in place of additive noise . We measured their performance on the validation set and characterized the cosine similarities between their internal representations . Full details of the CNN architecture and training procedure are in Supplement 9.3 . The results in Table 1 show small amounts of internal additive noise ( β ≤ 0.1 ) have a minimal impact on classification performance , while dropout strongly improves it . The histograms in Fig . 3 show that the noisy ( for small β ) and dropout models learn internal representations similar to the representations learned by the deterministic model . In this high-dimensional space , unrelated representations would create cosine similarity histograms with zero mean and standard deviation between 0.02–0.3 , so the observed values are quite large . As expected , dissimilarity increases as the noise increases , and similarity is lower for the internal layers ( 2 and 3 ) . Mutual Information : Noisy DNNs induce a stochastic map from X to the rest of the network , described by the conditional distribution PT1 , ... , TL|X . The corresponding PDF 2 is pT1 , ... , TL|X=x . Its marginals are denoted by keeping only the relevant variables in the subscript . Let X , { xi } i∈ [ m ] be 2PT1 , ... , TL|X=x is absolutely continuous with respect to ( w.r.t . ) the Lebesgue measure for all x ∈ X . the input dataset , and P̂X be its empirical distribution , described by the probability mass function ( PMF ) p̂X ( x ) = 1m ∑ i∈ [ m ] 1 { xi=x } , for x ∈ X . Since data sets typically contain no repetitions , we assume p̂X ( x ) = 1m , ∀x ∈ X . The input and the hidden layers are jointly distributed according to3 PX , T1 , ... , TL , P̂XPT1 , ... , TL|X , under which X − T1 − . . .− TL−1 − TL forms a Markov chain . For each ` ∈ [ L− 1 ] , we study the mutual information ( Supplement 7 explains this factorization ) I ( X ; T ` ) , ∫ X×Rd ` dPX , T ` log ( dPX , T ` dPX × PT ` ) = h ( pT ` ) − 1 m ∑ i∈ [ m ] h ( pT ` |X=xi ) , ( 1 ) where log ( · ) is with respect to the natural base . Although PT ` and PT ` |X are readily sampled from using the DNN ’ s forward pass , these distributions are too complicated ( due to the composition of Gaussian noises and nonlinearities ) to analytically compute I ( X ; T ` ) or even to evaluate their densities at the sampled points . Therefore , we must estimate I ( X ; T ` ) directly from the available samples .

Summary: This paper studied the information bottleneck principle for deep learning. In the paper by (Schwatz-Ziv & Tishby 17'), it is empirically shown that the mutual information I(X;T) between input X and internal layers T decreases, which is called a compression phase. In this paper, the author found that the compression phase is not always happening and the shape of the curve of I(X;T) highly depends on the "bining size" which is used for estimating mutual information by (Schwatz-Ziv & Tishby 17'). Then the authors proposed to use a noisy DNN to make sure the map X->T is stochastic, then proposed a guaranteed mutual information estimator. Then some empirical results are shown.