Robust Generalization despite Distribution Shift via Minimum Discriminating Information

1 Introduction . Developing machine learning-based systems for real world applications is challenging , particularly because the conditions under which the system was trained are rarely the same as when using the system . Unfortunately , a standard assumption in most machine learning methods is that test and training distribution are the same [ 78 , 59 , 12 ] . This assumption , however , rarely holds in practice , and the performance of many models suffers in light of this issue , often called dataset shift [ 52 ] or equivalently distribution shift . Consider building a model for diagnosing a specific heart disease , and suppose that most participants of the study are middle to high-aged men . Further suppose these participants have a higher risk for the specific disease , and as such do not reflect the general population with respect to age and gender . Consequently , the training data suffers from the so-called sample selection bias inducing a covariate shift [ 62 , 52 ] . Many other reasons lead to distribution shifts , such as non-stationary environments [ 67 ] , imbalanced data [ 52 ] , domain shifts [ 3 ] , label shifts [ 83 ] or observed contextual information [ 8 , 9 ] . A specific type of distribution shift takes center stage in off-policy evaluation ( OPE ) problems . Here , one is concerned with the task of estimating the resulting cost of an evaluation policy for a sequential decision making problem based on historical data obtained from a different policy known as behavioral policy [ 73 ] . This problem is of critical importance in various applications of reinforcement learning—particularly , when it is impossible or unethical to evaluate the resulting cost of an evaluation policy by running it on the underlying system . Solving a learning problem facing an arbitrary and unknown distribution shift based on training data in general is hopeless . Oftentimes , fortunately , partial knowledge about the distribution shift is available . In the medical example above , we might have prior information how the demographic attributes in our sample differ from the general population . Given a training distribution and partial knowledge about the shifted test distribution , one might ask what is the “ most natural '' distribution 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . shift mapping the training distribution into a test distribution consistent with the available structural information . Here , we address this question , interpreting “ most natural '' as maximizing the underlying Shannon entropy . This concept has attracted significant interest in the past in its general form , called principle of minimum discriminating information dating back to Kullback [ 37 ] , which can be seen as a generalization of Jaynes ’ maximum entropy principle [ 31 ] . While these principles are widely used in tasks ranging from economics [ 27 ] to systems biology [ 63 ] and regularized Markov decision processes [ 48 , 25 , 2 ] , they have not been investigated to model general distribution shifts as we consider in this paper . Irrespective of the underlying distribution shift , the training distribution of any learning problem is rarely known , and one typically just has access to finitely many training samples . It is well-known that models can display a poor out-of-sample performance if training data is sparse . These overfitting effects are commonly avoided via regularization [ 12 ] . A regularization technique that has become popular in machine learning during the last decade and provably avoids overfitting is distributionally robust optimization ( DRO ) [ 36 ] . Contributions . We highlight the following main contributions of this paper : • We introduce a new modelling framework for distribution shifts via the principle of minimum discriminating information , which encodes prior structural information on the resulting test distribution . • Using our framework and the available training samples , we provide generalization bounds via a DRO program and prove that the introduced DRO model is optimal in a precise statistical sense . • We show that the optimization problems characterizing the distribution shift and the DRO program can be efficiently solved by exploiting convex duality and recent accelerated first order methods . • We demonstrate the versatility of the proposed Minimum Discriminating based DRO ( MDI-DRO ) method on two distinct problem classes : Training classifiers on systematically biased data and the OPE for Markov decision processes . In both problems MDI-DRO outperforms existing approaches . 2 Related work . For supervised learning problems , there is a rich literature in the context of covariate shift adaptation [ 62 , 68 ] . A common approach is to address this distribution shift via importance sampling , more precisely by weighting the training loss with the ratio of the test and training densities and then minimize the so-called importance weighted risk ( IWERM ) , see [ 62 , 82 , 68 , 69 ] . While this importance weighted empirical risk is an unbiased estimator of the test risk , the method has two major limitations : It tends to produce an estimator with high variance , making the resulting test risk large . Further , the ratio of the training and test densities must be estimated which in general is difficult as the test distribution is unknown . There are modifications of IWERM reducing the resulting variance [ 15 , 13 , 65 ] , for example by exponentially flattening the importance ratios [ 62 ] . For the estimation of the importance weights several methods have been presented , see for example [ 81 ] . These methods , however crucially rely on having data from both training and test distribution . Liu and Ziebart [ 40 ] and Chen et al . [ 14 ] propose a minimax approach for regression problems under covariate shift . Similar to our approach taken in this paper , they consider a DRO framework , which however , optimizes over so-called moment-based ambiguity sets . Distribution shifts play a key role in causal inference . In particular , the connection between causal predictors and distributional robustness under shifts arising from interventions has been widely studied [ 58 , 42 , 56 , 66 ] . Oftentimes , a causal graph is used to represent knowledge about the underlying distribution shift induced by an intervention [ 49 , 50 ] . Distribution shifts have been addressed in a variety of different settings [ 35 ] , we refer the reader to the comprehensive textbook [ 52 ] and references therein . There is a vast literature on OPE methods which we will not attempt to summarize . In a nutshell , OPE methods can be grouped into three classes : a first class of approaches that aims to fit a model from the available data and uses this model then to estimate the performance of the given evaluation policy [ 41 , 1 , 38 ] . A second class of methods are based on invoking the idea of importance sampling to model the underlying distribution shift from behavioral to evaluation policy [ 51 , 29 , 74 ] . The third , more recent , class of methods combines the first two classes [ 24 , 32 , 76 , 33 ] . Key reasons for the popularity of DRO in machine learning are the ability of DRO models to regularize learning problems [ 36 , 60 , 61 ] and the fact that the underlying optimization problems can often be exactly reformulated as finite convex programs solvable in polynomial time [ 4 , 10 ] . Such reformulations hold for a variety of ambiguity sets such as : regions defined by moments [ 20 , 26 , 80 , 11 ] , φ-divergences [ 5 , 44 , 39 ] , Wasserstein ambiguity sets [ 43 , 36 ] , or maximum mean discrepancy ambiguity sets [ 64 , 34 ] . DRO naturally seems a convenient tool when analyzing “ small '' distribution shifts as it seeks models that perform well “ sufficiently close '' to the training sample . However , modelling a general distribution shift via DRO seems difficult , and recent interest has focused on special cases such as adversarial example shifts [ 23 ] or label shifts [ 83 ] . To the best of our knowledge , the idea of combining DRO with the principle of minimum discriminating information is new . 3 Problem statement and motivating examples . We study learning problems of the form min θ∈Θ R ( θ , P ? ) , ( 1 ) where R ( θ , P ? ) = EP ? [ L ( θ , ξ ) ] denotes the risk of an uncertain real-valued loss function L ( θ , ξ ) that depends on a parameter θ ∈ Θ ⊂ Rn to be estimated as well as a random vector ξ ∈ Ξ ⊂ Rm governed by the probability distribution P ? . In order to avoid technicalities , we assume from now on that Θ and Ξ are compact and L is continuous . In statistical learning , it is usually assumed that P ? is unknown but that we have access to independent samples from P ? . This paper departs from this standard scenario by assuming that there is a distribution shift . We first state our formal assumption about the shift and provide concrete examples below . Specifically , we assume to have access to samples from a distribution P 6= P ? and that P ? is only known to belong to the distribution family Π = { Q ∈ P ( Ξ ) : EQ [ ψ ( ξ ) ] ∈ E } ( 2 ) encoded by a measurable feature map ψ : Ξ→ Rd and a compact convex set E ⊂ Rd . In view of the principle of minimum discriminating information , we identify P ? with the I-projection of P onto Π . Definition 3.1 ( Information projection ) . The I-projection of P ∈ P ( Ξ ) onto Π is defined as Pf = f ( P ) = arg min Q∈Π D ( Q‖P ) , ( 3 ) where D ( Q‖P ) denotes the relative entropy of Q with respect to P. One can show that the I-projection exists whenever Π is closed with respect to the topology induced by the total variation distance [ 17 , Theorem 2.1 ] . As E is closed , this is the case whenever ψ is bounded . Note that f ( P ) = P if P ∈ Π . In the remainder , we assume that P 6∈ Π and that P is only indirectly observable through independent training samples ξ̂1 , . . . , ξ̂N drawn from P. Example 3.2 ( Logistic regression ) . Assume that ξ = ( x , y ) , where x ∈ Rm−1 is a feature vector of patient data ( e.g. , a patient ’ s age , sex , chest pain type , blood pressure , etc . ) , and y ∈ { −1 , 1 } a label indicating the occurrence of a heart disease . Logistic regression models the conditional distribution of y given x by a logistic function Prob ( y|x ) = [ 1 + exp ( −y · θ > x ) ] −1 parametrized by θ ∈ Rm−1 . The maximum likelihood estimator for θ is found by minimizing the empirical average of the logistic loss function L ( θ , ξ ) = log ( 1 + exp ( −y · θ > x ) ) on the training samples . If the samples pertain to a patient cohort , where elderly males are overrepresented with respect to the general population , then they are drawn from a training distribution P that differs from the test distribution P ? . Even if sampling from P ? is impossible , we may know that the expected age of a random individual in the population falls between 40 and 45 years . This information can be modeled as EP ? [ ψ ( ξ ) ] ∈ E , where E = [ ` , u ] , ` = 40 , u = 45 and ψ ( ξ ) projects ξ to its ‘ age ’ -component . Other available prior information can be encoded similarly . Inspired by the principle of minimum discriminating information , we then minimize the expected log-loss under the I-projection Pf of the data-generating distribution P onto the set Π defined in ( 2 ) . Example 3.3 ( Production planning ) . Assume that θ ∈ R and ξ ∈ R denote the production quantity and the demand of a perishable good , respectively , and that the loss function L ( θ , ξ ) represents the sum of the production cost and a penalty for unsatisfied demand . To find the optimal production quantity , one could minimize the average loss in view of training samples drawn from the historical demand distribution P. However , a disruptive event such as the beginning of a recession might signal that demand will decline by at least η % . The future demand distribution P ? thus differs from P and belongs to a set Π of the form ( 2 ) defined through ψ ( ξ ) = ξ and E = [ 0 , ( 1− η ) µ ] , where µ denotes the historical average demand . By the principle of minimum discriminating information it then makes again sense to minimize the expected loss under the I-projection Pf of P onto Π . Loosely speaking , the principle of minimum discriminating information identifies the I-projection Pf of P as the least prejudiced and thus most natural model for P ? in view of the information that P ? ∈ Π . The principle of minimum discriminating information is formally justified by the conditional limit theorem [ 18 ] , which we paraphrase below using our notation . Proposition 3.4 ( Conditional limit theorem ) . If the interior of the compact convex set E overlaps with the support of the pushforward measure P ◦ ψ−1 , the I-projection Pf = f ( P ) exists and the moment-generating function EPf [ etL ( θ , ξ ) ] is finite for all t in a neighborhood of 0 , then we have lim N→∞ EPN [ L ( θ , ξ1 ) | 1N ∑N i=1 ψ ( ξi ) ∈ E ] = EPf [ L ( θ , ξ ) ] ∀θ ∈ Θ . In the context of Examples 3.2 and 3.3 , the conditional limit theorem provides an intuitive justification for modeling distribution shifts via I-projections . More generally , the following proposition suggests that any distribution shift can be explained as an I-projection onto a suitably chosen set Π . Proposition 3.5 ( Every distribution is an I-projection ) . If P , Q ∈ P ( Ξ ) are such that Q is absolutely continuous with respect to P and if Π is a set of the form ( 2 ) defined through ψ ( ξ ) = log dQdP ( ξ ) and E = { D ( Q‖P ) } , then Q = f ( P ) . The modelling of arbitrary distribution shifts via the I-projection according to Proposition 3.5 has an interesting application in the off-policy evaluation problem for Markov decision processes ( MDPs ) . Example 3.6 ( Off-policy evaluation ) . Consider an MDP ( S , A , Q , c , s0 ) with finite state and action spaces S andA , respectively , transition kernel Q : S ×A → R , cost-per-stage function c : S ×A → R and initial state s0 . A stationary Markov policy π is a stochastic kernel that maps states to probability distributions over A . We use π ( a|s ) to denote the probability of selecting action a in state s under policy π . The long-run average cost generated by π can be expressed as Vπ = limT→∞ 1 T ∑T−1 t=0 Eπs0 [ c ( st , at ) ] . Each policy induces an occupation measure µπ on S×A defined through the state-action frequencies µπ ( x , a ) = limT→∞ 1 T ∑T−1 t=0 Pπs0 [ ( st , at ) = ( s , a ) ] ∀s ∈ S , a ∈ A , see [ 28 , Chapter 6 ] . One can additionally show that µπ belongs to the polytope M = { µ ∈ ∆S×A : ∑ a′∈A µ ( s ′ , a′ ) − ∑ s∈S ∑ a∈AQ ( s ′|s , a ) µ ( s , a ) = 0 ∀s′ ∈ S } , where ∆S×A represents the simplex of all probability mass functions over S ×A . Conversely , each occupation measure µ ∈M induces a policy πµ defined through πµ ( a|s ) = µ ( s , a ) / ∑ a′∈A µ ( s , a ′ ) for all s ∈ S and a ∈ A . Assuming that all parameters of the MDP except for the cost c are known , the off-policy evaluation problem asks for an estimate of the long-run average cost Vπe of an evaluation policy πe based on a trajectory of states , actions and costs generated by a behavioral policy πb . This task can be interpreted as a degenerate learning problem without a parameter θ to optimize if we define ξ = c ( s , a ) and set L ( θ , ξ ) = ξ . Here , a distribution shift emerges because we must evaluate the expectation of ξ under Q = µe ◦ c−1 given training samples from P = µb ◦ c−1 , where µb and µe represent the occupation measures corresponding to πb and πe , respectively . Note that P and Q are unknown because c is unknown . Moreover , as the policy πe generates different stateaction trajectories than πb , the costs generated under πe can not be inferred from the costs generated under πb even though πb and πe are known . Note also that Q coincides with the I-projection Pf of P onto the set Π defined in Proposition 3.5 . The corresponding feature map ψ as well as the set E can be computed without knowledge of c provided that c is invertible . Indeed , in this case we have ψ ( ξi ) = log dµe◦c−1 dµb◦c−1 ( ξi ) = log µe ( si , ai ) µb ( si , ai ) and E = { D ( µe ◦ c−1‖µb ◦ c−1 ) } for any si ∈ S , ai ∈ A and ξi = c ( si , ai ) . Note that as S and A are finite , c is generically invertible , that is , c can always be rendered invertible by an arbitrarily small perturbation . In summary , we may conclude that the off-policy evaluation problem reduces to an instance of ( 1 ) . Given N training samples ξ̂1 , . . . , ξ̂N , we henceforth use P̂N = 1N ∑N i=1 δξ̂i and P̂ f N to denote the empirical distribution on and its I-projection onto Π , respectively . As the true data-generating distribution P and its I-projection Pf are unknown , it makes sense to replace them by their empirical counterparts . However , the resulting empirical risk minimization problem is susceptible to overfitting if the number of training samples is small relative to the feature dimension . In order to combat overfitting , we propose to solve the DRO problem J ? N = min θ∈Θ R ? ( θ , P̂fN ) , ( 4 ) which minimizes the worst-case risk over all distributions close to P̂fN . Here , R ? is defined through R ? ( θ , P′ ) = supQ∈Π { R ( θ , Q ) : D ( P′‖Q ) ≤ r } ( 5 ) and thus evaluates the worst-case risk of a given parameter θ ∈ Θ in view of all distributions Q that have a relative entropy distance of at most r from a given nominal distribution P′ ∈ Π . In the remainder we use J ? N and θ ? N to denote the minimum and a minimizer of problem ( 4 ) , respectively . Main results . The main theoretical results of this paper can be summarized as follows . ( i ) Out-of-sample guarantee . We show that the optimal value of the DRO problem ( 4 ) provides an upper confidence bound on the risk of its optimal solution θ ? N . Specifically , we prove that PN ( R ( θ ? N , Pf ) > J ? N ) ≤ e−rN+o ( N ) , ( 6 ) where Pf = f ( P ) is the I-projection of P. If Ξ is finite , then ( 6 ) can be strengthened to a finite sample bound that holds for every N if the right hand side is replaced with e−rN ( N + 1 ) |Ξ| . ( ii ) Statistical efficiency . In a sense to be made precise below , the DRO problem ( 4 ) provides the least conservative approximation for ( 1 ) whose solution satisfies the out-of-sample guarantee ( 6 ) . ( iii ) Computational tractability . We prove that the I-projection P̂fN can be computed via a regularized fast gradient method whenever one can efficiently project onto E. Given P̂fN , we then show that θ ? N can be found by solving a tractable convex program whenever Θ is a convex and conic representable set , while L ( θ , ξ ) is a convex and conic representable function of θ for any fixed ξ .

Summary: In this paper, they aim at test and training distribution shift problem where partial knowledge of shifted test distribution is known. With the principle of minimum discriminating information, they introduce a new modeling framework catching prior structural information, and a specific MDI-DRO with theoretical generalization bound, satisfying the statistical guarantee of out-of-sample.  For computational tractability, they leverage convex duality and accelerated first-order methods to efficiently solve the DRO program. Experiments on two distinct problem classes prove the availability of the proposed approach.  ======== After reading the authors' response and other reviews, I would like to raise my score to 5, but still think the paper in its current form is not self-contained due to the lack of ablation study and important baselines.