EM And Moment Methods For Label-Shift Estimation

Estimating Target Prevalences From Unlabeled Predictions

Under label shift, the distribution of features conditional on class remains stable while the class prevalences change. This makes it possible to use labeled source data to estimate how a classifier behaves within each class, then use unlabeled target predictions to estimate the target class proportions. The key information is the classifier’s confusion matrix: if its class-specific prediction patterns are sufficiently distinct, their mixture in the target sample reveals the target prevalences.

There are two natural approaches. A moment method solves the linear equation relating target prediction frequencies to target class prevalences. An expectation-maximization (EM) method instead maximizes the likelihood of the unlabeled predicted labels under that same mixture model. Both rely on stable class-conditional prediction behavior, and both become unreliable when the confusion matrix has nearly indistinguishable columns.

Definition: Let $Y$ take values in $\{1,\ldots,K\}$, and let $\hat Y$ be a classifier’s predicted class. Define the source confusion matrix $C$ by $$ C_{ij}=P_S(\hat Y=i\mid Y=j). $$ Thus column $j$ describes the prediction distribution when the true class is $j$. Each column is a probability vector: $C_{ij}\geq 0$ and $\sum_{i=1}^K C_{ij}=1$. Write $\pi^T_j=P_T(Y=j)$ for the target prevalence of class $j$, and $q_i=P_T(\hat Y=i)$ for the target prediction frequency of class $i$.

Label shift means $P_T(X\mid Y)=P_S(X\mid Y)$, with the same classifier applied in source and target. It follows that $P_T(\hat Y=i\mid Y=j)=C_{ij}$, so the target prediction frequencies obey a mixture equation. In practice, $C$ is estimated from labeled source data and $q$ from unlabeled target predictions; the derivations below first treat $C$ as known.

The Confusion-Matrix Moment Equation

By conditioning on the target true class, the probability of a predicted class is a mixture over the unknown target prevalences:

q_i=\sum_{j=1}^K P_T(\hat Y=i\mid Y=j)P_T(Y=j) =\sum_{j=1}^K C_{ij}\pi^T_j.

In vector notation, this is $q=C\pi^T$. The population equation is exact under label shift. If $C$ is square and nonsingular, it has the unique solution $\pi^T=C^{-1}q$. More generally, identification requires that different prevalence vectors produce different prediction distributions; for a $K$-class square matrix, nonsingularity is the corresponding condition.

Result: If $C$ is a $K\times K$ column-stochastic matrix and is nonsingular, then $q=C\pi^T$ has the unique solution $\pi^T=C^{-1}q$. This solution is a valid prevalence vector when its entries are nonnegative; its entries then sum to one.

To verify the normalization, let $\mathbf{1}$ denote the vector of $K$ ones. Column-stochasticity gives $\mathbf{1}^{\mathsf T}C=\mathbf{1}^{\mathsf T}$. Hence, for an exact population solution,

\mathbf{1}^{\mathsf T}q =\mathbf{1}^{\mathsf T}C\pi^T =\mathbf{1}^{\mathsf T}\pi^T.

Since $q$ is a probability vector, both sides equal one. This establishes the sum constraint, but it does not ensure nonnegativity if an empirical frequency vector is substituted for $q$. Sampling noise can put the unconstrained inverse outside the simplex. In that case, a principled estimate solves a constrained problem, such as minimizing $\|\hat q-C\pi\|^2$ over $\pi_j\geq0$ and $\sum_j\pi_j=1$, or maximizing the prediction likelihood over that simplex. Renormalizing a vector with negative entries does not make it a valid solution to the moment equation.

Worked Example: Recovering Two Target Prevalences

Suppose the source confusion matrix and the target prediction frequencies are

C= \begin{pmatrix} 0.8&0.2\\ 0.2&0.8 \end{pmatrix}, \qquad \hat q= \begin{pmatrix} 0.62\\ 0.38 \end{pmatrix}.

Let $\pi^T_1=p$, so $\pi^T_2=1-p$. The first predicted-class frequency satisfies

0.62=0.8p+0.2(1-p)=0.2+0.6p, \qquad p=\frac{0.62-0.2}{0.6}=0.7.

The second prevalence is $1-0.7=0.3$. Checking both predicted frequencies gives $0.8(0.7)+0.2(0.3)=0.62$ and $0.2(0.7)+0.8(0.3)=0.38$. The estimate is normalized, nonnegative, and reproduces the observed prediction frequencies exactly.

EM for Unlabeled Predicted Classes

The same model gives a likelihood for the target predicted labels. Suppose $m$ unlabeled target observations produce predicted-class counts $n_1,\ldots,n_K$, with $\sum_i n_i=m$. For a candidate prevalence vector $\pi$ in the simplex, the probability of predicted class $i$ is $r_i(\pi)=\sum_j C_{ij}\pi_j$. Ignoring a multinomial coefficient that does not depend on $\pi$, the log-likelihood is

\ell(\pi)=\sum_{i=1}^K n_i\log r_i(\pi).

To derive EM, imagine that each predicted label also had an unobserved generating class $j$. Given an observation predicted as class $i$, its current conditional probability of coming from class $j$ is

\tau_{ij}(\pi) =P(Y=j\mid \hat Y=i;\pi) =\frac{C_{ij}\pi_j}{r_i(\pi)}.

The E-step computes these responsibilities using the current prevalence estimate. The M-step replaces each prevalence by its expected fraction of observations assigned to that class. Thus the update is

\pi_j^{\mathrm{new}} =\frac{1}{m}\sum_{i=1}^K n_i\tau_{ij}(\pi) =\pi_j\sum_{i=1}^K\frac{n_i}{m}\frac{C_{ij}}{r_i(\pi)}.

This is an update of the prevalence vector, not of the confusion matrix. Its normalization follows directly from the responsibilities:

\sum_{j=1}^K\pi_j^{\mathrm{new}} =\frac{1}{m}\sum_i n_i\sum_j\tau_{ij}(\pi) =\frac{1}{m}\sum_i n_i=1.

Each updated entry is nonnegative, so a simplex-valued starting estimate remains on the simplex. In applications, the empirical matrix $\hat C$ and frequencies $\hat q_i=n_i/m$ are substituted into the update. The method then estimates prevalences conditional on that estimated confusion matrix.

Worked Example: Two EM Iterations

Use the matrix and target frequencies from the previous example, and start at $\pi=(0.5,0.5)^{\mathsf T}$. Initially $r_1=r_2=0.5$. The first component after one EM update is

\pi^{\mathrm{new}}_1 =0.5\left(\frac{0.62(0.8)}{0.5} +\frac{0.38(0.2)}{0.5}\right) =0.5(0.992+0.152)=0.572.

The second component is $1-0.572=0.428$, as guaranteed by normalization. For the next update, $r_1=0.8(0.572)+0.2(0.428)=0.5432$ and $r_2=0.4568$. Therefore,

\pi^{\mathrm{new}}_1 =0.572\left(\frac{0.62(0.8)}{0.5432} +\frac{0.38(0.2)}{0.4568}\right) \approx 0.572(0.9131+0.1664) \approx 0.6175.

The first two updates move from $0.5$ to $0.572$ to approximately $0.6175$, toward the moment solution $0.7$. With finite data, the EM estimate is the likelihood-based estimate for the observed counts; it need not match an unconstrained moment inverse if empirical frequencies or the estimated confusion matrix make that inverse infeasible.

Why the EM Likelihood Does Not Decrease

The EM update has a general monotonicity property in this mixture model. Assume $C_{ij}>0$ for all $i,j$, and start from a prevalence vector with strictly positive entries. Then every $r_i(\pi)>0$, and every responsibility is defined. The M-step remains strictly inside the simplex.

Theorem: Under these conditions, each EM update satisfies $\ell(\pi^{\mathrm{new}})\geq\ell(\pi)$. Consequently, the sequence of log-likelihood values is nondecreasing and converges to a finite limit.

Proof. For each observed predicted class $i$, use the current responsibilities $\tau_{ij}=\tau_{ij}(\pi)$. The ratio of its new and old probabilities can be written as

\frac{r_i(\pi^{\mathrm{new}})}{r_i(\pi)} =\sum_j\tau_{ij}\frac{\pi_j^{\mathrm{new}}}{\pi_j}.

The logarithm is concave, so Jensen’s inequality gives

\ell(\pi^{\mathrm{new}})-\ell(\pi) \geq \sum_i n_i\sum_j\tau_{ij} \log\frac{\pi_j^{\mathrm{new}}}{\pi_j}.

By the M-step definition, $\sum_i n_i\tau_{ij}=m\pi_j^{\mathrm{new}}$. The right-hand side is therefore

m\sum_j\pi_j^{\mathrm{new}} \log\frac{\pi_j^{\mathrm{new}}}{\pi_j} =m\,\mathrm{KL}(\pi^{\mathrm{new}}\|\pi)\geq0.

The last inequality follows from $-\log x\geq1-x$: applying it to $x=\pi_j/\pi_j^{\mathrm{new}}$, multiplying by $\pi_j^{\mathrm{new}}$, and summing gives at least $\sum_j(\pi_j^{\mathrm{new}}-\pi_j)=0$. Finally, each $r_i$ is a probability and hence at most one, so $\ell(\pi)\leq0$. The nondecreasing log-likelihood sequence is bounded above and therefore converges. This proves monotonicity and convergence of the likelihood values; by itself, it does not establish a fast rate of convergence of the prevalence iterates. $\square$

Near-Singular Separation and Practical Limits

Identification in the population is not the same as stable estimation from finite data. When two columns of $C$ are close, the corresponding classes produce nearly the same distribution of predicted labels. Small changes in $\hat q$, caused by sampling variation, can then imply large changes in the recovered prevalences. In the square case, this is reflected by a small singular value of $C$: the inverse amplifies errors in directions associated with that singular value.

Worked Example: A Small Frequency Change, a Large Prevalence Change

Consider the binary confusion matrix

C= \begin{pmatrix} 0.51&0.49\\ 0.49&0.51 \end{pmatrix}.

If the first target prevalence is $p$, the first predicted-class frequency is $q_1=0.51p+0.49(1-p)=0.49+0.02p$. For $q_1=0.506$, inversion gives $p=(0.506-0.49)/0.02=0.8$. If the estimated frequency is instead $0.507$, the estimate is $p=(0.507-0.49)/0.02=0.85$. A change of only $0.001$ in the predicted frequency changes the estimated prevalence by $0.05$. In this matrix, the difference between the columns is small, so the target prevalence is weakly informed by the predictions.

At the limiting case $C_{ij}=0.5$ for every $i,j$ in the binary problem, the two columns are identical. Every prevalence produces $q=(0.5,0.5)^{\mathsf T}$, so the target prevalence is not identifiable from predicted labels. The EM update makes the same failure visible: for any interior starting prevalence, $r_i=0.5$, and the update returns $\pi_j^{\mathrm{new}}=\pi_j$. The likelihood contains no information that could move the estimate toward a uniquely determined answer.

Common Mistake: Treating Every Inverse as a Prevalence

A common mistake is to invert the empirical moment equation and assume the result must be a valid target prevalence. For example, with the well-separated binary matrix from the earlier calculations, an empirical frequency $\hat q=(0.90,0.10)^{\mathsf T}$ gives $\hat p=(0.90-0.20)/0.60=7/6$, and the second component is $-1/6$. Those values sum to one, but they are not probabilities. The error is to treat the exact population identity as if every noisy empirical frequency had to lie in the set of mixtures generated by $C$. A valid estimate must lie on the simplex; constrained estimation is needed when the unconstrained solution does not. Even a valid constrained answer may be unstable when the class-conditional prediction patterns are nearly indistinguishable.

Key takeaway: Moment inversion and EM use the same mixture relationship, $q=C\pi^T$. The moment method solves that relationship directly; EM maximizes the predicted-label likelihood while preserving a normalized prevalence vector. Check both feasibility and class separation: normalization does not guarantee stable or identifiable estimates.

Check Your Understanding

Use the column convention for the confusion matrix throughout.

Explain why each column of $C$, rather than each row, sums to one.
For a nonsingular column-stochastic $C$, show how the identity $\mathbf{1}^{\mathsf T}C=\mathbf{1}^{\mathsf T}$ checks the sum of the moment solution.
Why does an EM update preserve the sum of the prevalence components?
What does monotonicity of the EM likelihood establish, and what does it not establish by itself?
Describe what happens to prevalence estimation when two columns of $C$ become identical.