Conditional Probability Shift Model

• CPSM is a framework that addresses shifts in P(Y|Z) while keeping the downstream distribution P(W|Z,Y) invariant, bridging covariate and label shift scenarios.
• It employs multinomial logistic regression to parameterize target posteriors and uses an EM algorithm to estimate the local tilt even with missing target labels.
• Empirical evaluations show that CPSM significantly improves balanced accuracy and reduces approximation error compared to traditional label shift methods in various settings.

The Conditional Probability Shift Model (CPSM) addresses a nuanced regime of distributional shift in supervised learning, generalizing beyond the classical covariate shift and label shift frameworks. In CPSM, the conditional distribution of the class label given a subset of features (“conditioning” features) is allowed to change from source to target, while the conditional distribution of the remaining features, given the class and conditioning features, remains invariant. CPSM provides a principled methodology for recovering the posterior class probabilities on the target domain, relying on parameterization via multinomial logistic regression and parameter estimation through the Expectation-Maximization (EM) algorithm in the presence of missing target labels (Teisseyre et al., 4 Mar 2025).

1. Problem Definition and Conditional Probability Shift Assumption

Let $X \in \mathbb{R}^p$ denote the observed feature vector and $Y \in \{1, \ldots, K\}$ the class label. CPSM partitions $X = (Z, W)$, where $Z \in \mathbb{R}^d$ includes the “conditioning” features suspected of experiencing a shift in $P(Y|Z)$, and $W \in \mathbb{R}^{p-d}$ consists of all remaining features. The model presumes access to labeled samples from a source distribution $P_S(Z, W, Y)$ and only unlabeled samples from a target marginal $P_T(Z, W)$.

The central CPS assumption posits:

• $P_S(W|Z,Y) = P_T(W|Z,Y)$ (invariance of the “anti-causal” slice)
• $P_S(Y|Z) \neq P_T(Y|Z)$ (shift of the class-conditional on $Y \in \{1, \ldots, K\}$0)

Joint distributions factor as: $Y \in \{1, \ldots, K\}$1

$Y \in \{1, \ldots, K\}$2

Selecting $Y \in \{1, \ldots, K\}$3 recovers classic covariate shift ($Y \in \{1, \ldots, K\}$4, $Y \in \{1, \ldots, K\}$5). Setting $Y \in \{1, \ldots, K\}$6 trivializes to label shift ($Y \in \{1, \ldots, K\}$7, $Y \in \{1, \ldots, K\}$8). CPSM thus interpolates between these domains, uniquely capturing situations where $Y \in \{1, \ldots, K\}$9 varies while “downstream” structure remains stable.

2. Posterior Recovery and Model Parameterization

The inferential target is $X = (Z, W)$0. By Bayes' rule and the invariance of $X = (Z, W)$1, the target posterior is expressed as: $X = (Z, W)$2 where $X = (Z, W)$3 is the local “tilt” of the class-conditional.

This formulation reduces the adaptation task to that of learning the set of ratios $X = (Z, W)$4.

CPSM parameterizes $X = (Z, W)$5 as a multinomial logistic regression: $X = (Z, W)$6

$X = (Z, W)$7

for $X = (Z, W)$8 and some feature mapping $X = (Z, W)$9. The true $Z \in \mathbb{R}^d$0 is assumed to be in this family for some $Z \in \mathbb{R}^d$1.

3. Expectation-Maximization Estimation

Since target labels are unobserved, parameter estimation proceeds by maximizing the observed target marginal likelihood using the EM algorithm.

E-step: At iteration $Z \in \mathbb{R}^d$2, compute the pseudo-responsibilities

$Z \in \mathbb{R}^d$3

with $Z \in \mathbb{R}^d$4.

M-step: Update $Z \in \mathbb{R}^d$5 by maximizing

$Z \in \mathbb{R}^d$6

This coincides with the weighted multinomial logistic regression log-likelihood, solvable by Newton–Raphson or stochastic gradient methods. The procedure is initialized at $Z \in \mathbb{R}^d$7 (no shift).

Standard EM theory ensures that the observed likelihood is non-decreasing. Under regularity conditions, the procedure converges to a stationary point of the likelihood.

4. Integration with Probabilistic Classifiers

CPSM requires, from labeled source data:

• $Z \in \mathbb{R}^d$8 (“posterior model”)
• $Z \in \mathbb{R}^d$9 (“marginal model”)

These posteriors can be estimated using any probabilistic classifier (e.g., logistic regression, random forest, neural network) trained on source data. Typically:

• A classifier predicts $P(Y|Z)$0 from $P(Y|Z)$1 to provide $P(Y|Z)$2.
• A separate classifier predicts $P(Y|Z)$3 from $P(Y|Z)$4 only to estimate $P(Y|Z)$5.

These models supply the necessary terms for the E-step. Upon convergence of EM and estimation of $P(Y|Z)$6, the target posterior $P(Y|Z)$7 is reconstructed via the closed-form correction formula provided above.

5. Special Cases and Identifiability

• If $P(Y|Z)$8 is constant, $P(Y|Z)$9 reduces to $W \in \mathbb{R}^{p-d}$0, and CPSM collapses to the classical label shift model, with $W \in \mathbb{R}^{p-d}$1 modeling the class priors.
• If $W \in \mathbb{R}^{p-d}$2 (i.e., $W \in \mathbb{R}^{p-d}$3), $W \in \mathbb{R}^{p-d}$4, yielding the covariate shift regime with unchanged posterior mapping.
• Identifiability of $W \in \mathbb{R}^{p-d}$5 is ensured provided the feature transformation $W \in \mathbb{R}^{p-d}$6 spans $W \in \mathbb{R}^{p-d}$7 and classes are not perfectly separated. Under these mild assumptions, the EM estimator is consistent as $W \in \mathbb{R}^{p-d}$8.

6. Empirical Evaluation

Synthetic Studies

Two generative settings for $W \in \mathbb{R}^{p-d}$9 were examined: binary vectors and multivariate Gaussians. $P_S(Z, W, Y)$0 with $P_S(Z, W, Y)$1 constant and $P_S(Z, W, Y)$2 parameterized as a logistic function of $P_S(Z, W, Y)$3. By modulating parameters, the following regimes were generated: no shift, label shift, CPS only, and joint shift. Additional features $P_S(Z, W, Y)$4 were sampled conditionally according to the invariant $P_S(Z, W, Y)$5.

Metrics included balanced accuracy and the absolute error $P_S(Z, W, Y)$6. CPSM matched classical label shift methods (BBSC, MLLS) under pure label shift, but outperformed all methods—especially Sparse-Joint-Shift (SEES)—in all settings featuring conditional shifts ($P_S(Z, W, Y)$7 but no marginal label shift).

MIMIC Medical Case Study

On the MIMIC medical database, the tasks involved predicting four disease indicators from laboratory/vital data, with $P_S(Z, W, Y)$8 chosen as age category or gender. Source data was sampled to maintain constant $P_S(Z, W, Y)$9; the target was perturbed to induce a conditional shift $P_T(Z, W)$0 for $P_T(Z, W)$1. Marginal label distributions could remain unchanged or shift.

Logistic regression and three-layer neural networks were used as base classifiers. Competing approaches encompassed NAIVE (no correction), BBSC and MLLS (label shift correction), ExTRA (exponential tilt for joint shift), and SEES. Results established that:

• In pure conditional shift (no label shift): NAIVE, BBSC, and MLLS did not improve target performance.
• CPSM consistently improved balanced accuracy by 10–20 points and minimized approximation error in nearly all settings.
• CPSM outperformed SEES (the only other SJS-based method) in terms of approximation error and balanced accuracy.

7. Context, Extensions, and Implications

CPSM provides a middle ground between covariate and label shift adaptations, rigorously addressing shifts that are localized in conditional class distributions yet cannot be ascribed solely to priors or marginal features. Its generic EM-based implementation and compatibility with arbitrary probabilistic classifiers make it broadly applicable in machine learning pipelines facing distributional changes that violate stronger shift assumptions. Empirical evidence suggests distinct advantages over prior label shift and joint shift correction methods in settings where conditional structure dominates and prior-based corrections are insufficient. A plausible implication is increased robustness to “hidden stratification” in clinical or high-dimensional scientific datasets, where subpopulation-specific prevalence can change without altering aggregate class proportions (Teisseyre et al., 4 Mar 2025).