General Quantification of Covariate and Concept Shifts

General Quantification of Covariate and Concept Shifts§

Distribution shift is a fundamental obstacle in deploying ML models reliably. Existing learning bounds often assume ideal conditions (e.g., overlapping supports) and are not estimable from data. This paper introduces a novel framework to quantify and bound generalization error under broad distribution shifts, bridging theory and practice.

Core Contributions§

• Identifies a flaw in standard concept shift definition: When source and target supports mismatch (i.e., no overlap), the classical definition of concept shift $f_S(x) \neq f_T(x)$ becomes ill-defined because $f_T$ may not be defined on the source support.
• *Introduces $\gamma^$-concept shift*: Using entropic optimal transport (EOT), they define a relaxed notion that quantifies shift even with disjoint supports. The key idea is to couple source and target distributions via an entropic optimal plan $\pi^_\gamma$, and then measure the conditional label discrepancy under this plan.
• Derives a general error bound: The bound unifies covariate shift (shift in $P(X)$) and $\gamma^*$-concept shift (shift in $P(Y|X)$), applicable to arbitrary loss functions, label spaces (including stochastic labeling), and model classes.
• Develops estimators with concentration guarantees: They propose sample-based estimators for both shifts and the error bound, with finite-sample concentration bounds.
• DataShifts algorithm: A practical tool that takes source and target data, estimates the shifts, and outputs a rigorous error bound for a given model.

Mathematical Formulation§

Let $P_S$ and $P_T$ be source and target distributions over $\mathcal{X} \times \mathcal{Y}$. The entropic optimal transport plan between $P_S$ and $P_T$ is: $$\pi^_\gamma = \arg\min_{\pi \in \Pi(P_S, P_T)} \mathbb{E}_{(x,y)\sim\pi}[c((x,y),(x',y'))] - \gamma H(\pi),$$ where $c$ is a cost function (e.g., $\|x-x'\|^2 + \infty \cdot \mathbb{1}_{y\neq y'}$) and $H$ is entropy. The $\gamma^$-concept shift is then defined as the expected conditional label difference under $\pi^_\gamma$. The error bound decomposes as: $$\mathcal{E}_T(f) \leq \mathcal{E}_S(f) + \text{CovariateShift} + \gamma^\text{-ConceptShift} + \text{Approximation},$$ where each term is estimable from finite samples.

Implementation and Estimators§

Estimating EOT plans requires solving a regularized optimal transport problem, often via the Sinkhorn algorithm. The shift estimators leverage sample averages under the optimal coupling. The DataShifts algorithm (pseudocode) involves:

1. Compute EOT plan $\pi^*_\gamma$ between source and target samples using Sinkhorn.
2. Estimate covariate shift as the Wasserstein distance or discrepancy under $\pi^*_\gamma$.
3. Estimate $\gamma^$-concept shift as the conditional label discrepancy under $\pi^_\gamma$.
4. Combine with empirical source error to bound target error.

Below is a simplified PyTorch snippet illustrating EOT computation and shift estimation:

import torch
from geomloss import SamplesLoss  # requires geomloss library

def estimate_gamma_star_concept_shift(X_s, y_s, X_t, y_t, model, gamma=1.0):
    """
    Estimate gamma*-concept shift using entropic optimal transport.
    X_s, y_s: source samples; X_t, y_t: target samples
    model: trained model f
    """
    # Compute predictions
    f_s = model(X_s)
    f_t = model(X_t)
    
    # EOT plan via Sinkhorn (using geomloss)
    loss = SamplesLoss("sinkhorn", p=2, blur=gamma)
    pi = loss(X_s, X_t)  # returns OT plan implicitly; need custom to get plan
    # For demonstration, assume pi is a matrix of size (n_s, n_t)
    # In practice, use sinkhorn algorithm directly to get plan.
    
    # Compute conditional label discrepancy under plan
    # For regression: squared difference; for classification: cross-entropy
    diff = (f_s.unsqueeze(1) - f_t.unsqueeze(0)).pow(2)  # pairwise
    shift = (pi * diff).sum()
    return shift

Practical Implications§

The method enables practitioners to quantify distribution shift and obtain a certified upper bound on target error without target labels. It applies to diverse domains: domain adaptation, causal inference, and robust ML. Limitations include computational cost of EOT for large datasets (mitigated by stochastic solvers) and the need to choose the entropic regularization $\gamma$.