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

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

Problem Statement
The paper addresses the trade-off between communication efficiency and model performance in centralized distributed data-parallel training of Deep Neural Networks (DNNs). Standard synchronous gradient averaging (e.g., DDP) introduces significant communication bottlenecks due to frequent synchronization. While local update methods (e.g., LocalSGD) reduce communication by allowing workers to perform multiple local steps before synchronization, they often suffer from inferior generalization performance compared to fully synchronous methods. Existing strategies to mitigate this, such as soft-consensus methods (e.g., EASGD), often fail to prevent workers from collapsing into narrow, sharp minima in the loss landscape, which correlates with poor generalization. The core challenge is to develop a communication-efficient method that matches or exceeds the performance of synchronous averaging without incurring significant computational overhead.

Methodology
The authors propose Distributed Pull-Push Force (DPPF), a training strategy that incorporates a "push" mechanism to encourage the collaborative recovery of wide, flat minima.

1. Sharpness Measure (Inverse Mean Valley): The authors first introduce a new sharpness metric called Inverse Mean Valley (Inv. MV). This metric quantifies the average diameter of the valley surrounding converged worker parameters. Empirical analysis shows that Inv. MV strongly correlates with the generalization gap in distributed settings, outperforming existing sharpness metrics in these scenarios.
2. Relaxation and Regularization: Directly computing Inv. MV is computationally expensive as it requires line searches along worker directions. To make this efficient, the authors propose a lightweight relaxation: treating current worker parameters as approximate boundary points. This leads to a regularization term $R = -\lambda_r \sum \|x_i - x_A\|^2$, where $x_A$ is the average of worker parameters.
3. The Pull-Push Mechanism:
   • Pull: Standard distributed training involves a "pull" force where workers are moved toward a consensus variable (typically the average $x_A$) via a convex combination: $x_m \leftarrow (1-\alpha)x_m + \alpha x_A$.
   • Push: The relaxation of Inv. MV introduces a "push" force that counteracts the consensus step. The gradient of the regularization term exerts a force pushing workers away from the average.
   • Integration: These forces are combined into a single update step. The push force prevents the "valley collapse" phenomenon where workers converge too tightly into sharp minima, instead positioning them near the boundaries of wide, flat regions. The final update rule (when $x_C = x_A$) is:
     $$x_m \leftarrow x_m + (x_A - x_m) \left( \alpha - \frac{\lambda}{\|x_m - x_A\|} \right)$$
     Here, $\alpha$ is the pull strength and $\lambda$ is the push strength.

Key Contributions

• New Sharpness Metric: Introduction of the Inverse Mean Valley (Inv. MV), which demonstrates a strong correlation with generalization gaps in distributed training, particularly in settings with and without data augmentation.
• DPPF Algorithm: A novel algorithm that integrates a lightweight relaxation of Inv. MV as a regularizer. This creates a "push" force that counteracts the "pull" of consensus steps, enabling workers to collaboratively span wide valleys.
• Theoretical Analysis:
  • Valley Width Control: The authors prove that the asymptotic distance between workers and the global average (valley width) is governed by the ratio of push to pull strengths ($\lambda/\alpha$).
  • Generalization Guarantees: Using a PAC-Bayes framework, they show that increasing the valley width (by adjusting $\lambda/\alpha$) leads to tighter generalization bounds, provided certain technical conditions on the drift of the model parameters are met.
  • Convergence: The method is proven to converge in non-convex settings.
• Empirical Validation: Extensive experiments on CIFAR-10, CIFAR-100, and ImageNet using various architectures (ResNet, PyramidNet, ViT) and optimizers (SGD, AdamW).

Results

• Performance vs. Communication: DPPF consistently achieves lower test errors than baseline methods (DDP SGD, LocalSGD, and LocalSGD with Quadratic Synchronization Rule - QSR) while significantly reducing communication volume. For example, on ImageNet with ResNet-50, DPPF achieves a test error of 23.07% with only 25% of the communication rounds required by DDP SGD.
• Comparison with SAM: DPPF achieves performance comparable to Sharpness-Aware Minimization (SAM), a state-of-the-art local flatness-seeking method, but does so through distributed inter-worker repulsion rather than local perturbations. Combining DPPF with local SAM (DPPF SAM) yields further improvements.
• Robustness: The method performs well across different communication periods ($\tau$), worker counts (4 and 8 GPUs), and even in non-IID federated learning settings (when coupled with SCAFFOLD or FedLESAM).
• Visualization: Loss landscape visualizations confirm that DPPF prevents worker collapse (observed in SimpleAvg) and locates flatter minima where the training error remains stable over a wider area.

Significance and Claims
The paper claims that DPPF offers a substantial improvement in the trade-off between communication efficiency and model generalization. By explicitly modeling the "valley width" as a controllable hyperparameter via the push-pull ratio, DPPF allows practitioners to target specific flatness properties without the computational cost of second-order methods or extensive landscape sampling. The authors assert that their approach enables communication-efficient methods to match or exceed the performance of synchronous gradient averaging, effectively solving the performance degradation often seen in LocalSGD variants. The work suggests that preventing worker collapse via a repulsive force is essential for recovering wide minima in distributed settings, a mechanism that is distinct from and complementary to local flatness-seeking strategies like SAM. Future directions mentioned include applying this mechanism to ensemble learning and extending it to self-supervised settings.