Towards optimally abstaining from prediction with OOD test examples

TL;DR

This paper introduces an optimal abstention algorithm for machine learning models that accounts for out-of-distribution test examples, providing guarantees on prediction loss that adapt to distribution shifts and adversarial inputs.

Contribution

It presents a transductive abstention algorithm with loss guarantees that incorporate distributional distance, applicable to linear regression and binary classification, extending prior work.

Findings

Loss bounds match standard generalization when data is i.i.d.

Algorithm efficiently handles distribution shifts and adversarial examples.

Polynomial-time algorithm for linear regression based on optimization methods.

Abstract

A common challenge across all areas of machine learning is that training data is not distributed like test data, due to natural shifts, "blind spots," or adversarial examples; such test examples are referred to as out-of-distribution (OOD) test examples. We consider a model where one may abstain from predicting, at a fixed cost. In particular, our transductive abstention algorithm takes labeled training examples and unlabeled test examples as input, and provides predictions with optimal prediction loss guarantees. The loss bounds match standard generalization bounds when test examples are i.i.d. from the training distribution, but add an additional term that is the cost of abstaining times the statistical distance between the train and test distribution (or the fraction of adversarial examples). For linear regression, we give a polynomial-time algorithm based on Celis-Dennis-Tapia optimization algorithms. For binary classification, we show how to efficiently implement it using a proper agnostic learner (i.e., an Empirical Risk Minimizer) for the class of interest. Our work builds on a recent abstention algorithm of Goldwasser, Kalais, and Montasser (2020) for transductive binary classification.