Weighted Loss Methods for Robust Federated Learning under Data Heterogeneity

2025-06-11

AI Summary

In federated learning under data heterogeneity, distinguishing honest anomalous gradients from Byzantine malicious ones remains challenging. To address this, we propose the Worker Label Alignment Loss (WoLA), a novel weighted loss function that— for the first time—incorporates label alignment into gradient direction modeling to explicitly enhance consistency among honest gradients across heterogeneous clients. Our method integrates WoLA-based optimization, robust aggregation, and rigorous theoretical convergence analysis. Evaluated on multiple heterogeneous benchmarks, it achieves >95% Byzantine attack detection success rate and improves convergence stability by over 40%, significantly outperforming existing state-of-the-art approaches. Furthermore, we provide formal theoretical guarantees on both convergence and robustness under realistic heterogeneity and adversarial assumptions.

Technology Category

• Machine Learning: Adversarial Learning & Robustness
• Natural Language Processing: Safety and Robustness
• Computer Vision: Adversarial Attacks & Robustness

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization

Abstract

Federated learning (FL) is a machine learning paradigm that enables multiple data holders to collaboratively train a machine learning model without sharing their training data with external parties. In this paradigm, workers locally update a model and share with a central server their updated gradients (or model parameters). While FL seems appealing from a privacy perspective, it opens a number of threats from a security perspective as (Byzantine) participants can contribute poisonous gradients (or model parameters) harming model convergence. Byzantine-resilient FL addresses this issue by ensuring that the training proceeds as if Byzantine participants were absent. Towards this purpose, common strategies ignore outlier gradients during model aggregation, assuming that Byzantine gradients deviate more from honest gradients than honest gradients do from each other. However, in heterogeneous settings, honest gradients may differ significantly, making it difficult to distinguish honest outliers from Byzantine ones. In this paper, we introduce the Worker Label Alignement Loss (WoLA), a weighted loss that aligns honest worker gradients despite data heterogeneity, which facilitates the identification of Byzantines' gradients. This approach significantly outperforms state-of-the-art methods in heterogeneous settings. In this paper, we provide both theoretical insights and empirical evidence of its effectiveness.

Problem

• Distinguishing Byzantine gradients from honest ones in federated learning
• Addressing data heterogeneity in federated learning robustness
• Improving model convergence with weighted loss methods

Innovation

• Weighted loss aligns honest worker gradients
• WoLA facilitates Byzantine gradient identification
• Outperforms state-of-the-art in heterogeneous settings