# WASH: Train your Ensemble with Communication-Efficient Weight Shuffling, then Average

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

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 48 references that are in the map.
- First posted on arXiv 27 May 2024.

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

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [Linear Mode Connectivity and the Lottery Ticket Hypothesis](https://mlsanity.com/paper/1912.05671/) (2019)
- [Local SGD Converges Fast and Communicates Little](https://mlsanity.com/paper/1805.09767/) (2018)
- [Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs](https://mlsanity.com/paper/1802.10026/) (2018)
- [Federated Optimization: Distributed Machine Learning for On-Device Intelligence](https://mlsanity.com/paper/1610.02527/) (2016)
- [Federated Learning with Matched Averaging](https://mlsanity.com/paper/2002.06440/) (2020)
- [Deep Ensembles: A Loss Landscape Perspective](https://mlsanity.com/paper/1912.02757/) (2019)
- [SWAD: Domain Generalization by Seeking Flat Minima](https://mlsanity.com/paper/2102.08604/) (2021)
- [Bayesian Nonparametric Federated Learning of Neural Networks](https://mlsanity.com/paper/1905.12022/) (2019)
- [What is being transferred in transfer learning?](https://mlsanity.com/paper/2008.11687/) (2020)
- [Don't Use Large Mini-Batches, Use Local SGD](https://mlsanity.com/paper/1808.07217/) (2018)
- [EXTRA: An Exact First-Order Algorithm for Decentralized Consensus Optimization](https://mlsanity.com/paper/1404.6264/) (2014)
- [Efficient Deep Learning: A Survey on Making Deep Learning Models Smaller, Faster, and Better](https://mlsanity.com/paper/2106.08962/) (2021)
- [Diverse Weight Averaging for Out-of-Distribution Generalization](https://mlsanity.com/paper/2205.09739/) (2022)
- [D2: Decentralized Training over Decentralized Data](https://mlsanity.com/paper/1803.07068/) (2018)
- [SlowMo: Improving Communication-Efficient Distributed SGD with Slow Momentum](https://mlsanity.com/paper/1910.00643/) (2019)
- [Loss Surface Simplexes for Mode Connecting Volumes and Fast Ensembling](https://mlsanity.com/paper/2102.13042/) (2021)
- [No One Representation to Rule Them All: Overlapping Features of Training Methods](https://mlsanity.com/paper/2110.12899/) (2021)
- [Federated Dropout -- A Simple Approach for Enabling Federated Learning on Resource Constrained Devices](https://mlsanity.com/paper/2109.15258/) (2021)
- [ZeRO++: Extremely Efficient Collective Communication for Giant Model Training](https://mlsanity.com/paper/2306.10209/) (2023)
- [Consensus Control for Decentralized Deep Learning](https://mlsanity.com/paper/2102.04828/) (2021)
- [Cross-Gradient Aggregation for Decentralized Learning from Non-IID data](https://mlsanity.com/paper/2103.02051/) (2021)
- [Trade-offs of Local SGD at Scale: An Empirical Study](https://mlsanity.com/paper/2110.08133/) (2021)
- [Communication-efficient SGD: From Local SGD to One-Shot Averaging](https://mlsanity.com/paper/2106.04759/) (2021)

## Built on by (most cited first, 9 of 9)

- [Model Merging in LLMs, MLLMs, and Beyond: Methods, Theories, Applications and Opportunities](https://mlsanity.com/paper/2408.07666/) (2024)
- [Streaming DiLoCo with overlapping communication: Towards a Distributed Free Lunch](https://mlsanity.com/paper/2501.18512/) (2025)
- [Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCo](https://mlsanity.com/paper/2503.09799/) (2025)
- [Overcoming the Communication-Performance Tradeoff in LLM Pretraining](https://mlsanity.com/paper/2508.15706/) (2025)
- [Model Parallelism With Subnetwork Data Parallelism](https://mlsanity.com/paper/2507.09029/) (2025)
- [Probabilistic Token Alignment for Large Language Model Fusion](https://mlsanity.com/paper/2509.17276/) (2025)
- [AsyncMesh: Fully Asynchronous Optimization for Data and Pipeline Parallelism](https://mlsanity.com/paper/2601.22442/) (2026)
- [LOSCAR-SGD: Local SGD with Communication-Computation Overlap and Delay-Corrected Sparse Model Averaging](https://mlsanity.com/paper/2605.20866/) (2026)
- [Communication-Efficient Distributed Training for Collaborative Flat Optima Recovery in Deep Learning](https://mlsanity.com/paper/2507.20424/) (2025)

## Abstract (arXiv)

The performance of deep neural networks is enhanced by ensemble methods, which average the output of several models. However, this comes at an increased cost at inference. Weight averaging methods aim at balancing the generalization of ensembling and the inference speed of a single model by averaging the parameters of an ensemble of models. Yet, naive averaging results in poor performance as models converge to different loss basins, and aligning the models to improve the performance of the average is challenging. Alternatively, inspired by distributed training, methods like DART and PAPA have been proposed to train several models in parallel such that they will end up in the same basin, resulting in good averaging accuracy. However, these methods either compromise ensembling accuracy or demand significant communication between models during training. In this paper, we introduce WASH, a novel distributed method for training model ensembles for weight averaging that achieves state-of-the-art image classification accuracy. WASH maintains models within the same basin by randomly shuffling a small percentage of weights during training, resulting in diverse models and lower communication costs compared to standard parameter averaging methods.