Test-time Augmentation for Factual Probing (2023)

Abstract: Factual probing is a method that uses prompts to test if a language model “knows” certain world knowledge facts. A problem in factual probing is that small changes to the prompt can lead to large changes in model output. Previous work aimed to alleviate this problem by optimizing prompts via text mining or fine-tuning. However, such approaches are relation-specific and do not generalize to unseen relation types. Here, we propose to use test-time augmentation (TTA) as a relation-agnostic method for reducing sensitivity to prompt variations by automatically augmenting and ensembling prompts at test time. Experiments show improved model calibration, i.e., with TTA, model confidence better reflects prediction accuracy. Improvements in prediction accuracy are observed for some models, but for other models, TTA leads to degradation. Error analysis identifies the difficulty of producing high-quality prompt variations as the main challenge for TTA.

## introduction pre-trained language models (lms) such as bert @xcite and t5 @xcite implicitly encode world knowledge from the training corpus in their parameters @xcite . encoded knowledge can be retrieved from an lm via a suitable prompt @xcite . for example, a prompt such as "where did albert einstein die?" is designed to retrieve the fact that "albert einstein died in princeton." however, this capability is not robust since small changes to the prompt can lead to drastic output changes @xcite . if the model fails to answer correctly, it is thus difficult to distinguish if it did not learn the corresponding fact during pre-training or if it actually did but did not produce the correct answer with the given prompt. prior work has aimed at finding better prompts for factual probing foot_0 , typically by employing supervised learning to find an optimal input token sequence for a given relation @xcite @xcite . since these approaches require supervision for each relation, they do not generalize to unseen relation types, which reduces practical applicability. here, we propose to use test time augmentation (tta) as an unsupervised, relation-agnostic approach for improving prompt-based factual probing. tta originates in computer vision, where given an input image at test time, the idea is to 1) apply augmentations such as flipping, cropping, or contrast changes, 2) have the model make predictions for all augmented versions of the image, and then 3) aggregate these predictions to obtain the final prediction. tta has been found to increase robust-ness and to improve model calibration by reducing overconfidence in wrong predictions @xcite @xcite . the benefits are also desirable in factual probing, where lms should exhibit robustness to paraphrases and should generate wellcalibrated predictions. however, tta has found little use in nlp so far. while it is relatively easy to create feature-preserving augmentations of images (e.g., a flipped image of a cat is still an image of a cat), meaning-preserving augmentation of text is a challenging problem since even small edits, such as replacing a single token, can completely change the meaning of a sentence @xcite . in this paper, we empirically demonstrate the benefits and challenges of tta for factual probing. as a similar approach in the context of chainof-thought reasoning, @xcite prepare a single prompt as model input and aggregate multiple outputs ("reasoning paths") to improve model performance. tta differs from this method in that it automatically prepares multiple inputs, i.e., prompts. to apply tta to factual probing, we add a prompt augmenter and a prediction aggregator to the prediction process (figure 1 ). first, the input prompt is automatically augmented by the augmenter. the augmented prompts are then individually fed to an lm. the aggregator collects the outputs for each prompt and determines the final prediction. our evaluation of this setup consists of two parts: we 1) evaluated the overall prediction accuracy and investigated the impact of the number of augmented prompts on the accuracy, and 2) inspected the change in the confidence of model predictions. results showed that the greater the number of augmented prompts, the better the performance when implementing tta. tta was also effective in reducing the number of overconfident and incorrect outputs. however, in terms of overall accuracy, tta was only effective on smaller models. for each relation, we manually created an english prompt template, e.g., "where did {subject} die?" augmenter to yield paraphrases with variety, the augmenter uses three types of prompt augmentations. the first type is synonym replacement, which replaces words in the input prompt with a synonym. for instance, this type of augmentation replaced the word "buried" with "inhumed". synonyms are selected via word embedding similarity using glove @xcite or via word-net @xcite . the second augmentation type is back-translation, with french, russian, german, spanish, and japanese as target languages. the third augmentation type is stopword filtering, which removes stopwords from the prompts. from a single original prompt, the augmenter produces one prompt via stopword filtering and four prompts for each of the seven other augmentation methods, resulting in a total of 30 prompts. ## conclusions we applied the idea of test-time augmentation (tta) to language models, motivated by the observation that models are not robust to prompt variations. specifically, we examined the effect of tta on model accuracy and calibration in a factual probing setting. out of the seven models we investigated, tta had a positive effect on the accuracy of smaller models, namely t5-small, t5-3b, t0_3b. when controlling for the quality of augmentations, tta also improved the accuracy of one more model (t5-large). on other models, tta had a negative effect in terms of accuracy. in terms of model calibration, we observed a positive effect since tta reduced the number of high-confidence incorrect model answers. the main remaining question is why the effect of tta is inconsistent. we hypothesized that the inconsistent effect of tta is due to the poor quality of automatically augmented prompts, and our analysis showed that the high quality of paraphrases is one of the important conditions for tta to be effective. a related question we left to future work is how to enable tta for relations that are difficult to paraphrase automatically, such as the "follows" relation (see section 3.2). while the ultimate goal is arguably to make language models robust against paraphrases without extensions like tta, the results presented in this work show the potential benefit of tta, especially for smaller lms.