Transformers Provably Learn Two-Mixture of Linear Classification via Gradient Flow

Q1 Have you tried to measure this phenomenon and the training dynamics for any of the setting extensions I mentioned in W2?

A Please see our response in W2. We have conducted experiments on the multi-mixture setting, and the results can be found in Appendix A.4.

Q2 In Figure 1, we see that the neuron alignment grows with $c_k$ immediately from the beginning. This seems to be in contradiction with line 367, especially since I believe the weight initialization in experiments was the same as in the theory. What makes this difference?

A The neuron alignment for $c_k$ indeed stays flat at the beginning. Because the duration is not very long, it is not very obvious on the plot.

Q3 In figure 1c, how would the curves for $c_1$ being query look like? Can you provide a figure with these included?

A The shape of the curve for $c_1$ being query is similar to the curve for $v_1$ being query. The difference is that the magnitude of the scores for $c_1$ being query is a lot smaller than the scores for $v_1$ being the query. We have included a plot with $c_1$ being query in Figure 4 in Appendix section A.3 in the revision.

Q4 Where do the slight asymmetries in Theorem 4.1 come from when you consider the averaged initial value model?

A The asymmetries are defined in Lemma E.4 in the Appendix. In Lemma E.4, besides the scores, we also need symmetries between the self-score variables such as $v_1^\top W_{Q+}^{(T_2) \top} W_{Q+}^{(T_2)} v_1 \approx - v_1^\top W_{Q+}^{(T_2) \top} W_{Q+}^{(T_2)} v_2$. The slight asymmetries come from the initialization in both the averaged initial value model (original version) and the infinite-width transformer model (revised version), since $v_1^\top W_{Q+}^{(0) \top} W_{Q+}^{(0)} v_1 > 0$ whereas $v_1^\top W_{Q+}^{(0) \top} W_{Q+}^{(0)} v_2 = 0$. Thus, the exact symmetry cannot hold and the asymmetries in the self-score variables will propagate to the score variables during training. Fortunately, we are able to show that these asymmetries will not become very big after the training.

W3 Some aspects of the presentation should be improved. To provide a couple of examples:

• Q It could be clearer if authors re-named classes to something else as they are confused with the classification classes.
  A Thank you for your suggestion. We have changed ``class" to ``group".

• Q The assumption of averaged initial value model is not discussed enough. In particular it is not clear where any asymmetry at all comes from if we have this assumption.
  A Thank you for pointing this out. We also realized that the averaged initial value model is an unnatural assumption. In our revision, we were able to remove this assumption via mean-field technique by considering an infinite-width transformer, which is considered to be a more natural and easily justifiable technique in theoretical literature. With infinite-width transformer, the averaged initial value model assumption can be exactly satisfied at initialization, see Lemma C.1 in Appendix C. We have checked that replacing the original averaged initial value model by infinite-width transformer doesn't introduce much changes in the proof. The discretization to finite-width case can be handled easily in Appendix G.
  Furthermore, by adapting to the mean-field technique, we remove the $1/\sqrt{m}$ factor in the softmax attention and we initialize the attention weights by $\mathcal{N}(0, \omega^2 \frac{1}{m})$. One of the consequences of removing the $1/\sqrt{m}$ factor in the softmax attention is that now stage 2 of the training becomes shorter and only runs in $O(\log (1/\omega))$ time. We also updated the experiment plots in Section 3.1 accordingly to reflect this change.

• Q The proof in appendix should be more verbose and make some more explicit statements as intermediate steps. For instance, the separability of the consistent samples after phase 2 is only proved implicitly and is never explicitly stated in the proof. This makes it hard to make look-ups of specific proof parts mentioned in the main body, because the proof itself is an ocean of infinite formulas.
  A Thank you for your suggestions. First of all, the separability of the data can be seen by considering the neuron weight $w_+ = c_1 + c_2 + v_1 + v_2$ and $w_- = c_1 + c_2 - v_1 - v_2$ which can create a constant margin. We have made this statement more precise in the proof. In the meantime, we have added more explanation in various places in our proof. We will continue to polish our proof in the Appendix.

• Q The use of $g$ is somewhat peculiar. I suspect the authors to either make a few typos when using $g$ or not to define it properly.
  A The function $g$ is defined on line 175 as the negative gradient of the cross entropy loss.

Typo 1 The property 4 in Definition 2.1 is inconsistent with the verbal description in the rest of the paper as this way, you would leave most of the positions actually empty (zeros).

A Thank you for pointing it out! You are right. What we should write is for each $l_2 \in [L] \setminus \{l_0, l_1\}$, sample $k' \sim \textnormal{Uniform}([K] \setminus \{k\})$ and set $X_{l_2} = \epsilon v_{k'}$, where $\epsilon \stackrel{i.i.d.}{\sim} \textnormal{Uniform}(\{\pm 1\})$.

Typo 2 line 170: You should write $\nu$ instead of $X^{(i)}$.

A The input to the softmax function should be a vector and thus it should be $X$ (notice that $X^\top W_K^\top W_Q \mu$ is a vector and $l(\nu)$ will index the corresponding entry to $\nu$).

Typo 3 In lemma 5.2, the use of $g$ is somewhat weird. Please double-check what you meant there. The same holds for lines 427, 430.

A In Lemma 5.2, we have $g(X_{k,s,s'}) = -\ell'(f(X_{k,s,s'}))$. On line 427 and 430, we use $g^{(t)}$ to denote $g(f^{(t)}(X_{k,s,-s}))$ (by symmetry in the averaged initial value model (old version) and infinite-width transformer (revised version), $f^{(t)}(X_{k,s,-s})$ are the same for $k \in [2], s \in \{\pm 1\}$).

Typo 4 Line 408: based on my counting the numbers should be 16, 10, 10. Correct me if I am wrong.

A Yes! You are right about the counting. Thanks for pointing it out!

Typo 5 In line 60 you talk about class identifier getting all the signal from classification features, but later in your theory you consider the class identifier's role as the key to be crucial. I think it better to rewrite this sentence so that the roles are exchanged.

A Thank you for the suggestion! We have modified our introduction accordingly.

We thank the reviewer for providing the highly inspiring feedback. In the revised paper, we have made changes based on your review comments and have highlighted all revisions in blue fonts.

W1 While I appreciate the setting of the paper in general, there are a few questionable decisions when it comes to architecture. In particular, in the main part of the paper, it is a bit unusual that the value matrix basically maps to 1D outputs and the representations are averaged even before applying ReLU. I would expect averaging only after the application of the ReLU and also like to see what would happen if more than 1D value matrices would be used. However, bigger issue comes with Lemma 6.1. In particular, the order in which the operations are made in the transformer defined in the proof is so unusual I am not sure it can be interpreted as transformer anymore. The summation over $l_2$ should definitely come as the inner-most, otherwise we cannot even talk about attention. Then should come the summation over $k$, only after that should come the application of $\sigma$ and the summation over $l_1$ should be the outer-most. I would ask the authors to redo the statement and proof of Lemma 6.1 so that the summation over $l_2$ is the inner-most or otherwise justify their architectural choice.

A Thank you for this suggestion! We made a new construction in the revision in Lemma 6.1 and we highlight the new construction in blue.

W2 Although the presented phenomenon is intuitive, it is unclear if and how does it generalize to more general settings. In particular to multi-mixture of linear regression, where I would expect (almost) the same thing happen. It is unclear why the symmetry is broken and therefore it is also unclear whether the described phenomenon even holds the same way in multi-class setting. The same question holds for another natural extensions of the considered setting -- multi-class classification, presence of noise, non-orthogonality of features or class imbalance and more natural weight initialization. I wouldn't mind that the paper only studies the simplest setting if there was evidence that the knowledge obtained can be transferred to the more general settings but in this case it is not clear whether the same mechanism even works in extended settings.

A Great question! We first point out that the technique we used in our analysis can be applied to the settings the reviewer mentioned. We next discuss the extensions that the reviewer points out.

Multi-mixture of linear regression: We further investigated why the asymmetries between groups arise in the scores for 3 groups. We found out that the asymmetry is due to the randomness in the scores at initialization. We tried to increase the embedding dimension of $W_K, W_Q$ and the asymmetries between groups start to go away, see Figure 5(c) and Figure 5(d). (In the infinite-width transformer in our revision, the randomness will be reduced if we increase the width.) We further investigated the multi-group case and we found that with more groups, the magnitudes of between-group scores like $c_2^\top W_{K+}^{(t)\top} W_{Q+}^{(t)} v_1$ will become smaller and smaller with more groups and the magnitudes of same-group scores such as $c_1^\top W_{K+}^{(t)\top} W_{Q+}^{(t)} v_1$ remain big, see Figure 5(d), 6(b), 7(b) in Appendix A.4. On the other hand, our techniques are generalizable to study this setting.

Multi-class classification setting: This is indeed an interesting extension and our proof technique can be adapted to study this setting. In this case, the transformer model needs to be modified to output a vector instead of a single scalar in the binary case. We expect that the layer-wise learning behavior (i.e., first neurons learn and then the attention learns) will still occur in the multi-class classification setting. On the other hand, how the softmax attention will behave in this setting needs further investigation.

Presence of noise, non-orthogonality of features or class imbalance and more natural weight initialization: We expect that the training dynamics is similar when the data contains appropriate amount of Gaussian noise or we use Gaussian initialization with sufficiently small variance to initialize the neuron weight. On the other hand, non-orthogonality of the signals or class imbalance will definitely break the symmetry and the training dynamics will be certainly different. For example, (Huang, Cheng & Liang, 2024) studied the class imbalanced case in in-context learning and they found that class imbalance will introduce more phases in training dynamics. We also expect that more training phases will arise due to data dominance/under-representation.

Huang, Yu, Yuan Cheng, and Yingbin Liang. "In-context Convergence of Transformers." Forty-first International Conference on Machine Learning, 2024.

Thank you. There are a few unclear steps in this answer.

• As far as I understand Lemma E.1 does not prove an upper-bound on the score variables, instead it assumes they are bounded and proves some relative size bounds. How can we see the boundedness of the score variables from Lemma E.1?

• Next, in the penultimate paragraph you say that the base case can be established by the existence of $T_{1.5}.$ How do you prove that such a time exists?

• Finally, in the last paragraph you assume $||\delta x(t)||\le 1$ but that itself is waht you get in the end of the sentence and also what you need in your agrument in the first place. So it seems you are assuming what you want to prove.

Please clarify these points.

We thank Reviewer qpZd for the review. In the revised paper, we have made changes based on the reviewers' comments and have highlighted all revisions in blue fonts.

Q I had trouble understanding the setting of mixture of linear classification. Authors only spend one paragraph in the 2nd page discussing about the model. But I think it should be elaborated more in details to motivate the readers.

A Thank you for pointing this out. Below the formal definition of the mixture of linear classification task in Definition 2.1, we elaborate our explanation for the data model with an example. We also provide a high level description of the mixture of linear classification task in the introduction in our revision.

Q I might have missed the part while reading the paper, but I wonder how training dynamics discovered in the paper corresponds to the actual reading procedures of the human beings? For instance, using the example in the paper: “Bob is watching TV in the living room. Where is Bob?” How do people catch the correspondence between Where and Living room and how does this correspond to the training dynamics in the paper? I wonder if authors have any thoughts on this.

A Great question! For this example, the supervised loss function will be small if the answer to ``where" is ``living room". Our analysis of gradient flow suggests that the gradient of the loss with respect to the attention matrix will point to the direction to align the query token of ``where" with the key token of ``living" and ``room" during the training process, thus driving the attention matrix to align those two. This will ultimately derive the loss to be small and provide good answer. Such process is consistent with our theoretical analysis of the behavior of the attention matrix.

We point out that the example given in the introduction serves as an motivating example to justify our choice of mixture of linear classification task for studying how transformers utilize connections between words to solve certain tasks. In this example, there is a connection between the word ``where" and ``living room".

Q Is the training dynamics in the paper also observable in the vanilla transformer case? The authors consider two-headed symmetric transformers, but in real practice vanilla transformer is commonly used. If this is the case, I think this paper’s contribution is significant.

A Thank you for the question. In our work we observed layer-wise learning behavior (in our case, neuron weights learn first and then the softmax attention modules learn) during transformer training. Similar behavior has been observed in previous work such as [1], where the authors discovered that there is a layer-wise learning behavior for multi-layer transformer. Further, our insight about the training dynamics of the self-attention modules, where the attention matrices are driven by gradient flow to align relevant tokens for classifications, will also carry over to multi-layer transformers. Such an observation has been made in [2].

We would also like to point out that the purpose of our work is to understand how transformers especially the self-attention modules leverage the connections between tokens to solve certain task. At the same time, we aim to provide theoretical tools that can serve as building blocks and are generalizable to other settings when we analyze the training dynamics of transformers.

Q Do authors have any insights on the multiple layer cases?

A Our work currently focuses on a single layer transformer, and we discovered a layer-wise learning behavior in our setting. Similar layer-wise learning behavior has been observed previously in [1]. Further, our insight about the training dynamics of the self-attention modules, where the attention matrices are driven by gradient flow to align relevant tokens for classifications, will also carry over to multi-layer transformers. Such an observation has been made in [2]. We also believe that our techniques and tools developed in this work can be generalized to study more complicated architectures like multi-layer transformers.

[1] Nichani, Eshaan, Alex Damian, and Jason D. Lee. "How Transformers Learn Causal Structure with Gradient Descent." Forty-first International Conference on Machine Learning.

[2] Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, Simon Shaolei Du. JoMA: Demystifying Multilayer Transformers via Joint Dynamics of MLP and Attention, ICLR 2024.

We thank Reviewer Kkd8 for the review. In the revised paper, we have made changes based on the reviewers' comments and have highlighted all revisions in blue fonts.

Q The description of the experimental setup is missing, but I assume the results are all based on synthetic data. Therefore, justification with real data is necessary.

A Yes, the experiment results are based on the data distribution we introduced in Definition 2.1. We are actively developing code base to test our Algorithm 1 on real-world dataset. If our experiment results would not come out before the rebuttal deadline due to the complexity of the real world experiments, we commit to include these results in the final version of the paper (if it is accepted) or in the future version of the paper.

We also point out that the primary contribution of this paper lies in the theoretical analysis, which can help us understand and demonstrate the mechanism of neurons and attentions, and how they evolve during training.

Q The theoretical insights for guiding the training and deployment of Transformers are unclear in practical applications.

A Thank you for the comment. The main theoretical insight from our gradient flow analysis is that during training, the gradient flow (derived from the loss function) is able to effectively guide the transformer's weights toward learning the appropriate representations required to solve the underlying task. Specifically, the gradient of the attention weights actively steers the attention matrix to align with the correct feature directions for classification. This finding is particularly significant because, despite the highly nonconvex nature of the loss function, its gradient is strong enough to overcome the nonconvex landscape and consistently points toward the global minimum. Practically, this insight implies that to train transformers with specific desired properties in practice, we can design the training objective to shape the gradient accordingly, thereby directing the transformer's weights toward the desired solutions.

Another insight from our theoretical analysis is that the training process of transformer can often occur in phases. In our case, the neuron weights of the transformer learn first and after that the softmax attention modules start to learn. This implies that in some cases, we can train transformers layer-wise to reduce the training cost in practice.

In addition, we would like to point out to the reviewer that the purpose of our work is to understand how transformers especially the self-attention modules leverage the connections between tokens to solve certain task. At the same time, we aim to provide theoretical tools that can serve as building blocks and are generalizable to other settings when we analyze the training dynamics of transformers.

We thank Reviewer fKbv for the review. In the revised paper, we have made changes based on the reviewers' comments and have highlighted all revisions in blue fonts.

Q This paper only consider learning two-mixture of binary classification, and the data distribution is of size $d\times 3$. The main theorem is based on Algorithm 1, an algorithm designed from empirical observations on a synthetic dataset but not tested on any real-world datasets. So, I question whether analyzing Algorithm 1 and this specific choice of dataset is really helpful. I suggest the authors to apply Algorithm 1 on some real-world dataset, and shows that it works. (do not need to be better than standard training).

A Thank you for your suggestions. We would like to point out that although the proof idea for our theorem is inspired by a data distribution of size $d\times 3$, the actual proof technique is developed for data distributions of inputs with arbitrary sizes. The mathematical rigor of the proof ensures that the theorem is valid and applicable to more general distributions.

We are actively developing code base to test our Algorithm 1 on real-world dataset. If our experiment results would not come out before the rebuttal deadline due to the complexity of the real world experiments, we commit to include these results in the final version of the paper (if it is accepted) or in the future version of the paper.

We also point out that the primary contribution of this paper lies in the theoretical analysis of Algorithm 1, which can help us understand and demonstrate the mechanism of neurons and attentions, while serving as an accurate model to reflect how the weights in transformer evolve.

Q Although this paper's title is "two-mixture", there should be more experiments and discussions on $K\geq 3$ than current Appendix A.3.

A Thank you for your suggestions. We have included more experiment results and discussion for the case with $K\geq 3$ in Appendix A.4 in our revision. Our experiment results indicate that the training dynamics of 3-mixture, 4-mixture and 8-mixture models are all showing some similarities as that for 2-mixture model, demonstrating the broad applicability of our developed theoretical results.