An Information-theoretic Approach to Distribution Shifts

1 Introduction . One of the most common assumptions for machine learning models is that the training and test data are independently and identically sampled ( IID ) from the same distribution . In practice , this assumption does not hold in many practical scenarios ( Bengio et al. , 2020 ) . A machine learning model trained to recognize land usage from satellite images using pictures from the early 2000s may struggle to recognize the style of modern architectures ( Christie et al. , 2018 ) , data collected on a limited set of hospitals may not be representative of the variation introduced by the use of different machines or procedures ( Zech et al. , 2018 ; Beede et al. , 2020 ) . Other kinds of distribution shifts are more subtle and difficult to recognize despite having a noticeable impact on the model ’ s predictive performance . Examples include under-represented or over-represented population groups ( Popejoy & Fullerton , 2016 ; Buolamwini & Gebru , 2018 ) or biased annotations collected from crowd-sourcing services ( Zhang et al. , 2017 ; Xia et al. , 2020 ) . Different approaches in literature address these issues by using some external source of knowledge such as domain or environment annotations ( Wang & Deng , 2018 ) , protected attributes ( Mehrabi et al. , 2019 ) or sub-population groups ( Santurkar et al. , 2020 ) to reduce bias and minimize the model error outside of the training distribution . Despite the progress in the field of domain generalization literature , Gulrajani & Lopez-Paz ( 2020 ) has shown that the effectiveness of some of the most common algorithms heavily relies on the hyper-parameter tuning strategy , revealing limitations of 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . examined models when compared to a more traditional empirical risk minimization strategy . One of the major issues behind this observed behavior is a lack of clarity and applicability for the underlying assumptions regarding the problem statement and the data-generating process ( Zhao et al. , 2019 ; Rosenfeld et al. , 2020 ; Mahajan et al. , 2020 ) . With this work , we aim to present a new perspective on the problem to analyze and clarify the fundamental differences between some of the most common approaches by : 1 . Introducing a novel information-theoretical framework to describe the problem of distribution shift and connecting it to the test error and its components ( section 2 ) . 2 . Analyzing four main families of objectives and describing some of their guarantees and assumptions ( section 3 ) . 3 . Demonstrating that the effectiveness of different criteria is determined by the structure of the underlying data-generating process ( section 4 ) . 4 . Showing that the results obtained by popular models designed according to the aforementioned criteria can drastically differ from the theoretically expected performance ( section 4 ) . The analysis of the bottleneck , independence , sufficiency , and separation criteria reveals that some of the most popular models can systematically fail to reduce the test error even for simple datasets . No unique objective is simultaneously optimal for every problem , but additional knowledge about the selection procedure and better approximations can help to mitigate the bias . 1.1 Problem Statement . Consider x and y as the features and targets respectively with joint density p ( x , y ) for the predictive problem of interest . Let t be a binary random variable representing which data is selected for training ( t = 1 ) and which is not ( t = 0 ) . We will refer to p ( x , y|t = 1 ) as the joint Training distribution and p ( x , y|t = 0 ) as the Test distribution that are induced by the selection t. Throughout the paper , we will consider selections t that can be expressed as a function of the features x , targets y , independent noise and other variables e , which represent other factors that are affecting the data collection procedure ( e.g . geographic location , time intervals , population groups ) . Let q ( y|x ) represent a learnable model for the predictive distribution p ( y|x ) . We use the KullbackLeibler divergence to express the Train and Test error respectively1 : DKL ( p t=1 y|x||qy|x ) ︸ ︷︷ ︸ Train error : = DKL ( p ( y|x , t = 1 ) ||q ( y|x ) ) , ( 1 ) DKL ( p t=0 y|x||qy|x ) ︸ ︷︷ ︸ Test error : = DKL ( p ( y|x , t = 0 ) ||q ( y|x ) ) . ( 2 ) Although only the train distribution can be accessed at training time , we are interested in learning models q ( y|x ) that result in small test error . 1.2 Characterizing Distribution Shift . As a first step , we characterize the difference between training and test distribution as a function of the selection . The effect that the selection has in the joint distribution p ( x , y ) can be quantified by considering the mutual information between the features-target pair xy and the selection variable t , which can be also expressed as a Kullback-Leibler divergence : I ( xy ; t ) = DKL ( p ( x , y|t ) |p ( x , y ) ) . This measure of distribution shift intuitively represents how many bits the selection variable carries about the joint distribution of targets and features or , equivalently , how much the joint density p ( x , y ) has changed as a result of the selection . Using the chain rule of mutual information , one can express distribution shift as the sum of two separate components : I ( xy ; t ) ︸ ︷︷ ︸ Distribution shift = I ( x ; t ) ︸ ︷︷ ︸ Covariate shift + I ( y ; t|x ) ︸ ︷︷ ︸ Concept shift , ( 3 ) 1Further details regarding the convention used for conditional KL-divergence can be found in appendix A which can be interpreted as the amount of covariate shift ( Shimodaira , 2000 ) and concept shift ( Widmer & Kubat , 1996 ; Moreno-Torres et al. , 2012 ) , respectively . These two quantities refer to the changes in the predictive distribution p ( y|x ) and the marginal features distribution p ( x ) respectively , which add up to represent the changes in the joint distribution . Whenever the data selection is perfectly IID , the selection variable can be expressed as a function of some independent noise ( t = f ( ) ) and the corresponding distribution shift I ( xy ; t ) is zero . On the other hand , if the data collection procedure has been influenced by other factors , we do not have such a guarantee , even when the selection does not depend on features and targets directly ( t = f ( e , ) ) . 2 Information-theoretic Framework . Using the quantities described in the previous sections , we can show that the sum train and test error is lower-bounded by the amount concept shift : Proposition 1 . For any model q ( y|x ) and α : = min { p ( t = 0 ) , p ( t = 1 ) } : DKL ( p t=1 y|x||qy|x ) +DKL ( p t=0 y|x||qy|x ) ≥ 1 1− α I ( y ; t|x ) , ( 4 ) As a consequence , whenever the selection induces concept shift I ( y ; t|x ) > 0 , any sufficiently flexible model q̂ ( y|x ) trained using a maximum likelihood approach must incur in strictly positive test error . Intuitively , whenever the selection induces a change in the predictive distribution ( p ( y|x ) 6= p ( y|x , t = 1 ) ) fitting the model to the training distribution will incorporate the selection bias into the model prediction , necessarily resulting in errors when evaluated on the test distribution , as shown in the example reported in figure 1a . The different approaches analyzed in this work are based in the introduction of a latent representation z , which allows for the definition of different regularization strategies . Using the latent representation as an intermediate variable , the model q ( y|x ) can be re-parametrized with an encoder q ( z|x ) and a latent classifier2 q ( y|z ) : q ( y|x ) = Ez∼q ( z|x ) [ q ( y|z = z ) ] . ( 5 ) The effect of the introduction of a latent representation can be observed by expressing the training and test error as a function of the encoder and the classifier3 . 2The name classifier will be used for convenience to express both classification and regression problems . 3The expressions in equations 6 and 7 hold with equality for deterministic encoders as shown in appendix B . Proposition 2 . For any encoder q ( z|x ) and classifier q ( y|z ) : DKL ( p t=1 y|x||qy|x ) ≤ It=1 ( x ; y|z ) ︸ ︷︷ ︸ Train information loss +DKL ( p t=1 y|z ||qy|z ) ︸ ︷︷ ︸ Latent train error ( 6 ) DKL ( p t=0 y|x||qy|x ) ≤ It=0 ( x ; y|z ) ︸ ︷︷ ︸ Test information loss +DKL ( p t=0 y|z ||qy|z ) ︸ ︷︷ ︸ Latent test error . ( 7 ) The two terms It=1 ( x ; y|z ) and It=0 ( x ; y|z ) represent the amount of predictive information that is lost by encoding the features x into z on train and test distribution respectively , whileDKL ( pt=1y|z ||qy|z ) and DKL ( pt=0y|z ||qy|z ) refer to the train and test error when using z instead of the original observations as the predictive features , which will be referred to as latent training error and latent test error respectively . Test information loss and latent test error capture two intrinsically different kinds of error . The former indicates the increase in the prediction uncertainty as a result of the encoding procedure , while the latter represents the discrepancy between the model q ( y|z ) and the latent test predictive distribution p ( y|z , t = 0 ) . 2.1 Latent test error . Minimizing the latent test error makes the latent predictor q ( y|z ) approach the test latent predictive distribution p ( y|z , t = 0 ) . We can show that the Jensen-Shannon divergence between the two distributions is upper-bounded by a monotonic function of the latent training error and the amount of concept shift in the latent space ( latent concept shift I ( y ; t|z ) ) : Proposition 3 . For any q ( z|x ) , q ( y|z ) and any representation z that satisfies p ( z = z|t = 0 ) > 0 and p ( z = z|t = 1 ) > 0 : ( √ 1 2α I ( y ; t|z = z ) + √ 1 2 DKL ( p t=1 y|z=z||qy|z=z ) ) 2 ≥ DJSD ( pt=0y|z=z||qy|z=z ) . ( 8 ) Whenever the Jensen-Shannon divergence between the test predictive distribution p ( y|z , t = 0 ) and the classifier q ( y|z ) is small , the latent test error ( measured in terms of KL-divergence ) must also be small at least for the regions that have positive probability according to both train p ( x|t = 1 ) and test p ( x|t = 0 ) data distributions . Since the train predictive distribution p ( y|x , t = 1 ) is not defined for x that have zero probability on the train distribution , we have no guarantees regarding the model predictions in those regions unless other inductive biases are considered . The left hand side of the expression in proposition 3 can be minimized by addressing I ( y ; t|z ) and DKL ( p t=1 y|z ||qy|z ) with respect to the encoder and classifier respectively . In other words , we can minimize the latent test error by simultaneously fitting q ( y|z ) to the train predictive distribution p ( y|z , t = 1 ) and minimizing the latent concept shift induced by the encoder q ( z|x ) . When the latent concept shift and latent test error approach zero , the latent train predictive distribution p ( y|z , t = 1 ) approaches the true ( unselected ) predictive distribution p ( y|z ) , and so does the modeled classifier q ( y|z ) . In this ideal scenario , the only source of test error is the additional uncertainty that is due to the information lost in the encoding procedure .

Summary: This paper investigates distribution shift from an information-theoretic perspective. It outlines a theoretical framework for describing different sources of modeling error, including that due to distribution shift, in the presence of a latent representation. It uses the framework to categorise existing approaches to reducing distribution shift according to the type of regularization criteria they use. It explains the assumptions underlying the success of these approaches, and demonstrates empirically, on synthetic data, that when these assumptions fail to hold, performance drops.