Federated Learning Distributed Training Privacy and Communication Efficiency

# Federated Learning: Distributed Training, Privacy, and Communication Efficiency

## 1. Introduction & Motivation

Federated learning (FL) is a machine learning paradigm where training is distributed across multiple clients without centralizing data. Instead of sending raw data to a central server, clients perform local training and share only model updates (gradients). This approach addresses critical concerns in modern ML:

• Privacy: Raw data never leaves client devices
• Bandwidth efficiency: Only model updates are communicated
• Regulatory compliance: Enables training on sensitive data (healthcare, finance)
• Personalization: Models can be customized per-client

Federated learning is increasingly important for edge computing, mobile devices, and multi-institutional collaborations where data sharing is restricted by privacy regulations (GDPR, HIPAA) or competitive concerns.

The main challenge is the heterogeneous, non-IID (non-independent and identically distributed) nature of data across clients, which violates standard ML assumptions and causes slower convergence and worse generalization.

This article comprehensively covers federated learning algorithms, privacy guarantees, communication optimization, and practical deployment considerations.

## 2. Core Concepts & Theory

### 2.1 Federated Learning Framework

In federated learning, we have:
- Central server: Orchestrates training and maintains global model
- N clients: Each with local dataset D_i (private, stays on device)
- Objective: Minimize global loss

$$L_{ ext{global}} = \sum_{i=1}^{N} \frac{|D_i|}{|D_{ ext{total}}|} L_i(w)$$

where L_i is client i's local loss and w are shared model parameters.

### 2.2 Federated Averaging (FedAvg)

The standard federated learning algorithm:

Server side: Initialize model weights w_0. At round t: send w_t to selected client subset S_t, wait for updates, and aggregate w_{t+1} = sum_{k in S_t} (n_k / n) * w_k^t (where S_t is selected clients, n_k = |D_k|, and n = sum n_k).

Client side: For selected client k, download current model w_t, compute gradient on local data g_k = grad L_k(w_t), perform local updates w_k^t <- w_t - eta * g_k, and upload w_k^t to server.

### 2.3 Key Challenge: Non-IID Data

In federated settings, data is non-IID (clients have different label distributions):

$$p_i(y | x) eq p_j(y | x) \quad ext{for different clients } i, j$$

This causes:
- Slower convergence: Averaging dissimilar local models adds noise
- Worse generalization: Models overfit to client-specific distributions
- Instability: Local updates diverge before convergence

### 2.4 Convergence Analysis

For strongly convex losses with non-IID data:

$$ ext{Error} = O\left(\frac{1}{T} + \frac{\sigma^2}{m \cdot T} ight)$$

where T is number of rounds, m is batch size, sigma^2 is variance from non-IID data.

The non-IID term sigma^2 / (m * T) dominates convergence rate, explaining why federated learning is slower.

## 3. Mathematical Formulation

### 3.1 FedAvg Update Rule

Server aggregation in round t:

$$w_{t+1} = w_t - \eta \sum_{k \in S_t} \frac{n_k}{n} abla L_k(w_t)^t$$

where nabla L_k(w_t)^t denotes the gradient at local step t of client k.

For E > 1 local epochs:

$$w_k^{(e)} = w_k^{(e-1)} - \eta abla L_k(w_k^{(e-1)})$$

Multiple local epochs improve efficiency but increase divergence.

### 3.2 Variance Analysis

The variance of global gradient estimate:

$$ ext{Var}[\hat{g}_t] = ext{Var}_{ ext{sampling}}[\hat{g}_t] + ext{Var}_{ ext{non-IID}}[\hat{g}_t]$$

First term from stochastic sampling, second from data heterogeneity:

$$ ext{Var}_{ ext{non-IID}} = \sum_{k=1}^{N} \frac{n_k}{n} \| abla L_k - abla L_{ ext{global}}\|^2$$

### 3.3 Differential Privacy Guarantee

Adding Gaussian noise N(0, sigma_p^2 I) to gradients before aggregation provides (ε, δ)-differential privacy:

$$ ext{Pr}[ ext{output}_1 \in S] \leq e^\epsilon \cdot ext{Pr}[ ext{output}_2 \in S] + \delta$$

For two datasets differing by one sample. Privacy budget is typically epsilon in [0.1, 10] (smaller = more private).

### 3.4 Communication Cost

Average communication per round:

$$C_{ ext{round}} = N \cdot d \cdot b$$

where N is active clients, d is model dimension, b is bits per parameter.

Total communication over T rounds:

$$C_{ ext{total}} = T \cdot N \cdot d \cdot b$$

For large models ( d > 10^6 ), communication dominates computation.

## 4. Advanced Theory & Extensions

### 4.1 Personalized Federated Learning

Rather than single global model, learn personalized models:

$$\min_w \sum_{i=1}^{N} \ell_i(w_i) = \min_w \sum_{i=1}^{N} \left(L_i(w_i) + \frac{\lambda}{2}\|w_i - \bar{w}\|^2 ight)$$

where w_bar is shared global model, lambda controls personalization.

Algorithms:
- Per-FedAvg: Each client maintains local model regularized toward global
- Ditto: Explicitly decouple personalization and global model
- FedProx: Add proximal term to local objective

Benefits: 5-15% improvement on non-IID data at cost of model diversity.

### 4.2 Gradient Compression

Reduce communication by quantizing or sparsifying gradients:

Top-k sparsification:
- Send only top k gradient coordinates by magnitude
- Drop remaining coordinates (zero out)
- Compression ratio: 1 - k/d Quantization:
- Uniform quantization: Map [-alpha, alpha] to b-bit integers
- Error bounded by O(1/2^b) - 1-bit quantization reduces communication by 32x

Combined: Error feedback preserves small gradients:

$$\Delta_k = ext{quantize}(g_k + e_{k-1})$$
$$e_k = g_k - \Delta_k + e_{k-1}$$

Maintains convergence guarantee with 10-100x compression.

### 4.3 Asynchronous Federated Learning

Remove synchronization barrier (waiting for slowest client):

$$w_{t+1} = w_t - \eta \sum_k \alpha_k \Delta_k^{(t_k)}$$

where t_k is time client k last updated (can be old).

Benefits: Removes stragglers, enables continuous training

Drawbacks: Delayed gradients add staleness bias, slower convergence

### 4.4 Secure Multi-Party Computation

Protect aggregation step from honest-but-curious server:

1. Clients split gradient into shares: Delta_k = s_{k,1} + s_{k,2} + ... 2. Server never sees individual gradients
3. Aggregate: sum Delta_k = sum sum s_{k,j} Communication: O(N * d) compared to O(d) for unencrypted aggregation.

## 5. Computational Considerations

### 5.1 Computation Distribution

Total computation time = Client computation + Server computation + Communication

Client computation:
- Local training: O(E * |D_k| * T) where E is local epochs
- Gradient computation:

$$ O(E \cdot |D_k| / B) $$

forward-backward passes

Server computation:
- Aggregation: O(N * d) - negligible for large models
- Bottleneck: Client-side computation and communication

### 5.2 Bandwidth and Latency

Bandwidth requirement:
- Model size 1GB, 10K clients per round: 10TB aggregate
- With compression 100x: 100GB (feasible)

Latency:
- Wall-clock time = max(client time, communication time)
- Stragglers: Slow clients delay round completion
- Mitigation: Select fast clients or use asynchronous algorithms

### 5.3 Memory Constraints

Clients may have limited memory:
- Mobile devices: ~2-4GB RAM
- Edge devices: ~500MB-2GB
- Compression enables smaller models to fit

### 5.4 Optimization for Edge Devices

Techniques:
- FedProx: Proximal term reduces local model divergence
- Knowledge distillation: Smaller student models trained on teacher predictions
- Layer dropping: Reduce local model size while training global model

Benefits: Enable participation of resource-constrained devices.

## 6. Practical Implementation Strategies

### 6.1 Client Selection

Random selection:
- Select C * N random clients each round
- C is participation rate (typically 0.1)
- Simple but ignores client capabilities

Biased selection:
- Favors fast or large-data clients
- Reduces stragglers but hurts fairness
- Formula: P_i = (n_i + lambda * speed_i) / sum_j(n_j + lambda * speed_j) Clustering:
- Group similar clients, select cluster leaders
- Reduces communication and convergence time
- Useful for hierarchical FL

### 6.2 Local Training Configuration

Number of local epochs: E = 1-10 typical
- E = 1 : Minimum bias but slower convergence
- E = 5-10 : Good efficiency, manageable divergence
- Higher E helps with communication bottleneck

Local batch size: B = 32-128 typical
- Larger batches reduce communication rounds
- Smaller batches more stable (higher variance helps escape bad minima)

Local learning rate:

$$ \eta_k = \eta / \sqrt{E} $$

often works well
- Scale with local epochs to maintain stability
- Adaptive methods (Adam) less sensitive

### 6.3 Server-Side Aggregation

Weighted averaging (standard):

$$w_{t+1} = \sum_k \frac{n_k}{n} w_k^t$$

Accounts for different dataset sizes.

Clustering aggregation:

$$w^{(c)} = ext{avg}(w_k ext{ for } k \in ext{cluster } c)$$

Then meta-average clusters.

Importance weighting:

$$w_{t+1} = \sum_k p_k w_k^t$$

where

$$ p_k = \frac{ ext{update quality}_k}{\sum_j ext{update quality}_j} $$

### 6.4 Differential Privacy Configuration

DP-SGD for federated learning:

1. Clip gradient norm:

$$ ilde{g}_k = g_k \cdot \min(1, C / \|g_k\|) $$

2. Add Gaussian noise:

$$ \hat{g}_k = ilde{g}_k + N(0, C^2 \sigma^2 I) $$

3. Aggregate:

$$ w_{t+1} = w_t - \eta \sum_k \hat{g}_k / N $$

Privacy parameters:
- Clipping threshold: C = 1 standard
- Noise level: sigma = 1-5 (higher = more private)
- Privacy budget: epsilon = 1-10 (lower = more private)

## 7. Benchmark Datasets & Evaluation

### 7.1 Federated Datasets

FEMNIST (Federated MNIST):
- 3,550 writers (clients), 62M images
- N = 3550 , avg images per client = 17,500
- Non-IID degree (Dirichlet alpha = 0.1 ): High
- Baseline accuracy: 89% (IID), 84% (Non-IID)

Shakespeare (Character-Level Language Modeling):
- 715 writers (clients), 5M total characters
- Avg chars per client: ~7,000
- Baseline: ~30% accuracy (IID), ~20% (Non-IID)
- Useful for studying personalization

Google Federated Analytics:
- CIFAR-100 partitioned by client
- Stackexchange posts, Reddit comments
- Real-world non-IID distributions

### 7.2 Evaluation Metrics

Test Accuracy: Standard evaluation on held-out test set
- Measures generalization performance
- Averaged across clients (or aggregated)

Convergence Speed: Rounds to reach target accuracy
- Lower is better
- Accounts for communication and computation

Communication Efficiency: Bytes or model updates per accuracy gain
- Critical metric for bandwidth-constrained settings
-

$$ ext{efficiency} = \frac{ ext{accuracy}}{ ext{communication cost}} $$

Fairness: Variance of accuracy across clients
- High variance indicates some clients do poorly
- Min/mean/max accuracy tracking

### 7.3 Benchmark Results

FEMNIST (Non-IID):
- FedAvg: 83.6% accuracy, 400 rounds
- FedProx: 85.1% accuracy, 350 rounds (+1.5%, -12%)
- Personalized FL: 86.2% accuracy, 300 rounds (+2.6%, -25%)

Shakespeare:
- FedAvg: 19.2% accuracy, 500 rounds
- With differential privacy ( epsilon = 10 ): 15.8% accuracy (17% drop)
- Compression 10x: 18.9% accuracy (1.5% drop)

## 8. Key Challenges & Limitations

### 8.1 Non-IID Data and Convergence

The primary challenge: data heterogeneity across clients

• Causes divergence of local models
• Limits gradient information usefulness
• Slows convergence by 5-10x compared to IID case

Mitigation: Personalization, FedProx, increased local epochs (trade communication for variance reduction).

### 8.2 Privacy-Utility Trade-off

Adding privacy noise reduces model utility:

$$ ext{Accuracy}( ext{privacy level}) = f(\epsilon)$$

where f is decreasing. Typical impact:
- epsilon = 10 : ~2% accuracy drop
- epsilon = 1 : ~10-15% accuracy drop
- epsilon = 0.1 : Model useless (severe privacy protection)

Optimal epsilon depends on application requirements.

### 8.3 Communication Bottleneck

Even with small models (10MB), 10K clients = 100GB per round:

• Dominates computation cost
• Gradient compression helps (100x compression realistic)
• Still challenging for latency-sensitive applications

### 8.4 Model and System Heterogeneity

Real federated systems face:

Model heterogeneity:
- Different clients need different models
- Single global model may be suboptimal for specialized clients
- Personalization helps but adds complexity

System heterogeneity:
- Different computational power (phone vs. edge server)
- Different bandwidth capabilities
- Stragglers cause round delays

Solution: Asynchronous FL, client selection by capabilities.

## 9. Hyperparameter Tuning & Optimization

### 9.1 Global Algorithm Hyperparameters

Server learning rate: eta_server = 0.1-1.0 typical
- Not the same as local learning rate
- Controls aggregation step size
- Larger values faster convergence, may overshoot

Client participation rate: C = 0.01-0.2 typical
- Fraction of clients selected per round
- Smaller C: Faster rounds but noisier gradients
- Common: C * N = 50-100 clients

Number of rounds: T = 100-1000 typical
- Dataset and convergence criteria dependent
- Larger models/datasets need more rounds
- Monitor validation accuracy for early stopping

### 9.2 Local Training Hyperparameters

Local epochs: E = 1-10 typical
- E = 1 : Good for large N (communication dominates)
- E = 5-10 : Good balance between divergence and communication
- Higher E requires lower eta_local to maintain stability

Local batch size: B = 32-128 - Typical mini-batch sizes as in centralized setting
- Larger batches less noisy but may reduce regularization

Local learning rate:

$$ \eta_{ ext{local}} = \eta_{ ext{server}} / \sqrt{E} $$

works well
- Can also use adaptive methods (Adam with eta = 0.001 )
- Client-specific learning rates possible

### 9.3 Privacy Hyperparameters

Clipping threshold: C = 1 standard, tune if needed
- Controls norm of gradients
- Too small: Excessive clipping, slow convergence
- Too large: Less privacy protection

Noise scale: sigma = 1-5 typical for

$$ \epsilon \in [1, 10] $$

• Higher sigma : More privacy (lower epsilon )
• Privacy budget:

$$ \epsilon = C / \sigma $$

(simplified)
- Choose based on privacy requirements

### 9.4 Communication Hyperparameters

Compression ratio:

$$ k/d \in [0.01, 0.1] $$

for top-k sparsification
- 1-10% of gradients communicated
- Larger reduction acceptable with error feedback
- Combined with quantization for further compression

Quantization bits: b = 1-8 typical
- 1-bit: Extreme compression, large error
- 4-bit: Balanced (16x compression)
- 8-bit: Minimal error (4x compression)

## 10. Real-World Applications & Case Studies

### 10.1 Mobile Keyboard Prediction

Problem: Train on-device keyboard prediction model across million devices

Setup:
- Clients: Mobile phones with unique typing patterns
- Data: User keystrokes and completions (sensitive)
- Objective: Improve prediction without centralizing data

Deployment:
- Model: 1M parameters (small RNN)
- Local epochs: E = 1 (minimize on-device latency)
- Client selection:

$$ C = 1\% $$

(100K clients per round)
- Updates: ~1MB per client, 100GB aggregate

Results:
- Federated vs. centralized: 2-3% accuracy difference (federated slightly better due to local adaptation)
- Privacy protection: DP with epsilon = 1000 (very private)
- Deployment: Millions of devices

Key insights:
- Small models essential for mobile
- Personalization critical for heterogeneous users
- Privacy important for user acceptance

### 10.2 Healthcare: Federated Medical Imaging

Problem: Train tumor detection model across hospital networks

Setup:
- Clients: 50 hospitals with 100-1000 labeled CT scans each
- Data: Very sensitive (HIPAA regulated)
- Objective: State-of-the-art detection without centralizing data

Deployment:
- Model: ResNet50 (25M parameters, ~100MB)
- Local epochs: E = 5 (hospitals have computational resources)
- Client selection: All 50 hospitals per round (small N)
- Differential privacy: epsilon = 5 (strong privacy)

Results:
- Centralized baseline: 95.2% AUC
- Federated (no DP): 94.8% AUC (0.4% drop)
- Federated (DP epsilon = 5 ): 93.5% AUC (1.7% drop)
- Privacy-utility trade-off manageable for medical imaging

Lessons learned:
- Differential privacy acceptable cost for regulated data
- Hospital-level data heterogeneity manageable with FedAvg
- Communication limiting factor (bandwidth between hospitals)

### 10.3 IoT Sensor Networks

Problem: Anomaly detection on distributed IoT sensors

Setup:
- Clients: 10K sensors, continuous data streams
- Model: 1-layer MLP (10K parameters)
- Objective: Detect anomalies at edge without cloud centralization

Deployment:
- Rounds: Continuous (new data every hour)
- Local epochs: E = 1 (real-time requirements)
- Client participation: Asynchronous (nodes offline frequently)
- Communication: Gradient compression 10x

Results:
- Detection rate: 92% (vs. 94% centralized)
- False positive rate: 3% (vs. 2% centralized)
- Latency: <100ms per detection (enables real-time response)

Key techniques:
- Asynchronous aggregation (tolerates offline nodes)
- Gradient sparsification (bandwidth-constrained sensors)
- Local model adaptation for sensor drift

### 10.4 Federated Cross-Device Learning for NLP

Problem: Improve autocomplete across devices with language diversity

Setup:
- Clients: 1M devices, diverse languages and domains
- Model: 10M parameter LSTM
- Objective: Single multilingual model via federated training

Training protocol:
- Communication per round: 10GB (model + gradients)
- Local updates: 20 gradient steps
- Selective participation: Sample devices with available bandwidth

Results:
- Baseline (centralized): Perplexity 45 (English), 120 (low-resource languages)
- Federated: Perplexity 46 (English), 115 (low-resource languages)
- Federated with personalization: 44 (English), 105 (low-resource)

Challenges:
- Extreme non-IID (different languages, domains)
- Communication cost for large models
- Device heterogeneity (modern phones vs. older devices)

## 11. Integration with Other Methods

### 11.1 Transfer Learning in Federated Settings

Combine federated training with pretrained models:

1. Initialize from pretrained model (BERT, ResNet)
2. Fine-tune globally via federated learning
3. Allow local personalization

Benefits:
- Faster convergence (5-10x)
- Better performance on non-IID data
- Reduced communication

### 11.2 Multi-Task Federated Learning

Learn multiple related tasks across clients:

$$\min_{\{w_i\}} \sum_{i=1}^{N} L_i(w_i) + \frac{\lambda}{2}\|w_i - \bar{w}\|^2$$

• Each client solves own task i
• Global model w_bar captures shared structure
• Task-specific models capture personalizations

### 11.3 Federated Meta-Learning

Learn initial model weights useful for fast adaptation:

$$\min_w \sum_{i=1}^{N} L_i(w - \eta abla L_i(w))$$

• MAML (Model-Agnostic Meta-Learning) adapted to federated setting
• Each client adapts global model in 1-2 steps
• Useful when clients have limited data

### 11.4 Hybrid Federated-Centralized Training

Balance privacy and performance:

• Centralized phase: Train on less sensitive data
• Federated phase: Adapt to private client data
• Benefits: Faster convergence, privacy-utility trade-off

## 12. Future Research Directions

### 12.1 Efficient Communication

Research areas:
- Gradient compression beyond top-k (sketching, hashing)
- Learned compression (train encoder/decoder)
- Adaptive compression (different rates per layer)

Target: Reduce communication by another 100x while maintaining accuracy.

### 12.2 Handling System Heterogeneity

Current limitations:
- Stragglers delay each round
- Device dropouts interrupt training

Future directions:
- Asynchronous algorithms more robust to stragglers
- Fault tolerance mechanisms
- Incentive mechanisms for participation

### 12.3 Improving Non-IID Robustness

Current challenge:
- Federated learning 5-10x slower than centralized on non-IID data
- Personalization helps but adds complexity

Future directions:
- Better optimization algorithms (variance reduction techniques)
- Data-heterogeneity-aware aggregation
- Federated continual learning (streaming non-IID data)

### 12.4 Combining Privacy and Security

Research areas:
- Certified defenses against Byzantine failures
- Secure aggregation at scale
- Privacy-preserving model updates with formal guarantees

## 13. Summary & Key Takeaways

Core Algorithm (FedAvg):
- Clients perform local training, send model updates
- Server aggregates weighted by data size
- Multiple local epochs improve communication efficiency

Challenges:
- Non-IID data: 5-10x slower convergence than centralized training
- Communication bottleneck: Dominates training time
- Privacy-utility trade-off: DP adds 2-15% accuracy loss depending on privacy budget

Key Techniques:

*Addressing non-IID:*
- FedProx: Proximal term regularizes local updates
- Personalization: Per-client models with shared component
- Increased local epochs: Reduce communication at cost of divergence

*Communication efficiency:*
- Gradient compression: 10-100x reduction with minimal accuracy loss
- Quantization: 4-bit reduces communication by 8x
- Top-k sparsification: Send only top 1-10% of gradients

*Privacy protection:*
- Differential privacy: Add Gaussian noise to gradients
- Secure aggregation: Encrypt individual updates
- epsilon = 1-10 typical range (smaller = more private)

Practical Deployment:
- Client participation: 1-10% of total clients per round
- Local epochs: 5-10 (balances communication and convergence)
- Privacy budget: epsilon = 10-100 for deployment (strong privacy at epsilon = 1-10 )

Performance Characteristics:
- Communication-efficient FL: 10-100GB per round for large-scale systems
- Personalized FL: 1-2% accuracy improvement over global model
- Privacy-preserving FL: 2-15% accuracy loss depending on privacy strength

Current Status:
Federated learning is moving into production for privacy-critical applications (mobile, healthcare). Open challenges remain in communication efficiency, non-IID robustness, and privacy guarantees. Continued research in system heterogeneity handling and efficient communication protocols will further enable real-world deployment.

---

## Appendix: Practical Implementation Labs

### Lab 1: Basic FedAvg Implementation

import torch
import torch.nn as nn
from copy import deepcopy

class FedAvgServer:
 def __init__(self, model, num_clients, num_rounds):
 self.global_model = deepcopy(model)
 self.num_clients = num_clients
 self.num_rounds = num_rounds
 self.client_updates = []
 
 def send_model(self, client_id):
 """Send current global model to client"""
 return deepcopy(self.global_model)
 
 def receive_update(self, client_id, model, client_size):
 """Receive updated model from client"""
 self.client_updates.append((model, client_size))
 
 def aggregate(self, total_samples):
 """Aggregate client updates using FedAvg"""
 aggregated_state = None
 
 for model, client_size in self.client_updates:
 weight = client_size / total_samples
 
 if aggregated_state is None:
 aggregated_state = {}
 for key in model.state_dict():
 aggregated_state[key] = weight * model.state_dict()[key]
 else:
 for key in model.state_dict():
 aggregated_state[key] += weight * model.state_dict()[key]
 
 self.global_model.load_state_dict(aggregated_state)
 self.client_updates = []
 
 return self.global_model
 
 def train_round(self, clients, total_samples):
 """Execute one federated round"""
 for client_id, client in enumerate(clients):
 # Send model to client
 local_model = self.send_model(client_id)
 
 # Client trains locally
 client_size = client.train_local(local_model, epochs=5)
 
 # Receive update
 self.receive_update(client_id, local_model, client_size)
 
 # Aggregate
 self.aggregate(total_samples)

class FedAvgClient:
 def __init__(self, data_loader, device='cpu'):
 self.data_loader = data_loader
 self.device = device
 self.optimizer = None
 
 def train_local(self, model, epochs=5, lr=0.001):
 """Train model locally on client data"""
 model.to(self.device)
 optimizer = torch.optim.SGD(model.parameters(), lr=lr)
 criterion = nn.CrossEntropyLoss()
 
 model.train()
 for _ in range(epochs):
 for images, labels in self.data_loader:
 images, labels = images.to(self.device), labels.to(self.device)
 
 optimizer.zero_grad()
 outputs = model(images)
 loss = criterion(outputs, labels)
 loss.backward()
 optimizer.step()
 
 return len(self.data_loader.dataset)

# Test
# server = FedAvgServer(model, num_clients=10, num_rounds=100)
# for round in range(100):
# server.train_round(clients, total_samples=50000)

### Lab 2: Gradient Compression

import torch

def top_k_sparsify(gradients, k=0.1):
 """Keep top k% of gradients by magnitude"""
 flattened = torch.cat([g.flatten() for g in gradients])
 num_keep = max(1, int(len(flattened) * k))
 
 _, top_indices = torch.topk(torch.abs(flattened), num_keep)
 
 mask = torch.zeros_like(flattened, dtype=torch.bool)
 mask[top_indices] = True
 
 sparse_gradients = flattened * mask.float()
 return sparse_gradients.nonzero()[:, 0], flattened[mask]

def quantize_gradient(gradient, bits=4):
 """Quantize gradient to b-bit precision"""
 min_val = gradient.min()
 max_val = gradient.max()
 range_val = max_val - min_val + 1e-8
 
 # Quantize to [0, 2^bits - 1]
 quantized = ((gradient - min_val) / range_val * (2 ** bits - 1)).round()
 
 # Dequantize for reconstruction
 dequantized = quantized / (2 ** bits - 1) * range_val + min_val
 
 return quantized, (min_val, max_val), dequantized

# Test
gradient = torch.randn(1000)
indices, values = top_k_sparsify(gradient, k=0.1)
print(f"Sparsified: {len(values)} values out of 1000")

quantized, (min_val, max_val), dequantized = quantize_gradient(gradient)
error = torch.mean((gradient - dequantized) ** 2)
print(f"Quantization error: {error.item():.6f}")

### Lab 3: Differential Privacy in Federated Learning

import torch
from torch.nn.utils import clip_grad_norm_

def add_differential_privacy(gradients, clipping_threshold=1.0, noise_scale=0.5):
 """Add differential privacy to gradients"""
 # Clip gradient norms
 clipped_grads = []
 for g in gradients:
 norm = torch.norm(g)
 if norm > clipping_threshold:
 clipped_grads.append(g * (clipping_threshold / norm))
 else:
 clipped_grads.append(g)
 
 # Add Gaussian noise
 dp_grads = []
 for g in clipped_grads:
 noise = torch.randn_like(g) * noise_scale
 dp_grads.append(g + noise)
 
 return dp_grads

def compute_privacy_budget(num_rounds, noise_scale, delta=1e-5):
 """Compute epsilon for (epsilon, delta)-DP"""
 # Simplified: epsilon ≈ sqrt(log(1/delta) * num_rounds) / (2 * noise_scale)
 epsilon = (num_rounds ** 0.5 * (2 * torch.tensor(1e-5).log().abs()) ** 0.5) / (2 * noise_scale)
 return epsilon.item()

# Test
model = torch.nn.Linear(10, 5)
example_input = torch.randn(4, 10)
example_input.requires_grad = True
output = model(example_input)
loss = output.sum()
loss.backward()

gradients = [p.grad for p in model.parameters()]
dp_gradients = add_differential_privacy(gradients, clipping_threshold=1.0, noise_scale=1.0)

epsilon = compute_privacy_budget(num_rounds=1000, noise_scale=1.0, delta=1e-5)
print(f"Privacy budget: epsilon={epsilon:.4f}")

### Lab 4: Personalized Federated Learning (FedProx)

import torch
import torch.nn as nn

class FedProxClient:
 def __init__(self, data_loader, mu=0.01, device='cpu'):
 self.data_loader = data_loader
 self.mu = mu # Proximal term coefficient
 self.device = device
 
 def train_local(self, model, global_model, epochs=5, lr=0.001):
 """Train with proximal term regularizing toward global model"""
 model.to(self.device)
 global_model.to(self.device)
 
 optimizer = torch.optim.SGD(model.parameters(), lr=lr)
 criterion = nn.CrossEntropyLoss()
 
 model.train()
 global_model.eval()
 
 for epoch in range(epochs):
 for images, labels in self.data_loader:
 images, labels = images.to(self.device), labels.to(self.device)
 
 optimizer.zero_grad()
 
 # Standard cross-entropy loss
 outputs = model(images)
 loss = criterion(outputs, labels)
 
 # Add proximal term
 proximal_loss = self.mu / 2.0 * sum(
 torch.norm(p1 - p2) ** 2 
 for p1, p2 in zip(model.parameters(), global_model.parameters())
 )
 
 total_loss = loss + proximal_loss
 total_loss.backward()
 optimizer.step()
 
 return len(self.data_loader.dataset)

# Test
model = torch.nn.Linear(10, 5)
global_model = torch.nn.Linear(10, 5)
data_loader = [(torch.randn(4, 10), torch.randint(0, 5, (4,)))]

client = FedProxClient(data_loader, mu=0.01)
client.train_local(model, global_model, epochs=2)
print("FedProx training complete")