PICSC: Correcting Partitioning-Induced Covariate Shift

Official implementation of PICSC (Partitioning-Induced Covariate Shift Correction), a Fisher-information-based regularization framework that detects and corrects covariate shift induced by data partitioning — whether across batches/folds of a centrally held dataset, or across federated clients/nodes where data is natively non-colocated.

Data distributions are rarely stationary. PICSC continually re-estimates and corrects for the shift introduced by how data is partitioned, rather than assuming the distribution used at the start of training still holds.

Overview

Fragmenting a dataset for cross-validation, or training across distributed clients, induces a measurable covariate shift even when the underlying population is stationary. PICSC quantifies this shift using the Fisher Information Matrix (FIM), via a Cramér–Rao-anchored approximation, and adds it as a regularization term during training. The same estimator applies unmodified to both instantiations:

• Batch/fold instantiation — an unconditional, all-pairs penalty (Algorithm 1) suited to small, bounded partition counts.
• Client/node instantiation — a conditional detection-and-adaptation mechanism (Algorithm 2) that triggers correction only when a composite Fisher–KL signal exceeds a calibrated threshold, giving constant memory and communication cost per client.

Only fragment-local gradient statistics are required — never raw data — making the method directly compatible with federated privacy constraints.

Repository structure

.
├── picsc/                  # Core method: FIM estimation, penalty, conditional trigger
├── baselines/               # IW, IWCV, KMM, DIW, FedAvg, FedProx, SCAFFOLD
├── configs/                 # Experiment configs (batch/fold, k-fold, federated)
├── datasets/                 # Dataset loaders and partitioning utilities
├── experiments/              # Scripts reproducing paper tables (E1–E6)
├── results/                  # Logged outputs, checkpoints
└── requirements.txt

Installation

git clone https://github.com/<your-org>/picsc.git
cd picsc
pip install -r requirements.txt

Requires Python ≥3.9 and TensorFlow 2.11.

Usage

Batch/fold correction:

python experiments/run_batchwise.py --dataset cifar10 --num_batches 10 --lambda 0.1

k-fold correction:

python experiments/run_foldwise.py --dataset keel/australian --k 5 --lambda 0.1

Federated correction:

python experiments/run_federated.py --dataset femnist --clients 100 --alpha 0.9 --gamma 0.05

Reproducing paper results

Each script under experiments/ corresponds to an experiment (E1–E6) in the paper. See configs/ for exact hyperparameters used to produce the reported tables.

Citation

If you use this code, please cite:

@article{khan2026picsc,
  title   = {PICSC: Correcting Partitioning-Induced Covariate Shift},
  author  = {Khan, Behraj and Mirza, Behroz and Syed, Tahir},
  journal = {TBD},
  year    = {2026}
}

License

Contact

Behraj Khan — behrajkhan@gmail.com