Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis

Summary

This paper investigates the training dynamics of a single-layer transformer followed by a single MLP layer on a synthetic binary classification task, where the objective is to identify the co-occurrence of two specific tokens in the input sequence. They analyze the gradient flow dynamics for the case that all the attention parameters (key, query, value) and the linear layer all trainable and show that the model can achieve low loss despite the non-convexity of the objective. They identify two phases in the training, 1) the MLP aligns with the two target tokens at the start of the training and the model learns to classify all the samples correctly, 2) the attention and MLP parameters update to increase the classification margin and drive the loss to zero. They also run a small scale numerical experiment in their synthetic setup tp confirm their analysis.

Strengths

The paper makes no restricting assumptions on the weights of the transformer model and performs the analysis on the joint optimization of all the parameters. Although the paper and its proof are notation-heavy, the authors have broken down the complexity of the proof and notation in the main body to clarify the steps needed to prove the results.

Weaknesses

There are some restrictive assumptions on the synthetic data model: The vocabulary set $d$ is considered to be larger than the number of training tokens, which is not the case in realistic setups. Thus, some tokens are not visited at training time. Also, they assume, apart from the two target tokens, the remaining tokens appear at maximum once in the training set. The proof outline in the main body helps in understanding the high-level steps involved. However, it could still benefit from additional clarifications on some intermediate steps. For instance, in phases 1 and 2, it's mentioned how the alignment of the MLP with the target tokens $G^{(t)}(\mu_{1,2})$ behaves during training. However, it's not clear how this connects to the evolution of the attention scores in phase 2 as stated in Lemma 4.7.

Questions

1. As far as I understand, in your synthetic task, in the first phase of training, effectively only the MLP weights are learning the task. That is, the model can achieve 100% accuracy only by aligning the MLP weights with the relevant tokens $\mu_1,\mu_2$. So, the attention layer is not needed for identifying the co-occurrence of the tokens in this setup? 2. Can you also report validation and accuracy plots in your synthetic experiments? Does the validation loss decay at the same rate as the training loss as stated in Thm 3.3? 3. Regarding the proof sketch: a) The alignment of parameters with the target tokens is discussed in the main body. Can you also clarify how the gradients related to irrelevant tokens evolve? In particular, regarding the tokens that do not appear in the training set (since $nL\leq d$), does the model learn not to attend to them at test time? b) I find it confusing that the softmax output remains close to $1/L$ long in the training (line 320) and assigns uniform attention to all tokens in the sequence. Does this statement hold for all training samples? If yes, then how does the model learn to attend to the target tokens?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Summary

This paper studies the training dynamics of a single hidden layer transformer network (self-attention + linear MLP) trained on a binary word cooccurrence task. Specifically, given a data matrix $X \in R^{d \times L}$ representing L "words" (each column of X is a word vector of dimension d), the model must output +1 if words 1 and 2 both occur in X, and -1 otherwise. The paper shows that a transformer layer is able to learn this task, and that the training occurs in two stages: First, the linear MLP layer learns to classify data points correctly by positively aligning with the embeddings for words 1 and 2 (but without making large changes to attention matrices). Second, it drives the loss down further by using the attention matrices to positively correlate q,k,v for words 1 and 2, and anti-correlate the q,k,v for a common word (denoted word "3" in the paper) relative to words 1,2. After these phases, both the training and generalization losses go to zero (as long as embedding dimension is large enough). Overall, I found the results interesting and insightful, though not very surprising, and the practical implications of these results were not very clear to me. Thus, I currently recommend weak accept. Importantly, my primary research area is not learning theory, so my knowledge of the related work is relatively limited, and thus my review confidence is relatively low.

Strengths

- It is interesting to see that the training dynamics for this word cooccurrence task can be analyzed rigorously, with relatively few assumptions. - The theoretical results are validated with a nice synthetic experiments, that demonstrates that the two phases predicted by the theory do occur in practice.

Weaknesses

- This word cooccurrence task is very simple, and thus it is not surprising that a single transformer layer can easily learn it. - Only full gradients are considered, whereas transformers are typically trained with mini-batch Adam(W).

Questions

- What other tasks (beyond word cooccurrence) do you think could be analyzed with this methodology? - What are the implications of this result to more complex/realistic tasks, like next token prediction? - If mini-batch Adam is used during training, do the two phases still occur? - Can you add more details about the experimental setup to the main paper? - Can you add more discussion about the automatic balancing of gradients, and its significance, in a more central part of the text (e.g., section 3, not 4)?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, limitations have been discussed.

Summary

This article delves into the gradient flow dynamics for detecting word co-occurrence, demonstrating that the gradient flow approach can achieve minimal loss. The training process commences with random initialization and can be delineated into two distinct phases.

Strengths

- This article noticed an interesting phase transition during training in this special setting and demonstrates it with solid calculation and experiments. - A new property of gradient flow is noticed and contributes to prove near minimum training loss together with the analysis of softmax.

Weaknesses

The setting of empirical experiments is also simple and ideal and readers may have no idea if this is a general phenomenon during training for detecting word co-occurrence.

Questions

- In line 151 and 152, it is confusing why concentration theorems lead to the specific probability in (i) of Assumption 2.3. - Lack of explanation for $\langle w_{j_1}^{(t)},w_{j_2}^{(t)} \rangle$ in line 194. - It is not obvious why "the samples with only one target signal may be classified incorrectly as co-occurence" in line 282. - The notation in line 169 is somewhat misleading.

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

The semantic mechanism this article focuses on is not general enough in NLP and the the results cannot give instruction useful enough to guide the training process of Transformer.

Please respond to rebuttal

Dear reviewers, The author-reviewer discussion period will end soon. Please make sure you read the authors' rebuttal and respond to it. If you have additional questions after reading the rebuttal please discuss with the authors. For those who have done so, thank you! AC