# Transformers need glasses!
Information over-squashing in language tasks

Federico Barbero
University of Oxford
<EMAIL_ADDRESS>&Andrea Banino
Google DeepMind
<EMAIL_ADDRESS>&Steven Kapturowski
Google DeepMind
<EMAIL_ADDRESS>&Dharshan Kumaran
Google DeepMind
<EMAIL_ADDRESS>&João G.M. Araújo
Google DeepMind
<EMAIL_ADDRESS>&Alex Vitvitskyi
Google DeepMind
<EMAIL_ADDRESS>&Razvan Pascanu
Google DeepMind
<EMAIL_ADDRESS>&Petar Veličković
Google DeepMind
<EMAIL_ADDRESS>Work performed while at Google DeepMind.

###### Abstract

We study how information propagates in decoder-only Transformers, which are
the architectural backbone 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 sequences 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 also
points to simple solutions towards ameliorating these issues.

##  1 Introduction

In recent years the field of Natural Language Processing (NLP) has been
revolutionised through the introduction of Transformer-based architectures
[26]. 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 [9, 17],
understanding multi-modal inputs [1], and code completion [14]. 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.

(a) Representational Collapse (b) Over-squashing

Figure 1:  (a) Representational Collapse (Theorem 4.2). From top to bottom, we
have a series of sequences given to Transformer architectures, each comprising
repeated 1 tokens with a single 0 token at the end. The color and proximity of
the curved lines illustrate how these representations converge as sequence
length increases. (b) Over-squashing (Theorem 5.1). Due to the architecture of
decoder-only Transformers, tokens that are earlier in their input sequence
will have significantly more paths through which their data can reach the
representation used for next-token prediction, leading to ‘over-squashing’.
This effect is depicted here for an early token (blue) and later token (red)
in a five-token sequence.

However, despite the impressive performance of Transformers, recent works have
uncovered surprising failures that may point to fundamental issues in their
architecture. For instance, Transformer-based LLMs seem to be particularly
challenged by seemingly simple tasks requiring counting [27] or copying
elements along input sequences [15]. We find it important to study such
failure cases, as these operations are fundamental building blocks of
computation, and they are often necessary for solving reasoning tasks. A
common strategy to assist LLMs in solving such tasks is to supply them with
‘tools’ [22]. We argue that, while tool use will certainly help, it is still
important to improve base model capabilities in this regard, because
oftentimes, even producing accurate inputs to a tool may require complex
reasoning operations. Specifically, often we need to _copy_ some part of the
Transformer’s input into a tool—if the base model struggles with robust
copying, even this operation can be in peril.

Accordingly, we find it important to explain why decoder-only Transformers do
not perform well when it comes to such problems—not just as an intellectual
endeavour, but also to help guide further practical improvements. While many
works have studied the computational capabilities of Transformers [20, 16],
they often make assumptions which do not correspond to present practical
limitations, such as infinite floating-point precision or ‘hard attention’,
making their conclusions less directly practically applicable.

In this work, we take a different approach and study what information _can_ be
contained in the representation of the last token at the last layer, as this
is ultimately the information that will be used for next-token prediction —
the fundamental mechanism through which modern Transformer LLMs perform
training and inference. In particular, we show that for certain distinct
sequences, their last-token representations can become arbitrarily close to
each other. This leads to a _representational collapse_ , exacerbated by the
lower-precision floating point types typically used by modern LLM stacks. As a
result Transformers incorrectly produce the same tokens on these sequence
pairs — see Figure 1 (a).

Furthermore, we reveal that the computation graph employed by decoder-only
Transformers, with its unidirectional causal mask, contributes to the observed
representational collapse. This unidirectional flow of information, converging
at the final token, is in fact likely to lead to a loss of information due to
_over-squashing_ , an effect that is well studied in graph neural networks
(GNNs) [2, 25, 7, 3, 10], and related to vanishing gradients [13, 4, 12]. We
hope that this result will be of independent interest to the GNN community, as
a practical application of over-squashing results at scale. Finally, we
provide supporting empirical evidence that these issues are likely of
practical interest, and propose simple solutions—directly stemming from our
theoretical study—to help alleviate them.

In summary, our paper provides the following contributions:

  * •

Theoretical analysis of decoder-only Transformer limitations: we formalise the
concepts of ‘representational collapse’ (Section 4) and ‘over-squashing’
(Section 5) in the context of Transformer-based architectures.

  * •

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 real-world experiments conducted on contemporary LLMs,
demonstrating practical implications of the limitations we identified.

##  2 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
$\mathbf{Q},\mathbf{K},\mathbf{V}\in\mathbb{R}^{n\times d}$ be the query, key,
and value matrices respectively on $n$ tokens and $d$ dimensions. We denote
with $\mathbf{q}_{i},\mathbf{k}_{i},\mathbf{v}_{i}\in\mathbb{R}^{d}$ the
$d$-dimensional query, key, and value vectors of the $i$-th token. We let
$\mathbf{p}_{ij}\in\mathbb{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 [23, 26]. The Transformer
model we consider computes the values, for a single head, of the $i$-th token
at the $\ell$-th Transformer layer $\mathbf{v}^{(\ell)}_{i}$ as111Note that we
rely on an abuse of notation. We ignore the linear projections used to compute
the value $\mathbf{v}_{i}^{(l)}$ from the output of layer below $l-1$. This
will not change our derivations, but would otherwise make notations more
cumbersome.

$\displaystyle\mathbf{z}^{(\ell)}_{i}$ $\displaystyle=\sum_{j\leq
i}\alpha_{ij}^{(\ell)}\operatorname{norm}_{1}^{(\ell)}\left(\mathbf{v}^{(\ell)}_{i}\right)+\mathbf{v}^{(\ell)}_{i},\text{with
}\alpha_{ij}^{(\ell)}=\frac{\exp\left(k\left(\mathbf{q}_{i}^{(\ell)},\mathbf{k}_{j}^{(\ell)},\mathbf{p}_{ij}\right)\right)}{\sum_{w\leq
i}\exp\left(k\left(\mathbf{q}_{i}^{(\ell)},\mathbf{k}_{w}^{(\ell)},\mathbf{p}_{iw}\right)\right)}$
$\displaystyle\mathbf{v}_{i}^{(\ell+1)}$
$\displaystyle=\boldsymbol{\psi}^{(\ell)}\left(\operatorname{norm}_{2}^{(\ell)}\left(\mathbf{z}_{i}^{(\ell)}\right)\right)+\mathbf{z}_{i}^{(\ell)}\vspace{-15pt}$

for a function
$k:\mathbb{R}^{d}\times\mathbb{R}^{d}\times\mathbb{R}^{2e}\to\mathbb{R}$
mapping queries, key, and positional encoding information to a scalar value,
an MLP $\boldsymbol{\psi}:\mathbb{R}^{d}\to\mathbb{R}^{d}$, and normalization
functions at the $\ell$-th layer $\operatorname{norm}_{1}^{(\ell)}$ and
$\operatorname{norm}_{2}^{(\ell)}$. This specific interleaving of components
is often referred to as a Pre-LN Transformer [29]. We can view the output of
the $\ell$-th layer of a Transformer as a sequence of $d$-dimensional vectors
$\mathbf{v}^{(\ell)}=(\mathbf{v}_{1}^{(\ell)},\dots,\mathbf{v}_{n}^{(\ell)})$.
Importantly, due to the causal attention mechanism, the vector
$\mathbf{v}_{j}^{(\ell)}$, will only depend on elements
$\mathbf{v}_{i}^{(\ell-1)}$ for $i\leq j$. We can group the attention weights
into an _attention matrix_ at the $\ell$-th layer which we define element-wise
as ${\boldsymbol{\Lambda}}^{(\ell)}_{ij}=\alpha_{ij}^{(\ell)}$. 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:

$\mathbf{y}_{i}=\operatorname{norm}_{3}\left(\mathbf{v}^{(L)}_{i}\right)$

We note that the next-token prediction usually depends purely on
$\mathbf{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
Turing-complete under certain assumptions [5]. 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 [20]. 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. [28], 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 that model. Following that,
Delétang et al. [6] 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
[16]. Finally, Peng et al. [19] 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.

Figure 2: Results on simple copying tasks. (a). Gemini was prompted to predict
the last token (diamond) of a sequences ‘1…10’ or the first token (square) of
a sequence ‘01…1’. (b). Same as (a) but with hints (see 3.2 for details) (c).
Same as (a) but the sequences have interleaved 0s and 1s. See C.1 for extra
details

#### 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 [30].
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 [11]. 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 _over-squashing_ [2, 25, 7, 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_ [10] — 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 $\mathbf{x}_{v}^{(L)}$ of a node $v$ after $L$ GNN layers is to
the initial representation $\mathbf{x}_{u}^{(0)}$ of another node $u$. In
particular, the partial derivative
$\partial\mathbf{x}_{v}^{(L)}/\partial\mathbf{x}_{u}^{(0)}$ 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 [13, 4, 18, 12].

##  3 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
[9] as our frontier LLM (referred as Gemini) and later analyse the internal
representations of the open-sourced Gemma model [24]. The goal is to showcase
intriguing failure cases which will motivate our signal propagation analysis.

###  3.1 Copying

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\dots 10$’ or the first element of a sequence ‘$01\dots 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.

Figure 3: Gemini 1.5 being prompted to (i) sum $1+\dots+1$ (left),(ii) Count
the number of ones in a sequence of 1s (center), and (iii) Count the number of
ones in a sequence of ones and zeroes (the sequence is a Bernoulli sequence
with probability of sampling a one being 0.7) (right).

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 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.

###  3.2 Counting

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+\dots+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.

Figure 4: Frequency of different outputted values for Gemini 1.5 for the
counting 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 subitizing. 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.

##  4 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 $\mathbf{x},\mathbf{x}^{*}\in\mathbb{R}^{n}$ such that
$\lim_{n\to\infty}\left\lvert\mathbf{x}_{n}-\mathbf{x}^{*}_{n}\right\rvert=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).

######  Theorem 4.2 (Representational Collapse – informal).

Let $\mathbf{v}^{(0)}\in\mathbb{R}^{n\times d}$ be a sequence and
$\mathbf{v}^{*(0)}\in\mathbb{R}^{(n+1)\times d}$ be another sequence equal to
$\mathbf{v}^{(0)}$ with the last token of $\mathbf{v}^{(0)}$ repeated. Assume
that the positional encoding information decays to $0$ with the distance, e.g.
as in RoPE [23]. 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.

#### Measuring representational collapse.

We report experiments showcasing representational collapse by measuring the
internal representations of Gemma 7B [24]. For two sequences
$\mathbf{v}^{(0)}$ and $\mathbf{v}^{*(0)}$ we report their difference in
representation at the last layer
$\left\lVert\mathbf{v}^{(L)}-\mathbf{v}^{*(L)}\right\rVert_{\infty}$ 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) “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.

Figure 5: Representational collapse for counting (a, b) and copying (c, d)
tasks.

#### Quantisation and Tokenisation.

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 representations collapse for the sequence
of ones (c), adding commas every third digit (d) helps to keep the
representations well-separated.

##  5 Over-squashing in Language Tasks

In this Section, we study a more general phenomenon related to
representational collapse—over-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 $\partial\mathbf{y}_{n}/\partial\mathbf{v}_{i}^{(0)}$ 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_ [25, 7, 10], a problem related to the well-
known vanishing gradients problem in RNNs [13, 4, 10]. 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 $\beta_{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).

######  Theorem 5.1 (Over-squashing in Transformers).

Consider an input sequence $\mathbf{v}_{1}^{(0)},\dots,\mathbf{v}_{n}^{(0)}$.
Let $C>0$ be some constant and
$\bar{\alpha}_{i,j}^{(\ell)}=\frac{\alpha_{i,j}^{(\ell)}}{\beta_{2}}+\delta_{i,j}$,
then:

$\left\lVert\frac{\partial\mathbf{y}_{n}}{\partial\mathbf{v}_{i}^{(0)}}\right\rVert\leq
C\sum_{k_{1}\geq i}\dots\sum_{k_{L}\geq
k_{L-1}}\bar{\alpha}^{(L-1)}_{n,k_{L}}\prod_{\ell=2}^{L-1}\bar{\alpha}^{(\ell-1)}_{k_{\ell},k_{\ell-1}}\bar{\alpha}^{(0)}_{k_{1},i}$
(1)

Theorem 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\to 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 [12]. 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 $\mathbf{y}_{n}$ depends only on the starting input
token $\mathbf{v}_{1}^{(0)}$. 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\to\infty$ 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 _U-shape
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.

##  6 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 [20].
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 [8]. 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.

##  7 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
$\mathbf{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 information-propagation issues in
Transformer-based architectures, but we hope that the findings may help better
understand and improve language models available today.

##  8 Acknowledgements

We would like to thank Wojciech Marian Czarnecki (Google DeepMind), Simon
Osindero (Google DeepMind), and Timothy Nguyen (Google DeepMind) for the
valuable comments and suggestions regarding this work. We also thank
Constantin Kogler (University of Oxford) for helping to strengthen the theory
on representational collapse.

## References

  * [1] Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, Iain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, et al.  Flamingo: a visual language model for few-shot learning.  Advances in neural information processing systems, 35:23716–23736, 2022.
  * [2] Uri Alon and Eran Yahav.  On the bottleneck of graph neural networks and its practical implications.  In International Conference on Learning Representations, 2021.
  * [3] Federico Barbero, Ameya Velingker, Amin Saberi, Michael M. Bronstein, and Francesco Di Giovanni.  Locality-aware graph rewiring in GNNs.  In The Twelfth International Conference on Learning Representations, 2024.
  * [4] Y. Bengio, P. Simard, and P. Frasconi.  Learning long-term dependencies with gradient descent is difficult.  IEEE Transactions on Neural Networks, 5(2):157–166, 1994.
  * [5] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Łukasz Kaiser.  Universal transformers, 2019.
  * [6] Grégoire Delétang, Anian Ruoss, Jordi Grau-Moya, Tim Genewein, Li Kevin Wenliang, Elliot Catt, Chris Cundy, Marcus Hutter, Shane Legg, Joel Veness, and Pedro A. Ortega.  Neural networks and the chomsky hierarchy, 2023.
  * [7] Francesco Di Giovanni, Lorenzo Giusti, Federico Barbero, Giulia Luise, Pietro Lio, and Michael M Bronstein.  On over-squashing in message passing neural networks: The impact of width, depth, and topology.  In International Conference on Machine Learning, pages 7865–7885. PMLR, 2023.
  * [8] Andrew Joseph Dudzik, Tamara von Glehn, Razvan Pascanu, and Petar Veličković.  Asynchronous algorithmic alignment with cocycles.  In Learning on Graphs Conference, pages 3–1. PMLR, 2024.
  * [9] Team Gemini.  Gemini: a family of highly capable multimodal models.  arXiv preprint arXiv:2312.11805, 2023.
  * [10] Francesco Di Giovanni, T. Konstantin Rusch, Michael Bronstein, Andreea Deac, Marc Lackenby, Siddhartha Mishra, and Petar Veličković.  How does over-squashing affect the power of GNNs?  Transactions on Machine Learning Research, 2024.
  * [11] Michael Hahn.  Theoretical limitations of self-attention in neural sequence models.  Transactions of the Association for Computational Linguistics, 8:156–171, 2020.
  * [12] Luca Herranz-Celotti and Jean Rouat.  Stabilizing rnn gradients through pre-training, 2024.
  * [13] Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  Neural Computation, 9(8):1735–1780, 1997.
  * [14] Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, et al.  Competition-level code generation with alphacode.  Science, 378(6624):1092–1097, 2022.
  * [15] Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang.  Lost in the middle: How language models use long contexts.  Transactions of the Association for Computational Linguistics, 12:157–173, 2024.
  * [16] William Merrill and Ashish Sabharwal.  The expresssive power of transformers with chain of thought.  arXiv preprint arXiv:2310.07923, 2023.
  * [17] R OpenAI.  Gpt-4 technical report. arxiv 2303.08774.  View in Article, 2(5), 2023.
  * [18] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio.  On the difficulty of training recurrent neural networks.  In Sanjoy Dasgupta and David McAllester, editors, Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, pages 1310–1318, Atlanta, Georgia, USA, 17–19 Jun 2013. PMLR.
  * [19] Binghui Peng, Srini Narayanan, and Christos Papadimitriou.  On limitations of the transformer architecture, 2024.
  * [20] Jorge Pérez, Pablo Barceló, and Javier Marinkovic.  Attention is turing-complete.  Journal of Machine Learning Research, 22(75):1–35, 2021.
  * [21] S Unnikrishna Pillai, Torsten Suel, and Seunghun Cha.  The perron-frobenius theorem: some of its applications.  IEEE Signal Processing Magazine, 22(2):62–75, 2005.
  * [22] Timo Schick, Jane Dwivedi-Yu, Roberto Dessì, Roberta Raileanu, Maria Lomeli, Eric Hambro, Luke Zettlemoyer, Nicola Cancedda, and Thomas Scialom.  Toolformer: Language models can teach themselves to use tools.  Advances in Neural Information Processing Systems, 36, 2024.
  * [23] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu.  Roformer: Enhanced transformer with rotary position embedding.  Neurocomputing, 568:127063, 2024.
  * [24] Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al.  Gemma: Open models based on gemini research and technology.  arXiv preprint arXiv:2403.08295, 2024.
  * [25] Jake Topping, Francesco Di Giovanni, Benjamin Paul Chamberlain, Xiaowen Dong, and Michael M Bronstein.  Understanding over-squashing and bottlenecks on graphs via curvature.  2022\.
  * [26] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [27] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al.  Emergent abilities of large language models.  arXiv preprint arXiv:2206.07682, 2022.
  * [28] Gail Weiss, Yoav Goldberg, and Eran Yahav.  Thinking like transformers, 2021.
  * [29] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu.  On layer normalization in the transformer architecture.  In International Conference on Machine Learning, pages 10524–10533. PMLR, 2020.
  * [30] Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Joshua M Susskind.  Stabilizing transformer training by preventing attention entropy collapse.  In International Conference on Machine Learning, pages 40770–40803. PMLR, 2023.

##  Appendix A Broader impact

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.

##  Appendix B Proofs

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).

###  B.1 Representational Collapse

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 $\mathbf{a}\in\mathbb{R}^{n-1}$ and two scalars
$b,c\in\mathbb{R}$. Let $\mathbf{x}=[\mathbf{a}\ c]^{T}\in\mathbb{R}^{n}$ and
$\mathbf{x}^{*}=[\mathbf{a}\ b\ c]^{T}\in\mathbb{R}^{n+1}$ with all entries
bounded. Then,
$\operatorname{softmax}\left(\mathbf{x}\right)_{n}>\operatorname{softmax}\left(\mathbf{x}^{*}\right)_{n+1}$.
Moreover for any $p>0$ we can find large enough $n\in\mathbb{N}^{+}$ such that
$|\operatorname{softmax}\left(\mathbf{x}\right)_{n}-\operatorname{softmax}\left(\mathbf{x}^{*}\right)_{n+1}|$
< p.

###### Proof.

We directly compute:

$\displaystyle\operatorname{softmax}(\mathbf{x})_{n}$
$\displaystyle=\frac{\exp(c)}{\sum_{k=1}^{n-1}\exp(\mathbf{a}_{k})+\exp(c)}$
$\displaystyle\operatorname{softmax}(\mathbf{x}^{*})_{n+1}$
$\displaystyle=\frac{\exp(c)}{\sum_{k=1}^{n-1}\exp(\mathbf{a}_{k})+\exp(b)+\exp(c)}$

As we assume that the entries are bounded, we have that
$\sum_{j=1}^{n-1}\exp(\mathbf{a}_{j})+\exp(c)<\sum_{j=1}^{n-1}\exp(\mathbf{a}_{j})+\exp(b)+\exp(c)$,
therefore
$\operatorname{softmax}(\mathbf{x})_{n}>\operatorname{softmax}(\mathbf{x}^{*})_{n+1}$.

For the second part of the statement, we compute:

$\displaystyle\left\lvert\operatorname{softmax}(\mathbf{x})_{n}-\operatorname{softmax}(\mathbf{x}^{*})_{n+1}\right\rvert$
$\displaystyle=\left\lvert\frac{\exp(c)}{\sum_{k=1}^{n-1}\exp(\mathbf{a}_{k})+\exp(c)}-\frac{\exp(c)}{\sum_{k=1}^{n-1}\exp(\mathbf{a}_{k})+\exp(b)+\exp(c)}\right\rvert$
$\displaystyle\leq\left\lvert\frac{\exp(c)}{\sum_{k=1}^{n-1}\exp(\mathbf{a}_{k})+\exp(c)}\right\rvert+\left\lvert\frac{\exp(c)}{\sum_{k=1}^{n-1}\exp(\mathbf{a}_{k})+\exp(b)+\exp(c)}\right\rvert$
$\displaystyle<p$

for some $p>0$, as in the last step the summands tend to $0$ as $n\to\infty$.
We therefore have that
$\left\lvert\operatorname{softmax}(\mathbf{x})_{n}-\operatorname{softmax}(\mathbf{x}^{*})_{n+1}\right\rvert\to
0$ as $n\to\infty$ and for large enough $n$
$\left\lvert\operatorname{softmax}(\mathbf{x})_{n}-\operatorname{softmax}(\mathbf{x}^{*})_{n+1}\right\rvert<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 $\mu,\nu$
supported on the same space, we define their total variation $\delta(\mu,\nu)$
— or equivalently $L_{1}$ norm $\left\lVert\mu-\nu\right\rVert_{1}$, as:

$\delta(\mu,\nu)=\sum_{x}\left\lvert\mu(x)-\nu(x)\right\rvert.$ (2)

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, e.g.
as in RoPE [23], then the total variation between the two sequences goes to
$0$ with $n$. Such a result is presented in Lemma B.2

######  Lemma B.2.

Consider two sequences $\mathbf{x},\mathbf{x}^{*}\in\mathbb{R}^{n}$ such that
$\lim_{n\to\infty}\left\lvert\mathbf{x}_{n}-\mathbf{x}^{*}_{n}\right\rvert=0$,
with $\mathbf{x}_{i},\mathbf{x}^{*}_{i}$ bounded. Let
$\mathbf{y},\mathbf{y}^{*}\in\mathbb{R}^{n}$ be the softmax of $\mathbf{x}$
and $\mathbf{x}^{*}$ respectively. Then, as $n\to\infty$ the total variation
tends to $0$, i.e. $\lim_{n\to\infty}\delta(\mathbf{y},\mathbf{y}^{*})=0$.

###### Proof.

Let $Z=\sum_{i=1}^{n}e^{\mathbf{x}}_{i}$ and
$Z^{*}=\sum_{i=1}^{n}e^{\mathbf{x}^{*}_{i}}$ be the partition functions for
$\mathbf{x}$ and $\mathbf{x}^{*}$, respectively. We start by bounding the
quantity $\left\lvert Z-Z^{*}\right\rvert$. In particular, let $\epsilon>0$,
we first claim:

$\left\lvert Z-Z^{*}\right\rvert\leq\sum_{i=1}^{n}\left\lvert
e^{\mathbf{x}_{i}}-e^{\mathbf{x}^{*}_{i}}\right\rvert\leq\epsilon\min(Z,Z^{*}).$
(3)

Consider some $n_{0}\geq 1$ such that $\left\lvert
1-e^{\mathbf{x}^{*}_{i}-\mathbf{x}_{i}}\right\rvert\leq\epsilon/2$. We note
that this always possible as
$\left\lvert\mathbf{x}^{*}_{i}-\mathbf{x}_{i}\right\rvert\to 0$ by assumption.
We compute:

$\displaystyle\sum_{i=1}^{n}\left\lvert
e^{\mathbf{x}_{i}}-e^{\mathbf{x}^{*}_{i}}\right\rvert$
$\displaystyle=\sum_{i=1}^{n_{0}-1}\left\lvert
e^{\mathbf{x}_{i}}-e^{\mathbf{x}^{*}_{i}}\right\rvert+\sum_{i=n_{0}}^{n}\left\lvert
e^{\mathbf{x}_{i}}-e^{\mathbf{x}^{*}_{i}}\right\rvert$
$\displaystyle=\sum_{i=1}^{n_{0}-1}\left\lvert
e^{\mathbf{x}_{i}}-e^{\mathbf{x}^{*}_{i}}\right\rvert+\sum_{i=n_{0}}^{n}e^{\mathbf{x}_{i}}\left\lvert
1-e^{\mathbf{x}^{*}_{i}-\mathbf{x}_{i}}\right\rvert$
$\displaystyle\leq\sum_{i=1}^{n_{0}-1}\left\lvert
e^{\mathbf{x}_{i}}-e^{\mathbf{x}^{*}_{i}}\right\rvert+\frac{\epsilon}{2}\sum_{i=n_{0}}^{n}e^{\mathbf{x}_{i}}$
$\displaystyle\leq\sum_{i=1}^{n_{0}-1}\left\lvert
e^{\mathbf{x}_{i}}-e^{\mathbf{x}^{*}_{i}}\right\rvert+\frac{\epsilon}{2}Z$
$\displaystyle\leq\epsilon Z$

Where 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\epsilon/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:

$\displaystyle\delta(\mathbf{y},\mathbf{y}^{*})$
$\displaystyle=\sum_{i=1}^{n}\left\lvert\frac{e^{\mathbf{x}_{i}}}{Z}-\frac{e^{\mathbf{x}^{*}_{i}}}{Z^{*}}\right\rvert$
$\displaystyle=\sum_{i=1}^{n}\left\lvert\frac{Z^{*}e^{\mathbf{x}_{i}}-Ze^{\mathbf{x}^{*}_{i}}}{ZZ^{*}}\right\rvert$
$\displaystyle=\sum_{i=1}^{n}\left\lvert\frac{Z^{*}e^{\mathbf{x}_{i}}-Ze^{\mathbf{x}_{i}}+Ze^{\mathbf{x}_{i}}-Ze^{\mathbf{x}^{*}_{i}}}{ZZ^{*}}\right\rvert$
$\displaystyle\leq\sum_{i=1}^{n}\left\lvert\frac{Z^{*}e^{\mathbf{x}_{i}}-Ze^{\mathbf{x}_{i}}}{ZZ^{*}}\right\rvert+\left\lvert\frac{Ze^{\mathbf{x}_{i}}-Ze^{\mathbf{x}^{*}_{i}}}{ZZ^{*}}\right\rvert$
$\displaystyle\leq\frac{\epsilon\min(Z,Z^{*})}{ZZ^{*}}+\frac{\epsilon\min(Z,Z^{*})}{ZZ^{*}}$
$\displaystyle\leq 2\epsilon$

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, e.g. as in RoPE. The final token repetition is important as
otherwise the residual connection in the Transformer would not make the
representations necessarily converge.

######  Theorem B.3 (Representational Collapse).

Let $\mathbf{x}\in\mathbb{R}^{n-1\times d}$ be an underlying growing token
sequence. Let $\mathbf{v}^{(0)}=[\mathbf{v}\
\mathbf{v}_{a}]^{T}\in\mathbb{R}^{n\times d}$ and
$\mathbf{v}^{*(0)}=[\mathbf{v}\ \mathbf{v}_{a}\
\mathbf{v}_{a}]^{T}\in\mathbb{R}^{n+1\times d}$ be two sequences for a final
repeated token $\mathbf{x}_{a}\in\mathbb{R}^{d}$, with all token
representations bounded. Further, assume that the positional encodings decay
with distance to $0$, e.g. as in RoPE. Then, for large enough
$n\in\mathbb{N}^{+}$, we have that the representations are under any
$\epsilon$:

$||\mathbf{v}_{n}^{(L)}-\mathbf{v}_{n+1}^{*(L)}||_{1}<\epsilon.$

###### 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 $\mathbf{z}^{(0)}$ with
the $n+1$-th element of $\mathbf{z}^{*(0)}$:

$\displaystyle\left\lVert\mathbf{z}_{n}^{(0)}-\mathbf{z}_{n+1}^{*(0)}\right\rVert_{1}$
$\displaystyle=\left\lVert\sum_{i<n}\alpha^{(0)}_{n,i}\mathbf{v}_{i}^{(0)}+\alpha^{(0)}_{n,n}\mathbf{v}_{a}^{(0)}-\left(\sum_{i<n}\alpha^{*(0)}_{n+1,i}\mathbf{v}^{(0)}_{i}+\left(\alpha^{*(0)}_{n+1,n}+\alpha^{*(0)}_{n+1,n+1}\right)\mathbf{v}^{(0)}_{a}\right)\right\rVert_{1}$
$\displaystyle=\left\lVert\sum_{i<n}\left(\alpha^{(0)}_{n,i}-\alpha^{*(0)}_{n+1,i}\right)\mathbf{v}_{i}^{(0)}+\left(\alpha^{(0)}_{n,n}-\alpha^{*(0)}_{n+1,n}-\alpha^{*(0)}_{n+1,n+1}\right)\mathbf{v}^{(0)}_{a}\right\rVert_{1}$
$\displaystyle\leq\sum_{i<n}\left\lvert\alpha^{(0)}_{n,i}-\alpha^{*(0)}_{n+1,i}\right\rvert+\left\lvert\alpha^{(0)}_{n,n}-\alpha^{*(0)}_{n+1,n}-\alpha^{*(0)}_{n+1,n+1}\right\rvert$
$\displaystyle\leq\sum_{i<n}\left\lvert\alpha^{(0)}_{n,i}-\alpha^{*(0)}_{n+1,i}\right\rvert+\left\lvert\alpha^{(0)}_{n,n}-\alpha^{*(0)}_{n+1,n}\right\rvert+\left\lvert\alpha^{*(0)}_{n+1,n+1}\right\rvert$
$\displaystyle=\delta\left(\alpha^{(0)}_{n,:},\alpha^{*(0)}_{n,:n}\right)+\alpha^{*(0)}_{n+1,n+1}<\epsilon$

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 $\delta\left(\alpha^{(0)}_{n,:},\alpha^{*(0)}_{n,:n}\right)$ goes to $0$
with $n\to\infty$ thanks to Lemma B.2 and our assumptions on the positional
encodings. Similarly, the term $\alpha^{*(0)}_{n+1,n+1}$ goes to $0$ due to
Lemma B.1. We have also used the fact that
$\mathbf{v}_{n}^{(0)}=\mathbf{v}_{n+1}^{*(0)}$ by construction to ignore the
residual connection. As
$\mathbf{v}_{i}^{(\ell+1)}=\boldsymbol{\psi}^{(\ell)}\left(\operatorname{norm}_{2}^{(\ell)}\left(\mathbf{z}_{i}^{(\ell)}\right)\right)+\mathbf{v}_{i}^{(\ell)}$,
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 $\mathbf{x}=(1\ 0\ 1\ 0\dots)$ and $\mathbf{x}^{*}=(0\ 1\ 0\
1\dots)$, 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.

######  Lemma B.4.

Consider two sequences $\mathbf{x}=(1\ 0\ 1\ 0\dots)\in\mathbb{R}^{n}$ and
$\mathbf{x}^{*}=(0\ 1\ 0\ 1\dots)\in\mathbb{R}^{n}$ for $n$ even. Let
$\mathbf{y}$ and $\mathbf{y}^{*}$ be the softmax of $\mathbf{x}$ and
$\mathbf{x}^{*}$, respectively. Then the total variation
$\delta(\mathbf{y},\mathbf{y}^{*})$ does not tend to $0$ as $n\to\infty$.

###### Proof.

Let $Z=\sum_{i=1}^{n}e^{\mathbf{x}}_{i}$ and
$Z^{*}=\sum_{i=1}^{n}e^{\mathbf{x}^{*}_{i}}$ be the partition functions for
$\mathbf{x}$ and $\mathbf{x}^{*}$, respectively. We directly compute:

$\displaystyle\lim_{n\to\infty}\delta(\mathbf{y},\mathbf{y}^{*})$
$\displaystyle=\lim_{n\to\infty}\sum_{i=1}^{n}\left\lvert\frac{\mathbf{y}_{i}}{Z}-\frac{\mathbf{y}^{*}_{i}}{Z^{*}}\right\rvert$
$\displaystyle=\lim_{n\to\infty}\sum_{i=1}^{n}\left\lvert\frac{\mathbf{y}_{i}}{\frac{n}{2}e+\frac{n}{2}}-\frac{\mathbf{y}^{*}_{i}}{\frac{n}{2}e+\frac{n}{2}}\right\rvert$
$\displaystyle=\lim_{n\to\infty}\sum_{i=1}^{n}\left\lvert\frac{e-1}{\frac{n}{2}e+\frac{n}{2}}\right\rvert$
$\displaystyle=\lim_{n\to\infty}\frac{n\left(e-1\right)}{n\frac{e+1}{2}}$
$\displaystyle=2\frac{e-1}{e+1}>0.$

∎

###  B.2 Over-squashing

We 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.

######  Theorem B.5 (Over-squashing in Transformers).

Consider an input sequence $\mathbf{v}_{1}^{(0)},\dots,\mathbf{v}_{n}^{(0)}$
(including CoT). Let $\sigma_{\boldsymbol{\psi}}$ be the maximal Lipschitz
constant of any $\boldsymbol{\psi}^{(\ell)}$, and
$\bar{\alpha}^{(\ell)}_{j,i}=\frac{1}{\beta^{(\ell)}}\left(\alpha_{j,i}^{(\ell)}+\delta_{j,i}\right)$
the normalized attention coefficient, then:

$\left\lVert\frac{\partial\mathbf{y}_{n}}{\partial\mathbf{v}_{i}^{(0)}}\right\rVert\leq\sigma^{L}_{\boldsymbol{\psi}}\sum_{k_{1}\geq
i}\dots\sum_{k_{L}\geq
k_{L-1}}\bar{\alpha}^{(L-1)}_{n,k_{L}}\prod_{\ell=2}^{L-1}\bar{\alpha}^{(\ell-1)}_{k_{\ell},k_{\ell-1}}\bar{\alpha}^{(0)}_{k_{1},i}$
(4)

###### Proof.

Note that for $j\geq i$ we have:

$\displaystyle\left\lVert\frac{\partial\mathbf{v}_{j}^{(\ell+1)}}{\partial\mathbf{v}_{i}^{(\ell)}}\right\rVert$
$\displaystyle=\left\lVert\frac{\partial}{\partial\mathbf{v}_{j}^{(\ell)}}\left[\boldsymbol{\psi}^{(\ell)}\left(\operatorname{norm}_{2}^{(\ell)}\left(\mathbf{z}_{j}^{(\ell)}\right)\right)+\mathbf{z}_{j}^{(\ell)}\right]\right\rVert$
$\displaystyle\leq\left(\frac{\sigma_{\psi^{(\ell)}}}{\beta_{2}^{(\ell)}}+1\right)\frac{\partial\mathbf{z}_{j}^{(\ell)}}{\partial\mathbf{v}_{i}^{(\ell)}}$
$\displaystyle=\left(\frac{\sigma_{\psi^{(\ell)}}}{\beta_{2}^{(\ell)}}+1\right)\frac{\partial}{\partial\mathbf{v}_{i}^{(\ell)}}\left[\sum_{j\leq
i}\alpha_{ij}^{(\ell)}\operatorname{norm}_{1}^{(\ell)}\left(\mathbf{v}^{(\ell)}_{i}\right)+\mathbf{v}^{(\ell)}_{i}\right]$
$\displaystyle=\left(\frac{\sigma_{\psi^{(\ell)}}}{\beta_{2}^{(\ell)}}+1\right)\left(\frac{\alpha_{j,i}^{(\ell)}}{\beta_{1}^{(\ell)}}+\delta_{j,i}\right)$

where we let $\beta_{i}^{(\ell)}$ represent the effect of layer normalization
$i$ at the $\ell$-th layer and $\sigma_{\boldsymbol{\psi}^{(\ell)}}$ the
Lipschitz constant of $\boldsymbol{\psi}^{(\ell)}$. For the case when $j<i$
due to the causal mechanism we have that
$\partial\mathbf{v}_{j}^{(\ell)}/\partial\mathbf{v}_{i}^{(\ell-1)}=0$. We
compute the following bound:

$\displaystyle\left\lVert\frac{\partial\mathbf{y}_{n}}{\partial\mathbf{v}_{i}^{(0)}}\right\rVert$
$\displaystyle=\left\lVert\frac{1}{\beta_{3}}\sum_{k_{1}}\dots\sum_{k_{L}}\frac{\partial\mathbf{v}_{n}^{(L)}}{\partial\mathbf{v}_{k_{L}}^{(L-1)}}\prod_{\ell=2}^{L-1}\frac{\partial\mathbf{v}_{k_{\ell}}^{(\ell)}}{\partial\mathbf{v}_{k_{\ell-1}}^{(\ell-1)}}\frac{\partial\mathbf{v}_{k_{1}}^{(1)}}{\partial\mathbf{v}_{i}^{(0)}}\right\rVert$
$\displaystyle=\left\lVert\frac{1}{\beta_{3}}\sum_{k_{1}\geq
i}\dots\sum_{k_{L}\geq
k_{L-1}}\frac{\partial\mathbf{v}_{n}^{(L)}}{\partial\mathbf{v}_{k_{L}}^{(L-1)}}\prod_{\ell=2}^{L-1}\frac{\partial\mathbf{v}_{k_{\ell}}^{(\ell)}}{\partial\mathbf{v}_{k_{\ell-1}}^{(\ell-1)}}\frac{\partial\mathbf{v}_{k_{1}}^{(1)}}{\partial\mathbf{v}_{i}^{(0)}}\right\rVert$
$\displaystyle\leq\frac{1}{\beta_{3}}\prod_{\ell=1}^{L}\left(\frac{\sigma_{\boldsymbol{\psi}}}{\beta_{2}^{(\ell)}}+1\right)\sum_{k_{1}\geq
i}\dots\sum_{k_{L}\geq
k_{L-1}}\bar{\alpha}^{(L-1)}_{n,k_{L}}\prod_{\ell=2}^{L-1}\bar{\alpha}^{(\ell-1)}_{k_{\ell},k_{\ell-1}}\bar{\alpha}^{(0)}_{k_{1},i}$
$\displaystyle=C\sum_{k_{1}\geq i}\dots\sum_{k_{L}\geq
k_{L-1}}\bar{\alpha}^{(L-1)}_{n,k_{L}}\prod_{\ell=2}^{L-1}\bar{\alpha}^{(\ell-1)}_{k_{\ell},k_{\ell-1}}\bar{\alpha}^{(0)}_{k_{1},i}$

where we let
$\bar{\alpha}_{j,i}^{(\ell)}=\frac{\alpha_{j,i}^{(\ell)}}{\beta_{1}^{(\ell)}}+\delta_{j,i}$
and
$C=\frac{1}{\beta_{3}}\prod_{\ell=1}^{L}\left(\frac{\sigma_{\boldsymbol{\psi}}}{\beta_{2}^{(\ell)}}+1\right)$.
∎

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 $\mathbf{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.

#### Connection to the spectral theory of Markov chains.

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 $\mathbf{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 $\mathbf{A}$ cannot have eigenvalues $\lambda>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 $\lambda>1$ for some eigenvector $\phi$, we then have that
$\mathbf{A}\phi=\lambda\phi$. Consider $\phi_{i}=\max_{k}\phi_{k}>0$. Now
$(\mathbf{A}\phi)_{i}=\sum_{j\leq i}\mathbf{A}_{ij}\phi_{j}=\lambda\phi_{i}$.
As the sum is a convex combination, the result cannot be larger than the
already maximal element $\phi_{i}$. As $\lambda>1$, we however have that
$\lambda\phi_{i}>\phi_{i}$ which is a contradiction and we conclude that
$\lambda\leq 1$.

It is easy to find an eigenvector that always has eigenvalue $1$. Consider a
vector $\mathbf{x}$ which is a constant vector of $1$s. Then
$(\mathbf{A}\mathbf{x})_{i}=\sum_{j\leq i}\mathbf{A}_{ij}=\mathbf{x}_{i}$,
therefore $\mathbf{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 $\mathbf{y}_{i}>0$,
then we have that $(\mathbf{A}\mathbf{y})_{i}=\sum_{j\leq
i}\mathbf{A}_{ij}\mathbf{y}_{j}=\mathbf{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 $\mathbf{y}_{i}$. The condition that each
row has at least two non-zero entries is important as it means that the
condition $\mathbf{y}_{i}=\mathbf{y}_{j}$ is true for all tokens. ∎

######  Lemma B.7.

The product of two row-stochastic matrices is again row-stochastic. Moreover,
the product of two triangular row-stochastic matrices is a triangular row-
stochastic matrix.

###### Proof.

Let $\mathbf{A},\mathbf{B}\in\mathbb{R}^{n\times n}$ be two row-stochastic
matrices. We compute:

$\sum_{j}\left(\mathbf{A}\mathbf{B}\right)_{ij}=\sum_{j}\sum_{k}\mathbf{A}_{ik}\mathbf{B}_{kj}=\sum_{k}\mathbf{A}_{ik}\sum_{j}\mathbf{B}_{kj}=1$

The 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 $\mathbf{y}_{n}$ only depends on
the initial input token $\mathbf{v}_{1}^{(0)}$ as the number layers tends to
infinity, i.e. $L\to\infty$.

######  Proposition B.8.

Let $\beta^{(\ell)}_{1},\beta^{(\ell)}_{2}=1$, $\beta_{3}^{1/L}=4$,
$\sigma_{\psi}$ = 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\to\infty$ that $\partial\mathbf{y}_{n}/\partial\mathbf{v}_{i}^{(0)}=0$ when
$i\neq 1$ and $\partial\mathbf{y}_{n}/\partial\mathbf{v}_{i}^{(0)}=1$ when
$i=1$. In other words, $\mathbf{y}_{n}$ will only be sensitive to the first
token.

###### Proof.

Let the associated attention matrix be $\boldsymbol{\Lambda}$. We start by re-
writing the following:

$\displaystyle\left\lVert\frac{\partial\mathbf{y}_{n}}{\partial\mathbf{v}_{i}^{(0)}}\right\rVert$
$\displaystyle\leq\frac{1}{\beta_{3}}\prod_{\ell=1}^{L}\left(\frac{\sigma_{\boldsymbol{\psi}}}{\beta_{2}^{(\ell)}}+1\right)\sum_{k_{1}\geq
i}\dots\sum_{k_{L}\geq
k_{L-1}}\bar{\alpha}^{(L-1)}_{n,k_{L}}\prod_{\ell=2}^{L-1}\bar{\alpha}^{(\ell-1)}_{k_{\ell},k_{\ell-1}}\bar{\alpha}^{(0)}_{k_{1},i}$
$\displaystyle=\left(\prod_{\ell=1}^{L}\left[\frac{1}{\beta_{3}^{1/L}}\left(\frac{\sigma_{\boldsymbol{\psi}}}{\beta_{2}}+1\right)\left(\frac{1}{\beta_{1}}\boldsymbol{\Lambda}+\mathbf{I}\right)\right]\right)_{n,i}$
$\displaystyle=\left(\left[\frac{1}{2}\left(\boldsymbol{\Lambda}+\mathbf{I}\right)\right]^{L}\right)_{n,i}$

We now point out that
$\tilde{\boldsymbol{\Lambda}}=\frac{1}{2}\left(\boldsymbol{\Lambda}+\mathbf{I}\right)$
is row-stochastic and with our assumptions is diagonalizable into
$\tilde{\boldsymbol{\Lambda}}=\Psi\Sigma\Phi$. In particular, by Lemma B.7,
also $\tilde{\boldsymbol{\Lambda}}^{L}$ is row-stochastic and each entry is
non-negative. We now use the Perron-Frobenius theorem for non-negative
matrices [21], which guarantees us that all eigenvalues $\lambda_{k}$ of
$\tilde{\boldsymbol{\Lambda}}$ are bounded such that $|\lambda_{k}|\leq 1$. In
particular, thanks to Lemma B.6, we know that there is a unique eigenvector
(the constant eigenvector $\psi_{n}$) with eigenvalue $\lambda_{n}=1$. Denote
the left eigenvectors by $\psi_{k}$ and the right eigenvectors $\phi_{n}$, we
therefore have:

$\displaystyle\lim_{L\to\infty}\tilde{\boldsymbol{\Lambda}}^{L}=\sum_{k}\lambda_{k}^{L}\psi_{k}\phi_{k}^{T}=\psi_{n}\phi_{n}^{T}.$

In particular, one can check that $\phi_{n}^{T}=\left[1\ 0\dots\ 0\right]$,
meaning that $\psi_{n}\phi_{n}^{T}$ has as first column a constant vector of
$1$s and every other entry $0$. This completes the proof. ∎

###  B.3 Counting

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, $\mathbf{v}_{zero}^{(0)}$ and
$\mathbf{v}_{one}^{(0)}$, with $n_{0}$ and $n_{1}$ being the number of zeros
and ones respectively. We ignore in our calculations the MLPs $\psi$ and the
normalizations $\operatorname{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:

$\displaystyle\mathbf{z}_{zero}^{(1)}$
$\displaystyle=\sum_{j}\frac{\exp\left(\mathbf{q}^{(0)T}_{zero}\mathbf{k}^{(0)}_{j}\right)}{\sum_{w}\exp\left(\mathbf{q}^{(0)T}_{zero}\mathbf{k}^{(0)}_{w}\right)}\mathbf{v}_{j}^{(0)}+\mathbf{v}_{zero}^{(0)}$
$\displaystyle=\frac{n_{0}\exp\left(\mathbf{q}^{(0)T}_{zero}\mathbf{k}^{(0)}_{zero}\right)}{n_{0}\exp\left(\mathbf{q}^{(0)T}_{zero}\mathbf{k}^{(0)}_{zero}\right)+n_{1}\exp\left(\mathbf{q}^{(0)T}_{zero}\mathbf{k}^{(0)}_{one}\right)}\mathbf{v}_{zero}^{(0)}+\frac{n_{1}\exp\left(\mathbf{q}^{(0)T}_{zero}\mathbf{k}^{(0)}_{one}\right)}{n_{0}\exp\left(\mathbf{q}^{(0)T}_{zero}\mathbf{k}^{(0)}_{zero}\right)+n_{1}\exp\left(\mathbf{q}^{(0)T}_{zero}\mathbf{k}^{(0)}_{one}\right)}\mathbf{v}_{one}^{(0)}+\mathbf{v}_{zero}^{(0)}$
$\displaystyle=\frac{\exp\left(\mathbf{q}^{(0)T}_{zero}\mathbf{k}^{(0)}_{zero}\right)}{\exp\left(\mathbf{q}^{(0)T}_{zero}\mathbf{k}^{(0)}_{zero}\right)+\frac{n_{1}}{n_{0}}\exp\left(\mathbf{q}^{(0)T}_{zero}\mathbf{k}^{(0)}_{one}\right)}\mathbf{v}_{zero}^{(0)}+\frac{\exp\left(\mathbf{q}^{(0)T}_{zero}\mathbf{k}^{(0)}_{one}\right)}{\frac{n_{0}}{n_{1}}\exp\left(\mathbf{q}^{(0)T}_{zero}\mathbf{k}^{(0)}_{zero}\right)+\exp\left(\mathbf{q}^{(0)T}_{zero}\mathbf{k}^{(0)}_{one}\right)}\mathbf{v}_{one}^{(0)}+\mathbf{v}_{zero}^{(0)}$

Similarly, the ones will be mapped to:

$\displaystyle\mathbf{z}_{one}^{(1)}$
$\displaystyle=\sum_{j}\frac{\exp\left(\mathbf{q}^{(0)T}_{one}\mathbf{k}^{(0)}_{j}\right)}{\sum_{w}\exp\left(\mathbf{q}^{(0)T}_{one}\mathbf{k}^{(0)}_{w}\right)}\mathbf{v}_{j}^{(0)}+\mathbf{v}_{one}^{(0)}$
$\displaystyle=\frac{n_{0}\exp\left(\mathbf{q}^{(0)T}_{one}\mathbf{k}^{(0)}_{zero}\right)}{n_{0}\exp\left(\mathbf{q}^{(0)T}_{one}\mathbf{k}^{(0)}_{zero}\right)+n_{1}\exp\left(\mathbf{q}^{(0)T}_{one}\mathbf{k}^{(0)}_{one}\right)}\mathbf{v}_{zero}^{(0)}+\frac{n_{1}\exp\left(\mathbf{q}^{(0)T}_{one}\mathbf{k}^{(0)}_{one}\right)}{n_{0}\exp\left(\mathbf{q}^{(0)T}_{one}\mathbf{k}^{(0)}_{zero}\right)+n_{1}\exp\left(\mathbf{q}^{(0)T}_{one}\mathbf{k}^{(0)}_{one}\right)}\mathbf{v}_{one}^{(0)}+\mathbf{v}_{one}^{(0)}$
$\displaystyle=\frac{\exp\left(\mathbf{q}^{(0)T}_{one}\mathbf{k}^{(0)}_{zero}\right)}{\exp\left(\mathbf{q}^{(0)T}_{one}\mathbf{k}^{(0)}_{zero}\right)+\frac{n_{1}}{n_{0}}\exp\left(\mathbf{q}^{(0)T}_{one}\mathbf{k}^{(0)}_{one}\right)}\mathbf{v}_{zero}^{(0)}+\frac{\exp\left(\mathbf{q}^{(0)T}_{one}\mathbf{k}^{(0)}_{one}\right)}{\frac{n_{0}}{n_{1}}\exp\left(\mathbf{q}^{(0)T}_{zero}\mathbf{k}^{(0)}_{zero}\right)+\exp\left(\mathbf{q}^{(0)T}_{one}\mathbf{k}^{(0)}_{one}\right)}\mathbf{v}_{one}^{(0)}+\mathbf{v}_{one}^{(0)}$

Assuming that $\mathbf{z}_{zero}^{(1)}\neq\mathbf{z}_{one}^{(1)}$ (to avoid
the trivial case), we notice that the attention mechanism alongside the MLP
$\boldsymbol{\psi}$ 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. ∎

######  Corollary B.10.

Consider a task in which the goal is to count how many $\mathbf{v}_{a}$ tokens
there are in the sequence. Let $\mathbf{v}^{(0)}=[\mathbf{v}\
\mathbf{v}_{a}]^{T}\in\mathbb{R}^{n\times d}$ and
$\mathbf{v}^{*(0)}=[\mathbf{v}\ \mathbf{v}_{a}\
\mathbf{v}_{a}]\in\mathbb{R}^{(n+1)\times 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 $\mathbf{y}_{n}$ and $\mathbf{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. ∎

##  Appendix C Experiments

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.

###  C.1 Experimental Details

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.

For the CoT experiments, we supply examples of the form:

Let’s think step by step, showing me your reasoning. Here are a few examples:
Please perform the following sum: 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 + 1 +
1
We divide the sum into groups of 5. (1 + 1 + 1 + 1 + 1) + (1 + 1 + 1 + 1 + 1)
+ 1 + 1
The answer is then 2 * 5 + 2 = 12
The final answer to your maths question is: 12
Please perform the following sum: 1 + 1 + 1 + 1 + 1 + 1
We divide the sum into groups of 5.
(1 + 1 + 1 + 1 + 1) + 1
The answer is then 1 * 5 + 1 = 6
The final answer to your maths question is: 6

With similar strategies for the $4$ experiments.

#### 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.

###  C.2 Counting with Gemma

We report similar results for the counting experiments using Gemma [24] in
Figure 6 and 7. Compared to Gemini 1.5, Gemma seems to answer less accurately
on the counting prompts.

Figure 6: Gemma 7B LLMs being prompted to (i) sum $1+\dots+1$ (left), (ii)
Count the number of ones in a sequence of 1s (center), and (iii) Count the
number of ones in a sequence of ones and zeroes (the sequence is a Bernoulli
sequence with probability of sampling a one being 0.7) (right). Figure 7:
Frequency of different outputs for Gemma 7B

###  C.3 Synthetic Experiments on Representational Collapse

To provide further experimental evidence of representational collapse with
other positional encodings, we experiment using the original sinusoidal
embeddings from [26]. We sample key, query, and values from a Gaussian
distribution with variance $\sigma^{2}=1/d$, with $d$ the dimension of the
embeddings. We set $d=64$ and otherwise follow the exact structure of the
decoder-only 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 [26]. 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 $\mathbf{x}$ of length $n$ with values uniformly distributed in the
range $[0,1]$. We then create $\mathbf{x}^{*}$ by adding to the first $k=200$
elements of $\mathbf{x}$ noise which is uniformly sampled between $[0,0.1]$.
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 $[0,1]$ and then create a
related sequence by taking its first $k=200$ and adding to these elements
noise sampled uniformly from $[0,0.1]$. 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.