A Probabilistic Representation for Deep Learning: Delving into The Information Bottleneck Principle

The Information Bottleneck ( IB ) principle has recently attracted great attention to1 explaining Deep Neural Networks ( DNNs ) , and the key is to accurately estimate the2 mutual information between a hidden layer and dataset . However , some unsettled3 limitations weaken the validity of the IB explanation for DNNs . To address these4 limitations and fully explain deep learning in an information theoretic fashion , we5 propose a probabilistic representation for deep learning that allows the framework6 to estimate the mutual information , more accurately than existing non-parametric7 models , and also quantify how the components of a hidden layer affect the mutual8 information . Leveraging the probabilistic representation , we take into account the9 back-propagation training and derive two novel Markov chains to characterize the10 information flow in DNNs . We show that different hidden layers achieve different11 IB trade-offs depending on the architecture and the position of the layers in DNNs,12 whereas a DNN satisfies the IB principle no matter the architecture of the DNN.13 1 Introduction14 Deep learning [ 18 ] has already achieved great success in numerous applications . Deep Neural15 Networks ( DNNs ) , however , are still commonly viewed as ‘ black boxes ’ [ 27 ] . Considerable efforts16 have been devoted to explaining the internal mechanism of DNNs from various perspectives , such as17 mathematics [ 5 , 12 ] , statistics [ 14 , 20 , 23 ] , computer vision [ 37 , 21 ] , etc . Recently , the Information18 Bottleneck ( IB ) principle has attracted attention in opening the ‘ black boxes ’ of DNNs [ 30 , 33 ] .19 Given a joint distribution P ( X , Y ) , the IB principle posits a random variable T = f ( X ) obeying the20 Markov chain Y → X → T and optimizes T by the IB Lagrangian [ 32 , 31 ] 21 min P ( T |X ) I ( X ; T ) − βI ( Y ; T ) , ( 1 ) where f ( · ) is an arbitrary function , I ( · ; · ) denotes mutual information , and the Lagrange multiplier22 β > 0 controls the IB trade-off between compressing the input X and preserving the information23 of the label Y . In the seminal work [ 30 ] , Tishby et al . manifest the IB trade-off in every layer of24 DNNs = { x ; t1 ; · · · ; tI ; ŷ } via studying I ( X ; Ti ) and I ( Y ; Ti ) , where Ti is the random variable of25 the ith hidden layer ti . Especially , the authors ascribe DNN generalization to the compression [ 29 ] .26 In the context of deterministic DNNs , recent works reveal some limitations of the IB principle for27 explaining DNNs . Amjad et al . argue that the IB principle becomes an ill-posed optimization problem28 due to I ( X ; Ti ) =∞ [ 1 ] , and Kolchinsky et al . demonstrate that not every layer of DNNs satisfies a29 strict IB trade-off , i.e. , different layers only differ in I ( X ; Ti ) but I ( Y ; Ti ) keeps consistent in all30 layers [ 15 ] . In addition , Saxe et al . experimentally show that the compression does not occur in31 DNNs with non-saturating activation functions , e.g. , the popular ReLU function [ 28 ] , and Goldfeld32 et al . doubt the causality between the generalization of DNNs and the compression [ 10 , 7 ] . These33 unsettled limitations greatly weakens the validity of the IB explanations for DNNs.34 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . The key to examining the IB principle in DNNs is the accurate estimation of the mutual information.35 However , regarding DNNs as deterministic models hinders us from specifying the random variable36 Ti and the distribution P ( Ti ) , thus it is difficult to accurately estimate I ( X ; Ti ) and I ( Y ; Ti ) . More37 specifically , in the absence of a clear definition of Ti , simply assuming the activations of ti as the i.i.d.38 samples of Ti induces Ti being a continuous random variable and I ( X ; Ti ) = ∞ in deterministic39 DNNs ( see Appendix C in [ 28 ] ) . The complicated architecture of DNNs makes it challenging to40 specify P ( Ti ) . Therefore , most previous works have to indirectly estimate P ( Ti ) via non-parametric41 models [ 35 ] , such as the empirical distribution [ 30 ] , Kernel Density Estimation ( KDE ) [ 28 ] , and42 Gaussian convolution [ 10 ] . However , we experimentally confirm that classical non-parametric models43 derives poor mutual information estimation [ 24 , 22 ] in DNNs , and one reason is because activations44 do not satisfy the i.i.d . prerequisite of non-parametric models ( see Appendix G ) . In summary , the45 limitations mainly stem from the lack of an explicit probabilistic representation for deep learning.46 The IB principle only formulates the information flow in DNNs = { x , t1 , · · · , tI , ŷ } after training,47 and the corresponding Markov chain ( see Fig . 1 in [ 30 ] ) 48 Y → X → T1 · · · → TI → Ŷ ( 2 ) indicates that the information of Y transfers to Ti in the forward direction and Ti receives the49 information of Y only via X . However , training DNNs by the back-propagation [ 25 ] implies that the50 information of Y transfers to Ti in the backward direction during training and retains information51 in Ti after training . Notably , Zhang et al . show that a DNN can fit labels well even using Gaussian52 noise as input to train the DNN [ 38 ] , which implies that Ti can directly receive the information of Y .53 Hence , the IB principle does not comprehensively characterize the information flow in DNNs.54 To address the above limitations and comprehensively explain DNNs in an information theoretic55 fashion , we introduce the probability space ( ΩTi , F , PTi ) [ 6 ] for the ith hidden layer ti in DNNs.56 Compared to previous works , the probability space ( ΩTi , F , PTi ) enables us to : ( i ) accurately estimate57 I ( X ; Ti ) and I ( Y ; Ti ) via specifying Ti and P ( Ti ) , and ( ii ) quantify the effect of the architecture of58 ti and the back-propagation on I ( X ; Ti ) and I ( Y ; Ti ) via explicitly modeling all the ingredients of ti,59 such as the activation function and the weights in a probabilistic way . To the best of our knowledge,60 this is the first time the probability space of a hidden layer in DNNs is as defined.61 Leveraging ( ΩTi , F , PTi ) , we derive information theoretic explanations for DNNs as follows:62 • Two Markov chains1 characterize the information flow in DNNs = { x , t1 , · · · , tI , ŷ } 63 X̄ →T1 → · · · → TI → Ŷ T1 ← · · · ← TI ← Ŷ ← Y . ( 3 ) • Different hidden layers manifest different IB trade-offs depending on the architecture and64 the position of hidden layers in DNNs.65 • A DNN satisfies the IB principle no matter the architecture of the DNN.66 Preliminaries . P ( X , Y ) = P ( X ) P ( Y |X ) is an unknown joint distribution between X and Y . A67 dataset D = { ( xj , yj ) |xj ∈ RM , yj ∈ Z } Jj=1 consists of J i.i.d . samples generated from P ( X , Y ) 68 with finite L labels , i.e. , yj ∈ { 1 , · · · , L } . In the context of supervised learning , we focus on69 feedfworad fully connected DNNs = { x , t1 , · · · , tI , ŷ } , i.e. , Multi-Layer Perceptions ( MLPs ) [ 8 ] 70 for the image classification task . Without loss of generality , we use the MLP = { x , t1 , t2 , ŷ } with71 the cross-entropy loss ` CE for most theoretical derivations . In addition , H ( · ) denotes entropy.72 In the MLP , t1 and t2 have N and K neurons , respectively , and t1 = { t1n = σ1 [ 〈ω ( 1 ) n , x〉 ] } Nn=1,73 where 〈ω ( 1 ) n , x〉 = ∑M m=1 ω ( 1 ) mn · xm + b1n is the nth dot-product given the weight ω ( 1 ) mn and the bias74 b1n , and σ1 ( · ) denotes an activation function , e.g. , ReLU . Similarly , t2 = { t2k = σ2 [ 〈ω ( 2 ) k , t1〉 ] } Kk=1,75 where 〈ω ( 2 ) k , t1〉 = ∑N n=1 ω ( 2 ) nk · t1n + b2k . The output layer ŷ is softmax with L nodes76 ŷ = { ŷl = 1 ZY exp [ 〈ω ( 3 ) l , t2〉 ] = 1 ZY exp [ gl ( t2 ( t1 ( x ) ) ) ] } Ll=1 , ( 4 ) where 〈ω ( 3 ) l , t2〉 = ∑K k=1 ω ( 3 ) kl · t2k + byl and ZY = ∑L l=1 exp [ 〈ω ( 3 ) l , t2〉 ] is the partition function.77 1In which the virtual random variable X̄ has all the information of X except Y , namely H ( X̄ ) = H ( X|Y ) . A convolution layer . 2 A probabilistic representation for deep learning78 To accurately estimate I ( X ; Ti ) and I ( Y ; Ti ) , in this section , we specify the probability space [ 6 ] for79 a fully connected layer and derive the probabilistic explanations of the entire MLP.80 It is known that a convolution kernel ( namely the weights of convolution ) defines a local feature,81 and a convolution operation derives a feature map to measure the cross-correlation between the82 local feature and input in a receptive field ( Chapter 9.1 in [ 11 ] ) . Notably , a fully connected layer83 is equivalent to a convolution layer with the kernel size having the same dimension as input . Thus84 the weights of a neuron can be viewed as a global feature , and a fully connected layer with multiple85 neurons derives activations to measure the cross-correlation between the multiple global features and86 the input . The cross-correlation explanation for a fully connected layer is visualized in Figure 1.87 Assuming that a fully connected layer t consists ofN neurons { tn = σ [ 〈ωn , z〉 ] } Nn=1 , where z ∈ RM88 is the input of t , 〈ωn , z〉 = ∑M m=1 ωmn · zm + bn is the dot-product between z and ωn , and σ ( · ) is89 an activation function . Based on the cross-correlation explanation , the behavior of t is to measure90 the cross-correlations between z and the N possible features defined by the the weights { ωn } Nn=1.91 In the context of pattern recognition [ 34 ] , we define a virtual random process or ‘ experiment ’ as t92 recognizing one of the patterns/features with the largest cross-correlation to z from the N possible93 features . The experiment characterizes the behavior of t ( i.e. , before recognizing the features with94 the largest cross-correlation , t must measure the cross-correlations between z and all the N possible95 features ) while meets the requirement of the ‘ experiment ’ definition ( i.e. , only one outcome will96 occur on each trial of the experiment [ 6 ] ) . The probability space ( ΩT , F , PT ) is defined as follows:97 Definition 1 . ( ΩT , F , PT ) consists of three components : the sample space ΩT has N possible98 outcomes ( features ) { ωn = { ωmn } Mm=1 } Nn=1 defined by the weights2 of the N neurons ; the event99 space F is the σ-algebra ; and the probability measure PT is a Gibbs distribution [ 19 ] to quantify the100 probability of ωn being recognized as the feature with the largest cross-correlation to z.101 Taking into account the randomness of z , the conditional distribution PT |Z is formulated as102 PT |Z ( ωn|z ) = 1 ZT exp ( tn ) = 1 ZT exp [ σ ( 〈ωn , z〉 ) ] , ( 5 ) where Z is the random variable of z and ZT = ∑N n=1 exp ( fn ) is the partition function.103 ( ΩT , F , PT ) clearly explains all the ingredients of t in a probabilistic fashion . The nth neuron104 defines a global feature by the weights wn and the activation tn = σ ( 〈ωn , z〉 ) measures the cross-105 correlation between wn and z . The Gibbs distribution PT |Z indicates that if wn has the higher106 activation , i.e. , the larger cross-correlation to z , it has the larger probability being recognized as107 the feature with largest cross-correlation to z . For instance , if z ∈ R16 and t includes N = 2108 neurons , then ΩT = { ω1 , ω2 } defines two possible outcomes ( features ) , where ωn = { ωmn } 16m=1.109 F = { ∅ , { ω1 } , { ω2 } , { ω1 , ω2 } } means that neither , one , or both of the features are recognized110 by t given z , respectively . PT |Z ( ω1|z ) and PT |Z ( ω2|z ) are the probability of ω1 and ω2 being111 recognized as the feature with the largest cross-correlation to z , respectively.112 2We do not take into account the scalar value bn for defining ΩT , as it not affects the feature defined by ωn . ( ΩT , F , PT ) explains the representation ability of deep learning . Compared to Restricted Boltzmann113 Machines ( RBMs ) [ 26 ] simply using binary units to indicate features being recognized or not given114 input , the Gibbs distribution3 PT |Z ( ωn|z ) measures the probability of ωn being recognized with115 the largest cross-correlation to z , i.e. , it characterizes the relation between features and input more116 accurately . Moreover , Equation 5 shows that tn = σ ( 〈ωn , z〉 ) is the negative energy function [ 19 ] of117 the Gibbs distribution , thus PT |Z ( ωn|z ) can be derived as long as σ ( 〈ωn , z〉 ) are known because118 the energy function is the sufficient statistics [ 2 ] of the Gibbs distribution . That enables subsequent119 hidden layers to generate high-level features of input via directly processing the activations { tn } Nn=1,120 thus deep learning can form a hierarchical structure to represent much complex features.121 ( ΩT , F , PT ) answers a fundamental question : which component of a hidden layer contains the122 information of the layer ? Since ωn defines ΩT , the weights contain all the information of a layer . In123 particular , since the activation tn = σ ( 〈ωn , z〉 ) is a function of ωn , the data processing inequality124 [ 4 ] indicates that the information of tn is no more than the information of ωn . Simulations in Section125 4.2 demonstrate that if activations do not correctly characterize the cross-correlation between weights126 and input , activations contain less information than weights do.127 Based on ( ΩT , F , PT ) , we define the random variable T as follows:128 Definition 2 . Given the fully connected layer t , we define the random variable T : ΩT → ET as129 T ( ωn ) , n , ( 6 ) where the measurable space ET = { 1 , · · · , N } .130 Since ΩT is composed of finite N possible outcomes , T is a discrete random variable . Notably , the131 one-to-one correspondence between ωn and n indicates132 PT |Z ( ωn|z ) = PT |Z ( n|z ) . ( 7 ) If not considering the back-propagation training , the weights ( namely ΩTi ) of each layer are fixed.133 Thus Ti+1 entirely depends on Ti and the MLP = { x ; t1 ; t2 ; ŷ } forms a Markov chain134 X → T1 → T2 → Ŷ . ( 8 ) Based on the corresponding joint distribution P ( Ŷ , T2 , T1|X ) = P ( T1|X ) P ( T2|T1 ) P ( Ŷ |T2 ) and135 Definition 2 , we derive a probabilistic explanation for the entire MLP , which is summarized in136 Theorem 1 . The detailed derivation is presented in Appendix B.137 Theorem 1 . The MLP = { x ; t1 ; t2 ; ŷ } formulates a conditional Gibbs distribution138 PŶ |X ( l|x ) = K∑ k=1 N∑ n=1 P ( Ŷ = l , T2 = k , T1 = n|X = x ) = 1 ZMLP ( x ) exp [ gl ( t2 ( t1 ( x ) ) ) ] , ( 9 ) where ZMLP ( x ) = ∑L l=1 ∑K k=1 ∑N n=1 PŶ , T2 , T1|X ( l , k , n|x ) is the partition function.139 Since PŶ |X ( l|x ) exactly equals the output ŷl of the MLP , namely Equation ( 4 ) , we conclude that140 the entire architecture of the MLP forms a family of Gibbs distribution PŶ |X ( l|x ) . In general , the141 back-propagation updates a weight ω based on the gradient of ` CE with respect to ω,142 ω ( s+ 1 ) = ω ( s ) − α · ∂ ` CE ∂ω ( s ) = ω ( s ) − α · ∂KL [ P ( Y |X ) ||P ( Ŷ |X ) ] ∂ω ( s ) , ( 10 ) where s is the index of training iteration , α is the training rate , and KL [ ·||· ] is the KL-divergence.143 Figure 2 summarizes the probabilistic explanation for deep learning based on the MLP . In general,144 a single learning iteration , an epoch , consists of two phases : training and inference ( after training ) .145 During inference , the MLP bridges X and Ŷ via multiple intermediate features ΩT1 , ΩT2 , and ΩŶ146 defined by weights , and formulates the statistical relation between Ŷ and X as a family of conditional147 Gibbs distribution P ( Ŷ |X ) . During training , the back-propagation updates weights to learn optimal148 intermediate features for searching an optimal P ( Ŷ |X ) to accurately approximate P ( Y |X ) .149 3Recent works about Gibbs explanations for a hidden layer are discussed in Appendix A . 3 The information theoretic explanations for deep learning150 To address the limitations of existing IB explanations , this section proposes some novel information151 theoretic explanations for DNNs based on the proposed probabilistic representation.152 Proposition 1 . The mutual information between a fully connected layer and dataset is finite.153 I ( X ; T ) < ∞ . ( 11 ) Proof : Definition 2 shows ET = { 1 , · · ·N } . Thus T is a discrete random variable and H ( T ) < ∞,154 thereby I ( X ; T ) ≤ H ( T ) < ∞.155 Proposition 1 circumvents the infinite mutual information problem . In the absence of a clear definition156 T : ΩT → ET , most previous works [ 28 , 3 , 1 ] simply viewing the activation tn as the sample of T ,157 namely tn ∈ ET = R , implies T being continuous and gives rise to the infinite mutual information158 problem in deterministic DNNs . However , ( ΩT , F , PT ) indicates that tn actually is a variable159 measuring the cross-correlation between wn and z rather than the sample of T , namely tn /∈ ET .160 Theorem 2 . The information of Y flows into the MLP in the backward direction during training161 T1 ← T2 ← Ŷ ← Y . ( 12 ) Proof : First , since ΩT is defined by ω in ( ΩT , F , PT ) and Equation ( 10 ) shows that ω ( s + 1 ) is162 determined by all the previous gradients { ∂ ` CE∂ω ( s ) } S s=1 , and ω ( 0 ) is randomly initialized and α is a163 constant , we can derive that ΩT is determined by ∂ ` CE∂ω . Second , based on the back-propagation , the164 relation between gradients in two adjacent layers in the MLP = { x ; t1 ; t2 ; ŷ } is formulated as165 ∂ ` ? CE ∂ω ( 3 ) kl = [ PŶ |X ( l|x ) −PY |X ( l|x ) ] · t2k , ∂ ` CE ∂ω ( 2 ) nk = L∑ l=1 ∂ ` ? CE ∂ω ( 3 ) kl · ω ( 3 ) kl · σ′2 ( 〈ω ( 2 ) k , t1〉 ) f2k · t1n , ∂ ` CE ∂ω ( 1 ) mn = K∑ k=1 ∂ ` CE ∂ω ( 2 ) nk · ω ( 2 ) nk · σ′1 ( 〈ω ( 1 ) n , x〉 ) t1n · xm . ( 13 ) Equation 13 shows that ∂ ` CE ∂ω ( 3 ) is a function of PY |X ( l|x ) and ∂ ` CE∂ω ( i ) is a function of ∂ ` CE ∂ω ( i+1 ) , where166 ω ( 3 ) denotes the weight of ŷ . The two points above enable us to derive that ΩTi is a function of ΩTi+1167 and ΩŶ is a function of P ( Y |X ) . Based on Definition 2 , we can further derive that Ti is a function168 of Ti+1 and Ŷ is a function of Y , i.e. , T1 ← T2 ← Ŷ ← Y . ( See the detailed proof in Appendix C ) .169 Theorem 2 is consistent with the prevailing explanation for deep learning . LeCunn et al . show that170 deep learning exploits the hierarchical property of signals [ 18 ] , i.e. , the layers farther from output171 learn lower-level features , such as edges , whereas the layers closer to output assemble lower-level172 features into the higher-level features corresponding to labels ( see Figure 2 in [ 37 ] ) . Notably , since173 lower-level features commonly exist in signals with different labels ( e.g. , lower-level features , such174 as the edges of the vehicle frame and the circular contour of wheels , exist in both the car and the175 truck classes in the CIFAR-10 dataset [ 16 ] in Figure 2 ) , lower-level features do not contain much176 information of labels . Therefore , the layers farther from output do not have much information of177 labels , which is consistent with the Markov chain T1 ← T2 ← Ŷ ← Y .178 Since all the information of Y stems from X ( i.e. , H ( Y ) = I ( X ; Y ) proven in Appendix D ) ,179 Theorem 2 implies that partial information of X flows into the MLP in the backward direction during180 training . Equation ( 2 ) shows the information of X flowing into the MLP in the forward direction181 during inference . Overall , the information of X flows in the backward and forward directions during182 training and inference , respectively . As a result , the Markov chain , Equation ( 2 ) , proposed by recent183 works could not fully characterize the information flow of X in the MLP in each epoch . In other184 words , I ( X ; Ti ) is not necessarily greater than I ( X ; Ti+1 ) in the MLP in each epoch.185 Equation ( 2 ) shows that Ti receives the information of Y via X during inference . Theorem 2 shows186 that Ti also directly receives information of Y during training , because the back-propagation updates187 weights ( i.e. , ΩTi ) based on the label Y . Thus Equation ( 2 ) can not fully characterize the information188 flow of Y in the MLP in each epoch , when we take into account the back-propagation training.189 To fully characterize the information flow in the MLP in each epoch , we introduce Corollary 1.190 Corollary 1 . The information flow in the MLP can be characterized by two Markov chains as191 X̄ →T1 → T2 → Ŷ T1 ← T2 ← Ŷ ← Y . ( 14 ) The virtual random variable X̄ contains all the information of X except Y , i.e. , H ( X̄ ) = H ( X|Y ) .192 Proof of the first Markov chain : Since X̄ does not have any information of Y , it can only flow into193 the MLP in the forward direction during inference . Again since X̄ does not have any information of194 Y , the information flow of Y during training will not affect the information flow of X̄ . Therefore,195 X̄ → T1 → T2 → Ŷ characterizes the information flow of X̄ in both training and inference phases.196 Proof of the second Markov chain : Since the weights are fixed after training , the sample space and197 the distribution of hidden layers are fixed after training . Therefore , the information of Y transferred198 into hidden layers during training will retain there after training ( i.e. , during inference ) . In addition,199 Definition 1 indicates that a fully connected layer t = { tn = σ ( 〈ωn , z〉 ) } Nn=1 measures the cross-200 correlation between ω ( 1 ) n and z during inference , thus { ω ( 1 ) n } Nn=1 can be viewed as a representation201 of Z . As a result , even though Z has all the information of Y , the information of Y that t can learn202 from Z is determined by how much information of Y the representation { ω ( 1 ) n } Nn=1 has . Overall , the203 information flow of Y during inference will be the same as that during training . Based on Theorem 2,204 we conclude that T1 ← T2 ← Ŷ ← Y characterizes the information flow of Y in the MLP in both205 training and inference phases . Detailed derivations and explanations are presented in Appendix E.206 To quantify how much information of X and Y is learned by the MLP , we introduce Corollary 2.207 Corollary 2 . The mutual information between dataset and the entire MLP can be expressed as208 I ( X ; TMLP ) = I ( X̄ ; T1 ) + I ( Y ; Ŷ ) I ( Y ; TMLP ) = I ( Y ; Ŷ ) ( 15 ) where TMLP denotes a random variable corresponding to the entire architecture of the MLP.209 Proof : Since H ( Y ) = I ( X ; Y ) ( Appendix D ) , H ( X ) = H ( X̄ ) + I ( X ; Y ) = H ( X̄ ) + H ( Y ) .210 Hence , Corollary 2 can be derived by Corollary 1 and the chain rule . The proof is in Appendix F.211 4 Simulations212 In this section , we propose a mutual information estimator based on ( ΩT , F , PT ) and demonstrate the213 probabilistic representation and information theoretic explanations for deep learning on a synthetic214 dataset with known entropy . Additional experiments on benchmark datasets are in Appendix H.215 4.1 Setup216 Mutual information estimator . Based on the definition of mutual information , we have217 I ( X ; Ti ) = H ( Ti ) −H ( Ti|X ) . ( 16 ) Previous works simply estimate I ( X ; Ti ) = H ( Ti ) , because Ti is assumed to be entirely dependent218 on X in the Markov chain , Equation ( 2 ) , thereby H ( Ti|X ) = 0 . However , Corollary 1 shows that Ti219 depends on both X and Y if taking into account the training phase , thereby H ( Ti|X ) 6= 0.220 Image0 ( label [ 1,0 ] ) Image1 ( label [ 0,1 ] ) Image2 ( label [ 1,0 ] ) Image3 ( label [ 0,1 ] ) To accurately estimate I ( X ; Ti ) , we need to specify P ( Ti|X ) and P ( Ti ) . Based on ( ΩTi , F , PTi ) ,221 we formulate PTi|X ( n|xj ) of the three fully connected layers in the MLP as222 PT1|X ( n|xj ) = 1ZF1 exp [ σ1 ( 〈ω ( 1 ) n , xj〉 ) ] , PT2|X ( k|xj ) = 1ZF2 exp [ σ2 ( 〈ω ( 2 ) k , t1 ( x j ) 〉 ) ] , PTY |X ( l|xj ) = 1ZFY exp [ 〈ω ( 3 ) l , t2 ( t1 ( x j ) ) 〉 ] . ( 17 ) To derive the marginal distribution P ( Ti ) , we sum the joint distribution P ( Ti , X ) over x ∈ X ,223 P ( Ti = n ) = ∑ x∈X PX ( x ) PTi|X ( n|x ) ≈ ∑ xj∈D PX ( x j ) PTi|X ( n|xj ) = 1J ∑ xj∈D PTi|X ( n|xj ) , ( 18 ) where PX ( xj ) is estimated by the empirical distribution 1/J givenD . Finally , we can derive I ( X ; Ti ) 224 by Equation 16 , 17 , and 18 . Similarly , based on the definition of mutual information , we have225 I ( Y ; Ti ) = H ( Ti ) −H ( Ti|Y ) . ( 19 ) To estimate H ( Ti|Y ) , we reformulate P ( Ti|Y ) as226 PTi|Y ( n|l ) = ∑ x∈X PTi|X ( n|x ) PX|Y ( x|l ) ≈ 1 N ( l ) ∑ xj∈D , yj=l PTi|X ( n|xj ) , ( 20 ) where PX|Y ( xj |l ) is estimated by the empirical distribution 1/N ( l ) and N ( l ) denotes the number of227 samples with the label l in D. Finally , we can derive I ( Y ; Ti ) by Equation 18 , 19 , and 20.228 Synthetic dataset . The dataset consists of 512 gray-scale 32×32 images , which are evenly generated229 by rotating a deterministic image x̂ in four different orientations and adding Gaussian noise with230 expectation µ = E ( x̂ ) and variance σ2 = 1 , namely x = r ( x̂ ) + N ( µ , σ2 ) , where r ( · ) denotes231 the rotation method shown in Figure 3 . The reason for adding Gaussian noise is to avoid DNNs232 directly memorizing the deterministic image . In addition , the binary labels [ 1,0 ] and [ 0,1 ] evenly233 divide the synthetic dataset into two classes . As a result , the synthetic dataset has ( approximately ) 234 2 bits information and the labels have 1 bit information . Compared to popular benchmark dataset235 with unknown features and entropy , e.g. , MNIST [ 17 ] and Fashion-MNIST [ 36 ] , the features and the236 entropy of the synthetic dataset are clear and known , which enables us to examine the probabilistic237 representation and the mutual information estimator.238 Neural Networks . We train three MLPs , namely MLP1 , MLP2 and MLP3 , on the synthetic dataset239 by a variant of Stochastic Gradient Descent ( SGD ) method , namely Adam [ 13 ] , over 1000 epochs240 with the learning rate α = 0.03 . Table 1 summarizes the architecture of the three MLPs.241 4.2 Validating the probability space and the mutual information estimator242 We demonstrate the sample space ΩT by visualizing the weights4 of the eight neurons in t1 , i.e.,243 ω ( 1 ) n = { ω ( 1 ) mn } 1024m=1 , in 5 different epochs ( i.e. , 0,1,4,128,1000 ) in Figure 4 ( Left ) . As training244 continues , we observe that ω ( 1 ) n quickly learns all the spatial features of the synthetic dataset . For245 instance , ω ( 1 ) 2 has low magnitude at top-left positions and high magnitude at bottom-right positions,246 which correctly characterizes the spatial feature of Image0 . Similarly , ω ( 1 ) 3 , ω ( 1 ) 4 , and ω ( 1 ) 5 correctly247 characterize the spatial feature of Image1 , Image2 , and Image3 in Figure 3 , respectively.248 4We only show the learned weights in MLP1 because we observe that the learned weights in MLP1 and MLP2 are very similar , though they use different activation functions . We demonstrate that P ( T1|X ) correctly measures the probability of { ω ( 1 ) n } 8n=1 being recognized the249 feature with the largest cross-correlation to x in Table 2 . For instance , ω ( 1 ) 2 correctly characterizes250 the feature of Image0 and has the largest cross-correlation 〈ω ( 1 ) 2 , x〉 = 190.8 , thus it has the largest251 probability PReLUT1|X ( ω ( 1 ) 2 |Image0 ) = 1.0 being recognized as the feature with largest cross-correlation252 to Image0 . In contrast , since ω ( 1 ) 7 incorrectly characterizes the feature of Image0 and has the lowest253 cross-correlation 〈ω ( 1 ) 7 , x〉 = −210.0 , so it has the lowest probability PReLUT1|X ( ω ( 1 ) 7 |Image0 ) = 0.0254 being recognized as the feature with largest cross-correlation to Image0.255 We observe that an activation function ( abbr . ACT ) plays an important role in the distribution.256 Specifically , ReLU , a non-saturating ( unbounded ) ACT [ 9 ] , preserves the positive cross-correlations257 while resets all the negative ones as zero . PReLUT1|X ( ω ( 1 ) 2 |Image0 ) = 1.0 shows that ReLU derives the258 correct probability of ω ( 1 ) 2 being recognized as the feature with largest cross-correlation . In contrast,259 though ω ( 1 ) 2 has stronger cross-correlation to Image0 than ω ( 1 ) 4 , i.e. , 〈ω ( 1 ) 2 , x〉 > 〈ω ( 1 ) 4 , x〉 , Tanh , a260 saturating ( bounded ) ACT , derives fTanh12 ( x ) = f Tanh 14 ( x ) = 1.0 , and makes ω ( 1 ) 4 to incorrectly have261 the same probability 0.272 to ω ( 1 ) 2 being recognized as the feature with the largest cross-correlation262 to Image0 , i.e. , Tanh hinders t1 from correctly recognizing the features of input . The simulations for263 validating the probability space based on other synthetic images are presented in Appendix G.264 To validate the mutual information estimator , we follow recent works [ 30 , 28 ] to train the three265 MLPs with 50 different random initialization and study the average mutual information . Figure 4266 ( Right ) shows that I ( X ; T1 ) quickly increases to 1.81 and keeps stable in the MLP1 , i.e. , t1 learns267 most information of the dataset as H ( X ) = 2.0 . Notably , the result is consistent with the variation268 of the weights in Figure 4 ( Left ) , which shows that the weights correctly characterize the features269 of the dataset and keeps stable after the fourth epoch . As a comparison , we observe that I ( X ; T1 ) 270 keeps stable at 0.44 in the MLP2 , which confirms the statement that Tanh hinders t1 from correctly271 recognizing the features of input . In addition , Figure 4 ( Right ) shows that I ( X ; T1 ) ≈ 0.79 in MLP3272 is smaller than I ( X ; T1 ) ≈ 1.81 in MLP1 , which is consistent with Definition 1 , i.e. , a layer with273 fewer neurons would represent fewer possible features , thus it contains less information.274 In summary , we demonstrate the probability space ( ΩT , F , PT ) and show that if an ACT cannot275 preserve the cross-correlation between weights ( features ) and input , it would distort the distribution276 of a layer , thereby affecting the mutual information between the layer and data/labels . In addition,277 we show that the proposed mutual information estimator outperforms the existing non-parametric278 models , e.g. , empirical distribution [ 30 ] and KDE [ 28 ] , based on the synthetic dataset . Especially,279 activations do not satisfy the i.i.d . prerequisite of non-parametric models is an important reason for280 non-parametric models deriving inaccurate mutual information in DNNs . Due to limited space , the281 experimental comparison and study of non-parametric models are presented in Appendix G.282 4.3 Validating the information theoretic explanations for DNNs283 In Figure 5 , we observe I ( X ; Ti ) ≤ I ( X ; Ŷ ) in MLP2 and MLP3 , which confirms that the Markov284 chain proposed by previous works , Equation ( 2 ) , can not fully explain the information flow in MLPs,285 if taking into account the back-propagation training . As a comparison , the second and third row286 show I ( X̄ ; T1 ) ≥ I ( X̄ ; T2 ) ≥ I ( X̄ ; Ŷ ) and I ( Y ; T1 ) ≤ I ( Y ; T2 ) ≥ I ( Y ; Ŷ ) in all the three MLPs,287 which validates that Corollary 1 , i.e. , Equation ( 14 ) characterizes the information flow in MLPs.288 Figure 5 demonstrates that different hidden layers achieve different IB trade-offs depending on289 the architecture and the position of the layers in MLPs . In terms of architecture , I ( Y ; T1 ) > 0.8290 and I ( X̄ ; T1 ) > 0.75 in MLP1 indicate that t1 , with ReLU , achieves a good prediction without291 much compression , whereas I ( Y ; T1 ) < 0.5 and I ( X̄ ; T1 ) < 0.1 in MLP2 show that t1 , with Tanh,292 achieves a different IB trade-off . In addition , I ( Y ; T1 ) ≈ 0.45 and I ( X̄ ; T1 ) ≈ 0.25 in MLP3 show293 the effect of neuron numbers on the IB trade-off . In terms of position , I ( Y ; Ŷ ) = 1 and I ( X̄ ; Ŷ ) = 0294 in MLP1 means that ŷ has a different IB trade-off to t1 in MLP1.295 We demonstrate that a MLP satisfies the IB principle no matter what the architecture of the MLP296 is . Figure 5 visualizes I ( X ; TMLP ) and I ( Y ; TMLP ) based on Corollary 2 . It shows that all of three297 MLPs satisfy the IB principle , namely I ( X ; TMLP ) < H ( X ) = 2 and I ( Y ; TMLP ) = H ( Y ) = 1,298 though they have different architectures . Importantly , in contrast to previous work [ 28 ] claiming that299 the compression not exists in DNNs with non-saturating ACT , such as ReLU , Figure 5 clearly shows300 that the compression exists in all the MLPs , no matter the activation function of MLPs.301 We further demonstrate the information theoretic explanations for DNNs on the benchmark MNIST302 and Fashion-MNIST datasets . The experiments are presented in Appendix H.303 5 Conclusion and future work304 In this work , we ( 1 ) specify the probability space for a hidden layer for ( 2 ) accurately estimating the305 mutual information and ( 3 ) clearly explaining how the components of the layer affect the mutual306 information . We take into account the back-propagation training and derive two novel Markov chains307 to characterize the information flow in DNNs . Furthermore , we demonstrate that a DNN satisfies the308 IB principle no matter the architecture of the DNN . In contrast , different hidden layers show different309 IB trade-offs depending on the architecture and the position of the layers in DNNs . A potential310 direction is to study the generalization of DNNs based on the probabilistic representation.311 References312 [ 1 ] Rana Ali Amjad and Bernhard Claus Geiger . Learning representations for neural network-based classi-313 fication using the information bottleneck principle . IEEE transactions on pattern analysis and machine314 intelligence , 2019.315 [ 2 ] George Casella and Roger L Berger . Statistical inference . Cengage Learning , 2021.316 [ 3 ] Ivan Chelombiev , Conor Houghton , and Cian O ’ Donnell . Adaptive estimators show information compres-317 sion in deep neural networks . In International Conference on Learning Representations , 2019.318 [ 4 ] Thomas Cover and Joy Thomas . Elements of Information Theory . Wiley-Interscience , Hoboken , New319 Jersy , 2006.320 [ 5 ] Balázs Csanád Csáji et al . Approximation with artificial neural networks . Faculty of Sciences , Etvs Lornd321 University , Hungary , 24 ( 48 ) :7 , 2001.322 [ 6 ] Rick Durrett . Probability : theory and examples , volume 49 . Cambridge university press , 2019.323 [ 7 ] Marylou Gabrié , Andre Manoel , Clément Luneau , Nicolas Macris , Florent Krzakala , Lenka Zdeborová,324 et al . Entropy and mutual information in models of deep neural networks . In Advances in Neural325 Information Processing Systems , pages 1821–1831 , 2018.326 [ 8 ] Matt W Gardner and SR Dorling . Artificial neural networks ( the multilayer perceptron ) —a review of327 applications in the atmospheric sciences . Atmospheric environment , 32 ( 14-15 ) :2627–2636 , 1998.328 [ 9 ] Xavier Glorot and Yoshua Bengio . Understanding the difficulty of training deep feedforward neural329 networks . In Proceedings of the thirteenth international conference on artificial intelligence and statistics,330 pages 249–256 , 2010.331 [ 10 ] Ziv Goldfeld , Ewout Van Den Berg , Kristjan Greenewald , Igor Melnyk , Nam Nguyen , Brian Kingsbury,332 and Yury Polyanskiy . Estimating information flow in deep neural networks . In Proceedings of the 36th333 International Conference on Machine Learning , volume 97 , pages 2299–2308 , 2019.334 [ 11 ] Ian Goodfellow , Yoshua Bengio , Aaron Courville , and Yoshua Bengio . Deep learning , volume 1 . MIT335 press Cambridge , 2016.336 [ 12 ] Patrick Kidger and Terry Lyons . Universal approximation with deep narrow networks . In Conference on337 Learning Theory , pages 2306–2327 . PMLR , 2020.338 [ 13 ] Diederik P Kingma and Jimmy Ba . Adam : A method for stochastic optimization . arXiv preprint339 arXiv:1412.6980 , 2014.340 [ 14 ] Pang Wei Koh and Percy Liang . Understanding black-box predictions via influence functions . In341 International Conference on Machine Learning , pages 1885–1894 . PMLR , 2017.342 [ 15 ] Artemy Kolchinsky , Brendan D Tracey , and Steven Van Kuyk . Caveats for information bottleneck in343 deterministic scenarios . arXiv preprint arXiv:1808.07593 , 2018.344 [ 16 ] Alex Krizhevsky . Learning multiple layers of features from tiny images . Technique Report , 2009.345 [ 17 ] Y. LeCun , BE . Boser , and JS . Denker . Handwritten digit recognition with a back-propagation network . In346 NeurIPS , pages 396–494 , 1990.347 [ 18 ] Yann LeCun , Yoshua Bengio , and Geoffrey Hinton . Deep learning . nature , 521 ( 7553 ) :436–444 , 2015.348 [ 19 ] Yann LeCun , Sumit Chopra , Raia Hadsell , Marc ’ Aurelio Ranzato , and Fu Jie Huang . A tutorial on349 energy-based learning . MIT Press , 2006.350 [ 20 ] Jaehoon Lee , Yasaman Bahri , Roman Novak , Samuel S. Schoenholz , Jeffrey Pennington , and Jascha351 Sohl-Dickstein . Deep neural networks as gaussian processes . In ICLR , 2018.352 [ 21 ] Aravindh Mahendran and Andrea Vedaldi . Understanding deep image representations by inverting them.353 In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 5188–5196,354 2015.355 [ 22 ] David McAllester and Karl Stratos . Formal limitations on the measurement of mutual information . In356 International Conference on Artificial Intelligence and Statistics , pages 875–884 . PMLR , 2020.357 [ 23 ] Julian D Olden and Donald A Jackson . Illuminating the “ black box ” : a randomization approach for358 understanding variable contributions in artificial neural networks . Ecological modelling , 154 ( 1-2 ) :135–150,359 2002.360 [ 24 ] Liam Paninski . Estimation of entropy and mutual information . Neural computation , 15 ( 6 ) :1191–1253,361 2003.362 [ 25 ] David E. Rumelhart , Geoffrey E. Hinton , and Ronald J. Williams . Learning representations by back-363 propagating errors . Nature , 323:533–536 , October 1986.364 [ 26 ] Ruslan Salakhutdinov and Geoffrey Hinton . Deep boltzmann machines . In Artificial intelligence and365 statistics , pages 448–455 . PMLR , 2009.366 [ 27 ] Wojciech Samek , Thomas Wiegand , and Klaus-Robert Müller . Explainable artificial intelligence : Under-367 standing , visualizing and interpreting deep learning models . arXiv preprint arXiv:1708.08296 , 2017.368 [ 28 ] Andrew Saxe , Yamini Bansal , Joel Dapello , Madhu Advani , Artemy Kolchinsky , Brendan Tracey , and369 David Cox . On the information bottleneck theory of deep learning . In International Conference on370 Representation Learning , 2018.371 [ 29 ] Ohad Shamir , Sivan Sabato , and Naftali Tishby . Learning and generalization with the information372 bottleneck . Theoretical Computer Science , 411 ( 29-30 ) :2696–2711 , 2010.373 [ 30 ] Ravid Shwartz-Ziv and Naftali Tishby . Opening the black box of deep neural networks via information.374 arXiv preprint arXiv:1703.00810 , 2017.375 [ 31 ] Noam Slonim . The information bottleneck : Theory and applications . PhD thesis , Citeseer , 2002.376 [ 32 ] Naftali Tishby , Fernando C Pereira , and William Bialek . The information bottleneck method . arXiv377 preprint physics/0004057 , 2000.378 [ 33 ] Naftali Tishby and Noga Zaslavsky . Deep learning and the information bottleneck principle . In 2015 IEEE379 Information Theory Workshop ( ITW ) , pages 1–5 . IEEE , 2015.380 [ 34 ] Vladimir N Vapnik . An overview of statistical learning theory . IEEE transactions on neural networks,381 10 ( 5 ) :988–999 , 1999.382 [ 35 ] Larry Wasserman . All of nonparametric statistics . Springer Science & Business Media , 2006.383 [ 36 ] Han Xiao , Kashif Rasul , and Roland Vollgraf . Fashion-mnist : a novel image dataset for benchmarking384 machine learning algorithms . arXiv preprint arXiv:1708.07747 , 2017.385 [ 37 ] Matthew D Zeiler and Rob Fergus . Visualizing and understanding convolutional networks . In European386 conference on computer vision , pages 818–833 . Springer , 2014.387 [ 38 ] Chiyuan Zhang , Samy Bengio , Moritz Hardt , Benjamin Recht , and Oriol Vinyals . Understanding deep388 learning requires rethinking generalization . arXiv preprint arXiv:1611.03530 , 2016.389 Checklist390 1 . For all authors ... 391 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s contribu-392 tions and scope ? [ Yes ] 393 ( b ) Did you describe the limitations of your work ? [ Yes ] see Section 5394 ( c ) Did you discuss any potential negative societal impacts of your work ? [ N/A ] 395 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to them ? [ Yes ] 396 2 . If you are including theoretical results ... 397 ( a ) Did you state the full set of assumptions of all theoretical results ? [ Yes ] 398 ( b ) Did you include complete proofs of all theoretical results ? [ Yes ] 399 3 . If you ran experiments ... 400 ( a ) Did you include the code , data , and instructions needed to reproduce the main experimental401 results ( either in the supplemental material or as a URL ) ? [ Yes ] see the URL in Appendix G402 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they were chosen ) ? 403 [ Yes ] see Section 4.1 , Appendix G , and Appendix H404 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running experiments405 multiple times ) ? [ Yes ] see Section 4406 ( d ) Did you include the total amount of compute and the type of resources used ( e.g. , type of GPUs,407 internal cluster , or cloud provider ) ? [ Yes ] see Appendix G408 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 409 ( a ) If your work uses existing assets , did you cite the creators ? [ Yes ] 410 ( b ) Did you mention the license of the assets ? [ Yes ] 411 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ Yes ] see412 Appendix H413 ( d ) Did you discuss whether and how consent was obtained from people whose data you ’ re us-414 ing/curating ? [ N/A ] 415 ( e ) Did you discuss whether the data you are using/curating contains personally identifiable informa-416 tion or offensive content ? [ N/A ] 417 5 . If you used crowdsourcing or conducted research with human subjects ... 418 ( a ) Did you include the full text of instructions given to participants and screenshots , if applicable ? 419 [ N/A ] 420 ( b ) Did you describe any potential participant risks , with links to Institutional Review Board ( IRB ) 421 approvals , if applicable ? [ N/A ] 422 ( c ) Did you include the estimated hourly wage paid to participants and the total amount spent on423 participant compensation ? [ N/A ] 424

Summary: The paper points out that the IB principle based on the Markov chain assumption ($Y \rightarrow X \rightarrow T_1 \rightarrow \cdots \rightarrow \hat{Y}$) cannot fully characterize the information flow in DNNs, and hence several claims previously made on the assumption may not be accurate. For better quantification of the mutual information, this paper proposes a probabilistic framework of IB in deep neural networks (DNNs) by taking into account the back-propagation training. Under the probabilistic framework, the paper derives two Markov chains that characterize the information flow in DNNs and provides experimental evidence for supporting the validity of the Markov chain assumptions. The paper claims that the probabilistic approaches allow us to more accurately estimate the mutual information than existing non-parametric approaches and experimentally demonstrates how different activation functions and hidden layers achieve IB trade-offs.