DART: Diversify-Aggregate-Repeat Training Improves Generalization of
  Neural Networks
Generalization of neural networks is crucial for deploying them safely in the
real world. Common training strategies to improve generalization involve the
use of data augmentations, ensembling and model averaging. In this work, we
first establish a surprisingly simple but strong benchmark for generalization
which utilizes diverse augmentations within a training minibatch, and show that
this can learn a more balanced distribution of features. Further, we propose
Diversify-Aggregate-Repeat Training (DART) strategy that first trains diverse
models using different augmentations (or domains) to explore the loss basin,
and further Aggregates their weights to combine their expertise and obtain
improved generalization. We find that Repeating the step of Aggregation
throughout training improves the overall optimization trajectory and also
ensures that the individual models have a sufficiently low loss barrier to
obtain improved generalization on combining them. We shed light on our approach
by casting it in the framework proposed by Shen et al. and theoretically show
that it indeed generalizes better. In addition to improvements in In- Domain
generalization, we demonstrate SOTA performance on the Domain Generalization
benchmarks in the popular DomainBed framework as well. Our method is generic
and can easily be integrated with several base training algorithms to achieve
performance gains.
Samyak Jain, Sravanti Addepalli, Pawan Sahu, Priyam Dey, R. Venkatesh Babu
# DART: Diversify-Aggregate-Repeat Training

###### Abstract

Generalization of Neural Networks is crucial for deploying them safely in the real world. Common training strategies to improve generalization involve the use of data augmentations, ensembling and model averaging. In this work, we first establish a surprisingly simple but strong benchmark for generalization which utilizes diverse augmentations within a training minibatch, and show that this can learn a more balanced distribution of features. Further, we propose Diversify-Aggregate-Repeat Training (DART) strategy that first trains diverse models using different augmentations (or domains) to explore the loss basin, and further Aggregates their weights to combine their expertise and obtain improved generalization. We find that Repeating the step of Aggregation throughout training improves the overall optimization trajectory and also ensures that the individual models have a sufficiently low loss barrier to obtain improved generalization on combining them. We shed light on our approach by casting it in the framework proposed by Shen et al. [61] and theoretically show that it indeed generalizes better. In addition to improvements in In-Domain generalization, we demonstrate SOTA performance on the Domain Generalization benchmarks in the popular DomainBed framework as well. Our method is generic and can easily be integrated with several base training algorithms to achieve performance gains.

## 1 Introduction

Deep Neural Networks have outperformed classical methods in several fields and applications owing to their remarkable generalization. Classical Machine Learning theory assumes that test data is sampled from the same distribution as train data. This is referred to as the problem of In-Domain (ID) generalization [16, 19, 30, 34, 52], where the goal of the model is to generalize to samples within same domain as the train dataset. This is often considered to be one of the most important requirements and criteria to evaluate models. However, in several cases, the test distribution may be different from the train distribution. For example, surveillance systems are expected to work well at all times of the day, under different lighting conditions and when there are occlusions. However, it may not be possible to train models using data from all these distributions. It is therefore crucial to train models that are robust to distribution shifts, which is popularly referred to as Out-of-Domain (OOD) Generalization [26]. In this work, we consider the problems of In-Domain generalization and Out-of-Domain Generalization of Deep Networks. For the latter case, we consider the popular setting of Domain Generalization [6, 45, 24], where the training data is assumed to be composed of several source domains and the goal is to generalize to an unseen target domain.

The problem of generalization is closely related to the Simplicity Bias of Neural Networks, due to which models have a tendency to rely on simpler features that are often spurious correlations to the labels, when compared to the harder robust features [59]. For example, models tend to rely on weak features such as background, rather than more robust features such as shape, causing a drop in object classification accuracy when background changes [76, 23].

A common strategy to alleviate this is to use data augmentations [10, 11, 12, 28, 46, 56, 79, 81] or data from several domains during training [24], which can result in invariance to several spurious correlations, improving the generalization of models. Shen et al. [61] show that data augmentations enable the model to give higher importance to harder-to-learn robust features by delaying the learning of spurious features. We extend their observation by showing that training on a combination of several augmentation strategies (which we refer to as _Mixed_ augmentation) can result in the learning of a balanced distribution of diverse features. Using this, we obtain a strong benchmark for ID generalization as shown in Table-1. However, as shown in prior works [1], the impact of augmentations in training is limited by the capacity of the network in being able to
generalize well to the diverse augmented data distribution. Therefore, increasing the diversity of training data demands the use of larger model capacities to achieve optimal performance. This demand for higher model capacity can be mitigated by training specialists on each kind of augmentation and ensembling their outputs [13, 40, 83, 63], which results in improved performance as shown in Table-1. Another generic strategy that is known to improve generalization is model-weight averaging [74, 75, 33]. This results in a flatter minima, thereby improving the robustness to distribution shifts.

In this work, we aim to combine the benefits of the three strategies discussed above - diversification, specialization and model weight averaging, while also overcoming their individual shortcomings. We propose a **D**iversify-**A**gregate-**R**epeat **T**raining strategy dubbed DART (Fig.1), that first trains \(k\)_Diverse_ models after a few epochs of common training, and then _Aggregates_ their weights to obtain a single generalized solution. The aggregated model is then used to reinitialize the \(k\) models which are trained further post aggregation. This process is _Repeated_ over training to obtain improved generalization. The _Diversify_ step allows models to explore the loss basin and specialize on a fixed set of features. The _Aggregate_ (or Model Interpolation) step robustly combines these models, increasing the diversity of represented features while also suppressing spurious correlations. Repeating the _Diversify-Aggregate_ steps over training results in a more robust optimization trajectory and also ensures that the \(k\) diverse models remain in the same basin thereby permitting a fruitful combination of their weights. We justify our approach theoretically and empirically, and also show that the intermediate model aggregation not only ensures that the models are in the same basin, but also increases the learning time for spurious features, improving generalization. We present our key contributions below.

* We present a strong baseline termed Mixed-Training (MT) that uses a combination of diverse augmentations for different images in a training minibatch.
* We propose a novel algorithm DART, that learns specialized diverse models and aggregates their weights iteratively to improve generalization.
* We justify our method theoretically, and empirically on several In-Domain (CIFAR-10, CIFAR-100) and Domain Generalization (OfficeHome, PACS, VLCS, TerraIncognita, DomainNet) datasets.

## 2 Background: Mode Connectivity of Models

The overparameterization of Deep networks leads to the existence of multiple optimal solutions to any given loss function [49, 80, 35]. Prior works [50, 22, 15] have shown that all such solutions learned by SGD lie on a non-linear manifold, and are connected to each other by a path of low loss. Frankle _et al_. [20] further showed that converged models that share a common initial optimization path are linearly connected with a low loss barrier. This is referred to as the _linear mode connectivity_ between the models. Several optimal solutions that are linearly connected to each other are said to belong to a common _basin_ which is separated from other regions of the loss landscape with a higher loss barrier. Loss barrier between any two models \(\theta_{1}\) and \(\theta_{2}\) is defined as the maximum loss that can be attained by the models, \(\hat{\theta}=\alpha\cdot\theta_{1}+(1-\alpha)\cdot\theta_{2}\ \ \ \forall\ \ \alpha\in[0,1]\).

The linear mode connectivity of models facilitates the averaging of weights of different models in a common basin resulting in further gains. In this work, we leverage the linear mode connectivity of diverse models trained from a common initialization to improve generalization.

## 3 Related Works

### Generalization of Deep Networks

Prior works aim to improve the generalization of Deep Networks by imposing invariances to several factors of variation. This is achieved by using data augmentations during training to improve the In-Domain generalization [10, 11, 12, 27, 46, 69, 79, 81], or by training on a combination of multiple domains for Domain Generalization [42, 45, 42, 29, 9].

\begin{table}
\begin{tabular}{l c c c c} \hline \hline  & \multicolumn{4}{c}{Test} \\ \cline{2-5} Train & No Aug. & Cutout & Cutmix & AutoAugment \\ \hline Pad+Crop+HFlip (PC) & 78.51 & 67.04 & 56.52 & 58.33 \\ Cutout (CO) & 77.99 & 74.58 & 56.12 & 58.47 \\ Cutmix (CM) & 80.54 & 74.05 & **77.35** & 61.23 \\ AutoAugment (AA) & 79.18 & 71.26 & 60.97 & 73.91 \\ Mixed-Training (MT) & 81.43 & 77.31 & 73.20 & **74.73** \\ \hline Ensemble (CM+CO+AA) & **83.61** & **79.19** & 73.19 & 73.90 \\ \hline \hline \end{tabular}
\end{table}
Table 1: **Motivation:** Performance (%) on CIFAR100, ResNet-18 with ERM training for 200 epochs across different train and test augmentations. Mixed-Training (MT) outperforms individual augmentations, while ensembles perform best on average.

Figure 1: Schematic Diagram of the proposed method DART
Contrary to these methods, several works focused on utilizing domain-specific features [5, 14], while others try to disentangle the features into domain-specific and domain-invariant for better generalization [36, 37, 42, 53, 72]. Data augmentation has also been exploited for Domain Generalization [47, 54, 60, 62, 70, 71, 73, 77, 78, 85, 86] in order to increase the diversity of training data and simulate domain shift. Foret _et al_. [19] show that minimizing the maximum loss within an \(\ell_{2}\) norm ball of weights can result in a flatter minima thereby improving generalization. Gulrajani _et al_. [24] show that the simple strategy of ERM training on data from several source domains can indeed prove to be a very strong baseline for Domain Generalization. The authors also release DomainBed - which benchmarks several existing methods on some common datasets representing different types of distribution shifts. We empirically show that our method achieves SOTA on the popular Domain Generalization benchmarks and helps in further improving performance and generalization when used in conjunction with several other methods (Table-4) ascribing to its orthogonal nature.

### Averaging model weights across training

Recent works have shown that converging to a flatter minima can lead to improved generalization [16, 19, 30, 34, 52, 64]. Izmailov _et al_. [33] proposed Stochastic Weight Averaging (SWA) to average the model weights from the last few epochs such that the resulting model converges to a flatter minima, thus improving generalization. A variation of SWA, Exponential Moving Average (EMA) of model weights across training iterations is often used to boost model performance at no extra training cost. Cha _et al_. [6] theoretically show that converging to a flatter minima results in a smaller domain generalization gap. The authors propose SWAD that overcomes the limitations of SWA in the Domain Generalization setting and combines several models in the optimal solution basin to obtain a flatter minima with better generalization. We demonstrate that our approach effectively integrates with EMA and SWAD for In-Domain and Domain Generalization settings respectively to obtain further performance gains (Tables-2, 3).

### Averaging weights of fine-tuned models

While earlier works combined models generated from the same optimization trajectory, Tatro _et al_. [66] showed that for any two converged models with different random initializations, one can find a permutation of one of the models so that fine-tuning the interpolation of the second model and the permuted first model leads to improved generalized. On a similar note, Zhao _et al_. [84] proposed to achieve robustness to backdoor attacks by fine-tuning the linear interpolation of pretrained models. More recently, Wortsman _et al_. [75] proposed Model Soups and showed that in a transfer learning setup, fine-tuning and then averaging different models with same pre-trained initialization but with different hyperparameters such as learning rates, optimizers and augmentations can improve the generalization of the resulting model. The authors further note that this works best when the pre-trained model is trained on a large heterogeneous dataset. While all these approaches work only in a finetuning setting, the proposed method incorporates the interpolation of differently trained models in the regime of _training from scratch_, allowing the learning of models for longer schedules and larger learning rates.

### Averaging weights of differently trained models

Wortsman _et al_. [74] propose to average the weights of multiple models trained with different random initializations simultaneously by considering the loss of a combined model for optimization while performing gradient updates on the individual models. Additionally, they minimize the cosine similarity between model weights to ensure that the models learned are diverse. While this training formulation does learn diverse connected models, it leads to individual models having sub-optimal accuracy (Table-2) since their loss is not optimized directly. DART overcomes such issues since the individual models are trained directly to optimize their respective classification losses. Moreover, the step of intermediate interpolation ensures that the individual models also have better performance when compared to the baseline of standard ERM traning on the respective augmentations (Fig.8).

## 4 Proposed Method: DART

A series of observations from prior works [15, 20, 22, 50] have led to the conjecture that models trained independently with different initializations could be linearly connected

Figure 2: **Optimization trajectory** of the proposed approach DART when compared to independent ERM training on each augmentation. Axes represent the top two PCA directions obtained using the weights of DART training. The initial common point on the right represents the model obtained after 100 epochs of Mixed Training (MT). The trajectory shown is for an additional 100 epochs, with a total training budget of 200 epochs. The optimization trajectory of DART is better aligned with the negative gradient direction when compared to the baseline.
with a low loss barrier, when different permutations of their weights are considered, suggesting that _all solutions effectively lie in a common basin_[17]. Motivated by these observations (discussed in Section-2) and the above hypothesis, we aim at designing an algorithm that explores the basin of solutions effectively with a robust optimization path and combines the expertise of several diverse models to obtain a single generalized solution.

We show an outline of the proposed approach - _Diversify-Aggregate-Repeat Training_, dubbed DART, in Fig.1. Broadly, the proposed approach is implemented in four steps - i) ERM training for \(E^{\prime}\) epochs in the beginning, followed by ii) Training \(M\)_Diverse_ models for \(\lambda/M\) epochs each, iii) _Aggregating_ their weights, and finally iv) _Repeating_ the steps _Diversify-Aggregate_ for \(E-E^{\prime}\) epochs.

A cosine learning rate schedule is used for training the model for a total of \(E\) epochs with a maximum learning rate of \(\mathrm{LR}_{max}\). We present the implementation of DART in Algorithm-1, and discuss each step in detail below:

1. **Traversing to the _Basin_ of optimal solutions:** Since the goal of the proposed approach is to explore the _basin_ of optimal solutions, the first step is to traverse from a randomly initialized model upto the periphery of this basin. Towards this, the proposed _Mixed-Training_ strategy discussed in Section-1 is performed on a combination of several augmentations \(D^{*}\) for the initial \(E^{\prime}\) epochs (L4-L5 in Alg.1).
2. **Diversify - Exploring the _Basin_:** In this step, \(M\) diverse models \(f_{\theta^{k}}\) initialized from the Mixed-Training model (L8 in Alg.1), are trained using the respective datasets \(D^{k}\) (L10 in Alg.1). These are generated using diverse augmentations in the In-Domain setting, and from a combination of different domains in the Domain Generalization setting. To maintain the same compute as baselines, we set \(|D^{k}|=|D|/M\).
3. **Aggregate - Combining diverse experts:** Owing to the initial common training for \(E^{\prime}\) epochs, the \(k\) diverse models lie in the same basin, enabling an effective aggregation of their weights using simple averaging (L12 in Alg.1) to obtain a more generalized solution \(\theta\). Aggregation is done after every \(\lambda\) epochs.
4. **Repeat:** Next, all \(k\) models are reinitialized using the common model \(\theta\) (L13 of Alg.1), after which the individual models are trained for \(\lambda\) epochs on their respective datasets \(D^{k}\) as discussed in Step-2, and the process continues for a total of \(E-E^{\prime}\) epochs.

**Visualizing the Optimization Trajectory:** We compare the optimization trajectory of the proposed approach DART with independent training on the same augmentations in Fig.2 after a common training of \(E^{\prime}=100\) epochs on Mixed augmentations. We note that the models explore more in the initial phase of training, and lesser thereafter, which is a result of the cosine learning rate schedule and gradient magnitudes. The exploration in the initial phase helps in increasing the diversity of models, thereby improving the robustness to spurious features (as shown in Proposition-3) leading to a better optimization trajectory, while the smaller steps towards the end help in retaining the flatter optima obtained after Aggregation. The process of repeated aggregation also ensures that the models remain close to each other, allowing longer training regimes.

## 5 Theoretical Results

We use the theoretical setup from Shen _et al_. [61] to show that the proposed approach DART achieves robustness to spurious features, thereby improving generalization.

**Preliminaries and Setup:** We consider a binary classification problem with two classes \(\{-1,1\}\). We assume that the dataset contains \(n\) inputs and \(K\) orthonormal robust features which are important for classification and are represented as \(v_{1},v_{2},v_{3},\ldots,v_{K}\), in decreasing order of their frequency in the dataset. Let each input example \(x\) be composed of two patches denoted as \((x_{1},x_{2})\in R^{d\times 2},\) where each patch is characterized as follows: i) **Feature patch**: \(x_{1}=yv_{k}\), where \(y\) is the target label of \(x\) and \(k^{*}\in[1,K]\), ii) **Noisy patch**: \(x_{2}=\epsilon\) where \(\epsilon\sim\mathcal{N}\left(0,\frac{\sigma^{2}}{d}I_{d}\right)\).
We consider a single layer convolutional neural network consisting of C channels, with \(w=(w_{1},w_{2},w_{3},\ldots,w_{C})\in R^{d\times C}\). The function learned by the neural network (F) is given by \(F(w,x)=\sum\limits_{c=1}^{C}\sum\limits_{p=1}^{2}\phi(w_{c},x_{p})\), where \(\phi\) is the activation function as defined by Shen _et al_. [61].

**Weights learned by an ERM trained model:** Let \(K_{cut}\) denote the number of robust features learned by the model. Following Shen _et al_. [61], we assume the learned weights to be a linear combination of the two types of features present in the dataset as shown below:

\[w=\sum\limits_{k=1}^{K_{cut}}v_{k}+\sum\limits_{k>K_{cut}}y^{(k)}\epsilon^{(k)} \tag{1}\]

**Data Augmentations:** As defined by Shen _et al_. [61], an augmentation \(T_{k}\) can be defined as follows:

\[\forall\ k^{\prime}\ \in\ [1,K],\ \ \ \mathcal{T}_{k}(v_{k^{\prime}})\ =\ v_{((k^{\prime}+k-1)\ mod\ K)+1} \tag{2}\]

Assuming that \(K\) unique augmentation strategies are used (where \(K\) denotes the number of robust patches in the dataset), augmented data is defined as follows:

\[D_{train}^{(aug)}=D_{train}\ \cup\ \mathcal{T}_{1}(D_{train})..\cup\ \mathcal{T}_{K-1}(D_{train}) \tag{3}\]

where \(D_{train}\) is the training dataset. This ensures that each feature patch \(v_{i}\) appears \(n\) times in the dataset, thus making the distribution of all the feature patches uniform.

**Weight Averaging in DART:** In the proposed method, we consider that \(m\) models are being independently trained after which their weights are averaged as shown below:

\[w=\frac{1}{m}\sum\limits_{j=1}^{m}\sum\limits_{k=1}^{K_{cut_{j}}}v_{k_{j}}+ \frac{1}{m}\sum\limits_{j=1}^{m}\sum\limits_{k>K_{cut_{j}}}y_{j}^{(k)}\epsilon _{j}^{(k)} \tag{4}\]

Each branch is trained on the dataset \(D_{train}^{(k)}\) defined as:

\[D_{train}^{(k)}=\mathcal{T}_{k}(D_{train}),\ \ k\in[1,m] \tag{5}\]

**Propositions:** In the following propositions, we derive the convergence time for learning robust and noisy features, and compare with the bounds derived by Shen _et al_. [61] in Section-6. The proofs of all propositions are presented in the Section-9.

**Notation:** Let \(f_{\theta}\) denote a neural network obtained by averaging the weights of \(m\) individual models \(f_{\theta}^{k}\), \(k\in[1,m]\) which are represented as shown in Eq.1. \(n\) is the total number of data samples present in the original dataset \(D_{train}\). \(K\) is the number of orthonormal robust features present in the dataset. The weights \(w_{1},w_{2},\ldots,w_{C}\) of each model \(f_{\theta}^{k}\) are initialized as \(w_{c}\sim\mathcal{N}\left(0,\sigma_{0}^{2}I_{d}\right)\ \forall\ c\in[1,C]\), where C is the number of channels present in a single layer of the model. \(\frac{\sigma}{\sqrt{d}}\) is the standard deviation of the noise present in noisy patches, \(q\) is a hyperparameter used to define the activation (Details in the Section-9), where \(q\geq 3\) and \(d\) is the dimension of each feature patch and weight channel \(w_{c}\).

**Proposition 1**.: _The convergence time for learning any feature patch \(v_{i}\ \forall i\in[1,K]\) in at least one channel \(c\in C\) of the weight averaged model \(f_{\theta}\) using the augmentations defined in Eq.5, is given by \(O\left(\frac{K}{\sigma_{0}^{q-2}}\right)\), if \(\frac{\sigma^{q}}{\sqrt{d}}\ll\frac{1}{K}\), \(m=K\)._

**Proposition 2**.: _If the noise patches learned by each \(f_{\theta}^{k}\) are \(i.i.d\). Gaussian random variables \(\sim\mathcal{N}(0,\frac{\sigma^{2}}{d}I_{d})\) then with high probability, convergence time of learning a noisy patch \(\epsilon^{(j)}\) in at least one channels \(c\in[1,C]\) of the weight averaged model \(f_{\theta}\) is given by \(O\left(\frac{nm}{\sigma_{0}^{q-2}\sigma^{q}}\right)\), if \(d\gg n^{2}\)._

**Proposition 3**.: _If the noise learned by each \(f_{\theta}^{k}\) are \(i.i.d\). Gaussian random variables \(\sim\mathcal{N}\left(0,\frac{\sigma^{2}}{d}I_{d}\right)\), and model weight averaging is performed at epoch \(T\), the convergence time of learning a noisy patch \(\epsilon^{(j)}\) in at least one channels \(c\in[1,C]\) of the weight averaged model \(f_{\theta}\) is given by \(T+O\left(\frac{nm^{(q-2)}d^{(q-2)/2}}{\sigma^{(2q-2)}}\right)\) if \(d\gg n^{2}\)._

## 6 Analysis on the Theoretical Results

In this section, we present the implications of the theoretical results discussed above. While the setup in Section-5 discussed the existence of only two kinds of patches (feature and noisy), in practice, a combination of these two kinds of patches - termed as Spurious features - could also exist, whose convergence can be derived from the above results.

### Learning Diverse Robust Features

We first show that _using sufficiently diverse data augmentations during training generates a uniform distribution of feature patches, encouraging the learning of diverse and robust features by the network_. We consider the use of \(K\) unique augmentations (where \(K\) is the number of robust patches in the dataset) in Eq.3 which transform each feature patch into a different one using a unique mapping as shown in Eq.2. The mapping in Eq.2 can transform a skewed feature distribution to a more uniform distribution after performing augmentations. This results in \(K_{cut}\) being sufficiently large in Eq.1, which depends on the number of high frequency robust features, thereby encouraging the learning of a more balanced distribution of robust features.

Shen _et al_. showed that the time to learn any feature patch \(v_{k}\) by at least one weight channel \(c\in C\) is given by \(O\left(\frac{1}{\sigma_{0}^{q-2}\rho_{k}}\right)\) if \(\frac{\sigma^{q}}{\sqrt{d}}\ll\rho_{k}\), where \(\rho_{k}\) is the fraction of the frequency of occurrence of feature patch \(v_{k}\) divided by the total number of occurrences of all the feature patches in the dataset. The convergence time for learning feature patches
is thus limited by the one that is least frequent in the input data. Therefore, by making the frequency of occurrence of all feature patches uniform, the convergence time reduces. In Proposition-1, we show that the same holds true even for the proposed method DART where several branches are trained using diverse augmentations and their weights are finally averaged to obtain the final model. This justifies the improvements obtained in Mixed-Training (MT) (Eq.1) and in the proposed approach DART (Eq.4) as shown in Table-2.

### Robustness to Noisy Features

Firstly, the use of diverse augmentations in both Mixed-Training (MT) and DART results in better robustness to noisy features since the value of \(K_{cut}\) in Eq.1 and Eq.4 would be higher, resulting in the learning of more feature patches and suppressing the learning of noisy patches. _The proposed method DART indeed suppresses the learning of noisy patches further, and also increases the convergence time for learning noisy features as shown in Proposition-2_. When the augmentations used in each of the \(m\) individual branches of DART are diverse, the noise learned by each of them can be assumed to be \(i.i.d\). Under this assumption, averaging model weights at the end of training results in a reduction of noise variance, as shown in Eq.4. More formally, we show in Proposition-2 that the _convergence time of noisy patches increases by a factor of \(m\) when compared to ERM training_. We note that this does not hold in the case of averaging model weights obtained during a single optimization trajectory as in SWA [33], EMA or SWAD [6], since the noise learned by models that are close to each other in the optimization trajectory cannot be assumed to be \(i.i.d\).

### Impact of Intermediate Interpolations

We next analyse the impact of averaging the weights of the models at an intermediate epoch-T in addition to the interpolation at the end of training. The individual models are further reinitialized using the weights of the interpolated model as discussed in Algorithm-1. As shown in Proposition-3, averaging the weights of all branches at the intermediate epoch T helps in increasing the convergence time of noisy patches by a factor \(O\left(\frac{\sigma_{0}^{q-2}m^{q-3}d^{(q-2)/2}}{\sigma^{q-2}}\right)\) when compared to the case where models are interpolated only at the end of training as shown in Proposition-2. By assuming that \(q>3\) and \(d\gg n^{2}\) similar to Shen _et al_. [61], the lower bound on this can be written as \(O\left(\frac{\sigma_{0}n}{\sigma}\right)\). We note that in a practical scenario this factor would be greater than 1, demonstrating the increase in convergence time for noisy patches when intermediate interpolation is done.

## 7 Experiments and Results

In this section, we empirically demonstrate the performance gains obtained using the proposed approach DART on In-Domain (ID) and Domain Generalization (DG) datasets. We further attempt to understand the various factors that contribute to the success of DART.

**Dataset Details:** To demonstrate In-Domain generalization, we present results on CIFAR-10 and CIFAR-100 [38], while for Domain Generalization, we present results on the 5 real-world datasets on the DomainBed [24] benchmark - VLCS [18], PACS [42], OfficeHome [68], Terra Incognita [3] and DomainNet [51], which represent several types of domain shifts with different levels of dataset and task complexity. DomainNet is a large scale dataset with 6 domains, 345 categories and roughly 0.6 million images. More details on datasets are presented in Section-12.1.

**Training Details (ID):** We set the number of training epochs to 600 for the In-Domain experiments on CIFAR-10 and CIFAR-100. To enable a fair comparison, we select the best performing configuration amongst 200, 400 and 600 total training epochs for the ERM baselines and Mixed-Training, since they may be prone to overfitting. We use a cosine learning rate schedule with a maximum learning rate of 0.1 and weight decay of 5e-4. SGD with momentum of 0.9 is used for optimization of the cross-entropy (CE) loss. Interpolation frequency (\(\lambda\)) is set to be 50 epochs for CIFAR-100 and 40 epochs for CIFAR-10. We present results on ResNet-18 and WideResNet-28-10 architectures.

**Training Details (DG):** Following the setting in DomainBed [24], we use Adam [37] optimizer with a fixed learning rate of 5e-5. The number of training iterations are set to 15k for DomainNet (due to its higher complexity) and 10k for all other datasets with the interpolation frequency being set to 1k iterations. ResNet-50 [25] was used as the backbone, initialized with Imagenet [57] pre-trained weights. Best-model selection across training checkpoints was done based on validation results from the train domains itself, and no subset of the test domain was used. We present further details in Section-12.1.

**SOTA comparison - In Domain (ID) Generalization:** In Table-2, we compare our method against ERM training with several augmentations, and also the strong Mixed-Training benchmark (MT) obtained by using either AutoAugment [10], Cutout [12] or Cutmix [79] for every image in the training minibatch uniformly at random. We use

\begin{table}
\begin{tabular}{l l l} \hline \hline Method & CIFAR-10 & CIFAR-100 \\ \hline ERM+EMA (Pad+Crop+HFlip) & 96.41 & 81.67 \\ ERM+EMA (AutoAugment) & 97.50 & 84.20 \\ ERM+EMA (Cutout) & 97.43 & 82.33 \\ ERM+EMA (Cutmix) & 97.11 & 84.05 \\ Learning Subspaces [74] & 97.46 & 83.91 \\ \hline ERM+EMA (Mixed Training-MT) & 97.69 \(\pm\) 0.19 & 85.57 \(\pm\) 0.13 \\ DART (Ours) & **97.96**\(\pm\) 0.06 & **86.46**\(\pm\) 0.12 \\ \hline \hline \end{tabular}
\end{table}
Table 2: **In-Domain Generalization:** Performance (%) of DART when compared to baselines on WideResNet-28-10 model. Standard deviation for DART and MT are reported across 5 reruns.
the same augmentations in DART as well, with each of the 3 branches being trained on one of the augmentations. As discussed in Section-3, the method proposed by Wortsman _et al_. [74] is closest to our approach, and hence we compare with it as well. We utilize Exponential Moving Averaging (EMA) of weights for the ERM baselines and the proposed approach for a fair comparison. On CIFAR-10, we observe gains of 0.19% on using ERM-EMA (Mixed) and an additional 0.27% on using DART. On CIFAR-100, 1.37% improvement is observed with ERM-EMA (Mixed) and an additional 0.89% with the proposed method DART. The comparison of DART with ERM on ImageNet-1K and fine-grained datasets like Stanford-Cars and CUB200 is shown in Table-6. On ImageNet-1K we observe gains of 0.41% on using RandAugment across all the branches of the model and of 0.14% on using Pad-Crop, RandAugment and Cutout for different branches. We observe gains upto 1.5% on fine-grained datasets.

**SOTA comparison - Domain Generalization:** We present results on the DomainBed [24] datasets in Table-3. We compare only with ERM training (performed on data from a mix of all domains) and SWAD [6] in the main paper due to lack of space, and present a thorough comparison across all other baselines in Table-12.3.1 in Appendix. For the DG experiments, we consider 4 branches (\(M=4\)), with 3 branches being specialists on a given domain and the fourth being trained on a combination of all domains in equal proportion. For the DomainNet dataset, we consider 6 branches due to the presence of more domains. On average, we obtain 2.8% improvements over the ERM baseline without integrating with SWAD, and 1% higher accuracy when compared to SWAD by integrating our approach with it. We further note from Table-4 that the DART can be integrated with several base approaches - with and without SWAD, while obtaining substantial gains across the respective baselines. The best performance achieved is by integrating with Mixup, which yields 3.04% better accuracy over ERM without SWAD, and 1.23% improvement over SWAD by integrating with it. The proposed approach therefore is generic, and can integrated effectively with several algorithms. As observed in Table-7, we observe substantial gains (+2.6%) over MIRO [7] on incorporating DART with SWAD and MIRO and using CLIP [55] as the backbone.

**Evaluation without imposing diversity across branches:** While the proposed method imposes diversity across branches by using different augmentations, we show in Table-5 that the method works even without explicitly introducing diversity, by virtue of the randomness introduced by SGD and different ordering of input samples across models. We obtain an average improvement of 0.9% over the respective baselines, and maximum improvement of 1.82% using Cutout. This shows that the performance of DART is not dependent on data augmentations, although it achieves further improvements on using them.

**Accuracy across training epochs:** We show the accuracy across training epochs for the individual branches and the combined model in Fig.4 for two cases - (a) performing

\begin{table}
\begin{tabular}{l c c c c c} \hline \hline \multicolumn{1}{c}{**Searches (\%)**} & **\multicolumn{1}{c}{**\(\%\)**} & **\multicolumn{1}{c}{**\(\%\)**} & **\multicolumn{1}{c}{**\(\%\)**} & **\multicolumn{1}{c}{**\(\%\)**} & **\multicolumn{1}{c}{**\(\%\)**} \\ \hline
**Augmentation** & **SBM** & **Mixed** & **1.204** & **DART** & **EBM** & **DART** \\ \hline
**Single Augmentation** & 81.11 & **99.42** & 35.88 & **99.75** & 78.55 & **78.96** \\
**Model Augmentation** & 00.88 & **91.85** & 81.72 & **82.83** & 79.06 & **79.20** \\ \hline \hline \end{tabular}
\end{table}
Table 6: **DART on ImageNet-1K and finegrained datasets: Performance (%) of the DART when compared to ERM on ResNet-50, and Pad-Crop for finegrained datasets (_Single Augmentation_) and when Pad-Crop, RandAugment [11] and Cutout [12] are used for different branches (_Mixed Augmentation_), while instead of RandAugment, AutoAugment is used for finegrained datasets.**

\begin{table}
\begin{tabular}{l c c c c c} \hline \hline
**Algorithm** & **Vanilla** & **DART** (w/o SWAD) & **SWAD** & **DART** (\(\prec\) SWAD) \\ \hline ERM [67] & 66.5 & 70.31 & 70.60 & **72.28** \\ ARM [82] & 64.8 & 69.24 & 69.75 & **71.31** \\ SAM\({}^{\dagger}\)[19] & 67.4 & 70.39 & 70.26 & **71.55** \\ Cutmix\({}^{\dagger}\)[79] & 67.3 & 70.07 & 71.08 & **71.49** \\ Mixup [73] & 68.1 & 71.14 & 71.15 & **72.38** \\ DANN [21] & 65.9 & 70.32 & 69.46 & **70.85** \\ CDANN [44] & 65.8 & 70.75 & 69.70 & **71.69** \\ SagNet [48] & 68.1 & 70.19 & 70.84 & **71.96** \\ \hline \hline \end{tabular}
\end{table}
Table 4: **Combining DART with other DG methods:** OOD performance (%) on OfficeHome (trained using ResNet50) of the proposed method DART coupled with different algorithms against their vanilla and SWAD counterparts. Numbers represented with \({}^{\dagger}\) were reproduced while others are from Domained.

\begin{table}
\begin{tabular}{l c c c c c} \hline \hline
**Backbone** & **MIRO** & **MIRO+SWAD** & **DART** & **DART+MIRO** & **DART+MIRO+SWAD** \\ \hline ResNet-50 (NN-1k ink) & 70.50 & 72.40 & 70.10 & 72.50 & **72.70** \\ ViT-B16 (CLIP init.) & 83.35 & 84.80 & 80.95 & 86.14 & **87.37** \\ \hline \hline \end{tabular}
\end{table}
Table 7: **Integrating DART with MIRO:** OOD performance (%) on OfficeHome of the proposed method DART coupled with MIRO [7] and SWAD [6] on different pretrained backbones.

\begin{table}
\begin{tabular}{l c c c c} \hline \hline
**Algorithm** & **VALLS** & **PACS** & **OfficeHome** & **Terralac** & **DomainNet** & **Avg** \\ \hline ERM [67] & 77.5 \(\pm\) 0.4 & 85.5 \(\pm\) 0.2 & 66.5 \(\pm\) 0.3 & 46.1 \(\pm\) 1.8 & 40.9 \(\pm\) 0.1 & 63.3 \\ + DART (ours) & 78.5 \(\pm\) 0.7 & 87.3 \(\pm\) 0.5 & 70.1 \(\pm\) 0.2 & 48.7 \(\pm\) 0.8 & 45.8 \(\pm\) 0.61 \\ \hline SWAD [68] & 79.1 \(\pm\) 0.1 & 88.1 \(\pm\) 0.1 & 70.6 \(\pm\) 0.2 & 50.0 \(\pm\) 0.3 & 46.5 \(\pm\) 0.1 & 66.9 \\ + DART (ours) & **80.3 \(\pm\) 0.2** & **88.9 \(\pm\) 0.1** & **71.9 \(\pm\) 0.1** & **51.3 \(\pm\) 0.2** & **47.2** & **67.9** \\ \hline \hline \end{tabular}
\end{table}
Table 3: **Domain Generalization:** OOD accuracy(%) of DART trained using ResNet50 compared to the respective baselines on DomainBed datasets. Standard dev. across 3 reruns is reported.

\begin{table}
\begin{tabular}{l c c c c} \hline \hline Method & Pad+Crop+HFip & AutoAug. & Cutout & Cutmix & Mixed-Train. \\ \hline ERM & 81.48 & 83.93 & 82.01 & 83.02 & 85.54 \\ ERM + EMA & 81.67 & 84.20 & 82.33 & 84.05 & 85.57 \\ DART (Ours) & **82.31** & **85.02** & **84.15** & **84.72** & **86.13** \\ \hline \hline \end{tabular}
\end{table}
Table 5: **DART using same augmentation across all branches:** Performance (%) of the DART when compared to baselines across different augmentations on CIFAR-100 using WideResNet-28-10 architecture. DART is better than baselines in all cases.
interpolations from the beginning, and (b) performing interpolations after half the training epochs, as done in DART. It can be noted from (a) that the interpolations in the initial few epochs have poor accuracy since the models are not in a common basin. Further, as seen in initial epochs of (a), when the learning rate is high, SGD training on an interpolated model cannot retain the flat solution due to its implicit bias of moving towards solutions that minimize train loss alone. Whereas, in the later epochs as seen in (b), the improvement obtained after every interpolation is retained. We therefore propose a common training strategy for the initial half of epochs, and split training after that.

**Ablation experiments:** We note the following observations from the plots in Fig.3 (a-e):

1. **Effect of Compute:** Using DART, we obtain higher (or similar) performance gains as the number of training epochs increases, whereas the accuracy of ERM+EMA (Mixed) baseline starts reducing after 300 epochs of training. This can be attributed to the increase in convergence time for learning noisy (or spurious) features due to the intermediate aggregations as shown in Proposition-3, which prevents overfitting.
2. **Effect of Interpolation Frequency:** We note that an optimal range of \(\lambda\) or the number of epochs between interpolations is 10 - 80, and we set this value to 50. If there is no interpolation for longer epochs, the models drift apart too much, causing a drop in accuracy.
3. **Effect of Start Epoch:** We note that although the proposed approach works well even if interpolations are done from the beginning, by performing ERM training on mixed augmentations for 300 epochs, we obtain 0.22% improvement. Moreover, since interpolations do not help in the initial part of training as seen in Fig.4 (a), we propose to start this only in the second half.
4. **Effect of Number of branches:** As the number of branches increases, we note an improvement in performance due to higher diversity across branches, leading to more robustness to spurious features and better generalization as shown in Proposition-2.
5. **Effect of Interpolation epochs:** We perform a run with 50 epochs of common training followed by a single interpolation. We use a fixed learning rate and plot accuracy by varying the interpolation epoch. As this value increases, models drift far apart, reducing the accuracy after interpolation. At epoch-500, the accuracy even reaches 0, highlighting the importance of having a low loss barrier between models.

## 8 Conclusion

In this work, we first show that ERM training using a combination of _diverse_ augmentations within a training minibatch can be a strong benchmark for ID generalization, which is only outperformed by ensembling the outputs of individual experts. Motivated by this observation, we present DART - Diversify-Aggregate-Repeat Training, to achieve the benefits of both training diverse experts and combining their expertise throughout training. The proposed algorithm first trains several models on different augmentations (or domains) to learn a _diverse_ set of features, and further _aggregates_ their weights to obtain better generalization. We repeat the steps Diversify-Aggregate several times over training, and show that this makes the optimization trajectory more robust by suppressing the learning of noisy features, while also ensuring a low loss barrier between the individual models to enable their effective aggregation. We justify our approach both theoretically and empirically on several benchmark In-Domain and Domain Generalization datasets, and show that it integrates effectively with several base training algorithms. We hope our work motivates further research on leveraging the linear mode connectivity of models for better generalization.

Figure 4: **Accuracy of DART across training epochs** for CIFAR-100 on WideResNet-28-10 model: Each branch is trained on different augmentations, whose accuracy is also plotted. Model Interpolation is done (a) from the beginning, (b) after 300 epochs. Although model interpolation and reinitialization happens every 50 epochs, interpolated model accuracy is plotted every epoch.

Figure 3: **Ablations on CIFAR-100, WideResNet-28-10**: (a-d) Experiments comparing DART with the Mixed-Training baseline using the standard training settings. (e) Varying the interpolation epoch after 50 epochs of common training using a fixed learning rate of 0.1.