[Paper Review] Robust Generalization despite Distribution Shift via Minimum Discriminating Information

Tobias Sutter, Andreas Krause+1 more|arXiv (Cornell University)|Jun 8, 2021

Machine Learning and Algorithms4 citations

TL;DR

This paper proposes a distributionally robust optimization framework using minimum discriminating information to model distribution shifts when partial structural knowledge of the test distribution is available. By combining principles of minimum discriminating information with distributionally robust optimization and large deviation bounds, the method achieves provably robust generalization under distribution shift, outperforming existing approaches in both biased classification and off-policy evaluation tasks.

ABSTRACT

Training models that perform well under distribution shifts is a central challenge in machine learning. In this paper, we introduce a modeling framework where, in addition to training data, we have partial structural knowledge of the shifted test distribution. We employ the principle of minimum discriminating information to embed the available prior knowledge, and use distributionally robust optimization to account for uncertainty due to the limited samples. By leveraging large deviation results, we obtain explicit generalization bounds with respect to the unknown shifted distribution. Lastly, we demonstrate the versatility of our framework by demonstrating it on two rather distinct applications: (1) training classifiers on systematically biased data and (2) off-policy evaluation in Markov Decision Processes.

Motivation & Objective

• To address the challenge of model performance degradation due to distribution shift in real-world machine learning applications.
• To develop a principled method for modeling unknown test distributions using partial structural knowledge and the principle of minimum discriminating information.
• To provide explicit generalization bounds under distribution shift using large deviation theory.
• To demonstrate the framework's effectiveness in two distinct applications: biased data classification and off-policy evaluation in Markov decision processes.
• To ensure statistical optimality of the resulting distributionally robust optimization model through theoretical guarantees.

Proposed method

• The framework models the unknown shifted test distribution by minimizing the Kullback-Leibler divergence from the training distribution, subject to constraints encoding prior structural knowledge about the shift.
• It employs distributionally robust optimization (DRO) to minimize the worst-case expected loss over an ambiguity set defined by the minimum discriminating information principle.
• Generalization bounds are derived using large deviation results, ensuring performance guarantees under the unknown shifted distribution.
• The resulting optimization problems are solved efficiently via convex duality and accelerated first-order methods.
• The method is applied to two distinct settings: (1) training classifiers on systematically biased data with known demographic shifts, and (2) off-policy evaluation in Markov decision processes with known behavioral policy constraints.
• The framework is implemented as Minimum Discriminating Information-based DRO (MDI-DRO), which is shown to be statistically optimal and computationally efficient.

Experimental results

Research questions

• RQ1 How can we model an unknown shifted test distribution when only partial structural information about the shift is available?
• RQ2 Can we derive explicit generalization bounds for machine learning models under distribution shift using the principle of minimum discriminating information?
• RQ3 How does the proposed MDI-DRO framework compare to existing methods in terms of robustness and generalization performance on biased data and off-policy evaluation tasks?
• RQ4 What is the statistical optimality of the resulting DRO model in the presence of distribution shift and limited training samples?
• RQ5 Can the MDI-DRO framework be efficiently solved in practice, and what are the computational guarantees?

Key findings

• The MDI-DRO framework achieves provably robust generalization by combining minimum discriminating information with distributionally robust optimization, ensuring performance guarantees under unknown distribution shifts.
• The method provides explicit generalization bounds derived from large deviation theory, which hold for any unknown shifted distribution consistent with the prior structural constraints.
• In the synthetic biased classification task with m=6, MDI-DRO outperformed existing methods in terms of test accuracy and robustness across 1000 independent experiments.
• For off-policy evaluation in the inventory control model, MDI-DRO demonstrated superior performance compared to baseline methods, particularly in estimating long-run average costs under distribution shift.
• The optimization problems in the framework are efficiently solvable via convex duality and accelerated first-order methods, enabling practical deployment.
• The DRO predictor is continuous and well-behaved, as proven via Berge’s maximum theorem and continuity of the objective function.