Detecting and Correcting for Label Shift
with Black Box Predictors

arXiv:1802.03916v2 [cs.LG] 28 Feb 2018

Zachary C. Lipton1,2∗, Yu-Xiang Wang2∗, Alex Smola2
Carnegie Mellon University
Amazon AI
zlipton@cmu.edu, yuxiangw@amazon.com, smola@amazon.com
March 1, 2018
Abstract
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

Assume that in August we train a pneumonia predictor. Our covariates consist of chest X-rays
administered in the previous year (distribution P ) and the labels binary indicators of whether a
physician diagnoses the patient with pneumonia. We train a convnet f to predict pneumonia given
an X-ray image. Assume that in the training set .2% of patients have pneumonia. We deploy f in
the clinic and for several months, it reliably predicts roughly .2% positive.
Fast-forward to January (distribution Q): Running f on the last week’s data, we find that 5% of
patients are predicted to have pneumonia! Because f remains fixed, the shift must owe to a change
in the marginal p(x), violating the familiar iid assumption. Absent familiar guarantees, we wonder:
Is f still accurate? What’s the real current rate of pneumonia? Shouldn’t our classifier, trained
under an obsolete prior, underestimate pneumonia when uncertain? Thus, we might suspect that
the real prevalence is greater than 5%.
∗ Equal

contribution

1

Given only labeled training data, and unlabeled test data, we desire to: (i) detect distribution
shift, (ii) quantify it, and (iii) correct our model to perform well on the new data. Absent assumptions
on how p(y, x) changes, the task is impossible. The covariate shift assumption posits
q(y, x) = q(x)p(y|x).
In our case, covariate shift suggests that the distribution of radiologic findings P(x) changes, but
that the conditional probability of pneumonia P(y|x) remains constant. It’s easy to see how this can
go wrong: What if our only feature were a cough? Normally, cough might not (strongly) indicate
pneumonia. But during an epidemic, P(pneumonia|cough) might go up substantially. Instead, we
posit that the labels P(y) shift, i.e. the manifestations of a disease do not change, the prevalence of
the disease does:
q(y, x) = q(y)p(x|y)
This distinction is well-known: Schölkopf et al. [16] observe that covariate shift corresponds to
causal learning (predicting effects), while label shift corresponds to anticausal learning (predicting
causes). We focus on label shift, motivated by diagnosis (diseases cause symptoms) and recognition
tasks (objects cause sensory observations). During a pneumonia outbreak, P(y|x) might rise, but
the manifestations of the disease p(x|y) might not change. Moreover, sometimes the very purpose of
learning is to detect changes in p(y). Examples include monitoring public health [7] and wildlife
populations [1]. Despite its importance, label shift is comparatively under-investigated: given
covariates x from both P and Q quantifying q(x)/p(x) may seem more intuitive.
We introduce Black Box Shift Estimation (BBSE) to estimate label shift using a black box predictor
f . BBSE estimates the ratios wl = q(yl )/p(yl ) for each label l, requiring only that the expected
confusion matrix, on P is invertible1 . We estimate ŵ by solving a linear system Ax = b where A is
the confusion matrix of f on the training data and b is the average output of f calculated on the
test data. We make the following contributions:
1. Consistency and error bounds for BBSE.
2. We apply BBSE to statistical tests for detecting distribution shift (under the label shift
assumption)
3. Model correction through importance-weighted Empirical Risk Minimization.
4. A comprehensive empirical validation of BBSE.
BBSE requires only labeled training data, unlabeled test data, and a black box predictor whose
confusion matrices are invertible under the training distribution. In comparison to past approaches
based on Kernel Mean Matching (KMM) [22], EM [4], and Bayesian inference [18], BBSE offers the
following advantages:
1. Accuracy does not depend on the data dimensionality.
2. Works with arbitrary black box predictors, even biased, uncalibrated, or inaccurate models.
3. Exploits deep learning without losing theoretical guarantees: better predictors provably lower
sample complexity.
1 For degenerate confusion matrices, a variant using continuous predictor responses (e.g. softmax outputs) may
still work.

2

4. Owing to generality, BBSE could be a standard diagnostic and corrective tool for arbitrary
ML models.

2

Prior Work

Despite its wide applicability, learning under label shift with unknown q(y) remains curiously
under-explored. Noting the difficulty of the problem, Storkey [18] proposes placing a (meta-)prior
over p(y) and inferring the posterior distribution from the test covariates. This requires explicitly
estimating p(x|y), which may not be feasible in high-dimensional datasets. Chan & Ng [4] infer
q(y) using EM but their method also requires estimating p(x|y). Schölkopf et al. [16] articulates
connections between label shift and anti-causal learning and Zhang et al. [22] extend the kernel
mean matching approach due to [9] to the label shift problem. When q(y) is known, label shift
simplifies to the problem of changing base rates [3, 6].
Covariate shift, also called sample selection bias, is well-studied [20, 11, 19, 9]. Rosenbaum & Rubin
[15] introduce Propensity Scoring to design unbiased experiments. Shimodaira [17] is one of the
earliest to propose correcting models via weighting examples in ERM by q(x)/p(x). Later works
estimate importance weights from the available data, e.g., Gretton et al. [9] propose kernel mean
matching to re-weight training points.
The earliest relevant work to ours comes from econometrics and addresses the use of non-random
samples to estimate behavior. Heckman [10] addresses sample selection bias, while [13] investigates
estimating parameters under, choice-based or endogenous stratified sampling, cases analogous to a
shift in the label distribution. Crucially, unlike our setting, they assume that the test set marginal
distributions are both known. Also related to our work is the foundational research on propensity
scoring [15] Finally, we note a connection to cognitive science work demonstrating that humans
classify items differently depending on the set of unlabeled items in which they appear [23].
Previous work requires estimating q(x), q(x)/p(x), or p(x|y), often relying on kernel methods, which
scale poorly with dataset size and underperform on high-dimensional data. In contrast, BBSE uses
black-box predictors to reduce dimensionality. It exploits advances in discriminative modeling and
does not degrade, either in accuracy or tractability, as data or feature dimensions grow.

3

Problem setup

We use x ∈ X = Rd and y ∈ Y to denote the feature and label variables. For simplicity, we assume
that Y is a discrete domain equivalent to {1, 2, ..., k}. Let P, Q be the source and target distributions
defined on X × Y. We use p, q to denote the probability density function (pdf) or probability mass
function (pmf) associated with P and Q respectively. The random variable of interest is clear from
context. For example, p(y) is the p.m.f. of y ∼ P and q(x) is the p.d.f. of x ∼ Q. Moreover, p(y = i)
and q(y = i) are short for PP (y = i) and PQ (y = i) respectively, where P(S) := E[1(S)] denotes the
probability of an event S and E[·] denotes the expectation. Subscripts P and Q on these operators
make the referenced distribution clear.

3

In standard supervised learning, the learner observes training data (x1 , y1 ), (x2 , y2 ), ..., (xn , yn )
drawn iid from a training (or source) distribution P . We denote the collection of feature vectors by
X ∈ Rn×d and the label by y. Under Domain Adaptation (DA), the learner additionally observes
a collection of samples X 0 = [x01 ; ...; x0m ] drawn iid from a test (or target) distribution Q. Our
objective in DA is to predict well for samples drawn from Q.
In general this task is impossible, e.g., P and Q might not share support. This paper considers DA
under 3 extra assumptions:
A.1 The label shift (also known as target shift) assumption
p(x|y) = q(x|y) ∀ x ∈ X , y ∈ Y.
A.2 For every y ∈ Y with q(y) > 0 we require p(y) > 0.2
A.3 Access to a black box predictor f : X → Y where the expected confusion matrix Cp (f ) is
invertible.
CP (f ) := p(f (x), y) ∈ R|Y|×|Y|
We now comment on the assumptions. A.1 corresponds to anti-causal learning. This assumption is
strong but reasonable in many practical situations, including medical diagnosis, where diseases cause
symptoms. It also applies when classifiers are trained on non-representative class distributions: Note
that while visions systems are commonly trained with balanced classes [5], the true class distribution
for real tasks is rarely uniform.
Assumption A.2 addresses identifiability, requiring that the target label distribution’s support be a
subset of training distribution’s. For discrete Y, this simply means that the training data should
contain examples from every class.
Assumption A.3 requires that the expected predictor outputs for each class be linearly independent.
This assumption holds in the typical case where the classifier predicts class yi more often given
images actually belong to yi than given images from any other class yj . In practice, f could be a
neural network, a boosted decision-tree or any other classifier trained on a holdout training data
set. We can verify at training time that the empirical estimated normalized confusion matrix is
invertible. Assumption A.3 generalizes naturally to soft-classifiers, where f outputs a probability
distribution supported on Y. Thus BBSE can be applied even when the confusion matrix
is degenerate.
Objective We wish to estimate w(y) := q(y)/p(y) for every y ∈ Y with training data, unlabeled
test data and a predictor f . This estimate
enables DA techniques under the importance-weighted
Pn
ERM framework, which solves min i=1 wi `(yi , xi ), using wi = q(xi , yi )/p(xi , yi ). Under the label
shift assumption, the importance weight wi = q(yi )/p(yi ). This task isn’t straightforward because
we don’t observe samples from q(y).
2 Assumes the absolute continuity of the (hidden) target label’s distribution with respect to the source’s, i.e.
dq(y)/dp(y) exists.

4

4

Main results

We now derive the main results for estimating w(y) and q(y). Assumption A.1 has the following
implication:
Lemma 1. Denote by ŷ = f (x) the output of a fixed function f : X → Y. If Assumption A.1 holds,
then
q(ŷ|y) = p(ŷ|y).
The proof is simple: recall that ŷ depends on y only via x. By A.1, p(x|y) = q(x|y) and thus
q(ŷ|y) = p(ŷ|y).
Next, combine the law of total probability and Lemma 1 and we arrive at
X
q(ŷ) =
q(ŷ|y)q(y)
y∈Y

=

X

p(ŷ|y)q(y) =

y∈Y

X

p(ŷ, y)

y∈Y

q(y)
.
p(y)

(1)

We estimate p(ŷ|y) and p(ŷ, y) using f and data from source distribution P , and q(ŷ) with unlabeled
test data drawn from target distribution Q. This leads to a novel method-of-moments approach for
consistent estimation of the shifted label distribution q(y) and the weights w(y).
Without loss of generality, we assume Y = {1, 2, ..., k}. Denote by µy , µŷ , µ̂ŷ , ν y , ν ŷ , ν̂ ŷ , w ∈ Rk
moments of p and q and their plug-in estimates, defined via
[µy ]i = q(y = i)

[ν y ]i = p(y = i)

[µŷ ]i = q(f (x) = i)
P
0
j {f (xj ) = i}
[µ̂ŷ ]i =
m

[ν ŷ ]i = p(f (x) = i)
P
j {f (xj ) = i}
[ν̂ ŷ ]i =
n

and [w]i = q(y = i)/p(y = i). Lastly define the covariance matrices Cŷ,y , Cŷ|y and Ĉŷ,y in Rk×k
via
[Cŷ,y ]ij = p(f (x) = i, y = j)
[Cŷ|y ]ij = p(f (x) = i|y = j)
1X
[Ĉŷ,y ]ij =
I{f (xl ) = i and yl = j}
n
l

We can now rewrite Equation (1) in matrix form:
µŷ = Cŷ|y µy = Cŷ,y w
Using plug-in maximum likelihood estimates of the above quantities yields the estimators
ŵ = Ĉ−1
ŷ,y µ̂ŷ and µ̂y = diag(ν̂ y )ŵ,
where ν̂ y is the plug-in estimator of ν y .
Next, we establish that the estimators are consistent.
5

a.s.

Proposition 2 (Consistency). If Assumption A.1, A.2, A.3 are true, then as n, m → ∞, ŵ −→ w
a.s.
and µ̂y −→ µy .
The proof (see Appendix A) uses the First Borel-Cantelli Lemma to show that the probability that
the entire sequence of empirical confusion matrices with data size n + 1, ..., ∞ are simultaneously
invertible converges to 1, thereby enabling us to use the continuous mapping theorem after applying
the strong law of large numbers to each component.
We now address our estimators’ convergence rates.
Theorem 3 (Error bounds). Assume that A.3 holds robustly. Let σmin be the smallest eigenvalue
−2
of Cŷ,y . There exists a constant C > 0 such that for all n > 80 log(n)σmin
, with probability at least
−10
−10
1 − 3kn
− 2km
we have


C
kwk2 log n k log m
2
kŵ − wk2 ≤ 2
+
(2)
σmin
n
m
kµ̂y − µy k2 ≤

Ckwk2 log n
+ kν y k2∞ kŵ − wk22
n

(3)

The bounds give practical insights (explored more in Section 7). In (2), the square error depends
on the sample size and is proportional to 1/n (or 1/m). There is also a kwk2 term that reflects
how different the source and target distributions are. In addition, σmin reflects the quality of the
given classifier f . For example, if f is a perfect classifier, then σmin = miny∈Y p(y). If f cannot
distinguish between certain classes at all, then Cŷ,y will be low-rank, σmin = 0, and the technique is
invalid, as expected.
We now parse the error bound of µ̂y in (3). The first term kwk2 /n is required even if we observe
the importance weight w exactly. The second term captures the additional error due to the fact
that we estimate w with predictor f . Note that kν y k2∞ ≤ 1 and can be as small as 1/k 2 when p(y)
is uniform. Note that when f correctly classifies each class with the same probability, e.g. 0.5, then
2
kν y k2 /σmin
is a constant and the bound cannot be improved.
Proof of Theorem 3. Assumption A.2 ensures that w < ∞.
−1
ŵ = Ĉ−1
(µŷ + E2 )
ŷ,y µ̂ŷ = (Cŷ,y + E1 )
−1
= w + [(Cŷ,y + E1 )−1 − C−1
E2
ŷ,y ]µŷ + (Cŷ,y + E1 )

By completing the square and Cauchy-Schwartz inequality,
T
−1
kŵ − wk2 ≤ 2µTŷ [(Cŷ,y + E1 )−1 − C−1
− C−1
ŷ,y ] [(Cŷ,y + E1 )
ŷ,y ]µŷ

+ 2E2 [(Cŷ,y + E1 )−1 ]T (Cŷ,y + E1 )−1 E2 .
By the Woodbury matrix identity, we get that
−1
−1
−1
−1 −1 −1
Ĉ−1
Cŷ,y .
ŷ,y = Cŷ,y + Cŷ,y [E1 + Cŷ,y ]

Substitute into the above inequality and use (1) to get
n
oT
−1 T
−1
−1 −1
−1
kŵ − wk2 ≤2w [E1−1 + C−1
]
[Cŷ,y
] × C−1
w
ŷ,y
ŷ,y [E1 + Cŷ,y ]
+ 2E2T [(Cŷ,y + E1 )−1 ]T (Cŷ,y + E1 )−1 E2
6

(4)

We now provide a high probability bound on the Euclidean norm of E2 , the operator norm of
−1
E1 , which will give us an operator norm bound of [E1−1 + C−1
and (Cŷ,y + E1 )−1 under our
ŷ,y ]
assumption on n, and these will yield a high probability bound on the square estimation error.
Pn
Operator norm of E1 . Note that Ĉŷ,y = n1 i=1 ef (xi ) eTyi , where ey is the standard basis with
1 at the index of y ∈ Y and 0 elsewhere. Clearly, Eef (xi ) eTyi = Cŷ,y . Denote Zi := ef (xi ) eTyi − Cŷ,y .
Check that kZi k2 ≤ kZi kF ≤ kZi k1,1 ≤ 2, max{kE[Zi ZTi ]k, kE[ZTi Zi ]k} ≤ 1, by matrix Bernstein
inequality (Lemma 7) we have for all t ≥ 0
t2

P(kE1 k ≥ t/n) ≤ 2ke− n+2t/3 .
√
Take t = 20n log n and use the assumption that n ≥ 4 log n/9 (which holds under our assumption
on n since σmin < 1). Then with probability at least 1 − 2kn−10
r
20 log n
kE1 k ≤
.
n
Using the assumption on n, we have kE1 k ≤ σmin /2
−1
k[E1−1 + C−1
k ≤ 2kE1 k ≤
ŷ,y ]

Also, we have k(Cŷ,y + E1 )−1 k ≤

√
2 20 log n
√
.
n

2
σmin .

Pm
1
0
0
Euclidean norm of E2 . Note that [E2 ]l = m
i=1 1(f (xi ) = l) − q(f (xi ) = l). By the standard
Hoeffding’s inequality and union bound argument, we have that with probability larger than
1 − 2km−10
√
10k log m
√
kE2 k = kµŷ − µ̂ŷ k2 ≤
m
Substitute into Equation (4), we get
kŵ − wk2 ≤

80k log m
80 log n
kwk2 +
2
2 m ,
σmin n
σmin

(5)

which holds with probability 1 − 2kn−10 − 2km−10 . We now turn to µ̂y . Recall that µ̂y = diag(ν̂ y )ŵ.
Let the estimation error of ν̂ y be E0 .
µ̂y = µy + diag(E0 )w + diag(ν y )(ŵ − w) + diag(E0 )(ŵ − w).
q
n
By Hoeffding’s inequlaity kE0 k∞ ≤ 20 log
with probability larger than 1 − kn−10 . Combining
n
with (5) yields
1
20kwk2 log n
+ kν y k2∞ kŵ − wk2 + O( 2 )
kµ̂y − µy k2 ≤
n
n
which holds with probability 1 − 3kn−10 − 2km−10 .
7

0.8

1.0

1.0

0.8

0.8

0.6

0.6

0.4

0.4

0.2

0.2

0.0

Uniform
KS-test with MLP
oracle KS-test
MMD B-test on p(x)

0.0
0.0

0.2

0.4

0.6

0.8

(a) CDF of p-value at δ = 0

1.0

0.0

0.2

0.4

0.6

0.8

1.0

(b) CDF of p-value at δ = 0.6

Power of the test (1- Type II error)

Uniform
KS-test with MLP
oracle KS-test
MMD B-test on p(x)

1.0

KS-test with MLP (1 epoch)
KS-test with MLP (5 epoch)
MMD B-test on p(x)
Oracle KS-test

0.6

0.4

0.2

0.0

0.0

0.2

0.4

0.6

0.8

(c) Power at α = 0.05

Figure 1: Label-shift detection on MNIST. Pane 1a illustrates that Type I error is correctly controlled
absent label shift. Pane 1b illustrates high power under mild label-shift. Pane 1c shows increased
power for better classifiers. We compare to kernel two-sample tests [21] and an (infeasible) oracle
two sample test that directly tests p(y) = q(y) with samples from each. The proposed test beats
directly testing in high-dimensions and nearly matches the oracle.

5
5.1

Application of the results
Black Box Shift Detection (BBSD)

Formally, detection can be cast as a hypothesis testing problem where the null hypothesis is
H0 : q(y) = p(y) and the alternative hypothesis is that H1 : q(y) 6= p(y). Recall that we observe
neither q(y) nor any samples from it. However, we do observe unlabeled data from the target
distribution and our predictor f .
Proposition 4 (Detecting label-shift). Under Assumption A.1, A.2 and for each classifier f
satisfying A.3 we have that q(y) = p(y) if and only if p(ŷ) = q(ŷ).
Proof. Plug P and Q into (1) and apply Lemma 1 with assumption A.1. The result follows directly
from our analysis in the proof of Proposition 2 that shows p(ŷ, y) is invertible under the assumptions
A.2 and A.3.
Thus, under weak assumptions, we can test H0 by running two-sample tests on readily available
samples from p(ŷ) and q(ŷ). Examples include the Kolmogorov-Smirnoff test, Anderson-Darling
or the Maximum Mean Discrepancy. In all tests, asymptotic distributions are known and we can
almost perfectly control the Type I error. The power of the test (1-Type II error) depends on the
classifier’s performance on distribution P , thereby allowing us to leverage recent progress in deep
learning to attack the classic problem of detecting non-stationarity in the data distribution.
One could also test whether p(x) = q(x). Under the label-shift assumption this is implied by
q(y) = p(y). The advantage of testing the distribution of f (x) instead of x is that we only need to
deal with a one-dimensional distribution. Per theory and experiments in [14] two-sample tests in
high dimensions are exponentially harder.

8

Estimating error of w at = 0.1

Estimating error of w at = 1.0

Estimating error of w at = 10.0
100

101

MSE

100

MSE

MSE

100

10
10

1

BBSE-hard
BBSE-soft
KMM
KMM timeout

10

500 1000 2000 4000 8000 16000 32000

n
Target accuracy at = 0.1

2

0.98

BBSE-hard
BBSE-soft
KMM
KMM timeout

10

2

BBSE-hard
BBSE-soft
KMM
KMM timeout
500 1000 2000 4000 8000 16000 32000

n
Target accuracy at = 1.0

0.98

n
Target accuracy at = 10.0

0.96

0.96

0.92

Unweighted
BBSE-hard
BBSE-soft
KMM
KMM timeout

0.90
0.88
0.86

500 1000 2000 4000 8000 16000 32000

n

0.94
0.92

Unweighted
BBSE-hard
BBSE-soft
KMM
KMM timeout

0.90
0.88
0.86

Target accuracy

0.94

0.94

Target accuracy

Target accuracy

1

500 1000 2000 4000 8000 16000 32000

0.98
0.96

10

1

0.92
0.90

Unweighted
BBSE-hard
BBSE-soft
KMM
KMM timeout

0.88
0.86
0.84

500 1000 2000 4000 8000 16000 32000

n

500 1000 2000 4000 8000 16000 32000

n

Figure 2: Estimation error (top row) and correction accuracy (bottom row) vs dataset size on MNIST
data compared to KMM [22] under Dirichlet shift (left to right) with α = {.1, 1.0, 10.0} (smaller α
means larger shift). BBSE confidence interval on 20 runs, KMM on 5 runs due to computation;
n = 8000 is largest feasible KMM experiment.
One surprising byproduct is that we can sometimes use this approach to detect covariate-shift,
concept-shift, and more general forms of nonstationarity.
Proposition 5 (Detecting general nonstationarity). For any fixed measurable f : X → Y
P = Q =⇒ p(x) = q(x) =⇒ p(ŷ) = q(ŷ).
This follows directly from the measurability of f .
While the converse is not true in general, p(ŷ) = q(ŷ) does imply that for every measurable
S ⊂ Y,
q(x ∈ f −1 (S)) = p(x ∈ f −1 (S)).
This suggests that testing Ĥ0 : p(ŷ) = q(ŷ) may help us to determine if there’s sufficient statistical
evidence that domain adaptation techniques are required.

9

Estimating error of w at = 0.7

101

Estimating error of w at = 0.5

Estimating error of w at = 0.3

101
100

MSE

MSE

MSE

100
100

10

10

1

BBSE-hard
BBSE-soft
KMM
KMM timeout

10

1

500 1000 2000 4000 8000 16000 32000

BBSE-hard
BBSE-soft
KMM
KMM timeout

BBSE-hard
BBSE-soft
KMM
KMM timeout
10

500 1000 2000 4000 8000 16000 32000

n
Target accuracy at = 0.7

0.98

0.975

n
Target accuracy at = 0.5

Unweighted
BBSE-hard
BBSE-soft
KMM
KMM timeout

0.850
0.825
0.800

500 1000 2000 4000 8000 16000 32000

n

n
Target accuracy at = 0.3

0.94

0.92
0.90
0.88

Unweighted
BBSE-hard
BBSE-soft
KMM
KMM timeout

0.86
0.84
0.82

500 1000 2000 4000 8000 16000 32000

n

Target accuracy

0.875

Target accuracy

Target accuracy

0.900

500 1000 2000 4000 8000 16000 32000

0.96

0.94

0.925

2

0.98

0.96
0.950

1

0.92
0.90

Unweighted
BBSE-hard
BBSE-soft
KMM
KMM timeout

0.88
0.86
0.84

500 1000 2000 4000 8000 16000 32000

n

Figure 3: Label-shift estimation and correction on MNIST data with simulated tweak-one shift with
parameter ρ.

5.2

Black Box Shift Correction (BBSC)

Our estimator also points to a systematic method of correcting for label-shift via importance-weighted
ERM. Specifically, we propose the following algorithm:
Algorithm 1 Domain adaptation via Label Shift
input Samples from source distribution X, y. Unlabeled data from target distribution X 0 . A class
of classifiers F. Hyperparameter 0 < δ < 1/k.
1. Randomly split the training data into two X1 , X2 ∈ Rn/2×d and y 1 , y 2 Rn/2 .
2. Use X1 , y 1 to train the classifier and obtain f ∈ F.
3. On the
data set X2 , y 2 , calculate the confusion matrix Ĉŷ,y . If ,
 hold-out


if σmin Ĉŷ,y ≤ δ then
Set ŵ = 1.
else
Estimate ŵ = Ĉ−1
ŷ,y µ̂ŷ .
end if
4. Solve the importance weighted ERM on the X1 , y 1 with max(ŵ, 0) and obtain f˜.
output f˜

10

Note that for classes that occur rarely in the test set, BBSE may produce negative importance
weights. During ERM, a flipped sign would cause us to maximize loss, which is unbounded above.
Thus, we clip negative weights to 0.
Owing to its efficacy and generality, our approach can serve as a default tool to deal with domain
adaptation. It is one of the first things to try even when the label-shift assumption doesn’t hold. By
contrast, the heuristic method of using logistic-regression to construct importance weights [2] lacks
theoretical justification that the estimated weights are correct.
Even in the simpler problem of average treatment effect (ATE) estimation, it’s known that using
estimated propensity can lead to estimators with large variance [12]. The same issue applies in
supervised learning. We may prefer to live with the biased solution from the unweighted ERM
rather than suffer high variance from an unbiased weighted ERM. Our proposed approach offers a
consistent low-variance estimator under label shift.

6

Experiments

We experimentally demonstrate the power of BBSE with real data and simulated label shift. We
organize results into three categories — shift detection with BBSD, weight estimation with
BBSE, and classifier correction with BBSC. BBSE-hard denotes our method where f yields
classifications. In BBSE-soft, f outputs probabilities.
Label Shift Simulation To simulate distribution shift in our experiments, we adopt the following
protocols: First, we split the original data into train, validation, and test sets. Then, given
distributions p(y) and q(y), we generate each set by sampling with replacement from the appropriate
split. In knock-out shift, we knock out a fraction δ of data points from a given class from training
and validation sets. In tweak-one shift, we assign a probability ρ to one of the classes, the rest
of the mass is spread evenly among the other classes. In Dirichlet shift, we draw p(y) from a
Dirichlet distribution with concentration parameter α. With uniform p(y), Dirichlet shift is bigger
for smaller α.
Label-shift detection We conduct nonparametric two-sample tests as described in Section 5.1
using the MNIST handwritten digits data set. To simulate the label-shift, we randomly split the
training data into a training set, a validating set and a test set, each with 20,000 data points,
and apply knock-out shift on class y = 5 3 . Note that p(y) and q(y) differ increasingly as δ
grows large, making shift detection easier. We obtain f by training a two-layer ReLU-activated
Multilayer Perceptron (MLP) with 256 neurons on the training set for five epochs. We conduct a
two-sample test of whether the distribution of f (Validation Set) and f (Test Set) are the same using
the Kolmogorov-Smirnov test. The results, summarized in Figure 1, demonstrate that BBSD (1)
produces a p-value that distributes uniformly when δ = 0 4 (2) provides more power (less Type II
error) than the state-of-the-art kernel two-sample test that discriminates p(x) and q(x) at δ = 0.5,
and (3) gets better as we train the black-box predictor even more.
Weight estimation and label-shift correction We evaluate BBSE on MNIST by simulating
label shift and datasets of various sizes. Specifically, we split the training data set randomly in two,
3 Random
4 Thus

choice for illustration, method works on all classes.
we can control Type I error at any significance level.

11

Accuracy

0.9
0.8

Black Box Shift Correction on CIFAR10
BBSE-hard
Unweighted

1.0
0.9
0.8

Accuracy

1.0

0.7
0.6
0.5

Black Box Shift Correction on CIFAR10
BBSE-hard
Unweighted

0.7
0.6
0.5
0.4

0.4

0.3

0.3
0.0

0.2

0.4

0.6

0.8

1.0

1000

100

10

1

0.1

0.01

0.001

Figure 4: Accuracy of BBSC on CIFAR 10 with (left) tweak-one shift and (right) Dirichlet shift.
using first half to train f and the second half to estimate w. We use then use the full training set
for weighted ERM. As before, f is a two-layer MLP. For fair comparisons with baselines, the full
training data set is used throughout (since they do not need f without data splitting). We evaluate
our estimator ŵ against the ground truth w and by the prediction accuracy of BBSC on the test
set. To cover a variety of different types of label-shift, we take p(y) as a uniform distribution and
generate q(y) with Dirichlet shift for α = 0.1, 1.0, 10.0 (Figure 2).
Label-shift correction for CIFAR10 Next, we extend our experiments to the CIFAR dataset,
using the same MLP and this time allowing it to train for 10 epochs. We consider both tweak-one
and Dirichlet shift, and compare BBSE to the unweighted classifier under varying degrees of shift
(Figure 4). For the tweak-one experiment, we try ρ ∈ {0.0, 0.1, ..., 1.0}, averaging results over
all 10 choices of the tweaked label, and plotting the variance. For the Dirichlet experiments, we
sample 20 q(y) for every choice of α in the range {1000, 100, ..., .001}. Because the kernel-based
baselines cannot handle datasets of this size and dimensionality, here we compare only to unweighted
ERM.
Kernel mean matching (KMM) baselines We compare BBSE to the state-of-the-art kernel
mean matching (KMM) methods. For the detection experiments (Figure 1), our baseline is the
kernel B-test [21], an extension of the kernel max mean discrepancy (MMD) test due to Gretton
et al. [8] that boasts nearly linear-time computation and little loss in power. We compare BBSE to
a KMM approach Zhang et al. [22], that solves
min kCx|y (ν y ◦ w) − µx k2H ,
w

where we use operator Cx|y := E[φ(x)|ψ(y)] and function µx := EQ [φ(x)] to denote the kernel
embedding of p(x|y) and pQ (x) respectively. Note that under the label-shift assumption, Cx|y is
the same for P and Q. Also note that since Y is discrete, ψ(y) is simply the one-hot representation
of y, so ν y is the same as our definition before and Cx|y , ν y and µx must be estimated from finite
data. The proposal involves a constrained optimization by solving a Gaussian process regression
with automatic hyperparameter choices through marginal likelihood.

12

For fair comparison, we used the original authors’ implementations as baselines 5 and also used
the median trick to adaptively tune the RBF kernel’s hyperparameter. A key difference is that
BBSE matches the distribution of ŷ rather than distribution of x like [22] and we learn f through
supervised learning rather than by specifying a feature map φ by choosing a kernel up front.
Note that KMM, like many kernel methods, requires the construction and inversion of an n × n
Gram matrix, which has complexity of O(n3 ). This hinders its application to real-life machine
learning problems where n will often be 100s of thousands. In our experiments, we find that the
largest n for which we can feasibly run the KMM code is roughly 8, 000 and that is where we
unfortunately have to stop for the MNIST experiment. For the same reason, we cannot run KMM
for the CIFAR10 experiments. The MSE curves in Figure 2 for estimating w suggest that the
convergence rate of KMM is slower than BBSE by a polynomial factor and that BBSE better
handles large datasets.

7

Discussion

Constructing the training Set The error bounds on our estimates depend on the norm of the
true vector w(y) := q(y)/p(y). This confirms the common sense that absent any assumption on q(y),
and given the ability to select class-conditioned examples for annotations one should build a dataset
with uniform p(y). Then it’s always possible to apply BBSE successfully at test time to correct
f.
Sporadic Shift In some settings, p(y) might change only sporadically. In these cases, when no
label shift occurs, applying BBSC might damage the classifier. For these cases, we prose to combine
detection and estimation, correcting the classifier only when a shift has likely occurred.
Using known predictor In our experiments, f has been trained using a random split of the
data set, which makes BBSEto perform worse than baseline when the data set is extremely small.
In practice, especially in the context of web services, there could be a natural predictor f that
is currently being deployed whose training data were legacy and have little to do with the two
distributions that we are trying to distinguish. In this case, we do not lose that factor of 2 and we
do not suffer from the variance in training f with a small amount of data. This could allow us to
detect mild shift in distributions in very short period of time. Making it suitable for applications
such as financial market prediction.
BBSE with degenerate confusion matrices In practice, sometime confusion matrices will
be degenerate. For instance, when a class i is rare under P , and the features are only partially
predictive, we might find that p(f (x) = i) = 0. In these cases, two straightforward variations on the
black box method may still work: First, while our analysis focuses on confusion matrices, it easily
extends to any operator f , such as soft probabilities. If each class i, even if i is never the argmax
for any example, so long as p(ŷ = i|y = i) > p(ŷ = i|y = j) for any j 6= i, the soft confusion matrix
will be invertible. Even when we produce and operator with an invertible confusion matrix, two
options remain: We can merge c classes together, yielding a (k − c) × (k − c) invertible confusion
matrix. While we might not be able to estimate the frequencies of those c classes, we can estimate
the others accurately. Another possibility is to compute the pseudo-inverse.
5 https://github.com/wojzaremba/btest,

http://people.tuebingen.mpg.de/kzhang/Code-TarS.zip

13

Future Work As a next step, we plan to extend our methodology to the streaming setting. In
practice, label distributions tend to shift progressively, presenting a new challenge: if we apply BBSE
on trailing windows, then we face a trade-off. Looking far back increases m, lowering estimation
error, but the estimate will be less fresh. The use of propensity weights w on y makes BBSE
amenable to doubly-robust estimates, the typical bias-variance tradeoff, and related techniques,
common in covariate shift correction.

References
[1] Beijbom, O., Edmunds, P. J., Kline, D. I., Mitchell, B. G., & Kriegman, D. (2012). Automated
annotation of coral reef survey images. In Computer Vision and Pattern Recognition (CVPR).
[2] Bickel, S., Brückner, M., & Scheffer, T. (2009). Discriminative learning under covariate shift.
Journal of Machine Learning Research, 10 (Sep), 2137–2155.
[3] Bishop, C. M. (1995). Neural networks for pattern recognition. Oxford university press.
[4] Chan, Y. S., & Ng, H. T. (2005). Word sense disambiguation with distribution estimation. In
Proceedings of the 19th international joint conference on Artificial intelligence, (pp. 1010–1015).
Morgan Kaufmann Publishers Inc.
[5] Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., & Fei-Fei, L. (2009). Imagenet: A large-scale
hierarchical image database. In CVPR.
[6] Elkan, C. (2001). The foundations of cost-sensitive learning. In IJCAI .
[7] Ginsberg, J., Mohebbi, M. H., Patel, R. S., Brammer, L., Smolinski, M. S., & Brilliant, L.
(2009). Detecting influenza epidemics using search engine query data. Nature.
[8] Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., & Smola, A. (2012). A kernel
two-sample test. Journal of Machine Learning Research, 13 (Mar), 723–773.
[9] Gretton, A., Smola, A. J., Huang, J., Schmittfull, M., Borgwardt, K. M., & Schölkopf, B.
(2009). Covariate shift by kernel mean matching. Journal of Machine Learning Research.
[10] Heckman, J. J. (1977). Sample selection bias as a specification error (with an application to
the estimation of labor supply functions).
[11] Huang, J., Gretton, A., Borgwardt, K. M., Schölkopf, B., & Smola, A. J. (2007). Correcting
sample selection bias by unlabeled data. In Advances in neural information processing systems.
[12] Kang, J. D., & Schafer, J. L. (2007). Demystifying double robustness: A comparison of
alternative strategies for estimating a population mean from incomplete data. Statistical
science, 22 (4), 523–539.
[13] Manski, C. F., & Lerman, S. R. (1977). The estimation of choice probabilities from choice
based samples. Econometrica: Journal of the Econometric Society.
[14] Ramdas, A., Reddi, S. J., Póczos, B., Singh, A., & Wasserman, L. A. (2015). On the decreasing
power of kernel and distance based nonparametric hypothesis tests in high dimensions. In
AAAI .

14

[15] Rosenbaum, P. R., & Rubin, D. B. (1983). The central role of the propensity score in
observational studies for causal effects. Biometrika, 70 (1), 41–55.
[16] Schölkopf, B., Janzing, D., Peters, J., Sgouritsa, E., Zhang, K., & Mooij, J. (2012). On causal
and anticausal learning. In International Coference on International Conference on Machine
Learning (ICML-12), (pp. 459–466). Omnipress.
[17] Shimodaira, H. (2000). Improving predictive inference under covariate shift by weighting the
log-likelihood function. Journal of statistical planning and inference.
[18] Storkey, A. (2009). When training and test sets are different: characterizing learning transfer.
Dataset shift in machine learning.
[19] Sugiyama, M., Nakajima, S., Kashima, H., Buenau, P. V., & Kawanabe, M. (2008). Direct
importance estimation with model selection and its application to covariate shift adaptation.
In Advances in neural information processing systems.
[20] Zadrozny, B. (2004). Learning and evaluating classifiers under sample selection bias. In
Proceedings of the twenty-first international conference on Machine learning, (p. 114). ACM.
[21] Zaremba, W., Gretton, A., & Blaschko, M. (2013). B-test: A non-parametric, low variance
kernel two-sample test. In Advances in neural information processing systems, (pp. 755–763).
[22] Zhang, K., Schölkopf, B., Muandet, K., & Wang, Z. (2013). Domain adaptation under target
and conditional shift. In International Conference on Machine Learning, (pp. 819–827).
[23] Zhu, X., Gibson, B. R., Jun, K.-S., Rogers, T. T., Harrison, J., & Kalish, C. (2010). Cognitive
models of test-item effects in human category learning. In ICML.

15

A

Proofs

We present the proofs of Lemma 1 and Proposition 2 in this Appendix.
Proof of Lemma 1. By the law of total probability
X
X
q(ŷ|y) =
q(ŷ|x, y)q(x|y) =
q(ŷ|x, y)p(x|y)
y∈Y

=

X

y∈Y

pf (ŷ|x)p(x|y) =

y∈Y

X

p(ŷ|x, y)p(x|y) = p(ŷ|y).

y∈Y

We applied A.1 to the second equality, and used the conditional independence ŷ ⊥
⊥ y|x under P and
Q together with p(ŷ|x) being determined by f , which is fixed.
Proof of Proposition 2. A.2 ensures that w < ∞. By Assumption A.3, Cŷ,y is invertible. Let δ > 0
be its smallest singular value. We bound the probability that Ĉŷ,y is not invertible:

P(Ĉŷ,y is not invertible) ≤ P(σmin (Ĉŷ,y ) < δ/2)
δ
≤ P(kĈŷ,y − Cŷ,y k2 ≥ δ/2) ≤ P(kĈŷ,y − Cŷ,y kF ≥ √ )
↑
2 k
pigeon hole

≤ P(∃(i, j) ∈ [k]2 , s.t.|[Ĉŷ,y ]i,j − [Cŷ,y ]i,j | ≥
↑

nδ 2
δ
) ≤ 2e− 4k3 .
1.5
2k
↑

pigeon hole

Hoeffding

P

By the convergence of geometric series n P(Ĉŷ,y is not invertible) < +∞. This allows us to invoke
the First Borel-Cantelli Lemma, which shows
P(Ĉŷ,y is not invertible i.o.) = 0.

(6)

This ensures that as n → ∞, Ĉŷ,y is invertible almost surely. By the strong law of large numbers
a.s.
a.s.
a.s.
(SLLN), as n → ∞ Ĉŷ,y −→ Cŷ,y and ν̂ y −→ ν y . Similarly, as m → ∞, µ̂ŷ −→ µŷ . Combining
these with (6) and applying the continuous mapping theorem with the fact that the inverse of an
invertible matrix is a continuous mapping we get that
a.s.

a.s.

ŵ = [Ĉŷ,y ]−1 µ̂ŷ −→ w, and µ̂y = diag(ν̂ y )ŵ −→ µy .

B

Concentration inequalities

Lemma 6 (Hoeffding’s
Pn inequality). Let x1 , ..., xn be independent random variables bounded by
[ai , bi ]. Then x̄ = n1 i=1 xi obeys for any t > 0


2n2 t2
P
P(|x̄ − E[x̄]| ≥ t) ≤ 2 exp − n
.
2
i=1 (bi − ai )
16

Lemma 7 (Matrix Bernstein Inequality (rectangular case)). Let Z1 , ..., Zn be independent random
matrices with dimension d1 × d2 and each satisfy
EZi = 0 and kZi k ≤ R
almost surely. Define the variance parameter
X
X
σ 2 − max{k
E[Zi ZTi ]k, k
E[ZTi Zi ]k}.
i

i

Then for all t ≥ 0,
!
P k

X

Zi k ≥ t

−t2

≤ (d1 + d2 ) · e σ2 +Rt/3 .

i

17

