SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The proof of Theorem 2 is contained in Appendix B.2. The proof mostly uses the Rademacher complexity bound for the supremum of empirical process. The main non-trivial technical challenge lies in showing the concentration of  $\sup_{f\in \mathcal{V}_M}|\widehat{L}_D(f) - L_D(f)|$ , which cannot be simply controlled due to the unboundedness of the infinity norm of functions in the target function class  $\nu_{M}$ . We dealt with this subtlety by a carefully decomposition of  $\sup_{f\in \mathcal{V}_M}|\widehat{L}_D(f) - L_D(f)|$ . The seemingly unnatural constraint set (12) is used in bounding different terms in this decomposition.
3.3 Examples and comparison
We next give the sample complexity for learning several examples of target functions using the random-feature attention model. We will compare its sample complexity for learning these functions with that of the standard random-feature model [74] (thereafter, we call it the random-feature MLP model, in short RFMLP model). In the RFMLP model, we view  $\mathbf{x}_{0:N}$  as an input vector instead of a sequence of vectors denoted as  $\mathrm{vec}(\mathbf{x}_{0:N}) = [\mathbf{x}_0;\mathbf{x}_1;\dots ;\mathbf{x}_N;1]\in \mathbb{R}^{d(N + 1) + 1}$ . The RFMLP is given by
$$
f _ {M} ^ {\mathrm {M L P}} \left(\mathbf {x} _ {0: N}; \mathbf {v}\right) = \sum_ {m = 1} ^ {M} \sigma \left(\left\langle \mathbf {w} _ {m}, \operatorname {v e c} \left(\mathbf {x} _ {0: N}\right) \right\rangle\right) \cdot v _ {m}, \quad \left\{\mathbf {w} _ {m} \right\} _ {m \in [ M ]} \sim_ {\mathrm {i i d}} \mathsf {N} (\mathbf {0}, \mathbf {I} / (N + 2)). \tag {14}
$$
We choose the variance of random weights  $\mathbf{w}_m$  to be  $1 / (N + 2)$  to ensure that  $\langle \mathbf{w}_m, \mathrm{vec}(\mathbf{x}_{0:N}) \rangle \sim \mathsf{N}(0,1)$  has unit variance. The generalization and approximation properties of the random-feature MLP model have been well-studied in the literature, for example, [7, 8, 60].
We instantiate Theorem 2 on three concrete examples of target functions (calculations of the excess risks in Appendix B.4, where the result for RFMLP are adapted from Arora et al. [7]). In all three cases, the target functions are permutation invariant with respect to  $\{\mathbf{x}_i\}_{i\in [N]}$ , by which we naturally expect RFA to achieve better sample complexity than RFMLP in accordance with this structure. Although the comparison between RFA and RFMLP is based on comparing upper bounds on the sample complexity of both models, existing work has also derived lower bounds on the sample complexity of RFMLP, which aligns with the upper bound for RFMLP we used. We do not invoke these lower bounds, as they apply to a special case with a uniform distributional assumption on the input tokens.
Example 1 (Functions of  $\mathbf{x}_0$ ): We consider functions of  $\mathbf{x}_0$  (no dependence on  $\mathbf{x}_{1:N}$ ) of the form
$$
f _ {\star} \left(\mathbf {x} _ {0: N}\right) = \sum_ {k = 0} ^ {\infty} \left\langle \mathbf {x} _ {0} ^ {\otimes k}, \mathbf {A} _ {k} \right\rangle , \quad \mathbf {A} _ {k} \in \mathbb {R} ^ {d ^ {k}}, \quad \text {w i t h} B \left(f _ {\star}\right) = \sum_ {k = 0} ^ {\infty} C _ {k} \| \mathbf {A} _ {k} \| _ {\mathrm {F r}} ^ {2} \text {b y (1 0)}.
$$
By Theorem 2, setting  $M = \Theta (d^{2}n)$ , the excess risk bound gives  $\tilde{\mathcal{O}} (\sqrt{\sum_{k = 0}^{\infty}k^{4.5}A^{k}\|\mathbf{A}_{k}\|_{\mathrm{Fr}}^{2} / n})$ .
As a special case, consider  $f_{\star}(\mathbf{x}_{0:N}) = (\boldsymbol{\beta}^{\top}\mathbf{x}_0)^p$ , which corresponds to taking  $\mathbf{A}_k = \boldsymbol{\beta}^{\otimes p}$  for  $k = p$  and  $\mathbf{A}_k = \mathbf{0}$  for  $k \neq p$ . The above excess risk of RFA model and the RFMLP model scales as
$$
\mathrm {R F A}: \widetilde {\mathcal {O}} \Big (\operatorname {P o l y} (p) \sqrt {4 ^ {p} \| \boldsymbol {\beta} \| _ {2} ^ {2 p} / n} \Big), \qquad \mathrm {R F M L P}: \widetilde {\mathcal {O}} \Big (\operatorname {P o l y} (p) \sqrt {(N + 2) ^ {p} \| \boldsymbol {\beta} \| _ {2} ^ {2 p} / n} \Big).
$$
Compared to the RFMLP model, the RFA model significantly reduces the necessary sample size by a factor of  $(N / 4)^{p}$ .
Example 2 (Average of functions of  $\mathbf{x}_i$ ): We consider average of functions of  $\mathbf{x}_i$  of the form
$$
f _ {\star} \left(\mathbf {x} _ {0: N}\right) = \frac {1}{N} \sum_ {i = 1} ^ {N} \sum_ {k = 0} ^ {\infty} \left\langle \mathbf {x} _ {i} ^ {\otimes k}, \mathbf {A} _ {k} \right\rangle , \quad \mathbf {A} _ {k} \in \mathbb {R} ^ {d ^ {k}}, \quad \text {w i t h} B \left(f _ {\star}\right) = \sum_ {k = 0} ^ {\infty} C _ {k} \| \mathbf {A} _ {k} \| _ {\mathrm {F r}} ^ {2} \text {b y (1 0)}.
$$
Theorem 2 then gives an  $\widetilde{\mathcal{O}} (\sqrt{\sum_{k = 0}^{\infty}k^{4.54^k}\|\mathbf{A}_k\|_{\mathrm{Fr}}^2 / n})$  excess risk, same as Example 1.
As a specific example, consider  $f_{\star} = \frac{1}{N}\sum_{i = 1}^{N}\psi (\langle \beta ,\mathbf{x}_i\rangle)$  with  $\psi (z) = z\arctan (z / \eta)$  for some  $\eta >2,\| \beta \| _2 = 1$ . Using the power series expansion of  $\psi$ , the excess risk bound of RFA model and the RFMLP model scale as
$$
\mathrm {R F A}: \widetilde {\mathcal {O}} \Big (\sqrt {\sum_ {k = 1} ^ {\infty} k ^ {4 . 5} (2 / \eta) ^ {2 k} / n} \Big) = \widetilde {\mathcal {O}} (\sqrt {1 / n}), \quad \mathrm {R F M L P}: \widetilde {\mathcal {O}} \Big (\sqrt {\sum_ {k = 1} ^ {\infty} k ^ {4 . 5} [ (N + 2) / (2 \eta) ] ^ {2 k} / n} \Big).
$$
<sup>5</sup>By deriving the corresponding results for Random Features instead of Neural Tangent Kernels.
Figure 1: Visualization of weight matrices of the 2nd, 5th, 8th, and 11th layers of the BERT-Base model. Each row contains weight matrices of a layer. All matrices are clipped to the top-left  $32 \times 32$  block. Lighter color indicates a larger absolute value.
The latter diverges whenever  $\eta \leq (N + 2) / 2$ , in which case the bound is meaningless.
Example 3 (Correlation-weighted functions):  $f_{\star}$  is the following function:
$$
f _ {\star} (\mathbf {x} _ {0: N}) = \frac {1}{N} \sum_ {i = 1} ^ {N} F (\langle \mathbf {x} _ {0}, \mathbf {S x} _ {i} \rangle) G (\mathbf {x} _ {i}), \quad F (t) = \sum_ {k = 0} ^ {\infty} a _ {k} \cdot t ^ {k}, \quad G (\mathbf {x} _ {i}) = \sum_ {k = 0} ^ {\infty} \left\langle \mathbf {x} _ {i} ^ {\otimes k}, \mathbf {G} _ {k} \right\rangle ,
$$
for  $\mathbf{S} \in \mathbb{R}^{d \times d}$ ,  $\{a_k\}_{k \geq 0} \subseteq \mathbb{R}$ ,  $\mathbf{G}_k \in \mathbb{R}^{d^k}$ . This target function fully exploits the representation power of the attention layer. Eq. (10) gives  $B(f_\star) = \mathcal{O}\big(\sum_{k=0}^\infty C_k \left( \sum_{r+s=k} a_r^2 \| \mathbf{S} \|_{\mathbb{F}_r}^{2r} \| \mathbf{G}_s \|_{\mathbb{F}_r}^2 \right)\big)$ .
As a specific example, consider  $f_{1,\star} = \frac{1}{N}\sum_{i = 1}^{N}\langle \mathbf{x}_0,\mathbf{x}_i\rangle^p$ , corresponding to taking  $\mathbf{S} = \mathbf{I}_d$ ,  $F(t) = t^p$ , and  $G\equiv 1$ . The excess risk bound of RFA (by Theorem 2) and RFMLP scale as
$$
\mathrm {R F A}: \widetilde {\mathcal {O}} \Big (\operatorname {P o l y} (p) \sqrt {(4 d) ^ {p} / n} \Big), \qquad \mathrm {R F M L P}: \widetilde {\mathcal {O}} \Big (\operatorname {P o l y} (p) \sqrt {[ (N + 2) d ] ^ {p} / n} \Big).
$$
As another example, consider  $f_{2,\star} = \frac{1}{N}\sum_{i=1}^{N}\cos(\langle\mathbf{x}_0,\mathbf{x}_i\rangle)\langle\mathbf{x}_i^{\otimes p},\mathbf{G}\rangle$  with  $\| \mathbf{G}\|_{\mathrm{Fr}} = 1$ . Then the excess risk bound of RFA and RFMLP scale as
$$
\mathrm {R F A}: \widetilde {\mathcal {O}} \Big (\operatorname {P o l y} (p d) \sqrt {e ^ {4 \sqrt {d}} 4 ^ {p} / n} \Big), \quad \mathrm {R F M L P}: \widetilde {\mathcal {O}} \Big (\operatorname {P o l y} (p N d) \sqrt {e ^ {2 (N + 2) \sqrt {d}} N ^ {p} / n} \Big).
$$
RFA reduces the required sample size by factors of  $(N / 4)^p$  for  $f_{1,\star}$  and  $\exp (N\sqrt{d})$  for  $f_{2,\star}$ .
4 Expressivity of biased random-feature attention model
We now move beyond the Gaussian weight assumption by exploring alternative possibilities for the weight distribution in the attention heads. We observe empirically that the weight matrices in transformer architectures learned in practice are often more similar to the identity matrix than a mean-zero matrix (Figure 1; see the details in Appendix D.1. This is also observed in a recent and concurrent work [84]).
Towards understanding this effect, we consider an alternative attention model with biased random weights, where the bias is a fixed matrix  $\mathbf{W}_0\in \mathbb{R}^{(d + 1)\times (d + 1)}$
$$
f _ {M} ^ {\mathbf {W}, \mathbf {W} _ {0}} \left(\mathbf {x} _ {0: N}; \mathbf {V}\right) = \sum_ {m = 1} ^ {M} \frac {1}{N} \sum_ {i = 1} ^ {N} \sigma \left(\left\langle \mathbf {W} _ {0} + \mathbf {W} _ {m}, \widetilde {\mathbf {x}} _ {0} \widetilde {\mathbf {x}} _ {i} ^ {\top} \right\rangle\right) \left\langle \mathbf {v} _ {m}, \widetilde {\mathbf {x}} _ {i} \right\rangle . \tag {15}
$$
Here  $\{\mathbf{W}_m\}_{m\in [M]}$  are again Gaussian random matrices sampled according to (4). The biased random-feature attention model is similar to (5) except that a bias weight  $\mathbf{W}_0$  is added. Motivated by our observation, we choose  $\mathbf{W}_0 = [\mathbf{I}_{d\times d},\mathbf{0}_{d\times 1};\mathbf{0}_{1\times d},0]\in \mathbb{R}^{(d + 1)\times (d + 1)}$ , so that the diagonal elements of  $\mathbf{W}_0 + \mathbf{W}_m$  will be on average larger than the off-diagonal elements.
4.1 Expressivity of biased random-feature attention
Given the formulation of biased random-feature attention models (thereafter, we call it the biased random-feature attention model, in short BRFA model), a natural conjecture is that this model can better fit functions that are the average of function of  $\langle \mathbf{x}_0,\mathbf{x}_i\rangle$ . We here show that this is indeed the case. In particular, we consider a broad class of target functions  $g_{\star}:\mathcal{X}\to \mathbb{R}$  that take forms
$$
g _ {\star} \left(\mathbf {x} _ {0: N}\right) = \frac {1}{N} \sum_ {i = 1} ^ {N} F \left(\left\langle \mathbf {x} _ {0}, \mathbf {x} _ {i} \right\rangle\right) G \left(\mathbf {x} _ {0}, \mathbf {x} _ {i}\right)
$$
$$
F (t) = \sum_ {k = 0} ^ {\infty} a _ {k} t ^ {k}, \quad G (\mathbf {x} _ {0}, \mathbf {x} _ {i}) = \left\langle \widetilde {\mathbf {x}} _ {i} ^ {\otimes 3} \otimes \widetilde {\mathbf {x}} _ {0} ^ {\otimes 2}, \mathbf {A} _ {\star} \right\rangle . \tag {16}
$$
Here the scalars  $\{a_k\}_{k\geq 0}\subseteq \mathbb{R}$  and the tensor  $\mathbf{A}_{\star}\in \mathbb{R}^{d^5}$  parameterizes  $g_{\star}$ . As we will explain in Section 4.3, confining  $\bar{G}$  to be a degree-(3,2) polynomial in  $(\mathbf{x}_i,\mathbf{x}_0)$  is essential to our theoretical results. Our next theorem provides the excess risk of learning target function  $g_{\star}$  using the BRFA model (15).
Theorem 3. Given the same setting and assumptions as in Theorem 2, when the population risk minimizer gives  $f_{\star} = g_{\star}$ , with probability at least  $1 - \delta$ , we have
$$
L _ {D} \left(\widehat {f} _ {M} ^ {\mathbf {W}, \mathbf {W} _ {0}}\right) - L _ {D} \left(g _ {\star}\right) = \widetilde {\mathcal {O}} \left(\inf  _ {L} \left[ \sqrt {B \left(g _ {\star} , L\right)} \left(\sqrt {\frac {1}{n}} + \sqrt {\frac {d ^ {2} \delta^ {- 1}}{M}}\right) + \varepsilon_ {L} \| g _ {\star} \| _ {\infty} \right]\right), \tag {17}
$$
where  $\varepsilon_{L} = 1 / [2^{L + 1}(L + 1)!]$  and
$$
B \left(g _ {\star}, L\right) = \left\| \mathbf {A} _ {\star} \right\| _ {\mathrm {F r}} ^ {2} \cdot \left(\sum_ {k = 0} ^ {\infty} \left| a _ {k} \right| \cdot C _ {k}\right) ^ {2}, \quad \text {w i t h} C _ {k} = (2 L + k) ^ {(k + 3) / 2} 8 ^ {L + k / 2}. \tag {18}
$$
The proof of Theorem 3 is contained in Appendix C. We provide the intuitions of the result and an overview of the proof technique in Section 4.3.
4.2 Examples and comparison
Compared to the target functions (7) discussed in Section 3.1, functions in (16) may not express the average of arbitrary functions of  $\mathbf{x}_0$  and  $\mathbf{x}_i$ , but are well-suited to express functions of correlations. Consequently, we anticipate that the BRFA model will outperform the RFA model in learning functions of correlations. We will now present three concrete examples of target functions (16), and compare the excess risk of the BRFA model to that of the RFA model. The proof of excess risk is contained in Appendix C.3.
Example 4 (Low degree polynomials): Consider average of polynomials of  $\mathbf{x}_i$  and  $\mathbf{x}_0$ ,
$$
g _ {\star} = \frac {1}{N} \sum_ {i = 1} ^ {N} \left\langle \mathbf {x} _ {i} ^ {\otimes 3} \otimes \mathbf {x} _ {0} ^ {\otimes 2}, \mathbf {A} \right\rangle , \quad \text {w i t h} B (g _ {\star}, L) = \| \mathbf {A} \| _ {\mathrm {F r}} ^ {2} L ^ {3} 8 ^ {2 L} \quad \text {b y} \tag {18}.
$$
For any  $\eta > 0$ , if we take  $n \geq \exp(\exp(\Theta(1/\eta))$ ,  $L = \Theta((1 + \log \log n)^{-1} \log n)$ , and  $M = \Theta(d^2 n)$ , the excess risk will scale as  $\widetilde{\mathcal{O}}(\sqrt{\|\mathbf{A}\|_{\mathrm{Fr}}^2 / n^{1 - \eta}})$ .
Compared with the excess risk of the RFA model as detailed in Example 2, the excess risk bound of the BRFA model loses a factor of  $n^{-\eta /2}$ .
Example 5 (Functions of correlations): Consider a special case of functions of correlations,
$$
g _ {\star} = \frac {1}{N} \sum_ {i = 1} ^ {N} \left\langle \mathbf {x} _ {0}, \mathbf {x} _ {i} \right\rangle^ {p} \left\langle \boldsymbol {\beta}, \mathbf {x} _ {i} \right\rangle , \quad \boldsymbol {\beta} \in \mathbb {S} ^ {d - 1}, \text {w i t h} B (g _ {\star}, L) = (2 L + p) ^ {p + 3} 8 ^ {2 L + p} \quad \text {b y} \tag {18}.
$$
For any  $\eta > 0$ , choosing the same parameters  $(n, L, M)$  as Example 4, the excess risk bound scales as  $\widetilde{\mathcal{O}}(\sqrt{(\log n + p)^{(p+3)}8^p / n^{1 - \eta}})$ .
Consider the required sample size  $n_{\star}$  to reach an accuracy of 0.01. The BRFA model requires  $n_{\star} = \widetilde{\mathcal{O}}((8p + 48)^{p+3})$ , whereas the RFA model requires  $n_{\star} = \widetilde{\mathcal{O}}((4d)^p)$ . Thus, in comparison to the RFA model, the BRFA model can reduce the required sample size by a factor of  $\widetilde{\mathcal{O}}([d/(2p + 12)]^p)$ .
Example 6 (Correlation-weighted functions): Consider the function
$$
g _ {\star} = \frac {1}{N} \sum_ {i = 1} ^ {N} \cos (\langle \mathbf {x} _ {0}, \mathbf {x} _ {i} \rangle) \left\langle \mathbf {x} _ {i} ^ {\otimes 3}, \mathbf {G} \right\rangle , \quad \mathrm {w i t h} \| \mathbf {G} \| _ {\mathsf {F r}} ^ {2} \leq 1 \mathrm {a n d} B (g _ {\star}, L) = \Theta ((8 e) ^ {2 L}),
$$
where  $B(g_{\star},L)$  is bounded through the Taylor expansion of  $\cos (t)$  and (18). For any  $\eta >0$  , choosing the same parameters as Example 4, the excess risk bound scales as  $\widetilde{\mathcal{O}} (\sqrt{1 / n^{1 - \eta}})$  .
Consider the required sample size  $n_{\star}$  to reach an accuracy of 0.01. The BRFA model requires  $n_{\star} = \widetilde{\mathcal{O}}(1)$ , whereas the RFA model requires  $n_{\star} = \widetilde{\mathcal{O}}(\mathrm{Poly}(d)\exp(\sqrt{d}))$ . Thus, in comparison to the RFA model, the BRFA model can reduce the required sample size by a factor of  $\widetilde{\mathcal{O}}(\mathrm{Poly}(d)\exp(\sqrt{d}))$ .
4.3 Overview of techniques
Here we provide the intuition and an overview of the technique of Theorem 3, with the proof details in Appendix C. To show the sample complexity of learning with the BRFA model, the first step is to derive the kernel  $K_{\mathrm{BRFA}}(\mathbf{x}_{0:N}, \mathbf{x}_{0,N}^{\prime})$  associated with the infinite-width BRFA model. This kernel has a natural feature map, given by  $\{\Psi_k : \mathcal{X} \to \mathbb{R}^{d^{2k+1}}\}_{k \geq 0}$ , where
$$
\Psi_ {k} (\mathbf {x} _ {0: N}) = \sum_ {i = 1} ^ {N} \phi (\langle \mathbf {x} _ {0}, \mathbf {x} _ {i} \rangle) \cdot \mathrm {H e} _ {k - 2} (\langle \mathbf {x} _ {0}, \mathbf {x} _ {i} \rangle) \cdot \widetilde {\mathbf {x}} _ {i} ^ {\otimes k + 1} \otimes \widetilde {\mathbf {x}} _ {0} ^ {\otimes k}, \forall k \geq 2.
$$
Here  $\phi(t) = (2\pi)^{-1/2} e^{-t^2/2}$  is the Gaussian density function, and  $\mathrm{He}_k(z)$  denotes the  $k$ -th probabilist's Hermite polynomial, with detailed expression and properties given in Appendix A.1. This feature map implies the learnability of the following target function class by the BRFA model,
$$
\widetilde {g} _ {\star} \left(\mathbf {x} _ {0: N}\right) = \frac {1}{N} \sum_ {i = 1} ^ {N} \phi \left(\left\langle \mathbf {x} _ {0}, \mathbf {x} _ {i} \right\rangle\right) \sum_ {k = 2} ^ {\infty} \operatorname {H e} _ {k - 2} \left(\left\langle \mathbf {x} _ {0}, \mathbf {x} _ {i} \right\rangle\right) \left\langle \widetilde {\mathbf {x}} _ {i} ^ {\otimes k + 1} \otimes \widetilde {\mathbf {x}} _ {0} ^ {\otimes k}, \mathbf {A} _ {k} \right\rangle , \tag {19}
$$
whose RKHS norm associated with kernel  $K_{\mathrm{BRFA}}$  is bounded by  $B(\widetilde{g}_{\star}) = \sum_{k=2}^{\infty}(k-2)!k^{2}4^{k}\|\mathbf{A}_{k}\|_{\mathrm{Fr}}^{2}$ . Notice that  $\widetilde{g}_{\star}$  bears similarities to, but also distinct differences from,  $g_{\star}$  as presented in (16). The key difference lies in the  $\phi(\langle \mathbf{x}_0,\mathbf{x}_i\rangle)$  factor in  $\widetilde{g}_{\star}$ , which is hard to interpret and analyze. To obtain the excess risk bound for learning  $g_{\star}$ , we can use  $\widetilde{g}_{\star}$  to approximate  $g_{\star}$  in the  $L^{\infty}$  norm. The excess risk for learning  $g_{\star}$  can be bounded by the summation of the excess risk for learning  $\widetilde{g}_{\star}$  and the approximation error. Acquiring this approximation error bound necessitates a truncation argument of the Taylor expansion of  $1 / \phi(\cdot)$ .
We test our theory by experimentally approximating two types of target functions using the three models under investigation RFA (5), BRFA (15), and RFMLP (14). We choose the target functions to be of form
$$
f _ {1, p} \left(\mathbf {x} _ {0: N}\right) = \frac {1}{N} \sum_ {i = 1} ^ {N} \left\langle \boldsymbol {\beta}, \mathbf {x} _ {i} \right\rangle^ {p}, \quad p \in \mathbb {N}, \quad \boldsymbol {\beta} \in \mathbb {S} ^ {d - 1}, \tag {20}
$$
$$
f _ {2, q} \left(\mathbf {x} _ {0: N}\right) = \frac {1}{N} \sum_ {i = 1} ^ {N} \left\langle \mathbf {x} _ {0}, \mathbf {x} _ {i} \right\rangle^ {q} \left\langle \boldsymbol {\beta}, \mathbf {x} _ {i} \right\rangle , \quad q \in \mathbb {N}, \quad \boldsymbol {\beta} \in \mathbb {S} ^ {d - 1}. \tag {21}
$$
The first target function (20) is a specific instance of Example 2, whereas the second target function (21) has been considered in both Example 3 and 5.
In our experimental setup, we set the input dimension as  $d = 16$  and the number of tokens as  $N = 16$ . We fix the width of RFA and BRFA to be  $M_{\mathrm{RFA}} = M_{\mathrm{BRFA}} = M = 1000$ , whereas the width of RFMLP is set as  $M_{\mathrm{RFMLP}} = M(d + 1) = 17000$ . This configuration ensures an equal number of parameters across all three models. To further accentuate the test risk difference between the BRFA and RFA, in BRFA we use a bias matrix of  $\mathbf{W}_0 = 4[\mathbf{I}_{d\times d},\mathbf{0}_{d\times 1};\mathbf{0}_{1\times d},0]\in \mathbb{R}^{(d + 1)\times (d + 1)}$ , which is four times the matrix investigated in our theory. The input distribution is selected as  $\{\mathbf{x}_i\}_{0\leq i\leq N}\sim_{\mathrm{iid}}\mathrm{Unif}(\mathbb{S}^{d - 1})$  and we take  $y = f_{\star}(\mathbf{x}_{0:N})$  without any noise. We consider three representative target functions:  $f_{1,p}$  for  $p = 2,4$ , and  $f_{2,p}$  for  $p = 3$ , as per (20) and (21). We examine a list of sample sizes  $n$  from  $2^4$  to  $2^{12}$ . Prior to training with RF models, we standardize the  $y_i$ 's to have zero mean and unit standard deviation, ensuring that the trivial risk equals 1. We train the RF models using square loss with ridge regularization, selecting the ridge parameter to minimize the test error. The experimental results are displayed in Figure 2.
The left and middle panels of Figure 2 demonstrate a noticeable separation between RFMLP and the other two random-feature attention models for learning these target functions. RFMLP can hardly approximate the target function, whereas RFA and BRFA exhibit significantly better performance. This observation is consistent with our sample complexity analysis detailed in Example 2, where the
Figure 2: Test error of three RF models for learning  $f_{1,2}$  (left),  $f_{1,4}$  (mid), and  $f_{2,3}$  (right), as per (20) and (21). We set  $d, N = 16$ ,  $M_{\mathrm{RFA}} = M_{\mathrm{BRFA}} = 1000$ , and  $M_{\mathrm{RFMLP}} = 17000$ . We train the RF models using square loss with ridge regularization, with the ridge parameter selected to minimize test error. The test error is calculated using  $n_{\mathrm{test}} = 1000$  fresh samples. The figure reports the mean and normalized standard error of the test error, based on 5 independent experimental instances.
sample complexity bound of RFMLP for learning average of functions of  $\mathbf{x}_i$  is found to be  $\mathcal{O}((N / 4)^p)$  times greater than that of RFA.
The performance comparison between RFA and BRFA depends on the target functions. RFA outperforms BRFA in learning  $f_{1,2}$  and  $f_{1,4}$ , whereas BRFA outperforms RFA in learning  $f_{2,3}$ . The latter phenomenon is as we expected: as demonstrated in Example 3 and 5, BRFA is more powerful than RFA in approximating the correlation-weighted functions.
We have conducted further experiments with various other target functions, detailed in Appendix D.
In this work, we introduced and examined the expressivity of two random-feature attention models, namely RFA (5) and BRFA (15). For general classes of functions that are invariant to the permutation of key tokens  $\mathbf{x}_{1:N}$ , the excess risk of RFA (5) can avoid the dependence on sequence length, in contrast to the standard random-feature model RFMLP (14). Moreover, for specific functions that adopt the form of correlation-weighted polynomials (6), the excess risk of BRFA can avoid the polynomial dependence on the dimension. These insights enhance our understanding of the attention mechanism within a simplified context. Finally, our work left open many interesting questions for future work, such as the expressivity of softmax attention, the influence of positional encoding in expressivity, and the expressivity of multi-layer transformers.
S. Mei is supported in part by NSF DMS-2210827 and NSF CCF-2315725.
[1] E. Abbe, E. B. Adsera, and T. Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pages 4782-4887. PMLR, 2022.
[2] E. Abbe, E. Boix-Adsera, and T. Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. arXiv preprint arXiv:2302.11055, 2023.
[3] E. Akyurek, D. Schuurmans, J. Andreas, T. Ma, and D. Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661, 2022.
[4] Z. Allen-Zhu and Y. Li. What can resnet learn efficiently, going beyond kernels? Advances in Neural Information Processing Systems, 32, 2019.