A phase transition between positional and semantic learning in a solvable model of dot-product attention

Summary

The authors consider a simplified attention network with shared Query and Key matrices trained with an MSE loss and show a sharp phase transition exists when training this network. In the high-dimensional limit, this paper provides a closed form solution to the training and test loss and shows that a phase transition exists in terms of sample complexity where the model goes from solving the solution with positional information (locations of the tokens in the sentence) to semantic information (content of the tokens in the sentence). The authors show that this theoretical result shows the advantage of the attention mechanism over a fully-connected network for this task with sufficient data.

Strengths

Understanding the properties that lead to phase transitions in neural networks, and more broadly understanding training dynamics in transformer models is an important area of research. This work provides the first theoretical result showing phase transitions existing in attention mechanisms from learning, and thus opens the door for more work on the learning dynamics of transformers. I think the findings of this paper are important for interpretability research, but since it is not my area, I'm not able to strongly recommend it one way or the other. The design for the task and properties resulting in the phase transition are clear and simple. The authors are able to empirically test their results and find that models training on data nearly match the theory. The purely positional baseline provides a nice comparison for the dot product attention mechanism which *can* modulate between positional and semantic information as a result of the phase transition. Other empirical results are well justified and presented clearly.

Weaknesses

Because the results rely on the model reaching the minimum, it is unclear how well these results extend to randomly initialized networks. This is more a problem for studying training dynamics, though, which this paper does not aim to address. some parts of the paper are not well motivated or presented. In particular 4.1 and 4.2 may only reach a small audience without more description. The authors make many simplifications on the attention mechanism

Questions

Do the authors have speculations about how changes to the current architecture would affect the phase transition? For example, if the value matrix was not the identity, would the model be likely to transition to the semantic solution faster?

Rating

6

Confidence

2

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes

Summary

This paper introduces a simplifed model of attention and analyzes it theoretically, showing that there exists a phase transition between a paradigm where attention is based mostly on position to one where it is not (which they call "semantic"). I will confess to not being an expert on the methods used and so did not follow the main results (which take about a page just to state) and proofs in detail. They strike me, however, as genuinely useful and insightful, albeit with a caveat or two about some of the assumptions needed to get them to work (e.g. sequential independence).

Strengths

* Provides theoretical analyses of a model of self-attention, finding a closed-form solution. * Asymptotic analysis demonstrates a phase-shift between two minima, one which relies on position and one which does not. * The first analysis of this type to an attention layer, instead of just a feed-forward layer.

Weaknesses

* Very dense mathematically, so hard to follow for a reader not intimately familiar with the literature to which it contributes. * While toy models are indeed useful objects to study in general, there are unclear connections between some of the assumptions (e.g. independent samples of individual tokens, and the low-rank attention) and actual language modeling practice.

Questions

* How much do you think the results depend on the nature of the data? I'm thinking in particular of the fact that words are drawn _independently_: could this be part of why positional information becomes irrelevant, since the distribution at each position is the same? Do you have any expectations for slightly more realistic settings (even, e.g. independent $n$-grams instead of unigrams)?

Rating

7

Confidence

2

Soundness

4

Presentation

3

Contribution

3

Limitations

Yes

Summary

The authors state an asymptotic result characterizing the test MSE and training loss in a simplified single-layer model of dot product attention. They apply this result to study a special case in which the target attention function contains a tradeoff parameterized by $\omega$ between positional (i.e. dependent only on index location) and semantic (i.e. input-dependent) terms. Analysis of the solution characterized in the theoretical result shows that for a fixed $\omega$, there is a sharp boundary in terms of the sample complexity $\alpha$ (ratio of sample size to embedding dimension) between a semantic vs. positional parameter as the global minimum. These results are corroborated by an empirical analysis illustrating the distinct minima and comparing the empirical difference in train loss at semantic vs positional minima to the theoretical prediction as a function of $\alpha$ and $\omega$.

Strengths

The paper contributes an original result on the learning theory of attention models, along with a creative and insightful application of this result to a setting that contrasts positional versus semantic solutions. The result appears significant as a theoretical lens through which to characterize the loss landscape of an attention model, and it may have applications beyond the specific positional-vs-semantic target model studied in this paper. The exposition and empirical illustration of the result are clear.

Weaknesses

As noted by the authors in the limitations, the stated result applies to a simplified model both in terms of the structure of the attention function and in terms of the input data. It is unclear how the solutions of (7) were found in practice, as discussed in Section 5, or how it was determined that the global minimum was among this pair of fixed points for the values of $(\alpha, \omega)$ studied in the experiments.

Questions

Can the authors comment on why the simplifications in the attention model (value weights set to identity, key and query weights tied) are required for their result? The positional-to-semantic phase transition detailed in Section 5 is discovered as a consequence of Theorem 4.2 applied to a specific model (Eq (14)). Have the authors considered other models or aspects of attention-based learning that could be studied through the same lens?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations are identified and discussed in the paper.

Thank you for the reply

Thank you for the detailed reply. I appreciate the pointer to appendix D.3 which is helpful. Besides that, I definitely agree this would be out of the scope of the current work. P2: Thank you the extra explanation will be helpful. Again, I am sympathetic that it simply isn't possible to catch everyone up in such a short space, I have found that adding the extra explanations has been worth it, though. I think the authors properly address the concerns, and I'm a bit more confident after the followup discussion and reading the other reviews