Title: Designing Robust Transformers using Robust Kernel Density Estimation
Venue: NeurIPS 2023
Accepted: yes
Keywords: Transformers

Kernel Density Estimation

Robustness
Abstract: Transformer-based architectures have recently exhibited remarkable successes across different domains beyond just powering large language models. However, existing approaches typically focus on predictive accuracy and computational cost, largely ignoring certain other practical issues such as robustness to contaminated samples. In this paper, by re-interpreting the self-attention mechanism as a non-parametric kernel density estimator, we adapt classical robust kernel density estimation methods to develop novel classes of transformers that are resistant to adversarial attacks and data contamination. We first propose methods that down-weight outliers in RKHS when computing the self-attention operations. We empirically show that these methods produce improved performance over existing state-of-the-art methods, particularly on image data under adversarial attacks. Then we leverage the median-of-means principle to obtain another efficient approach that results in noticeably enhanced performance and robustness on language modeling and time series classification tasks. Our methods can be combined with existing transformers to augment their robust properties, thus promising to impact a wide variety of applications.

Designing Robust Transformers using Robust Kernel Density Estimation

Abstract

Transformer-based architectures have recently exhibited remarkable successes across different domains beyond just powering large language models. However, existing approaches typically focus on predictive accuracy and computational cost, largely ignoring certain other practical issues such as robustness to contaminated samples. In this paper, by re-interpreting the self-attention mechanism as a non-parametric kernel density estimator, we adapt classical robust kernel density estimation methods to develop novel classes of transformers that are resistant to adversarial attacks and data contamination. We first propose methods that down-weight outliers in RKHS when computing the self-attention operations. We empirically show that these methods produce improved performance over existing state-of-the-art methods, particularly on image data under adversarial attacks. Then we leverage the median-of-means principle to obtain another efficient approach that results in noticeably enhanced performance and robustness on language modeling and time series classification tasks. Our methods can be combined with existing transformers to augment their robust properties, thus promising to impact a wide variety of applications.

1 Introduction

Attention mechanisms and transformers (Vaswani et al., 2017a) have drawn lots of attention in the machine learning community (Lin et al., 2021; Tay et al., 2020; Khan et al., 2021). Now they are among the best deep learning architectures for a variety of applications, including those in natural language processing (Devlin et al., 2019; Al-Rfou et al., 2019; Dai et al., 2019; Child et al., 2019; Raffel et al., 2020; Baevski & Auli, 2019; Brown et al., 2020; Dehghani et al., 2019), computer vision (Dosovitskiy et al., 2021; Liu et al., 2021b; Touvron et al., 2021a; Ramesh et al., 2021; Radford et al., 2021; Fan et al., 2021; Liu et al., 2022), and reinforcement learning (Chen et al., 2021; Janner et al., 2021). They are also known for their effectiveness in transferring knowledge from various pretraining tasks to different downstream applications with weak supervision or no supervision (Radford et al., 2018, 2019; Devlin et al., 2019; Yang et al., 2019; Liu et al., 2019). 37th Conference on Neural Information Processing Systems (NeurIPS 2023).

While there have been notable advancements, the robustness of the standard attention module remains an unresolved issue in the literature. In this paper, our goal is to reinforce the attention mechanism and construct a comprehensive framework for robust transformer models. To achieve this, we first revisit the interpretation of self-attention in transformers, viewing it through the prism of the Nadaraya-Watson (NW) estimator (Nadaraya, 1964) in a non-parametric regression context. Within the transformer paradigm, the NW estimator is constructed based on the kernel density estimators (KDE) of the keys and queries. However, these KDEs are not immune to the issue of sample contamination (Kim & Scott, 2012). By conceptualizing the KDE as a solution to the kernel regression problem within a Reproducing Kernel Hilbert Space (RKHS), we can utilize a range of state-of-theart robust KDE techniques, such as those based on robust kernel regression and median-of-mean estimators. This facilitates the creation of substantially more robust self-attention mechanisms. The resulting suite of robust self-attention can be adapted to a variety of transformer architectures and tasks across different data modalities. We carry out exhaustive experiments covering vision, language modeling, and time-series classification. The results demonstrate that our approaches can uphold comparable accuracy on clean data while exhibiting improved performance on contaminated data. Crucially, this is accomplished without introducing any extra parameters.

Related Work on Robust Transformers: Vision Transformer (ViT) models (Dosovitskiy et al., 2020; Touvron et al., 2021b) have recently demonstrated impressive performance across various vision tasks, positioning themselves as a compelling alternative to CNNs. A number of studies (e.g., Subramanya et al., 2022; Paul & Chen, 2022; Bhojanapalli et al., 2021; Mahmood et al., 2021; Mao et al., 2022; Zhou et al., 2022) have proposed strategies to bolster the resilience of these models against common adversarial attacks on image data, thereby enhancing their generalizability across diverse datasets. For instance, Mahmood et al. (2021) provided empirical evidence of ViT's vulnerability to white-box adversarial attacks, while demonstrating that a straightforward ensemble defense could achieve remarkable robustness without compromising accuracy on clean data. Zhou et al. (2022) suggested fully attentional networks to enhance self-attention, achieving state-of-the-art accuracy on corrupted images. Furthermore, Mao et al. (2022) conducted a robustness analysis on various ViT building blocks, proposing position-aware attention scaling and patch-wise augmentation to enhance the model's robustness and accuracy. However, these investigations are primarily geared toward vision-related tasks, which restricts their applicability across different data modalities. As an example, the position-based attention from Mao et al. (2022) induces a bi-directional information flow, which is limiting for position-sensitive datasets such as text or sequences. These methods also introduce additional parameters. Beyond these vision-focused studies, robust transformers have also been explored in fields like text analysis and social media. Yang et al. (2022) delved into table understanding and suggested a robust, structurally aware table-text encoding architecture to mitigate the effects of row and column order perturbations. Liu et al. (2021a) proposed a robust end-to-end transformer-based model for crisis detection and recognition. Furthermore, Li et al. (2020) developed a unique attention mechanism to create a robust neural text-to-speech model capable of synthesizing both natural and stable audios. We have noted that these methods vary in their methodologies, largely due to differences in application domains, and therefore limiting their generalizability across diverse contexts.

Other Theoretical Frameworks for Attention Mechanisms: Attention mechanisms in transformers have been recently studied from different perspectives. Tsai et al. (2019) show that attention can be derived from smoothing the inputs with appropriate kernels. Katharopoulos et al. (2020); Choromanski et al. (2021); Wang et al. (2020) further linearize the softmax kernel in attention to attain a family of efficient transformers with both linear computational and memory complexity. These linear attentions are proven in Cao (2021) to be equivalent to a Petrov-Galerkin projection (Reddy, 2004), thereby indicating that the softmax normalization in dot-product attention is sufficient but not necessary. Other frameworks for analyzing transformers that use ordinary/partial differential equations include Lu et al. (2019); Sander et al. (2022). In addition, the Gaussian mixture model and graph-structured learning have been utilized to study attentions and transformers (Tang & Matteson, 2021; Gabbur et al., 2021; Zhang & Feng, 2021; Wang et al., 2018; Kreuzer et al., 2021). Nguyen et al. (2022c) has linked the self-attention mechanism with a non-parametric regression perspective, which offers enhanced interpretability of Transformers. Our approach draws upon this viewpoint, but focuses instead on how it can lead to robust solutions.

2 Self-Attention Mechanism from a Non-parametric Regression Perspective

Assume we have the key and value vectors { k j , v j } j ∈ [ N ] that is collected from the data generating process v = f ( k ) + ε , where ε is some noise vectors with E [ ε ] = 0 , and f is the function that we want to estimate. We consider a random design setting where the key vectors { k j } j ∈ [ N ] are i.i.d. samples from the distribution p ( k ) , and we use p ( v , k ) to denote the joint distribution of ( v , k ) defined by the data generating process. Our target is to estimate f ( q ) for any new queries q . The NWestimator provides a non-parametric approach to estimating the function f , the main idea is that

where the first equation comes from the fact that E [ ε ] = 0 , the second equation comes from the definition of conditional expectation, and the last equation comes from the definition of conditional density. To provide an estimation of f , we just need to obtain estimations for both the joint density function p ( v , k ) and the marginal density function p ( k ) . KDE is commonly used for the density estimation problem (Rosenblatt, 1956; Parzen, 1962), which requires a kernel k σ with the bandwidth parameter σ satisfies ∫ R D k σ ( x -x ′ ) d x = 1 , ∀ x ′ , and estimate the density as

where [ v , k ] denotes the concatenation of v and k . Specifically, when k σ is the isotropic Gaussian kernel, we have ˆ p σ ( v , k ) = 1 N ∑ j ∈ [ N ] k σ ( v -v j ) k σ ( k -k j ) . Combine this with Eq. (1) and Eq. (2), we can obtain the NW estimator of the function f as

Furthermore, it is not hard to show that if the keys { k j } j ∈ [ N ] are normalized, the selfattention mechanism ̂ f σ ( q i ) in Eq. (3) is exactly the standard Softmax attention ̂ f σ ( q i ) = ∑ j ∈ [ N ] softmax ( q ⊤ k j /σ 2 ) v j . Such an assumption on the normalized key { k j } j ∈ [ N ] can be mild, as in practice we always have a normalization step on the key to stabilizing the training of the transformer (Schlag et al., 2021). If we choose σ 2 = √ D , where D is the dimension of q and k j , then ̂ f σ ( q i ) = h i . As a result, the self-attention mechanism in fact performs a non-parametric regression with NW-estimator and isotropic Gaussian kernel when the keys are normalized.

KDE as a Regression Problem in RKHS We start with the formal definition of the RKHS. The space H k = { f | f : X → R } is the RKHS associated with kernel k , where k : X × X → R , if it is a Hilbert space with inner product ⟨· , ·⟩ H k and following properties:

k ( x , · ) ∈ H k , ∀ x ∈ X ;

∀ f ∈ H k , f ( x ) = ⟨ f, k ( x , · ) ⟩ H k . Aka the reproducing property.

With slightly abuse of notation, we define k σ ( x , x ′ ) = k σ ( x -x ′ ) . By the definition of the RKHS and the KDE estimator, we know ˆ p σ = 1 N ∑ j ∈ [ N ] k σ ( x j , · ) ∈ H k σ , and can be viewed as the optimal solution of the following least-square regression problem in RKHS:

…

…

…

Clean Data

Brightness

Image Patches

Attention Weight Factor

Later that year, she auctioned a breakfast in Mayfair,

London, where she raised around £4000 for the Pratham

NGO, which helps children's primary education.

Later that year, she auctioned a breakfast in Mayfair,

London, where she AAA around £4000 for the Pratham

AAA, which helps children's primary education.

̸

Note that, in Eq. (4), the same weight factor 1 /N is applied uniformly to each error term ∥ k σ ( x j , · ) -p ∥ 2 H kσ . This approach functions effectively if there are no outliers in the set { k σ ( x j , · ) } j ∈ [ N ] . However, when outliers are present (for instance, when there is some j such that ∥ k σ ( x j , · ) ∥ H kσ ≫∥ k σ ( x i , · ) ∥ H kσ , ∀ i ∈ [ N ] , i = j ), the error attributable to these outliers will overwhelmingly influence the total error, leading to a significant deterioration in the overall density estimation. We illustrate the robustness issue of the KDE in Figure 1. The view that KDE is susceptible to outliers, coupled with the non-parametric understanding of the self-attention mechanism, implies a potential lack of robustness in Transformers when handling outlier-rich data. We now offer a fresh perspective on this robustness issue, introducing a universal framework that is applicable across diverse data modalities.

3 Robust Transformers that Employ Robust Kernel Density Estimators

Drawing on the non-parametric regression formulation of self-attention, we derive multiple robust variants of the NW-estimator and demonstrate their applicability in fortifying existing Transformers. We propose two distinct types of robust self-attention mechanisms and delve into the properties of each, potentially paving the way for Transformer variants that are substantially more robust.

3.1 Down-weighting Outliers in RKHS

Inspired by robust regression (Fox & Weisberg, 2002), a direct approach to achieving robust KDE involves down-weighting outliers in the RKHS. More specifically, we substitute the least-square loss in Eq. (4) with a robust loss function ρ , resulting in the following formulation:

Examples of the robust loss function ρ include the Huber loss (Huber, 1992), Hampel loss (Hampel et al., 1986), Welsch loss (Welsch & Becker, 1975) and Tukey loss (Fox & Weisberg, 2002). We empirically evaluate different loss functions in our experiments. The critical step here is to estimate the set of weights ω = ( ω 1 , · · · , ω N ) ∈ ∆ N , with each ω j ∝ ψ ( ∥ k σ ( x j , · ) -ˆ p robust ∥ H kσ ) , where ψ ( x ) := ρ ′ ( x ) x . Since ˆ p robust is defined via ω , and ω also depends on ˆ p robust, one can address this circular definition problem via an alternative updating algorithm proposed by Kim & Scott (2012). The algorithm starts with randomly initialized ω (0) ∈ ∆ n , and performs alternative updates between ˆ p robust and ω until the optimal ˆ p robust is reached at the fixed point (see details in Appendix A).

However, while this technique effectively diminishes the influence of outliers, it also comes with noticeable drawbacks. Firstly, it necessitates the appropriate selection of the robust loss function, which may entail additional effort to understand the patterns of outliers. Secondly, the iterative

Algorithm 1 Procedure of Computing Attention Vector of Transformer-RKDE/SPKDE/MoM

1: Input : Q = { q i } i ∈ [ N ] , K = { k j } j ∈ [ N ] , V = { v l } l ∈ [ N ] , initial weights ω (0)

2: Normalize K = { k j } j ∈ [ N ] along the head dimension.

3: Compute kernel function between each pair of sequence: k σ ( Q , K ) = { k σ ( q i -k j ) } i,j ∈ [ N ] .

4: (Optional) apply attention mask on k σ ( Q , K ) .

5: [MoM] Randomly sample B subsets I 1 , . . . , I B of size S , obtain the median block I l such that 1 S ∑ j ∈ I l k σ ( q i -k j ) = median { 1 S ∑ j ∈ I 1 k σ ( q i -k j ) , . . . , 1 S ∑ j ∈ I B k σ ( q i -k j ) }

6: [RKDE] Update weights ω (0) for marginal/joint density by ω (1) j = ψ ( ∥ ∥ ∥ k σ ( k j , · ) -ˆ p ( k ) robust ( k ) ∥ ∥ ∥ H kσ ) ∑ j ∈ [ N ] ψ ( ∥ ∥ ∥ k σ ( k j , · ) -ˆ p ( k ) robust ( k ) ∥ ∥ ∥ H kσ ) .

7: [SPKDE] Obtain optimal weights ω for marginal/joint density via solving equation (7).

8: [RKDE, SPKDE] Obtain robust self-attention vector

9: [MoM] Obtain attention vector ̂ h i = ∑ j ∈ I l v j k σ ( q i -k j ) ∑ j ∈ I l k σ ( q i -k j ) .

updates might not successfully converge to the optimal solution. A better alternative is to assign higher weights to high-density regions and reduce the weights for atypical samples. The original KDEis scaled and projected to its nearest weighted KDE according to the L 2 norm. Similar concepts have been studied by Scaled and Projected KDE (SPKDE) (Vandermeulen & Scott, 2014), which offer an improved set of weights that better defend against outliers in the RKHS space. Specifically, given the scaling factor β > 1 , and let C N σ be the convex hull of k σ ( x 1 , · ) , . . . , k σ ( x N , · ) ∈ H k σ , i.e., the space of weighted KDEs, the optimal density ˆ p robust is given by

which is guaranteed to have a unique minimizer since we are projecting in a Hilbert space and C N σ is closed and convex. Notice that, by definition, ˆ p robust can also be represented as ˆ p robust = ∑ j ∈ [ N ] ω j k σ ( x j , · ) , ω ∈ ∆ N , which is same as the formulation in Eq. (5). Then Eq. (6) can be written as a quadratic programming (QP) problem over ω :

where G is the Gram matrix of { x j } j ∈ [ N ] with k σ and q = G 1 β N . Since k σ is a positive-definite kernel and each x i is unique, the Gram matrix G is also positive-definite. As a result, this QP problem is convex, and we can leverage commonly used solvers to efficiently obtain the solution and the optimal density ˆ p robust .

Robust Self-Attention Mechanism We now introduce the robust self-attention mechanism that down-weights atypical samples. We consider the density estimator of the joint distribution and the marginal distribution when using isotropic Gaussian kernel:

Following the non-parametric regression formulation of self-attention in Eq. (3), we obtain the robust self-attention mechanism as

where ω joint and ω marginal are obtained via either alternative updates or the QP solver. We term Transformers whose density from the non-parametric regression formulation of self-attention employs Eq. (5) and Eq. (6) as Transformer-RKDE and Transformer-SPKDE, respectively. Figure 2 presents an example of the application of the attention weight factor during the learning process from image and text data. The derived weight factor can potentially emphasize elements relevant to the class while reducing the influence of detrimental ones. Note that, the computation of { ω marginal j } j ∈ [ N ]

̂ h i = ∑ j ∈ [ N ] v j ω joint j k σ ( q i -k j ) ∑ j ∈ [ N ] ω marginal j k σ ( q i -k j )

and { ω joint j } j ∈ [ N ] are separate as ω joint j involves both keys and values vectors. During the empirical evaluation, we concatenate the keys and values along the head dimension to obtain the weights for the joint density ˆ p robust ( v , k ) and only use the key vectors for obtaining the set of weights for the marginal ˆ p robust ( k ) . In addition, ω marginal , ω joint ∈ R j × i for i, j = 1 , . . . , N are 2-dimensional matrices that include the pairwise weights between each position of the sequence and the rest of the positions. The weights are initialized uniformly across a certain sequence length dimension. For experiments related to language modeling, we can leverage information from the attention mask to initialize the weights on the unmasked part of the sequence.

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 } N j =1 into B subsets I 1 , . . . , I B of equal size, namely, | I 1 | = | I 2 | = . . . = | I B | = S . Then, the robust estimator of p ( k ) takes the following form:

where we define ˆ p σ,I l ( k ) = 1 S ∑ j ∈ I l k σ ( k -k j ) for any l ∈ [ B ] . Similarly, the robust estimator of p ( v , k ) is as follows:

where ˆ p σ,I l ( v , k ) = 1 S ∑ j ∈ I l k σ ( v -v j ) k σ ( k -k j ) for any l ∈ [ B ] . We now propose the selfattention 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:

̸

where I l is the block such that ˆ p σ ( q i -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 ∩ I l 2 = ∅ for any 1 ≤ l 1 = l 2 ≤ 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.

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, TransformerMoM 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.

4 Experimental Results

In this section, we provide empirical validation of the benefits of integrating our proposed robust KDEattention 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 (Merity 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-

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 TransformerRKDE 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) 1 . 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

1

2

3

4

5

6

Purturbation Budget

10

20

30

40

50

60

Top-1 Accuracy

FGSM

ViT

DeiT

Fourier

RVT

F-MoM

RKDE

SPKDE

ViT-MoM

(a)

1

2

3

4

5

6

Purturbation Budget

0

10

20

30

40

Top-1 Accuracy

PGD

(b)

1

2

3

4

5

6

Purturbation Budget

20

30

40

50

Top-1 Accuracy

SPSA

(c)

approaches do not alter the model architecture, each model employs 5 . 7 M 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 . 2 M 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 ϵ = 1 / 255 under l ∞ norm, while the PGD attack uses 20 steps with a step size of α = 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-

EfficiencyandMemoryComparison

Methods

14000

RKDE-1-step

DeiT

RVT

12000

MoM

MaximumGPUMemory

Size

Tiny

10000

Small

Base

8000

6000

X

×

4000

2000

0.35

0.40

0.45

0.50

0.55

0.60

0.65

Averaged Computation Speed (it/sec)

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.

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.

5 Conclusion and Future Work

In this work, we explored the link between the dot-product self-attention mechanism and nonparametric 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.

6 Acknowledgment

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.

References

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.

Bhojanapalli, S., Chakrabarti, A., Glasner, D., Li, D., Unterthiner, T., and Veit, A. Understanding robustness of transformers for image classification. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pp. 10231-10241, 2021.

Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Cao, S. Choose a transformer: Fourier or galerkin. Advances in Neural Information Processing Systems , 34, 2021.

Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems , 34:15084-15097, 2021.

Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 , 2019.

Choromanski, K. M., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J. Q., Mohiuddin, A., Kaiser, L., Belanger, D. B., Colwell, L. J., and Weller, A. Rethinking attention with performers. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=Ua6zuk0WRH.

Dai, Z., Yang, Z., Yang, Y., Carbonell, J., Le, Q. V ., and Salakhutdinov, R. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint arXiv:1901.02860 , 2019.

Dehghani, M., Gouws, S., Vinyals, O., Uszkoreit, J., and Kaiser, L. Universal transformers. In International Conference on Learning Representations , 2019. URL https://openreview.net/forum? id=HyzdRiR9Y7.

Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pp. 248255. Ieee, 2009.

Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pp. 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/N19-1423. URL https://aclanthology.org/N19-1423.

Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 , 2020.

Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=YicbFdNTTy.

Fan, H., Xiong, B., Mangalam, K., Li, Y., Yan, Z., Malik, J., and Feichtenhofer, C. Multiscale vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pp. 6824-6835, 2021.

Fox, J. and Weisberg, S. Robust regression. An R and S-Plus companion to applied regression , 91, 2002.

Gabbur, P., Bilkhu, M., and Movellan, J. Probabilistic attention for interactive segmentation. Advances in Neural Information Processing Systems , 34, 2021.

Goodfellow, I. J., Shlens, J., and Szegedy, C. Explaining and harnessing adversarial examples. arXiv preprint arXiv:1412.6572 , 2014.

Hampel, F. R., Ronchetti, E. M., Rousseeuw, P., and Stahel, W. A. Robust statistics: the approach based on influence functions . Wiley-Interscience; New York, 1986.

Hendrycks, D. and Dietterich, T. Benchmarking neural network robustness to common corruptions and perturbations. arXiv preprint arXiv:1903.12261 , 2019.

Hendrycks, D., Basart, S., Mu, N., Kadavath, S., Wang, F., Dorundo, E., Desai, R., Zhu, T., Parajuli, S., Guo, M., et al. The many faces of robustness: A critical analysis of out-of-distribution generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pp. 8340-8349, 2021a.

Hendrycks, D., Zhao, K., Basart, S., Steinhardt, J., and Song, D. Natural adversarial examples. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 15262-15271, 2021b.

Huber, P. J. Robust estimation of a location parameter. In Breakthroughs in statistics , pp. 492-518. Springer, 1992.

Huber, P. J. Robust statistics. In International encyclopedia of statistical science , pp. 1248-1251. Springer, 2011.

Humbert, P., Le Bars, B., and Minvielle, L. Robust kernel density estimation with median-of-means principle. In International Conference on Machine Learning , pp. 9444-9465. PMLR, 2022.

Janner, M., Li, Q., and Levine, S. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems , 34:1273-1286, 2021.

Jerrum, M. R., Valiant, L. G., and Vazirani, V. V. Random generation of combinatorial structures from a uniform distribution. Theoretical computer science , 43:169-188, 1986.

Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are rnns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning , pp. 51565165. PMLR, 2020.

Khan, S., Naseer, M., Hayat, M., Zamir, S. W., Khan, F. S., and Shah, M. Transformers in vision: A survey. ACM Computing Surveys (CSUR) , 2021.

Kim, J. and Scott, C. Robust kernel density estimation. The Journal of Machine Learning Research , 13:2529-2565, 2012.

Kreuzer, D., Beaini, D., Hamilton, W., L´ etourneau, V., and Tossou, P. Rethinking graph transformers with spectral attention. Advances in Neural Information Processing Systems , 34, 2021.

Li, N., Liu, Y., Wu, Y., Liu, S., Zhao, S., and Liu, M. Robutrans: A robust transformer-based textto-speech model. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 34, pp. 8228-8235, 2020.

Lin, T., Wang, Y., Liu, X., and Qiu, X. A survey of transformers. arXiv preprint arXiv:2106.04554 , 2021.

Liu, J., Singhal, T., Blessing, L. T., Wood, K. L., and Lim, K. H. Crisisbert: a robust transformer for crisis classification and contextual crisis embedding. In Proceedings of the 32nd ACM Conference on Hypertext and Social Media , pp. 133-141, 2021a.

Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., and Stoyanov, V. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692 , 2019.

Liu, Z., Lin, Y ., Cao, Y ., Hu, H., Wei, Y ., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pp. 10012-10022, 2021b.

Liu, Z., Ning, J., Cao, Y., Wei, Y ., Zhang, Z., Lin, S., and Hu, H. Video swin transformer. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2022.

Lu, Y., Li, Z., He, D., Sun, Z., Dong, B., Qin, T., Wang, L., and Liu, T.-Y. Understanding and improving transformer from a multi-particle dynamic system point of view. arXiv preprint arXiv:1906.02762 , 2019.

Madry, A., Makelov, A., Schmidt, L., Tsipras, D., and Vladu, A. Towards deep learning models resistant to adversarial attacks. arXiv preprint arXiv:1706.06083 , 2017.

Mahmood, K., Mahmood, R., and Van Dijk, M. On the robustness of vision transformers to adversarial examples. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pp. 7838-7847, 2021.

Mao, X., Qi, G., Chen, Y., Li, X., Duan, R., Ye, S., He, Y., and Xue, H. Towards robust vision transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 12042-12051, 2022.

Merity, S., Xiong, C., Bradbury, J., and Socher, R. Pointer sentinel mixture models. arXiv preprint arXiv:1609.07843 , 2016.

Morris, J. X., Lifland, E., Yoo, J. Y., Grigsby, J., Jin, D., and Qi, Y. Textattack: A framework for adversarial attacks, data augmentation, and adversarial training in nlp. arXiv preprint arXiv:2005.05909 , 2020.

Nadaraya, E. A. On estimating regression. Theory of Probability & Its Applications , 9(1):141-142, 1964.

Nguyen, T., Nguyen, T., Do, H., Nguyen, K., Saragadam, V., Pham, M., Nguyen, K., Ho, N., and Osher, S. Improving transformer with an admixture of attention heads. In Advances in Neural Information Processing Systems , 2022a.

Nguyen, T., Nguyen, T., Le, D., Nguyen, K., Tran, A., Baraniuk, R., Ho, N., and Osher, S. Improving transformers with probabilistic attention keys. In International Conference on Machine Learning , 2022b.

Nguyen, T., Pham, M., Nguyen, T., Nguyen, K., Osher, S. J., and Ho, N. Fourierformer: Transformer meets generalized Fourier integral theorem. Advances in Neural Information Processing Systems , 2022c.

Parzen, E. On estimation of a probability density function and mode. The annals of mathematical statistics , 33(3):1065-1076, 1962.

Paul, S. and Chen, P.-Y. Vision transformers are robust learners. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 36, pp. 2071-2081, 2022.

Radford, A., Narasimhan, K., Salimans, T., and Sutskever, I. Improving language understanding by generative pre-training. OpenAI report , 2018.

Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning , pp. 8748-8763. PMLR, 2021.

Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research , 21(140):1-67, 2020. URL http://jmlr.org/papers/v21/20-074.html.

Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Radford, A., Chen, M., and Sutskever, I. Zeroshot text-to-image generation. In International Conference on Machine Learning , pp. 8821-8831. PMLR, 2021.

Reddy, J. An introduction to the finite element method , volume 1221. McGraw-Hill New York, 2004.

Rosenblatt, M. Remarks on some nonparametric estimates of a density function. The annals of mathematical statistics , pp. 832-837, 1956.

Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M., et al. Imagenet large scale visual recognition challenge. International journal of computer vision , 115(3):211-252, 2015.

Sander, M. E., Ablin, P., Blondel, M., and Peyr´ e, G. Sinkformers: Transformers with doubly stochastic attention. In International Conference on Artificial Intelligence and Statistics , pp. 3515-3530. PMLR, 2022.

Schlag, I., Irie, K., and Schmidhuber, J. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning , pp. 9355-9366. PMLR, 2021.

Subramanya, A., Saha, A., Koohpayegani, S. A., Tejankar, A., and Pirsiavash, H. Backdoor attacks on vision transformers. arXiv preprint arXiv:2206.08477 , 2022.

Tang, B. and Matteson, D. S. Probabilistic transformer for time series analysis. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems , 2021. URL https://openreview.net/forum?id=HfpNVDg3ExA.

Tay, Y., Dehghani, M., Bahri, D., and Metzler, D. Efficient transformers: A survey. arXiv preprint arXiv:2009.06732 , 2020.

Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., and J´ egou, H. Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning , pp. 10347-10357. PMLR, 2021a.

Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., and J´ egou, H. Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning , pp. 10347-10357. PMLR, 2021b.

Tsai, Y.-H. H., Bai, S., Yamada, M., Morency, L.-P., and Salakhutdinov, R. Transformer dissection: An unified understanding for transformer's attention via the lens of kernel. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) , pp. 4344-4353, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v1/D19-1443. URL https://aclanthology.org/D19-1443.

Uesato, J., O'donoghue, B., Kohli, P., and Oord, A. Adversarial risk and the dangers of evaluating against weak attacks. In International Conference on Machine Learning , pp. 5025-5034. PMLR, 2018.

Vandermeulen, R. A. and Scott, C. Robust kernel density estimation by scaling and projection in hilbert space. Advances in Neural Information Processing Systems , 27, 2014.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. In Advances in neural information processing systems , pp. 5998-6008, 2017a.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polosukhin, I. Attention is all you need. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017b. URL https://proceedings.neurips.cc/paper/ 2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.

Wang, H., Ge, S., Lipton, Z., and Xing, E. P. Learning robust global representations by penalizing local predictive power. Advances in Neural Information Processing Systems , 32, 2019.

Wang, S., Li, B., Khabsa, M., Fang, H., and Ma, H. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768 , 2020.

Wang, X., Girshick, R., Gupta, A., and He, K. Non-local neural networks. In Proceedings of the IEEE conference on computer vision and pattern recognition , pp. 7794-7803, 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 X = [ x 1 , . . . , x N ] ⊤ ∈ R N × D x of N feature vectors, the self-attention mechanism transforms it into another sequence H := [ h 1 , · · · , h N ] ⊤ ∈ R N × D v as follows:

The vectors q i , k j and v j are the query, key and value vectors, respectively. They are computed as follows:

where W Q , W K ∈ R D × D x , W V ∈ R D v × D x are the weight matrices. Equation (13) can be written in the following equivalent matrix form:

where the softmax function is applied to each row of the matrix ( QK ⊤ ) / D . Equation (15) is also called the 'softmax attention'. Assume we have the key and value vectors { k j , v j } j ∈ [ N ] that is collected from the data generating process

where ε is some noise vectors with E [ ε ] = 0 , and f is the function that we want to estimate. If { k j } j ∈ [ N ] are i.i.d. samples from the distribution p ( k ) , and p ( v , k ) is the joint distribution of ( v , k ) defined by equation (16), we have

We need to obtain estimations for both the joint density function p ( v , k ) and the marginal density function p ( k ) to obtain function f , one popular approach is the kernel density estimation:

where [ v , k ] denotes the concatenation of v and k . k σ could be isotropic Gaussian kernel: k σ ( x -x ′ ) = exp ( -∥ x -x ′ ∥ 2 / (2 σ 2 ) ) , we have

Combining equations (19), (20), and (17), we obtain the NW estimator of the function f as

Now we show how the self-attention mechanism is related to the NW estimator. If the keys { k j } j ∈ [ N ] are normalized

Then estimating the softmax attention is equivalent to estimating ̂ f σ ( q ) .

B Details on Leveraging Robust KDE on Transformers

For simplicity, we use the Huber loss function as the demonstrating example, which is defined as follows:

where a is a constant. The solution to this robust regression problem has the following form:

Proposition 1. Assume the robust loss function ρ is non-decreasing in [0 , ∞ ] , ρ (0) = 0 and lim x → 0 ρ ( x ) x = 0 . Define ψ ( x ) := ρ ′ ( x ) x and assume ψ (0) = lim x → 0 ρ ′ ( x ) x exists and finite. Then the optimal ˆ p robust can be written as

where ω = ( ω 1 , · · · , ω N ) ∈ ∆ N , with each ω j ∝ ψ ( ∥ k σ ( x j , · ) -ˆ p robust ∥ H kσ ) . Here ∆ n denotes the n -dimensional probability simplex.

Proof. The proof of Proposition 1 is mainly adapted from the proof in Kim & Scott (2012). Here, we provide proof of completeness. For any p ∈ H k σ , we denote

Then we have the following lemma regarding the Gateaux differential of J and a necessary condition for ˆ p robust to be optimal solution of the robust loss objective function in equation (5).

Lemma 1. Given the assumptions on the robust loss function ρ in Proposition 1, the Gateaux differential of J at p ∈ H k σ with incremental h ∈ H k σ , defined as δJ ( p ; h ) , is

where the function V : H k σ →H k σ is defined as:

A necessary condition for ˆ p robust is V (ˆ p robust ) = 0 .

The proof of Lemma 1 can be found in Lemma 1 of Kim & Scott (2012). Based on the necessary condition for ˆ p robust in Lemma 1, i.e., V (ˆ p robust ) = 0 , we have

Direct algebra indicates that ˆ p robust = ∑ j ∈ [ N ] ω j k σ ( x j , · ) where ω = ( ω 1 , · · · , ω N ) ∈ ∆ N , and ω j ∝ ψ ( ∥ k σ ( x j , · ) -ˆ p robust ∥ H kσ ) . As a consequence, we obtain the conclusion of the proposition.

For the Huber loss function, we have that

Hence, when the error ∥ k σ ( x j , · ) , · -ˆ p robust ∥ H kσ is over the threshold a , the final estimator will down-weight the importance of k σ ( x j , · ) . This is in sharp contrast with the standard KDE method, which will assign uniform weights to all of the k σ ( x j , · ) . As we mentioned in the main paper, the estimator provided in Proposition 1 is circularly defined, as ˆ p robust is defined via ω , and ω depends on ˆ p robust. Such an issue can be addressed by estimating ω with an iterative algorithm termed as kernelized iteratively re-weighted least-squares (KIRWLS). The algorithm starts with randomly initialized ω (0) ∈ ∆ n , and perform the following iterative updates between two steps:

Note that, the optimal ˆ p robust is the fixed point of this iterative update, and the KIRWLS algorithm converges under standard regularity conditions. Furthermore, one can directly compute the term ∥ ∥ ∥ k σ ( x j , · ) -ˆ p ( k ) robust ∥ ∥ ∥ H kσ via the reproducing property:

Therefore, the weights can be updated without mapping the data to the Hilbert space.

C Fourier Attention with Median of Means

We introduce the Fourier Attention coupled with the Median of Means (MoM) principle and show how this is robust to outliers. For any given function ϕ : R → R and radius R , we randomly divide the keys { k i } i ∈ [ N ] into B subsets I 1 , . . . , I B of equal size where | I 1 | = | I 2 | = · · · = | I B | = S . Define ˆ p R,I m ( q l ) = 1 S ∑ i ∈ I m ∏ D j =1 ϕ ( sin( R ( q lj -k ij )) R ( q lj -k ij ) ) , then the MoM Fourier attention is defined as

where I m is the block such that ˆ p R ( q l , k ) achieves its median value. To shed light into the robustness of Transformers that use Eq. (27) as the attention mechanism, we demonstrate that the estimator ˆ p R ( q ) = median { ˆ p R,I 1 ( q ) , . . . , ˆ p R,I B ( q ) } is a robust estimator of the density function p ( q ) of the keys. We first introduce a few notations that are useful for stating this result. Denote C = { 1 ≤ i ≤ N : k i is clean } and O = { 1 ≤ i ≤ N : k i is outlier } . Then, we have C ∩ O = ∅ and C ∪ O = { 1 , 2 , . . . , N } . The following result establishes a high probability upper bound on the sup-norm between ̂ p R ( q ) and p ( q ) .

Theorem 1. Assume that the function ϕ satisfies ∫ ϕ (sin( z ) /z ) z j dz = 0 for all 1 ≤ j ≤ m and ∫ | ϕ (sin( z ) /z ) || z | m +1 dz < ∞ for some m ∈ N . Furthermore, the density function p ( q ) satisfies sup q | p ( q ) | < ∞ . The number of blocks B and the number of outliers |O| are such that B > (2 + δ ) |O| where δ is the failure probability. Then, with ∆ = 1 2+ δ -|O| B for the radius R sufficiently large and δ sufficiently small, with probability at least 1 -exp( -2∆ 2 B ) we find that

where C is some universal constant.

Remark 1 . The result of Theorem 1 indicates by choosing R = O ( N -1 2( m +1)+ D ) , the rate of ˆ p R to p under the supremum norm is O ( N -m +1 2( m +1)+ D ) . With that choice of R , when N approaches infinity, the MoM estimator ˆ p R is a consistent estimator of the clean distribution p of the keys. This confirms the validity of using ˆ p R to robustify p and similarly the usage of MoM Fourier attention Eq. (27) as a robust attention for Transformers.

Proof. From the formulation of the MoM estimator ̂ p R ( q ) , we obtain the following inequality

This bound indicates that to bound P ( ∥ ˆ p R ( q ) -p ( q ) ∥ ∞ ≥ ϵ ) , it is sufficient to bound P ( { sup q ∑ B b =1 1 {| ˆ p R,I b ( q ) -p ( q ) |≥ ϵ } ≥ B 2 } ) . Indeed, for each 1 ≤ b ≤ B , we find that

Therefore, we have

which leads to sup q ∑ B b =1 1 {| ˆ p R,I b ( q ) -p ( q ) |≥ ϵ } ≤ ∑ B b =1 1 { sup q {| ˆ p R,I b ( q ) -p ( q ) |≥ ϵ } . This inequality shows that

To ease the presentation, we denote W b = 1 { sup q {| ˆ p R,I b ( q ) -p ( q ) |≥ ϵ } and B = { 1 ≤ b ≤ B : I b ∩ O = ∅} . Then, the following inequalities hold

where we assume without loss of generality that 1 ∈ B , which is possible due to the assumption that B > (2 + δ ) |O| . We now prove the following uniform concentration bound:

Lemma 2. Assume that ϕ ( z ) ≤ C for all | z | ≤ 1 for some universal constant C . We have

Proof of Lemma 2. By the triangle inequality, we have

To bound sup q | ˆ p R,I 1 ( q ) -E [ˆ p R,I 1 ( q )] | , we use Bernstein's inequality along with the bracketing entropy under L 1 norm in the space of queries q . In particular, we denote Y i = R D A D ∏ D j =1 ϕ ( sin( R ( q j -k ij )) R ( q j -k ij ) ) where A = ∫ R ϕ ( sin( z ) z ) dz for all i ∈ I 1 . Since sin( R ( q j -k ij )) / ( R ( q j -k ij )) ≤ 1 for all 1 ≤ j ≤ D , we obtain that | Y i | ≤ C D R D /A D where C is the constant such that ϕ ( z ) ≤ C when | z | ≤ 1 . Furthermore, E [ | Y i | ]

By choose ϵ = C ( 1 R m +1 + √ R D log R log(2 /δ ) | I 1 | )) , then we find that

Collecting the above inequalities leads to

where ∆ = 1 2+ δ -|O| B . As a consequence, we obtain the conclusion of the theorem.

D Dataset Information

WikiText-103 The dataset 1 contains around 268 K words and its training set consists of about 28 K articles with 103 M tokens, this corresponds to text blocks of about 3600 words. The validation set and test sets consist of 60 articles with 218 K and 246 K tokens respectively.

ImageNet We use the full ImageNet dataset that contains 1 . 28 M training images and 50 K validation images. The model learns to predict the class of the input image among 1000 categories. We report the top-1 and top-5 accuracy on all experiments. The following ImageNet variants are test sets that are used to evaluate model performance.

ImageNet-C For robustness on common image corruptions, we use ImageNet-C (Hendrycks & Dietterich, 2019) which consists of 15 types of algorithmically generated corruptions with five levels of severity. ImageNet-C uses the mean corruption error (mCE) as a metric: the smaller mCE means the more robust the model under corruption.

ImageNet-A This dataset contains real-world adversarially filtered images that fool current ImageNet classifiers. A 200-class subset of the original ImageNet-1K's 1000 classes is selected so that errors among these 200 classes would be considered egregious, which cover most broad categories spanned by ImageNet-1K.

ImageNet-O This dataset contains adversarially filtered examples for ImageNet out-ofdistribution detectors. The dataset contains samples from ImageNet-22K but not from ImageNet1K, where samples that are wrongly classified as an ImageNet-1K class with high confidence by a ResNet-50 are selected. We use AUPR (area under precision-recall) as the evaluation metric.

ImageNet-R This dataset contains various artistic renditions of object classes from the original ImageNet dataset, which is discouraged by the original ImageNet. ImageNet-R contains 30,000 image renditions for 200 ImageNet classes, where a subset of the ImageNet-1K classes is chosen.

ImageNet-Sketch This dataset contains 50,000 images, 50 images for each of the 1000 ImageNet classes. The dataset is constructed with Google Image queries 'sketch of xxx', where xxx is the standard class name. The search is only performed within the 'black and white' color scheme.

E Ablation Studies

In this section, we provide additional results and ablation studies that focus on different design choices for the proposed robust KDE attention mechanisms. The detailed experimental settings can be found in the caption of each table.

(a)

(b)

(c)

(d)

-2

0

0.060

15

0.060

-Z

Z

00

0

D

-2

0.030

0.015

003

-2

Q

O

O

0200

00

0.010

-6

-2

2

Z-

0

Z-

0

-6

-2

0

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this work the authors present novel adaptation of a nonstandard self attention module, like those used in transformers. Previous works have demonstrated that the weighting matrix in self attention can be evaluated using a positive kernel rather than the inner product, which ends up looking like a kernel regressor, which itself can be understood through kernel density estimation. In this work the authors propose using robust versions KDEs in place of KDEs for this method. The authors introduce and describe 3 existing versions of robust KDEs for this and apply their method to several experimental settings on a few data types. \n\nOn image datasets the authors find that, while their method does not improve the standard performance on clean datasets, the methods improve performance when robustness is advantageous, eg. adversarial samples, dataset shift, etc.. On time series and NLP datasets the authors observe some improvement on the clean data as well as unclean data.\n\nSTRENGTHS:\nThe core contribution of this paper is providing a new deep learning methodology that yields substantial performance improvements. The method seems to work on several datatypes and robustness scenarios indicating a fundamental improvement to the general transformer methodology, rather than something that is very application/task/architecture specific. The results also give some insight as to origins of non-robustness of transformer architecture. Its not obvious that key attention values being \"heavy-tailed\" may make a transformer sensitive to word swapping.\n\nAn additional strength of this paper is its novelty. This work is brings methods far outside of what one usually sees in deep learning research, which is also a plus. \n\nOverall the paper is fairly clearly written with a nice overview of a few robust KDEs and introduction to using KDEs for transformers.\n\nWEAKNESSES:\nThe primary weakness, in my opinion, is that the proposed method is simply swapping out a component of a transformer with a existing robust version of that component, and is somewhat trivial.\n\nQUESTIONS:\nThe omega values in the robust KDEs depend on the values k_i (they are functions of $k_i$). During training do the gradients take into account how these omega values will be modified as the functions for k_i are modified? Its a bit opaque to me how the robustness plays into the optimization.\n\n Could one just train using the standard KDE and then at test time use a robust KDE?\n\nLIMITATIONS:\nNone", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper re-interpret the self-attention mechanism in transformer architectures as a non-parametric kernel density estimator (KDE). Leveraging existing robustness results on classical KDE, the authors develop new variants of transformers; empirical results on a variety of datasets demonstrate how these variants are resistant to adversarial attacks and data contamination.\n\nSTRENGTHS:\nThis paper provides a simple yet powerful re-interpretation of the self-attention mechanism through the lens of KDE. This re-interpretation allows to leverage results from robust KDE to design better transformer architectures. Kernel methods are notably amenable to theoretical analysis; therefore, links between transformers and kernels as the one proposed in this paper can help sheding lights on justify the empirical successess of transformer architectures. The paper is in general well-written and its structure is easy to follow.\n\nWEAKNESSES:\nLinks between Transformers and kernel methods have been proposed in prior literature, see for example [1]. Although an extensive literature review is proposed by the authors, some important references such as [1] are omitted. \n\n**References**\n[1] Tsai, Yao-Hung Hubert, et al. \"Transformer dissection: a unified understanding of transformer's attention via the lens of kernel.\" arXiv preprint arXiv:1908.11775 (2019).\n\nQUESTIONS:\nIn my view, the paper makes a strong contribution. Could the authors please address my questions below.\n\n- In line 93 when the kernel $k_{\\sigma}$ is introduced, it should be specified what set does this kernel act on and the fact that the considered class of kernels is translation invariant $\\tilde k_\\sigma(x,y) = k_\\sigma(x-y)$. For example, equation (2) is ambiguous because the kernels used in the joint density and marginal density estimations are different (indexed on different spaces).\n- In line 95-96, it would be good to add an expression to define the mentioned isotropic Gaussian kernel.\n- In lines 186-187 it is mentioned that the contruction of attention weight factors necessitate iterative algorithms to calculate the set of weights when computing self-attention at each layer; however, if (sometimes) the weight estimation boils down to a finite dimensional quadratic optimisation problem with linear constraints such as equation (7), can't the weights be determined by simply solving a linear system of equations?\n- Kernel regression often involves regularisation; this setup is not considered in the paper. Can the authors comment on the potential role of regularisation?\n\nLIMITATIONS:\nThe authors have discussed limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper addresses robustness of transformer architecture in both natural language and vision contexts.\nMotivated by classical nonparametric estimation problem, the authors propose a procedure to robustly estimate the self-attention values. While not loosing too much in performance on the original data, the method obtains improvements on corrupted data, with several benchmarks considered. The paper proposes a number of methods to address an important problem of robustness to adversarial contaminations and outliers.\n\nPros:\n- sound experimental results. the proposed method beats sota on a number of benchmarks.\n- ablation study is included in the appendix\n- comparison of time\n\nCons:\n- I did not understand from reading the introduction whether this method is a post-hoc modification or requires in-processing training\n- the theoretical contribution is questionable, I have some questions for Theorem 1 in the appendix, see below\n- lack of originality, the strongest part of the paper is experimental results.\n\nminor comments:\n- line 84: it is best to refer to Nguyen et al to motivate using this form of regression\n- line 84: usually conditional mean is equal to zero E[eps | k] = 0\n- line 86: I suggest to make a precise statement that the derivations only hold for Gaussian kernel earlier.\n- line 256 \"We follow...\" correct the grammar in this sentence\n- theorem 1 in appendix: the sentence \"by adapting Lemma 1 to uniform case\". How exactly do you adapt this to uniform case without any loss in the upper bound? I don't think it's that easy to turn a pointwise bound into a uniform one, so I reckon you do not have a proof of the claimed theorem.\n\n\nQuestions:\n\n- do the methods require fine-tunning the parameters?\n\n- how do you choose hyper-parameters, such as number of bins in MoM, cut-off parameter in Huber loss, etc?\n\n- it is completely unclear from the paper whether robust modification only concerns with inference time, or it changes the training as well?\n\n- robust estimators such as MoM are typically robust towards a corruptions of finite amount of data points. Does it mean that you can be robust to one pixel modifications? Or even one 16x16 patch modifications?\n\n- in the experiments, does the hyperparameter choice procedure require the access to contaminated data?\n\nSTRENGTHS:\n-\n\nWEAKNESSES:\n-\n\nQUESTIONS:\n-\n\nLIMITATIONS:\n-", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

comments: SUMMARY:
This paper introduces fortified transformers that employ robust KDE as an alternative to dot-product attention, self-attention, mitigating the impacts of attacked images. Especially, this paper proposed two variants of robust self-attention mechanisms based on kernel density estimators (KDE). Extensive experimental results demonstrate the effectiveness of the proposed approach on diverse tasks.

STRENGTHS:
- This paper is well-written.
- This paper analyzes the characteristics of the proposed methods, especially their limitation.
- Extensive experiments are conducted on diverse tasks with multiple modalities.

WEAKNESSES:
- The paper offers a comprehensive study of several methods, but a dominant strategy does not emerge as the results seem to depend largely on the task at hand, as evidenced in the data tables provided.
- Although the stated limitations are acknowledged, a more in-depth analysis of these shortcomings would enhance the persuasive power. 
   - It is acknowledged that RKDE and SPKDE possess computational inefficiencies. Then, how computationally expensive they are than the original ones which are not robust? It would be better to provide the experimental results.
   - Suboptimality is a possible limitation of MoM, then how much the extent of MoM's accessibility of MoM is restricted compared to the originals? How those restricted accessibility affects the performances?

QUESTIONS:
N/A

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

comments: SUMMARY:
The paper describes and leverages the connection between the self-attention mechanism and KDE estimation to motivate replacing dot-product attention with a robust KDE approach to enable more robust transformer models. The experiments illustrate that the proposed approach can maintain performance on clean data while improving performance on contaminated data.

STRENGTHS:
* Given the wide use of transformer models, the paper tackles an important topic, as robustness to training data contamination can improve their performance.
* The proposed approach is well motivated by the connection to non-parametric regression.
* The experiments show very promising results on both clean data and contaminated data, across 
* The proposed approach does not introduce additional parameters.

WEAKNESSES:
Practicality of the approach for large datasets: the proposed approach incurs additional computational complexity even with the proposed approximations. The experiments are performed on small versions of the models. It's unclear how well this can be applied to large scale training of transformer models.

QUESTIONS:
N/A

LIMITATIONS:
The authors accurately describe limitations and potential solutions throughout the paper and in the Conclusion and Future Work section.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-11