Towards Analyzing Self-attention via Linear Neural Network

The paper analyzes the gradient flow training dynamics of a simplified linear transformer on the histogram task. The method reduces the training of a simplified transformer to that of a linear neural network with two layers where the first layer is a diagonal matrix. The theoretical results of the paper are based on one assumption, which is experimentally justified.

This paper simplifies the training dynamics problem of a 1-layer linear self-attention layer into the joint optimization problem with two matrix variables that minimize loss like l_{df}(Q, v) = 0.5 * |Q*Diag(v) - M|_F^2. They show that this loss will decrease in exponential speed. They also try histogram tasks and show that the learned attention maps match their expectation.

This paper aims to understand the training dynamics of self-attention networks. In particular, the paper focuses on a simplified single-layer self-attention network without softmax, MLP, layer normalization, and positional embeddings. It restricts itself to a specific class of learning tasks, namely histogram-like tasks. In its simplest form, given an $N$ length sequence, this task requires the network to produce an $N$ length output sequence where $i$th output element contains the frequency of the input element at the $i$th position in the input sequence. The paper reduces the problem of learning the simplified self-attention model to learning a two-layer linear network. Subsequently, the paper analyzes the gradient flow for learning the two-layer linear network.

We thank the reviewers for their responses.

A common concern of the reviewers is that the problem considered is too simple because (1) histogram-like tasks are very simple, (2) the self-attention model strips many features of transformers like softmax.

Theoretical analysis of deep learning lags far behind empirical model. Even one-hidden-layer MLPs are not well-understood. This has necessitated that theoretical analyses simplify the problem to its core, e.g., by considering toy learning problems and simplifying the transformer, or by introducing multiple assumptions.

For example, recently the problem of analysing linear regression using one-layer transformers has attracted much attention, e.g. [4, 5, 6], and these papers study simplified models similar to ours.

1. Similarly, in the papers suggested by Reviewer 4ez9, either the results are not as strong as ours, or the proofs require stronger assumptions than ours:
   1. In [1], there are no guarantees on convergence or convergence rates. Furthermore, all the results assume that some of the parameters are fixed during training. In contrast, we provide results under no such assumptions.
   2. In [2], the results rely on the assumptions that the input sequence is long enough, and that the learning rates for decoder layer and attention layer are different, and one is much faster than the other.
   3. In [3], there are multiple assumptions and it studies a related but distinct problem, namely prompt tuning.

While these results are incomparable and so is ours, we think the problem of understanding transformer training dynamics, and more generally of neural networks, is a difficult and multifaceted one, and these papers are attacking different facets. Thus not being comparable remark in our paper was not meant as a criticism of prior work.

Linear neural networks Apart from the analysis of self-attention, our work also contributes to the analysis of linear neural networks as discussed in our paper. Perhaps surprisingly, even the analysis of two-layer linear neural networks is currently incomplete despite considerable work in the last few year. Further simplified models such as diagonal linear networks are studied in the literature, e.g., [Vaskevicius et al. (2019); Woodworth et al. (2020); HaoChen et al. (2021); Pesme et al. (2021); Berthier (2023); Boix-Adsera et al. (2023)] as cited in our paper.

Furthermore, recent work suggests that analyzing linear transformers (without softmax) is helpful even for studying more general transformers [5].

[1] Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding, 2023.

[2] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer, 2023.

[3] Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. arXiv preprint arXiv:2306.03435, 2023.

[4] J. von Oswald et al. Transformer learn in-context by gradient descent. ICML 2023.

[5] Ahn, K., Cheng, X., Daneshmand, H. and Sra, S., 2023. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.

[6] R. Zhang et al. Trained Transformers Learn Linear Models In-Context. https://arxiv.org/abs/2306.09927

[7] Ahn, K., Cheng, X., Song, M., Yun, C., Jadbabaie, A. and Sra, S., 2023. Linear attention is (maybe) all you need (to understand transformer optimization). arXiv preprint arXiv:2310.01082.

All reviewers find the paper making too many simplifications to Attention architecture as well as the training task for the presented analysis to be of interest. Given the existing work on analyzing transformer dynamics, I agree with reviewers that the setting of this work is not very interesting and hence recommend rejection.

Reject