Usable Information and Evolution of Optimal Representations During Training

1 INTRODUCTION . An important open question for the theory of deep learning is why highly over-parametrized neural networks learn solutions that generalize well even though the model can in principle memorize the entire training set . Some have speculated that neural networks learn minimal but sufficient representations of the input through implicit regularization of Stochastic Gradient Descent ( SGD ) ( Shwartz-Ziv & Tishby , 2017 ; Achille & Soatto , 2018 ) , and that the minimality of the representations relates to generalizability . Follow-up work has disputed the validity of some of these claims when using deterministic deep networks ( Saxe et al. , 2018 ) , leading to an ongoing debate on the notion of optimality of representations and how they are learned during training . Part of the disagreement stems from the use of information-theoretic quantities : most previous studies in deep learning have analyzed the amount of information that the learned representation contains about the inputs using Shannon ’ s mutual information . However , when the mapping from input to representation is deterministic , the mutual information between the representation and input is degenerate ( Saxe et al. , 2018 ; Goldfeld et al. , 2018 ) . Rather than study the mutual information in a neural network , here we instead define and study the “ usable information ” in the network , which measures the amount of information that can be extracted from the representation by a learned decoder , and is scalable to high dimensional realistic tasks . We use this notion to quantify how relevant and irrelevant information is represented across layers of the network throughout the training process , and how this is affected by the optimization algorithms and the network pretraining . In particular , we propose to study a simple task inspired by decision-making tasks in neuroscience , where inputs and outputs are carefully designed to probe specific information processing phenomena . We then extend our findings to standard image classification tasks trained with state-of-the-art models . Our neuroscience-inspired task is the checkerboard ( CB ) task ( Chandrasekaran et al. , 2017 ; Kleinman et al. , 2019 ) . In the CB task , one discerns the dominant color of a checkerboard filled with red and green squares . The subject then makes a reach to a left or right target whose color matches the dominant color in the checkerboard ( Fig 1a ) . This task therefore involves making two binary choices : a color decision ( i.e. , reach to the red or green target ) and a direction decision ( i.e. , reach to left or right ) . Critically , the color of the targets ( red left , green right ; or green left , red right ) is random on every trial . The direction decision output is conditionally independent of the color decision , as detailed further in Fig 1b and Section B.6 , even though the color information needs to be used to solve the task . This task allows us to evaluate how both of these components of information are represented through training and across layers . We used this task and extensions to study the evolution of minimal representations during training . If a representation is sufficient and minimal , we refer to this representation as optimal ( Achille & Soatto , 2018 ) . Our contributions are the following . ( 1 ) We introduce a notion of usable information for studying representations and training dynamics in deep networks ( Section 3 ) . ( 2 ) We used this notion to characterize the transient training dynamics in deep networks by studying the amount of usable relevant and irrelevant information in deep network layers and across training epochs . We first use the CB task to gain intuition of the training dynamics in a simplified setting . We find that training with SGD is critical to bias the network toward learning minimal representations in intermediate layers ( Section 4.1 ) . This adds to the literature suggesting that SGD results in minimal representations of input information ( Achille & Soatto , 2018 ; Shwartz-Ziv & Tishby , 2017 ) while avoiding some of the pitfalls . ( 3 ) We used the intuition gained from the simple task , evaluating our findings on CIFAR-10 and CIFAR-100 task using modern architectures . Remarkably , we find that the networks increased usable information about an irrelevant component of information early in training and discarded it later on in training to arrive at a minimal sufficient solution , consistent with a proposed ( Shwartz-Ziv & Tishby , 2017 ) though controversial theory ( Saxe et al. , 2018 ) . 2 RELATED WORK . Some efforts to understand why neural networks generalize focus on representation learning , that is , how deep networks learn optimal ( i.e. , minimal and sufficient ) representations of inputs in order to solve a task . Typically , representation learning is focused on studying the properties of the asymptotic representations after training ( Achille & Soatto , 2018 ) . Recent work suggests that these asymptotic representations contain minimal but sufficient input information for performing a task ( Achille & Soatto , 2018 ; Shwartz-Ziv & Tishby , 2017 ) . Implicit regularization coming from SGD , and in particular from the use of large learning rates and small batch sizes , is believed to play an important role in forming these minimal sufficient representations . How does the training process lead to these minimal but sufficient asymptotic representations ? Shwartz-Ziv & Tishby ( 2017 ) propose that there are two distinct phases of training : an empirical risk minimization phase where the network minimizes the loss on the training set , and a “ compression ” phase where the network discards information about the inputs that do not need to be represented to solve the task . Recently , Saxe et al . ( 2018 ) challenged this view , arguing that the observed compression was dependent on the activation function and the mutual information estimator used in Shwartz-Ziv & Tishby ( 2017 ) . These works highlight the challenges of estimating mutual information to study how representations emerge through training . In general , estimating mutual information from samples is challenging for high-dimensional random variables ( Paninski , 2003 ) . The primary difficulty in estimating mutual information is estimating a high-dimensional probability distribution from the samples , since generally the number of samples required scales exponentially with the dimension . This is impractical for realistic deep learning tasks where the representations are high dimensional . To estimate the mutual information , Shwartz-Ziv & Tishby ( 2017 ) used a binning approach , discretizing the activations into a finite number of bins . While this approximation is exact in the limit of infinitesimally small bins , in practice , the size of the bin affects the estimator ( Saxe et al. , 2018 ; Goldfeld et al. , 2018 ) . In contrast to binning , other approaches to estimate mutual information include entropic-based estimators ( e.g. , Goldfeld et al . ( 2018 ) ) and a nearest neighbours approach ( Kraskov et al. , 2004 ) . Although mutual information is difficult to estimate , it is an appealing quantity to summarily characterize key aspects of the transient neural network training behavior because of its invariance to smooth and invertible transformations . In this work , rather than estimate the mutual information directly , we instead define and study the “ usable information ” in the network , which corresponds to a variational approximation of the mutual information ( Barber & Agakov , 2003 ; Poole et al. , 2019 ) ( see Sections 3 and A.1 ) . Recently , such variational approximations to mutual information have been viewed as a meaningful characterization of representations in deep networks , and the theoretical underpinnings of this approach are beginning to be investigated ( Xu et al. , 2020 ; Dubois et al. , 2020 ) . Research into the training dynamics of deep networks , and how they represent relevant and irrelevant task information , is nascent . A related study by Achille et al . ( 2019 ) found that early periods of training were critical for determining the asymptotic network behavior . Additionally , it was found that the timing of regularization was important for determining asymptotic performance ( Golatkar et al. , 2019 ) , with regularization during this “ critical period ” having the most influential effect . Notably , both of these studies found an initial increase in the amount of information that weights encode about the dataset ( as measured by the Fisher information ) , that coincides with the critical period of learning . This phase is followed later in training by a “ forgetting ” phase where the network discards unnecessary information . This suggests that a similar dynamic to the one we study can be observed in weight space instead of representation space . 3 USABLE INFORMATION IN A REPRESENTATION . A deep neural network consists of a set of ` layers , with each layer forming a successive representation of the input . A representation Z ` may store information in a variety of ways . It may be that a complex transformation is required to read out the information , or it may be that a simple linear decoder could read out the information . In both cases , from an information-theoretic perspective , the same information is contained in the representation , however , there is an important distinction regarding how “ usable ” this information is . Information is usable if later layers , which comprise affine transformations and element-wise nonlinearities , can easily extract it to solve the task . Equivalently , usable information should be decodable by a separate neural network also employing affine transformations and element-wise nonlinearities . Formally , we define the usable information that a representation Z contains about a quantity Y , which may refer to the output or a component of the input , as : Iu ( Z ; Y ) = H ( Y ) − LCE ( p ( y|z ) , q ( y|z ) ) . ( 1 ) Here , H ( Y ) is the entropy , or uncertainty , of Y , and LCE is the cross-entropy loss on the test set of a discriminator network q ( y|z ) trained to approximate the true distribution p ( y|z ) . Our definition is motivated in the following manner . The test set cross-entropy loss approximates how much uncertainty there is in the output Y given Z and the discriminator . A low loss implies that there is low uncertainty in Y given Z , or that the discriminator can extract a lot of “ information ” about Y from Z . If the logarithm in the cross-entropy loss is in base 2 , it is measured in bits . If the value of Y were approximately the same for any Z , there would be little uncertainty in Y to begin with , so it is important to know the amount of uncertainty in Y given Z with respect to the initial uncertainty in Y . What is most relevant is the amount of remaining uncertainty in Y given Z . Thus we use the difference in uncertainty H ( Y ) −LCE as the amount of “ usable information ” that Z contains about Y , as shown in our definition in Equation 1 . This definition is appealing to study representations , in part , because it can be computed from samples of Z and Y , and is a quantity that is comparable across network training . We estimate LCE using a small neural network that learns a distribution q ( y|z ) . To train the network , we sample activations Z and the quantity Y and learn q ( y|z ) by minimizing the cross-entropy loss on a training set . We then evaluate the LCE on the test set ( Equation 1 ) . We provide details about the neural network and the training we used for decoding in Appendix B.3 and C.2 . We also show in the Appendix that the usable information is a lower bound on the mutual information ( Appendix A.1 ) . Importantly , usable information also is not constrained by the data processing inequality ; that is , the information can be made more “ usable ” by transformation to later layers , consistent with the representation learning view that later layers are forming improved representations of the inputs ( Xu et al. , 2020 ) .

Summary: Broadly, this work is an attempt to understand how neural networks can form generalizable representations while being severely overparameterized. This work proposes an information theoretic measure, called the "usable information", and use it to quantify the amount of relevant information in different layers of a neural network during training. The key idea is that, in order for the information represented in one layer to be "usable" by the next layer, it should be decodable by a simple transformation (affine + element-wise nonlinearity).