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)

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.