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

We thank the reviewer for their review.

W1: It is important to note that our goal in this work is not to achieve better Ensemble accuracy, but to achieve better performance for the Averaged model, or at least comparable performance for a much smaller communication volume.

As discussed with reviewer QJmV, we have provided novel ImageNet experiments in this rebuttal, and that to obtain models in the same loss basins with PAPA, we find that it is necessary to perform EMA steps every 1 or 2 training steps. In this case, we note in particular that the Ensemble accuracy is around 74.7%, very close to that obtained in WASH. Note again that our ImageNet experiments on PAPA are reproduced from the initial codebase of PAPA, and we were not able to retrieve their results compared to CIFAR-100.

W3: As mentioned above, we found that the EMA step frequency with PAPA had to be less than 2 to be in the same loss basin. Therefore, the improvement in communication volume is actually between 10 and 40 times that of PAPA. This volume improvement is not negligible and can lead to significant improvements in training speed.

The actual training time improvement will depend on a myriad of factors (computation time, number of workers, communication hardware, and node sizes...), but we refer to [A] for a study showing the impact of communication time on training time for different factors (note that the communication bandwidth in localSGD is directly analogous to DART and PAPA, and thus also to WASH). In these practical settings, the communication time can easily overshadow the computation time. Thus, reducing the communication volume (and thus the communication bandwidth) by a factor of at least $10$ in these cases will result in a significant increase in training speed. Note that we are proposing to shuffle the weights at every step, but only a very small percentage of them. Thus, the communication bandwidth required by WASH is very small, since each worker only needs to send/receive a fraction $p$ of its weights.

We also added a novel experiment in Figure 9, measuring the training computation speed of different models compared to the communication speed for varying numbers of GPUs and communication volumes. We observe that communication, as reported by [A], can slow down training, even in clusters with high-speed connections. However, by limiting the communication volume, the communication time is similarly reduced.

W4: We agree with the reviewer that our ImageNet baseline is a few points shorter than it could be. However, our aim here was to compare ourselves to a reproducible baseline. Thus we started from the official code base of PAPA and ran their code with their reported hyperparameters for the ImageNet experiments as it is our main point of comparison.

W5: Our aim in presenting the two WASH variants is to show several points: that WASH can be adapted in several ways without affecting performance much; and that the effect of the optimizer's weight shuffling is mainly between the heterogeneous and homogeneous settings. Presenting the results of the two variants side by side makes it easier to observe these points. We will try to improve the clarity of the presentation of WASH+Opt in the final version.

W2/6: We agree that more extensive experiments on other datasets and networks could further confirm the effectiveness of our method. Our goal in this work was to present the simplest form of our method and to demonstrate its effectiveness despite its simplicity, as well as to provide a deeper understanding of its mechanisms. Nevertheless, we argue that providing experiments on image classification, in particular on ImageNet, has been sufficient to confirm the results of many other methods in the field, such as in [B-G].

[A] Trade-offs of Local SGD at Scale: An Empirical Study, Ortiz et al., 2021.

[B] Averaging weights leads to wider optima and better generalization, Izmailov et al., 2019.

[C] SWAD : Domain Generalization by Seeking Flat Minima, Cha et al., 2021.

[D] PopulAtion Parameter Averaging (PAPA), Jolicoeur-Martineau et al., 2023.

[E] Learning Neural Network Subspaces, Wortzman et al., 2021.

[F] Loss Surface Simplexes for Mode Connecting Volumes and Fast Ensembling, Benton et al., 2021.

[G] REPAIR: REnormalizing Permuted Activations for Interpolation Repair, Jordan et al., 2023.

After reading the author's response, I have additional questions. I believe there may have been some miscommunications. To clarify, I acknowledge that the paper's objective is not to achieve better ensemble accuracy but to enhance the performance of the averaged model. One of my main concerns is the reliability of the experiments.

I think it would be beneficial to paraphrase my questions as follows:

• Why are all the results (including separate training and PAPA) on ImageNet training lower than the standalone training result of 76%? I believe more than a 1pp decrease in accuracy on ImageNet is significant. I initially thought this degradation stemmed from training without data augmentation, but if I understand correctly from the response, that was a typo and simple data augmentations are used. Is that right?
• If there are some experimental configuration differences, would it be possible to achieve higher accuracy than the standalone training result by using the standalone training configuration or using the modern configurations? Considering that the goal of model ensembling is to improve accuracy and robustness, the inability to use a specific configuration suggests that the performance is limited, which implies that the method's impact is constrained.

We thank the reviewer for their thorough review.

W1: First, note that we have performed additional experiments on PAPA's ImageNet (for $N=3$) by varying the frequency of EMA steps. We find that the models converge to the same loss basin only when EMA steps are performed every 1 or 2 steps. Thus, the effective communication volume improvement of WASH over PAPA is actually between 10x and 40x (depending on the number of steps and whether or not WASH+Opt is used).

In this paper, we report the theoretical communication volume improvement as an objective way to measure the improvement of our method. In comparison, the effective training speed improvement depends on many factors, such as computational speed (model/batch size, hardware...) and communication speed (number of nodes, inter/intra node connections...). For these reasons, we decided to focus on a value agnostic to the training framework.

However, we agree with the reviewer that it remains important to measure this type of improvement in practice. We first refer the reviewer to another paper that measures the impact of communication on training throughput in the case of localSGD (which is functionally equivalent to, for example, DART/PAPA) [A]. In their framework, they find that the time taken by communication is much higher than the time taken by computation, even in cases with only 16 workers. Note also that the recorded communication time increases linearly with the number of workers, and decreases similarly with fewer averaging steps (and thus a smaller communication volume).

We also added a novel experiment in Figure 9, measuring the training computation speed of different models compared to the communication speed for varying number of GPUs and communication volumes. We observe that communication, as reported by [A], can slow down training, even in clusters with high-speed connections. However, by limiting the communication volume, the communication time is similarly reduced, hiding once again the communication time.

W2: As pointed out by the reviewer, we report the Averaged and GreedySoup accuracies for the Baseline models, trained separately. Our aim here is twofold. Firstly, reporting the Averaged accuracy aims to clarify that in all cases considered here, even with $p=3$, homogeneous models starting from the same initialisation, the converged models do not end up in the same loss basin. This is not a novelty, of course, but it serves to make this important point clear. Adding the GreedySoup accuracy gives some additional clarifying information for the rest of the Table: it shows that standard merging approaches provide worse accuracies than methods like PAPA or WASH, which explicitly train models to be in the same loss basin. By retaining only the accuracy of a single model, it is also analogous to reporting the top accuracy of a single model trained separately, an important baseline for comparing the Averaged models trained by PAPA and WASH. We argue that this presentation is not a weakness of our paper, but a way of effectively presenting: that the separately trained models are in different loss basins and therefore unmergeable, with individual model accuracies that are outperformed by the Averaged models.

W3: Indeed, in this work, we have not directly compared ourselves with methods such as SWA or the Lookaround optimizer, as these methods are orthogonal and fully complementary to our approach. These techniques can be used to improve the generalisation of each model during the training phase by finding flatter minima in their optimisation path. Note that SWA in particular was compared to PAPA in Appendix A.11 of their paper. They showed that SWA provides a similar improvement to PAPA, and that combining the two approaches leads to even further performance gains. We expect WASH to behave similarly, as our mechanism works similarly to the EMA steps used in PAPA. Our goal in this work was not to find the most effective combination of techniques (such as REPAIR, which we find improves performance, or other optimizers such as SWA or Lookaround) to achieve state-of-the-art performance, but to introduce a novel distributed training algorithm that greatly decreases communication costs and compare ourselves with similar approaches such as PAPA or DART.

Nevertheless, we agree that we could increase the accuracy of our ImageNet experiments by using a more modern training framework, such as that of [C]. Despite this small accuracy gap with the state of the art, we argue that our experiments still show that WASH performs similarly to PAPA, for a fraction of the communication volume.

The concerns regarding the results for ResNet50 on ImageNet are not about achieving high performance but also about the impact on the reliability of the experimental results presented in the paper. The top-1 accuracy of the baseline reported in the PAPA paper is 76.8%, a reasonable improvement of around 1%p over the expected 76% range with standard augmentation and 90 epochs of training, probably thanks to the use of advanced data augmentation techniques. However, the baseline performance reported in this paper is around 74% with the same advanced data augmentation (and around 73% in a homogeneous setting with standard augmentation), suggesting a significant issue with the experimental setup, which calls into question the validity of any experimental comparisons in this context.

We thank the reviewer for their review.

W1/Q1: Our aim in this submission was to provide a novel method that improves on previous ones for a lower communication volume, rather than precisely measure an improvement in training speed. Such an improvement will depend on many factors, such as the size of the model and the batch, the hardware speed, the communication speed (which also depends on the number of nodes, the type of inter- and intra-node connections, etc.), and more. For these reasons, in this work, we preferred to focus on the performance of the method (with further ablations and experiments) and on the communication volume, a simple value to compare methods, that is agnostic to other variables such as the type of hardware.

Nevertheless, we fully agree with the reviewer that it remains important to measure such improvements in practice. We refer the reviewer first to another paper that measures the impact of communication on training throughput in the case of localSGD (which is equivalent to PAPA, for example) [A]. Note that in their case the communication time is higher than the computation time, even for cases with 16 workers. Note also that the communication time increases linearly with the number of workers, and decreases significantly as fewer averaging steps are needed (and thus a smaller communication volume).

We also added a novel experiment in Figure 9, measuring the training computation speed of different models compared to the communication speed for varying number of GPUs and communication volumes. We observe that communication, as reported by [A], can slow down training, even in clusters with high-speed connections. However, by limiting the communication volume, the communication time is similarly reduced, hiding once again the communication time.

Finally, as discussed with reviewer QJmV, note that we added further experiments on ImageNet for PAPA and found comparable accuracies to WASH only with EMA steps every 1 or 2 steps, resulting in an even higher communication volume required.

Q2: Other schedules have been considered. We tried using different fixed values in different parts of the networks, for example dividing it into quarters. Instead of decreasing to $0$, we also tried decreasing to higher values, with worse results and higher communication volumes. In general, these experiments all converged on the same trend: more permutations early in the networks and fewer for the later layers. Our chosen schedule was thus chosen to be simple and applicable to a wide variety of networks, without the need for other hyperparameters.

The idea suggested by the reviewer is quite interesting and seems to be a way to effectively choose more precise schedules for different architectures, as well as noting which components require more shuffling. It would also be consistent with our intuition. Note, however, that knowledge of parameter variance is not applicable during training, as it would require communication between workers to obtain this information. We emphasize that contrary to PAPA, our method WASH never performs an averaging operation requiring an all-reduce communication and only permutes parameters; thus accessing the variance is not possible without adding a significant communication overhead. It remains a very interesting idea to explore and we thank the reviewer.

Q3: First, note that larger models do not directly require a higher $p$ value: using ResNets-50 with the same $p$ as smaller models also leads to the same loss basin when training on CIFAR-100 (see Table 6).

The need for a higher $p$ value for ImageNet compared to CIFAR-100 seems analogous to the need for more frequent EMA steps in PAPA, as we have found (see our changes and discussion with reviewer QJmV). In both cases, more actions are needed to keep the models in the same loss basin. There are several possible reasons for this: sharper local minima, more local minima with larger data sets, etc. Nevertheless, it suggests that more complex datasets require more work to keep the models in the same loss basin.

However, we do not believe that it will be necessary to reach values of $p$ close to $1$. The most communication required for a model trained with approaches such as PAPA or DART is to require an EMA step at each step, which is already the case for ImageNet here. Even in this case, we achieve the same accuracy for a small value of $p$. In a case where $p$ becomes too high, a simple workaround could be to combine weight shuffling steps with periodic and infrequent EMA or averaging steps. This approach would allow us to maintain a small communication volume while keeping the models in the same loss basin.

Q4: The distance metric reported here is a simple L2 distance between parameters. The consensus is computed and we compute the average distance between the weights of each model and the weights of the consensus.

[A] Trade-offs of Local SGD at Scale: An Empirical Study, Ortiz et al., 2021

We thank the reviewer for their review.

W1: As the reviewer points out, the results reported in the DART paper seem to outperform the EMA method of PAPA in some cases. However, as the PAPA authors point out in Section 4.5.2, DART does not provide an implementation and the hyperparameters used are unclear. Furthermore, they find that in their own re-implementation, the accuracy of DART is lower than reported. Notably, it performs worse than PAPA-all (which is fundamentally the same method). Since the authors of PAPA report that their PAPA-all method generally performs worse than PAPA, we have chosen to report only PAPA over DART and PAPA-all in order to keep only the best (reproducible) performing method.

W2: Indeed, the improvement gap appears mainly in the most difficult cases, on larger datasets and with a higher number of models. These are precisely the cases we are interested in, and in particular where the communication volume can become a bottleneck for faster training. Note also that the remaining gap with the Ensemble may be due to the fact that all models are in the same loss basin, which directly limits their diversity (see line 353).

W3: We agree with the reviewer and have extended our experiments and discussion regarding the reimplementation in Appendix C.4, reporting results on ImageNet for different EMA frequencies in Table 6. We found that training on ImageNet required an EMA step every 1 or 2 steps, rather than every 10 steps as reported. We found accuracies very close to those of WASH, for a communication volume 40 times higher.

Q1: This is an interesting point, and we thank the reviewer for discussing it. First, note that the randomness is predetermined by a random seed that is known to all workers in advance, so no communications are required to know which weights are being shuffled and between whom. Then note also that the only computations that need to be done for the shuffling step (other than the communications) are: randomly choosing which parameters to permute (which can be done efficiently with a random Bernoulli mask, for example), and choosing permutations for each one. This can be a bit slower, but remember that only a very small proportion of the parameters are shuffled at each step. To be absolutely sure that no computational overhead is caused by this entire step, it is even possible to do it in parallel with the training computations, as it requires a very small amount of memory (only about one layer in memory at a time for efficient computations, and each worker only needs to store the few shuffled parameter indexes and their associated previous/next workers in the permutation).