FedTK: A Method for Non-IID Fed Learning via Trace-Kernel and Robust Aggregation

With data privacy issues coming to the forefront, federated learning has arisen as a distributed machine learning approach. It enables clients to conduct training using their own local data and jointly train models by uploading model parameters. This way, model performance can be sustained while safeguarding privacy. However, federated learning encounters the issue of model performance decline because of the Non-Independent and Identical Distribution of Data (non-iid) and malicious attacks that impact the global model performance. To address the data heterogeneity problem, existing methods suffer from high computational cost, limited applicability, or inability to adequately capture data features; in the model aggregation phase, the traditional weighted average aggregation method is sensitive to anomalous clients, and is still facing challenges despite multiple improved algorithms. This paper proposes the FedTK method, which is designed to alleviate the model dimension collapse problem caused by non-iid data by combining the trace and kernel paradigms to approximate the singular value variance. And at the same time, we introduce a robust aggregation strategy based on the geometric median to reduce the impact of anomalous clients on the global model, and to enhance the model stability and reliability. Experiments are conducted on MNIST, CIFAR10 and CIFAR100 datasets to mimic varying degrees of non-iid data and data poisoning and omniscient poisoning situations. The outcomes reveal that the proposed approach attains greater model training accuracy in contrast to other methods across different scenarios. It also effectively enhances the generalization capacity of the federated learning model under non-iid data, rendering federated learning more practical.