repro_bundle/README.md

# Reproduction bundle — DFMPQ

Toy-scale reproduction of *"No Retraining at Edge: Efficient
Resource-Aware Mixed-Precision Quantization via Federated Supernet Learning"*
(ICML 2026, OpenReview `nHbvQF35ch`).

This is now a faithful reproduction of the paper's three core mechanisms
(re-verified against the paper PDF, which was provided after the first version
of this bundle shipped; the bundle was then upgraded to implement CCSA, SAHA,
and SAGS directly per the paper's equations):

  * **CCSA** (Class-Conditional Semantic Alignment, §4.1, Eqs. 4-6)
    — a linear-kernel MMD regularizer that pulls each client's per-class
    feature centroid toward the global full-precision prototype; added to
    the local loss with `lambda = 0.1` (paper default).
  * **SAHA** (Semantic-Aware Hybrid Aggregation, §4.2, Eqs. 7-8)
    — crucial-layer (stage-output) parameters are aggregated with the
    semantic-confidence weight `w_k ~ N_k * exp(-Σ_c ||mu_k,c - mu_g,c|| / tau)`;
    non-crucial layers use standard volume-weighted FedAvg.
  * **SAGS** (Sensitivity-Aware Greedy Search, §4.3, Eqs. 9-10)
    — a per-layer sensitivity score `Omega_l = (1/|W_l|) * ||grad(W_l) ⊙ W_l||_1`
    is computed ONCE from the FP supernet using a small server buffer,
    then we greedily downgrade the top-K least-sensitive layers down by one
    bit until the memory budget is met. This is the retraining-free step.

## What's in here

```
repro_dfmpq/
  dfmpq.py          STE quantizer + QResNet8 / QMobileTiny supernets +
                    Dirichlet non-IID partition + CCSA regularizer +
                    SAHA semantic aggregation + SAGS search + FedAvg.
  run.py            End-to-end pipeline (HF parquet cifar mirror) +
                    flags --no-ccsa / --no-saha / --lambda-ccsa / --saha-tau.
  derive_only.py    Claim-1 demonstration: load checkpoints, derive subnets
                    via SAGS at every memory budget, ZERO further training.
  make_plots.py     Plotly Pareto-frontier figures + aggregate metrics.
  outputs_run.log   Verbatim stdout of the upgraded run that produced the
                    published numbers.

outputs/
  results.csv           every (dataset × arch × budget × method) row
  results.json          structured copy of results.csv + per-pair timings
  summary.json          headline aggregate metrics
  claim1_derive.csv     retraining-free derivation run (Claim 1)
  pareto_*.html         one interactive Plotly Pareto frontier per pair
  supernet_*.pt         the 4 federally-trained supernet checkpoints
```

## How to rerun (≤ 25 minutes on one consumer GPU)

```bash
mkdir -p /tmp/opencode/data
curl -sL -o /tmp/opencode/data/cifar10_train.parquet   https://huggingface.co/datasets/uoft-cs/cifar10/resolve/main/plain_text/train-00000-of-00001.parquet
curl -sL -o /tmp/opencode/data/cifar10_test.parquet    https://huggingface.co/datasets/uoft-cs/cifar10/resolve/main/plain_text/test-00000-of-00001.parquet
curl -sL -o /tmp/opencode/data/cifar100_train.parquet  https://huggingface.co/datasets/uoft-cs/cifar100/resolve/main/cifar100/train-00000-of-00001.parquet
curl -sL -o /tmp/opencode/data/cifar100_test.parquet   https://huggingface.co/datasets/uoft-cs/cifar100/resolve/main/cifar100/test-00000-of-00001.parquet

pip install torch torchvision pyarrow pillow plotly

python repro_dfmpq/run.py --rounds 15 --n-clients 5 --lr 0.1 \
    --lambda-ccsa 0.1 --saha-tau 1.0 \
    --datasets cifar10,cifar100 --archs resnet8,mobiletiny --out outputs

python repro_dfmpq/derive_only.py --ckpt-dir repro_dfmpq/outputs \
    --out-csv repro_dfmpq/outputs/claim1_derive.csv

python repro_dfmpq/make_plots.py outputs
```

## Scope vs paper

* Paper: ResNet-18 / MobileNetV2 / EfficientNet-Lite0 on CIFAR-10/100 +
  Tiny-ImageNet, 10 clients, Dirichlet α in {0.1,0.5,1.0}, 50 global rounds,
  5 local epochs, batch 256, lr 0.1, on A100 (80 GB).
* This reproduction: ResNet-8 / tiny-MobileNetV2 on CIFAR-10/100, **5
  clients**, α=0.5, **15 rounds, 2 local epochs**, batch 128, on a single
  4 GB RTX 3050 laptop GPU. Total wall ~20 min training + ~1 min SAGS
  derivation across 4 supernets.
* Note on the accuracy gap: the gap from the per-budget-retrained reference
  is *larger* than at paper scale (paper report ~80% CIFAR-10 with ResNet-18
  @ 3MP, ours ~52–60% with ResNet-8 @ mixed 2/4/8-bit). The CCSA/SAHA
  mechanisms are designed for paper-scale (10+ clients, 50+ rounds) where
  semantic drift and aggregation variance actually bite; at toy scale they
  over-regularise and under-tune. We document this in the logbook's Scope &
  cost table and Claim 2 page.