Agree to Disagree: Diversity through Disagreement for Better Transferability

TL;DR

The paper tackles the brittleness of gradient-trained models under distribution shifts by addressing the simplicity bias that leads to shortcut learning. It introduces D-BAT, a diversity-inducing regularizer that trains ensembles to agree on in-distribution data but disagree on OOD data, leveraging generalized discrepancy theory. The approach promotes learning diverse predictive features, improves transferability, and yields better uncertainty estimates and OOD detection across multiple datasets. This has practical impact for building robust AI systems capable of reliable performance under domain shifts and unknown downstream conditions.

Abstract

Gradient-based learning algorithms have an implicit simplicity bias which in effect can limit the diversity of predictors being sampled by the learning procedure. This behavior can hinder the transferability of trained models by (i) favoring the learning of simpler but spurious features -- present in the training data but absent from the test data -- and (ii) by only leveraging a small subset of predictive features. Such an effect is especially magnified when the test distribution does not exactly match the train distribution -- referred to as the Out of Distribution (OOD) generalization problem. However, given only the training data, it is not always possible to apriori assess if a given feature is spurious or transferable. Instead, we advocate for learning an ensemble of models which capture a diverse set of predictive features. Towards this, we propose a new algorithm D-BAT (Diversity-By-disAgreement Training), which enforces agreement among the models on the training data, but disagreement on the OOD data. We show how D-BAT naturally emerges from the notion of generalized discrepancy, as well as demonstrate in multiple experiments how the proposed method can mitigate shortcut-learning, enhance uncertainty and OOD detection, as well as improve transferability.

Figures (17)

• Figure 1: Example of applying D-BAT on a simple 2D toy example similar to the LMS-5 dataset introduced by shah2020pitfalls. The two classes, red and blue, can easily be separated by a vertical boundary decision. Other ways to separate the two classes --- with horizontal lines for instance --- are more complex., i.e. they require more hyperplanes. The simplicity bias will push models to systematically learn the simpler feature, as in the second column (b). Using D-BAT, we are able to learn the model in column (c), relying on a more complex boundary decision, effectively overcoming the simplicity bias. The ensemble $h_{ens}(x) = h_1(x) + h_2(x)$, in column (d), outputs a flat distribution at points where the two models disagree, effectively maximizing the uncertainty at those points. In this experiments the samples from $\mathcal{D}_\text{ood}$ were obtained through computing adversarial perturbations, see App. \ref{['app:details-2d-exp']} for more details.
• Figure 2: Illustration of how D-BAT can promote learning diverse features. Consider the task of classifying bird pictures among several classes. The red color represents the attention of a first model $h_1$. This model learnt to use some simple yet discriminative feature to recognise an African Crowned Crane on the left. Now suppose we use the top image $\mathcal{D}_\text{ood}$ on which the models must disagree. $h_2$ cannot again use the same feature as $h_1$ since then it will not disagree on $\mathcal{D}_\text{ood}$. Instead, $h_2$ would look for other distinctive features of the crane which are not present on the right e.g. using its beak and red throat pouch.
• Figure 3: If $h_1$ is computed by minimizing the training loss on ${\mathcal{D}}$, its loss on the OOD task ${\mathcal{D}}_\text{ood}$ may be very large i.e. $h_1$ may be very far from the optimal OOD model $h_\text{ood}$ as measured by ${\mathcal{L}}_{{\mathcal{D}}_\text{ood}}(h_1, h_\text{ood})$ (left). To mitigate this, we propose to learn a diverse ensemble $\{h_1, \dots, h_4\}$ which is maximally 'spread-out' (with distance measured using ${\mathcal{L}}_{{\mathcal{D}}_\text{ood}}(\cdot, \cdot)$) and cover the entire space of possible solutions ${\mathcal{H}}_t^\star$. This minimizes the distance between the unknown $h_\text{ood}$ and our learned ensemble, ensuring we learn transferable features with good performance on ${\mathcal{D}}_\text{ood}$.
• Figure 4: All results are in the "$\mathcal{D}_\text{ood} =$ test data" setting. (a) and (b): Test accuracies as a function of the ensemble size for both D-BAT and Deep Ensembles (ERM ensembles). We observe a significant advantage of D-BAT on both the Waterbirds and the Office-Home datasets. The difference is especially visible on the Waterbirds dataset, which has a stronger spurious correlation. Results have been obtained averaging over $3$ seeds for the Waterbirds dataset and $6$ seeds for the Office-Home dataset. (c): Comparison of D-BAT with several other methods on the Camelyon17, results except D-BAT are taken from extwilds.
• Figure 5: Entropy of ensembles of two models trained with and without D-BAT (deep-ensemble), for inputs $x$ taken from along line $t \cdot \text{1} + (1-t) \cdot \text{0}$ for $t \in [-1,2]$. In-distribution samples are obtained for $t\in\{0,1\}$. All ensembles have a similar test accuracy of $99\%$. Unlike deep ensembles, D-BAT ensembles are able to correctly give high uncertainty values for points far away from the decision boundary. The standard deviations have been omitted here for clarity, but can be seen in App. \ref{['app:stdev-mnist']}.

Theorems & Definitions (1)

• Theorem 3.1: D-BAT favors diversity