Learning and Transferring Sparse Contextual Bigrams with Linear Transformers

Summary

The paper studies the training dynamics of a one-layer linear transformer on the sparse contextual bigram task. Using $\ell_1-$regularization and proximal gradient descent, they show that the training process goes through a sample-intensive phase and then a sample-efficient phase. They also extend the results to transfer learning.

Strengths

1. The paper is well-written and organized. 2. The theory is solid. The discussion of the two phases in training dynamics is interesting.

Weaknesses

1. Could the authors comment more about the choice of the initialization? Could you show, at least in experiments, the consequence of another choice of initializations? 2. The loss curve does not look like converging in Figure 1. It would be important to show the loss curve with more training rounds.

Questions

1. Could the authors comment more about the choice of the setting? How does sparsity affect the results? I do not find the parameter Q to have a special effect on the results besides serving as an extra parameter. 2. Can authors give some implications of their results on the training dynamics with the complete transformer structure?

Rating

6

Confidence

3

Soundness

4

Presentation

3

Contribution

2

Limitations

Although the theory is solid, and some findings are interesting, I do not find any important intuitions that can be gained from the author's discussion. In particular, 1. The current setting makes too many modifications to the original transformer structure and training practice. Such simplification is OK. However, the results do not facilitate understanding transformers' training under more realistic settings. 2. The paper lacks experiments under a more realistic setting with the standard transformer structure and training method.

Summary

This paper proposes a new data model, Sparse Contextual Bigram (SCB), to study the training dynamics and sample complexity of transformers. The paper analyzes a one-layer linear transformer trained on data generated by the SCB model using a gradient-based algorithm. They prove convergence guarantees and provide bounds on the sample complexity, showing that the training process can be split into two stages: an initial sample-intensive stage and a more sample-efficient stage. The paper also includes results on transfer learning, showing that pretraining on a related task can improve sample efficiency.

Strengths

* The proposed SCB model is a novel contribution that simplifies the analysis of transformer learning while still capturing essential aspects of language modeling. Especially the fact that training can be split into two stages with different data requirements and the implications for transfer learning are quite interesting. * Although I am not familiar enough with these type of theoretical analyses to fully verify the correctness of the proofs, I found the paper relatively easy to follow.

Weaknesses

* The analysis focuses on a simplified linear transformer model and a specific gradient-based algorithm. The generalizability of the results to more complex architectures and optimization methods is unclear. This limitation is acknowledged in the paper, and I don't think it's reasonable to expect a theoretical analysis without any simplifications. Hence, this might not actually be a weakness, although it is a limitation. * The paper could benefit from more intuitive explanations of the proofs to make the theoretical results more accessible to a wider audience.

Questions

* How do you think the results would change if a more complex transformer architecture (e.g., with multiple layers or non-linear activations) were used? * Can the results on transfer learning be extended to cases where the pretraining and downstream tasks have different structures or distributions? * Could you provide more intuition behind the proofs, either in the main text or in the appendix?

Rating

7

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have adequately addressed the limitations of their work in Appendix A. They acknowledge that the model and algorithm studied are simplified compared to those used in practice and that the empirical evaluation is limited to synthetic data. They also mention the need for further research to extend the results to more complex architectures and real-world datasets.

Summary

The topic of this paper is understanding transformers. Since this is an ambitious goal, the authors make some reasonable simplications. Specifically, they study the training of a **one-layer linear** transformer on synthetic datasets generated by a novel data-generating model, called **Sparse Contextual Bigram (SCB)**. SCB extends traditional bigram by also conditions the next token generation on prior earlier positions determined by the last token. Such design of SCB allows study on both contextual (last token conditioned early token dependence Q) and global (transition matrix P) knowledge. Based on this data-generation model, their analysis shows the the one-layer linear transformer training process can be divided into two stages: an initial sample-intensive stage followed by a more efficient one. In transfer learning scenarios, the first stage can be skipped if there's significant correlation between the pretraining and downstream tasks.

Strengths

- The paper is well-scoped. The structure is easy to follow. The authors clearly state their simplification in order to analyze the thus complicated transformer training process. - The proposed data-generating model, SCB, is novel to me. It extends the classic bigram in a smart way so that it encompasses both contextual and global knowledge.

Weaknesses

- Architectural simplification. The simplification of transformer into a one-layer linear version could make the conclusions less useful or even invalid. Particularly, the softmax over the attention weights and the layernorm before the output are known to be important for the impressive empirical performance of transformer. It would be great if the authors can explain how their simplification of architecture can impact their conclusions, either theoretically or empirically. - Experimental setup. The experiments in Sec 5 are with a vocabulary size of N=3. Such a small vocab size is far from practical ones (~32K-250K). I doubt the conclusions on the stages would still holds across different orders of vocab size. And it is unclear to me how the hyper-parameters are chosen from line 321-line328. It would be great if the authors can explain why they use this specific setup instead of a more systematic investigation of hyper-parameters. - Complexity. What is the complexity of the SCB data-generation model? How does it differ from a real natural language distribution?

Questions

- See weakness. - What's $I_T$ in equation 2?

Rating

5

Confidence

4

Soundness

2

Presentation

4

Contribution

2

Limitations

They discuss the limitations in appendix A

Summary

The paper studies a natural problem that generalizes the bigram model for language generation. In the bigram model, each token $x_t$ is only dependent on the previous token $x_{t-1}$. Therefore, to learn to predict the next token, one only needs to look at the bigram frequencies $P[i,j] = \Pr[x_t = i \mid x_{t-1}= j]$ which is simple to learn. The authors generalize this in two steps: 1) First there is a mapping $Q: [K] \to [T]$ that maps each token value to a sparse distribution over the indices. The algorithm picks an index and looks at the value at this index. 2) Then they use a bigram model $P$ to predict the next token based on the value chosen in step 1. This is an interesting model and can capture natural tasks performed by transformers as they process long contexts. It requires the ability to learn the transition matrix $P$ as well as the mapping $Q$. The authors try and show is that simplified linear model can also learn this efficiently in $N,Q$. This effectively means that they take one layer of the attention matrix and and remove soft-max and the non-linear activation in the final MLP layer. Finally they use the one hot shot embedding which greatly simplifies the model further. This results in a very simple form: $x_{T+2} \approx VXa^{x_{T+1}}$ where $V$ is the trainable vector and $a$ is the $k^{th}$ column of a trainable matrix $A$. The key is to understand how this will be transformed using standard gradient descent and a mean squared error loss. The main result is that with the right initialization as well as a staged use of proximal gradient descent, one can learn this model using linear transformers.

Strengths

The paper is well written and the assumptions are clearly stated. The analysis is mostly in the appendix but from what I have glanced at, it looks thorough. The bulk of the analysis is to deal with the noise introduced by the empirical gradient from having a small batch. The model itself is a nice tractable model that still captures a nontrivial task that we observe performed by modern transformers.

Weaknesses

Naively, one can imagine a much simpler model such as Hidden Markov Model or other sequence being capable of learning the sparse bigram model. The main interest in the paper is that a sufficiently simpler version of transformer can "provably" learn such a model. The assumptions do seem necessary to make the problem sufficiently simple and hence approachable. However this means that the actual algorithm and model is quite far from what is used in practice. I would have liked to see a bit more empirical work to see if the introduction of non-linearity greatly reduces the sample complexity. In particular, I would have liked to see some empirical work showing that the actual attention matrix has patterns that are similar to what is learnt in the linear model itself.

Questions

NA

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

NA