02. Distribution Shift: A Practical Introduction

• Primary References:

1. 1
   Moreno-Torres, J. G., Raeder, T., Alaiz-Rodriguez, R., Chawla, N. V., & Herrera, F. (2012). A Unifying View on Dataset Shift in Classification. Pattern Recognition, 45(1), 521–530.
2. 2
   Shimodaira, H. (2000). Improving Predictive Inference Under Covariate Shift by Weighting the Log-Likelihood Function. Journal of Statistical Planning and Inference, 90(2), 227–244.
3. 3
   Lipton, Z. C., Wang, Y. X., & Smola, A. (2018). Detecting and Correcting for Label Shift with Black Box Predictors (BBSE). In ICML 2018.
4. 4
   Liu, W., Wang, X., Owens, J., & Li, Y. (2020). Energy-based Out-of-Distribution Detection. In NeurIPS 2020.
5. 5
   Koh, P. W., et al. (2021). WILDS: A Benchmark of in-the-Wild Distribution Shifts. In ICML 2021.

• Why I Wrote This Note: While developing models for smart grid stability prediction (IEEE ICCIT 2024) and chest X-ray lung segmentation (TIG-UNet++, ICICT 2025), I repeatedly encountered accuracy drops when testing across unseen operating regimes or external hospital datasets. I needed a unified diagnostic note that clarifies which factor of \(p(x, y)\) shifted and which mathematical tool can actually fix it.

• Joint Factorization Dictates the Taxonomy: Any joint distribution factors either as \(p(x, y) = p(y \mid x)p(x)\) (causal/discriminative factorization \(\mathcal{X} \to \mathcal{Y}\)) or as \(p(x, y) = p(x \mid y)p(y)\) (anti-causal/generative factorization \(\mathcal{Y} \to \mathcal{X}\)).
• Covariate Shift (\(p(x)\) changes, \(p(y\mid x)\) invariant): Input feature statistics drift (e.g., X-ray scanner exposure or lighting in textile defect cameras), but the ground-truth diagnostic label for a given physical state remains unchanged.
• Label / Prior Shift (\(p(y)\) changes, \(p(x\mid y)\) invariant): Class prevalence changes between train and test (e.g., tuberculosis prevalence in a tertiary referral center vs. community screening), while class-conditional symptom manifestation \(p(x\mid y)\) stays constant.
• Concept Shift (\(p(y\mid x)\) or \(p(x\mid y)\) changes): The underlying functional mapping between inputs and targets changes over time (e.g., consumer demand dynamics or grid control policies after adding renewable microgrids). No unlabeled test-time trick can fully recover performance without new supervision or structural priors.

4.1 Formal Taxonomy via Joint Factorizations

Let \(p_S(x, y)\) denote the source (training) joint distribution and \(p_T(x, y)\) denote the target (deployment) joint distribution. Using Bayes' rule, we examine two canonical factorizations:

flowchart TD
    J["Joint Shift: p_S(x, y) ≠ p_T(x, y)"] --> F1["Factorization 1: p(y|x) p(x)"]
    J --> F2["Factorization 2: p(x|y) p(y)"]
    F1 --> CS["Covariate Shift<br/>p_S(x) ≠ p_T(x)<br/>p_S(y|x) = p_T(y|x)"]
    F1 --> CON1["Concept Drift (Discriminative)<br/>p_S(y|x) ≠ p_T(y|x)<br/>p_S(x) = p_T(x)"]
    F2 --> LS["Label (Prior) Shift<br/>p_S(y) ≠ p_T(y)<br/>p_S(x|y) = p_T(x|y)"]
    F2 --> CON2["Class-Conditional Shift<br/>p_S(x|y) ≠ p_T(x|y)<br/>p_S(y) = p_T(y)"]

1. Covariate Shift

Defined by \(p_S(x) \neq p_T(x)\) while the conditional label relationship remains invariant: \(p_S(y \mid x) = p_T(y \mid x)\). If a model \(f_\theta\) were Bayes-optimal and had infinite capacity, covariate shift would not require retraining. However, for misspecified or finite-capacity neural networks trained via ERM, the source risk \(\mathcal{R}_S(\theta)\) does not match the target risk \(\mathcal{R}_T(\theta)\). Importance Weighting (Shimodaira, 2000) restores unbiased target risk estimation provided \(\text{supp}(p_T) \subseteq \text{supp}(p_S)\):

2. Label Shift (Target Prior Shift)

Defined in anti-causal settings (\(\mathcal{Y} \to \mathcal{X}\), where diseases cause symptoms) by \(p_S(y) \neq p_T(y)\) while \(p_S(x \mid y) = p_T(x \mid y)\). Notice that under label shift, the discriminative posterior \(p(y\mid x)\) necessarily changes via Bayes' rule! Given a calibrated source classifier \(p_S(y \mid x)\) and an estimate of the target prior \(p_T(y)\) (obtained via Black-Box Shift Estimation, BBSE, using the source confusion matrix \(C_{\hat{y}, y}\)), we recalibrate predictions at test time without retraining:

3. Concept Shift (Concept Drift)

Defined by \(p_S(y \mid x) \neq p_T(y \mid x)\) even for identical feature vectors \(x\). The ground-truth decision boundary itself has moved. For example, in smart grid stability prediction, if new inverter-based solar generation alters the damping dynamics of the grid, the same load and frequency vector \(x\) that was previously stable (\(y=1\)) may now lead to oscillatory instability (\(y=0\)).

4.2 Diagnosing Covariate & OOD Shift in Representation Space

In high-dimensional vision or language models, estimating density ratios \(w(x) = p_T(x)/p_S(x)\) directly in pixel/token space fails catastrophically. Instead, I combine two lightweight feature/logit diagnostics:

1. 1
   Free Energy Score (Liu et al., 2020): Viewing logits \(f_\theta(x) \in \mathbb{R}^C\) through an energy-based model lens, the Helmholtz free energy is:

Unlike softmax confidence (which normalizes away logit magnitude), \(E(x; f_\theta)\) is strictly lower (more negative) for in-distribution inputs and higher for covariate-shifted / OOD inputs.

1. 1
   Feature-Space Mahalanobis Distance (Lee et al., 2018): Fitting class-conditional empirical means \(\hat{\mu}_c \in \mathbb{R}^D\) and a tied covariance matrix \(\hat{\Sigma} \in \mathbb{R}^{D \times D}\) on penultimate layer embeddings \(h(x)\):

import torch

def energy_ood_score(logits: torch.Tensor, temperature: float = 1.0) -> torch.Tensor:
    # Higher energy (-T * logsumexp) indicates greater covariate / OOD shift
    return -temperature * torch.logsumexp(logits / temperature, dim=-1)

def correct_label_shift_bbse(
    source_confusion: torch.Tensor,  # shape (C, C) where C[i, j] = P_S(y_hat=i | y=j)
    target_pred_hist: torch.Tensor,  # shape (C,) empirical P_T(y_hat=i) on unlabeled test batch
    source_priors: torch.Tensor,     # shape (C,) P_S(y=j)
    source_probs: torch.Tensor,      # shape (N, C) P_S(y | x) on test samples
) -> torch.Tensor:
    # Estimate target prior ratio w = p_T(y) / p_S(y) and apply Bayes prior adjustment
    target_priors = torch.linalg.pinv(source_confusion) @ target_pred_hist
    target_priors = target_priors.clamp_min(1e-6)
    target_priors = target_priors / target_priors.sum()

    importance_weights = target_priors / source_priors.clamp_min(1e-6)  # shape (C,)
    adjusted = source_probs * importance_weights.unsqueeze(0)
    return adjusted / adjusted.sum(dim=-1, keepdim=True)

Every domain I have published in exhibits a distinct member of this taxonomy:

• Medical Imaging (TIG-UNet++ Lung Segmentation, ICICT 2025 & Tuberculosis Screening, ICEEICT 2024): Cross-hospital deployment simultaneously introduces covariate shift (portable vs. fixed X-ray detectors changing pixel intensity histograms \(p(x)\)) and label shift (screening clinic TB prevalence of \(1.5\%\) vs. balanced \(50\%\) training datasets).
• Industrial Computer Vision (YOLOv8 + CNN Textile Defect Detection, ICCA 2024): Changing fabric dye lots and loom illumination creates pure covariate shift, where feature-space normalization and thresholded preprocessing stabilize recall.
• Smart Power Grids (Smart Grid Stability Prediction, ICCIT 2024): Evolving consumer price elasticity and distributed generation introduce concept drift, requiring rolling-window drift monitors rather than static batch normalization tweaks.

• Compound Shifts Are the Rule, Not the Exception: Academic benchmarks isolate covariate shift (e.g., CIFAR-10-C) or label shift (long-tailed LT benchmarks) separately. In clinical deployments, a rural clinic uses different imaging hardware (\(p(x \mid y)\) shifts) and sees a different patient demographic (\(p(y)\) shifts).
• Energy Outperforms Softmax for Covariate Drift: Because softmax is shift-invariant (\(\text{softmax}(\mathbf{z} + c\mathbf{1}) = \text{softmax}(\mathbf{z})\)), a sample whose feature norm collapses near the origin can still yield a peaked softmax if one logit is slightly less negative than the others. The Helmholtz energy \(E(x; f_\theta)\) immediately flags the collapsed logit magnitude.

• Implement Energy Score and Mahalanobis distance monitors for penultimate CNN/ViT embeddings.Verified
• Validate BBSE prior shift adjustment on imbalanced medical classification splits.Verified
• Run a compound shift ablation (simultaneous intensity corruption + prior skew) on chest radiograph benchmarks.Pending
• Coupling Energy-based covariate shift scores with Test-Time Adaptation gating so adaptation only triggers on mild covariate drift, not severe OOD noise.Pending

"Before asking how to make a neural network robust to distribution shift, we must first ask whether the world is causal (\(\mathcal{X} \to \mathcal{Y}\)) or anti-causal (\(\mathcal{Y} \to \mathcal{X}\))—because the remedy for one can actively harm the other."

Bayes Prior Adjustment Under Label Shift:

• Connected Research Notes:
• Understanding Predictive Uncertainty — Why ECE degrades under covariate shift and how epistemic mutual information reacts to OOD inputs.
• Test-Time Adaptation — How to adapt batch normalization statistics and affine parameters online when \(p_T(x)\) drifts.
• Risk-Constrained AI Decision-Making — Using energy and Mahalanobis scores to abstain when test inputs leave the training support.