SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\mathbf{W}_Q, \mathbf{W}_K \in \mathbb{R}^{D \times D_x}, \mathbf{W}_V \in \mathbb{R}^{D_v \times D_x}$  are the weight matrices. Equation (13) can be written in the following equivalent matrix form:
$$
\mathbf {H} = \operatorname {s o f t m a x} \left(\frac {\boldsymbol {Q} \boldsymbol {K} ^ {\top}}{\sqrt {D}}\right) \boldsymbol {V}, \tag {15}
$$
where the softmax function is applied to each row of the matrix  $(\mathbf{Q}\mathbf{K}^{\top}) / \sqrt{D}$ . Equation (15) is also called the "softmax attention". 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 (\boldsymbol {k}) + \varepsilon , \tag {16}
$$
where  $\varepsilon$  is some noise vectors with  $\mathbb{E}[\varepsilon] = 0$ , and  $f$  is the function that we want to estimate. If  $\{\pmb{k}_j\}_{j \in [N]}$  are i.i.d. samples from the distribution  $p(\pmb{k})$ , and  $p(\mathbf{v}, \pmb{k})$  is the joint distribution of  $(\mathbf{v}, \pmb{k})$  defined by equation (16), we have
$$
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 {17}
$$
We need to obtain estimations for both the joint density function  $p(\mathbf{v}, \mathbf{k})$  and the marginal density function  $p(\mathbf{k})$  to obtain function  $f$ , one popular approach is the kernel density estimation:
$$
\hat {p} _ {\sigma} (\mathbf {v}, \boldsymbol {k}) = \frac {1}{N} \sum_ {j \in [ N ]} k _ {\sigma} \left(\left[ \mathbf {v}, \boldsymbol {k} \right] - \left[ \mathbf {v} _ {j}, \boldsymbol {k} _ {j} \right]\right) \tag {18}
$$
$$
\hat {p} _ {\sigma} (\boldsymbol {k}) = \frac {1}{N} \sum_ {j \in [ N ]} k _ {\sigma} (\boldsymbol {k} - \boldsymbol {k} _ {j}), \tag {19}
$$
where  $[\mathbf{v},\pmb{k}]$  denotes the concatenation of  $\mathbf{v}$  and  $\pmb{k}$ .  $k_{\sigma}$  could be isotropic Gaussian kernel:  $k_{\sigma}(\pmb{x}-\pmb{x}^{\prime}) = \exp \left(-\|\pmb{x} - \pmb{x}^{\prime}\|^{2} / (2\sigma^{2})\right)$ , we have
$$
\hat {p} _ {\sigma} (\mathbf {v}, \boldsymbol {k}) = \frac {1}{N} \sum_ {j \in [ N ]} k _ {\sigma} (\mathbf {v} - \mathbf {v} _ {j}) k _ {\sigma} (\boldsymbol {k} - \boldsymbol {k} _ {j}). \tag {20}
$$
Combining equations (19), (20), and (17), we obtain the NW estimator of the function  $f$  as
$$
\begin{array}{l} \widehat {f} _ {\sigma} (\boldsymbol {k}) = \int_ {\mathbb {R} ^ {D}} \frac {\mathbf {v} \cdot \hat {p} _ {\sigma} (\mathbf {v} , \boldsymbol {k})}{\hat {p} _ {\sigma} (\boldsymbol {k})} d \mathbf {v} (21) \\ = \int_ {\mathbb {R} ^ {D}} \frac {\mathbf {v} \cdot \sum_ {j \in [ N ]} k _ {\sigma} (\mathbf {v} - \mathbf {v} _ {j}) k _ {\sigma} (\boldsymbol {k} - \boldsymbol {k} _ {j})}{\sum_ {j \in [ N ]} k _ {\sigma} (\boldsymbol {k} - \boldsymbol {k} _ {j})} d \mathbf {v} \\ = \frac {\sum_ {j \in [ N ]} k _ {\sigma} (\boldsymbol {k} - \boldsymbol {k} _ {j}) \int \mathbf {v} \cdot k _ {\sigma} (\mathbf {v} - \mathbf {v} _ {j}) d \mathbf {v}}{\sum_ {j \in [ N ]} k _ {\sigma} (\boldsymbol {k} - \boldsymbol {k} _ {j})} \\ = \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})}. (22) \\ \end{array}
$$
Now we show how the self-attention mechanism is related to the NW estimator. If the keys  $\{\pmb{k}_j\}_{j\in [N]}$  are normalized
$$
\begin{array}{l} \widehat {f} _ {\sigma} (\boldsymbol {q}) = \frac {\sum_ {j \in [ N ]} \mathbf {v} _ {j} \exp \left(- \| \boldsymbol {q} - \boldsymbol {k} _ {j} \| ^ {2} / 2 \sigma^ {2}\right)}{\sum_ {j \in [ N ]} \exp \left(- \| \boldsymbol {q} - \boldsymbol {k} _ {j} \| ^ {2} / 2 \sigma^ {2}\right)} \\ = \frac {\sum_ {j \in [ N ]} \mathbf {v} _ {j} \exp \left[ - \left(\| \boldsymbol {q} \| ^ {2} + \| \boldsymbol {k} _ {j} \| ^ {2}\right) / 2 \sigma^ {2} \right] \exp \left(\boldsymbol {q} ^ {\top} \boldsymbol {k} _ {j} / \sigma^ {2}\right)}{\sum_ {j \in [ N ]} \exp \left[ - \left(\| \boldsymbol {q} \| ^ {2} + \| \boldsymbol {k} _ {j} \| ^ {2}\right) / 2 \sigma^ {2} \right] \exp \left(\boldsymbol {q} ^ {\top} \boldsymbol {k} _ {j} / \sigma^ {2}\right)} \\ = \sum_ {j \in [ N ]} \frac {\exp \left(\boldsymbol {q} ^ {\top} \boldsymbol {k} _ {j} / \sigma^ {2}\right)}{\sum_ {j \in [ N ]} \exp \left(\boldsymbol {q} ^ {\top} \boldsymbol {k} _ {j} / \sigma^ {2}\right)} \mathbf {v} _ {j} \\ = \sum_ {j \in [ N ]} \operatorname {s o f t m a x} \left(\boldsymbol {q} ^ {\top} \boldsymbol {k} _ {j} / \sigma^ {2}\right) \mathbf {v} _ {j}. \tag {23} \\ \end{array}
$$
Then estimating the softmax attention is equivalent to estimating  $\widehat{f}_{\sigma}(\pmb {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:
$$
\rho (x) := \left\{ \begin{array}{c c} x ^ {2} / 2, & 0 \leq x \leq a \\ a x - a ^ {2} / 2, & a <   x, \end{array} \right. \tag {24}
$$
where  $a$  is a constant. The solution to this robust regression problem has the following form:
Proposition 1. Assume the robust loss function  $\rho$  is non-decreasing in  $[0,\infty]$ ,  $\rho(0) = 0$  and  $\lim_{x\to 0}\frac{\rho(x)}{x} = 0$ . Define  $\psi(x) \coloneqq \frac{\rho'(x)}{x}$  and assume  $\psi(0) = \lim_{x\to 0}\frac{\rho'(x)}{x}$  exists and finite. Then the optimal  $\hat{p}_{\text{robust}}$  can be written as
$$
\hat {p} _ {\text {r o b u s t}} = \sum_ {j \in [ N ]} \omega_ {j} k _ {\sigma} (\boldsymbol {x} _ {j}, \cdot),
$$
where  $\omega = (\omega_{1},\dots ,\omega_{N})\in \Delta_{N}$ , with each  $\omega_{j}\propto \psi$ $(\| k_{\sigma}(\pmb{x}_{j},\cdot) - \hat{p}_{\text{robust}}\|_{\mathcal{H}_{k_{\sigma}}})$ . Here  $\Delta_{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 \in \mathcal{H}_{k_{\sigma}}$ , we denote
$$
J (p) = \frac {1}{N} \sum_ {j \in [ N ]} \rho \left(\| k _ {\sigma} (\boldsymbol {x} _ {j}, \cdot) - p \| _ {\mathcal {H} _ {k _ {\sigma}}}\right).
$$
Then we have the following lemma regarding the Gateaux differential of  $J$  and a necessary condition for  $\hat{p}_{\mathrm{robust}}$  to be optimal solution of the robust loss objective function in equation (5).
Lemma 1. Given the assumptions on the robust loss function  $\rho$  in Proposition 1, the Gateaux differential of  $J$  at  $p\in \mathcal{H}_{k_{\sigma}}$  with incremental  $h\in \mathcal{H}_{k_{\sigma}}$ , defined as  $\delta J(p;h)$ , is
$$
\delta J (p; h) := \lim  _ {\tau \rightarrow 0} \frac {J (p + \tau h) - J (p)}{\tau} = - \langle V (p), h \rangle_ {\mathcal {H} _ {k _ {\sigma}}},
$$
where the function  $V:\mathcal{H}_{k_{\sigma}}\to \mathcal{H}_{k_{\sigma}}$  is defined as:
$$
V (p) = \frac {1}{N} \sum_ {j \in [ N ]} \psi \left(\| k _ {\sigma} (\boldsymbol {x} _ {j}, \cdot) - p \| _ {\mathcal {H} _ {k _ {\sigma}}}\right) \left(k _ {\sigma} (\boldsymbol {x} _ {j}, \cdot) - p\right).
$$
A necessary condition for  $\hat{p}_{\text{robust}}$  is  $V(\hat{p}_{\text{robust}}) = 0$ .
The proof of Lemma 1 can be found in Lemma 1 of Kim & Scott (2012). Based on the necessary condition for  $\hat{p}_{\mathrm{robust}}$  in Lemma 1, i.e.,  $V(\hat{p}_{\mathrm{robust}}) = 0$ , we have
$$
\frac {1}{N} \sum_ {j \in [ N ]} \psi \left(\| k _ {\sigma} (\boldsymbol {x} _ {j}, \cdot) - \hat {p} _ {\text {r o b u s t}} \| _ {\mathcal {H} _ {k _ {\sigma}}}\right) \left(k _ {\sigma} (\boldsymbol {x} _ {j}, \cdot) - \hat {p} _ {\text {r o b u s t}}\right) = 0.
$$
Direct algebra indicates that  $\hat{p}_{\mathrm{robust}} = \sum_{j\in [N]}\omega_jk_\sigma (\boldsymbol {x}_j,\cdot)$  where  $\omega = (\omega_{1},\dots ,\omega_{N})\in \Delta_{N}$ , and  $\omega_{j}\propto \psi \left(\| k_{\sigma}(\boldsymbol{x}_{j},\cdot) - \hat{p}_{\mathrm{robust}}\|_{\mathcal{H}_{k_{\sigma}}}\right)$ . As a consequence, we obtain the conclusion of the proposition.
For the Huber loss function, we have that
$$
\psi (x) := \left\{ \begin{array}{c c} 1, & 0 \leq x \leq a \\ a / x, & a <   x. \end{array} \right.
$$
Hence, when the error  $\| k_{\sigma}(\pmb{x}_j,\cdot),\cdot -\hat{p}_{\mathrm{robust}}\|_{\mathcal{H}_{k_\sigma}}$  is over the threshold  $a$ , the final estimator will down-weight the importance of  $k_{\sigma}(\pmb{x}_j,\cdot)$ . This is in sharp contrast with the standard KDE method, which will assign uniform weights to all of the  $k_{\sigma}(\pmb{x}_j,\cdot)$ . As we mentioned in the main paper, the estimator provided in Proposition 1 is circularly defined, as  $\hat{p}_{\mathrm{robust}}$  is defined via  $\omega$ , and  $\omega$  depends on  $\hat{p}_{\mathrm{robust}}$ . Such an issue can be addressed by estimating  $\omega$  with an iterative algorithm termed as kernelized iteratively re-weighted least-squares (KIRWLS). The algorithm starts with randomly initialized  $\omega^{(0)}\in \Delta_n$ , and performs the following iterative updates between two steps:
$$
\hat {p} _ {\text {r o b u s t}} ^ {(k)} = \sum_ {j \in [ N ]} \omega_ {i} ^ {(k - 1)} k _ {\sigma} \left(\boldsymbol {x} _ {j}, \cdot\right), \quad \omega_ {j} ^ {(k)} = \frac {\psi \left(\left\| k _ {\sigma} \left(\boldsymbol {x} _ {j} , \cdot\right) - \hat {p} _ {\text {r o b u s t}} ^ {(k)} \right\| _ {\mathcal {H} _ {k \sigma}}\right)}{\sum_ {j \in [ N ]} \psi \left(\left\| k _ {\sigma} \left(\boldsymbol {x} _ {j} , \cdot\right) - \hat {p} _ {\text {r o b u s t}} ^ {(k)} \right\| _ {\mathcal {H} _ {k \sigma}}\right)}. \tag {25}
$$
Note that, the optimal  $\hat{p}_{\mathrm{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  $\left\| k_{\sigma}(\boldsymbol{x}_j,\cdot) - \hat{p}_{\mathrm{robust}}^{(k)}\right\|_{\mathcal{H}_{k_\sigma}}$  via the reproducing property:
$$
\begin{array}{l} \left\| k _ {\sigma} \left(\boldsymbol {x} _ {j}, \cdot\right) - \hat {p} _ {\text {r o b u s t}} ^ {(k)} \right\| _ {\mathcal {H} _ {k _ {\sigma}}} ^ {2} = - 2 \sum_ {m \in [ N ]} \omega_ {m} ^ {(k - 1)} k _ {\sigma} \left(\boldsymbol {x} _ {m}, \boldsymbol {x} _ {j}\right) + k _ {\sigma} \left(\boldsymbol {x} _ {j}, \boldsymbol {x} _ {j}\right) \\ + \sum_ {m \in [ N ], n \in [ N ]} \omega_ {m} ^ {(k - 1)} \omega_ {n} ^ {(k - 1)} k _ {\sigma} \left(\boldsymbol {x} _ {m}, \boldsymbol {x} _ {n}\right). \tag {26} \\ \end{array}
$$
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  $\phi : \mathbb{R} \to \mathbb{R}$  and radius  $R$ , we randomly divide the keys  $\{\pmb{k}_i\}_{i \in [N]}$  into  $B$  subsets  $I_1, \ldots, I_B$  of equal size where  $|I_1| = |I_2| = \dots = |I_B| = S$ . Define  $\hat{p}_{R,I_m}(\pmb{q}_l) = \frac{1}{S} \sum_{i \in I_m} \prod_{j=1}^{D} \phi\left(\frac{\sin(R(q_{l_j} - k_{ij}))}{R(q_{l_j} - k_{ij})}\right)$ , then the MoM Fourier attention is defined as
$$
\hat {\mathbf {h}} _ {l} = \frac {\frac {1}{S} \sum_ {i \in I _ {m}} \mathbf {v} _ {i} \prod_ {j = 1} ^ {D} \phi \left(\frac {\sin \left(R \left(q _ {l j} - k _ {i j}\right)\right)}{R \left(q _ {i j} - k _ {l j}\right)}\right)}{\operatorname {m e d i a n} \left\{\hat {p} _ {R , I _ {1}} \left(\boldsymbol {q} _ {l}\right) , \dots , \hat {p} _ {R , I _ {B}} \left(\boldsymbol {q} _ {l}\right) \right\}}, \tag {27}
$$
where  $I_{m}$  is the block such that  $\hat{p}_{R}(\boldsymbol{q}_{l},\boldsymbol{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  $\hat{p}_R(\boldsymbol{q}) = \mathrm{median}\{\hat{p}_{R,I_1}(\boldsymbol{q}),\dots,\hat{p}_{R,I_B}(\boldsymbol{q})\}$  is a robust estimator of the density function  $p(\boldsymbol{q})$  of the keys. We first introduce a few notations that are useful for stating this result. Denote  $\mathcal{C} = \{1\leq i\leq N: k_i \text{ is clean}\}$  and  $\mathcal{O} = \{1\leq i\leq N: k_i \text{ is outlier}\}$ . Then, we have  $\mathcal{C}\cap \mathcal{O} = \emptyset$  and  $\mathcal{C}\cup \mathcal{O} = \{1,2,\ldots ,N\}$ . The following result establishes a high probability upper bound on the sup-norm between  $\widehat{p}_R(\boldsymbol{q})$  and  $p(\boldsymbol{q})$ .
Theorem 1. Assume that the function  $\phi$  satisfies  $\int \phi (\sin (z) / z)z^{j}dz = 0$  for all  $1\leq j\leq m$  and  $\int |\phi (\sin (z) / z)||z|^{m + 1}dz < \infty$  for some  $m\in \mathbb{N}$ . Furthermore, the density function  $p(\pmb {q})$  satisfies  $\sup_q|p(\pmb {q})| < \infty$ . The number of blocks  $B$  and the number of outliers  $|\mathcal{O}|$  are such that  $B > (2 + \delta)|\mathcal{O}|$  where  $\delta$  is the failure probability. Then, with  $\Delta = \frac{1}{2 + \delta} -\frac{|\mathcal{O}|}{B}$  for the radius  $R$  sufficiently large and  $\delta$  sufficiently small, with probability at least  $1 - \exp (-2\Delta^2 B)$  we find that
$$
\| \hat {p} _ {R} - p \| _ {\infty} \leq C \left(\frac {1}{R ^ {m + 1}} + \sqrt {\frac {B R ^ {D} \log R \log (2 / \delta)}{N}}\right)
$$
where  $C$  is some universal constant.
Remark 1. The result of Theorem 1 indicates by choosing  $R = \mathcal{O}\big(N^{-\frac{1}{2(m + 1) + D}}\big)$ , the rate of  $\hat{p}_R$  to  $p$  under the supremum norm is  $\mathcal{O}\big(N^{-\frac{m + 1}{2(m + 1) + D}}\big)$ . With that choice of  $R$ , when  $N$  approaches infinity, the MoM estimator  $\hat{p}_R$  is a consistent estimator of the clean distribution  $p$  of the keys. This confirms the validity of using  $\hat{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  $\widehat{p}_R(\boldsymbol{q})$ , we obtain the following inequality
$$
\{\sup  _ {\boldsymbol {q}} | \hat {p} _ {R} (\boldsymbol {q}) - p (\boldsymbol {q}) | \geq \epsilon \} \subset \{\sup  _ {\boldsymbol {q}} \sum_ {b = 1} ^ {B} \mathbf {1} _ {\{| \hat {p} _ {R, I _ {b}} (\boldsymbol {q}) - p (\boldsymbol {q}) | \geq \epsilon \}} \geq \frac {B}{2} \}
$$
This bound indicates that to bound  $\mathbb{P}(\| \hat{p}_R(\boldsymbol {q}) - p(\boldsymbol {q})\|_{\infty}\geq \epsilon)$ , it is sufficient to bound  $\mathbb{P}(\{\sup_q\sum_{b = 1}^B\mathbf{1}_{\{| \hat{p}_{R,I_b}(\boldsymbol {q}) - p(\boldsymbol {q})|\geq \epsilon \}}\geq \frac{B}{2}\})$ . Indeed, for each  $1\leq b\leq B$ , we find that
$$
\mathbf {1} _ {\{| \hat {p} _ {R, I _ {b}} (\boldsymbol {q}) - p (\boldsymbol {q}) | \geq \epsilon \}} \leq \mathbf {1} _ {\{\sup  _ {\boldsymbol {q}} \{| \hat {p} _ {R, I _ {b}} (\boldsymbol {q}) - p (\boldsymbol {q}) | \geq \epsilon \}}.
$$
$$
\sum_ {b = 1} ^ {B} \mathbf {1} _ {\{| \hat {p} _ {R, I _ {b}} (\boldsymbol {q}) - p (\boldsymbol {q}) | \geq \epsilon \}} \leq \sum_ {b = 1} ^ {B} \mathbf {1} _ {\left\{\sup  _ {\boldsymbol {q}} \left\{| \hat {p} _ {R, I _ {b}} (\boldsymbol {q}) - p (\boldsymbol {q}) | \geq \epsilon \right\}, \right.}
$$
which leads to  $\sup_q\sum_{b = 1}^B\mathbf{1}_{\{| \hat{p}_{R,I_b}(\pmb {q}) - p(\pmb {q})|\geq \epsilon \}}\leq \sum_{b = 1}^B\mathbf{1}_{\{\sup_q\{| \hat{p}_{R,I_b}(\pmb {q}) - p(\pmb {q})|\geq \epsilon \}}$ . This inequality shows that
$$
\mathbb {P} \left(\left. \sum_ {\boldsymbol {q}} \sum_ {b = 1} ^ {B} \mathbf {1} _ {\{| \hat {p} _ {R, I _ {b}} (\boldsymbol {q}) - p (\boldsymbol {q}) | \geq \epsilon \}} \geq \frac {B}{2} \right\}\right) \leq \mathbb {P} \left(\sum_ {b = 1} ^ {B} \mathbf {1} _ {\left\{\sup  _ {\boldsymbol {q}} \left\{\left| \hat {p} _ {R, I _ {b}} (\boldsymbol {q}) - p (\boldsymbol {q}) \right| \geq \epsilon \right\}\right)}\right).
$$
To ease the presentation, we denote  $W_{b} = \mathbf{1}_{\{\sup_{\boldsymbol{q}}\{|\hat{p}_{R,I_{b}}(\boldsymbol{q}) - p(\boldsymbol{q})| \geq \epsilon \}}$  and  $\mathcal{B} = \{1 \leq b \leq B : I_{b} \cap \mathcal{O} = \emptyset\}$ . Then, the following inequalities hold
$$
\begin{array}{l} \sum_ {b = 1} ^ {B} \mathbf {1} _ {\{\sup _ {\boldsymbol {q}} \{| \hat {p} _ {R, I _ {b}} (\boldsymbol {q}) - p (\boldsymbol {q}) | \geq \epsilon \}} = \sum_ {b \in \mathcal {B}} W _ {b} + \sum_ {b \in \mathcal {B} ^ {c}} W _ {b} \\ \leq \sum_ {b \in \mathcal {B}} W _ {b} + | \mathcal {O} | \\ \leq \sum_ {b \in \mathcal {B}} \left(W _ {b} - \mathbb {E} \left[ W _ {b} \right]\right) + B \cdot \mathbb {P} \left(\sup  _ {\boldsymbol {q}} \left| \hat {p} _ {R, I _ {1}} (\boldsymbol {q}) - p (\boldsymbol {q}) \right| > \epsilon\right) + | \mathcal {O} |, \\ \end{array}
$$
where we assume without loss of generality that  $1 \in \mathcal{B}$ , which is possible due to the assumption that  $B > (2 + \delta)|\mathcal{O}|$ . We now prove the following uniform concentration bound:
Lemma 2. Assume that  $\phi (z)\leq C$  for all  $|z|\leq 1$  for some universal constant  $C$ . We have
$$
\mathbb {P} \left(\sup  _ {\boldsymbol {q}} | \hat {p} _ {R, I _ {1}} (\boldsymbol {q}) - p (\boldsymbol {q}) | \geq C \left(\frac {1}{R ^ {m + 1}} + \sqrt {\frac {R ^ {D} \log R \log (2 / \delta)}{| I _ {1} |}}\right)\right) \leq \delta .
$$
Proof of Lemma 2. By the triangle inequality, we have
$$
\sup  _ {\boldsymbol {q}} | \hat {p} _ {R, I _ {1}} (\boldsymbol {q}) - p (\boldsymbol {q}) | \leq \sup  _ {\boldsymbol {q}} | \hat {p} _ {R, I _ {1}} (\boldsymbol {q}) - \mathbb {E} [ \hat {p} _ {R, I _ {1}} (\boldsymbol {q}) ] | + \sup  _ {\boldsymbol {q}} | \mathbb {E} [ \hat {p} _ {R, I _ {1}} (\boldsymbol {q}) ] - p (\boldsymbol {q}) |.
$$
To bound  $\sup_q |\hat{p}_{R,I_1}(\boldsymbol{q}) - \mathbb{E}[\hat{p}_{R,I_1}(\boldsymbol{q})]|$ , we use Bernstein's inequality along with the bracketing entropy under  $\mathbb{L}_1$  norm in the space of queries  $\boldsymbol{q}$ . In particular, we denote  $Y_i = \frac{R^D}{A^D} \prod_{j=1}^{D} \phi\left(\frac{\sin(R(q_j - k_{ij}))}{R(q_j - k_{ij})}\right)$  where  $A = \int_{\mathbb{R}} \phi\left(\frac{\sin(z)}{z}\right) dz$  for all  $i \in I_1$ . Since  $\sin(R(q_j - k_{ij})) / (R(q_j - k_{ij})) \leq 1$  for all  $1 \leq j \leq D$ , we obtain that  $|Y_i| \leq C^D R^D / A^D$  where  $C$  is the constant such that  $\phi(z) \leq C$  when  $|z| \leq 1$ . Furthermore,  $E[|Y_i|]$
By choose  $\epsilon = C\left(\frac{1}{R^{m + 1}} +\sqrt{\frac{R^D\log R\log(2 / \delta)}{|I_1|}}\right)$ , then we find that
$$
\mathbb {P} \left(\sup  _ {\boldsymbol {q}} \left| \hat {p} _ {R, I _ {1}} (\boldsymbol {q}) - p (\boldsymbol {q}) \right| > \epsilon\right) \leq \frac {\delta}{2 (2 + \delta)}
$$
Collecting the above inequalities leads to
$$
\mathbb {P} (\{\sup _ {\boldsymbol {q}} \sum_ {b = 1} ^ {B} \mathbf {1} _ {\{| \hat {p} _ {R, I _ {b}} (\boldsymbol {q}) - p (\boldsymbol {q}) | \geq \epsilon \}} \geq \frac {B}{2} \}) \leq \exp (- 2 B \Delta^ {2}),
$$
where  $\Delta = \frac{1}{2 + \delta} -\frac{|\mathcal{O}|}{B}$ . As a consequence, we obtain the conclusion of the theorem.
D Dataset Information
WikiText-103 The dataset<sup>1</sup> contains around 268K words and its training set consists of about 28K articles with 103M tokens, this corresponds to text blocks of about 3600 words. The validation set and test sets consist of 60 articles with 218K and 246K tokens respectively.
ImageNet We use the full ImageNet dataset that contains  $1.28M$  training images and  $50K$  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-of-distribution detectors. The dataset contains samples from ImageNet-22K but not from ImageNet-1K, 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.
1www.salesforce.com/products/einstein/ai-research/the-wikitext-dependency-language-modeling-dataset/
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.
Table 5: Perplexity (PPL) and negative likelihood loss (NLL) of our methods (lower part) and baselines (upper part) on WikiText-103 using a medium version of Transformer. The best results are highlighted in bold font and the second best are highlighted in underline. On clean data, Transformer-SPKDE achieves better PPL and NLL than other baselines. Under random swap with outlier words, Transformers with MoM self-attention show much better performance.
Method (median version)	Clean Data	Word Swap
Valid PPL/Loss	Test PPL/Loss	Valid PPL/Loss	Test PPL/Loss
Transformer (Vaswani et al., 2017b)	27.90/3.32	29.60/3.37	65.36/4.31	68.12/4.36
Performer (Choromanski et al., 2021)	27.34/3.31	29.51/3.36	64.72/4.30	67.43/4.34
Transformer-MGK (Nguyen et al., 2022b)	27.28/3.31	29.24/3.36	64.46/4.30	67.31/4.33
FourierFormer (Nguyen et al., 2022c)	26.51/3.29	28.01/3.33	63.74/4.28	65.27/4.31
Transformer-RKDE (Huber)	26.12/3.28	27.89/3.32	49.37/3.85	51.22/3.89
Transformer-RKDE (Hampel)	25.87/3.27	27.44/3.31	48.62/3.83	51.03/3.88
Transformer-SPKDE	25.76/3.27	27.35/3.31	46.91/3.79	49.14/3.84
Transformer-MoM	28.26/3.34	29.98/3.38	45.35/3.75	47.92/3.81
FourierFormer-MoM	27.13/3.31	29.02/3.36	43.23/3.71	44.97/3.74
Table 6: Test PPL/NLL loss versus the parameter  $a$  of Huber loss function defined in Eq. (24) (upper) and Hampel loss function (Kim & Scott, 2012) (lower; we use  $2 \times a$  and  $3 \times a$  as parameters  $b$  and  $c$  ) on original and word-swapped Wiki-103 dataset. The best results are highlighted in bold font and the second best are highlighted in underline. We choose  $a = {0.4}$  in rest of the experiments.
Robust Loss Parameter	0.1	0.2	0.4	0.6	0.8	1
Clean Data	32.92/3.48	32.87/3.48	32.29/3.47	32.38/3.48	32.46/3.48	32.48/3.48
Word Swap	55.82/3.99	55.97/3.99	55.68/3.99	56.89/4.01	57.26/4.01	57.37/4.01
Clean Data	32.67/3.48	32.32/3.48	32.35/3.48	32.47/3.48	32.53/3.48	32.58/3.48
Word Swap	58.02/4.03	57.86/4.03	57.92/4.03	58.24/4.04	58.37/4.04	58.43/4.04
Table 7: Top-1 classification accuracy on ImageNet versus the parameter  $a$  of Huber loss function defined in Eq. (24) under different settings. The best results are highlighted in bold font and the second best are highlighted in underline. We choose  $a = 0.2$  in rest of the experiments.
Huber Loss Parameter	0.1	0.2	0.4	0.6	0.8	1
Clean Data	71.45	72.83	71.62	71.07	70.65	70.34
FGSM	56.72	55.83	55.34	54.87	54.02	52.98
PGD	46.37	44.15	43.87	43.25	42.69	41.96
SPSA	52.38	52.42	51.69	51.34	50.97	48.22
Imagenet-C	45.37	45.58	45.63	45.26	44.63	43.76