Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Tolga Dimlioglu, Anna Choromanska

TL;DR

This work tackles the challenge of communication-efficient distributed training by reexamining flat minima and introducing a practical sharpness proxy, Inverse Mean Valley (Inv. MV). It embeds a lightweight relaxed regularizer into the training objective to realize Distributed Pull-Push Force (DPPF), which couples a pull toward consensus with a push that encourages exploration of wide, flat valleys. The authors prove that the final valley width scales with the pull/push ratio $\lambda/\alpha$ and provide a PAC-Bayes generalization bound linked to valley width, then validate the approach across CNNs and ViTs, showing improved generalization and reduced communication relative to strong baselines and SAM-like performance. Collectively, DPPF offers a theoretically grounded, practically effective strategy for achieving flat minima under communication constraints, with broad applicability and robust empirical gains.

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.

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

TL;DR

This work tackles the challenge of communication-efficient distributed training by reexamining flat minima and introducing a practical sharpness proxy, Inverse Mean Valley (Inv. MV). It embeds a lightweight relaxed regularizer into the training objective to realize Distributed Pull-Push Force (DPPF), which couples a pull toward consensus with a push that encourages exploration of wide, flat valleys. The authors prove that the final valley width scales with the pull/push ratio and provide a PAC-Bayes generalization bound linked to valley width, then validate the approach across CNNs and ViTs, showing improved generalization and reduced communication relative to strong baselines and SAM-like performance. Collectively, DPPF offers a theoretically grounded, practically effective strategy for achieving flat minima under communication constraints, with broad applicability and robust empirical gains.

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.

Paper Structure

This paper contains 36 sections, 6 theorems, 80 equations, 45 figures, 10 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminary
  4. Mean Valley Measure
  5. Comparison with Other Measures
  6. Distributed Pull-Push Force
  7. Theoretical Analysis
  8. Empirical Results
  9. Push Mechanism with Soft-Consensus Methods
  10. Comparison with Other Communication-Efficient Methods
  11. DPPF Achieves SAM-like Performance While Being Communication-Efficient
  12. Ablation Studies
  13. Analysis of the Pull-Push Mechanism
  14. Error Landscape Visualizations
  15. Closing Remarks
  16. ...and 21 more sections

Key Result

Theorem 1

Consider $M$ workers running DPPF with pull strength $\alpha$, push strength $\lambda$, communication period $\tau$, and local step size $\eta$. Let $x_{m,k}^t$ be worker $m$’s parameters at local iteration $t$ of communication round $k$, and $x_{A,k}$ the average over workers before the round. Defi where the expectation is over stochastic randomness. Letting $\eta\to0$ and $M\gg1$ yields $\lim_{k

Figures (45)

  • Figure 1: Communication volume (ratio between the number of communication rounds till convergence to the total number of local iterations; lower is better) vs. test error (%) for baseline distributed training methods (LocalSGD, QSR and DDP SGD) and our approach, DPPF. (W: Worker)
  • Figure 2: Comparison of training with weaker pulling force and using the pull-push mechanism.
  • Figure 3: Analysis of the tug-of-war between pull and push forces.
  • Figure 4: 2D contour plots of training error (%) around the average variable ($x_A$). For SimpleAVG, the error quickly rises to 100%. In contrast, the error remains stable at approximately 0.03% around the minima found by DPPF.
  • Figure 5: 3D train error (%) plots.
  • ...and 40 more figures

Theorems & Definitions (13)

  • Theorem 1: DPPF: Final Valley Width
  • Theorem 2: DPPF: Wider Valley Tightens the Generalization Gap
  • Remark 1
  • Remark 2
  • Remark 3
  • Lemma 1
  • proof
  • Lemma 2
  • proof
  • Theorem 3: Valley Width under DPPF
  • ...and 3 more