Technical Summary: Byzantine-Robust Federated Learning with Learnable Aggregation Weights

Problem Statement

Federated Learning (FL) allows multiple clients to collaboratively train a global model without sharing private data. However, the system is vulnerable to Byzantine clients—malicious actors that submit arbitrary or corrupted model updates to compromise the global model. This challenge is exacerbated in heterogeneous settings where data distributions vary significantly across clients. In such scenarios, distinguishing between benign updates (influenced by data heterogeneity) and malicious updates is difficult.

Existing Byzantine-robust FL methods typically follow a two-step process: first, they identify and remove malicious clients, and second, they assign uniform aggregation weights to the remaining benign clients. The authors argue that this uniform weighting approach is suboptimal in heterogeneous environments. After removing malicious clients, the remaining benign clients may still exhibit imbalanced data distributions (e.g., missing certain labels). Assigning uniform weights fails to adapt to this imbalance, potentially degrading model accuracy by not adequately prioritizing underrepresented classes.

Methodology

The paper proposes FedLAW (Federated Learning with Learnable Aggregation Weights), a novel framework that treats aggregation weights as learnable parameters rather than fixed values or simple post-filtering assignments.

1. Problem Formulation

The authors formulate the FL objective as a joint optimization problem over both the global model parameters $\theta$ and the aggregation weights $w$:
$$\min_{\theta \in \mathbb{R}^d, w \in \Delta^+_{t,\ell_0}} \sum_{i=1}^n w_i f_i(\theta)$$
where $f_i(\theta)$ is the loss function of client $i$. The weight vector $w$ is constrained to a sparse unit-capped simplex $\Delta^+_{t,\ell_0}$, defined by:

• $\sum w_i = 1$ and $w_i \geq 0$ (standard simplex constraints).
• $w_i \leq t$ (capping individual client influence).
• $\|w\|_0 \leq s$ (sparsity constraint, limiting the number of non-zero weights).

By setting $t = 1/(n - b_f)$ and $s = n - b_f$ (where $b_f$ is the number of Byzantine clients), the feasible region forces the exclusion of up to $b_f$ clients (setting their weights to zero) while distributing the remaining weight equally among the selected $n - b_f$ clients. This transforms the binary problem of "detecting and removing" into a continuous optimization problem.

2. Proposed Algorithm

To solve this non-convex joint optimization, the authors develop an alternating minimization algorithm (Algorithm 2) based on a nested optimization structure:

1. Inner Step (Model Update): For fixed weights $w$, the server approximates the client loss functions with quadratic surrogates and computes a model update $\theta_{k+1}(w) = \theta_k - \alpha G_k w$, where $G_k$ contains the client gradients.
2. Outer Step (Weight Update): For the updated model, the algorithm optimizes the weights $w$ to minimize the global objective. This step involves minimizing a quadratic approximation of the objective function with respect to $w$, followed by a projection onto the sparse unit-capped simplex.

The weight update leverages both loss information (to detect data poisoning) and gradient alignment (to detect model poisoning/inverse gradient attacks). Malicious clients, whose gradients deviate from the coherent cluster of benign clients, are naturally down-weighted or excluded by the sparsity constraint.

3. Communication Protocol

The algorithm requires two communication rounds per training epoch:

1. The server broadcasts $\theta_k$; clients return gradients and losses.
2. The server computes a provisional update $\tilde{\theta}_{k+1}$, broadcasts it, and clients return updated gradients and losses to compute the final weight update and the next model state.
   The authors note that while this adds a round, the accelerated convergence of the global model often offsets this cost, resulting in comparable or lower total communication rounds to reach a target accuracy.

Key Contributions

1. Novel Optimization Formulation: The paper introduces a joint optimization framework where aggregation weights are learnable parameters, embedding Byzantine defense directly into the learning objective rather than treating it as a separate filtering step.
2. Algorithm Development: An alternating minimization algorithm is proposed with a specific projection operator for the sparse unit-capped simplex, ensuring efficient computation of robust weights.
3. Theoretical Guarantees:
   • Byzantine Resilience: Theorem 2 establishes that the proposed aggregator is $(\zeta, b_f)$-Byzantine resilient with high probability, providing an explicit error bound that decomposes into sources like loss heterogeneity, gradient heterogeneity, and sampling noise.
   • Convergence: Theorem 3 proves that the algorithm converges to a neighborhood of the optimum for both non-convex and strongly convex settings, with the error radius scaling with the adversary's influence.
4. Empirical Validation: Extensive experiments on MNIST and CIFAR-10 under five attack types (Label Flipping, Inverse Gradient, Backdoor, Double Attack, and LIE) and varying levels of data heterogeneity demonstrate the method's superiority.

Experimental Results

The authors evaluated FedLAW against state-of-the-art baselines (Krum, Trimmed Mean, Bulyan, CwMed, RFA, CClip, Huber Aggregator) across diverse scenarios:

• Performance: FedLAW consistently outperformed existing methods, particularly in high-heterogeneity ($q=0.9$) and high-malicious-fraction (up to 40%) settings. For instance, under an inverse gradient attack with 40% malicious clients on MNIST, FedLAW achieved 3.6% higher test accuracy than the next-best defense.
• Stability: Unlike other methods that diverged or suffered sharp accuracy drops under severe attacks (e.g., Double Attack or high LIE fractions), FedLAW maintained robust performance with graceful degradation.
• Weight Dynamics: Visualizations showed that benign clients quickly converged to stable, non-trivial weights, while malicious clients were consistently suppressed to near-zero weights within the first 20 epochs.
• Detection: The method demonstrated high precision and recall in identifying malicious clients, effectively neutralizing both data poisoning and model poisoning attacks.

Significance and Claims

The paper claims that FedLAW represents a significant advancement in secure federated learning by addressing the limitations of uniform weighting in heterogeneous environments. By integrating adaptive weighting into the learning process, the method achieves:

• Enhanced Robustness: It effectively neutralizes malicious clients without requiring prior knowledge of the number of attackers or complex pre-filtering mechanisms.
• Fairness and Adaptability: It adaptively balances contributions from benign clients, ensuring that data heterogeneity does not lead to the neglect of specific data distributions.
• Theoretical Soundness: The work provides rigorous convergence and resilience guarantees, linking the static resilience of the aggregator to the dynamic stability of the learning algorithm.

The authors conclude that their findings highlight the benefits of integrating adaptive aggregation into the learning process, paving the way for more secure and equitable FL deployments in real-world, heterogeneous, and adversarial settings.