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: This paper proposes a distributionally robust optimization (DRO) framework for learning models that can generalize well under distribution shifts, when partial knowledge about the nature of the shift is available. As opposed to classical domain adaptation methods, the proposed method does not require data from the target distribution. For an uncertainty set of distributions defined to satisfy a set of moment conditions (the partial knowledge about the shift), the authors propose to train a model to minimize risk with respect to the distribution in the uncertainty set which minimizes the relative entropy from the training distribution. Using the empirical distribution of the training data as the source distribution, the authors show that the solution to this minimax problem upper bounds the optimal risk and provide finite sample guarantees. The authors then develop a two step learning procedure: First, they estimate the I-projection (using a fast gradient-like method, Algorithm 1). Then, they can efficiently compute the model parameters. They also consider simple extensions to the case of off policy evaluation in reinforcement learning. Synthetic and real data experiments demonstrate the effectiveness of the proposed method given a good choice of the relative entropy radius hyperparameter.