Source-linked AI summary
Detecting and Correcting for Label Shift with Black Box Predictors
Zachary C. Lipton, Yu-Xiang Wang, Alex Smola
TL;DR
The paper addresses how to detect, quantify, and correct distribution shift without target labels when label frequencies change but p(x|y) remains stable. It introduces BBSE, which uses black-box predictors and invertible confusion matrices to estimate target label ratios, establishes theoretical guarantees, and applies the estimates to detection and correction. The method is designed to remain useful with biased or uncalibrated predictors and in high-dimensional settings, subject to its assumptions and practical handling of degenerate confusion matrices.
Problem
The paper asks how to detect, quantify, and correct distribution shift without target labels under label shift, where p(y) changes but p(x|y) does not.
Method
BBSE uses a black-box predictor’s training confusion matrix and average target predictions to estimate target-to-source label ratios through a linear system.
Results
BBSE has consistency and error bounds, supports statistical shift detection and importance-weighted classifier correction, and is empirically validated on high-dimensional data.
Takeaways & Limitations
BBSE can exploit arbitrary predictors, including biased, inaccurate, or uncalibrated models, while better predictors provide tighter estimates and lower sample complexity.
Takeaways & Limitations
BBSE requires an invertible confusion matrix; degenerate cases may require soft predictions or merging classes.
Abstract
from arXiv · showhide
Faced with distribution shift between training and test set, we wish to detect and quantify the shift, and to correct our classifiers without test set labels. Motivated by medical diagnosis, where diseases (targets) cause symptoms (observations), we focus on label shift, where the label marginal $p(y)$ changes but the conditional $p(x| y)$ does not. We propose Black Box Shift Estimation (BBSE) to estimate the test distribution $p(y)$. BBSE exploits arbitrary black box predictors to reduce dimensionality prior to shift correction. While better predictors give tighter estimates, BBSE works even when predictors are biased, inaccurate, or uncalibrated, so long as their confusion matrices are invertible. We prove BBSE's consistency, bound its error, and introduce a statistical test that uses BBSE to detect shift. We also leverage BBSE to correct classifiers. Experiments demonstrate accurate estimates and improved prediction, even on high-dimensional datasets of natural images.
1. Introduction
The paper frames label shift as a setting where target label frequencies change while class-conditional features remain stable, and introduces BBSE to estimate and correct for this shift using black-box predictors.
- Problem motivation: Unlabeled target data make shift detection, quantification, and correction impossible without assumptions about how source and target distributions are related.The paper focuses on label shift as one tractable assumption.
- Problem motivation: Label shift changes p(y) while preserving p(x|y), unlike covariate shift, which preserves p(y|x).The setting is motivated by diagnosis, where diseases cause observed symptoms.
- BBSE: BBSE estimates target-to-source label ratios using a black-box predictor and an invertible training confusion matrix.It solves a linear system using the predictor’s confusion matrix and average test predictions.
- Limitations: Degenerate confusion matrices may require soft predictions or merging classes before applying the black-box method.The paper identifies these as practical variants when ordinary confusion matrices are unusable.
- Contributions: BBSE supports consistency and error bounds, statistical tests for label-shift detection, and importance-weighted empirical-risk correction.These are presented as the paper’s main theoretical, diagnostic, corrective, and empirical contributions.
- BBSE: BBSE can use biased, inaccurate, or uncalibrated predictors, while better predictors reduce sample complexity and accuracy does not depend on data dimensionality.Its central requirement is an invertible confusion matrix.
2. Prior Work
Prior work studied related distribution-shift and sample-selection problems, but the paper positions BBSE as extending label-shift estimation to high-dimensional modern machine-learning settings with theoretical guarantees.
- Label-shift estimation: Earlier label-shift methods based on Bayesian inference and EM require explicitly estimating p(x|y), which may be infeasible for high-dimensional data.These methods motivate reducing the estimation problem through a predictor.
- Covariate shift: Covariate-shift research corrects models by estimating or applying importance weights based on q(x)/p(x), including kernel mean matching.This literature addresses a different invariant conditional assumption.
- Related sampling problems: Econometric work on non-random, choice-based, and endogenous samples studies problems analogous to shifts in label distributions.The cited work concerns estimating behavior or parameters under non-representative sampling.
- Related connections: Propensity scoring connects the paper to methods for designing unbiased experiments, while cognitive-science work links classification to surrounding items.These connections broaden the related-work context beyond machine learning.
- Positioning BBSE: Prior epidemiological estimators addressed related problems, but the paper distinguishes BBSE through theoretical guarantees and application to massively higher-dimensional x than y.The paper presents dimensionality reduction as especially valuable in this modern ML setting.
3. Problem setup
The problem setup considers domain adaptation from labeled source data and unlabeled target data, under label shift, support coverage, and an invertible confusion-matrix condition.
- Domain adaptation: Domain adaptation provides labeled source samples and unlabeled target samples, with the objective of predicting well on the target distribution.Source data are drawn from P and target data from Q.
- Assumptions: Without shared support, adapting from source to target is generally impossible, so the paper imposes three additional assumptions.These assumptions define the scope of the proposed estimation problem.
- Label shift: Under label shift, the target distribution factorizes as q(y,x)=q(y)p(x|y), with the class-conditional feature distribution shared across domains.This is also called target shift and is framed as anti-causal learning.
- Assumptions: Target-label support must be contained in source-label support, meaning the training data include examples from every target class.For discrete labels, this is the identifiability requirement described by A.2.
- Assumptions: The predictor’s expected confusion matrix must be invertible, with linearly independent expected outputs across classes.This condition can be checked empirically, and it extends naturally to soft classifiers.
- Estimation objective: BBSE estimates w(y)=q(y)/p(y), which becomes the importance weight in weighted empirical-risk minimization under label shift.The target label distribution is otherwise unobserved because target data are unlabeled.
4. Main results
BBSE estimates label-shift weights from predictor outputs by combining source confusion statistics with unlabeled target predictions. Under the stated assumptions, the estimators are consistent and admit finite-sample error bounds whose quality depends on sample sizes, shift magnitude, and predictor distinguishability.
- Estimation: BBSE estimates source-to-target label ratios by solving a linear system built from the predictor’s source confusion matrix and unlabeled target predictions.The method uses source data to estimate predictor-label moments and target data to estimate predictor outputs.
- Consistency: Under Assumptions A.1–A.3, the estimated label-shift weights converge as source and target sample sizes grow.The proof establishes eventual invertibility of empirical confusion matrices and then applies the strong law of large numbers with the continuous mapping theorem.
- Error bounds: The error bounds decrease with sample size and include terms reflecting the magnitude of the shift and the predictor’s quality.The bounds contain sample-size dependence proportional to 1/n or 1/m, a ∥w∥2 term, and dependence on the smallest eigenvalue σmin.
- Conditions: If the predictor cannot distinguish some classes, the confusion matrix is low-rank, σmin = 0, and BBSE is invalid.Invertibility is therefore a necessary operational condition for the estimator’s use.
- Error bounds: The bound for estimating target moments includes sampling error even with exact weights plus additional error from estimating weights through the predictor.The predictor-dependent component is governed by the predictor-related terms in the bound and cannot always be improved.
5. Application of the results
The paper applies black-box prediction outputs to detect label shift, test broader nonstationarity, estimate correction weights, and retrain classifiers through importance-weighted ERM.
- Black Box Shift Detection: Label-shift detection tests whether predicted-label distributions differ between source and target data, providing evidence for q(y) ≠ p(y).The approach uses two-sample tests on predictor outputs rather than observing target labels directly.
- Black Box Shift Detection: Figure 1 reports controlled Type I error without shift, high power under mild shift, and power gains from better classifiers; the test nearly matches an oracle and beats direct high-dimensional testing.The figure compares BBSD with kernel two-sample tests and an infeasible oracle test.
- Black Box Shift Detection: Testing predictor outputs reduces the comparison to a one-dimensional distribution, avoiding the exponential difficulty of high-dimensional two-sample testing.Under label shift, equality of the feature distributions implies equality of predictor-output distributions.
- Black Box Shift Correction: BBSE estimates label-shift ratios by solving a linear system based on an invertible confusion matrix, then uses the estimates for importance-weighted ERM.The correction algorithm trains a predictor, estimates its confusion matrix on held-out data, solves for weights, clips negative weights, and fits a weighted classifier.
- Black Box Shift Correction: The method is presented as a general domain-adaptation tool, although weighted estimation can involve a bias-variance tradeoff and sporadic shifts may make correction harmful.The paper proposes combining detection and estimation so correction is applied when shift is likely.
6. Experiments
Experiments evaluate BBSD, BBSE, and BBSC on simulated shifts in MNIST and CIFAR10. Results compare black-box approaches with kernel baselines where feasible and examine predictor quality, dataset size, and shift magnitude.
- Experimental setup: The experiments cover shift detection, weight estimation, and classifier correction using hard- and soft-output black-box predictors.MNIST and CIFAR10 experiments use simulated knock-out, tweak-one, and Dirichlet label shifts.
- MNIST detection: MNIST detection experiments train a two-layer ReLU MLP and test increasingly strong knock-out shifts, with Type I error controlled at any significance level.The data are split into training, validation, and test sets, with 20,000 points in each split.
- MNIST detection: BBSD has greater power than the kernel two-sample test at δ = 0.5, and power improves as the black-box predictor receives more training.This connects detection performance to predictor quality rather than requiring a calibrated predictor.
- CIFAR10 correction: CIFAR10 experiments compare BBSC with unweighted classifiers under tweak-one and Dirichlet shifts, using varying shift degrees and multiple sampled target label distributions.Kernel baselines are omitted because they cannot handle datasets of this size or dimensionality.
- Baseline comparison: KMM requires O(n^3) Gram-matrix computation, is feasible only to n = 8000 on MNIST, and cannot run on CIFAR10; its weight-estimation convergence is slower than BBSE.CIFAR10 comparisons therefore use unweighted ERM rather than kernel-based baselines.
7. Discussion
The discussion addresses dataset construction, sporadic shifts, predictor availability, degenerate confusion matrices, and future streaming extensions.
- Practical considerations: Uniform source label distributions make BBSE applicable across unknown target label distributions because the weight norm remains controlled.The paper recommends uniform class sampling when class-conditioned annotation can be selected.
- Practical considerations: When label shift occurs sporadically, applying correction without detection can damage the classifier, so the paper combines detection with estimation.Correction is proposed only when a shift is likely to have occurred.
- Practical considerations: Randomly training the predictor on a data split makes BBSE worse than the baseline on extremely small datasets, whereas a known deployed predictor avoids that factor-of-two cost.The discussion suggests this setting may support rapid mild-shift detection.
- Degenerate confusion matrices: Degenerate confusion matrices can be addressed with soft predictions, merging classes, or a pseudo-inverse, though some class frequencies may then remain unestimable.Soft confusion matrices may remain invertible when hard predictions never select a class.
- Future work: Streaming deployment remains future work because trailing-window BBSE trades fresher estimates against larger sample sizes and lower estimation error.Progressively shifting label distributions create this window-selection tradeoff.
A. Additional discussion
The discussion addresses practical use of BBSE, including testing label-shift fit, selecting predictors, interpreting singular values, and reusing data with a bias trade-off.
- Kernel two-sample tests can assess whether label shift is a good finite-sample approximation despite unobserved variables affecting both x and y.The test compares reweighted source features with unlabeled target features using an arbitrary feature map and kernel.
- Under label shift, the reweighted source expectation and target sample mean should differ only by statistical error when the assumption is valid or approximately valid.The error depends on sample sizes and the accuracy of the estimated label ratios.
- Model selection criterion and the choice of f: The recommended model-selection heuristic chooses the black-box predictor whose estimated confusion matrix has the largest smallest singular value.This quantity is estimable using labeled source data alone.
- Model selection criterion and the choice of f: Figure 5 plots the smallest singular value of the estimated confusion matrix against classifier training epochs under a source distribution removing 80% of the fifth class.The experiment uses the same multilayer perceptron and dataset as Figure 1c.
- Model selection criterion and the choice of f: The one-epoch predictor has much lower test power than the five-epoch predictor, and their singular-value gap is qualitatively predictive of that difference.
- Reusing the same labeled data to train f and estimate w improves data efficiency but introduces a small, difficult-to-quantify bias.The paper presents data reuse as a potentially useful practical heuristic, especially for large n.
B. Proofs
The proofs establish BBSE’s asymptotic validity by proving eventual invertibility of the empirical confusion matrix and applying convergence results to its inverse-based estimator.
- The appendix presents proofs of Lemma 1 and Proposition 2.
- Conditional independence of predicted and true labels given x, together with the fixed predictor, supports the confusion-matrix identities used in Lemma 1.The proof also invokes the law of total probability.
- If the population confusion matrix is invertible with smallest singular value δ, concentration bounds control the probability that its empirical estimate is non-invertible.The proof bounds entrywise deviations using Hoeffding’s inequality.
- The non-invertibility probabilities are summable, so the First Borel-Cantelli Lemma implies that the empirical confusion matrix is eventually invertible almost surely.This yields almost-sure invertibility as n approaches infinity.
- The strong law of large numbers and continuous mapping through matrix inversion establish almost-sure convergence of the BBSE estimates.The argument uses continuity of inversion for invertible matrices.
C. Concentration inequalities
This section introduces scalar and matrix concentration tools used to control empirical estimation errors.
- Hoeffding’s inequality applies to independent random variables bounded within individual intervals [a_i, b_i].
- These inequalities supply concentration results for the scalar and matrix quantities appearing in the paper’s proofs.
- The matrix Bernstein inequality provides concentration for independent rectangular random matrices under a variance-parameter condition.