Distributionally Robust Recurrent Decoders with Random Network Distillation (2022)

Abstract: Neural machine learning models can successfully model language that is similar to their training distribution, but they are highly susceptible to degradation under distribution shift, which occurs in many practical applications when processing out-of-domain (OOD) text. This has been attributed to “shortcut learning””:" relying on weak correlations over arbitrary large contexts. We propose a method based on OOD detection with Random Network Distillation to allow an autoregressive language model to automatically disregard OOD context during inference, smoothly transitioning towards a less expressive but more robust model as the data becomes more OOD, while retaining its full context capability when operating in-distribution. We apply our method to a GRU architecture, demonstrating improvements on multiple language modeling (LM) datasets.

## introduction neural language models have become the main component of modern natural language processing systems, with larger and larger models being used as feature extractors for downstream tasks @xcite , as probability estimators for ranking and ensembling @xcite or as language generators @xcite @xcite . despite their success, neural machine learning models can suffer large performance degradation when they are applied to out-of-domain data which is substantially different than their training data @xcite @xcite . unlike the older statistical language models, recurrent lms (rnnlms) @xcite and their successors transformers lms @xcite can consider the entire prefix of a sentence when predicting or generating the next token. by being able to relate a very high-dimensional input to the output, these models can learn many subtle correlations which are highly useful as long as the input is in-distribution, unfortunately these correlations tend to be brittle to distribution shift, causing a model that depends on them to go astray. this phenomenon is known as "shortcut learning" @xcite and it has been found to also occur in humans and animals, but it is especially prevalent in artificial neural networks. research on this problem has explored models invariant or equivariant w.r.t. certain transformations by means of compositional representations @xcite @xcite , causal modeling @xcite , or both @xcite , but these works focus on classification tasks often on synthetic datasets and can't be straightforwardly applied to black-box language models. approaches specific to lms have focused on robustness where the data domains are known and represented in the training data @xcite . in this work we propose a method that uses random network distillation (rnd) @xcite to dynamically adapt the amount of context that the model relies upon during inference based on an estimate of how much this context is out-of-distribution (ood). this way the model can still make use of all available context when operating within a familiar context space, exploiting long-distance weak correlations, but it reduces to a less expressive and more robust model when operating ood, relying only on the strongest correlations. as a proof of concept we implement our approach on a gru recurrent language model @xcite . while transformer decoders outperform rnns when trained on large training sets, rnns remain competitive on smaller datasets (< 10 7 tokens) where ood phenomena are easier to measure, furthermore they are easier to optimize, simplifying architecture and hyperparameter 1 search. we evaluate our method on language modeling tasks on english datasets, obtaining improvements when evaluating on eight ood domains.we report additional preliminary sequence-to-sequence results on transformer-rnn models @xcite in appendix a. we leave extensions of our method to full transformers as future research. random network distillation in order to estimate how much the state of our rnnlm has deviated from the training distribution we choose the random network distillation (rnd) approach @xcite . given a representation h, we define an ood detector as see @xcite for an extensive analysis. we chose this method because it can be applied to internal representations, is completely unsupervised and does not require any ood tuning data. rnd has been proposed initially in the context of reinforcement learning where the ood signal can be used as a "curiosity" reward to stimulate exploration, and it has been subsequently studied in the context of ood estimation for image classification. to our knowledge, we are the first to apply it to nlp, and to use it to actively compensate for distribution shift rather than just measure it. ## proposed approach our approach consists of estimating how much outof-distribution the state of the model is and scaling it towards the all-zero initial state accordingly, effectively purging the ood context out of the memory of the model and forcing it to rely only on the strongest, usually short-distance, correlations that survive the purge. as the state is pushed towards zero, the model also becomes more conservative in its predictions, avoiding the typical overconfidence of neural networks in ood conditions. specifically, for our language modeling experiments, we train a gru rnnlm as usual, then we freeze it and train a rnd ood estimator on the rnnlm states on the same training set. then during inference we modify the recurrence relation @xcite to perplexity estimation we investigate ood performance with two standard corpora, penn treebank and wikitext2. we evaluate each of the models both in-distribution, on the default test set of its training corpus, and out-of-distribution, on the test set of the other corpus. we also use additional test sets adapted from machine translation robustness evaluations, specifically the english sides of the de-en test sets of @xcite , which is a collection of corpora from diffent domains (i.t., koran, law, medical and movie subtitles) and the english sides of the mtnt ja-en and fr-en test sets of @xcite , which are corpora scraped from reddit and have been used for the wmt-19 robustness shared task @xcite . we report the results in tables 1 and 2. we find that for the word-level models trained on penn treebank our approach improves the perplexity consistently both in-distribution and out-of-distribution for all the test sets we considered. for the wordlevel models trained on wikitext-2 our approach preserves perplexity in-distribution and improves it on most ood test sets, namely the penn treebank test set, the ja-en test set of the mtnt corpus and all the tests sets of @xcite the subtitles test set. the penn treebank results are somewhat anomalous in that the perplexity of some ood test sets is lower than the perplexity of the in-distribution test sets (and in fact the perplexity of the wikitext-2 test set is even lower that the perplexity of the same test set evaluated by its own in-domain model). this effect is caused by the limited vocabulary of the penn treebank training set which causes many of the tokens of the ood test sets to be replaced by unks, which are easy to predict. to avoid this artifact, we evaluate bpe-level models, which are open vocabulary and hence do not introduce any unks. for the bpelevel models we find that for both the baselines and the rnd approach the perplexities on the ood datasets are much higher than the perplexities on the in-domain test sets. comparing our approach to the baselines, we observe a minimal degradation of perplexity in-distribution and substantial improvements on all ood test sets when training on penn treebank, while when training on wikitext-2 we observe more mixed results. in order to analyse if the model is learning sensible values for scaling the out-of-distribution states, we compute the ood scores estimated by the rnd ood detectors, averaged over the two gru layers and over all the tokens in each test set. we report these scores in ablation one could hypothesize that the improvements obtained by our model are due to just increasing the entropy of the output distribution rather than dropping unnecessary context from the rnn state. we evaluate a variant of our model where we apply the ood scaling only on the output of the top-layer rnn but not to the internal states. this increases the output entropy without affecting the context remembered by the model between time steps. this ablation generally improves over the baseline but performs worse than our full model except for the model trained on wikitext-2 bpe where the results are mixed. ## conclusions and future work we proposed a method to improve the robustness of language models to distribution shift caused by train/test domain mismatch. our model contracts the rnn state based on an unsupervised out-of-distribution estimator in order to reduce the model dependency on weak long-distance correlations, which are useful in-distribution but tend to be spurious in out-of-distribution conditions. we obtain perplexity improvements on multiple out-ofdomain test sets without substantial degradation on in-domain test sets. while our approach is based on recurrent decoders, its general principles may be applicable to other neural architectures. for instance, the selfattention heads of a transformer might modulated by an ood detector in order to avoid attending to out-of-distribution parts of a sentence. we anticipate that extending our method to these kind of models will be a promising research direction.