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: The paper studies the distributional shift setting, in which test and training distribution are different.  The distribution shift is described in terms of mutual information of the data with a binary selection variable $t$ (train or test). Several propositions describe relations between the distribution of test and train covariates and labels and a latent space. Based on those observations, several regularization criteria are suggested. The experiments compare various settings of distributional shifts and popular methods.