Communication-Efficient Distributed Training for Collaborative Flat Optima Recovery in Deep Learning

TL;DR

This paper introduces a communication-efficient distributed training method, DPPF, that encourages deep neural networks to find flatter minima, leading to better generalization and improved performance over existing approaches.

Contribution

The paper proposes the DPPF algorithm, incorporating a novel sharpness measure as a regularizer, with theoretical analysis and empirical validation demonstrating its effectiveness in distributed deep learning.

Findings

DPPF outperforms other communication-efficient methods in generalization.

DPPF effectively locates flatter minima in loss landscapes.

Theoretical analysis confirms convergence and stability of DPPF.

Abstract

We study centralized distributed data parallel training of deep neural networks (DNNs), aiming to improve the trade-off between communication efficiency and model performance of the local gradient methods. To this end, we revisit the flat-minima hypothesis, which suggests that models with better generalization tend to lie in flatter regions of the loss landscape. We introduce a simple, yet effective, sharpness measure, Inverse Mean Valley, and demonstrate its strong correlation with the generalization gap of DNNs. We incorporate an efficient relaxation of this measure into the distributed training objective as a lightweight regularizer that encourages workers to collaboratively seek wide minima. The regularizer exerts a pushing force that counteracts the consensus step pulling the workers together, giving rise to the Distributed Pull-Push Force (DPPF) algorithm. Empirically, we show that DPPF outperforms other communication-efficient approaches and achieves better generalization performance than local gradient methods and synchronous gradient averaging, while maintaining communication efficiency. In addition, our loss landscape visualizations confirm the ability of DPPF to locate flatter minima. On the theoretical side, we show that DPPF guides workers to span flat valleys, with the final valley width governed by the interplay between push and pull strengths, and that its pull-push dynamics is self-stabilizing. We further provide generalization guarantees linked to the valley width and prove convergence in the non-convex setting.