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 contributes theoretically to the information bottleneck (IB) principle. In particular, the author(s) provided theoretical reasoning on the phase transition phenomenon: when the beta parameter of IB varies, the generalization performance changes in a stepwise manner rather than continuously. The core result is given by theorem 1: the phase transition betas necessarily satisfy an equation, where the LHS is expressed in terms of an optimal perturbation of the encoding function X->Z.