Implicit Gradient Alignment in Distributed and Federated Learning

TL;DR

This paper introduces GradAlign, a novel algorithm that promotes gradient alignment in distributed and federated learning, leveraging implicit regularization to improve generalization even with large batch sizes.

Contribution

The paper proposes GradAlign, a new method that induces implicit regularization to align gradients across clients, enhancing generalization in large-batch distributed and federated learning.

Findings

GradAlign improves test accuracy in federated learning scenarios.

Gradient alignment correlates with better generalization performance.

The method enables large-batch training without losing regularization benefits.

Abstract

A major obstacle to achieving global convergence in distributed and federated learning is the misalignment of gradients across clients, or mini-batches due to heterogeneity and stochasticity of the distributed data. In this work, we show that data heterogeneity can in fact be exploited to improve generalization performance through implicit regularization. One way to alleviate the effects of heterogeneity is to encourage the alignment of gradients across different clients throughout training. Our analysis reveals that this goal can be accomplished by utilizing the right optimization method that replicates the implicit regularization effect of SGD, leading to gradient alignment as well as improvements in test accuracies. Since the existence of this regularization in SGD completely relies on the sequential use of different mini-batches during training, it is inherently absent when training with large mini-batches. To obtain the generalization benefits of this regularization while increasing parallelism, we propose a novel GradAlign algorithm that induces the same implicit regularization while allowing the use of arbitrarily large batches in each update. We experimentally validate the benefits of our algorithm in different distributed and federated learning settings.