The effectiveness of layer-by-layer training using the information bottleneck principle

1 INTRODUCTION . Deep neural nets ( DNN ) have shown unparalleled success in many fields , yet the theoretical understanding of these models is lagging behind . Among the various attempts to shed light on this field , Tishby & Zaslavsky ( 2015 ) have recently suggested a link between DNN training and the Information Bottleneck ( IB ) principle of Tishby et al . ( 1999 ) . Their theory claims that optimal performance is attained if each layer simultaneously attempts to maximize its mutual information ( MI ) with the target space , while minimizing its MI with the input space . The logic is that the DNN layers should compress the latent representation so as to allow for good generalization , while retaining only the information essential for target prediction . Moreover , it was hypothesized that training DNNs through minimization of the popular cross-entropy loss , implicitly achieves this goal . More formally , it was suggested that during training , each DNN layer approaches the maximal value of the IB Langragian LIB = I ( Y ; Li ) − βI ( X ; Li ) , ( 1 ) where I denotes MI , X and Y are the input and target random variables , Li is the latent representation at the output of the ith hidden layer , and β is a trade-off parameter . This elegant connection between DNN training and the IB principle was later supported by empirical experiments conducted by Shwartz-Ziv & Tishby ( 2017 ) , as well as by others ( Alemi et al. , 2018 ; Belghazi et al. , 2018 ; Kolchinsky et al. , 2017 ; Saxe et al. , 2018 ) . These works either ( i ) showed the effect of optimizing the IB functional in ( 1 ) for classification tasks , or ( ii ) analyzed the dynamics of each layer in the information plane ( defined by axes I ( X ; L ) , I ( Y ; L ) ) during training with the cross-entropy loss . And indeed , these works have shown such a scheme is beneficial in terms of classification accuracy , and that the IB functional tends to increase during training with the crossentropy loss . Yet in all these demonstrations , nets were trained in an end-to-end fashion , where the IB objective was either enforced only on a single layer Li , or not at all ( trained only with a crossentropy loss ) . Therefore , it remains unclear whether directly optimizing the IB functional of each layer , as suggested by Tishby & Zaslavsky ( 2015 ) , would indeed yield useful representations . The goal of this paper is to experimentally examine this hypothesis . In addition to the lack of explicit validation , criticism on the rationale behind the IB theory of deep learning has surfaced lately . One major concern highlighted in the recent works of Saxe et al . ( 2018 ) and Amjad & Geiger ( 2018 ) , is that for deterministic DNNs with continuous inputs , the term I ( X ; Li ) is always infinite so that the IB principle is in fact meaningless . The authors suggest to resolve this by resorting to stochastic DNNs via the introduction of noise after each layer , which ensures that I ( X ; Li ) is finite . And indeed , all attempts to train DNNs with the IB objective thus far employed stochastic nets ( Alemi et al. , 2018 ; Belghazi et al. , 2018 ; Kolchinsky et al. , 2017 ) . Yet , deterministic DNNs are far more popular for most tasks , raising questions on the applicability of the IB theory for these models . In this paper , we focus on validating ( a regularized version of ) the IB theory to deterministic DNNs trained on high-dimensional data . To alleviate the explosion of MI with the input space , we introduce noise only for quantifying the MI between the input X and hidden layer Li ( as also suggested by Saxe et al . ( 2018 ) ) , and not into the DNN itself . Namely , we replace the first term in ( 1 ) by I ( X ; Li + ε ) , where ε is noise . As we show , the resulting term can be interpreted as a weight decay penalty , which aligns with common practice in DNN training . To estimate MI , we use an auxiliary net , similarly to the suggestion by Belghazi et al . ( 2018 ) , and in contrast to applying a variational lower-bound as in ( Alemi et al. , 2018 ; Kolchinsky et al. , 2017 ) . However , as opposed to the original mutual information neural estimator ( MINE ) of Belghazi et al . ( 2018 ) , which fails to accurately estimate I ( X ; Li + ε ) in our setting1 , we tailor the architecture of the auxiliary net specifically for our case in which Li is a known deterministic function of X . This modification results in significantly more accurate estimates , as we show in Section 4 . Equipped with this estimation strategy , we directly test the validity of the IB theory of deep learning for deterministic models . This is accomplished by training a DNN while guaranteeing that each layer optimizes the IB functional in ( 1 ) . Starting from the first hidden layer , each layer is trained independently , and frozen before moving on to the next layer . We find that such a training process is not only useful , but surprisingly consistently leads to slightly better performance than end-toend training with the cross-entropy loss . This is thus the first explicit confirmation of the theory suggested by Tishby & Zaslavsky ( 2015 ) , subject to adapting the definition of I ( X ; L ) so as to make this quantity meaningful . Please note that whether layer-wise training with the IB principle is generally beneficial over other layer-wise training schemes , or over other regularized end-to-end schemes , is out of this paper ’ s scope . Accompanying the original theory were several key observations , the most notable being the emergence of two-phase training dynamics ( Shwartz-Ziv & Tishby , 2017 ) . It was hypothesized that the second “ compression ” phase is responsible for enhancing the generalization capability . Yet Saxe et al . ( 2018 ) contradict this claim , showing that this phase may not exist when training with a crossentropy loss . Our experiments align with those of Saxe et al . ( 2018 ) , though we do observe these two phases when training explicitly with the IB functional . Surprisingly , when training with crossentropy and weight decay regularization , the two-phase dynamics are also apparent . This suggests that this two-phase pattern is induced by regularization , and is thus indeed linked to generalization . 2 RELATED WORK . Deep neural nets and The Information Bottleneck The Information Bottleneck ( IB ) technique for summarizing a random variable X while maintaining maximal mutual information with a desirable output Y , was first introduced by Tishby et al . ( 1999 ) . It is designed to find the optimal trade-off between prediction accuracy and representation complexity , by compressing X while retaining the essential information for predicting Y . The IB optimization problem can be directly solved in discrete settings , as well as in certain simple families of continuous distributions , like jointly Gaussian random variables ( Chechik et al. , 2005 ) . However , in general high-dimensional continuous scenarios , exact optimization becomes impractical due to the intractable continuous MI . The IB principle was recently linked to DNNs by Tishby & Zaslavsky ( 2015 ) , which hypothesized that DNN layers converge to the optimal IB curve during training . The subsequent work by Shwartz- 1For example , it does not tend to infinity when the noise is taken to zero . Ziv & Tishby ( 2017 ) attempted to experimentally verify this by analyzing DNNs trained with the cross-entropy loss ( with no regularization ) . In this work , the MI estimation was made tractable by binning ( discretization ) of the latent representations Li , which works for “ toy ” examples but does not scale to real-world scenarios . Among several claims , the authors report that two training phases emerge : a “ fitting ” phase followed by a “ compression ” phase , which to their understanding is linked to increased generalization . Yet Saxe et al . ( 2018 ) contradict these claims , by linking the compression phase to the activation type and discretization strategy , and questioning the connection between this compression phase and generalization . Moreover , these authors as well as Amjad & Geiger ( 2018 ) , also recognize that the term I ( X ; L ) in the IB functional is theoretically infinite for deterministic DNNs with a continuous input X , and thus the attempt to measure it is meaningless . Both works propose to remedy this by adding noise , which ensures this term is finite . Despite these obstacles , DNNs were in fact trained on real high-dimensional data for classification tasks with the IB functional . The difficulties arising from the intractable and infinite term I ( X ; L ) were overcome by ( i ) training stochastic DNNs which ensure it is finite , and ( ii ) using a variational approximation of the MI which makes it tractable ( Alemi et al. , 2018 ; Kolchinsky et al. , 2017 ; Chalk et al. , 2016 ; Achille & Soatto , 2018 ; Belghazi et al. , 2018 ) . These schemes all rely on some form of injected stochasticity , and in fact , most enforce the IB objective only on a single “ bottleneck layer ” . An attempt to optimize the IB functional for deterministic DNNs in a layer-wise fashion , as in the original theory , has yet to appear due to these practical difficulties . Here , we report on such a training scheme and present its results in the following sections . We note that layer-wise training with a related information-theoretic objective has been studied by Ver Steeg & Galstyan ( 2015 ) and Gao et al . ( 2018 ) , which have shown an effective and practical method to compose and analyze hierarchical representations . These works consider an unsupervised setting , whereas we analyze supervised DNN training ( specifically ) with the IB functional . Estimating mutual information Quantifying the MI between distributions is inherently difficult ( Paninski , 2003 ) , and is tractable only in discrete settings or for a limited family of problems . In other more general settings the exact computation is impossible , and known approximations do not scale well with dimension and sample size ( Gao et al. , 2015 ) . Most recently , Belghazi et al . ( 2018 ) proposed the MI Neural Estimator ( MINE ) for approximating the MI between continuous high-dimensional random variables via back-prop over a DNN . The core idea is to estimate the Kullback-Leibler ( KL ) divergence that is used to define MI , through the maximization of the dual representation of Donsker & Varadhan ( 1983 ) . Minimizing the MI between the input X and hidden layer Li of a DNN using this neural MI estimator can be efficiently accomplished , by formulating a minmax objective between the examined DNN and the estimator ( an auxiliary net ) , similar to adversarial training ( Goodfellow et al. , 2014 ) ( see Belghazi et al . ( 2018 ) ) . We use this strategy to enforce the IB objective on DNN layers during training . Note that the original authors also demonstrate the IB principle with this estimator , however they only do so on a single “ bottleneck layer ” , by using the cross-entropy loss as an approximation for the MI with the desired output Y , and in an end-to-end manner ( as described above ) .

Summary: This work is about layer-wise training of networks by way of optimizing the IB cost function, which basically measures the compression of the inputs under the constraint that some degree of information with respect to the targets must be preserved. Both terms of the IB cost function are formalized as mutual informations, but since in neural nets, the latent "compression" is a deterministic function of the inputs, a severe technical problems arises: the joint distribution between p-dimensional inputs X and the q-dimensional latent compression L is degenerate in that  its support lies in a space of dimension p (and not p+q as it would be in the non-degenerate case). As a consequence, no p.d.f. exists (with respect to the Lebesgue measure of R^{p+q}). Thus, defining mutual information is cumbersome. The paper attempts to overcome this problem by using a noisy version of the latent compression, i.e. L' = L + \epsilon, which can be seen as an "ad hoc" fix of this problem. Not too surprising, this additive noise works as a ridge-type (or weight-decay) regularizer, just as a Gaussian prior in regression.