Robust Optimization for Federated Learning Under Client Drift

Abstract

Federated learning (FL) enables collaborative model training without requiring clients to share their raw data with a central server. However, the statistical heterogeneity that characterizes many federated environments can cause local optimization trajectories to diverge from the direction of the global objective. Client drift has always been a major concern in federated learning, and it becomes more so when the client performs several local optimization iterations before communicating. The popular approaches to counter this, namely FedProx and SCAFFOLD, rely on proximal regularization and control-variate-driven drift correction, respectively. In the literature, it is believed that most of the existing methods implicitly do not make a distinction between persistent client disagreement and temporary deviations.Persistent Client-Drift Aware Proximal Federated Learning (PCD-Prox) is a computationally efficient, adaptive regularization technique that takes into account the time-invariant nature of client drift. Rather than a single static proximal coefficient for all clients, PCD-Prox retains an exponential average of each client's direction of deviation from the global update, and uses this to scale the strength of their local proximal regularization.The method preserves the typical fed-averaging scheme, with the addition of one scalar state variable per user. In order to verify the method, is to perform a small numerical validation with 50 heterogeneous users, partial participation, long-term and short-term user shifts, and 20 seeds.There are four methods evaluated in this test: Fed Avg, Fed Prox, a proximal version depending on the instantaneous drift and PCD-Prox. The results indicate that PCD-Prox exhibits the smallest final value for the objective for both persistent and transient drifting. It performs better than Fed Avg by almost 0.24% for the persistent drifting and by 0.08% in the case of transient drifting.The sensitivity analysis provides additional evidence of the importance of the parameter for temporal persistence on the accumulative drift estimate as well as the final optimization performance. These findings suggest that temporal information in the clients' update pattern can be used to guide adaptive local regularization without changing the existing aggregation framework. Further experiments with real non-IID data are necessary to establish the approach more generally.