"""Weight aggregation algorithms for Federated Learning. Implements various aggregation methods beyond simple FedAvg: - FedProx (proximal term) - SCAFFOLD (variance reduction) - FedOpt (adaptive optimization) """ from abc import ABC, abstractmethod from typing import Any, Dict, List, Optional, Tuple import numpy as np import torch class Aggregator(ABC): """Base class for federated aggregation algorithms.""" @abstractmethod def aggregate( self, parameters: List[List[np.ndarray]], weights: List[float], client_metrics: List[Dict], ) -> List[np.ndarray]: """Aggregate client parameters. Args: parameters: List of client parameters weights: Weight for each client (usually proportional to data size) client_metrics: Metrics from each client Returns: Aggregated parameters """ pass class FedAvgAggregator(Aggregator): """Federated Averaging (FedAvg) - standard weighted average.""" def aggregate( self, parameters: List[List[np.ndarray]], weights: List[float], client_metrics: List[Dict], ) -> List[np.ndarray]: """Weighted average of client parameters.""" if not parameters: raise ValueError("No parameters to aggregate") if len(parameters) == 1: return parameters[0] # Normalize weights total_weight = sum(weights) normalized_weights = [w / total_weight for w in weights] # Weighted average aggregated = None for params, weight in zip(parameters, normalized_weights): if aggregated is None: aggregated = [p * weight for p in params] else: aggregated = [ a + p * weight for a, p in zip(aggregated, params) ] return aggregated class FedProxAggregator(Aggregator): """FedProx with proximal term for handling heterogeneity.""" def __init__(self, mu: float = 0.01): """ Args: mu: Proximal term coefficient """ self.mu = mu def aggregate( self, parameters: List[List[np.ndarray]], weights: List[float], client_metrics: List[Dict], global_parameters: Optional[List[np.ndarray]] = None, ) -> List[np.ndarray]: """Aggregate with proximal term regularization.""" if global_parameters is None: # Fall back to FedAvg if no global params return FedAvgAggregator().aggregate(parameters, weights, client_metrics) # Weighted average with proximal correction total_weight = sum(weights) normalized_weights = [w / total_weight for w in weights] aggregated = None for params, weight, global_params in zip( parameters, normalized_weights, global_parameters ): # Proximal term: add regularization toward global model corrected_params = [ p + self.mu * (g - p) for p, g in zip(params, global_params) ] if aggregated is None: aggregated = [p * weight for p in corrected_params] else: aggregated = [ a + p * weight for a, p in zip(aggregated, corrected_params) ] return aggregated class SCAFFOLDAggregator(Aggregator): """SCAFFOLD: Stochastic Controlled Averaging for Federated Learning.""" def __init__(self, global_control: Optional[List[np.ndarray]] = None): self.global_control = global_control or None def set_global_control(self, control: List[np.ndarray]) -> None: """Set global control variates.""" self.global_control = control def aggregate( self, parameters: List[List[np.ndarray]], weights: List[float], client_metrics: List[Dict], client_controls: Optional[List[List[np.ndarray]]] = None, ) -> List[np.ndarray]: """Aggregate using SCAFFOLD algorithm.""" if client_controls is None or self.global_control is None: return FedAvgAggregator().aggregate(parameters, weights, client_metrics) total_weight = sum(weights) normalized_weights = [w / total_weight for w in weights] # Compute weight updates (gradient-like terms) aggregated = None for params, weight, client_ctrl, global_ctrl in zip( parameters, normalized_weights, client_controls, self.global_control ): # Direction: client_params - global_params + global_control - client_control delta = [ p - g + gc - cc for p, g, gc, cc in zip(params, parameters[0], self.global_control, client_ctrl) ] if aggregated is None: aggregated = [d * weight for d in delta] else: aggregated = [ a + d * weight for a, d in zip(aggregated, delta) ] # Add back global parameters if aggregated: aggregated = [ g + self.mu * a if hasattr(self, 'mu') else g + 0.001 * a for g, a in zip(self.global_control if self.global_control else parameters[0], aggregated) ] return aggregated @property def mu(self) -> float: """Learning rate for SCAFFOLD.""" return 0.001 class FedOptAggregator(Aggregator): """FedOpt: Adaptive Federated Optimization using server-side optimizer.""" def __init__( self, server_lr: float = 1.0, beta_1: float = 0.9, beta_2: float = 0.99, epsilon: float = 1e-4, ): self.server_lr = server_lr self.beta_1 = beta_1 self.beta_2 = beta_2 self.epsilon = epsilon self.m_t: Optional[List[np.ndarray]] = None self.v_t: Optional[List[np.ndarray]] = None self.t = 0 def aggregate( self, parameters: List[List[np.ndarray]], weights: List[float], client_metrics: List[Dict], ) -> List[np.ndarray]: """Aggregate using FedOpt (server-side Adam).""" if not parameters: raise ValueError("No parameters to aggregate") # FedAvg as base base_aggregate = FedAvgAggregator().aggregate(parameters, weights, client_metrics) # Compute delta from base to weighted average if self.m_t is None: self.m_t = [np.zeros_like(p) for p in base_aggregate] self.v_t = [np.zeros_like(p) for p in base_aggregate] delta = [ ba - p0 for ba, p0 in zip(base_aggregate, parameters[0]) ] self.t += 1 # Update momentum and second moment self.m_t = [ self.beta_1 * m + (1 - self.beta_1) * d for m, d in zip(self.m_t, delta) ] self.v_t = [ self.beta_2 * v + (1 - self.beta_2) * (d ** 2) for v, d in zip(self.v_t, delta) ] # Bias correction m_hat = [m / (1 - self.beta_1 ** self.t) for m in self.m_t] v_hat = [v / (1 - self.beta_2 ** self.t) for v in self.v_t] # Apply update aggregated = [ p0 + self.server_lr * m / (np.sqrt(v) + self.epsilon) for p0, m, v in zip(parameters[0], m_hat, v_hat) ] return aggregated def create_aggregator( method: str = "fedavg", **kwargs, ) -> Aggregator: """Factory function to create aggregator. Args: method: Aggregation method ("fedavg", "fedprox", "scaffold", "fedopt") **kwargs: Additional arguments for the aggregator Returns: Aggregator instance """ if method == "fedavg": return FedAvgAggregator() elif method == "fedprox": return FedProxAggregator(mu=kwargs.get("mu", 0.01)) elif method == "scaffold": return SCAFFOLDAggregator() elif method == "fedopt": return FedOptAggregator( server_lr=kwargs.get("server_lr", 1.0), beta_1=kwargs.get("beta_1", 0.9), beta_2=kwargs.get("beta_2", 0.99), ) else: raise ValueError(f"Unknown aggregation method: {method}") if __name__ == "__main__": print("Available aggregators: fedavg, fedprox, scaffold, fedopt") # Example usage agg = create_aggregator("fedavg") print(f"Created aggregator: {type(agg).__name__}")