Phase Transitions for the Information Bottleneck in Representation Learning

1 INTRODUCTION . The Information Bottleneck ( IB ) objective ( Tishby et al. , 2000 ) : IBβ [ p ( z|x ) ] : = I ( X ; Z ) − βI ( Y ; Z ) ( 1 ) explicitly trades off model compression ( I ( X ; Z ) , I ( · ; · ) denoting mutual information ) with predictive performance ( I ( Y ; Z ) ) using the Lagrange multiplier β , where X , Y are observed random variables , and Z is a learned representation of X . The IB method has proved effective in a variety of scenarios , including improving the robustness against adversarial attacks ( Alemi et al. , 2016 ; Fischer , 2018 ) , learning invariant and disentangled representations ( Achille & Soatto , 2018a ; b ) , underlying information-based geometric clustering ( Strouse & Schwab , 2017b ) , improving the training and performance in adversarial learning ( Peng et al. , 2018 ) , and facilitating skill discovery ( Sharma et al. , 2019 ) and learning goal-conditioned policy ( Goyal et al. , 2019 ) in reinforcement learning . From Eq . ( 1 ) we see that when β → 0 it will encourage I ( X ; Z ) = 0 which leads to a trivial representation Z that is independent of X , while when β → +∞ , it reduces to a maximum likelihood objective1 that does not constrain the information flow . Between these two extremes , how will the IB objective behave ? Will prediction and compression performance change smoothly , or do there exist interesting transitions in between ? In Wu et al . ( 2019 ) , the authors observe and study the learnability transition , i.e . the β value such that the IB objective transitions from a trivial global minimum to learning a nontrivial representation . They also show how this first phase transition relates to the structure of the dataset . However , to answer the full question , we need to consider the full range of β . 1For example , in classification , it reduces to cross-entropy loss . Motivation . To get a sense of how I ( Y ; Z ) and I ( X ; Z ) vary with β , we train Variational Information Bottleneck ( VIB ) models ( Alemi et al. , 2016 ) on the CIFAR10 dataset ( Krizhevsky & Hinton , 2009 ) , where each experiment is at a different β and random initialization of the model . Fig . 1 shows the I ( X ; Z ) , I ( Y ; Z ) and accuracy vs. β , as well as I ( Y ; Z ) vs . I ( X ; Z ) for CIFAR10 with 20 % label noise ( see Appendix I for details ) . From Fig . 1 ( b ) ( c ) , we see that as we increase β , instead of going up smoothly , both I ( X ; Z ) and I ( Y ; Z ) show multiple phase transitions , where the slopes dI ( X ; Z ) dβ and dI ( Y ; Z ) dβ are discontinuous and the accuracy has discrete jumps . The observation lets us refine our question : When do the phase transitions occur , and how do they depend on the structure of the dataset ? These questions are important , since answering them will help us gain a better understanding of the IB objective and its close interplay with the dataset and the learned representation . Moreover , the IB objective belongs to a general form of two-term trade-offs in many machine learning objectives : L = Prediction-loss + β · Complexity , where the complexity term generally takes the form of regularization . Usually , learning is set at a specific β . Many more insights can be gained if we understand the behavior of the prediction loss and model complexity with varying β , and how they depend on the dataset . The techniques developed to address the question in the IB setting may also help us understand the two-term tradeoff in other learning objectives . Contributions . In this work , we begin to address the above question in IB settings . Specifically : • We identify a qualitative change of the IB loss landscape w.r.t . p ( z|x ) for varying β as IB phase transitions ( Section 3 ) . • Based on the definition , we introduce a quantity G [ p ( z|x ) ] and use it to prove a theorem giving a practical condition for IB phase transitions . We further reveal the connection between G [ p ( z|x ) ] and the Fisher information matrix when p ( z|x ) is parameterized by θ ( Section 3 ) . • We reveal the close interplay between the IB objective , the dataset and the learned represen- tation , by showing that in IB , each phase transition corresponds to learning a new nonlinear component of maximum correlation between X and Y , orthogonal to the previously-learned Z , and each with decreasing strength ( Section 4 ) . To the best of our knowledge , our work provides the first theoretical formula to address IB phase transitions in the most general setting . In addition , we present an algorithm for iteratively finding the IB phase transition points ( Section 5 ) . We show that our theory and algorithm give tight matches with the observed phase transitions in categorical datasets , predict the onset of learning new classes and class difficulty in MNIST , and predict prominent transitions in CIFAR10 experiments ( Section 6 ) . 2 RELATED WORK . The Information Bottleneck Method ( Tishby et al. , 2000 ) provides a tabular method based on the Blahut-Arimoto ( BA ) Algorithm ( Blahut , 1972 ) to numerically solve the IB functional for the optimal encoder distribution P ( Z|X ) , given the trade-off parameter β and the cardinality of the representation variable Z . This work has been extended in a variety of directions , including to the case where all three variables X , Y , Z are multivariate Gaussians ( Chechik et al. , 2005 ) , cases of variational bounds on the IB and related functionals for amortized learning ( Alemi et al. , 2016 ; Achille & Soatto , 2018a ; Fischer , 2018 ) , and a more generalized interpretation of the constraint on model complexity as a Kolmogorov Structure Function ( Achille et al. , 2018 ) . Previous theoretical analyses of IB include Rey & Roth ( 2012 ) , which looks at IB through the lens of copula functions , and Shamir et al . ( 2010 ) , which starts to tackle the question of how to bound generalization with IB . We will make practical use of the original IB algorithm , as well as the amortized bounds of the Variational Informormation Bottleneck ( Alemi et al. , 2016 ) and the Conditional Entropy Bottleneck ( Fischer , 2018 ) . Phase transitions , where key quantities change discontinuously with varying relative strength in the two-term trade-off , have been observed in many different learning domains , for multiple learning objectives . In Rezende & Viola ( 2018 ) , the authors observe phase transitions in the latent representation of β-VAE for varying β. Strouse & Schwab ( 2017b ) utilize the kink angle of the phase transitions in the Deterministic Information Bottleneck ( DIB ) ( Strouse & Schwab , 2017a ) to determine the optimal number of clusters for geometric clustering . Tegmark & Wu ( 2019 ) explicitly considers critical points in binary classification tasks using a discrete information bottleneck with a non-convex Pareto-optimal frontier . In Achille & Soatto ( 2018a ) , the authors observe a transition on the tradeoff of I ( θ ; X , Y ) vs. H ( Y |X , θ ) in InfoDropout . Under IB settings , Chechik et al . ( 2005 ) study the Gaussian Information Bottleneck , and analytically solve the critical values βci = 1 1−λi , where λi are eigenvalues of the matrix Σx|yΣ−1x , and Σx is the covariance matrix . This work provides valuable insights for IB , but is limited to the special case that X , Y and Z are jointly Gaussian . Phase transitions in the general IB setting have also been observed , which Tishby ( 2018 ) describes as “ information bifurcation ” . In Wu et al . ( 2019 ) , the authors study the first phase transition , i.e . the learnability phase transition , and provide insights on how the learnability depends on the dataset . Our work is the first work that addresses all the IB phase transitions in the most general setting , and provides theoretical insights on the interplay between the IB objective , its phase transitions , the dataset , and the learned representation . 3 FORMULA FOR IB PHASE TRANSITIONS . 3.1 DEFINITIONS . Let X ∈ X , Y ∈ Y , Z ∈ Z be random variables denoting the input , target and representation , respectively , having a joint probability distribution p ( X , Y , Z ) , with X × Y × Z its support . X , Y and Z satisfy the Markov chain Z−X−Y , i.e . Y and Z are conditionally independent given X . We assume that the integral ( or summing if X , Y or Z are discrete random variables ) is on X × Y × Z . We use x , y and z to denote the instances of the respective random variables . The above settings are used throughout the paper . We can view the IB objective IBβ [ p ( z|x ) ] ( Eq . 1 ) as a functional of the encoding distribution p ( z|x ) . To prepare for the introduction of IB phase transitions , we first define relative perturbation function and second variation , as follows . Definition 1 . Relative perturbation function : For p ( z|x ) , its relative perturbation function r ( z|x ) is a bounded function that maps X × Z to R and satisfies Ez∼p ( z|x ) [ r ( z|x ) ] = 0 . Formally , define QZ|X : = { r ( z|x ) : X × Z → R ∣∣ Ez∼p ( z|x ) [ r ( z|x ) ] = 0 , and ∃M > 0 s.t . ∀X ∈ X , Z ∈ Z , |r ( z|x ) | ≤ M } . We have that r ( z|x ) ∈ QZ|X iff r ( z|x ) is a relative perturbation function of p ( z|x ) . The perturbed probability ( density ) is p′ ( z|x ) = p ( z|x ) ( 1 + · r ( z|x ) ) for some > 0 . Definition 2 . Second variation : Let functional F [ f ( x ) ] be defined on some normed linear space R. Let us add a perturbative function · h ( x ) to f ( x ) , and now the functional F [ f ( x ) + · h ( x ) ] can be expanded as ∆F [ f ( x ) ] = F [ f ( x ) + · h ( x ) ] − F [ f ( x ) ] = ϕ1 [ · h ( x ) ] + ϕ2 [ · h ( x ) ] + ϕr [ · h ( x ) ] ‖ · h ( x ) ‖2 such that lim →0 ϕr [ · h ( x ) ] = 0 , where ‖·‖ denotes the norm , ϕ1 [ · h ( x ) ] = dF [ f ( x ) ] d is a linear functional of ·h ( x ) , and is called the first variation , denoted as δF [ f ( x ) ] . ϕ2 [ ·h ( x ) ] = 12 2 d 2F [ f ( x ) ] d 2 is a quadratic functional of · h ( x ) , and is called the second variation , denoted as δ2F [ f ( x ) ] . We can think of the perturbation function · h ( x ) as an infinite-dimensional “ vector ” ( x being the indices ) , with being its amplitude and h ( x ) its direction . With the above preparations , we define the IB phase transition as a change in the local curvature on the global minimum of IBβ [ p ( z|x ) ] . Definition 3 . IB phase transitions : Let r ( z|x ) ∈ QZ|X be a perturbation function of p ( z|x ) , p∗β ( z|x ) denote the optimal solution of IBβ [ p ( z|x ) ] at β , where the IB functional IB [ · ] is defined in Eq . ( 1 ) . The IB phase transitions βci are the β values satisfying the following two conditions : ( 1 ) ∀r ( z|x ) ∈ QZ|X , δ2IBβ [ p ( z|x ) ] ∣∣ p∗β ( z|x ) ≥ 0 ; ( 2 ) lim β′→β+ inf r ( z|x ) ∈QZ|X δ2IBβ′ [ p ( z|x ) ] ∣∣ p∗β ( z|x ) = 0− . Here β+ and 0− denote one-sided limits . We can understand the δ2IBβ [ p ( z|x ) ] as a local “ curvature ” of the IB objective IBβ ( Eq . 1 ) w.r.t . p ( z|x ) , along some relative perturbation r ( z|x ) . A phase transition occurs when the convexity of IBβ [ p ( z|x ) ] w.r.t . p ( z|x ) changes from a minimum to a saddle point in the neighborhood of its optimal solution p∗β ( z|x ) as β increases from βc to βc + 0+ . This means that there exists a perturbation to go downhill and find a better minimum . We validate this definition empirically below .

Summary: This paper studies the phase transition problem in the information bottleneck (IB) objective and derives a formula for IB phase transitions. Based on the theory developed in the paper, an algorithm is developed to find phase transition points. The interesting observation is phase transition can correspond to learning new class and this paper conjectures in IB for classification, the number of phase transitions is at most C-1, C the number of classes.  This observation deserves to be further explored and may be a key to a deeper understanding of neural networks.