STTA: enhanced text classification via selective test-time augmentation

Main contributions

• We analyzed and empirically verified why TTA is sensitive to some data augmentation methods and revealed why some data augmentation methods lead to erroneous predictions.

• We proposed a simple yet effective online TTA method, which selectively aggregates augmented predictions based on risk and reward criteria, thus effectively reduces the bias caused by abnormal transformation samples and enhances the robustness of the model.

• Our proposed method is “plug-and-play”, can be applied to any existing model without the need for hyperparameter tuning or model modification, and can seamlessly collaborate with other robustness methods.

Test-time augmentation

TTA is a technique applied to a trained model during testing, where multiple augmented samples are generated for each original sample. The average prediction over the augmented samples is then used as the aggregated result to improve the final prediction of the model. Although data augmentation is typically applied during model training, it can also be used during prediction. TTA has been widely demonstrated to enhance model accuracy and robustness (Krizhevsky, Sutskever & Hinton, 2012; Matsunaga et al., 2017b; Cohen, Rosenfeld & Kolter, 2019b), address distribution shift issues (Zhang, Levine & Finn, 2022), and defend against adversarial attacks (Prakash et al., 2018; Gao et al., 2020). Researchers have proposed various TTA methods in different domains, including image segmentation (Moshkov et al., 2020), text grammar correction (Yang et al., 2022), text classification (Lu et al., 2022), audio-text retrieval (Kim et al., 2022), theoretical research (Kimura, 2021; Kim, Kim & Kim, 2020), and uncertainty estimation (Conde et al., 2023; Conde & Premebida, 2022).

Although TTA has been widely studied in many tasks (Shanmugam et al., 2020; Lyzhov et al., 2020; Guo et al., 2017; Kimura, 2021), the fact remains that the source data is assumed to be accessible and there is a lack of in-depth insights into data augmentation policies, which is often impractical in reality. For example, both Shanmugam et al. (2020); Lyzhov et al. (2020) utilize a labeled source dataset to learn aggregation. On the other hand, the standard TTA method and the Max method select the maximum logit across all augmented samples, and Smote (Fernández et al., 2018) only considers the nearest sample interpolation to generate new samples, which is prone to interference from noisy augmented samples, resulting in significant bias in aggregation prediction. In this work, we propose a simple yet effective TTA method, called STTA, which effectively combines confidence and similarity to select reliable augmented samples for aggregation.

Ensembling

Ensemble deep learning models use multiple neural networks instead of a single one to compute predictions to improve the performance of various machine learning problems. Typically, ensembling involves obtaining a set of trained neural network models, each of which has a different algorithm or variant, and averaging the predictions for each test object. The various approaches ranging from traditional methods such as Bagging (Breiman, 1996), Boosting (Schapire, 2003), Stacking (Ting & Witten, 1997), to the latest methods homogeneous ensembles (Ganaie et al., 2022) and heterogeneous ensembles (van Rijn et al., 2018; Fang et al., 2021), have resulted in better performance ensemble models.

Sub-ensemble selection

Ensembling has been demonstrated to enhance overall performance and robustness. However, training multiple models for ensembling requires additional computational overhead (Shen, He & Xue, 2019). Even though a single model is used for TTA, it makes sense to view TTA as an ensemble of different models. This is because each sample of the augmentation sub-policy in TTA can be regarded as a new test sample. More specifically, multiple models generate multiple different predictions of the same sample, and TTA employs a single model to generate multiple different predictions for multiple augmented samples. When the predictions of multiple augmented samples in TTA are aggregated, it approximates the ensemble of multiple models. A related instance of this concept is found in the work of Fern & Lin (2008), who proposed to select the optimal clustering result from multiple clustering models for aggregation. This demonstrates the notion of TTA and similar aggregations of diverse augmented samples to improve overall prediction accuracy.

Motivation

As a matter of fact, TTA is still in its infancy in the NLP field, and its effectiveness remains an unresolved issue. It mainly involves two obstacles: (1) It is still unclear which data augmentation methods preserve labels at test time, so simply averaging the predictions of all samples is prone to prediction bias. (2) Previous TTA methods attempt to learn aggregation through neural networks. On the one hand, they assume that the source data is accessible and labeled, as shown in Table 1. On the other hand, previous TTA methods may fail due to inappropriate text augmentation methods. In short, the urgent need has prompted us to propose a selective TTA method.

Data augmentation on test time

We first consider m types of augmentation to obtain N augmented samples, where N = n × m and n represents the number of augmented samples for each type of augmentation. During the inference at time step t, for a given test sample x_t, we can obtain N augmented samples ${\tilde{x}}_{t,i}$ by applying m types of augmentation functions $a_i\in \mathcal{A}$ to x_t. Then, we can obtain the predicted probability distribution p_t,i for each augmented sample ${\tilde{x}}_{t,i}$ as follows: (2)${\displaystyle p_{t,i}=f_{\theta }\left({\tilde{x}}_{t,i}\right)\in {\Delta }^C}$ where $f_{\theta }\left({\tilde{x}}_{t,i}\right)$ represents the logit associated with the ith augmented sample ${\tilde{x}}_{t,i}$. Then, we can define ∀i ∈ [1, n × m], (3)${\displaystyle p_{t,i}=\left[p_{t,i}^1,\dots ,p_{t,i}^C\right]\in {\mathbb{R}}^C}$ and then the augmented logits matrix P^N×C can be defined as follows: (4)${\displaystyle {\mathbf{P}}^{N\times C}={\left[p_{t,i}\right]}_{i=1,\dots ,N}\left[\begin{array}{cccc}\hfill p_{t,0}^1\hfill & \hfill p_{t,0}^2\hfill & \hfill \cdots \hfill & \hfill p_{t,0}^C\hfill \\ \hfill p_{t,1}^1\hfill & \hfill p_{t,1}^2\hfill & \hfill \cdots \hfill & \hfill p_{t,1}^C\hfill \\ \hfill ⋮\hfill & \hfill ⋮\hfill & \hfill \ddots \hfill & \hfill ⋮\hfill \\ \hfill p_{t,N}^1\hfill & \hfill p_{t,N}^2\hfill & \hfill \cdots \hfill & \hfill p_{t,N}^C\hfill \\ \hfill \hfill \end{array}\right]}$

Similarity-based recognized

Suppose we have a set of augmented samples ${\tilde{x}}_{t,i}$, where i ∈ [1, N]. Let g_ϕ:𝒳 → ℝ^d be the model encoder to map the augmented samples ${\tilde{x}}_{t,i}$ to a d-dimensional feature space ℝ^d. Then, we can obtain the feature representation of the augmented samples ${\tilde{x}}_{t,i}$ as follows: (5)${\displaystyle {\mathbf{F}}^{N\times d}={\left[g_{\varphi }\left({\tilde{x}}_{t,i}\right)\right]}_{i=1,\dots ,N}}$ where $g_{\varphi }\left({\tilde{x}}_{t,i}\right)\in {\mathbb{R}}^d$ represents the feature representation of the ith augmented sample ${\tilde{x}}_{t,i}$.

Then we feed the feature representation F^N×d into the HNSW algorithm (Malkov & Yashunin, 2018) to calculate the similarity distance $\mathcal{S}$ between the augmented samples and the original sample as shown in Algorithm 2 .

 
 
Algorithm 2 Similarity-based Recognized 
Require: Augmented feature matrix FN×d 
Ensure: Similarity distance set S 
 1:  Initialize the graph G with a single node v0 and set v0 as the entry point 
  2:  Current level l ← 0 
  3:  for i = 1 to N do 
 4:       vi ← HNSW (FN×d, v0) 
  5:       Add vi to G and connect it to v0 
  6:       if i2l = 0 then 
 7:            l ← l + 1 
  8:       end if 
 9:  end for 
10:  Query feature q ← gϕ(xt) 
11:  Initialize the priority queue PQ with v0 
12:  Initialize the set of visited nodes V with v0 
13:  search ← 0 
14:  efSearch ← 100 
15:  while search < efSearch do 
16:       v ← PQ.pop() 
17:       for vi ∈ v do 
18:            if vi / ∈ V then 
19:                  V ← V ∪ vi 
20:                  PQ ← PQ ∪ vi 
21:            end if 
22:       end for 
23:       search ← search + 1 
24:  end while 
25:  S←{d(q,vi)|vi ∈ V } 
26:  return S    

$\mathcal{S}$ indicates the similarity distance between x_t and each augmented sample ${\tilde{x}}_{t,i}$, so we can divide the augmented samples into quartile intervals by similarity distance. (6)${\displaystyle {\mathcal{Q}}_{sim}\left({\mathbf{P}}^{N\times C}\right):\left\{\bigcup _{q=1}^4{\mathcal{S}}_q\right\}\to \left\{\begin{array}{c}{\mathcal{S}}_1,{\mathcal{S}}_2,{\mathcal{S}}_3,{\mathcal{S}}_4\hfill \end{array}\right\}}$ where ${\mathcal{S}}_q$ represents the set of augmented samples in the q-th quartile interval based on similarity distance.

Confidence-based recognized

The confidence of the model prediction is defined as the maximum probability value of the predicted probability distribution of the model. For each augmented sample ${\tilde{x}}_{t,i}$, the confidence score can be obtained using the following formula: (7)${\displaystyle conf\left(x_{t,i}\right)=\underset{j=1,\dots ,C}{max}{p}_{t,i}^j}$

Then, we can obtain the set of confidence scores for augmented samples, $\mathcal{C}$, by sorting the confidence scores of the augmented samples and dividing them into quartile intervals. (8)${\displaystyle {\mathcal{Q}}_{conf}\left({\mathbf{P}}^{N\times C}\right):\left\{\bigcup _{g=1}^4{\mathcal{C}}_g\right\}\to \left\{\begin{array}{c}{\mathcal{C}}_1,{\mathcal{C}}_2,{\mathcal{C}}_3,{\mathcal{C}}_4\hfill \end{array}\right\}}$ where ${\mathcal{C}}_g$ represents the set of augmented samples in the q-th quartile interval based on confidence score.

Sample roles recognition

Based on the two perspectives mentioned above, we can divide the all augmented samples into four different roles according to the high and low degrees of confidence and similarity. The details are as follows:

• W_gold refers to samples with high confidence and are the most similar to the original sample.These samples can effectively evaluate and improve the model’s generalization ability, making them valuable for aggregation. It can be formulated as $\mathit{candidates}=\left\{x_{t_i}|x_{t_i}\in \left({\mathcal{C}}_1\cup {\mathcal{C}}_2\right)\cap \left({\mathcal{S}}_1\cup {\mathcal{S}}_2\right)\right\}$

• W_bonus refers to samples with high confidence and low similarity which can provide additional feature information for aggregation. It can be formulated as $\mathit{candidates}=\left\{x_{t_i}|x_{t_i}\in \left({\mathcal{C}}_1\cup {\mathcal{C}}_2\right)\cap \left({\mathcal{S}}_3\cup {\mathcal{S}}_4\right)\right\}$

• W_potential refers to samples with high similarity to the original sample and low confidence which may have some deviation in the augmentation process but still retain the features of the original sample. It can be formulated as $\mathit{candidates}=\left\{x_{t_i}|x_{t_i}\in \left({\mathcal{C}}_3\cup {\mathcal{C}}_4\right)\cap \left({\mathcal{S}}_1\cup {\mathcal{S}}_2\right)\right\}$

• W_risk refers to samples with low confidence and similarity to the original sample which lose the features of the original sample and are less valuable for aggregation. It can be formulated as $\mathit{candidates}=\left\{x_{t_i}|x_{t_i}\in \left({\mathcal{C}}_3\cup {\mathcal{C}}_4\right)\cap \left({\mathcal{S}}_3\cup {\mathcal{S}}_4\right)\right\}$

Note that although W_gold theoretically has a greater value contribution. However, this does not imply that W_risk and W_potential do not have the ability to contribute.In practical applications, the selection can be made based on specific circumstances. Here, we have selected W_gold, W_bonus, and W_potential as candidate samples in all experiments.

Base on the above analysis, we can obtain the candidate set $W_{\mathit{candidates}}^K$ for aggregation: (9)${\displaystyle W_{\mathit{candidates}}^K=\left\{x_{t_i}|x_{t_i}\in \left(W_{\text{gold}}\cup W_{\text{bonus}}\right)\cup W_{\text{potential}}\right\}}$ where K is the number of samples in the candidate set, the matrix of logits for the selected samples, denoted as P^K×C, is defined as follows: (10)${\displaystyle {\mathbf{P}}^{N\times C}\to \left[AugmentedSamples\right]\text{Select}{\mathbf{P}}^{K\times C}={\left[p_{t,i}\right]}_{i=1,\dots ,K}\left[\begin{array}{cccc}\hfill p_{t,0}^1\hfill & \hfill p_{t,0}^2\hfill & \hfill \cdots \hfill & \hfill p_{t,0}^C\hfill \\ \hfill p_{t,1}^1\hfill & \hfill p_{t,1}^2\hfill & \hfill \cdots \hfill & \hfill p_{t,1}^C\hfill \\ \hfill ⋮\hfill & \hfill ⋮\hfill & \hfill \ddots \hfill & \hfill ⋮\hfill \\ \hfill p_{t,K}^1\hfill & \hfill p_{t,K}^2\hfill & \hfill \cdots \hfill & \hfill p_{t,K}^C\hfill \\ \hfill \hfill \end{array}\right]}$

The final prediction probability vector $\widetilde{p_t}$ instead of p_t is obtained by aggregating the logits matrix P^K×C and the original prediction probability vector p₀ as follows: (11)${\displaystyle {\tilde{p}}_t=\omega \cdot p_{t,0}+\left(1-\omega \right)\cdot \frac{1}{K}\sigma \left({\mathbf{P}}^{K\times C}\right)}$ (12)${\displaystyle \text{with}\omega =max\left\{\right.\omega \in \left[0,\alpha \right]:arg\underset{i\in \left\{1,\dots ,k\right\}}{max}{\tilde{p}}_t\left(\omega \right)=arg\underset{i\in \left\{1,\dots ,k\right\}}{max}{p}_{t,0}\left\}\right.}$ where ω is the aggregation weight, σ:P^K×C → Δ_k represents the softmax function, and α←0.5 since we consider both the original prediction and the aggregated prediction equally important. Although more advanced techniques exist to determine the optimal value of α. Nonetheless, our primary goal is to demonstrate the ability of our approach to mitigate bias through the memory bank.

Setup

Effect of augmentation policies

The factors that affect TTA include data augmentation methods and how to aggregate predictions effectively. In this section, we will conduct extensive discussions on different data augmentation methods to demonstrate the effectiveness of our method.

Note that although more complex data augmentation methods may be more effective in theory, but they are not necessarily representative in practice. Due to the large number of possible combinations of the above data augmentation methods, we mainly consider the most representative data augmentation methods and then combine them to form six data augmentation policies based on insights from Lu et al. (2022), as shown in Table 2, where RPI, RWSR, RWI, RWS, and RWD represent random punctuation insertion, random word substitution with replacement, random word insertion, random word swap, and random word deletion, respectively. And when the data augmentation method is Aug. 6 (RWSR + RWI + RWS + RWD), the standard data augmentation method is equivalent to the TTA method given by Lu et al. (2022).

Figure 3 illustrates a boxplot displaying the distribution of accuracy for different data augmentation policies. It can be observed that our method effectively adapts to various data transformations, even performing well in cases where semantic changes are likely to occur. This indicates that our approach successfully identifies anomalous samples from the batch data, preserving valuable contributions and thereby enhancing the accuracy of the model. In contrast, other methods lacking this capability experience a greater decline in performance when faced with data augmentation policies such as RWD.

Additionally, we observe that in other methods, the outliers in the boxplot tend to be located in the region where the model’s net gain decreases. However, in our method, the outliers are generally situated in the region where the model’s net gain increases. This suggests that our approach has even greater potential to further improve model gains.

Although the gains achieved by our method may not be significant for some data augmentation policies, overall, our approach consistently demonstrates absolute improvements across all employed data augmentation policies.

Main results

Table 3 displays the absolute accuracy improvement results of all TTA methods on representative benchmark datasets for text classification. In general, most TTA methods fail to generate positive performance gains and, to varying extents, impair the model’s performance. Contrarily, our approach achieves significant performance improvements across all datasets, with an average increase of 0.50%. Notably, our approach demonstrates impressive performance gains of 0.72% and 0.70% on RTE and MRPC, respectively, indicating its effectiveness in handling diverse types of datasets.

In contrast, Partial-LR utilizes a learnable network to allocate weights to different samples, but it introduces an additional assumption that the source data is accessible. However, its performance improvement falls short of expectations. As discussed in ‘Related Work’, the unreasonable data augmentation operations may destroy the sample label, making it challenging for Partial-LR to learn the correct weight of each augmented sample.

In the context of uncertain predictions stemming from noisy samples, the baseline model may prove challenging in making robust predictions. Consequently, overconfident and erroneous predictions may arise. Such a scenario presents difficulty in selecting the appropriate high-confidence sample from a set of augmented samples when employing the Max method. This, in turn, results in a notable decline of 10.60% in performance on TREC Fine. And Hard Vote also faces this obstacle. Furthermore, it is worth noting that while Smote achieves the second-best performance on TREC Fine, i.e., 0.40%, its interpolation operation on the original samples may also be impaired by the presence of additional noisy samples. As a result, the problem of erroneous predictions may still persist.

It is worth noting that we observe that when the baseline model performs poorly, such as on the RTE and MRPC datasets, TTA achieves a greater improvement. As the baseline performance improves, when the model performs well, such as SUBJ, TREC Coarse and TREC Fine datasets with an accuracy of over 90%, the improvement of TTA will decrease. Specifically, when we turn our attention to the TREC Coarse dataset, which has a good baseline performance, the improvement of some TTA methods is even 0, such as Smote. This finding suggests that TTA is best applied when the baseline model performs poorly. Additionally, Mean can achieve the highest performance improvement on TREC Fine with a simple average, but it performs poorly on other datasets, which suggests that there is great potential for TTA, but there is a lack of reasonable ways to identify those noisy samples caused by inappropriate data augmentation methods. Overall, although our method does not significantly improve the datasets on which the model is already performing well, such as SUBJ, our method does not compromise any model performance. These results demonstrate that selectively aggregating predictions based on sample roles is effective.

Necessity of both similarity and confidence

To verify the importance of both confidence and similarity in STTA, we labeled high (low) confidence samples as C_h(C_l) and high (low) similarity samples as S_h(S_l). Then we conducted the following ablation experiments:

• STTA w/o Conf.: When discriminating the roles of the augmented samples, we only use similarity to divide them into two types: S_h and S_l.

• STTA w/o Sim.: When discriminating the roles of the augmented samples, we only use confidence scores to divide them into two types: C_h and C_l.

• STTA with Conf.&Sim.: When discriminating the roles of the augmented samples, we use both similarity and confidence scores to divide them.

From Fig. 5, the experimental results indicate that conducting experiments solely based on similarity or confidence both lead to a decline in model profit gained, while STTA w/o Conf. is worse than STTA w/o Sim., which indicates that similarity plays a greater role than confidence in aggregating predictions. Furthermore, the results also show that STTA with Conf. & Sim. indicates the best performance, which demonstrates the necessity of combining similarity and confidence to define the different roles of augmented samples.

Effect of different roles

To verify the necessity of different sample roles (W_gold, W_bonus, W_potential and W_risk) and the combination of different sample roles, we conducted experiments separately. As shown in Fig. 6, we observed that the gains of W_gold, W_bonus, W_potential and W_risk exhibit a linear downward trend, which confirms the effectiveness of our method in distinguishing valid samples. However, any single sample recognition role is not as good as the method (ours) that combines different roles, which indicates that it is necessary to combine different roles.

Effect of number of augmentations

Our analysis of the results presented in Fig. 7 aims to explore the impact of different quantities of data augmentation on the performance of TTA. Notably, our proposed method, STTA, exhibits a distinct linear positive correlation between the number of augmentations and model performance. Conversely, the standard TTA method does not yield any discernible advantage with increasing the number of augmentation, while the Max method exacerbates the decline in model performance. Although Partial-LR mitigates the detrimental effects as the number of augmentations increases, it still falls short of achieving positive gains.