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 provides a principled way to examine the compression phrase, i.e, I(X;T) in deep neural networks. To achieve this, the authors provides an theoretical sounding entropy estimator to estimate mutual information.  Empirically, the paper did observe this compression phrase across both synthetic and real-world data and relates this compression behavior with geometric clustering.