Label Shift Prior Correction

Label shift occurs when the marginal distribution of labels P(Y) changes between the training (source) environment and the testing (target) environment, while the class-conditional distribution of features P(X|Y) remains invariant. If a model predicts a calibrated posterior q = P_train(Y=1|X) under the training distribution, we can correct it to reflect the target distribution by reweighting it based on the changed class priors.

Worked Example

Suppose we have a binary classification task. Let the training prior for class 1 be π_train = 0.5, and the target prior be π_target = 0.25. Thus, the prior for class 0 changes from 0.5 in training to 0.75 in the target domain. Assume an observation has a calibrated training posterior q = 0.75 for class 1.

Calculate Class Weights: We define weights as the ratio of target to training priors for each class.
w₁ = π_target / π_train = 0.25 / 0.5 = 0.5
w₀ = (1 - π_target) / (1 - π_train) = 0.75 / 0.5 = 1.5
Calculate the Normalizer: The normalizer adjusts the evidence based on the new class weights.
Normalizer = w₁ · q + w₀ · (1 - q)
Normalizer = 0.5 · 0.75 + 1.5 · 0.25 = 0.375 + 0.375 = 0.75
Calculate Target Posterior: The corrected posterior probability for class 1 is the weighted training posterior divided by the normalizer.
Target Posterior = (w₁ · q) / Normalizer = 0.375 / 0.75 = 0.5

Experiment

You can test the effect of changing class priors on the posterior using the Label Shift Prior Correction module.

Set Training prior to 0.5.
Set Target prior to 0.25.
Set Training posterior to 0.75.
Observe that the calculated Target posterior is 0.5, which matches our worked example, showing that the reduced prevalence of class 1 in the target domain decreases its probability for this observation.

Limitations

This correction strictly assumes that the class-conditional distributions P(X|Y) are invariant between the training and target environments, and that the original training posterior is perfectly calibrated. If P(X|Y) changes (e.g., in covariate shift or concept drift), this method is insufficient. Furthermore, if a class has zero probability in the target domain, the normalizer can become zero for certain inputs (where the only possible class given the training posterior has a zero target prior). In such cases, the observation is impossible under the target distribution, and the posterior is undefined rather than simply clipped to 0 or 1.

References

The following reference provides a rigorous formulation of label shift and its correction via marginal reweighting:

Lipton, Z. C., Wang, Y. X., & Smola, A. (2018). Detecting and Correcting for Label Shift with Black Box Predictors. Read paper. This paper formally discusses label shift, demonstrating how to use the target class marginals to correct black-box predictions.

Label Shift Prior Correction