fed-clma 0.1.0

Collaborative Model Adaptation: federated learning under distributed concept and covariate drift

CLMA — Collaborative Model Adaptation

Federated learning under distributed concept and covariate drift, for any data modality.

A reference implementation of Collaborative Drift Compensation (Adarsh N L, Madapu Amarlingam, Divyasheel Sharma — CODS-COMAD '24), generalised from the paper's tabular case to work with tabular, image, text or any other data.

---

The problem

Standard federated learning assumes client data is homogeneous and stationary. In practice — especially in industrial settings — it is neither. Different clients drift in different ways at different times, and FedAvg blends all of it into one global model. A single drifted client can poison every other client's model.

Existing drift-aware approaches either assume all clients drift the same way, or cluster clients by their loss — a scalar, so two clients with completely different drifts can look identical.

The approach

CLMA infers drift structure from model parameters, which carry far more information than a scalar loss:

1. Detect — each client flags drift locally when its loss rises by more than γ over the previous round. One bit, sent alongside the weights it was already sending.
2. Quarantine — flagged clients are detached from their cohort so they cannot contaminate stable clients.
3. Cohort — the server runs a signed KS test of each quarantined client's parameters against the reference model to determine the direction of the shift, projects each direction group onto its principal components, and clusters them. Clients that drifted the same way end up together.
4. Adapt — each cohort trains its own model.

No extra data, metrics, or round-trips are required from clients. The entire cohorting procedure runs server-side on parameters that were being transmitted anyway.

Install

pip install clma                  # core
pip install "clma[torch]"         # + PyTorch helpers
pip install "clma[tensorflow]"    # + TensorFlow/Keras helpers
pip install "clma[simulation]"    # + Flower simulation runtime

Requires Python ≥ 3.10 and flwr >= 1.32 (the message-based Strategy API).

Usage

CLMA is a drop-in Flower strategy. Swap FedAvg for CLMA:

from clma import CLMA
from flwr.serverapp import ServerApp

app = ServerApp()

@app.main()
def main(grid, context):
    strategy = CLMA()
    result = strategy.start(
        grid=grid,
        initial_arrays=ArrayRecord(model.state_dict()),
        num_rounds=10,
    )
    # Which client sat in which cohort, per round (the data behind Figs. 2-3):
    print(strategy.cohort_history)

On the client, the only new requirement is reporting the drift flag. DriftReporter handles both the detection and the reply:

from clma.client import DriftReporter
from flwr.clientapp import ClientApp

app = ClientApp()

@app.train()
def train(msg, context):
    model.load_state_dict(msg.content["arrays"].to_torch_state_dict())
    loss, num_examples = train_one_round(model, trainloader)

    return DriftReporter(context).build_reply(
        msg,
        arrays=ArrayRecord(model.state_dict()),
        loss=loss,
        num_examples=num_examples,
    )

That is the entire contract. Your model, data loader, optimiser and modality are untouched — CLMA never sees a raw feature.

Any modality

Because the algorithm only ever reads parameter vectors, it does not care what produced them. The one thing worth tuning for larger models is which parameters feed the drift signal:

from clma import CLMA, LastLayers
from clma.adapters.torch import trainable_selector

# Tabular / small CNN — the paper's setting. Defaults are fine.
CLMA()

# Vision or language backbone — restrict the signal to the head, so the KS test
# isn't swamped by millions of near-static backbone weights.
CLMA(param_selector=LastLayers(2))

# Fine-tuning a frozen backbone — ask the model which parameters actually move.
CLMA(param_selector=trainable_selector(model))

The default selector uses every learnable parameter but excludes tracked buffers (BatchNorm running statistics, num_batches_tracked). Those track input scale rather than the learned function, so including them lets covariate shift drown out concept drift.

Post-deployment drift (Algorithm 4)

from clma.client import InferenceDriftMonitor

monitor = InferenceDriftMonitor()
for batch in inference_stream:
    if monitor.observe(loss_fn(model(batch.x), batch.y)):
        trigger_clma_round()   # your transport
        monitor.reset()

Results on CIFAR-10

examples/cifar10_drift compares CLMA against FedAvg on real data: 10 clients, IID CIFAR-10 shards, a 545k-parameter CNN, and concept drift (label rotation) injected at rounds 3, 5, 7 and 9 into 8 of 10 clients, split into two opposite-direction families. Identical client code in both arms — only the strategy differs.

	CLMA	FedAvg
Stationary clients, final round	0.649	0.058
Stationary clients, mean rounds 4–10	0.632	0.176
Drifted clients, final round	0.514	0.342
Drift families correctly separated	yes, every event	n/a

The two clients whose data never changes end at 0.649 under CLMA versus 0.058 under FedAvg — the latter below the 0.10 random baseline, because averaging across incompatible concepts teaches the global model the rotated label mapping. CLMA's stationary clients improve monotonically through every drift event, as if the drifting clients were not there.

The two drift families are separated correctly at every drift event, even though both lose a similar amount of accuracy and are therefore near-indistinguishable by scalar loss — the signal comes entirely from parameter space.

This example also exercises the scalability fix on a real model: the drift signal is 545,290-dimensional, so the paper's literal eig(SᵀS) would need a 545,290 × 545,290 matrix (~2.4 TB). The Gram trick makes it 10 × 10.

Configuration

Option	Default	Purpose
cohorter	CLMACohorter()	Algorithm 3. Swap for your own Cohorter.
param_selector	DefaultSelector()	Which parameters carry the drift signal.
merge_cohorts_threshold	None (off)	Merge cohorts whose models reconverge.

And on CLMACohorter:

Option	Default	Purpose
clusterer	AutoKMeans()	K chosen by silhouette; use fixed_k to pin it.
variance_threshold	0.95	Variance the retained components must explain.
center	True	Mean-center before PCA. Set False for the paper's exact form.
min_ks_magnitude	0.0 (off)	Pool clients with no distinguishable shift.

To reproduce the paper's algorithm exactly:

from clma import CLMA, AllParams
from clma.cohort import AutoKMeans, CLMACohorter

CLMA(
    param_selector=AllParams(),
    cohorter=CLMACohorter(center=False, clusterer=AutoKMeans(fixed_k=2)),
)

Every stage is an extension point: DriftDetector, ParamSelector, Cohorter and Clusterer are all ABCs with working defaults.

Implementation notes

Three places where a literal transcription of the paper needed a decision. Each is called out in the source and covered by tests.

1. The Gram trick (the one that matters). Algorithm 3 forms eig(SᵀS), where S is n_clients × n_params. That matrix is n_params × n_params — fine for the paper's small CNN, but a ResNet-18 (11M parameters) would need roughly 500 TB. Since n_clients ≪ n_params always, CLMA eigendecomposes the n_clients × n_clients matrix SSᵀ instead. The two share non-zero eigenvalues and the projection follows as S V = U Λ^{1/2}, so V is never formed. This is an exact reformulation, not an approximation — tests/test_projection.py checks it against the literal eig(SᵀS) — and it is what makes the method usable beyond tabular data at all.

2. Signed KS statistic. Algorithm 3 branches on if T > 0, but the standard two-sample KS statistic is non-negative by construction, which would leave S₋ permanently empty. CLMA uses a signed form — sign(D⁺ − D⁻) · max(D⁺, D⁻) — whose sign records the direction of the shift (the stated purpose of the test in §3) and whose magnitude is exactly the classical two-sided D, verified against SciPy in the test suite.

3. Centering before PCA, and choosing K. The paper's eig(SᵀS) is uncentered and does not specify k. Both defaults changed, for one connected reason.

Every client starts a round from the same global model, so their parameter vectors share a large common component. Uncentered, that component absorbs nearly all the variance: the 95% rule collapses the projection to a single dimension, and 1-D data always looks clusterable. Measured on eight identically drifted clients, the best silhouette score is 0.46 uncentered — a confident, entirely spurious 3-way split — versus 0.03 centered, while genuinely distinct drifts score 1.00 either way. So CLMA centers by default (center=False restores the paper's form), and picks k by silhouette with a min_silhouette=0.25 floor, comfortably between those regimes. Pass AutoKMeans(fixed_k=...) to pin k instead.

One addition beyond the paper: Algorithm 1 only ever appends cohorts, so a long-running federation fragments toward singletons even after drifts reconcile. merge_cohorts_threshold merges converged cohorts. It is off by default so the paper's behaviour is reproduced exactly.

Development

pip install -e ".[dev]"
pytest
ruff check .

Citing

@inproceedings{nl2024collaborative,
  title     = {Collaborative Drift Compensation},
  author    = {N L, Adarsh and Amarlingam, Madapu and Sharma, Divyasheel},
  booktitle = {Proceedings of the 8th International Conference on Data Science
               and Management of Data (CODS-COMAD)},
  year      = {2024},
  doi       = {10.1145/3703323.3703341}
}

License

Apache-2.0.