We thank the reviewer for their insightful comments. We respond to the questions as follows:
>Rather it uses these two components in parallel. The reason for this choice is not clear and should be better motivated.
The motivation for the parallel structure comes from the interventions in the 1-layer BB model (Figure 3(a)), which show that the attention and MLP layers have separable utilities. Even when trained with a sequential structure, the attention layer predicts the next token only for trigger tokens, having no impact on non-trigger tokens. Conversely, the MLP layer predicts the next token exclusively for non-trigger tokens. As a result, the outputs of the attention and MLP layers belong to different subspaces. This indicates that the trained sequential structure is functionally equivalent to a parallel structure.
>In addition, the trained models use pre-layer norm and it is not clear if residual connections are used and how.
We use residual connections and sequential structure for the trained models. We will add more details in the revised version.
>In the derivation of the loss(alpha, beta), it is not clear to me how the stable distribution is defined. Could you clarify?
The stable distribution is defined as the marginal distribution of tokens in the BB task, which means the frequency of each token in the sequence generated by the BB task.
>Can you provide more intuition on the role of the softmax?
The SoftMax requires the total attention weights from each token to previous tokens to sum up to 1. However, sometimes the previous token information is unnecessary for the next word prediction. In this situation, more attention weights will concentrate on tokens with small value states, leading to the extreme-token phenomena.
>The natural question then is whether this component is actually needed in practice and why modern LLMs use it. Could you elaborate on this?
We agree that understanding whether softmax is essential is an important question, one that warrants a comprehensive analysis in future work. -. Some work has shown that a SoftMax transformer may be stabler in the training and has better performance [Shen et al., 2023; Hua et al., 2022].
>First, are the authors now using a standard transformer architecture with attention layers interleaved with mlps? Second, do the authors have any intuition on why at least 3 layers are needed?
For the first question: Yes. For the second: Our intuition is that the “OV” structure in the attention plus the MLP of the first layer is essential to the rise of massive residual states. The third layer MLP, however, removes the massiveness of the residual states so that it doesn’t hurt the read-out layer. We do not include them as we lack appropriate experiments to support the intuition. A more fine-grained analysis is an interesting future direction.
>In the construction of the token embeddings, the authors are assuming them to have dimension V
Since the output of the transformer should have dimension $V$, we set the hidden dimension $d=V$ to avoid complications. All results can generalize to arbitrary embedding dimension $d$ if we include the $V*d$ read-out layer in the theoretical model.
> it is shown that a head can have sinks yet being non-dormant (due to trigger tokens). Is this case contemplated by the theory?
Yes. Our theory only considers non-trigger tokens and focuses on explaining the sinks. The trigger tokens would have all the attention weights concentrated on the previous token.
>How did you practically identify the heads exhibiting sink tokens in real LLMs?
One way to identify the heads which exhibit sink tokens for a particular sample is to first identify which tokens have the potential to become sink/extreme tokens in a particular head. We can do this by examining the residual state peak phenomenon: the residual state norms of the (potential) sink tokens are significantly greater (usually > 10x greater) than others. Then, we examine heads which have large mean attention score on the sink tokens, i.e., compute the average attention score $\frac{1}{n-i}\sum_{j = i}^{n}A_{hij}$ for each token $i$ at head $h$ and examine the attention score on the potential sink tokens identified in the previous step. If these scores are large then the head has sink tokens, otherwise it does not. This is a straightforward and intuitive method to find such heads, but we do not claim it is optimal in any way; we leave improving this mechanism to future work.
References
Shen, Kai, Junliang Guo, Xu Tan, Siliang Tang, Rui Wang, and Jiang Bian. "A study on relu and softmax in transformer." arXiv preprint arXiv:2302.06461 (2023).
Hua, Weizhe, Zihang Dai, Hanxiao Liu, and Quoc Le. "Transformer quality in linear time." In International conference on machine learning, pp. 9099-9117. PMLR, 2022.