SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Limitations While constructing attention weight factors proves effective, they necessitate iterative algorithms to calculate the set of weights when computing self-attention at each layer, resulting in increased overall complexity. Moreover, the foundational contamination model for these methods is the classical Huber contamination model (Huber, 2011), which requires assumptions about contamination distributions and parameters that may not be universally applicable, especially in discrete settings. To address these limitations, we propose a novel approach that circumvents these computational constraints while effectively fostering robust self-attention mechanisms.
3.2 Robust Self-Attention via Median-of-Means Principle
The Median-of-Means (MoM) principle (Jerrum et al., 1986; Alon et al., 1996) is one other way to construct robust estimators. Rather than taking the average of all the observations, the sample is split into several blocks over which the median is computed. The MoM principle also improved the robustness of KDE and demonstrated its statistical performance under a less restrictive outlier framework (Humbert et al., 2022). More importantly, it can be easily adapted to self-attention. Specifically, we randomly divide the keys  $\{k_j\}_{j=1}^N$  into  $B$  subsets  $I_1, \ldots, I_B$  of equal size, namely,  $|I_1| = |I_2| = \ldots = |I_B| = S$ . Then, the robust estimator of  $p(k)$  takes the following form:
$$
\hat {p} _ {\text {r o b u s t}} (\boldsymbol {k}) \propto \operatorname {m e d i a n} \left\{\hat {p} _ {\sigma , I _ {1}} (\boldsymbol {k}), \dots , \hat {p} _ {\sigma , I _ {B}} (\boldsymbol {k}) \right\}, \tag {10}
$$
where we define  $\hat{p}_{\sigma, I_l}(\boldsymbol{k}) = \frac{1}{S} \sum_{j \in I_l} k_{\sigma}(\boldsymbol{k} - \boldsymbol{k}_j)$  for any  $l \in [B]$ . Similarly, the robust estimator of  $p(\mathbf{v}, \boldsymbol{k})$  is as follows:
$$
\hat {p} _ {\text {r o b u s t}} (\mathbf {v}, \boldsymbol {k}) \propto \operatorname {m e d i a n} \left\{\hat {p} _ {\sigma , I _ {1}} (\mathbf {v}, \boldsymbol {k}), \dots , \hat {p} _ {\sigma , I _ {B}} (\mathbf {v}, \boldsymbol {k}) \right\}, \tag {11}
$$
where  $\hat{p}_{\sigma, I_l}(\mathbf{v}, \mathbf{k}) = \frac{1}{S} \sum_{j \in I_l} k_{\sigma}(\mathbf{v} - \mathbf{v}_j) k_{\sigma}(\mathbf{k} - \mathbf{k}_j)$  for any  $l \in [B]$ . We now propose the self-attention mechanism utilizing the median-of-means principle.
Median-of-Means Self-Attention Mechanism Given the robust estimators in Eq. (10) and (11), we can consider the following robust estimation of the attention:
$$
\widehat {\boldsymbol {h}} _ {i} = \frac {\frac {1}{S} \sum_ {j \in I _ {l}} v _ {j} k _ {\sigma} \left(\boldsymbol {q} _ {i} - \boldsymbol {k} _ {j}\right)}{\text {m e d i a n} \left\{\hat {p} _ {\sigma , I _ {1}} \left(\boldsymbol {q} _ {i} - \boldsymbol {k}\right) , \dots , \hat {p} _ {\sigma , I _ {B}} \left(\boldsymbol {q} _ {i} - \boldsymbol {k}\right) \right\}}, \tag {12}
$$
where  $I_{l}$  is the block such that  $\hat{p}_{\sigma}(\boldsymbol{q}_i - \boldsymbol{k})$  achieves its median value in equation (11). In this context, the random subsets apply to input sequences rather than individual data points, distinguishing this approach from stochastic batches. It's worth noting that our proposed attention mechanism assumes that key and query vectors achieve their median on the same block. Consequently, we apply the median block  $I_{l}$  obtained from the denominator into the numerator, rather than considering the median over all potential blocks, which leads to a process that is faster than computing median blocks on both sides. Moreover, the original MoM principle mandates that each subset be non-overlapping, i.e.  $I_{l_1} \cap I_{l_2} = \emptyset$  for any  $1 \leq l_1 \neq l_2 \leq B$ . However, for structured, high-dimensional data, dividing into non-overlapping blocks may result in the model only gaining a partial perspective of the dataset, leading to sub-optimal performance. Therefore, we construct each subset by sampling with replacement from the original dataset, maintaining the sequential relationship thereafter. In particular, we have found this strategy to be effective in discrete contexts, such as identifying and filtering out aberrant words in a sentence. As illustrated in the bottom right segment of Figure 2, under a word swap attack, the MoM robust attention retains the subsequence that is most relevant to the content, while discarding unhelpful parts. The downside of MoM self-attention is also apparent: since the attention mechanism only accesses a portion of the sequence, it is likely to result in suboptimal performance on clean datasets. The simplicity of MoM allows for easy integration with many existing models. We initially demonstrate that the MoM self-attention mechanism can enhance the recent state-of-the-art FourierFormer (Nguyen et al., 2022c). The theoretical explanation for the ability to remove outliers using MoM-Fourier attention can be found in Appendix C.
Table 1: Perplexity (PPL) and negative likelihood loss (NLL) of our methods (lower part) and baselines (upper part) on WikiText-103. The best results are highlighted in bold font and the second best are highlighted in underline. On clean data, Transformer-SPKDE achieves better PPL and NLL than other baselines. Under random swap with outlier words., Transformers with MoM self-attention show much better performance.
Method (small version)	Clean Data	Word Swap
Valid PPL/Loss	Test PPL/Loss	Valid PPL/Loss	Test PPL/Loss
Transformer (Vaswani et al., 2017b)	33.15/3.51	34.29/3.54	72.28/4.45	74.56/4.53
Performer (Choromanski et al., 2021)	32.35/3.48	33.49/3.51	71.64/4.42	73.48/4.49
Transformer-MGK (Nguyen et al., 2022b)	32.28/3.47	33.21/3.51	69.78/4.38	71.03/4.41
FourierFormer (Nguyen et al., 2022c)	31.86/3.44	32.85/3.49	65.76/4.32	68.33/4.36
Transformer-RKDE (Huber)	31.22/3.42	32.29/3.47	52.14/3.92	55.68/3.99
Transformer-RKDE (Hampel)	31.24/3.42	32.35/3.48	55.61/3.98	57.92/4.03
Transformer-SPKDE	31.05/3.41	32.18/3.46	51.36/3.89	54.97/3.96
Transformer-MoM	33.56/3.52	34.68/3.55	48.29/3.82	52.14/3.92
FourierFormer-MoM	32.26/3.47	33.14/3.50	47.66/3.81	50.96/3.85
3.3 Incorporating Robust Self-Attention Mechanisms into Transformers
Computational Efficiency The two types of robust attention mechanisms proposed above have their own distinct advantages. To expedite the computation for Transformer-RKDE and obtain the attention weight factor in a more efficient manner, we employ a single-step iteration on the alternative updates to approximate the optimal set of weights. Empirical results indicate that this one-step iteration can produce sufficiently precise results. For Transformer-SPKDE, as the optimal set of weights is acquired via the QP solver, it demands more computation time but yields superior performance on both clean and contaminated data. As an alternative to weight-based methods, Transformer-MoM offers significantly greater efficiency while providing competitive performance, particularly with text data. The complete procedure for computing the attention vector for Transformer-RKDE, Transformer-SPKDE, and Transformer-MoM is detailed in Algorithm 1.
Training and Inference We incorporate our robust attention mechanisms into both training and inference stages of Transformers. Given the uncertainty about data cleanliness, there's a possibility of encountering contaminated samples at either stage. Therefore, it is worthwhile to defend against outliers throughout the entire process. In the training context, our methods modify the computation of attention vectors across each Transformer layer, making them less susceptible to contamination from outliers. This entire process remains nonparametric, with no introduction of additional model parameters. However, the resulting attention vectors, whether shaped by re-weighting or the median-of-means principle, diverge from those generated by the standard softmax attention. This distinction influences the model parameters learned during training. During inference, the test data undergoes a similar procedure to yield robust attention vectors. This ensures protection against potential outlier-induced disruptions within the test sequence. However, if we assume the availability of a clean training set — where contamination arises solely from distribution shifts, adversarial attacks, or data poisoning during inference — it is sufficient to restrict the application of the robust attention mechanism to just the inference phase. This could considerably reduce the computational time required for robust attention vector calculation during training. In our empirical evaluation, we engaged the robust attention mechanism throughout both phases and recorded the associated computational time during training. We found that on standard datasets like ImageNet-1K, WikiText-103, and the UEA time-series classification, infusing the robust attention led to a drop in training loss. This suggests that training data itself may contain inherent noise or outliers.
In this section, we provide empirical validation of the benefits of integrating our proposed robust KDE attention mechanisms (Transformer-RKDE/SPKDE/MoM) into Transformer base models. We compare these with the standard softmax Transformer across multiple datasets representing different modalities. These include language modeling on the WikiText-103 dataset (Merit et al., 2016) (Section 4.1) and image classification on ImageNet (Russakovsky et al., 2015; Deng et al., 2009). Furthermore, we assess performance across multiple robustness benchmarks, namely ImageNet-C (Hendrycks & Dietterich, 2019), ImageNet-A (Hendrycks et al., 2021b), ImageNet-O (Hendrycks et al., 2021b), ImageNet-R (Hendrycks et al., 2021a), and ImageNet-Sketch (Wang et al., 2019) (Section 4.2), as well as UEA time-series classification (Section 4.3). Our proposed robust trans
Table 2: Top-1 and top-5 accuracy (%) on ImageNet. The best results are highlighted in bold font and the second best are highlighted in underlines. RVT (Mao et al., 2022) and DeiT (Touvron et al., 2021b) achieve better results on clean data; meanwhile, Transformers incorporating robust self-attention hold stronger defense under different adversarial attacks while still achieving competitive performance on the original ImageNet.
Method	Clean Data	FGSM	PGD	SPSA
Top 1	Top 5	Top 1	Top 5	Top 1	Top 5	Top 1	Top 5
ViT (Dosovitskiy et al., 2020)	72.23	91.13	52.61	82.26	41.84	76.49	48.34	79.36
DeiT (Touvron et al., 2021b)	74.32	93.72	53.24	84.07	41.72	76.43	49.56	80.14
RVT (Mao et al., 2022)	74.37	93.89	53.67	84.11	43.39	77.26	51.43	80.98
FourierFormer (Nguyen et al., 2022c)	73.25	91.66	53.08	83.95	41.34	76.19	48.79	79.57
ViT-RKDE (Huber)	72.83	91.44	55.83	85.89	44.15	79.06	52.42	82.03
ViT-RKDE (Hampel)	72.94	91.63	55.92	85.97	44.23	79.16	52.48	82.07
ViT-SPKDE	73.22	91.95	56.03	86.12	44.51	79.47	52.64	82.33
ViT-MoM	71.94	91.08	55.76	85.23	43.78	78.85	49.38	80.02
FourierFormer-MoM	72.58	91.34	53.25	84.12	41.38	76.41	48.82	79.68
formers are compared with state-of-the-art models, including Performer (Choromanski et al., 2021), MGK (Nguyen et al., 2022a), RVT (Mao et al., 2022), and FourierFormer (Nguyen et al., 2022c). All experiments were conducted on machines with 4 NVIDIA A-100 GPUs. For each experiment, Transformer-RKDE/SPKDE/MoM were compared with other baselines under identical hyperparameter configurations.
4.1 Robust Language Modeling
WikiText-103 is a language modeling dataset that contains collection of tokens extracted from good and featured articles from Wikipedia, which is suitable for models that can leverage long-term dependencies. We follow the standard configurations in Merity et al. (2016); Schlag et al. (2021) and splits the training data into  $L$ -word independent long segments. During evaluation, we process the text sequence using a sliding window of size  $L$  and feed into the model with a batch size of 1. The last position of the sliding window is used for computing perplexity except in the first segment, where all positions are evaluated as in Al-Rfou et al. (2019); Schlag et al. (2021).
In our experiments, we utilized the small and medium (shown in Appendix E) language models developed by Schlag et al. (2021). We configured the dimensions of key, value, and query to 128, and set the training and evaluation context length to 256. We contrasted our methods with Performer (Choromanski et al., 2021), Transformer-MGK (Nguyen et al., 2022a) and FourierFormer (Nguyen et al., 2022c), which have demonstrated competitive performance. For self-attention, we allocated 8 heads for our methods and Performer, and 4 for Transformer-MGK. The dimension of the feedforward layer was set to 2048, with the number of layers established at 16. To prevent numerical instability, we used the log-sum-exp trick in equation (3) when calculating the attention probability vector through the Gaussian kernel. We employed similar tactics when computing the attention weight factor of Transformer-RKDE, initially obtaining the weights in log space, followed by the log-sum-exp trick to compute robust self-attention as outlined in equation (9). For TransformerMoM, the sampled subset sequences constituted  $80\%$  of the length of the original sequence.
Table 1 presents the validation and test perplexity (PPL) for several methods. Both Transformer-RKDE and SPKDE models exhibit an improvement over baseline PPL and NLL on both validation and test sets. However, MoM-based models display slightly higher perplexity, a result of using only a portion of the sequence. When the dataset is subjected to a word swap attack, which randomly substitutes selected keywords with a generic token "AAA" during evaluation, our method, particularly MoM-based robust attention, yields significantly better results. This is particularly evident when filtering out infrequent words, where the median trick has proven its effectiveness. We also noticed superior robustness in RKDE/SPKDE-based robust attention compared to other baseline methods that were not protected from the attack. Our implementation of the word swap is based on the publicly available TextAttack code by Morris et al. (2020)<sup>1</sup>. We employed a greedy search method with constraints on stop-word modifications provided by the TextAttack library.
4.2 Image Classification under Adversarial Attacks and Data Corruptions
For initial simplicity, we utilize the original ViT (tiny) as our base model. However, our methods are versatile and can be integrated with more advanced base models to enhance their robustness. As our
<sup>1</sup>Implementation available at github.com/QData/TextAttack
Table 3: We evaluated the performance of proposed models across multiple robustness benchmarks, using appropriate evaluation metrics for each. In the majority of cases, our methods outperformed the baselines.
Dataset	ImageNet-C	ImageNet-A	ImageNet-O	ImageNet-R	ImageNet-Sketch
Metric	mCE↓	Top-1 Acc↑	AUPR↑	Top-1 Err Rate↓	Top-1 Acc↑
ViT	71.14	0.18	18.31	96.89	37.13
DeiT	70.26	0.73	19.56	94.23	41.68
RVT	68.57	9.45	22.14	63.12	50.07
FourierFormer	71.07	0.69	18.47	94.15	38.36
ViT-RKDE (Huber)	68.69	6.98	25.61	64.33	45.63
ViT-RKDE (Hampel)	68.55	6.34	26.14	62.16	45.76
ViT-SPKDE	68.34	8.29	28.42	61.39	50.13
ViT-MoM	69.11	2.29	28.08	76.44	36.92
FourierFormer-MoM	70.62	2.42	28.86	74.15	39.44
(b)
Figure 3: The top-1 classification accuracy v.s. perturbation budget  $\times 255$  curves on ImageNet against three untargeted attack methods under the  $l_{\infty}$  norm. The proposed set of ViT with robust self-attention mechanisms shows stronger defense under all attack methods with different perturbation budgets.
approaches do not alter the model architecture, each model employs  $5.7M$  parameters. We have also implemented leading-edge methods, including DeiT with hard distillation (Touvron et al., 2021b), FourierFormer (Nguyen et al., 2022c), and Robust Vision Transformer (RVT) (Mao et al., 2022), as our baselines. It's important to note that, for a fair comparison with RVT, we only incorporated its position-aware attention scaling without further architectural modifications. Consequently, the resulting RVT model comprises approximately  $7.2M$  parameters. To evaluate adversarial robustness, we utilized adversarial examples generated by untargeted white-box attacks, which included the single-step attack method FGSM (Goodfellow et al., 2014), multi-step attack method PGD (Madry et al., 2017), and score-based black-box attack method SPSA (Uesato et al., 2018). These attacks were applied to the entire validation set of ImageNet. Each attack distorts the input image with a perturbation budget  $\epsilon = 1/255$  under  $l_{\infty}$  norm, while the PGD attack uses 20 steps with a step size of  $\alpha = 0.15$ . In addition, we assessed our methods on multiple robustness benchmarks, which include images derived from the original ImageNet through algorithmic corruptions or outliers.
Table 2 presents the results under adversarial attack. On clean ImageNet, our performance aligns closely with RVT and DeiT, the leading performers. Notably, our methods outperform RVT under several adversarial attack types, particularly when employing the ViT-SPKDE method. Figure 3 illustrates the relationship between accuracy and perturbation budget across three attack methods. We observe that transformers equipped with robust self-attention mechanisms offer significantly enhanced defense capabilities across different perturbation budgets, with their advantages amplifying as the level of perturbation increases, as expected. We provide more ablation studies in Appendix E that explore dif
Figure 4: Comparison of averaged computation speed (measured by iteration per second) and maximum GPU memory (measured by the CUDA max_memory_allocated function) on ImageNet classification task. The results are measured under the Transformer base models with different capacities: Tiny (5.7M parameters), Small (22M), and Base (86M).
ferent design choices for each proposed robust KDE attention. Table 3 displays the results across multiple robustness benchmarks, employing appropriate evaluation metrics for each. In most instances, the best performance is achieved by our proposed methods, which clearly improve upon existing baselines. Additionally, Figure 4 demonstrates the scalability of our methods as model size increases. We have excluded the SPKDE-based method from this analysis due to its heightened computational demands. The results indicate that both the computation speed and memory increase as model capacity expands for all methods. The MoM-based robust attention closely mirrors the vanilla Transformer, while the RKDE-based robust attention, with one-step approximation, also demonstrates scalability with larger models.
Table 4: A comparison of classification accuracy using the UEA Time Series Classification Archive. Transformers that incorporate our proposed robust attention mechanisms outperform the existing baselines.
Model/Dataset	Ethanol	Heart	PEMS-SF	Spoken	UWave
Transformer	33.70	75.77	82.66	99.33	84.45
FourierFormer	36.12	76.42	86.70	99.00	86.66
Transformer-RKDE (Huber)	34.72	75.84	84.28	99.28	86.49
Transformer-SPKDE	36.09	76.29	86.02	99.36	88.14
Transformer-MoM	38.41	73.24	86.75	97.64	82.97
FourierFormer-MoM	39.89	74.11	87.63	98.12	85.43
4.3 UEA Time Series Classification
Lastly, we conducted experiments using five datasets from the UEA Time-Series Classification Archive (Bagnall et al., 2018), and compared the outcomes across various methodologies (Table 4). The baseline implementation and datasets were adapted from Wu et al. (2022). Our findings indicate that our proposed approaches can notably enhance classification accuracy.
In this work, we explored the link between the dot-product self-attention mechanism and non-parametric kernel regression. This led to the development of a family of fortified transformers, which leverage robust KDE as an alternative to dot-product attention, mitigating the impacts of contaminated samples. We proposed two variants of robust self-attention mechanisms designed to either down-weight or filter out potential corrupted data, both of which can be seamlessly integrated into commonly used transformer models. As our ongoing effort, we are exploring more efficient techniques for estimating the weight set for robust KDE in extremely large models, and incorporating regularization strategies to mitigate the instability of kernel regression with outliers.
Xing Han and Joydeep Ghosh acknowledge support from Intuit Inc. Nhat Ho acknowledges support from the NSF IFML 2019844 and the NSF AI Institute for Foundations of Machine Learning.
Al-Rfou, R., Choe, D., Constant, N., Guo, M., and Jones, L. Character-level language modeling with deeper self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 3159–3166, 2019.
Alon, N., Matias, Y., and Szegedy, M. The space complexity of approximating the frequency moments. In Proceedings of the twenty-eighth annual ACM symposium on Theory of computing, pp. 20-29, 1996.
Baevski, A. and Auli, M. Adaptive input representations for neural language modeling. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id=ByxZX20qFQ.
Bagnall, A., Dau, H. A., Lines, J., Flynn, M., Large, J., Bostrom, A., Southam, P., and Keogh, E. The uea multivariate time series classification archive, 2018. arXiv preprint arXiv:1811.00075, 2018.
Welsch, R. E. and Becker, R. A. Robust non-linear regression using the dogleg algorithm. Technical report, National Bureau of Economic Research, 1975.
Wu, H., Wu, J., Xu, J., Wang, J., and Long, M. Flowformer: Linearizing transformers with conservation flows. In International Conference on Machine Learning, 2022.
Yang, J., Gupta, A., Upadhyay, S., He, L., Goel, R., and Paul, S. Tableformer: Robust transformer modeling for table-text encoding. arXiv preprint arXiv:2203.00274, 2022.
Yang, Z., Dai, Z., Yang, Y., Carbonell, J., Salakhutdinov, R., and Le, Q. V. Xlnet: Generalized autoregressive pretraining for language understanding. arXiv preprint arXiv:1906.08237, 2019.
Zhang, S. and Feng, Y. Modeling concentrated cross-attention for neural machine translation with Gaussian mixture model. In Findings of the Association for Computational Linguistics: EMNLP 2021, pp. 1401-1411, Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021-findings-emnlp.121. URL https://aclanthology.org/2021-findings-emnlp.121.
Zhou, D., Yu, Z., Xie, E., Xiao, C., Anandkumar, A., Feng, J., and Alvarez, J. M. Understanding the robustness in vision transformers. In International Conference on Machine Learning, pp. 27378-27394. PMLR, 2022.
Supplementary Material of "Designing Robust Transformers using Robust Kernel Density Estimation"
A The Non-parametric Regression Perspective of Self-Attention
Given an input sequence  $\mathbf{X} = [\pmb{x}_1, \dots, \pmb{x}_N]^\top \in \mathbb{R}^{N \times D_x}$  of  $N$  feature vectors, the self-attention mechanism transforms it into another sequence  $\mathbf{H} := [h_1, \dots, h_N]^\top \in \mathbb{R}^{N \times D_v}$  as follows:
$$
\boldsymbol {h} _ {i} = \sum_ {j \in [ N ]} \operatorname {s o f t m a x} \left(\frac {\boldsymbol {q} _ {i} ^ {\top} \boldsymbol {k} _ {j}}{\sqrt {D}}\right) \boldsymbol {v} _ {j}, \text {f o r} i = 1, \dots , N. \tag {13}
$$
The vectors  $\pmb{q}_i, \pmb{k}_j$  and  $\pmb{v}_j$  are the query, key and value vectors, respectively. They are computed as follows:
$$
[ \boldsymbol {q} _ {1}, \boldsymbol {q} _ {2}, \dots , \boldsymbol {q} _ {N} ] ^ {\top} := \boldsymbol {Q} = \boldsymbol {X} \boldsymbol {W} _ {Q} ^ {\top} \in \mathbb {R} ^ {N \times D},
$$
$$
\left[ \boldsymbol {k} _ {1}, \boldsymbol {k} _ {2}, \dots , \boldsymbol {k} _ {N} \right] ^ {\top} := \boldsymbol {K} = \boldsymbol {X} \boldsymbol {W} _ {K} ^ {\top} \in \mathbb {R} ^ {N \times D}, \tag {14}
$$
$$
\left[ \boldsymbol {v} _ {1}, \boldsymbol {v} _ {2}, \dots , \boldsymbol {v} _ {N} \right] ^ {\top} := \boldsymbol {V} = \boldsymbol {X} \boldsymbol {W} _ {V} ^ {\top} \in \mathbb {R} ^ {N \times D _ {v}},
$$