SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xing Han
Department of ECE  
University of Texas at Austin  
aaronhan223@utexas.edu
Tan Minh Nguyen  
Department of Mathematics  
University of California, Los Angeles  
tanmnguyen89@ucla.edu
Joydeep Ghosh  
Department of ECE  
University of Texas at Austin  
jghosh@utexas.edu
Tongzheng Ren  
Department of Computer Science  
University of Texas at Austin  
tongzheng@utexas.edu
Khai Nguyen  
Department of Statistics and Data Sciences  
University of Texas at Austin  
khainb@utexas.edu
Nhat Ho  
Department of Statistics and Data Sciences  
University of Texas at Austin  
minhnhat@utexas.edu
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.
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-the-art 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 audiios. 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  $\{\pmb{k}_j, \mathbf{v}_j\}_{j \in [N]}$  that is collected from the data generating process  $\mathbf{v} = f(\pmb{k}) + \varepsilon$ , where  $\varepsilon$  is some noise vectors with  $\mathbb{E}[\varepsilon] = 0$ , and  $f$  is the function that we want to estimate. We consider a random design setting where the key vectors  $\{\pmb{k}_j\}_{j \in [N]}$  are i.i.d. samples from the distribution  $p(\pmb{k})$ , and we use  $p(\mathbf{v}, \pmb{k})$  to denote the joint distribution of  $(\mathbf{v}, \pmb{k})$  defined by the data generating process. Our target is to estimate  $f(\pmb{q})$  for any new queries  $\pmb{q}$ . The NW estimator provides a non-parametric approach to estimating the function  $f$ , the main idea is that
$$
f (\boldsymbol {k}) = \mathbb {E} [ \mathbf {v} | \boldsymbol {k} ] = \int_ {\mathbb {R} ^ {D}} \mathbf {v} \cdot p (\mathbf {v} | \boldsymbol {k}) d \mathbf {v} = \int_ {\mathbb {R} ^ {D}} \frac {\mathbf {v} \cdot p (\mathbf {v} , \boldsymbol {k})}{p (\boldsymbol {k})} d \mathbf {v}, \tag {1}
$$
where the first equation comes from the fact that  $\mathbb{E}[\varepsilon] = 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(\mathbf{v}, \mathbf{k})$  and the marginal density function  $p(\mathbf{k})$ . KDE is commonly used for the density estimation problem (Rosenblatt, 1956; Parzen, 1962), which requires a kernel  $k_{\sigma}$  with the bandwidth parameter  $\sigma$  satisfies  $\int_{\mathbb{R}^D} k_{\sigma}(\mathbf{x} - \mathbf{x}') dx = 1, \forall \mathbf{x}'$ , and estimate the density as
$$
\hat {p} _ {\sigma} (\mathbf {v}, \boldsymbol {k}) = \frac {1}{N} \sum_ {j \in [ N ]} k _ {\sigma} ([ \mathbf {v}, \boldsymbol {k} ] - [ \mathbf {v} _ {j}, \boldsymbol {k} _ {j} ]) \quad \hat {p} _ {\sigma} (\boldsymbol {k}) = \frac {1}{N} \sum_ {j \in [ N ]} k _ {\sigma} (\boldsymbol {k} - \boldsymbol {k} _ {j}), \tag {2}
$$
where  $[\mathbf{v},\pmb{k}]$  denotes the concatenation of  $\mathbf{v}$  and  $\pmb{k}$ . Specifically, when  $k_{\sigma}$  is the isotropic Gaussian kernel, we have  $\hat{p}_{\sigma}(\mathbf{v},\pmb{k}) = \frac{1}{N}\sum_{j\in [N]}k_{\sigma}(\mathbf{v} - \mathbf{v}_j)k_{\sigma}(\pmb{k} - \pmb{k}_j)$ . Combine this with Eq. (1) and Eq. (2), we can obtain the NW estimator of the function  $f$  as
$$
\widehat {f} _ {\sigma} (\boldsymbol {k}) = \frac {\sum_ {j \in [ N ]} \mathbf {v} _ {j} k _ {\sigma} (\boldsymbol {k} - \boldsymbol {k} _ {j})}{\sum_ {j \in [ N ]} k _ {\sigma} (\boldsymbol {k} - \boldsymbol {k} _ {j})}. \tag {3}
$$
Furthermore, it is not hard to show that if the keys  $\{\pmb{k}_j\}_{j\in [N]}$  are normalized, the self-attention mechanism  $\widehat{f}_{\sigma}(\pmb{q}_i)$  in Eq. (3) is exactly the standard Softmax attention  $\widehat{f}_{\sigma}(\pmb{q}_i) = \sum_{j\in [N]} \text{softmax}(\pmb{q}^\top \pmb{k}_j / \sigma^2) \mathbf{v}_j$ . Such an assumption on the normalized key  $\{\pmb{k}_j\}_{j\in [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  $\sigma^2 = \sqrt{D}$ , where  $D$  is the dimension of  $\pmb{q}$  and  $\pmb{k}_j$ , then  $\widehat{f}_{\sigma}(\pmb{q}_i) = \pmb{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  $\mathcal{H}_k = \{f\mid f:\mathcal{X}\to \mathbb{R}\}$  is the RKHS associated with kernel  $k$  , where  $k:\mathcal{X}\times \mathcal{X}\rightarrow \mathbb{R},$  if it is a Hilbert space with inner product  $\langle \cdot ,\cdot \rangle_{\mathcal{H}_k}$  and following properties:
-  $k(\pmb{x},\cdot)\in \mathcal{H}_k,\forall \pmb{x}\in \mathcal{X};$
-  $\forall f\in \mathcal{H}_k,f(\pmb {x}) = \langle f,k(\pmb {x},\cdot)\rangle_{\mathcal{H}_k}$  .Aka the reproducing property.
With slightly abuse of notation, we define  $k_{\sigma}(\boldsymbol{x}, \boldsymbol{x}') = k_{\sigma}(\boldsymbol{x} - \boldsymbol{x}')$ . By the definition of the RKHS and the KDE estimator, we know  $\hat{p}_{\sigma} = \frac{1}{N}\sum_{j \in [N]} k_{\sigma}(\boldsymbol{x}_j, \cdot) \in \mathcal{H}_{k_{\sigma}}$ , and can be viewed as the optimal solution of the following least-square regression problem in RKHS:
Figure 1: The contour plots illustrate the density estimation of the two-dimensional query vector embedding within a transformer's attention layer. The left plot employs the regular KDE method, as defined in Eq. (4), whereas the right plot utilizes a robustified version of the KDE method, which enhances KDE's robustness against outliers.
$$
\hat {p} _ {\sigma} = \underset {p \in \mathcal {H} _ {k _ {\sigma}}} {\arg \min } \sum_ {j \in [ N ]} \frac {1}{N} \| k _ {\sigma} (\boldsymbol {x} _ {j}, \cdot) - p \| _ {\mathcal {H} _ {k _ {\sigma}}} ^ {2}. \tag {4}
$$
Figure 2: The application of Transformers with robust KDE attention on image and text is shown. (Left) The robust KDE self-attention generates varying weight factors for image patch embeddings under adversarial attacks or data corruption. The adversely impacted regions that would otherwise lead to incorrect predictions are down-weighted, ensuring enhanced accuracy and robustness. (Right) In the field of language modeling, the weight factors lend significance to essential keywords (highlighted in red). In the face of word swap attacks, the fortified self-attention mechanism, particularly when utilizing the medians-of-means principle, is proficient in disregarding or reducing the importance of less consequential words (marked in green). Consequently, this results in a more resilient procedure during self-attention computations.
Note that, in Eq. (4), the same weight factor  $1 / N$  is applied uniformly to each error term  $\| k_{\sigma}(\pmb{x}_j,\cdot) - p\|_{\mathcal{H}_{k_\sigma}}^2$ . This approach functions effectively if there are no outliers in the set  $\{k_{\sigma}(\pmb{x}_j,\cdot)\}_{j\in [N]}$ . However, when outliers are present (for instance, when there is some  $j$  such that  $\| k_{\sigma}(\pmb{x}_j,\cdot)\|_{\mathcal{H}_{k_\sigma}}\gg \| k_{\sigma}(\pmb{x}_i,\cdot)\|_{\mathcal{H}_{k_\sigma}},\forall i\in [N],i\neq 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  $\rho$ , resulting in the following formulation:
$$
\hat {p} _ {\text {r o b u s t}} = \arg \min  _ {p \in \mathcal {H} _ {k _ {\sigma}}} \sum_ {j \in [ N ]} \rho \left(\| k _ {\sigma} \left(\boldsymbol {x} _ {j}, \cdot\right) - p \| _ {\mathcal {H} _ {k _ {\sigma}}}\right) = \sum_ {j \in [ N ]} \omega_ {j} k _ {\sigma} \left(\boldsymbol {x} _ {j}, \cdot\right). \tag {5}
$$
Examples of the robust loss function  $\rho$  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  $\omega = (\omega_{1},\dots ,\omega_{N})\in \Delta_{N}$ , with each  $\omega_{j}\propto \psi \left(\| k_{\sigma}(\pmb{x}_{j},\cdot) - \hat{p}_{\mathrm{robust}}\|_{\mathcal{H}_{k_{\sigma}}}\right)$ , where  $\psi (x)\coloneqq \frac{\rho'(x)}{x}$ . Since  $\hat{p}_{\mathrm{robust}}$  is defined via  $\omega$ , and  $\omega$  also depends on  $\hat{p}_{\mathrm{robust}}$ , one can address this circular definition problem via an alternative updating algorithm proposed by Kim & Scott (2012). The algorithm starts with randomly initialized  $\omega^{(0)}\in \Delta_n$ , and performs alternative updates between  $\hat{p}_{\mathrm{robust}}$  and  $\omega$  until the optimal  $\hat{p}_{\mathrm{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:  $\mathbf{Q} = \{\pmb {q}_i\}_{i\in [N]}$ $\mathbf{K} = \{\pmb {k}_j\}_{j\in [N]}$ $\mathbf{V} = \{\mathbf{v}_l\}_{l\in [N]}$ , initial weights  $\omega^{(0)}$
2:Normalize  $\mathbf{K} = \{\pmb {k}_j\}_{j\in [N]}$  along the head dimension.
3: Compute kernel function between each pair of sequence:  $k_{\sigma}(\mathbf{Q}, \mathbf{K}) = \{k_{\sigma}(\pmb{q}_i - \pmb{k}_j)\}_{i,j \in [N]}$ .
4: (Optional) apply attention mask on  $k_{\sigma}(\mathbf{Q},\mathbf{K})$
5: [MoM] Randomly sample  $B$  subsets  $I_1, \ldots, I_B$  of size  $S$ , obtain the median block  $I_l$  such that  $\frac{1}{S} \sum_{j \in I_l} k_{\sigma}(\pmb{q}_i - \pmb{k}_j) = \text{median}\{\frac{1}{S} \sum_{j \in I_1} k_{\sigma}(\pmb{q}_i - \pmb{k}_j), \ldots, \frac{1}{S} \sum_{j \in I_B} k_{\sigma}(\pmb{q}_i - \pmb{k}_j)\}$
6: [RKDE] Update weights  $\omega^{(0)}$  for marginal/joint density by  $\omega_{j}^{(1)} = \frac{\psi\left(\left\|k_{\sigma}(\boldsymbol{k}_{j},\cdot) - \hat{p}_{\mathrm{robust}}^{(k)}(\boldsymbol{k})\right\|_{\mathcal{H}_{k_{\sigma}}}\right)}{\sum_{j\in[N]}\psi\left(\left\|k_{\sigma}(\boldsymbol{k}_{j},\cdot) - \hat{p}_{\mathrm{robust}}^{(k)}(\boldsymbol{k})\right\|_{\mathcal{H}_{k_{\sigma}}}\right)}.$
7: [SPKDE] Obtain optimal weights  $\omega$  for marginal/joint density via solving equation (7).
8: [RKDE, SPKDE] Obtain robust self-attention vector  $\widehat{\pmb{h}}_i = \frac{\sum_{j\in[N]}\mathbf{v}_j\omega_j^{\mathrm{joint}}k_\sigma(\pmb{q}_i - \pmb{k}_j)}{\sum_{j\in[N]}\omega_j^{\mathrm{marginal}}k_\sigma(\pmb{q}_i - \pmb{k}_j)}.$
9: [MoM] Obtain attention vector  $\widehat{\pmb{h}}_i = \frac{\sum_{j\in I_l}v_jk_\sigma(\pmb{q}_i - \pmb{k}_j)}{\sum_{j\in I_l}k_\sigma(\pmb{q}_i - \pmb{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 KDE is 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  $\beta > 1$ , and let  $\mathcal{C}_{\sigma}^{N}$  be the convex hull of  $k_{\sigma}(\pmb{x}_1,\cdot),\dots,k_{\sigma}(\pmb{x}_N,\cdot)\in \mathcal{H}_{k_\sigma}$ , i.e., the space of weighted KDEs, the optimal density  $\hat{p}_{\mathrm{robust}}$  is given by
$$
\hat {p} _ {\text {r o b u s t}} = \arg \min  _ {p \in \mathcal {C} _ {\sigma} ^ {N}} \left\| \frac {\beta}{N} \sum_ {j \in [ N ]} k _ {\sigma} \left(x _ {j}, \cdot\right) - p \right\| _ {\mathcal {H} _ {k _ {\sigma}}} ^ {2}, \tag {6}
$$
which is guaranteed to have a unique minimizer since we are projecting in a Hilbert space and  $\mathcal{C}_{\sigma}^{N}$  is closed and convex. Notice that, by definition,  $\hat{p}_{\mathrm{robust}}$  can also be represented as  $\hat{p}_{\mathrm{robust}} = \sum_{j\in [N]}\omega_jk_\sigma (x_j,\cdot)$ ,  $\omega \in \Delta^{N}$ , which is same as the formulation in Eq. (5). Then Eq. (6) can be written as a quadratic programming (QP) problem over  $\omega$ :
$$
\min  _ {\omega} \omega^ {\top} G \omega - 2 q ^ {\top} \omega , \quad \text {s u b j e c t} \omega \in \Delta^ {N}, \tag {7}
$$
where  $G$  is the Gram matrix of  $\{\pmb{x}_j\}_{j\in [N]}$  with  $k_{\sigma}$  and  $q = G1\frac{\beta}{N}$ . Since  $k_{\sigma}$  is a positive-definite kernel and each  $\pmb{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  $\hat{p}_{\mathrm{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:
$$
\hat {p} _ {\text {r o b u s t}} (\mathbf {v}, \boldsymbol {k}) = \sum_ {j \in [ N ]} \omega_ {j} ^ {\text {j o i n t}} k _ {\sigma} \left(\left[ \mathbf {v} _ {j}, \boldsymbol {k} _ {j} \right], \left[ \mathbf {v}, \boldsymbol {k} \right]\right), \quad \hat {p} _ {\text {r o b u s t}} (\boldsymbol {k}) = \sum_ {j \in [ N ]} \omega_ {j} ^ {\text {m a r g i n a l}} k _ {\sigma} \left(\boldsymbol {k} _ {j}, \boldsymbol {k}\right). \tag {8}
$$
Following the non-parametric regression formulation of self-attention in Eq. (3), we obtain the robust self-attention mechanism as
$$
\widehat {\boldsymbol {h}} _ {i} = \frac {\sum_ {j \in [ N ]} \mathbf {v} _ {j} \omega_ {j} ^ {\text {j o i n t}} k _ {\sigma} \left(\boldsymbol {q} _ {i} - \boldsymbol {k} _ {j}\right)}{\sum_ {j \in [ N ]} \omega_ {j} ^ {\text {m a r g i n a l}} k _ {\sigma} \left(\boldsymbol {q} _ {i} - \boldsymbol {k} _ {j}\right)}, \tag {9}
$$
where  $\omega^{\mathrm{joint}}$  and  $\omega^{\mathrm{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  $\{\omega_{j}^{\mathrm{marginal}}\}_{j\in [N]}$
and  $\{\omega_{j}^{\mathrm{joint}}\}_{j\in [N]}$  are separate as  $\omega_{j}^{\mathrm{joint}}$  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  $\hat{p}_{\mathrm{robust}}(\mathbf{v},\mathbf{k})$  and only use the key vectors for obtaining the set of weights for the marginal  $\hat{p}_{\mathrm{robust}}(\mathbf{k})$ . In addition,  $\omega^{\mathrm{marginal}},\omega^{\mathrm{joint}}\in \mathbb{R}^{j\times i}$  for  $i,j = 1,\dots ,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.