Not Eliminate but Aggregate: Post-Hoc Control over Mixture-of-Experts to Address Shortcut Shifts in Natural Language Understanding (2024)

Abstract: Recent models for natural language understanding are inclined to exploit simple patterns in datasets, commonly known as shortcuts. These shortcuts hinge on spurious correlations between labels and latent features existing in the training data. At inference time, shortcut-dependent models are likely to generate erroneous predictions under distribution shifts, particularly when some latent features are no longer correlated with the labels. To avoid this, previous studies have trained models to eliminate the reliance on shortcuts. In this study, we explore a different direction: pessimistically aggregating the predictions of a mixture-of-experts, assuming each expert captures relatively different latent features. The experimental results demonstrate that our post-hoc control over the experts significantly enhances the model’s robustness to the distribution shift in shortcuts. Additionally, we show that our approach has some practical advantages. We also analyze our model and provide results to support the assumption.1

## introduction the datasets for natural language understanding (nlu) often contain simple patterns correlated with target labels, which are unintentionally introduced by annotators' simple heuristics, preferences, etc. @xcite . more fundamentally, the compositional nature of natural language inherently introduces tokens that correlate with target labels individually @xcite . for example, word overlap @xcite and specific vocabulary, such as negations @xcite , are known to have such correlations. however, these correlations are not guaranteed to hold in general and are therefore called spurious correlations @xcite . the simple patterns are easy to exploit, so recent nlu models are inclined to take advantage of them. this exploitation or the exploited patterns themselves are called shortcuts @xcite @xcite . 2 at inference time, shortcuts often result in inaccurate predictions under relevant distribution shifts. the shifts can occur, for example, when test data are collected from annotators with different heuristics or preferences @xcite @xcite . data from the same distribution as the training data is referred to as in-distribution (id) data, while data from a distribution shifted relative to the training data is referred to as out-of-distribution (ood) data. figure 1 shows the examples of id and ood data. a simple solution to this problem is to eliminate reliance on shortcuts, which is the mainstream approach, including recent studies in nlu @xcite @xcite . typically, those methods up-weight training instances where some known shortcuts cannot predict labels correctly and down-weight the others. a practical deficiency of this approach arises in a performance trade-off between id and ood data. it deviates models from id data by eliminating shortcuts, which are valid features in id data. due to this trade-off, another practical problem arises where the hyperparameter search has to be made using ood test or validation data, as noticed as the limitation of previous work @xcite @xcite @xcite @xcite . figure 1: an illustrative example of shortcuts in the task of natural language inference. p and h denote the premise and hypothesis sentence, respectively. {a i } are latent features related to x. the value on the right-hand side of ŷ shows the @xmath0.0, 1.0] of the prediction. a i is correctly predictive of label y in the training and id data but not in the ood data where the association between a i and y changed. a * is an ideal latent feature predictive of y across distributions. however, such a * is generally difficult for models to rely on. this figure illustrates the common case where predictions are not based on a * . even when using ood validation data, its distribution is the same as that of test data. thus, in other words, the approach requires knowing the test-time distribution to tune hyperparameters, which is impractical in testing ood robustness. in this paper, we opt not to pursue training to eliminate shortcuts. instead, we propose to aggregate predictions of a mixture model during inference. the problem with shifts in shortcuts is that some latent features in the training data are no longer associated with the labels. we hypothesize that this ood situation can be addressed by effectively aggregating predictions, assuming that the predictions are based on relatively different latent features. we propose a mixture model and its training strategy to encourage such modeling of latent features. at inference time, we perform theoretically grounded risk minimization strategies through post-hoc control for the predictions in the event of potential shifts in shortcuts. the experimental results demonstrate that our method significantly enhances the model's robustness when faced with shifts in shortcuts. moreover, our method shows two other practical benefits that address the problems of previous methods. first, the mixture weights of our model can be used to detect shifts in latent features during inference. this opens up the possibility of adaptive post-hoc control to address the performance trade-off between id and ood data. second, hyperparameters can be tuned with id data only, removing the need to tune hyperparameters with ood data. we also analyze our mixture model and provide results supporting the assumption of modeling latent features. ## background this section first overviews shortcuts. then, we describe how previous approaches have addressed shortcuts and outline how we approach them. below, x and y denote the input instance space and the entire class of target labels, respectively. ## methods the proposed method consists of two parts. the first part is a training method using a mixture model. the second part is a test-time operation, which aggregates the mixture model's predictions to make robust predictions when facing distribution shifts. figure 3 in appendix a shows the overview of our method. ## experiments our goal is to achieve predictions robust to distribution shifts related to shortcuts. in this section, we test whether the proposed post-hoc control improves performance on those ood tests and analyze the mechanism based on our assumption. ## related work as seen in section 2.1, datasets for nlu tasks are known to have multiple shortcuts due to the simple heuristics, preferences, etc., possessed by annotators @xcite , or more fundamentally, the compositional nature of natural language @xcite . a number of studies have addressed the problem of shortcuts in nlu, but their primary difference lies in prior knowledge of shortcuts. known shortcut setting. this setting allows models to know the existence and details of shortcuts in advance. previous studies used this prior knowledge to mitigate the identified shortcuts. reweighting is the basic strategy of previous methods. they used shortcut-dependent models that only take shortcut features as input, e.g., word overlap @xcite @xcite . these shortcut-dependent models let main models know which training instances cannot be predicted correctly via shortcuts and thus should be up-weighted. @xcite showed that the performance of these methods was further enhanced by calibrating the uncertainty of the shortcut-dependent models. utama et al. (2020a) additionally employed a teacher model to adjust the weights so that a main model would not deviate too much from the distribution of training data. @xcite and @xcite trained a main model adversarially to a shortcut-dependent classifier. @xcite and @xcite first trained a model on id data and then re-trained its last classification layer on a small amount of ood data, showing that this small parameter update for the id-fitted model is enough to improve ood performance. several approaches used the counterfactual framework of causal inference. to make counterfactual predictions unaffected by shortcuts, tian et al. (2022) and niu et al. (2021) combined predictions of a main model and a shortcut-dependent model. wang and culotta (2020) classified features into genuine or spurious and selected genuine features for predictions. others utilized identified spurious features to train model predictions to be invariant to interventions on the spurious features @xcite @xcite . the above methods effectively mitigate shortcuts but require the significant cost of careful analysis to achieve the prior knowledge of shortcuts. unknown shortcut setting. the existence and details of shortcuts are generally unknown. another line of studies has sought a way to mitigate shortcuts without the cost of manual identification. the basic strategy is reweighting, just as in the known shortcut setting. to estimate the weights, previous methods have utilized the heuristics that weak models are likely to exploit shortcuts. the weak models include models with limited capacity @xcite , models trained with a limited number of data (utama et al., 2020b), a single epoch @xcite , or shallow layers @xcite . while these methods used continuous weights, @xcite used binary weights that take the value of 1 only for training instances that a weak model misclassified. other approaches applied reweighting when learning to prune fully trained models @xcite @xcite . some studies addressed shortcuts in the feature space by removing redundancy @xcite or correlations @xcite in the space. these studies reported the best scores at different epochs for each of the id validation data and ood test data, so their results are not directly comparable to the other studies. the above methods still require ood test data related to pre-identified shortcuts to tune hyperparameters, as described in section 2.3. our method is different from them in that the training and tuning can be conducted solely on id data. in section 4.2, we demonstrated that in the setting of fully unknown shortcuts where only id data are available, our method improves the performance on ood data significantly better than the previous methods. ## conclusion this study proposed a conceptually novel approach to address the shortcuts problem by pessimistically aggregating the mixture model's predictions at inference time. we introduced the moe-based model, a penalty term to encourage different experts to capture different latent features, and post-hoc control for the mixture weights that is theoretically grounded in risk minimization. the experimental results show that our method not only significantly enhances the model's robustness to shifts in shortcuts but also provides additional benefits to address the previous methods' problems: the performance trade-off between id and ood data and the need for ood test or validation data to tune hyperparameters. our analyses provided results supporting the assumption: different experts capture different latent features to some extent. however, we also noted the limitations in the encoding method (analysis 1), the tested features and interpretability (analysis 2), and the theoretical guarantee of identifiability (analysis 4). future work includes improving the encoding method to capture latent features more accurately, analyzing the instances assigned to the same expert to interpret what it captures and further support the assumption, and theoretically accounting for how the penalty term enhances identifiability. while the focus of this study is on shortcuts, another future direction is extending our method to address a broader range of ood problems (see section 2.1). we believe these are interesting future research departing from this study.