Beware of Model Collapse! Fast and Stable Test-time Adaptation for Robust Question Answering (2023)

Abstract: Although pre-trained language models (PLM) have achieved great success in question answering (QA), their robustness is still insufficient to support their practical applications, especially in the face of distribution shifts. Recently, test-time adaptation (TTA) has shown great potential for solving this problem, which adapts the model to fit the test samples at test time. However, TTA sometimes causes model collapse, making almost all the model outputs incorrect, which has raised concerns about its stability and reliability. In this paper, we delve into why TTA causes model collapse and find that the imbalanced label distribution inherent in QA is the reason for it. To address this problem, we propose Anti-Collapse Fast test-time adaptation (Anti-CF), which utilizes the source model‘s output to regularize the update of the adapted model during test time. We further design an efficient side block to reduce its inference time. Extensive experiments on various distribution shift scenarios and pre-trained language models (e.g., XLM-RoBERTa, BLOOM) demonstrate that our method can achieve comparable or better results than previous TTA methods at a speed close to vanilla forward propagation, which is 1.8 \times to 4.4 \times speedup compared to previous TTA methods.

## introduction pre-trained language models (plms) have achieved great success on many nlp tasks @xcite @xcite @xcite @xcite . however, their success is based on the assumption that the test distribution is consistent with the training distribution. in many scenarios, this assumption is not true, such as adversarial attack @xcite , cross-lingual @xcite , cross-domain @xcite , and so on. this situation is known as distribution shift. unfortunately, even the most advanced models currently available, such as chatgpt, do not perform well under the distribution shift @xcite . to address this problem, researchers have proposed many approaches such as adversarial training @xcite , data augmentation @xcite . these methods improve the robustness of the model by changing the training strategy, but according to the no free lunch theorem @xcite , a fixed model still cannot perform perfectly in all distribution-shifted scenarios. therefore, some works @xcite @xcite explore how to update the model during the testing phase to adapt it to the distribution shifts of the test samples, called test time adaptation (tta). a typical approach @xcite uses the shannon entropy of the probability given by the model as the loss to update itself. however, due to the unreliable output of the model, tta may accumulate erroneous information learned in test samples, leading to model collapse and a sharp decline in model performance, which makes tta extremely unstable and unreliable in practical applications. to solve this problem, we take qa task as an example and investigate why tta causes the model collapse. our experiments indicate that the main reason for the model collapse is the imbalanced label distribution of the test data. in contrast to the direct inference, tta exacerbates this imbalanced distribution, making all outputs of the model to be a specific class. therefore, we propose anti-collapse fast test-time adaptation (anti-cf), which utilizes the output of the source model as a soft label to regularize the update of the adapted model during test time to ensure that the adapted model will not deviate too far from the source model, thus avoiding model collapse. however, to obtain the output of the source model and the adapted model, we need to keep the parameters of two models and conduct forward propagation twice, which will bring a lot of additional costs in practical applications. therefore, we freeze the source model and add an efficient side block as the adapted model to reduce the cost of additional forward propagation and back propagation. extensive experiments on various distribution shift scenarios and plms demonstrate that our method can achieve comparable or better results than previous tta methods at a speed close to vanilla forward propagation, which is 1.8× to 4.4× speedup compared to previous tta methods. overall, our contributions in this work include: ## preliminary in this section, we begin by introducing extractive question answering and the application of tta to enhance its robustness. subsequently, we focus on tent @xcite as a case to discuss the potential risks of tta, specifically the occurrence of model collapse. additionally, we conduct an analysis of generative qa in section 5.5. ## method in this section, we propose anti-collapse fast testtime adaptation (anti-cf). anti-cf consists of two strategies. (1) entropy minimization with source constraints (section 3.1) seeks to ensure that the adapted model does not deviate too far from the source model, thus avoiding the occurrence of model collapse. (2) efficient side block (section 3.2) aims to reduce the inference time by building a small network next to the backbone. ## related work test-time adaptation test-time adaptation (tta) is a promising paradigm to deal with the distribution shift. tta uses self-supervised signals to update the model at the inference stage. it has achieved surprising performance in various tasks @xcite . depending on whether or not modifying the objectives of the training phase and accessing the training data, tta can be divided into test-time training and fully test-time adaptation. test-time training @xcite @xcite @xcite is dedicated to designing a distribution shift-aware auxiliary task during the training phase and using this task to update the model during the testing phase. on the other line, the fully test time adaptation @xcite updates the model with test data using a self-training strategy without modifying the training process. however, fully test-time adaptation faces performance degradation due to the low quality of pseudo or soft labels, so current works @xcite @xcite @xcite focus on improving the quality of pseudo-labels, designing more robust and efficient tuning methods. in nlp, tta has also attracted researchers' attention. @xcite first evaluate the effectiveness of fully test-time adaptation in the qa task and introduce an imitation learning paradigm to ensure the quality of pseudo-labels. robustness nlp training a sufficiently robust model is a prerequisite for the practical application of a trustworthy nlp system. @xcite divide robustness in nlp into two categories: adversarial robustness under artificial attacks and naturally occurring distribution shifts. for adversarial robustness, researchers @xcite @xcite @xcite @xcite have proposed many methods for constructing adversarial attack samples to evaluate the model's adversarial robustness. in nlp, naturally occurring distribution shifts also have rich real-world scenarios, such as cross-lingual @xcite @xcite @xcite , cross-domain @xcite @xcite , and on different styles of text. many robustness tuning methods are proposed to improve the robustness of nlp models, such as data augmentation @xcite @xcite @xcite and adversarial training @xcite @xcite . ## conclusion in this paper, we attempt to improve the robustness of qa models by testing time adaptation (tta) but find that tta causes the models collapse. we thoroughly investigate why previous tta methods cause the model collapse and find that the imbalanced label distribution is the main reason. we address this problem by adding constraints between the source and adapted model during the tta process. we also design an efficient side block to speed up the inference time. sufficient experimental results show that our proposed method is effective and efficient, making tta a big step closer to being applied in real-world scenarios.