Metadata-Version: 2.4
Name: fed-clma
Version: 0.1.0
Summary: Collaborative Model Adaptation: federated learning under distributed concept and covariate drift
Project-URL: Homepage, https://github.com/adarshnl/clma
Project-URL: Paper, https://doi.org/10.1145/3703323.3703341
Project-URL: Issues, https://github.com/adarshnl/clma/issues
Author-email: Adarsh N L <adarshnanjaiya@gmail.com>
License: Apache-2.0
License-File: LICENSE
Keywords: clustered-federated-learning,concept-drift,covariate-shift,federated-learning,flower,flwr
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: flwr<2.0,>=1.32
Requires-Dist: numpy>=1.24
Requires-Dist: scikit-learn>=1.3
Provides-Extra: dev
Requires-Dist: mypy>=1.8; extra == 'dev'
Requires-Dist: pytest-cov>=4.1; extra == 'dev'
Requires-Dist: pytest>=7.4; extra == 'dev'
Requires-Dist: ruff>=0.4; extra == 'dev'
Provides-Extra: simulation
Requires-Dist: flwr[simulation]<2.0,>=1.32; extra == 'simulation'
Provides-Extra: tensorflow
Requires-Dist: tensorflow>=2.13; extra == 'tensorflow'
Provides-Extra: torch
Requires-Dist: torch>=2.0; extra == 'torch'
Description-Content-Type: text/markdown

# CLMA — Collaborative Model Adaptation

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

A reference implementation of [*Collaborative Drift Compensation*][paper]
(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.

[paper]: https://doi.org/10.1145/3703323.3703341

---

## 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

```bash
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`:

```python
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:

```python
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:

```python
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)

```python
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`](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:

```python
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

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

## Citing

```bibtex
@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.
