FedGraM: Defending Against Untargeted Attacks in Federated Learning via Embedding Gram Matrix

TL;DR

FedGraM tackles untargeted attacks in federated learning by computing embedding Gram matrices from a server-held auxiliary dataset to infer each client's generalization capability. By pruning the top-C high-norm Gram matrices and averaging the rest, FedGraM provides a robust aggregation that remains effective under non-IID data distributions with limited server data. Empirical evaluations across CIFAR-10, SVHN, and CIFAR-100 show FedGraM outperforming state-of-the-art defenses in most scenarios, with ablations supporting the choice of C and auxiliary data. The approach is complemented by an analysis of its limitations under adaptive attacks and potential gains from combining with other robust defenses for stronger protection in practice.

Abstract

Federated Learning (FL) enables geographically distributed clients to collaboratively train machine learning models by sharing only their local models, ensuring data privacy. However, FL is vulnerable to untargeted attacks that aim to degrade the global model's performance on the underlying data distribution. Existing defense mechanisms attempt to improve FL's resilience against such attacks, but their effectiveness is limited in practical FL environments due to data heterogeneity. On the contrary, we aim to detect and remove the attacks to mitigate their impact. Generalization contribution plays a crucial role in distinguishing untargeted attacks. Our observations indicate that, with limited data, the divergence between embeddings representing different classes provides a better measure of generalization than direct accuracy. In light of this, we propose a novel robust aggregation method, FedGraM, designed to defend against untargeted attacks in FL. The server maintains an auxiliary dataset containing one sample per class to support aggregation. This dataset is fed to the local models to extract embeddings. Then, the server calculates the norm of the Gram Matrix of the embeddings for each local model. The norm serves as an indicator of each model's inter-class separation capability in the embedding space. FedGraM identifies and removes potentially malicious models by filtering out those with the largest norms, then averages the remaining local models to form the global model. We conduct extensive experiments to evaluate the performance of FedGraM. Our empirical results show that with limited data samples used to construct the auxiliary dataset, FedGraM achieves exceptional performance, outperforming state-of-the-art defense methods.

Figures (6)

• Figure 1: The norm distribution and the accuracy distribution across clients at the beginning stage of the training. The Gram matrix norm is calculated as we shown in Section \ref{['sec:method']} with only 10 data samples. The accuracy is estimated with 100 data samples. The detailed setting of this demo experiment are refereed to Section \ref{['supp_demoexpsetting']} in appendix. In each figure, we arrange the clients in descending order based on the corresponding values. Accordingly, with limited data, the Gram matrix norm is a better indicator to capture the generalization divergence between local models.
• Figure 2: The overview of FedGraM. For each uploaded local model, the server feeds the auxiliary dataset to its representation layers to extract the corresponding embeddings. The server then calculates the product of the normalized embedding matrix and its transpose to get Gram Matrix. After that, the server calculates the norm of each Gram Matrix and removes the model with the highest norm. At the end, the server averages the remaining models to generate the global model.
• Figure 3: The experiment results of the comparison between FedGraM and SOTA defenses. We draw a spider chart for each defense to represent its performance in defending against untargeted attacks. Each dimension of the map represents the performance of the corresponding defense in defending a certain type of attack. From the top row to the bottom row, we present the results in CIFAR10, SVHN, and CIFAR100.
• Figure 4: From lest to right, we show the experiment results of Gram matrix norm distribution, Fidelity of FedGraM and Impact of C.
• Figure 5: The experiment results of the impact of the quality of the auxiliary dataset on the performance of FedGraM.