# AB-Training: A Communication-Efficient Approach for Distributed Low-Rank Learning

arXiv: https://arxiv.org/abs/2405.01067 · HTML: https://mlsanity.com/paper/2405.01067/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 30 of its 43 references that are in the map.
- First posted on arXiv 2 May 2024.

## Builds on (most cited first, 30 of 30)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Decoupled Weight Decay Regularization](https://mlsanity.com/paper/1711.05101/) (2017)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Channel Pruning for Accelerating Very Deep Neural Networks](https://mlsanity.com/paper/1707.06168/) (2017)
- [ThiNet: A Filter Level Pruning Method for Deep Neural Network Compression](https://mlsanity.com/paper/1707.06342/) (2017)
- [Entropy-SGD: Biasing Gradient Descent Into Wide Valleys](https://mlsanity.com/paper/1611.01838/) (2016)
- [Soft Filter Pruning for Accelerating Deep Convolutional Neural Networks](https://mlsanity.com/paper/1808.06866/) (2018)
- [PowerSGD: Practical Low-Rank Gradient Compression for Distributed Optimization](https://mlsanity.com/paper/1905.13727/) (2019)
- [Don't Use Large Mini-Batches, Use Local SGD](https://mlsanity.com/paper/1808.07217/) (2018)
- [Demystifying Parallel and Distributed Deep Learning: An In-Depth Concurrency Analysis](https://mlsanity.com/paper/1802.09941/) (2018)
- [Channel Pruning via Automatic Structure Search](https://mlsanity.com/paper/2001.08565/) (2020)
- [On the Generalization Benefit of Noise in Stochastic Gradient Descent](https://mlsanity.com/paper/2006.15081/) (2020)
- [Stable Low-rank Tensor Decomposition for Compression of Convolutional Neural Network](https://mlsanity.com/paper/2008.05441/) (2020)
- [Provable Benefit of Orthogonal Initialization in Optimizing Deep Linear Networks](https://mlsanity.com/paper/2001.05992/) (2020)
- [Pufferfish: Communication-efficient Models At No Extra Cost](https://mlsanity.com/paper/2103.03936/) (2021)
- [Learning Low-rank Deep Neural Networks via Singular Vector Orthogonality Regularization and Singular Value Sparsification](https://mlsanity.com/paper/2004.09031/) (2020)
- [Lightweight and Efficient End-to-End Speech Recognition Using Low-Rank Transformer](https://mlsanity.com/paper/1910.13923/) (2019)
- [Low-rank lottery tickets: finding efficient low-rank neural networks via matrix differential equations](https://mlsanity.com/paper/2205.13571/) (2022)
- [Compressing Neural Networks: Towards Determining the Optimal Layer-wise Decomposition](https://mlsanity.com/paper/2107.11442/) (2021)
- [Cuttlefish: Low-Rank Model Training without All the Tuning](https://mlsanity.com/paper/2305.02538/) (2023)
- [Learned Gradient Compression for Distributed Deep Learning](https://mlsanity.com/paper/2103.08870/) (2021)
- [Low Rank Regularization: A Review](https://mlsanity.com/paper/1808.04521/) (2018)
- [Play and Prune: Adaptive Filter Pruning for Deep Model Compression](https://mlsanity.com/paper/1905.04446/) (2019)
- [Dimensionality Reduced Training by Pruning and Freezing Parts of a Deep Neural Network, a Survey](https://mlsanity.com/paper/2205.08099/) (2022)
- [Greenformer: Factorization Toolkit for Efficient Deep Neural Networks](https://mlsanity.com/paper/2109.06762/) (2021)
- [Massively Parallel Genetic Optimization through Asynchronous Propagation of Populations](https://mlsanity.com/paper/2301.08713/) (2023)
- [Accelerating Neural Network Training with Distributed Asynchronous and Selective Optimization (DASO)](https://mlsanity.com/paper/2104.05588/) (2021)
- [Harnessing Orthogonality to Train Low-Rank Neural Networks](https://mlsanity.com/paper/2401.08505/) (2024)

## Abstract (arXiv)

Communication bottlenecks severely hinder the scalability of distributed neural network training, particularly in high-performance computing (HPC) environments. We introduce AB-training, a novel data-parallel method that leverages low-rank representations and independent training groups to significantly reduce communication overhead. Our experiments demonstrate an average reduction in network traffic of approximately 70.31% across various scaling scenarios, increasing the training potential of communication-constrained systems and accelerating convergence at scale. AB-training also exhibits a pronounced regularization effect at smaller scales, leading to improved generalization while maintaining or even reducing training time. We achieve a remarkable 44.14 : 1 compression ratio on VGG16 trained on CIFAR-10 with minimal accuracy loss, and outperform traditional data parallel training by 1.55% on ResNet-50 trained on ImageNet-2012. While AB-training is promising, our findings also reveal that large batch effects persist even in low-rank regimes, underscoring the need for further research into optimized update mechanisms for massively distributed training.