Skip to content

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

1 Commit
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

EFFEKT

Efficient Federated Knowledge Transfer to Foundation Models

Python PyTorch Flower License Paper

TL;DR: EFFEKT lets a server-side Foundation Model (FM) learn new, private domains from lightweight client-side proxy models without ever seeing client data. At every federated round, the knowledge collected from the clients is distilled into domain-specific LoRA adapters attached to the FM (Clients-to-Server, C2S distillation), and the FM is then used to realign the client proxies to the updated feature space (Joint Alignment, JA distillation). This bidirectional cross-distillation scheme replaces plain weight-averaging aggregation and consistently outperforms it, while keeping client-side compute limited to a MobileNetV3-Small proxy.

👥 Authors

Matteo Caligiuri1,2 Francesco Barbato2 Pietro Zanuttigh2 Francesco Restuccia1

1 Northeastern University, Boston (MA), United States 2 University of Padua, Padua, Italy

📊 Graphical Abstract

EFFEKT Graphical Abstract

📑 Citation

If you use this code in your research, please cite our paper:

@article{caligiuri2026effekt,
      title={{EFFEKT}: Efficient Federated Knowledge Transfer to Foundation Models},
      author={Matteo Caligiuri and Francesco Barbato and Pietro Zanuttigh and Francesco Restuccia},
      journal={Transactions on Machine Learning Research},
      year={2026},
      url={https://openreview.net/forum?id=jpUDUJfE1K},
}

🌟 Key Features

🚀 Core Contributions

  • Clients-to-Server (C2S) Distillation: Instead of averaging client weights, the knowledge collected from the active clients' proxy heads is distilled server-side into a small set of domain-specific LoRA adapters attached to the frozen FM, using a reversed (mode-covering) logit-level KD loss computed on public, task-specific pretraining data.
  • Joint Alignment (JA) Distillation: After the LoRA update, a realignment step retrains the aggregated client model and the LoRA-adapted FM jointly (cross-entropy + bidirectional feature/logit KD), restoring feature-space compatibility before the next federated round starts.
  • Multi-Domain, Domain-Asynchronous Federation: Each domain owns its own LoRA adapter and classification head, so new domains can be learned incrementally without touching previously-learned ones or the shared proxy encoder.
  • Prototype-Based Domain Discriminator: A frozen, few-shot MobileNetV3-Small classifier picks the right domain (and thus LoRA/head) at inference time when the client does not supply it, with no retraining needed when a new domain is added.
  • Privacy-Preserving by Design: The server never sees raw client data — aggregation is done via LoRA/head distillation on public pretraining data — and an optional Local Differential Privacy mode (norm clipping + Gaussian noise) is available for a tunable privacy/accuracy trade-off.
  • Real-World, On-Device Validation: Full federated pipeline deployed via Flower on a heterogeneous cluster of Raspberry Pi 4/5 and NVIDIA Jetson Nano/Orin-Nano boards, with power and network usage monitoring.

🔧 Technical Features

  • Heterogeneous Architectures: Server-side DINOv2 (ViT-L/14 with registers) foundation model paired with a lightweight client-side proxy (MobileNetV3-Small + linear feature translator by default; EfficientNet-B0 and TinyViT-5M are also supported for ablations).
  • LoRA-Based Server Adaptation: Rank-configurable LoRA adapters injected into the Q/V projections of every attention block of the FM (fedlib/models/lora_qkv.py), trained through C2S/JA instead of full fine-tuning (~194x fewer trainable server parameters at rank 16).
  • ICP Client Regularization: Client-side Inactive-Classes-Preservation training that keeps the proxy encoder frozen and preserves latent-space alignment while only the classification head is trained locally.
  • Configurable Distillation Schedule: The effekt strategy can run C2S/JA every round, at a fixed frequency, or after a warm-up period of plain FedAvg rounds.
  • Multiple FL Baselines Included: FedAvg, FedAvg+EMA, FedProx, MOON, and a similarity-reweighted aggregation strategy (used to reproduce FedHEAL), alongside EFFEKT (ICP client + C2S/JA server distillation), all runnable from the same codebase for apples-to-apples comparison.
  • Differential Privacy: Client-side Local DP (norm clipping + Gaussian noise, moments-accountant based) and server-side fixed/adaptive clipping DP wrappers around the aggregation strategies.
  • Hydra Configuration: Every component (models, datasets, partitioners, client algorithms, server strategies, losses, transforms) is a composable Hydra config group, so new experiments are pure CLI overrides.
  • Experiment Tracking: MLflow (default) or TensorBoard logging, plus label-distribution and heatmap visualizations for the federated partitions.

📊 Comprehensive Evaluation

  • 5 Public → Private Domain Pairs: StanfordCars→CompCars, Food101→UECFOOD256, CUB-200→NABirds, FGVCAircraft→MilitaryAircraft, ImageNetPets→OxfordPets.
  • Multi-Domain Inference: Concatenating domain-specific LoRAs/heads on the FM and routing queries through the prototype-based domain discriminator, evaluated with extras/scripts/multi_domain_eval.py.
  • Real-Device Deployment: 3× Raspberry Pi 4, 3× Raspberry Pi 5, 5× Jetson Nano, 1× Jetson Orin-Nano, matched against the simulated results to validate the simulation setup.

🏗️ Architecture Overview

EFFEKT tackles the CA-FKT (Cross-Architecture Federated Knowledge Transfer) task through four stages:

  1. Pretraining: The client proxy encoder is distilled from the frozen FM on a public, domain-specific dataset so that both share a common feature space (and a shared classification head can be attached to either).
  2. Federated (Client) Round: Active clients receive the current proxy model, train only the classification head locally with ICP regularization on their private data, and upload the updated head.
  3. Clients-to-Server (C2S) Distillation: The server updates the domain's LoRA adapter by distilling the received (frozen) client heads' logits into the LoRA-adapted FM, using public pretraining data only.
  4. Joint Alignment (JA) Distillation: The server aggregates the client heads (FedAvg), then jointly fine-tunes the LoRA-adapted FM, the aggregated client model, and the pretraining head to restore feature/logit alignment before the next round.
foundation-fed/
├── main.py                    # Entry point: pretraining + federated pipeline
├── statistical_analysis.py    # Wilcoxon / Friedman significance tests over seeds
├── data_downloader.py         # CLI to fetch released checkpoints, embeddings & prototypes
├── extras/
│   └── scripts/               # create_embeddings.py, domain_proto_extractor.py, multi_domain_eval.py
├── conf/                      # Hydra configuration groups
│   ├── base.yaml              # Main pretraining + federated config
│   ├── base_deploy.yaml       # Config used for real-device (Flower deployment) runs
│   ├── server_model/          # DINOv2 variants
│   ├── client_model/          # MobileNetV3, EfficientNet-B0, TinyViT proxy encoders
│   ├── classifier/            # Classification head configs
│   ├── federated/
│   │   ├── client_type/       # base, ema, fedprox, moon, icp (ours)
│   │   ├── strategy/          # fedavg, fedalign, fedhpa, effekt (EFFEKT, ours)
│   │   ├── dts/               # Federated (private, target-domain) datasets
│   │   ├── partitioner/       # Dirichlet non-IID partitioners
│   │   ├── client_dp/         # Local DP client mod
│   │   └── server_dp/         # Fixed/adaptive clipping server-side DP
│   ├── pretraining/           # Pretraining datasets, losses, optimizers, schedulers
│   └── transforms/            # Data augmentation / preprocessing building blocks
├── fedlib/                    # Core implementation
│   ├── models/                # DINOv2, MobileNetV3, EfficientNet, TinyViT, LoRA-QKV
│   ├── federated/
│   │   ├── clients/           # Flower ClientApps (base, ICP, EMA, FedProx, MOON) + LocalDP mod
│   │   ├── strategies/        # Effekt (C2S+JA, ours), FedAlign, FedHPA, DP wrappers
│   │   └── visualization/     # Label-distribution / heatmap plots
│   ├── multi_domain/          # Domain prototypes, multi-domain loader & inference task
│   ├── datasets/              # Dataset handlers & Dirichlet partitioners
│   ├── trainers/              # Pretraining and federated trainer orchestration
│   ├── losses/                # Cross-entropy, cosine, reconstruction, MSIW losses
│   └── logging/               # MLflow/TensorBoard logger + MQTT power monitor
├── data/                      # Datasets, checkpoints, embeddings, prototypes (gitignored)
└── outputs/                   # Hydra experiment outputs and logs

🚀 Quick Start

1. Environment Setup

The project uses uv with a committed uv.lock for reproducible installs (Python 3.11):

# Clone the repository
git clone https://github.com/LTTM/EFFEKT.git
cd EFFEKT

# Install the dependencies (defaults to the "simulation" group: Ray + cuML acceleration
# for running server and simulated clients side-by-side on a single GPU machine)
uv sync

# Verify installation
uv run python -c "import torch; print(f'PyTorch: {torch.__version__}')"
uv run python -c "import flwr; print(f'Flower: {flwr.__version__}')"

pyproject.toml defines separate install profiles for the two roles in a real, multi-device deployment (see Section 5) — these are [project.optional-dependencies] extras, installed with --extra, not --group:

  • Server (GPU machine hosting the FM): uv sync --extra server — adds xformers/cuml-cu12 acceleration plus paho-mqtt for optional power monitoring.
  • Client (Raspberry Pi / Jetson boards): uv sync --extra client — installs the base dependencies only; it intentionally does not pull in the CUDA-only xformers/cuml-cu12 packages, which won't install on those boards anyway.

Conda environment files are also provided under extras/ for GPU/CPU/Wayland setups (env_linux_cuda.yml, env_win_cuda.yml, env_cpu.yml) if you prefer conda over uv.

2. Data Preparation

# Set your data directory in conf/base.yaml
# IMPORTANT: The trailing slash (/) is essential and cannot be removed!
# All datasets will be located in a 'datasets' folder inside this path
# Example: if dts_root_dir: /home/user/data/, datasets will be in /home/user/data/datasets/
dts_root_dir: /path/to/your/data/

# Some datasets are automatically downloaded on first use for supported datasets

3. Pre-computed Embeddings, Checkpoints & Prototypes (Optional)

To speed up DINOv2 evaluation/distillation and to reproduce the paper's results exactly, released checkpoints, embeddings, and domain prototypes (needed by the multi-domain discriminator) can be downloaded with data_downloader.py:

# Interactive mode: the script guides you through the available download options
uv run python data_downloader.py

# Non-interactive examples:
uv run python data_downloader.py --list-files           # List available files without downloading
uv run python data_downloader.py --pretrain             # Pretraining (teacher-distillation) checkpoints
uv run python data_downloader.py --federated            # Federated (proxy) checkpoints
uv run python data_downloader.py --all-checkpoints      # Pretraining + federated checkpoints
uv run python data_downloader.py --embeddings           # Cached DINOv2 embeddings
uv run python data_downloader.py --prototypes           # Domain prototypes for multi-domain eval
uv run python data_downloader.py --everything           # Everything above
uv run python data_downloader.py --data-dir /custom/path --everything

# Files are placed under ./data/{checkpoints,embeddings,domain_prototypes}/

Alternative: Create Embeddings / Prototypes Locally

Both scripts live under extras/scripts/ and, like main.py, are Hydra-driven and must be run from the repository root:

# Cache DINOv2 embeddings for one, several, or all supported datasets (Hydra-configurable,
# see conf/create_embeddings.yaml — override e.g. datasets_to_load=[stanfordcars,compcars])
uv run python extras/scripts/create_embeddings.py

# Build the per-domain feature prototypes used by the domain discriminator
# (see conf/proto_extractor.yaml)
uv run python extras/scripts/domain_proto_extractor.py

4. Basic Usage

Pretraining Only

uv run python main.py \
    resume_pretraining=True \
    pretraining_epochs=60 \
    pretraining_batch_size=64 \
    server_model=dinov2_vit_large14_reg \
    client_model=mobilenetv3_small \
    classifier=single \
    skip_fed=True \
    pretraining_checkpoint=null \
    federated_checkpoint=null

Federated Learning Only (EFFEKT: ICP client + C2S/JA server distillation)

uv run python main.py \
    server_model=dinov2_vit_large14_reg \
    client_model=mobilenetv3_small \
    classifier=single \
    resume_pretraining=False \
    federated/dts=compcars \
    federated/partitioner=dirichlet_compcars \
    federated/client_type=icp \
    federated/strategy=effekt \
    num_clients=100 \
    num_clients_per_round_fit=10 \
    num_rounds=500

Local Differential Privacy

uv run python main.py \
    federated/dts=compcars \
    federated/client_type=icp \
    federated/strategy=effekt \
    federated/client_dp=local_dp_compcars

5. On-Device / Real Deployment (Flower)

EFFEKT has been validated on a heterogeneous cluster of Raspberry Pi 4/5 and Jetson Nano/Orin-Nano boards using the Flower SuperLink/SuperNode deployment runtime, instead of the single-machine simulation used in the commands above.

Step 1 — Set up the environment on every device, cloning the repository as in Section 1 and then installing the role-appropriate extra (see above):

# On the server (GPU machine)
uv sync --extra server

# On each client board (Raspberry Pi / Jetson)
# --no-default-groups skips the "simulation" group (xformers/cuml-cu12), which is
# CUDA-only and will not install on these boards
uv sync --extra client --no-default-groups

On resource-constrained boards install the CPU/board-appropriate PyTorch build for that platform if the pinned torch/torchvision wheels aren't available for your architecture. Every device also needs the conf/ directory available locally; the deployment runtime always loads conf/base_deploy.yaml (not conf/base.yaml), so tune the federated setup (num_clients, num_rounds, learning rate, etc.) there to match your real cluster.

Step 2 — Start the SuperLink on the server:

export FEDLIB_CONFIG_DIR="/path/to/effekt/conf"
export FEDLIB_ENABLE_POWER_MONITORING="1"           # optional: MQTT-based power logging
export FEDLIB_MQTT_BROKER_URL="mqtt://<your-mqtt-broker-url>"
flower-superlink --insecure --fleet-api-address <server-ip>:9092

Step 3 — Start a SuperNode on each client device:

export FEDLIB_CONFIG_DIR="/path/to/effekt/conf"
flower-supernode --superlink <server-ip>:9092 --insecure --node-config "partition-id=0"

Step 4 — Launch the run from the server:

flwr run . deployment

Step 5 — Inspect the logs (on the server):

flwr ls .                      # list runs and get the <run_id>
flwr log <run_id> .            # full log
flwr log <run_id> . | grep ERROR  # errors only

6. Individual Model Testing

Model definitions can also be exercised standalone for debugging:

# DINOv2 foundation model (server side)
uv run python fedlib/models/dinov2.py

# MobileNetV3 proxy model (client side)
uv run python fedlib/models/mobilenetv3.py

📊 Supported Models

Foundation Models (Server)

  • DINOv2: dinov2_vit_small14_reg, dinov2_vit_large14_reg (default, ViT-L/14 with registers)
  • LoRA adapters (fedlib/models/lora_qkv.py) injected into the Q/K/V attention projections, rank- and projection-configurable

Client Proxy Models

  • MobileNetV3-Small (+ linear feature translator) — default, used in the paper's main results
  • MobileNetV3-Large
  • EfficientNet-B0 — ablation, see paper Fig. A.7/A.8
  • TinyViT-5M — ablation, transformer-based proxy

🔧 Federated Learning Algorithms

Client Types (federated/client_type=)

  • icp: Inactive-Classes-Preservation client (ours) — frozen encoder, head-only local training
  • base: Standard FedAvg client
  • fedprox: Proximal-term regularized client (FedProx)
  • moon: Model-contrastive client (MOON)
  • ema: Exponential Moving Average client

Server Strategies (federated/strategy=)

  • effekt: EFFEKT — Clients-to-Server (C2S) LoRA distillation + Joint Alignment (JA) realignment (ours)
  • fedavg: Standard federated averaging baseline
  • fedhpa: Similarity-reweighted aggregation, used to reproduce the FedHEAL baseline
  • fedalign: Unsupervised server-side fine-tuning of the aggregated model on public pretraining data

Note

EFFEKT = ICP (client) + effekt (server, C2S + JA distillation).

📂 Datasets & Features

Supported Domain Pairs (public pretraining → private federated target)

Pretraining (public) Federated target (private)
StanfordCars CompCars
Food101 UECFOOD256
CUB-200 NABirds
FGVCAircraft MilitaryAircraft
ImageNetPets OxfordPets

Data is split across nk clients using a Dirichlet distribution (concentration α, default 1) to simulate non-IID label skew. Additional pretraining-only datasets are available (ImageNet-1k, iNaturalist 2017/2018, Flowers-102, Central Asian Food, MNIST) for out-of-domain / ablation studies.

Custom Datasets

  • Method 1 — Existing PyTorch datasets: add a config file under conf/federated/dts/ or conf/pretraining/dts/ following the existing examples.

  • Method 2 — Custom dataset class:

    from fedlib.decorators import add_dts
    
    @add_dts
    class MyCustomDataset(ImageFolder, DefDataset):
        # Your custom dataset implementation

    then add the matching config file (see fedlib/datasets/compcars_handler.py for reference).

Data Partitioning

  • Dirichlet distribution (fedlib/datasets/partitioners/dirichlet.py) for non-IID simulation
  • Custom partitioners can be implemented following the Flower partitioner docs

🔒 Privacy Features

  • Client-side Local DP (federated/client_dp=local_dp_compcars): gradient clipping + Gaussian noise, with optional per-round clipping-norm scheduling, based on the moments accountant (Abadi et al., 2016).
  • Server-side DP (federated/server_dp=fixed_clipping / adaptive_clipping): fixed or adaptive client-side clipping wrappers around the aggregation strategy.
  • Even without DP, EFFEKT only ever shares classification heads and ~1.5M LoRA parameters (<1% of the FM) between clients and server — the FM encoder weights and all raw data stay local.

📈 Multi-Domain Inference & Evaluation

uv run python extras/scripts/multi_domain_eval.py

Concatenates the domain-specific heads/LoRAs learned independently for each domain and attaches them to the shared FM encoder, so a single server can recognize classes from every trained domain. When the domain of a query is unknown, the prototype-based domain discriminator (built with extras/scripts/domain_proto_extractor.py) selects the right LoRA/head automatically (see conf/multi_domain_eval.yaml for per-dataset checkpoint/prototype paths).

Partition Visualization

uv run python main.py plot_label_distribution=True

Generates bar/heatmap visualizations of the per-client label distribution to inspect the non-IID setup.

🔬 Extending the Framework

Since the codebase is built on Flower, custom client algorithms and server strategies can be added by following the Flower documentation (Client API, Strategy API).

Custom Client Implementation

from fedlib.federated.clients.base_client_app import BaseClient

class MyCustomClient(BaseClient):
    def fit(self, parameters, config):
        # Your custom training logic
        return super().fit(parameters, config)

Custom Strategy Implementation

from flwr.server.strategy import FedAvg

class MyCustomStrategy(FedAvg):
    def aggregate_fit(self, server_round, results, failures):
        # Your custom aggregation logic
        return super().aggregate_fit(server_round, results, failures)

📝 License

This project is licensed under the MIT License - see the LICENSE file for details.

🙏 Acknowledgments

  • Flower Framework for the federated learning infrastructure
  • Meta AI Research for DINOv2
  • PyTorch for the deep learning framework
  • This work was partially supported by the European Union under the Italian National Recovery and Resilience Plan (NRRP) of NextGenerationEU, partnership on "Telecommunications of the Future" (PE00000001-program "RESTART"), by the NSF under grants CNS-2312875 and OAC-2530896, by AFOSR under grant FA9550-23-1-0261, by ONR under grant N00014-23-1-2221, and by DARPA under Cooperative Agreement D25AC00374-00.

📞 Support

For questions and support:

About

Efficient multi-domain federated learning via server-side LoRA adapters and bi-directional cross-distillation between a Foundation Model and lightweight client proxies.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Contributors

Languages