DRAG: Divergence-based Adaptive Aggregation in Federated learning on Non-IID Data

TL;DR

This paper introduces DRAG, a divergence-based adaptive aggregation method in federated learning that mitigates client-drift on non-IID data, achieves sublinear convergence, and enhances robustness against Byzantine attacks without extra communication overhead.

Contribution

The paper proposes a novel divergence metric and an adaptive aggregation algorithm that improves convergence and robustness in federated learning on heterogeneous data.

Findings

DRAG outperforms state-of-the-art algorithms in managing client-drift.

DRAG achieves a proven sublinear convergence rate.

DRAG enhances resilience against Byzantine attacks.

Abstract

Local stochastic gradient descent (SGD) is a fundamental approach in achieving communication efficiency in Federated Learning (FL) by allowing individual workers to perform local updates. However, the presence of heterogeneous data distributions across working nodes causes each worker to update its local model towards a local optimum, leading to the phenomenon known as ``client-drift" and resulting in slowed convergence. To address this issue, previous works have explored methods that either introduce communication overhead or suffer from unsteady performance. In this work, we introduce a novel metric called ``degree of divergence," quantifying the angle between the local gradient and the global reference direction. Leveraging this metric, we propose the divergence-based adaptive aggregation (DRAG) algorithm, which dynamically ``drags" the received local updates toward the reference direction in each round without requiring extra communication overhead. Furthermore, we establish a rigorous convergence analysis for DRAG, proving its ability to achieve a sublinear convergence rate. Compelling experimental results are presented to illustrate DRAG's superior performance compared to state-of-the-art algorithms in effectively managing the client-drift phenomenon. Additionally, DRAG exhibits remarkable resilience against certain Byzantine attacks. By securely sharing a small sample of the client's data with the FL server, DRAG effectively counters these attacks, as demonstrated through comprehensive experiments.