Skip to content

Repository files navigation

GradGuard: Gradient Inversion Attacks and Defenses in Federated Learning

Overview

This project demonstrates that gradients shared in Federated Learning (FL) are not as privacy-preserving as commonly assumed. We implement a gradient inversion attack (iDLG) that reconstructs a client's private training image purely from its shared gradient, then evaluate two defense strategies — PRECODE (variational bottleneck) and Homomorphic Encryption (HE) — measuring their privacy protection and cost tradeoffs.

Attack and defense are a synergistic relationship: a powerful means of attack will give birth to targeted defense technology, and solid defense technology will promote the means of attack upgrade.

Built as a reference implementation alongside the survey:

Rao et al., "Privacy Inference Attack and Defense in Centralized and Federated Learning: A Comprehensive Survey," IEEE TAI, 2025.

Other references:


How It Works

Client trains locally → sends gradient to server
                              ↓
                     [ATTACK] iDLG reconstructs
                     the original training image
                     from the gradient alone
                              ↓
              [DEFENSE A] PRECODE — stochastic bottleneck
              makes gradient a moving target, attack fails
                              ↓
              [DEFENSE B] HE — gradient encrypted before
              sending, server never sees plaintext gradient

Background: Key Concepts

Federated Learning (FL)

In traditional machine learning, everyone sends their data to one central server which trains the model. Federated Learning flips this: the model goes to the data instead.

Traditional:   Client → sends raw data → Server trains model
 
Federated:     Server → sends model → Client trains locally
                                            ↓
                              Client → sends gradient back
                              (gradient = what the model learned,
                               not the raw data itself)

The promise: your private data (medical records, photos, messages) never leaves your device. Only the gradient update is shared. This is used in production systems like Google's Gboard keyboard prediction.


The Attack: iDLG (Improved Deep Leakage from Gradients)

The problem is that gradients carry more information than expected.

In plain terms: If you know exactly how a model changed after seeing one image, you can mathematically work backwards to figure out what that image looked like — without ever seeing it directly.

iDLG does this in two steps:

Step 1 — Label extraction (the "i" improvement over original DLG):

The gradient of the final layer has a mathematical property: the row corresponding to the true class label has the most negative values. So the attacker can read the label directly from the gradient — no guessing needed.

Step 2 — Image reconstruction:

Start with:  random noise image
             ↓
Repeatedly:  "what gradient would THIS image produce?"
             compare to the intercepted real gradient
             adjust the image to make them match
             ↓
After ~300 iterations: random noise → recognisable reconstruction

The attacker never touched the original data. They only saw a gradient. Yet they reconstructed the image. This is the threat GradGuard defends against.


Defense A: PRECODE (Variational Bottleneck)

In plain terms: PRECODE makes the gradient a moving target.

Every time the model processes an image, PRECODE injects a small amount of structured randomness into the intermediate representation. This means the gradient produced is slightly different every single time, even for the exact same input image.

Without PRECODE:
  same image → always same gradient → attacker can converge on reconstruction
 
With PRECODE:
  same image → different gradient each time → attacker chases a moving target
                                               optimization never converges

Technically, PRECODE inserts a Variational Information Bottleneck between the convolutional feature extractor and the classifier. It compresses the feature vector into a probability distribution (mean μ, variance σ), then samples from that distribution before passing it forward. Because sampling is random, the gradient is never reproducible.

The cost: A small accuracy drop (~5-10%) because the model must learn to be robust to this internal noise. No formal privacy proof — it's an empirical defense that works against specific attacks like iDLG.


Defense B: Homomorphic Encryption (HE)

In plain terms: The server does the maths on a locked box without ever opening it.

Standard encryption works like this: lock the data, send it, unlock it, do maths. The problem for FL is that the server needs to add up gradients from many clients — and you can't add locked data without unlocking it first.

Homomorphic Encryption solves this with a special mathematical property:

Normal encryption:
  encrypt(5) + encrypt(3) = gibberish   ✗
 
Homomorphic encryption:
  encrypt(5) + encrypt(3) = encrypt(8)  ✓
  server computed the sum without seeing 5 or 3

In GradGuard, we use the CKKS scheme (via TenSEAL), which supports floating-point arithmetic — necessary because model weights are floats, not integers.

Each client:   encrypts their gradient → sends ciphertext to server
Server:        adds ciphertexts together (never decrypts)
               result = encrypt(gradient_1 + gradient_2 + ...)
Clients:       decrypt the aggregate → get the average gradient
 
The server performed the entire aggregation
without ever seeing any individual gradient.
The attack has nothing to invert.

The cost: Mathematically perfect privacy, but computationally expensive. Every weight in the model must be individually encrypted and decrypted. For a CNN with millions of parameters, this can be 10-100x slower than unencrypted aggregation. This is why we reduce to 2 clients for the HE experiments.


Why Neither Defense is Strictly Better

PRECODE:  cheap, fast, small accuracy cost
          but no formal privacy guarantee
          works against iDLG, might not work against future attacks
 
HE:       mathematically provable privacy
          zero accuracy cost (no noise added to gradients)
          but very slow — impractical for large models or many clients
 
The right choice depends on your deployment:
  Mobile / edge devices     → PRECODE (can't afford HE compute)
  High-security servers     → HE (compute is available, proof is needed)
  Research / benchmarking   → both, which is what GradGuard does

Defense Comparison

Defense Privacy Mechanism Reconstruction (SSIM / PSNR) Label Protected Test Accuracy Speed Overhead
Baseline None High (~0.97 / ~26.5 dB) No ~33% 1.0x
PRECODE Stochastic variational bottleneck Near zero (~−0.007 / ~6.2 dB) No ~23% ~1.03x
HE TenSEAL CKKS encryption Blocked (ciphertext only) Yes ~33% ~1.6x+

Key insight: PRECODE degrades reconstruction quality at minimal accuracy cost. HE blocks the attack structurally but imposes significant computational overhead. Neither is strictly better — the right choice depends on deployment constraints.


Repository Structure

gradguard/
├── attack/
│   └── idlg.py               # iDLG gradient interception and reconstruction
├── defense/
│   ├── precode.py             # VariationalBottleneck + KL loss term
│   └── he.py                  # TenSEAL CKKS context + secure aggregation
├── src/
│   ├── models.py              # CNN with optional PRECODE bottlenecktoggle
│   └── fl_simulation.py       # FedAvg loop with defense_mode hook
├── experiments/
│   └── run_experiment.py      # Unified experiment script (all three conditions)
├── analysis/
│   ├── diagnostics.py         # MSE, PSNR, SSIM, accuracy metric functions
│   └── visualize.py           # Comparison plots and trilemma chart
├── results/                   # Generated outputs (gitignored)
│   ├── baseline/
│   ├── precode/
│   ├── he/
│   └── comparison/
├── .gitlab-ci.yml             # CI/CD pipeline
└── requirements.txt           # Python dependencies

Quickstart (Local)

1. Install dependencies

git clone <repository-url>
cd gradguard
pip install -r requirements.txt

Requires Python 3.10+ and a CUDA-capable GPU (or MPS on Apple Silicon). The device is selected automatically:

# src/models.py — get_device() picks the best available:
cudampscpu

2. Run experiments

# Baseline — attack with no defense
DEFENSE_MODE=none python experiments/run_attack_baseline.py

# PRECODE defense
DEFENSE_MODE=precode USE_PRECODE=true python experiments/run_precode.py

# HE defense
DEFENSE_MODE=he python experiments/run_he.py

3. Generate comparison figures

python analysis/visualise_metrics.py

Outputs:

  • results/comparison/comparison_grid.png — original vs reconstructed per defense
  • results/comparison/trilemma.png — privacy / accuracy / speed radar chart
  • results/comparison/metrics_summary.json — all metrics in one file

The "Endpoint Leak" Scenario (HE Vulnerability)

In plain terms: HE is a mathematically perfect shield, but only while the data is in transit or on the server. If the shield is bypassed at the source, you have exactly zero protection.

To demonstrate this, GradGuard includes an endpoint leakage simulation, triggered by setting the CI/CD variable RUN_ENDPOINT_LEAK=true.

Unlike PRECODE (which permanently alters the gradient by injecting noise during the forward pass), Homomorphic Encryption calculates the exact, perfect gradient first, and then encrypts it. If an attacker compromises the client device itself (endpoint) and intercepts the plaintext gradient right before it gets encrypted by TenSEAL, the iDLG attack will reconstruct the victim's image flawlessly.

This experiment proves that HE acts as a binary, "all-or-nothing" defense:

  • Server-Side Attack (Normal HE): 100% protected. The server sees only ciphertext. The attack is mathematically blocked.
  • Endpoint-Side Attack (Leaked HE): 0% protected. The gradient is unmodified, allowing complete pixel-perfect image reconstruction.

This highlights a critical architectural tradeoff: HE defends the aggregation phase perfectly, but leaves the client node vulnerable to local malware or memory extraction.


Evaluation Metrics

To objectively evaluate the "Privacy vs. Utility vs. Efficiency" trilemma, GradGuard measures the following metrics for every experiment:

1. Privacy Guarantee (Reconstruction Metrics)

These metrics measure how successfully the iDLG attack reconstructed the victim's private image.

  • SSIM (Structural Similarity Index Measure): Evaluates the perceived visual similarity between the original image and the attacker's reconstruction. Scores range from -1.0 to 1.0. A score close to 1.0 means the attacker successfully reconstructed the exact structures (shapes, edges) of the image. A score near 0 or lower means the defense succeeded in reducing the image to static.
  • PSNR (Peak Signal-to-Noise Ratio): Measured in decibels (dB), this indicates the ratio of the original image's signal to the noise of the reconstruction. Higher values (e.g., 25+ dB) indicate severe privacy leakage. Lower values mean the attacker only recovered noise.
  • MSE (Mean Squared Error): The mathematical pixel-by-pixel difference between the original and reconstructed images.
  • Label Leakage: A binary metric tracking whether the attacker successfully extracted the correct ground-truth classification label directly from the gradient.

2. Model Utility

  • Test Accuracy: The primary metric for the global model's performance on the unseen CIFAR-10 test dataset. This measures the "cost" of the defense. If a defense protects privacy but drops accuracy from 35% to 10%, it destroys the model's utility.

3. Computational Efficiency

  • Pure FL Time: The raw execution time (in seconds) it takes to complete the federated training rounds and model aggregation.
  • Total Execution Time: Includes context initialization, federated training, and the 300 iterations of the iDLG attack phase.
  • Speed Overhead: We use the Baseline run as a 1.0x benchmark to calculate exactly how much extra latency a defense adds (e.g., HE's heavy cryptographic chunking computations).

CI/CD Pipeline

The pipeline runs automatically on every push using a self-hosted runner on a Windows machine with an NVIDIA GPU (tag: gradguard-gpu).

Pipeline Stages

validate    → import_check + gpu_check     (confirms env and GPU)
experiment  → baseline_attack              (train + attack, no defense)
             precode_defense               (train + attack + PRECODE)
             he_defense [commented]        (train + attack + HE, uncomment when ready)

Configurable Variables

All experiment parameters are set via CI/CD environment variables. Override them when triggering a pipeline manually.

Variable Default Description
DEFENSE_MODE none Which defense to run: none, precode, he
USE_PRECODE false Set true to activate PRECODE bottleneck in CNN
NUM_CLIENTS 2 Number of simulated FL clients
GLOBAL_EPOCHS 2 Number of FL communication rounds
ATTACK_ITERS 50 LBFGS iterations for iDLG reconstruction
VICTIM_IDX 1 Which client index the attacker targets
PYTHON_EXE D:\Anaconda3\envs\gradguard\python.exe Path to Python on the runner

Trigger via GitLab Web UI

  1. Go to CI/CD → Pipelines in your GitLab project
  2. Click Run pipeline (top right)
  3. Select your branch (e.g., main)
  4. Under Variables, add any overrides:
Key Value Effect
DEFENSE_MODE precode Run PRECODE experiment
GLOBAL_EPOCHS 10 Full training run
ATTACK_ITERS 300 Presentation quality reconstruction
  1. Click Run pipeline

Trigger via GitLab CLI (glab)

# Quick test — baseline only, mini settings
glab ci run -b main

# PRECODE defense, full quality
glab ci run -b main \
  -v DEFENSE_MODE:precode \
  -v USE_PRECODE:true \
  -v GLOBAL_EPOCHS:10 \
  -v ATTACK_ITERS:300

# HE defense
glab ci run -b main \
  -v DEFENSE_MODE:he \
  -v GLOBAL_EPOCHS:10 \
  -v ATTACK_ITERS:300

# Full presentation run (all three, highest quality)
glab ci run -b main \
  -v GLOBAL_EPOCHS:10 \
  -v ATTACK_ITERS:300

Artifacts

After each job completes, results are available under Pipeline → Job → Browse artifacts or in the MR sidebar:

Job Artifact Contents
baseline_attack results/baseline/ original.png, reconstructed.png, metrics.json
precode_defense results/precode/ original.png, reconstructed.png, metrics.json
he_defense results/he/ original.png, reconstructed.png, metrics.json

Team

Module Owner
Attack (iDLG) Ayoub
PRECODE Defense + FL Simulation + Diagnostics & Evaluation Nihal
HE Defense Dioran

References

Rao, B., Zhang, J., Wu, D., Zhu, C., Sun, X., & Chen, B. (2025). Privacy Inference Attack and Defense in Centralized and Federated Learning: A Comprehensive Survey. IEEE Transactions on Artificial Intelligence, 6(2), 333–353.

Zhu, L., Liu, Z., & Han, S. (2019). Deep Leakage from Gradients. NeurIPS.

Zhao, B., Mopuri, K. R., & Bilen, H. (2020). iDLG: Improved Deep Leakage from Gradients. arXiv:2001.02610.

Scheliga, D., Mäder, P., & Seeland, M. (2022). PRECODE — A Generic Model Extension to Prevent Deep Gradient Leakage. WACV.

About

Gradient inversion attacks and privacy defenses in federated learning using iDLG, PRECODE, and homomorphic encryption.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages