Summary
This paper studies in-context nonparametric learning using transformers. In the setting used in the main result of the paper, the transformer is trained on a dataset consisting of multiple sequences/tasks. For each task, the target function $F_\beta$ is drawn from the span of a certain countable set of functions - in the setting used for the main result of the paper, the target function is drawn from a Besov space spanned by a B-spline wavelet basis. The task then consists of multiple in-context pairs $(x_k, y_k)$, where $y_k$ is $F_\beta(x_k) + \xi$, where $\xi$ is random noise. Certain technical assumptions are made on the basis functions $\psi_j$ and the coefficients $\beta_j$: the $\psi_j$ are assumed to have a property that is similar to linear independence/orthonormality, while the $\beta_j$ are assumed to decay at a certain rate.
The model in this paper's setting consists of a feature map $\phi$ applied to all of the $x_k$, followed by a linear attention layer applied to the $\phi(x_k)$ and the $y_k$. The feature map $\phi$ is assumed to be expressive enough to approximate the $\psi_j$. In the main result of this paper, $\phi$ is chosen from the class of deep neural networks (DNNs) with a logarithmic number of layers (logarithmic in the size of the feature map), and $O(1)$ width and $O(1)$ entries per layer. The class of feature maps $\phi$ is denoted as $\mathcal{F}_N$, and the overall model class is denoted as $\mathcal{T}_N$.
The main result of the paper, Theorem 4.5, is that with a sufficient number of tasks during pretraining, in-context learning with transformers will achieve the optimal minimax risk. This is shown as follows. Theorem 3.1 first gives a bound on the expected test loss in terms of the covering number of the model class $\mathcal{T}_N$, and the minimum test loss achievable by some member of $\mathcal{T}_N$. Next, the paper gives a particular construction for a set of parameters which achieves low test loss. The weight matrix for the linear attention layer is chosen similarly to the optimal weight matrix given in Zhang, et al. (2023) and nonlinear feature map outputs a subset of the basis functions.
Additional results are also given. This paper shows how the curse of dimensionality can be avoided when the target function is drawn from an anisotropic Besov space. Also, under certain assumptions, even if the target function is drawn from a Besov space with smoothness $\tau < \alpha$, it is possible to achieve the same minimax rate as in the case where the target function is drawn from a Besov space with smoothness $\alpha$. Finally, the paper gives similar results in the setting where the tokens are themselves sequences of unbounded length, and gives minimax lower bounds that match the upper bounds.
Weaknesses
I think the main weakness is that the techniques used to show Theorem 4.5 are somewhat standard.
- Theorem 3.1, which shows that the empirical risk minimizer achieves good test loss, seems to be using relatively standard arguments from previous work, based on covering numbers.
- The construction used to show that low approximation error can be achieved also seems somewhat straightforward. The weight matrix for the attention layer is chosen similarly to Zhang, et al. (2023), and the feature map outputs a subset of the basis vectors.
Questions
- In Lemma 4.4, how important is the assumption that the width/sparsity of the deep neural network is $O(1)$? Would the guarantee obtained in Theorem 4.5 change if the deep neural network is allowed to have more nonzero entries?
- Could the setting in this paper be considered similar to the linear regression setting studied by previous works such as Zhang, et al. (2023), with the difference being that the linear regression weight vector is replaced by the vector $\beta$ of coefficients for the basis functions?