SADDLe: Sharpness-Aware Decentralized Deep Learning with Heterogeneous Data

Decentralized training enables learning with distributed datasets generated at different locations without relying on a central server. In realistic scenarios, the data distribution across these sparsely connected learning agents can be significantly heterogeneous, leading to local model over-fitting and poor global model generalization. Another challenge is the high communication cost of training models in such a peer-to-peer fashion without any central coordination. In this paper, we jointly tackle these two-fold practical challenges by proposing SADDLe, a set of sharpness-aware de-centralized deep learning algorithms. SADDLe leverages Sharpness-Aware Minimization (SAM) to seek a flatter loss landscape during training, resulting in better model generalization as well as enhanced robustness to communication compression. We present two versions of our approach and demonstrate its effectiveness through extensive experiments on various Computer Vision datasets (CIFAR-10, CIFAR-100, Imagenette, and ImageNet), model architectures, and graph topologies. Our results show that SADDLe leads to 1-20% improvement in test accuracy as compared to existing techniques while incurring a minimal accuracy drop (~ 1%) in the presence of up to 4× compression.