Transformers need glasses! Information over-squashing in language tasks

Authors: federico  barbero
andrea  banino
steven  kapturowski
dharshan  kumaran
joão g m araújo
google  deepmind
alex  vitvitskyi
razvan  pascanu
petar  veličković
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 17:17:45

Abstract: We study how information propagates in decoder-only Transformers, which are the architectural foundation of most existing frontier large language models (LLMs). We rely on a theoretical signal propagation analysis---specifically, we analyse the representations of the last token in the final layer of the Transformer, as this is the representation used for next-token prediction. Our analysis reveals a representational collapse phenomenon: we prove that certain distinct pairs of inputs to the Transformer can yield arbitrarily close representations in the final token. This effect is exacerbated by the low-precision floating-point formats frequently used in modern LLMs. As a result, the model is provably unable to respond to these sequences in different ways---leading to errors in, e.g., tasks involving counting or copying. Further, we show that decoder-only Transformer language models can lose sensitivity to specific tokens in the input, which relates to the well-known phenomenon of over-squashing in graph neural networks. We provide empirical evidence supporting our claims on contemporary LLMs. Our theory points to simple solutions towards ameliorating these issues.

Main paper: 
 introduction: 
 In recent years the field of Natural Language Processing (NLP) has been revolutionised through the introduction of Transformer-based architectures [30] . Large Transformers trained on some version of next-token prediction, known as Large Language Models (LLMs), have demonstrated impressive performance across different tasks, including conversational agents [10, 19] , understanding multi-modal inputs [1] , and code completion [16] . Most contemporary LLMs specifically focus on the decoder part of the original Transformer architecture, and are commonly referred to as decoder-only Transformers. Consequently, we focus primarily on such models in this paper.However, despite the impressive performance of Transformers, recent works have uncovered surprising failures that may point to fundamental issues in their architecture. For instance, ˚Work performed while at Google DeepMind.• Impact of floating point precision: we explore how low floating-point precision exacerbate the identified theoretical issues, causing them to manifest even in relatively short input sequences. • Empirical validation of theoretical analysis: our theoretical findings are supported by realworld experiments conducted on contemporary LLMs, demonstrating practical implications of the limitations we identified.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: We provide extensive evaluations and analysis to support our claims. Also we specify which assumption we made and justified them either with ablation or explanations or both. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations.A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
background: 
 In this work, we study a class of Transformers which we believe forms the basis for a large number of current LLMs. We let Q, K, V P R nˆd be the query, key, and value matrices respectively on n tokens and d dimensions. We denote with q i , k i , v i P R d the d-dimensional query, key, and value vectors of the i-th token. We let p ij P R 2e be the 2e-dimensional positional encoding information between tokens i and j. We focus on the case in which the positional encodings are bounded, which is the case for the large majority of positional encodings used in practice [26, 30] . . This specific interleaving of components is often referred to as a Pre-LN Transformer [34] . We can view the output of the ℓ-th layer of a Transformer as a sequence of d-dimensional vectors v pℓq " pv pℓq 1 , . . . , v pℓq n q. Importantly, due to the causal attention mechanism, the vector v pℓq j , will only depend on elements v pℓ´1q i for i ď j. We can group the attention weights into an attention matrix at the ℓ-th layer which we define element-wise as Λ pℓq ij " α pℓq ij . This is a row-stochastic triangular matrix that can also be interpreted as a probabilistic directed graph. After the last transformer block a normalization is applied to the token representations:y i " norm 3 ´vpLq iWe note that the next-token prediction usually depends purely on y n -the final representation of the last token.Existing theory on Transformers. The theoretical representational capacity of Transformers has become a popular area of study, providing interesting results on what classes of problems they are able to model. For instance, it has been pointed out that Transformers are not Turing-complete, but one can apply modifications which make Transformers Turingcomplete under certain assumptions [6] . Works have also shown that Transformers using 'hard attention' which replaces softmax with one-hot vectors alongside the use of infinite precision makes Transformers Turing-complete [22] . This contrasts with our work, which focuses on the more standard setting of Transformers using soft-attention and finite precision, and shows the limitations imposed by it.Works have also tried to study transformers capabilities through the lense of formal languages, such as Weiss et al. [33] , which develops a computational model of what transformers can represent in an analogous way to how Recurrent Neural Networks are associated with finite automata, and then derive an implementable programming language that represents 1 for extra details that model. Following that, Delétang et al. [7] place transformers within the Chomsky Hierarchy, showing that they are quite limited and cannot learn the decision problem for simple languages, which prompted authors to show that Transformer LLMs can perform substantially better if they generate a number of decoding tokens linear in the problem input size, through scratch-pad, Chain-of-Thought (CoT) or similar [18] . Finally, Peng et al. [21] show that the Transformer block with finite precision is fundamentally limited in its ability to represent compositional functions and solve simple problems that require it. Our work similarly analyses the Transformer's inability to solve simple computational tasks, and proves that even with techniques like Chain-of-Thought that inability persists as it is inherent to the combination of architecture, next-token prediction, and limited floating point precision.Decay in attention mechanisms. Works have also studied the limitations of self-attention by showing that it can reach pathological states that limit what transformers are able to learn. For instance, it has been show how a great reduction in the attention entropy can lead to unstable training if occurring early, but even when occurring later in training it can still lead to significantly lower performance [35] . Further, it has been shown that specific tokens can strongly concentrate attention, leading to transformers being unable to learn to process simple languages, like PARITY and DYCK [12] . Our work will similarly focus on showing how Transformers end-up effectively ignoring many tokens in their input which leads them to fail to solve simple computational problems, studying such a phenomenon by directly analysing the representational capacity.Over-squashing. Graph neural networks (GNNs) are neural networks designed to operate over graph structures. Importantly, Transformers, may be seen as types of attention-based GNNs operating over specific types of graphs. The difficulties of propagating information over a graph have been thoroughly analysed, with a notable phenomenon being that of oversquashing [2, 29, 8, 3] . Over-squashing refers to the fact that propagating information over certain graphs that exhbit 'bottlenecks' is likely to induce a 'squashing' of information. This can be made more precise by studying this effect via the notion of a commute time [11] -the expected number of steps that a random walk takes to travel from a node to another node and back. Information travelling between nodes with higher commute time will be squashed more.A common way to measure over-squashing is by looking at how sensitive the representation x pLq v of a node v after L GNN layers is to the initial representation x p0q u of another node u. In particular, the partial derivative Bx pLq v {Bx p0q u may be shown to decay, especially for nodes with high commute times between them. Our work may be seen as acting as a bridge between the well-studied phenomenon of over-squashing in GNNs and the loss of information we analyse in decoder-only Transformers specifically for language tasks. Note that this type of derivation is typical in the study of vanishing gradients for recurrent models as well [14, 5, 20, 13] .Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: Limitations are explained in the appropriate section. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper. • The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speechto-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
motivating examples: 
 This section presents a series of experiments focused on copying and counting tasks. These experiments reveal surprising failure cases in modern decoder-only Transformer architectures, providing concrete evidence that motivates the theoretical analysis presented in the following sections.We start by providing motivating examples that show surprisingly simple failure cases of frontier LLMs specifically on copying (Section 3.1) and counting (Section 3.2) tasks. By copying we specifically mean tasks that involve the 'copying' or 'recalling' of a single or multiple tokens from the prompt. Instead, by counting, we mean the task of counting how many times a specific token appears in a sequence. We focus our evaluation on Gemini 1.5 [10] as our frontier LLM (referred as Gemini) and later analyse the internal representations of the open-sourced Gemma model [27] . The goal is to showcase intriguing failure cases which will motivate our signal propagation analysis.In this Section, we present surprising results on simple copying tasks. In particular, we focus on tasks that involve the copying of a single token -i.e. what is the token occurring at a particular position? The copy of a single token is in principle the most straightforward type of copying task, but still requires the LLM to accurately identify the token based on a prompt and to then propagate its information correctly.Importantly, we study cases in which the LLM is prompted to copy tokens either at the start or at the end of a sequence. We avoid tasks that involve the copy of tokens at the 'n-th' position as most frontier LLMs do not have absolute positional information, making it very challenging for them to solve tasks that require absolute position. We focus on tasks that involve sequences of 'zeros' and 'ones' growing in length with specific patterns.In Figure 2 (a), we prompt Gemini to copy the last element of a sequence '1 . . . 10' or the first element of a sequence '01 . . . 1'. The answer for both is zero, but we progressively grow the number of ones. We observe how the task seems considerably easier when asked to return the first rather than the last element. Surprisingly, already at a sequence length of only 300 elements, Gemini incorrectly starts to output 'one' when trying to copy the last element. In Figure 2 (b), we show that providing hints in the form of: " *Hint* It's not necessarily a 1, check carefully", helps significantly with the performance. Finally, in Figure 2 (c), we show that replacing the constant sequence of ones with alternating ones and zeros seems to also help. We refer to the Appendix (Section C.1) for further details on the experiments. These three motivating experiments seem to point towards a type of vanishing of information, caused by the growing number of ones dominating the sequence. Interestingly, such a The large density at 100 suggests that Gemini is likely not counting, but instead possibly performing some crude form of subitising. vanishing of information effect (a) seems to depend on the position in the sequence, (b) seems to be affected by the prompting, and (c) by the items that make up the sequence. We will later argue how all three of such observations can explained by our theoretical analysis.We now turn our attention to counting problems, i.e. tasks of the form -given a specific sequence, how many times does a particular token appear? Such problems are related to copying in the sense that they also require careful consideration of individual tokens in the sequence as ignoring even a single token may potentially lead to an incorrect output.We consider four different tasks: (i) Summing 1 `¨¨¨`1, (ii) Counting the number of ones in a sequence of ones, (iii) Counting the number of ones in a sequence of ones and zeros, with ones being sampled with 70% probability, and (iv) Counting the number of times a specific word appears in a sentence. We consider predictions of an LLM which (1) Is instructed to only output the answer, (2) Is prompted to break down the problem (CoT-no-shot), and (3) Is prompted to break down the problem with few-shot in-context examples (CoT-few-shot). We refer to the Appendix (Section C.1) for a more detailed description of the tasks.Results are presented in Figure 3 . It is clear that the performance rapidly deteriorates with the sequence length. It is also interesting to see that the error seems to increase with the sequence very rapidly. For instance in task (i), the LLM is quite likely to predict the value of '100' once the sequence reaches a size around or larger than 100. Such an observation provides motivating evidence for the argument that Transformers may not be in fact mechanically counting but rather perform a type of crude subitising. This explains why arguably 'common' numbers such as 100 are much more likely to be outputted by the LLM and why in tasks such as (i) and (ii) the values near 100 have relatively lower error. This does not happen in task (iii) as the response should actually be around 70% of the sequence length due to the sequence sampling procedure, explaining why the absolute error actually seems to increase around a sequence length of 100. Figure 4 further showcases this issue, more clearly showing how 100 is by far the most common response.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean.  
representational collapse: 
 We start our theoretical analysis by showcasing a type of loss of information which we call representational collapse. More precisely, we show that under certain conditions, we can find distinct sequences such that their final representations of the last token at the last layer become arbitrarily close as the sequence length increases. As Transformer models operate over finite machine precision, this points to a fundamental representational incapacity of Transformers to distinguish certain prompts if the sequence is long enough.The key intuition is that if two sequences are similar everywhere except at the last token, as the sequences get larger, their final representations will become closer and closer until they reach a critical point which is below floating point precision. In other words, solving certain tasks would require infinite floating point precision. We will later show how this phenomenon is not only theoretical, but also occurs in practice on sequences of reasonable length. In the Appendix (Section 4), we relate representational collapse to the L 1 distance -or total variation -between the softmax distributions of the two sequences. We start by presenting a result that shows that the L 1 difference tends to 0 as the sequence length grows, under some assumption on the sequences. We point to the Appendix (Lemma B.2) for the complete statement. Lemma 4.1 (Informal). Consider two sequences x, x ˚P R n such that lim nÑ8 |x n ´xn | " 0. Then, the L 1 difference of their softmax tends to 0.We now show, using Lemma 4.1, that we can find distinct sequences that will have arbitrarily close final representations. In particular, as language models often operate in low floating regimes, i.e. bf16, this can practically become catastrophic. The result is summarised in Theorem 4.2, which describes what we call representational collapse in this work. The complete statement is reported in the Appendix (Theorem B.3).Let v p0q P R nˆd be a sequence and v ˚p0q P R pn`1qˆd be another sequence equal to v p0q with the last token of v p0q repeated. Assume that the positional encoding information decays to 0 with the distance. Then, their representations become arbitrarily close as n increases.Theorem 4.2 shows that it becomes increasingly challenging for a Transformer to distinguish two sequences that only differ via a repeated last token. We note that the repetition of the last token is a technical consideration to show this direct representational collapse. As we will later show in Section 5.1, it is particularly problematic in general to depend on the last token due to a type of topological 'squashing' present in decoder-only Transformers.We report experiments showcasing representational collapse by measuring the internal representations of Gemma 7B [27] . For two sequences v p0q and v ˚p0q we report their difference in representation at the last layer v pLq ´v˚pLq 8 averaged out over each head, alongside the minimum and maximum over each head. Figure 5 shows the collapse occuring on (a) prompting the model to count the number of ones in a sequence of ones, with one having an additional one, and (b) prompting the model to count the number of ones for a sequences with digits sampled uniformly ending with either a single one or two ones. The repeated digits seem to make the collapse occur much sooner with a sequence length of around 50 being near machine precision, while varying the digits seems to delay such a collapse, but a downward trend is maintained with respect to the sequence length.A common technique used to speedup the inference of an LLM is that of quantisation, a process that constructs an approximate version of an LLM that operates over lower precision datatypes. This helps drastically improve the inference speed of LLMs as modern accelerators produce significantly more FLOPs over lower precision datatypes. Of course quantisation usually comes at a cost. Our theoretical analysis points towards a potentially catastrophic loss in representation due to quantisation. In particular, a lower machine precision will mean that the convergence of representations in Theorem 4.2 will occur much sooner, and consequently the LLM will not be able to distinguish even shorter sequences.In practice, the direct application of theoretical results is made more complicated due to tokenisation. In particular, a sequence of repeated tokens '11111' for instance may not be necessarily tokenised into 5 distinct '1' tokens. In principle, this should help alleviate the direct collapse of the representations. Tokenisation in general makes it more of a challenge to study such phenomena as it adds an additional layer of complexity to the experimental analysis. In our experiments, we took tokenisation into consideration and attempted to mitigate its effects.A simple solution to representational collapse. An important consequence of Theorem 4.2 is that it is challenging for a Transformer to deal with a long sequence of repeated tokens. A practical solution is to this issue is to introduce additional tokens throughout the sequence to help keep the representations distant. We provide direct evidence of this in Figure 5 (c, d ), where we prompt the model on a simple copying task of a long string of ones. While the (a) "How many ones are in the following sequences?" Followed by a sequence of ones.(b) "How many ones are in the following sequences?" Followed by sampled digits.(c) "Can you copy the following number?" Followed by a sequence of ones.(d) "Can you copy the following number?" Followed by a sequence of ones with commas.  
over-squashing in language tasks: 
 In this Section, we study a more general phenomenon related to representational collapseover-squashing. In particular, we are interested in analysing how information from the input sequence affects the information contained within the representation of the last token in the final layer-the representation ultimately used for next-token prediction. For this reason, we study the quantity By n {Bv p0q i which measures how sensitive is the final token to an input token at position i.In graph neural network theory, the decay of such a partial derivative is often associated with the 'squashing' of information, leading to the phenomenon of over-squashing [29, 8, 11] , a problem related to the well-known vanishing gradients problem in RNNs [14, 5, 11] . The over-squashing analysis we carry out in this work is particularly challenging due to the flexible nature of the attention mechanism and the many components that are part of decoder-only Transformers. Consequently, we make two simplifying assumptions in our analysis: (i) We summarise the effect of layer normalisation via a constant β i for the i-th layer norm component, and (ii) the attention weights are treated as independent of the input. Such simplifications are not strictly necessary for our analysis, but they greatly simplify the resulting bound we derive and do not detract from the two key takeaways: (1) the sensitivity to an input token depends on its position in the sequence and (2) the sensitivity to an input token depends on the attention weights. The result is summarised in Theorem 5.1. The full statement is reported in the Appendix (Theorem B.5). Let C ą 0 be some constant and ᾱpℓq i,j "α pℓq i,jβ2 `δi,j , then:EQUATIONTheorem 5.1 provides intuition on how information propagates in a decoder-only Transformer.In particular, there is a topological aspect present in the bound which is directly controlled by the attention mechanism. More concretely, the sensitivity depends on the sum of the weighted paths between the token i at the input and the final layer. In other words, for tokens coming sooner in the sequence, there will be more opportunity for their information to be preserved. This is clear for instance for the last token, which will only be preserved by attention mechanism if the attention n Ñ n is large at every layer L, i.e. there is only one path. The paths instead grow very quickly for tokens coming sooner in the sequence. A related observation, in terms of path counting, was also made for deep RNNs [13] . We note that such a bound explains the better performance when copying elements at the start of the sequence in Figure 2 (a), why hints help in Figure 2 (b), and why repeating the final elements within the sequence also helps in Figure 2 (c ).This analysis leads to an interesting limiting case described in Proposition 5.2, that shows a type of exponential vanishing that can occur in some degenerate cases in which y n depends only on the starting input token v p0q 1 . Fortunately, there are many mechanisms which prevent this from happening, but believe this to be an interesting limiting case which is a direct consequence of the topology of the causal attention mechanism. Further, it provides an interesting connection between the spectral theory of directed graphs and causal attention mechanisms. We report the formal statement in the Appendix (Proposition B.8). Proposition 5.2 (Informal). Under certain assumptions on the effect of the normalisation and on the attention weights, in the limit of layers L Ñ 8 the output representation will only depend on the first input token.U-shape effect. Theorem 5.1 in part also helps to explain the empirically observed Ushape effect-the observation that LLMs seem to perform better at retrieval tasks when the information to be retrieved is located either near the start or the end of the sequence. In fact, due to the topology of the causal mechanism, we find from Theorem 5.1 that tokens at the start of the sequence have more opportunity for the information to be maintained at the end. The final tokens being also easier instead can be explained from the recency bias that is learnt by the attention mechanism during training. In auto-regressive next-token prediction, it is in fact reasonable to assume that tokens that are closer to the end will be more important and this is likely a bias that is learnt during training by the LLM. 
counting: 
 We finally highlight another representational problem that arises specifically in counting problems. Our analysis points to a fundamental difficulty that emerges from the normalisation of the softmax. In particular, the normalisation of the softmax makes it hard for a model to take into account the length of a sequence. This is exacerbated by the fact that positional encodings are often normalised and thus relative, meaning that they also do not hold absolute positional information. Intuitively, counting is a problem that requires some notion of 'unboundedness' of the representations, whilst the normalisations used inside a Transformer work against this.We start by showing that without causal masking and positional embeddings, a Transformer is immediately unable to count the number of tokens in a sequence, highlighting a pathological issue which stems directly from the softmax normalisation. We note that similar issues have been already pointed out [e.g. 22]. We show the result in Proposition 6.1 and report the full statement in the Appendix (Proposition B.9). Proposition 6.1. A Transformer without positional encodings and a causal attention mechanism is immediately unable to count.While causal mechanisms and positional encodings help to break such representational issues, they break the permutation invariance of the Transformer, meaning that the representations will be heavily miss-aligned with the task, something which has been shown to hinder performance [9] . As permutations grow factorially with sequence length, this makes it practically very challenging for a decoder-only Transformer to learn such a property simply from the data. This explains the extreme incapacity of counting highlighted in Section 3. Further, as a corollary of Theorem 4.2, we have that even if a model would be able to generalise perfectly, the problem of representational collapse points to an impossibility result in counting regardless. The result is summarised in Corollary 6.2, with the full statement in the Appendix (Corollary B.10). Corollary 6.2 (Informal). Counting in certain situations becomes impossible due to representational collapse and finite floating point precision. Corollary 6.2 shows how our main result on representational collapse points to practical issues when it comes to certain styles of prompts. When paired with low floating point arithmetic precision, representation collapse becomes problematic. 
conclusion and future work: 
 In this work, we first presented surprising failure cases of LLMs on simple copying and counting tasks. We then discussed how such failure cases can be explained by studying what can be contained inside the representation y n and in particular how information may be lost. This lead to the unvealing of two phenomena : representational collapse and over-squashing. We showed how we can measure these phenomena in practice and proposed simple solutions to help alleviate such information loss.We believe that this work uncovers an interesting framework which can be used to study failure cases of Transformers and LLMs more generally. We believe that our analysis could be extended in many practical different directions, for instance by understanding how to directly measure over-squashing or how to best use this newly-found understanding to improve current Transformer models. In our work, we focused on pointing out informationpropagation issues in Transformer-based architectures, but we hope that the findings may help better understand and improve language models available today. 
 
 Appendix: This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here. In particular the limitations highlighted in the work can pose some issue in terms of reliability of LLMs, so highlighting these issues can help fixing them.We provide formal statements and proofs for the results shown in the main text. We follow the order in which they are presented in the main text. In Section B.1, we present the proofs on representational collapse (Section 4), in Section B.2 the proofs on over-squashing (Section 5) over-squashing, and finally in Section B.3 the proofs on counting (Section 6).We start by showing that adding a new element to a sequence results in the softmax value of a specific token to decrease. In particular, we consider the case in which the tokens are bounded and show that we can use this to construct an upper bound on the softmax value for any token.Lemma B.1. Consider a vector a P R n´1 and two scalars b, c P R. Let x " ra cs T P R n and x ˚" ra b cs T P R n`1 with all entries bounded. Then, softmax pxq n ą softmax px ˚qn`1 . Moreover for any p ą 0 we can find large enough n P N `such that | softmax pxq n śoftmax px ˚qn`1 | < p.Proof. We directly compute:softmaxpxq n " exppcq ř n´1 k"1 exppa k q `exppcq softmaxpx ˚qn`1 " exppcq ř n´1 k"1 exppa k q `exppbq `exppcqAs we assume that the entries are bounded, we have that ř n´1 j"1 exppa j q `exppcq ă ř n´1 j"1 exppa j q `exppbq `exppcq, therefore softmaxpxq n ą softmaxpx ˚qn`1 . For the second part of the statement, we compute:|softmaxpxq n ´softmaxpx ˚qn`1 | " exppcq ř n´1 k"1 exppa k q `exppcq ´exppcq ř n´1 k"1 exppa k q `exppbq `exppcq ď exppcq ř n´1 k"1 exppa k q `exppcq ` exppcq ř n´1k"1 exppa k q `exppbq `exppcq ă p for some p ą 0, as in the last step the summands tend to 0 as n Ñ 8. We therefore have that |softmaxpxq n ´softmaxpx ˚qn`1 | Ñ 0 as n Ñ 8 and for large enough n |softmaxpxq n ´softmaxpx ˚qn`1 | ă p for any p ą 0 due to the previous statement.Total variation. We now study a quantity known as the total variation between two distributions of interest. Given two categorical distributions µ, ν supported on the same space, we define their total variation δpµ, νq -or equivalently L 1 norm ∥µ ´ν∥ 1 , as:EQUATION)The total variation is a distance between probabilty distributions. We note that oftentimes the quantity above is strictly called the L 1 norm, while 1{2 of such a quantity the total variation. For the scope of our work, the factor of 1{2 is not important so we ignore it and use the two terms synonymously. Interestingly, the total variation is intimately related to the KL-divergence, pointing towards potential connections to information theory. We leave such a connection to future work.We now study how the total variation between two softmax distributions behaves in the limit of the sequence length. In particular, we show that if the positional encoding information decays to 0 with the sequence length, then the total variation between the two sequences goes to 0 with n. Such a result is presented in Lemma B.2Lemma B.2. Consider two sequences x, x ˚P R n such that lim nÑ8 |x n ´xn | " 0, with x i , x i bounded. Let y, y ˚P R n be the softmax of x and x ˚respectively. Then, as n Ñ 8 the total variation tends to 0, i.e. lim nÑ8 δpy, y ˚q " 0.ř n i"1 e x i and Z ˚" ř n i"1 e x i be the partition functions for x and x ˚, respectively. We start by bounding the quantity |Z ´Z˚| . In particular, let ϵ ą 0, we first claim:EQUATIONConsider some n 0 ě 1 such that 1 ´ex i ´xi ď ϵ{2. We note that this always possible as |x i ´xi | Ñ 0 by assumption. We compute:n ÿ i"1 e xi ´ex i " n0´1 ÿ i"1 e xi ´ex i `n ÿ i"n0 e xi ´ex i " n0´1 ÿ i"1 e xi ´ex i `n ÿ i"n0 e xi 1 ´ex i ´xi ď n0´1 ÿ i"1 e xi ´ex i `ϵ 2 n ÿ i"n0 e xi ď n0´1 ÿ i"1 e xi ´ex i `ϵ 2 Z ď ϵZWhere the last step comes from the observation that the first sum is fixed and Z is unbounded with n. Therefore, for n large enough, we can also bound the left summand by Zϵ{2. The same argument also holds when bounding with Z ˚instead of Z, leading to the claim in Equation 3.We now proceed with the following computation:δpy, y ˚q "n ÿ i"1 e xi Z ´ex i Z ˚ " n ÿ i"1 Z ˚exi ´Ze x i ZZ ˚ " n ÿ i"1 Z ˚exi ´Ze xi `Ze xi ´Ze x i ZZ ˚ ď n ÿ i"1 Z ˚exi ´Ze xi ZZ ˚ ` Ze xi ´Ze x i ZZ ˚ " n ÿ i"1 |Z ˚´Z | e xi ZZ ˚`Z |e xi ´ex i | ZZ " |Z ˚´Z | ZZ ˚Z `n ÿ i"1 |e xi ´ex i | Z ď ϵ minpZ, Z ˚q Z ˚`ϵ minpZ, Z ˚q Z ď 2ϵwhich concludes the proof.We are now ready to show the main result on representational collapse. In particular, we show that given two sequences of length n and n `1 where the second sequence is the same as the first with a final repeated token, their representations become arbitrarily close. Importantly, we require that the information from the positional encodings decays to 0 as the distance grows between tokens. The final token repetition is important as otherwise the residual connection in the Transformer would not make the representations necessarily converge.). Let x P R n´1ˆd be an underlying growing token sequence. Let v p0q " rv v a s T P R nˆd and v ˚p0q " rv v a v a s T P R n`1ˆd be two sequences for a final repeated token x a P R d , with all token representations bounded. Further, assume that the positional encodings decay with distance to 0. Then, for large enough n P N `, we have that the representations are under any ϵ:||v pLq n ´v˚pLq n`1 || 1 ă ϵ.Proof. We note that since the sequences are identical up to the n-th element, it is sufficient to only check the representations of the final elements in both sequences. We therefore compare the n-th element of z p0q with the n `1-th element of z ˚p0q :z p0q n ´z˚p0q n`1 1 " ÿ iăn α p0q n,i v p0q i `αp0q n,n v p0q a ´˜ÿ iăn α ˚p0q n`1,i v p0q i `´α ˚p0q n`1,n `α˚p0q n`1,n`1 ¯vp0q a ¸ 1 " ÿ iăn ´αp0q n,i ´α˚p0q n`1,i ¯vp0q i `´α p0q n,n ´α˚p0q n`1,n ´α˚p0q n`1,n`1 ¯vp0q a 1 ď ÿ iăn α p0q n,i ´α˚p0q n`1,i ` α p0q n,n ´α˚p0q n`1,n ´α˚p0q n`1,n`1 ď ÿ iăn α p0q n,i ´α˚p0q n`1,i ` α p0q n,n ´α˚p0q n`1,n ` α ˚p0q n`1,n`1 " δ ´αp0q n,: , α ˚p0q n,:n ¯`α ˚p0q n`1,n`1 ă ϵWe assume for simplicity that the values are unit norm. This is not crucial as otherwise one would equivalently just need to consider additional constant factors as we assume the token representations are bounded. We note that the term δ ´αp0q n,: , α ˚p0q n,:n ¯goes to 0 with n Ñ 8 thanks to Lemma B.2 and our assumptions on the positional encodings. Similarly, the term α ˚p0q n`1,n`1 goes to 0 due to Lemma B.1. We have also used the fact that v p0q n " v ˚p0q n`1 by construction to ignore the residual connection. As v pℓ`1q i " ψ pℓq ´norm pℓq 2 ´zpℓq i ¯¯`v pℓq i , the sequences will have arbitrarily close final token representations when entering the next layer. The result then follows via a simple inductive approach on the layers.We also highlight a negative decay result which highlights why the assumption on the positional encodings is important. In particular, we show that given two sequences x " p1 0 1 0 . . . q and x ˚" p0 1 0 1 . . . q, the total variation of the softmax does not decay to 0. This implies that there may be solutions to representational collapse depending on how the positional encodings are chosen.Proposition B.4. Consider two sequences x " p1 0 1 0 . . . q P R n and x ˚" p0 1 0 1 . . . q P R n for n even. Let y and y ˚be the softmax of x and x ˚, respectively. Then the total variation δpy, y ˚q does not tend to 0 as n Ñ 8.ř n i"1 e x i and Z ˚" ř n i"1 e x i be the partition functions for x and x ˚, respectively. We directly compute: lim nÑ8 δpy, y ˚q " lim nÑ8 n ÿ i"1 y i Z ´yi Z ˚ " lim nÑ8 n ÿ i"1 y i n 2 e `n 2We now present our results on over-squashing. In our derivations, we assume that the attention coefficients are independent of the values and that we can summarise the effect of the layer norms via a constant factor. These assumptions are not necessary for the same derivation process to hold, but they greatly simplify the obtained bound and help more clearly point out the main takeaways. (including CoT). Let σ ψ be the maximal Lipschitz constant of any ψ pℓq , and ᾱpℓq j,i " 1 β pℓq ´αpℓq j,i `δj,i ¯the normalized attention coefficient, then:EQUATIONProof. Note that for j ě i we have: represent the effect of layer normalization i at the ℓ-th layer and σ ψ pℓq the Lipschitz constant of ψ pℓq . For the case when j ă i due to the causal mechanism we have that Bv pℓq j {Bv pℓ´1q i " 0. We compute the following bound:By n Bv p0q i " 1 β 3 ÿ k1 . . . ÿ k L Bv pLq n Bv pL´1q k L L´1 ź ℓ"2 Bv pℓq k ℓ Bv pℓ´1q k ℓ´1 Bv p1q k1 Bv p0q i " 1 β 3 ÿ k1ěi . . . ÿ k L ěk L´1 Bv pLq n Bv pL´1q k L L´1 ź ℓ"2 Bv pℓq k ℓ Bv pℓ´1q k ℓ´1 Bv p1q k1 Bv p0q i ď 1 β 3 L ź ℓ"1 ˜σψ β pℓq 2 `1¸ÿ k1ěi . . . ÿ k L ěk L´1 ᾱpL´1q n,k L L´1 ź ℓ"2 ᾱpℓ´1q k ℓ ,k ℓ´1 ᾱp0q k1,i " C ÿ k1ěi . . . ÿ k L ěk L´1 ᾱpL´1q n,k L L´1 ź ℓ"2 ᾱpℓ´1q k ℓ ,k ℓ´1 ᾱp0q k1,iwhere we let ᾱpℓq j,i "α pℓq j,i β pℓq 1 `δj,i and C " 1 β3 ś L ℓ"1 ˆσψ β pℓq 2 `1˙.We note that in this derivation, we use simplifying assumptions on the layer norms and attention coefficients, more specifically we assume that they are independent of the v i s. Of course, there is nothing stopping us from avoiding such assumptions and pushing the partial derivatives inside these components as well. The drawback is that this would add a great deal of additional complexity to the result and potentially distract from what we believe are the two key takeaways: (1) the position of the token matters, and (2) the attention coefficients matter.We now show some results on the spectral theory of matrices which relate to causal attention mechanisms. We emphasize that in this work, we view causal attention mechanisms as triangular row-stochastic matrices. We show that these matrices have interesting spectral properties. Lemma B.6. A row-stochastic triangular matrix A has 1 as its largest eigenvalue. Moreover, such eigenvalue has multiplicity 1 if each row except the first has at least 2 non-zero entries.Proof. We start by showing that A cannot have eigenvalues λ ą 1. We then provide an eigenvector with eigenvalue 1. We finally show that such an eigenvector is unique if each row has at least 2 non-zero entries.Assume λ ą 1 for some eigenvector ϕ, we then have that Aϕ " λϕ. Consider ϕ i " max k ϕ k ą 0. Now pAϕq i " ř jďi A ij ϕ j " λϕ i . As the sum is a convex combination, the result cannot be larger than the already maximal element ϕ i . As λ ą 1, we however have that λϕ i ą ϕ i which is a contradiction and we conclude that λ ď 1.It is easy to find an eigenvector that always has eigenvalue 1. Consider a vector x which is a constant vector of 1s. Then pAxq i " ř jďi A ij " x i , therefore x is an eigenvector with eigenvalue 1.Finally, we show that when each row is non-zero, the only eigenvector is the constant-valued eigenvector. Consider the largest entry y i ą 0, then we have that pAyq i " ř jďi A ij y j " y i . Again, as this defines a convex combination, we must have that all tokens that i points to (i.e. the non-zero entries) are also equal to y i . The condition that each row has at least two non-zero entries is important as it means that the condition y i " y j is true for all tokens.Proof. Let A, B P R nˆn be two row-stochastic matrices. We compute:ÿ j pABq ij " ÿ j ÿ k A ik B kj " ÿ k A ik ÿ j B kj " 1The final statement follows immediately from the fact that the product of two triangular matrices is triangular.We now show that under specific conditions, our over-squashing bound converges to a steady state in which the final token y n only depends on the initial input token v p0q 1 as the number layers tends to infinity, i.e. L Ñ 8.Proposition B.8. Let β pℓq 1 , β pℓq 2 " 1, β 1{L 3" 4, σ ψ = 1. Furthermore, for simplicity, let the attention coefficients be equal at each layer and such that each row except the first of the causal mechanism has at least two non-zero elements. Then, we have as L Ñ 8 that By n {Bv p0q i " 0 when i ‰ 1 and By n {Bv p0q i " 1 when i " 1. In other words, y n will only be sensitive to the first token.Proof. Let the associated attention matrix be Λ. We start by re-writing the following:By n Bv p0q i ď 1 β 3 L ź ℓ"1 ˜σψ β pℓq 2 `1¸ÿ k1ěi . . . ÿ k L ěk L´1 ᾱpL´1q n,k L L´1 ź ℓ"2 ᾱpℓ´1q k ℓ ,k ℓ´1 ᾱp0q k1,i " ˜L ź ℓ"1 « 1 β 1{L 3 ˆσψ β 2 `1˙ˆ1 β 1 Λ `I˙ff ¸n,i " ˜" 1 2 pΛ `Iq ȷ L ¸n,iWe now point out that Λ " 1 2 pΛ `Iq is row-stochastic and with our assumptions is diagonalizable into Λ " ΨΣΦ. In particular, by Lemma B.7, also ΛL is row-stochastic and each entry is non-negative. We now use the Perron-Frobenius theorem for non-negative matrices [23] , which guarantees us that all eigenvalues λ k of Λ are bounded such that |λ k | ď 1. In particular, thanks to Lemma B.6, we know that there is a unique eigenvector (the constant eigenvector ψ n ) with eigenvalue λ n " 1. Denote the left eigenvectors by ψ k and the right eigenvectors ϕ n , we therefore have:lim LÑ8 ΛL " ÿ k λ L k ψ k ϕ T k " ψ n ϕ T n .In particular, one can check that ϕ T n " r1 0 . . . 0s, meaning that ψ n ϕ T n has as first column a constant vector of 1s and every other entry 0. This completes the proof.We finally show in this section our final results that apply specifically to counting tasks. We start by highlighting a potential difficulty that the softmax layer encounters when counting, namely that the normalisation used makes it hard for it to preserve a notion of magnitude present in the sequence. Proposition B.9. A Transformer without positional encodings and a causal attention mechanism is immediately unable to solve the counting problem.Proof. We show this statement by demonstrating that the only information preserved about the 'count' by the attention mechanism will be the ratio of the elements present in a sequence. In particular, sequences with the same ratio of tokens will be assigned the exact same representation -this applies as we specifically study an attention mechanism without positional encodings and causal masking. Of course, having the same ratio of elements does not mean that the count will be the same, for instance the sequences '10' and '1100' have the same ratio of digits but clearly different counts.Consider a sequence of two values, v p0q zero and v p0q one , with n 0 and n 1 being the number of zeros and ones respectively. We ignore in our calculations the MLPs ψ and the normalizations norm as these don't affect the argument. As this specific attention mechanism is permutation equivariant, the initial zero tokens will all be mapped to: one (to avoid the trivial case), we notice that the attention mechanism alongside the MLP ψ define an isomorphism between sequences at different layers, updating all zeros and ones to a different value vector. The critical fact is that the representations only depend on the ratio between n 0 and n 1 , meaning that sequences of different lengths (therefore different counts) will have the exact same representation. This is respected at each layer, meaning that the LLM after L layers will assign the same representation to different sequences as long as they have the same ratio. This points to a loss of representation for the counting problem.z p1q zero " ÿ j exp ´qp0qT zero k p0q j řw exp ´qp0qT zero k p0q w ¯vp0q j `vp0q zero " n 0 exp ´qp0qT zero k p0q zero n0 exp ´qp0qT zero k p0q zero ¯`Corollary B.10. Consider a task in which the goal is to count how many v a tokens there are in the sequence. Let v p0q " rv v a s T P R nˆd and v ˚p0q " rv v a v a s P R pn`1qˆd . Due to representational collapse, at least one sequence will be given the wrong count for large enough finite n.Proof. This statement is a direct consequence of representational collapse. In particular, as y n and y n will be indistinguishable for large enough n, the Transformer will be forced to make a mistake for at least one of them. This points to an impossibility result of counting on certain sequences due to floating point error. This holds regardless of the positional encodings used (as long as they satisfy the required decay conditions) and causal mechanism.The prompting done on Gemini 1.5 in our work does not require custom resources as we use hosted Gemini instances. We run a local version of Gemma 7B on modest hardware to analyse the internal representations.We detail the way in which we execute the prompting for the various experiments.Counting experiments. For the sum experiment we prompt as: Please perform the following sum: seq. Please give the answer on the final line exactly as 'The final answer to your maths question is: xxxx', where 'xxxx' is your answer.. For the ones and zero sequences, we similarly prompt as Please count the number of ones in the following sequence:seq. Please give the answer on the final line exactly as 'The final answer to your maths question is: xxxx', where 'xxxx is your answer. For the word counting experiment, we prompt as Please count the number of times 'car' appears in the following sentence: 'seq'. Please give the answer on the final line exactly as 'The final answer to your maths question is: xxxx', where 'xxxx' is your answer. Copying experiments. For the copying experiments, we use the following prompt: Consider the following sequence: seq. What is the last digit in this sequence? Please answer exactly as 'The answer to your question is: <ANSWER>' and change appropriately the sequence as described.We commit to releasing the code we have used to generate the prompts in the near future.We report similar results for the counting experiments using Gemma [27] in Figure 6 and 7 . Compared to Gemini 1.5, Gemma seems to answer less accurately on the counting prompts. To provide further experimental evidence of representational collapse with other positional encodings, we experiment using the original sinusoidal embeddings from [30] . We sample key, query, and values from a Gaussian distribution with variance σ 2 " 1{d, with d the dimension of the embeddings. We set d " 64 and otherwise follow the exact structure of the decoderonly Transformer presented in the original Transformer paper. We experiment with a single attention layer and check the convergence of the representations of the final token between a sequence of length n and a sequence of length n `1 in which we simply copy the final token. We present the results in Figure 8 . We see how also for sinusoidal PEs with key, queries, and values randomly sampled, the convergence still occurs.Figure 8 : Convergence behaviour with a synthetic Transformer experiment. We sample the key, query, and values from a Gaussian distribution and apply the traditional sinusoidal PEs from [30] . We apply a logarithmic scale on the y-axis.Finally, we test the decay of the total variation of the softmax distributions of two growing sequences, to experimentally verify Lemma B.2. We sample a sequence x of length n with values uniformly distributed in the range r0, 1s. We then create x ˚by adding to the first k " 200 elements of x noise which is uniformly sampled between r0, 0.1s. In Figure 9 , we show how the total variation between their respective softmax distributions decays with the sequence length.Figure 9 : Total variation decay of softmax distributions with growing sequence length. We sample n elements uniformly from r0, 1s and then create a related sequence by taking its first k " 200 and adding to these elements noise sampled uniformly from r0, 0.1s. We measure the total variation between their softmax distributions. It is clear how the total variation decays with length, in accordance with Lemma B.2. Error bars show minimum and maximum over 5 seeds.We include a synthetic experiment where we verify the occurrence of the representational collapse phenomenon with different Positional Encodings, namely: Alibi [24] , the original Absolute Positional Encodings (APE) [30] , and No Positional Encodings (NoPE) [15] . In our synthetic experiment, we sample n queries and keys independently directly from a standard Gaussian. We then construct a related sequence of length n `1 by repeating its last element. We report the L 1 distance between the two sequences after a single decoder-Transformer layer, as done for the other representational collapse experiments. We consider a Transformer with a hidden dimension of 64, a single attention head, and we apply normalisations to simulate layer norm. The Transformer is not trained, but only used to simulate the propagation of information of queries and keys sampled from a Gaussian distribution. The results are shown in Figure 10 . Representational collapse seems to occur with all 4 positional encodings, with the convergence of the representations happening at a similar sequence lengths.We would like to highlight that our condition on the decay of RoPE necessary to fulfill the requirement of Theorem 4.2 is inspired by claims of the decay of RoPE coming from the original work by Su et al. [26] . However, recent work has shown that such claims may not be strictly always upheld [4] , as the original claims relied on very specific conditions on the queries and keys. We are of the opinion; however, that the range of synthetic and real-world experiments in this work support our representational collapse claims in practice with RoPE. A more precise mathematical treatment of RoPE specifically is therefore left as future work.Figure 10 : We sample n queries, keys, and values independently from a standard Gaussian, applying different positional encodings. We then construct sequences of length n `1, by repeating the n-th token. We report the L 1 distance between the last tokens of the sequences of length n and n `1 after one decoder-only Transformer layer. We set the hidden dimension to 64, use a single attention head, and normalise appropriately to simulate the effects of LayerNorm. The y-axis is shown in log-scale.We ablate the prompt structure specifically for the copying task. In particular, we consider the prompts: (Type 1) "What is the last digit of the following sequence? Please answer exactly as 'The answer to your question is: <ANSWER>' ". Here is the sequence: {seq} and (Type 2) "Please answer exactly as 'The answer to your question is: <ANSWER>'. What is the last digit of the following sequence? {seq}". The results are presented in Figure 11 . We find that the prompt indeed does affect the performance on the task, as the prompt affects the distribution of the attention over the layer. However, for both types of prompts, the model ends up failing, in accordance with our theory. We also show for completeness, in Figure 12 , that representational collapse occurs in Gemma 7B also for the 'Type 1' prompt.A fundamental limitation of an attention mechanism that leverages the softmax function is that it cannot remain sharp, especially as the sequence length grows [31] . This is in fact a key intuition that we exploit to show our result on representational collapse. A good way to address representational collapse and the related phenomenon of over-squashing is then that of limiting the spread of the softmax function, by directly limiting the amount of tokens the attention mechanism pays attention to. This mechanism is often referred to as a local sliding window and is a major architectural change present in Gemma 2 [28] . We believe that such an architectural change elegantly addresses representational collapse and over-squashing at the source as it avoids the issues that come with growing token sequences -something which our theory often exploits.Note that we rely on an abuse of notation. We ignore the linear projections used to compute the value v plq i from the output of layer below l ´1. This will not change our derivations, but would otherwise make notations more cumbersome. 


Decision: The paper indicates some major shortcomings suffered by Transformer architectures via both empirical studies and theoretical analysis. The work could be further improved if the authors could explain their discoveries with more practical examples (in contrast to the toy tasks used in the paper), and it would be even greater if the authors could provide some insights for how to remedy the shortcomings. Despite these weakness, all reviewers are positive to the work, and I also agree that the work sheds light on the future studies of Transformers. Therefore, I would like to recommend an acceptance.

Reviews: review: paper_summary: paper_summary: In the paper, the authors first discuss a phenomenon occurring in LLMs that they call "representational collapse". They provide empirical evidence of the phenomenon in state-of-the-art language models and they provide a theoretical justification for it. They then show that decoder-only transformers exhibit what is known in graph network theory as "over-squashing", which may cause loss of information in the final prediction.
strength_weakness: strength_weakness: The observed phenomenon, while interesting, has limited scope. Also, the representational collapse result is hardly surprising to me. It feels quite trivial that in the limit of a very long sequence, the importance of a single token out of many repeated ones becomes negligible.
questions: questions: 1) While the paper claims that the transformer architecture that they analyze is the most popular in state-of-the-art LLMs, it seems to me that they only consider relative positional embeddings and not absolute ones. To the best of my knowledge, current LLMs like GPT and Llama all use absolute positional embeddings. Would your results also extend to this case? In any case, I would highlight the use of relative positional embeddings more prominently in the main text.
2) In the actual prompts used for the experiments, the sequence does not appear at the end of the prompt but somewhere in the middle. Would your theoretical results somehow generalize if the repeated token is not the last one? Would anything change in the experiments if the sequence appears exactly at the end?
limitations: limitations: N/A
paper_summary: paper_summary: The paper provides a theoretical and empirical analysis of the final representations of transformers, revealing the phenomenon of representational collapse.
strength_weakness: strength_weakness: - There is no strong solution to this problem. Although the authors provide a simple approach, it is challenging to understand how to apply this approach in a practical setting.
- The problems seem to be tied to the positional embeddings, as the authors stated. NoPE (https://arxiv.org/abs/2305.19466) and no causal mask would make this problem impossible. However, there are no experiments on how certain positional embeddings might be better than others.
questions: questions: - What do the authors believe is the breadth of tasks affected by this fundamental issue in transformers?
- What are the positional embeddings used for the experimental settings? Can these affect the representational collapse?
- Minor: Some figures, like Figure 2, have a very small font size, making them hard to read.
limitations: limitations: See Questions and Weaknesses.
paper_summary: paper_summary: This paper presents an in-depth analysis of decoder-only Transformers, focusing on their limitations in handling information propagation. The authors identify two key phenomena: "representational collapse" due to "over-squashing" (line in Graph Neural Networks). These issues lead to a significant loss of information, especially in tasks involving counting and copying, which require precise handling of individual tokens. The paper combines theoretical analysis with empirical evidence to demonstrate these problems.
strength_weakness: strength_weakness: 1. The experiments are primarily focused on specific artificial tasks (counting and copying), but the study lacks the analysis of the real-world texts. It would be beneficial to include some statistics of difference of tokens' representations on the standard text corpora. 
2. Authors provide only one artificial solution to split the sequence by different token.
3. The theoretical analysis makes some simplifying assumptions, such as treating attention weights as independent of the input. While the authors justify these assumptions, it would be beneficial to explore their impact on the results more thoroughly.
questions: questions: 1. Could you please include some statistics of difference of tokens' representations on the standard text corpora? It would be interesting to see does representational collapse happen in natural texts. 
2. How can you deal with representational collapse in real-world scenarios? Wouldn't it be worse to include random token for the information? 
3. Could this theoretical result be generalized to non-causal language modeling (different attention mask)?

Also here are some small remarks:

Line 331 typo: summariesd

Line 339 space in "phenomena : representational"
limitations: limitations: 1. Theoretical analysis relies on several simplifying assumptions, which are justified but may not fully capture the complexities of real-world models.
2. Empirical validation is conducted only on two specific models

Comments: title: Thank you!
Thank you!
comment: Thank you for acknowledging our responses and appreciating our work! We are delighted to hear our comments answered your concerns. Should you have any other questions, please feel free to let us know.
Thank you for your quick and very detailed reply. My concerns were answered, I understand now the real-world scenarios, and I admire the filler-token solution supported by references. I believe the score is already high enough, but you have fully answered all the questions. I am looking forward to see the detailed analysis of several types of attention in appendix.

Best wishes to your paper, and let me know if I can be of any help.
We would like to thank you for acknowledging our responses and increasing your score. We remain available should you have any additional questions during the discussion period!
Thank you for your satisfactory comments to my questions. I will raise my score to 6.
We would like to thank you once again for your time and consideration. We are pleased that you wish to maintain your acceptance score for our paper. We are very much excited about the potential of our work to contribute to stronger and scalable solutions in the future.
I thank the authors for their reply. 

After reading their comments, I will keep my score.  I do believe that this paper helps us understand a problem in LLM that is directly seen in a couple of problems such as copying and counting that has an impact in many applications. However, I did not increase my score further because I do not believe that a strong solution was proposed to this problem.