Official Review · NEURIPS 2024

This paper tackles the problem of task learning with representations using transformers. It presents a theoretical proof for convergence of training of a single-layer multi-head softmax attention transformer to the global minimum of the loss. Specifically, given $N$ example tokens and their corresponding task outputs, the transformer learns to predict the task output for unseen tokens $K-N$, where $K$ is the size of the dictionary. The paper shows that when trained for sufficiently long, the transformers can perform of ridge regression on the example tokens to label the unseen tokens. (1) The paper shows that the parameters of a single-layer multi-head softmax attention transformers can be trained at linear convergence rate (Theorem 1) to the global minimizer of the training loss. This result elucidates that the remarkably efficient training that a transformer architecture offers. (2) The paper proves that once trained for sufficiently long, the transformer performs ridge regression (Theorem 2) on the example tokens to label the query tokens during inference. This is a very promising step towards understanding the generalization capabilities of transformers, which is quite different from earlier works that require the number of example tokens $N$ to be large. (3) The paper is written crisp and clean, making it very easy to follow. (1) The optimal transformer configuration $(\boldsymbol{\theta})$ is not explicitly stated; only the functional form of the output is given, which makes it difficult to the interpret the result in the parameter $(\boldsymbol{\theta})$ space. (2) The result in this paper is about the wide single-layer transformers (more heads), rather than the deep ones (more layers), and relies on Assumption 3, which holds when the number of heads, $H\ge N$ (Proposition 2). This is restrictive and there might be other ways in which the assumption can hold even if $N > H$. Some empirical demonstrations with shallow transformers on simple experiments might support this better. (1) Proposition 1 provides a way to make (almost) sure that Assumption 3 holds, but requiring $H \ge N$ seems very demanding in terms of compute. Can you comment on how to obtain similar results by simpler assumption? (2) I am curios about the motivation behind your problem formulation, i.e., task-based labelling of a finite dictionary? Is there a real-world problem where this might be of use?