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

TL;DR

WASH is a novel distributed training method that improves ensemble performance and efficiency by shuffling weights during training to keep models in the same loss basin, reducing communication costs.

Contribution

WASH introduces a new weight shuffling technique during training to enhance ensemble accuracy and efficiency with lower communication overhead.

Findings

Achieves state-of-the-art image classification accuracy.

Maintains models in the same loss basin through weight shuffling.

Reduces communication costs compared to existing methods.

Abstract

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.