Venue: NAACL 2024

# Transformers Can Represent \(n\)-gram Language Models

Anej Svate   Ryan Cotterell

{asvete, ryan.cotterell}@inf.ethz.ch

###### Abstract

Existing work has analyzed the representational capacity of the transformer architecture by means of formal models of computation. However, the focus so far has been on analyzing the architecture in terms of language _acceptance_. We contend that this is an ill-suited problem in the study of _language models_ (LMs), which are definitionally _probability distributions_ over strings. In this paper, we focus on the relationship between transformer LMs and \(n\)-gram LMs, a simple and historically relevant class of language models. We show that transformer LMs using the hard or sparse attention mechanisms can exactly represent any \(n\)-gram LM, giving us a concrete lower bound on their probabilistic representational capacity. This provides a first step towards understanding the mechanisms that transformer LMs can use to represent probability distributions over strings.

+
Footnote †: _Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers)_, pages 6845–6881

June 16-21, 2024 02024 Association for Computational Linguistics

[https://github.com/rycolab/transformer-mgrams](https://github.com/rycolab/transformer-mgrams)

## 1 Introduction

Neural language models (LMs) have become the backbone of many NLP systems. Their widespread adoption has prompted a plethora of theoretical work investigating what they can and cannot do by studying their representational capacity. Most state-of-the-art LMs are based on the transformer architecture Vaswani et al. (2017), whose theoretical abilities and limitations have been studied extensively; see, e.g., the survey by Strobl et al. (2023). But, many questions remain unanswered. Most existing work studies the architecture in terms of binary language recognition. This introduces a category error between the object of study--an LM, which is definitionally a _distribution_ over strings--and the theoretical abstraction--a _set_ of strings. To amend this discrepancy, we ask: What classes of probability distributions over strings can transformer LMs represent?

Formal models of probabilistic computation provide a natural, well-understood, and precise framework for studying the classes of probability distributions language models can represent. Traditionally, the representational capacity of neural networks, both in terms of lower bounds (what they can provably do) as well as upper bounds (what they can provably _not_ do), has been studied in terms of Boolean sequential models of computation, such as finite-state automata and Turing machines (e.g., Kleene, 1956; Minsky, 1954; Siegelmann and Sontag, 1992; Hao et al., 2018; Merrill, 2019; Merrill et al., 2020, 2022; Merrill and Tsilivis, 2022). Recent work has extended this paradigm to work with probabilistic models of computation Svete and Cotterell, 2023; Nowak et al., 2023), but so far only for LMs based on recurrent neural networks.

However, the sequential nature of classical models makes the connection to the inherently _parallelizable_ transformer architecture less straightforward and has resulted in a number of results upper-bounding their representational capacity Hahn (2020); Bhattamishra et al. (2020); Chiang and Cholak (2022); Hao et al. (2022); Merrill and Sabharwal (2023). We connect transformer LMs to a classical class of LMs that lend themselves particularly well to parallelized computations: \(n\)-gram LMs. We show that both hard as well as sparse attention transformer LMs can represent any \(n\)-gram LM (Theorems 3.1 and 4.1).1 This gives

Figure 1: A transformer LM can simulate a \(4\)-gram LM using \(3\) heads. The stronger arrows from the heads to the symbols show where the heads focus their attention.



us a concrete lower bound on their probabilistic representational capacity. We also study the role of the number of heads (Theorem 3.1) and the number of layers (Theorem 3.2), illustrating a trade-off between the number of heads, layers, and the complexity of the non-linear transformations required for the simulation of \(n\)-gram LMs. Altogether, these results offer a step towards understanding the probabilistic representational capacity of transformer LMs and the mechanisms they might employ to implement formal models of computation.

Footnote 1: The reader is referred to as “_a_-gram LM in the sense of Definition 2.4_”.

## 2 Preliminaries

Let \(\Sigma\) be an alphabet, i.e., a finite, non-empty set of symbols, and \(\Sigma^{*}\) the (infinite) set of all strings formed from symbols of \(\Sigma\). Most modern LMs define \(p\left(\boldsymbol{y}\right)\) for \(\boldsymbol{y}\in\Sigma^{*}\) autoregressively--as a product of conditional probability distributions:

\[p\left(\boldsymbol{y}\right)\stackrel{{\text{def}}}{{=}}p\left( \textsc{eos}~{}|~{}\boldsymbol{y}\right)\prod_{t=1}^{\left|\boldsymbol{y} \right|}p\left(y_{t}~{}|~{}\boldsymbol{y}_{<t}\right). \tag{1}\]

Here, \(\textsc{eos}\notin\Sigma\) is a distinguished end-of-string symbol. The eos symbol enables us to define the probability of a string purely based on the conditional distributions. Such a factorization can be done without loss of generality Du et al. (2023). We define \(\overline{\Sigma}\stackrel{{\text{def}}}{{=}}\Sigma\cup\left\{ \textsc{eos}\right\}\). Further, the conditional probability distributions \(p\left(y_{t}~{}|~{}\boldsymbol{y}_{<t}\right)\) are usually defined based on _vectorial_ representations of \(\boldsymbol{y}_{<t}\) computed by some function enc: \(\Sigma^{*}\rightarrow\mathbb{R}^{D}\). This leads us to the definition of representation-based LMs below.

**Definition 2.1**.: _Let \(\Sigma\) be an alphabet and enc\(\colon\Sigma^{*}\rightarrow\mathbb{R}^{D}\) a **representation function** encoding strings as \(D\)-dimensional representations. Let \(\mathbf{E}\in\mathbb{R}^{\lvert\widetilde{\Sigma}\rvert\times D}\) be an **output matrix**. A **representation-based**LM \(p\) defines the conditional probability distributions \(p\left(y_{t}~{}|~{}\boldsymbol{y}_{<t}\right)\) as2_

Footnote 2: One could, more generally, swap the softmax for any other normalization function, such as the sparesmax Martins and Astudillo (2016). Here, however, we focus on the softmax for conciseness.

\[p\left(y_{t}~{}|~{}\boldsymbol{y}_{<t}\right)\stackrel{{\text{ def}}}{{=}}\mathrm{softmax}\!\left(\mathbf{E}~{}\mathsf{enc}\left( \boldsymbol{y}_{<t}\right)\right)_{yt}. \tag{2}\]

At a high level, we are interested in encoding an arbitrary \(n\)-gram LM using a transformer LM. To do so, we need a notion of equivalence between language models. In this paper, we will work with the following simple definition.

**Definition 2.2**.: _Two LMs \(p\) and \(q\) over \(\Sigma^{*}\) are weakly equivalent if \(p\left(\boldsymbol{y}\right)=q\left(\boldsymbol{y}\right)\) for all \(\boldsymbol{y}\in\Sigma^{*}\)._

This paper precisely explains and proves the following theorem, stated informally below.

**Theorem 2.1** (Informal).: _For every \(n\)-gram LM, there exists a weakly equivalent ({hard, sparse} attention) transformer LM._

### An Aside about Boolean Recognition

Fundamentally, Theorem 2.1 is about weak equivalence (Definition 2.2) between two LMs. In this subsection, we make our case against treating LMs as _recognizers_. The most common manner of analyzing a language model as a recognizer is based on using its representations as an input to a classifier Merrill (2019); Merrill et al. (2020). We recapitulate a common definition below.

**Definition 2.3**.: _Let \(p\) be a representation-based LM with the representation function enc\(\colon\Sigma^{*}\rightarrow\mathbb{R}^{D}\) and let \(g\colon\mathbb{R}^{D}\rightarrow\{0,1\}\) be a classifier. The **binary language** of \(p\) with \(g\) is defined as_

\[\mathcal{L}_{g}\left(p\right)\stackrel{{\text{def}}}{{=}}\{ \boldsymbol{y}\in\Sigma^{*}~{}|~{}g\left(\mathsf{enc}\left(\boldsymbol{y} \right)\right)=1\}. \tag{3}\]

Related is the notion of truncated recognition.

**Definition 2.4** (Hewitt et al. (2020), Definition 4).: _Let \(p\) be a language model over \(\Sigma^{*}\) and \(\alpha>0\). The \(\alpha\)-**truncated language** of \(p\) is defined as_

\[\mathcal{L}_{\alpha}\left(p\right)\stackrel{{\text{def}}}{{=}} \{\boldsymbol{y}\in\Sigma^{*}~{}|~{}p\left(\textsc{eos}~{}|~{} \boldsymbol{y}\right)\geq\alpha \tag{4}\] \[\qquad\quad\mbox{and}~{}p\left(y_{t}~{}|~{}\boldsymbol{y}_{<t} \right)\geq\alpha\quad\forall t\in[\left|\boldsymbol{y}\right|]\}.\]

There are many results in the literature treating transformers' ability to recognize languages in the sense of the two definitions above. For instance, transformers are unable to recognize the Dyck language with more than one bracket type and the Parity language in the sense of Definition 2.3 Hahn (2020), but _can_ recognize _bounded_ Dyck languages in the sense of Definition 2.4 Yao et al. (2021). Our indictment of analyzing \(\mathcal{L}_{g}\left(p\right)\) and \(\mathcal{L}_{\alpha}\left(p\right)\) is that proceeding in such a manner disregards the probabilities assigned to strings by \(p\), which we view as essential to language modeling. Moreover, Definition 2.3 depends on the form of the classifier \(g\) while Definition 2.4 depends on the hyperparameter \(\alpha\). For example, positively classified strings from a language could have their conditional probabilities only slightly above the classification threshold and the negatively classified ones only slightly below the threshold Hahn (2020), which hides the true distribution defined by the 

LM. In this context, our contention is that \({\cal L}_{g}\left(p\right)\) and \({\cal L}_{\alpha}\left(p\right)\) are not useful definitions for studying the representational capacity of LMs. Instead, we advocate for analyzing LMs as _probabilistic_ formal languages.

### Language Modeling with _n-_grams

The next-symbol probabilities in _n-_gram LMs are computed under the _n-_gram assumption.

**Assumption 2.1**.: _The **n-gram assumption** states that the conditional probability of the symbol \(y_{t}\) given \(\boldsymbol{y}_{t+1}\) only depends on n \(-1\) previous symbols \(\boldsymbol{y}^{t-1}_{t-n+1}\stackrel{{\text{\tiny def}}}{{=}} y_{t-1},\ldots,y_{t-n+1}\):_

\[p\left(y_{t}~{}|~{}\boldsymbol{y}_{<t}\right)=p\left(y_{t}~{}|~{}\boldsymbol{ y}^{t-1}_{t-n+1}\right). \tag{5}\]

_We will refer to \(\boldsymbol{y}^{t-1}_{t-n+1}\) as the **history** of \(y_{t}\)._

Padding.Eq. (5) assumes the existence of _n \(-1\)_preceding symbols that define the conditional distribution \(p\left(y_{t}~{}|~{}\boldsymbol{y}_{<t}\right)\). To ensure this is the case even at the beginning of the string, it is standard to _pad_ the input string with _n \(-1\)_beginning-of-string symbols bos. For ease of notation, we index the _n \(-1\)_bos tokens with indices \(-n+2,\dots,0\) so that _n-_gram LMs conveniently fit the autoregressive factorization from Eq. (1). We also define \(\underline{\Sigma}\stackrel{{\text{\tiny def}}}{{=}}\Sigma \cup\left\{\textsc{bos}\right\}\).

Despite their simplicity, _n-_gram LMs have a storied place in language modeling Shannon (1948); Baker (1975a, b); Jelinek (1976); Bahl et al. (1983); Belgio et al. (2000, 2003, 2006); Schwenk (2007); Heafield (2011); Heafield et al. (2013). Because the conditional probabilities of _n-_gram LMs only depend on the previous _n \(-1\)_symbols, different parts of the string can be processed independently, i.e., in parallel. This facilitates a natural connection to transformer LMs since parallelizability is a prevalent feature of the architecture and one of its main advantages over other neural LMs such as RNN LMs Vaswani et al. (2017).

### Transformer Language Models

Transformer LMs are LMs whose conditional distributions \(p\left(y_{t}~{}|~{}\boldsymbol{y}_{<t}\right)\) are computed by a **transformer**. A transformer is a composition of multiple transformer **layers**, each of which implements the **attention mechanism**. We give definitions of these building blocks in what follows.

Notation.We use bold, unitalicized letters such as \(\mathbf{x}\in\mathbb{R}^{D}\) to denote real-valued vectors and italcized letters \(x_{j}\in\mathbb{R}\) for their entries. Capital bold letters such as \(\mathbf{X}\in\mathbb{R}^{N\times D}\) denote matrices. All vectors are _column_ vectors unless transposed. We define the vertically stacking operator \(\left(\cdot~{};\cdots~{};~{}\cdot\right)\), which denotes the vertical concatenation of the \(D\)-dimensional _column_ vectors \(\mathbf{x}_{1},\dots,\mathbf{x}_{N}\) into a \(ND\)-dimensional vector \(\left(\mathbf{x}_{1};\cdots;\mathbf{x}_{N}\right)\in\mathbb{R}^{ND}\) and the concatenation of the \(D\)-dimensional _row_ vectors \(\mathbf{x}_{1}^{\top},\dots,\mathbf{x}_{N}^{\top}\) into a matrix \(\mathbf{X}\in\mathbb{R}^{N\times D}\) with \(N\) rows and \(D\) columns. Given the matrix \(\mathbf{X}=\left(\mathbf{x}_{1}^{\top};\cdots;\mathbf{x}_{N}^{\top}\right)\), we write \(\mathbf{X}_{n}=\left(\mathbf{x}_{1}^{\top};\cdots;\mathbf{x}_{n}^{\top}\right)\) for the submatrix composed of the first \(n\) rows. We call a function \(f\colon\mathbb{R}^{D}\times\mathbb{R}^{D}\to\mathbb{R}\) whose purpose is to evaluate the compatibility of two vectors a **scoring function**. A **normalization function**\(\pi\colon\mathbb{R}^{N}\to\mathbf{\Delta}^{N-1}\) maps vectors in \(\mathbb{R}^{N}\) to \(N\) probabilities. Here, \(\mathbf{\Delta}^{N-1}\stackrel{{\text{\tiny def}}}{{=}}\left\{ \mathbf{x}\in\left[0,1\right]^{N}\mid\sum_{n=1}^{N}x_{n}=1\right\}\) is the \(\left(N-1\right)\)-dimensional probability simplex. This notation is summarized in Tab. 1.

The Attention Mechanism.The attention mechanism works as follows. It takes a **query** vector \(\mathbf{q}\in\mathbb{R}^{D}\) and two matrices: The matrix \(\mathbf{K}\in\mathbb{R}^{N\times D}\) of **keys** and the matrix \(\mathbf{V}\in\mathbb{R}^{N\times D}\) of **values** and computes a weighted average of the value vectors based on the compatibilities of the key vectors and the query vector, as determined by a scoring function \(f\). A formal definition is given below.

**Definition 2.5** (Attention Mechanism).: _Let \(f\) be a scoring function and \(\pi\) a normalization function. Let \(\mathbf{q}\in\mathbb{R}^{D}\) be a query vector and let \(\mathbf{K}=\left(\mathbf{k}_{1}^{\top};\cdots;\mathbf{k}_{N}^{\top}\right)\in \mathbb{R}^{N\times D}\) and \(\mathbf{V}=\left(\mathbf{v}_{1}^{\top};\cdots;\mathbf{v}_{N}^{\top}\right)\in \mathbb{R}^{N\times D}\) be matrices of keys and values, respectively. An **attenion mechanism** Att: \(\mathbb{R}^{D}\times\mathbb{R}^{N\times D}\times\mathbb{R}^{N\times D}\to \mathbb{R}^{D}\) is defined as_

\[\texttt{Att}\left(\mathbf{q},\mathbf{K},\mathbf{V}\right)\stackrel{{ \text{\tiny def}}}{{=}}\sum_{n=1}^{N}s_{n}\mathbf{v}_{n}, \tag{6}\]

_where_

\[\mathbf{s}\stackrel{{\text{\tiny def}}}{{=}}\mbox{$\pi$}\left(f \left(\mathbf{q},\mathbf{k}_{1}\right),\dots,f\left(\mathbf{q},\mathbf{k}_{N} \right)\right) \tag{7}\]

Figure 2: An illustration of how a \(4\)-gram LM computes the probability of a string. All conditional probabilities can be computed in parallel and then multiplied into the probability of the entire string.

 

is the vector of normalized scores between the query \({\bf q}\) and the keys in \({\bf K}\)._

The most standard implementation of the scoring function \(f\) is the (scaled) inner product \(f\left({\bf q},{\bf k}\right)\stackrel{{\text{def}}}{{=}}\langle{ \bf q},{\bf k}\rangle\). Some of our results rely on this standard formulation. However, some also rely on the more general, but still simple and just as efficiently computable scoring functions.

**The Transformer Architecture.** A transformer layer uses the attention mechanism to compute augmented representations \({\bf z}_{t}=\operatorname*{\sf Att}\left({\bf q}_{t},{\bf K}_{t},{\bf V}_{t}\right)\) of the input representations \({\bf X}_{t}=\left({\bf x}_{1};\cdots;{\bf x}_{t}\right)\). The query \({\bf q}_{t}\), the keys \({\bf K}_{t}\), and values \({\bf V}_{t}\) are all transformations of the input representations \({\bf X}_{t}\).

**Definition 2.6**.: _Let \(Q,K,V,O\colon\mathbb{R}^{D}\to\mathbb{R}^{D}\) be the query, key, value, and **output** functions. A **transformer layer** is a function \(\mathcal{L}\colon\mathbb{R}^{T\times D}\to\mathbb{R}^{T\times D}\) that computes_

\[\mathcal{L}\left({\bf x}_{1}^{\top};\ldots;{\bf x}_{T}^{\top}\right)=\left({ \bf z}_{1}^{\top};\ldots;{\bf z}_{T}^{\top}\right)\in\mathbb{R}^{T\times D} \tag{8}\]

_for \(t\in[T]\) where_

\[{\bf a}_{t} \stackrel{{\text{def}}}{{=}}\operatorname*{\sf Att }\left({\bf q}_{t},{\bf K}_{t},{\bf V}_{t}\right)+{\bf x}_{t}\in\mathbb{R}^{D} \tag{9a}\] \[{\bf z}_{t} \stackrel{{\text{def}}}{{=}}O\left({\bf a}_{t}\right) +{\bf a}_{t} \in\mathbb{R}^{D}. \tag{9b}\]

_Here, we define_

\[{\bf q}_{t} \stackrel{{\text{def}}}{{=}}Q\left({\bf x}_{t}\right) \in\mathbb{R}^{D} \tag{10a}\] \[{\bf K}_{t} \stackrel{{\text{def}}}{{=}}\left(K\left({\bf x}_{1 }\right)^{\top};\cdots;K\left({\bf x}_{t}\right)^{\top}\right) \in\mathbb{R}^{t\times D}\] (10b) \[{\bf V}_{t} \stackrel{{\text{def}}}{{=}}\left(V\left({\bf x}_{1 }\right)^{\top};\cdots;K\left({\bf x}_{t}\right)^{\top}\right) \in\mathbb{R}^{t\times D}. \tag{10c}\]

Note: _For simplicity, we do not include layer normalization._

Without further modification, the transformations applied by the transformer layer are position-invariant, which necessitates the addition of explicit positional information.

**Definition 2.7**.: _A position-augmented symbol **representation function**\({\bf r}\colon\Sigma\times\mathbb{N}\to\mathbb{R}^{D}\) is a function representing symbols and their positions as \(D\)-dimensional vectors._

Position-augmented symbol representation functions are often implemented as an addition or concatenation of separate symbol-only and position-only representation functions Vaswani et al. (2017). Here, we define it more generally as any function of the symbol and its position.

**Definition 2.8**.: _A **static encoding**\(\mathcal{R}\) is a function \(\mathcal{R}\colon\Sigma^{T}\to\mathbb{R}^{T\times D}\) defined for any \(T\in\mathbb{N}\) as_

\[\mathcal{R}\left(\boldsymbol{y}\right)\stackrel{{\text{def}}}{{=} }\left({\bf r}\left({y}_{1},1\right)^{\top};\cdots;{\bf r}\left({y}_{T},T \right)^{\top}\right). \tag{11}\]

Multiple transformer layers are stacked into a transformer, which computes the (deep) contextual representations of all symbols in the string.

**Definition 2.9**.: _For \(L\in\mathbb{N}\), let \(\mathcal{L}_{\ell}\) for \(\ell\in[L]\) be transformer layers. Let \(\mathcal{R}\) be a static encoding. An \(L\)-layer **transformer**\(\mathcal{T}\) is defined as_

\[\mathcal{T}\left(\mathcal{R}\right)\stackrel{{\text{def}}}{{=}} \mathcal{L}_{L}\circ\cdots\circ\mathcal{L}_{1}\circ\mathcal{R}. \tag{12}\]

A transformer computes the contextual representations of the symbols \(\boldsymbol{y}=y_{1}\cdots y_{T}\) as

\[\left({\bf x}_{1}^{L\top};\cdots;{\bf x}_{T}^{L\top}\right)\stackrel{{ \text{def}}}{{=}}\mathcal{T}\left(\mathcal{R}\right)\left(\boldsymbol{y} \right). \tag{13}\]

\begin{table}
\begin{tabular}{l l l} \hline \hline Symbol & Type & Meaning \\ \hline \([N]\) & \(\subset\mathbb{N}\) & The set \(\{1,\dots,N\}\) for \(N\in\mathbb{N}\). \\ \(y\) & \(\in\Sigma\) & A symbol, element of \(\Sigma\). \\ \(\Sigma,\underline{\Sigma},\overline{\Sigma}\) & alphabet & \(\Sigma\) is a set of symbols, \(\underline{\Sigma}\stackrel{{\text{def}}}{{=}}\Sigma\cup\{\texttt{bos}\}\), \(\overline{\Sigma}\stackrel{{\text{def}}}{{=}}\Sigma\cup\{\texttt{ Ecos}\}\) \\ \(\boldsymbol{y}\) & \(\in\Sigma\)* & A string over \(\Sigma\). \\ \(\boldsymbol{y}_{j}^{i}\) & \(\in\Sigma\)* & A substring of \(\boldsymbol{y}\), a string. \\ \([y]\) & \(\in\{0,1\}^{|\Sigma|}\) & One-hot encoding of the symbol \(y\in\Sigma\). \\ \(D\) & \(\in\mathbb{N}\) & Size of the contextual representations in the transformer. \\ \(\mathbf{\Delta}^{N-1}\) & \(\subseteq\mathbb{R}^{N}\) & The \(N-1\)-dimensional probability simplex. \\ \(f\) & \(\mathbb{R}^{D}\times\mathbb{R}^{D}\to\mathbb{R}\) & A scoring function. \\ \(\pi\) & \(\mathbb{R}^{N}\to\mathbf{\Delta}^{N-1}\) & A normalization function. \\ \(Q,K,V,O\) & \(\mathbb{R}^{D}\to\mathbb{R}^{D}\) & The query, key, value, and output functions. \\ \(F\) & \(\mathbb{R}^{D}\to\mathbb{R}^{D}\) & The final transformer LM transformation function. \\ enc & \(\Sigma\)* \(\to\mathbb{R}^{D}\) & The string representation function. \\ \({\bf r}\) & \(\underline{\Sigma}\times\mathbb{N}\to\mathbb{R}^{D}\) & The position-augmented representation function. \\ \(L\) & \(\in\mathbb{N}\) & Number of layers. \\ \(H\) & \(\in\mathbb{N}\) & Number of heads. \\ \(\mathcal{H}\) & \(\mathbb{R}^{HD}\to\mathbb{R}^{D}\) & The head combining function. \\ \((\cdot;\cdots;\cdot)\) & Vertical concatenation operator of vectors or matrices. \\ \hline \hline \end{tabular}
\end{table}
Table 1: A summary of the notation used in the paper.

 

If \({\cal R}\) is clear from the context or arbitrary, we will omit it as an argument to \({\cal T}\) and just write \({\cal T}\left(\mathbf{y}\right)\).

**Definition 2.10**.: _Let \({\cal T}\) be a transformer, \(F\colon\mathbb{R}^{D}\to\mathbb{R}^{D}\) the final representation transformation function, and \(\overline{\mathbf{y}}\in\Sigma^{*}\) with \(|\mathbf{y}|=T\). We define_

\[\mathsf{enc}\left(\mathbf{y}\right)\stackrel{{\text{def}}}{{=}}F \left(\mathbf{x}_{T}^{L}\right), \tag{14}\]

_where \(\mathbf{x}_{T}^{L}\) is the representation of the \(T^{\text{th}}\) symbol in \(\mathbf{y}\) computed by \({\cal T}\), i.e., \(\left(\mathbf{x}_{1}^{L\top};\cdots;\mathbf{x}_{T}^{L\top}\right)={\cal T} \left(\mathbf{y}\right)\)._

Transformer Language Models.So far, we have only defined how the transformer architecture can be used to compute the contextual representations of the symbols. To complete the definition, we define a transformer _language model_ as follows.

**Definition 2.11**.: _A **transformer LM**\(p_{\mathcal{T}}\) is the representation-based autoregressive LM with the representation function \(\mathsf{enc}\) from Eq. (14). That is, \(p_{\mathcal{T}}\) defines the conditional probability distributions_

\[p_{\mathcal{T}}\left(y_{t}~{}|~{}\mathbf{y}_{<t}\right)\stackrel{{ \text{def}}}{{=}}\mathrm{softmax}(\mathbf{E}\,\mathsf{ enc}\left(\mathbf{y}_{<t}\right))_{y_{t}}. \tag{15}\]

#### 2.3.1 Variants of the Attention Mechanism

In this subsection, we discuss many common variants of the attention mechanism. First, **multi-headed** attention uses \(H\)**attention heads** to compute \(H\) representations of the symbols in the string. The representations constructed by the different attention heads are concatenated into a long vector and projected down to the output size of a single head with a head-combinner function \({\cal H}\).

**Definition 2.12**.: _For \(L,H\in\mathbb{N}\), let \({\cal L}_{\ell}^{h}\colon\mathbb{R}^{T\times D}\to\mathbb{R}^{T\times D},\ell \in[L],h\in[H]\) be transformer layers. Define \({\cal L}_{\ell}\colon\mathbb{R}^{T\times D}\to\mathbb{R}^{T\times HD}\) as_

\[{\cal L}_{\ell}\left(\mathbf{X}\right)\stackrel{{\text{def}}}{{=} }\left({\cal L}_{\ell}^{1}\left(\mathbf{X}\right)^{\top};\cdots;{\cal L}_{\ell }^{H}\left(\mathbf{X}\right)^{\top}\right)^{\top}. \tag{16}\]

_Furthermore, let \({\cal H}\colon\mathbb{R}^{HD}\to\mathbb{R}^{D}\). An \(L\)-layer transformer with \(H\) heads computes:_

\[{\cal T}\left({\cal R}\right)\stackrel{{\text{def}}}{{=}}{\cal L }_{L}\circ{\cal H}\circ\cdots\circ{\cal H}\circ{\cal L}_{1}\circ{\cal R}, \tag{17}\]

_where \({\cal H}\) is applied row-wise to project the representations of \(H\) heads to \(\mathbb{R}^{D}\)._

Attention types.Attention weights are computed by normalizing the scores \(f\left(\mathbf{q},\mathbf{k}_{1}\right),\ldots,f\left(\mathbf{q},\mathbf{k}_{ t}\right)\). The choice of the projection function \(\mathbf{\pi}\) determines the type of attention and has concrete implications on representational capacity Hao et al. (2022a).

**Definition 2.13**.: _Hard attention is computed with the hardmax projection function:_

\[\operatorname{hardmax}\left(\mathbf{x}\right)_{d}\stackrel{{ \text{def}}}{{=}}\begin{cases}\frac{1}{m}&\textbf{if $d\in \operatorname{argmax}\left(\mathbf{x}\right)$}\\ 0&\textbf{otherwise}\end{cases} \tag{18}\]

_for \(d\in[D]\), where \(\mathbf{x}\in\mathbb{R}^{D}\) and \(m\stackrel{{\text{def}}}{{=}}|\operatorname{argmax}\left(\mathbf{ x}\right)|\) is the cardinality of the argmax set._

We also introduce _sparse_ attention, which uses the sparsemax normalization function to compute the attention weights.

**Definition 2.14**.: _Sparse attention is computed with the \(\operatorname{sparsemax}\) projection function:_

\[\operatorname{sparsemax}(\mathbf{x})\stackrel{{\text{def}}}{{=}} \operatornamewithlimits{argmin}_{\mathbf{p}\in\boldsymbol{\Delta}^{D-1}}\| \mathbf{p}-\mathbf{x}\|_{2}^{2}. \tag{19}\]

## 3 Hard Attention Transformer LMs

This section presents a set of results describing the representational capacity of hard attention transformer LMs. Concretely, we show that transformer LMs with hard attention can represent \(n\)-gram LMs, either using \(n-1\) heads (Theorem 3.1) or \(n-1\) layers (Theorem 3.2). Simulation is possible even with a single head and a single layer (Theorem 3.3) but might require a more elaborate set of non-linear transformations and positional encodings whose precision scales linearly with the string length.

**Theorem 3.1**.: _For any \(n\)-gram LM, there exists a weakly equivalent single-layer hard attention transformer LM with \(n-1\) heads._

Proof intuition.Given an \(n\)-gram LM \(p\), we can construct a weakly equivalent LM \(p_{\mathcal{T}}\) defined by a transformer \({\cal T}\) that looks back at the preceding \(n-1\) positions using \(n-1\) heads, each of them uniquely attending to exactly one position. The symbols attended to can be used to identify the full history, which can be used to access the conditional distribution over the next symbol. This is illustrated in Fig. 1. See Appendix B.2 for the full proof. \(\blacksquare\)

Theorem 3.1 shows that transformer LMs with hard attention can represent \(n\)-gram LMs, establishing, to the best of our knowledge, the first concrete relationship between transformer LMs and probabilistic languages. A natural follow-up question then is whether \(n-1\) heads are _necessary_ to correctly simulate an \(n\)-gram LM. Besides aiming to illuminate different mechanisms enabling the implementation of classical LMs, this 

question also follows the line of inquiry about the _uniqueness_ and _interpretability_ of the representations of formal models by neural LMs Liu et al. (2023). The following two theorems show that the intuitive construction using \(n-1\) heads is far from unique: Theorem 3.2 shows that a similarly simple simulation is possible with \(n-1\) layers and a single head, while Theorem 3.3 shows that even a transformer LM with a single head and a single layer can simulate an \(n\)-gram LM, albeit with more complex position invariant transformation \(F\). This suggests that there is no canonical way of determining whether a transformer LM has learned an \(n\)-gram LM by looking at individual components (e.g., positions attended to by the different heads).

**Theorem 3.2**.: _For any n-gram LM, there exists a weakly equivalent \((n-1)\)-layer hard attention transformer LM with a single head._

Proof intuition..: Whereas the transformer LM constructed in Theorem 3.1 used \(n-1\) heads to look at all the \(n-1\) positions of interest, an \(n-1\)-layer transformer LM can use the \(n-1\) layers to look back at the _immediately_ preceding position and copy it forward \(n-1\) times (keeping the current symbol there as well). After \(n-1\) layers of such transformations, the entire history can be read from the current contextual representation. See Appendix B.2 for the full proof. 

Apart from using hard attention, both transformer LMs used in Theorems 3.1 and 3.2 rely on modeling assumptions often found in practical implementations of the transformer: The transformations \(Q,K\), and \(V\) are linear functions, the scoring function is implemented as a dot-product and positional encodings are bounded. This makes the results comparable to practical implementations. The following theorem, in contrast, shows that if we permit the use of less standard components, transformer LMs can identify the history of interest using only a single head and a single layer.

**Theorem 3.3**.: _For any n-gram LM, there exists a weakly equivalent single-layer hard attention transformer LM with a single head._

Proof intuition..: The bulk of this construction lies in the encoding \({\boldsymbol{y}}^{t-1}_{t-n+1}\) in a vector that can be constructed by a single attention head in one layer. This is done by an attention head that _[1]_ puts non-zero attention on only the previous \(n-1\) symbols and _[2]_ encodes the identities and the positions of symbols in a \(|\Sigma|\)-dimensional value vector. The value vector can then be decoded into a one-hot encoding of \({\boldsymbol{y}}^{t-1}_{t-n+1}\) by an \(n-1\)-layer MLP that defines \(F\), which allows us to match the conditional probabilities of the \(n\)-gram LM as in Theorems 3.1 and 3.2. See Appendix B.2 for the full proof. 

## 4 Sparse Attention Transformer LMs

While the results in SS3 concretely characterize the abilities of hard attention transformer LMs, the assumption of hard attention is somewhat removed from practical implementations of the model. Those most often rely on differentiable normalization functions, such as the softmax.3 However, the full support of the softmax function makes the connection to formal models of computation difficult Hahn (2020). To bring the theoretical models closer to practical implementations yet still be able to make clear analogies to formal models of computation, we now consider _sparse_ attention transformers, which use the sparsemax normalization function. The sparsity allows sparse attention transformers to simulate \(n\)-gram LMs just like hard attention transformers while relying on differentiable operations.

Footnote 3: As noted in §1, the analysis of soft attention transformers requires a different type of analysis in terms of approximation of the probabilities. A complete study would have to consider the approximation over _arbitrarily_ long strings (since \(\Sigma*\) is an infinite set), which is difficult by simply scaling model parameters to a large constant. We focus on exact simulation here, but conjecture that soft attention transformers can approximate LMs whose _average_ string length is finite.

**Theorem 4.1**.: _For any n-gram LM, there exists a weakly equivalent single-layer sparse attention transformer LM with \(n-1\) heads._

Proof intuition..: The intuition behind the simulation with sparse attention is similar to the hard attention one; each head attends to a single position, as illustrated in Fig. 1. Effectively, the construction results in a sparse attention transformer that simulates hard attention. In contrast to Theorems 3.1 and 3.2, we here require a model with _unbounded_ positional encodings and a non-linearly transformed dot-product scoring function. Intuitively, the unbounded positional encodings are required to scale the unnormalized attention scores to differ enough for the sparsemax to focus on a single position. The rest of the proof follows that of Theorem 3.1; see Appendix C for the details. 

Theorem 4.1 (representing an \(n\)-gram LM with \(n-1\) heads) could naturally be extended to analogs 

of Theorem 3.2 (representing an \(n\)-gram LM with \(n-1\) layers) and Theorem 3.3 (representing an \(n\)-gram LM with a single head and a single layer) using a similar adaptation of the construction from the hard attention case to the sparse attention one as in Theorem 4.1.

## 5 Space Complexity

In SS3 and SS4, we describe lower bounds that tell us what types of probability distributions transformer LMs _can_ represent, but do not say how _efficiently_ they can do so. The space complexity of simulating \(n\)-gram LMs is discussed in this section. We focus on hard-attention transformer LMs with multiple heads or multiple layers (Theorems 3.1 and 3.2) since their modeling assumptions (bar hard attention) are closest to practical implementations. The constructive proofs of Theorems 3.1 and 3.2 allow us to directly analyze the space requirements for the simulation of \(n\)-gram LMs, both in terms of _(1)_ the size of the contextual representations \({\bf X}^{h}_{\ell}\) as well as _(2)_ the number of bits required to represent the individual entries of the vectors \({\bf x}^{h}_{\ell,t}\).

### Scaling with Respect to the Number of Computational Steps

We first address the second point. Specifically, we are interested in how the number of bits scales with respect to \(t\), the number of computational steps performed during the generation of a string \(\boldsymbol{y}\in\Sigma^{*}\). As summarized by Tab. 2, the models constructed in the proofs of Theorems 3.1 and 3.2 use positional encodings with entries of the form \(\sqrt{\frac{1}{t}}\) for \(t\in\mathbb{N}\), i.e., they contain square roots of rational numbers. This makes the scaling of the space complexity difficult, as square roots of rational numbers are not in general representable with a finite number of bits. While this might seem discouraging, we emphasize that these specific positional encodings were only used to keep the contextual representations bounded and the scoring function in line with the original formulation Vaswani et al. (2017) and concurrent work Merrill and Sabharwal (2023). A closer look at the constructions in the proof of Theorems 3.1, 3.2 and 4.1 reveals that simpler (but unbounded) positional encodings with a less standard scoring function can be used to the same effect. In particular, we can use positional encodings that contain entries of the form \(t\), which only require a _logarithmic_ number of bits with respect to \(t\). Since such scaling is required to uniquely encode the positional information in general Merrill and Sabharwal (2023c), this represents an asymptotically optimal scaling of the space complexity of the contextual representations.4

Footnote 4: The construction in Theorem 3.3, in contrast, relies on encoding the entire preceding string in a single dimension with one digit per position in the string, which requires a number of bits that scales linearly with respect to the string length. A simplification to a logarithmic number of bits does not seem as straightforward.

Importantly, \(n\)-gram LMs are _real-time_: they, by definition, generate a symbol at each step of the computation. This means that the scaling of the space complexity with respect to the number of computation steps coincides with its scaling with respect to the length of the generated string \(\boldsymbol{y}\)--the scaling is logarithmic in \(|\boldsymbol{y}|\). This is in contrast to non-real-time models of computation which might not generate a symbol at each step of the computation; while those might still require an asymptotically optimal scaling with respect to \(t\), the additional computational steps that do not emit any symbol might mean that the space complexity is _unbounded_ with respect to the length of the generated string. An example of such a model is a transformer LM simulating a (probabilistic) Turing machine, which would require the model to not emit symbols at some points of the computation Nowak et al. (2023, 2024).

### The Dimensionality of the Contextual Representations

We now discuss the size of the contextual representations required for the simulation of \(n\)-gram LMs. From a high level, we have to consider two stages: _(1)_ the contextual representations \({\bf x}^{h}_{\ell,t}\) of the different layers and heads and _(2)_ the size of the final representation \(\mathsf{enc}\left(\boldsymbol{y}\right)\). The contextual representations \({\bf x}^{h}_{\ell,t}\) in stage _(1)_ are composed of the symbol and positional encodings. The symbol representations include (two copies of) the one-hot encodings while the positional encodings include between two and \(2n\) dimensions encoding positional information. This means that the per-head and per-layer space complexity scales with \(|\Sigma|\) and \(n\). For stage _(2)_ we use the one-hot encodings of the entire _history_\(\boldsymbol{y}^{l-1}_{t-n+1}\), with which we index the matrix of \(|\underline{\Sigma}|^{n-1}\) conditional probabilities defined by the \(n\)-gram LM. While the size of \(\mathsf{enc}\left(\boldsymbol{y}\right)\) is theoretically only lower-bounded by \(|\overline{\Sigma}|\) Yang et al. (2018); Svate and Cotterell (2023),6 reducing its size requires a lower 

rank decomposition of the matrix of conditional probabilities and a corresponding reparametrization of the contextual symbol representation. This might require a blow-up in the number of bits required to represent individual dimensions, or, in general, result in a real-valued vector that could not be represented on a finite-precision system. Altogether, most of the space complexity of the contextual representations comes from the one-hot encodings in \(\mathsf{enc}\left(\mathbf{y}\right)\), which require \(|\Sigma|^{n-1}\) dimensions.

The discussion in this section can be summarized by the following theorem.

**Theorem 5.1**.: _Let \(p\) be an \(n\)-gram LM over the alphabet \(\Sigma\). There exists a weakly equivalent hard-attention transformer LM with contextual representations \(\mathbf{x}\) of size \(\mathcal{O}\left(n|\Sigma|\right)\) and representation \(\mathsf{enc}\left(\mathbf{y}\right)\) of size \(\mathcal{O}\left(|\Sigma|^{n-1}\right)\). Each of \(\mathbf{x}\)'s entries can be represented with \(\mathcal{O}\left(\log_{2}\left(|\mathbf{y}|\right)\right)\) bits for \(\mathbf{y}\in\Sigma^{*}\)._

## 6 Discussion and Related Work

To the best of our knowledge, SS3 and SS4 provide the first results on the probabilistic representational capacity of transformer LMs.

### The relevance of \(n\)-gram LMs to modern LMs.

One might rightfully question the utility of connecting the state-of-the-art language modeling architecture to \(n\)-gram LMs. LMs based on the \(n\)-gram assumption indeed constitute some of the simplest and least expressive classes of probability distributions. Nevertheless, \(n\)-gram LMs provide a useful playground and theoretical foundation for contextualizing and interpreting the inner workings of modern LMs. For example, \(n\)-gram LMs have been found to comprise a crucial component of _in-context learning_, where attention heads in different layers of the model together identify individual \(n\)-grams and base their predictions on the presence of such \(n\)-grams Olsson et al. (2022); Akyurek et al. (2024). The presence of specific \(n\)-grams might therefore present the foundations of in-context-based knowledge of transformer LMs. As such, understanding the requirements for correct simulation of \(n\)-gram LMs is important for a thorough grasp of the abilities of LMs to learn in context. Encouraging \(n\)-gram-LM-based learning has also been observed to aid in-context learning abilities Akyurek et al. (2024). Existing work has also linked \(n\)-gram LMs to other neural network architectures such as one-dimensional convolutional neural networks Merrill (2019). Moreover, the theoretical grounding in classical formal language theory makes precise statements about \(n\)-gram LMs possible while their simplicity makes them inherently interpretable and easy to analyze. \(n\)-gram LMs also have a cognitive interpretation Jager and Rogers (2012). Most importantly, however, \(n\)-gram LMs lend themselves to _paralellized_ processing, which affords a succinct and natural connection to transformer LMs and has been suggested to be a crucial part of any theoretical treatment of transformer LMs Strobl et al. (2023); Merrill and Sabharwal (2023).

#### Understanding the limitations and abilities of hard attention transformer LMs.

Our results are the strongest in the hard attention setting, which follows the trend of using hard attention in theoretical treatments. Hard attention makes singling out the important aspects of the string (in our case, the history) possible. Concretely, we showcase three different mechanisms that make it possible for transformer LMs to implement \(n\)-gram LMs with hard attention and investigate the role of the number of heads and layers. Particularly, our constructions suggest a possible mechanism in which the different transformer heads or layers can specialize in focusing on different positions in the string, a feature that has previously been suggested as an explanation of how transformer LMs process strings Elhabe et al. (2021) and has been observed in trained transformer LMs Olsson et al. (2022); Akyurek et al. (2024).6 Our presentation of multiple orthogonal mechanisms that can simulate \(n\)-gram LMs equivalently well is another confirmation of the observation that algorithmic principles learned or implemented by neural LMs do not always correspond to a _single_ intuitive implementation of a formal model of computation, which has concrete implications on interpretability methods for the architecture. This has been observed in practical scenarios and warns us that focusing on individual components of the model (that is, using myopic interpretability methods) might result in misleading interpretability results Wen et al. (2023).

Footnote 6: Note that the constructions presented in this paper are purely meant to showcase the existence of a mechanism that can be used to simulate \(n\)-gram LMs; we do not suggest that the same mechanisms will be employed by models used in practice, which is also why do not present any empirical results. For example, the use of the sparse one-hot encodings differs from the standard dense representations of symbols.

#### Probabilistic representational capacity.

We augment existing literature by providing an explicit connection between transformer L 

language models. In line with work comparing transformers to circuits (Merrill and Sabharwal, 2023c,a; Weiss et al., 2021; Hao et al., 2022b; Merrill and Sabharwal, 2023b; Merrill et al., 2022; Chiang et al., 2023; Angluin et al., 2023), we also show the utility of analyzing transformers with parallelizable models of computation, which go hand in hand with the parallelizable nature of the transformer architecture. Moreover, the formalization of an LM with the output matrix indexed by the contextual representation (cf. Definition 2.11) makes it easy to connect existing results on the expressivity of the transformer architecture with the probabilistic setting by defining a mapping from the contextual representation to the conditional probability distribution through the output matrix. We suppose other simple and parallelizable classes of distributions might lend themselves well to similar probabilistic treatments; probabilistic analysis may be particularly interesting in the context of circuit complexity with sigmoid-activated circuits (Maass et al., 1991), (Smolensky et al., 1996, Chapter 4).

Connection to (sub-)regular languagesThis work focuses on connecting transformer LMs to _n_-gram LMs, which are a special instance of the more general class of _sub-regular LMs_. In words, sub-regular LMs are LMs that can be described without using the full power of the probabilistic finite-state automata (Jager and Rogers, 2012). In this sense, sub-regular LMs fall below regular languages in the Chomsky hierarchy and themselves define their own hierarchy (Jager and Rogers, 2012; Heinz and Rogers, 2013; Avcu et al., 2017). For example, some classes of sub-regular languages do not require sequential processing of the input string that is usually required by finite-state automata for correct recognition. The intuitive connection between transformer LMs and   _n_-gram LMs encourages further work on the connections between other classes of (sub-)regular LMs and transformer LMs. Transformers have been linked to (sub)-regular languages by Yao et al. (2021) and Liu et al. (2023). Yao et al. (2021) study the ability of transformers to generate bounded hierarchical languages using a transformer but do not extend their analysis to the fully probabilistic case.7Liu et al. (2023) connect transformers with both hard and soft attention to general finite-state automata (FSAs), which define a strictly larger set of languages than   _n_-gram languages. This additional generality, however, comes with some caveats. Liu et al. (2023, Theorem 1)--their most general result--relies on a model whose depth _scales_ (logarithmically) with the string length. As such, no particular finite-size transformer construction can simulate FSAs on strings of _arbitrary_ length (which is required for equivalence). This is in contrast to our results, which feature transformers of _fixed size_ with respect to the string length. While Liu et al. (2023, Theorem 2) also provides a result using a finite-depth transformer for a subset of FSAs, that construction results in a network much bigger and deeper than ours. For example, their construction would result in \(\left(|\Sigma|^{n-1}\right)^{2}(n-1)\log|\Sigma|\) layers in contrast to \(n-1\) layers with a single head in our case. Our focus on  _n_-gram LMs thus allows for a more compact and simpler representation. Worth noting is that despite being close to our analysis, none of the related work treats probability distributions over strings but rather focuses on the binary decision of whether a string belongs to a language or not based on the conditional probabilities output by the model, or, in the case of Liu et al. (2023), only the computation of state sequences. Transformer LMs are studied probabilistically by Xie et al. (2022), who provide a Bayesian interpretation of their ability to learn _in context_ by studying the learning abilities of hidden Markov models (HMMs). While HMMs are equivalent to probabilistic finite-state automata, Xie et al. (2022) do not connect the in-context learning ability to concrete representational capacity results. Rather, they seek to _explain_ the behavior of in-context learning as implicit Bayesian inference.

Footnote 7: The natural connection of the transformer architecture to both bounded hierarchical languages as well as to   _n_-gram LMs is interesting since those two classes of LMs can both be represented particularly _efficiently_ by _recurrent_ neural LMs (Svete et al., 2024).

## 7 Conclusion

We study the representational capacity of transformer LMs with   _n_-gram LMs. We show how the parallelizable nature of   _n_-gram LMs is easy to capture with the transformer architecture and provide multiple lower bounds on the probabilistic representational capacity of transformer LMs. Concretely, we show that transformer LMs can represent   _n_-gram LMs both with hard and sparse attention, exhibiting multiple mechanisms transformer LMs can employ to simulate   _n_-gram LMs. Altogether, our results reinforce the utility of non-sequential models of computation for the study of transformers, 

particularly in the language modeling setting.

### Limitations

We connect transformer LMs to _n_-gram LMs because of their parallelizable nature and their traditional popularity in NLP. However, _n_-gram LMs describe a very simple class of LMs, meaning that the lower bounds are somewhat less relevant than the characterization in terms of more expressive formal models of computation would be. Accordingly, we expect that the lower bounds are somewhat loose and that transformer LMs can represent more than _n_-gram LMs, which is also in line with the empirical success of transformer LMs. We leave it to future work to tighten the established lower bounds.

As with most theoretical investigations of transformers, our results are strongest and the most precise in the hard attention setting. However, hard attention is not used in practice, which limits the applicability of the results. The constructions presented in this paper are also purely meant to showcase the existence of a mechanism that can be used to simulate _n_-gram LMs. They do not suggest that the same mechanisms will be learned by models used in practice. Indeed, the very sparse representations are not in line with the common dense contextual representations usually learned by trained models.

We also only focus on _lower bounds_ of the representational capacity. We do not consider any upper bounds and existing results for similar models to ours suggest that the lower bound is indeed somewhat loose (Yao et al., 2021).8 That is, we expect that transformer LMs can represent much more than _n_-gram LMs, and expect that many of the existing results on the computational power of such models can be extended to the probabilistic setting.

Footnote 8: For example, we cannot say that the lower bounds imply any limitations of hard attention transformer LMs.

While we present a comprehensive analysis of transformer LMs in the context of _n_-gram LMs, we do not consider various aspects of the relationship that could be interesting. This is done to keep the presentation focused and concise. For example, we do not consider whether such simulations can be _learned_ from data, an interesting avenue for future research. Lastly, note that while we focus here specifically on the commonly deployed transformer-based _language models_, there are many other interesting applications of transformers, such as encoder-only acceptors of unweighted languages. These applications are better covered by existing work.

## Ethics Statement

The paper provides a way to theoretically analyze language models. To the best knowledge of the authors, there are no ethical implications of this paper.

### Acknowledgements

Ryan Cotterell acknowledges support from the Swiss National Science Foundation (SNSF) as part of the "The Forgotten Role of Inductive Bias in Interpretability" project. Anej Svate is supported by the ETH AI Center Doctoral Fellowship. We thank Tim Vieira and Shannon Veitch for helpful discussions and comments on the manuscript.

## References

* Akyurek et al. (2024) Ekin Akyurek, Bailin Wang, Yoon Kim, and Jacob Andreas. 2024. In-context language learning: Architectures and algorithms. _arXiv preprint arXiv:2401.12973_.
* Angluin et al. (2023) Dana Angluin, David Chiang, and Andy Yang. 2023. Masked hard-attention transformers and boolean rasp recognize exactly the star-free languages. _arXiv preprint arXiv:2310.13897_.
* Avcu et al. (2017) Enes Avcu, Chihiro Shibata, and Jeffrey Heinz. 2017. Subregular complexity and deep learning. _arXiv preprint arXiv:1705.05940_.
* Bahl et al. (1983) Lalit R. Bahl, Frederick Jelinek, and Robert L. Mercer. 1983. A maximum likelihood approach to continuous speech recognition. _IEEE Transactions on Pattern Analysis and Machine Intelligence_, PAMI-5(2):179-190.
* Baker (1975a) J. Baker. 1975a. The DRAGON system-An overview. _IEEE Transactions on Acoustics, Speech, and Signal Processing_, 23(1):24-29.
* Baker (1975b) James K. Baker. 1975b. _Stochastic Modeling as a Means of Automatic Speech Recognition_. Ph.D. thesis, Carnegie Mellon University, USA. AAI7519843.
* Bengio et al. (2000) Yoshua Bengio, Rejean Ducharme, and Pascal Vincent. 2000. A neural probabilistic language model. In _Advances in Neural Information Processing Systems_, volume 13. MIT Press.
* Bengio et al. (2003) Yoshua Bengio, Rejean Ducharme, Pascal Vincent, and Christian Janvin. 2003. A neural probabilistic language model. _J. Mach. Learn. Res._, 3:1137-1155.
* Bengio et al. (2006) Yoshua Bengio, Holger Schwenk, Jean-Sebastien Senecal, Frederic Morin, and Jean-Luc Gauvain. 2006. _Neural Probabilistic Language Models_, pages 137-186. Springer Berlin Heidelberg, Berlin, Heidelberg.

 

Satwik Bhattacharya, Kabir Ahuja, and Navin Goyal. 2020. On the ability and limitations of transformers to recognize formal languages. In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_, pages 7096-7116, Online. Association for Computational Linguistics.
* Chiang and Cholak [2022] David Chiang and Peter Cholak. 2022. Overcoming a theoretical limitation of self-attention. In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_, pages 7654-7664, Dublin, Ireland. Association for Computational Linguistics.
* Chiang et al. [2023] David Chiang, Peter Cholak, and Anand Pillay. 2023. Tighter bounds on the expressivity of transformer encoders. In _Proceedings of the 40th International Conference on Machine Learning_, ICML'23. JMLR.org.
* Du et al. [2023] Li Du, Lucas Torroba Hennigen, Tiago Pimentel, Clara Meister, Jason Eisner, and Ryan Cotterell. 2023. A measure-theoretic characterization of tight language models. In _Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_, pages 9744-9770, Toronto, Canada. Association for Computational Linguistics.
* Elhage et al. [2021] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova Das Sarma, Dawn Drain, Deep Ganguli, Zacq Hathfeld-Dodds, Danry Hernandez, Andy Jones, Jackson Kerionion, Llane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. 2021. A mathematical framework for transformer circuits. _Transformer Circuits Thread_.
* Hahn [2020] Michael Hahn. 2020. Theoretical limitations of self-attention in neural sequence models. _Transactions of the Association for Computational Linguistics_, 8:156-171.
* Hao et al. [2022a] Yiding Hao, Dana Angluin, and Robert Frank. 2022a. Formal language recognition by hard attention transformers: Perspectives from circuit complexity. _Transactions of the Association for Computational Linguistics_, 10:800-810.
* Hao et al. [2022b] Yiding Hao, Dana Angluin, and Robert Frank. 2022b. Formal language recognition by hard attention transformers: Perspectives from circuit complexity. _Transactions of the Association for Computational Linguistics_, 10:800-810.
* Hao et al. [2018] Yiding Hao, William Merrill, Dana Angluin, Robert Frank, Noah Amsel, Andrew Benz, and Simon Mendelsohn. 2018. Context-free transductions with neural stacks. In _Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_, pages 306-315, Brussels, Belgium. Association for Computational Linguistics.
* Heafield [2011] Kenneth Heafield. 2011. KenLM: Faster and smaller language model queries. In _Proceedings of the Sixth Workshop on Statistical Machine Translation_, pages 187-197, Edinburgh, Scotland. Association for Computational Linguistics.
* Heafield et al. [2013] Kenneth Heafield, Ivan Pozyrevsky, Jonathan H. Clark, and Philipp Koehn. 2013. Scalable modified Kneser-Ney language model estimation. In _Proceedings of the 51st Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers)_, pages 690-696, Sofia, Bulgaria. Association for Computational Linguistics.
* Heinz and Rogers [2013] Jeffrey Heinz and James Rogers. 2013. Learning subregular classes of languages with factored deterministic automata. In _Proceedings of the 13th Meeting on the Mathematics of Language (MoL 13)_, pages 64-71, Sofia, Bulgaria. Association for Computational Linguistics.
* Hewitt et al. [2020] John Hewitt, Michael Hahn, Surya Ganguli, Percy Liang, and Christopher D. Manning. 2020. RNNs can generate bounded hierarchical languages with optimal memory. In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_, pages 1978-2010, Online. Association for Computational Linguistics.
* Jager and Rogers [2012] Gerhard Jager and James Rogers. 2012. Formal language theory: Refining the Chomsky hierarchy. _Philos Trans R Soc Lond B Biol Sci_, 367(1598):1956-1970.
* Jelinek [1976] F. Jelinek. 1976. Continuous speech recognition by statistical methods. _Proceedings of the IEEE_, 64(4):532-556.
* Jelinek [1990] F. Jelinek. 1990. _Self-Organized Language Modeling for Speech Recognition_, page 450-506. Morgan Kaufmann Publishers Inc., San Francisco, CA, USA.
* Kleene [1956] S. C. Kleene. 1956. Representation of events in nerve nets and finite automata. In C. E. Shannon and J. McCarthy, editors, _Automata Studies. (AM-34), Volume 34_, pages 3-42. Princeton University Press, Princeton.
* Liu et al. [2023] Bingbin Liu, Jordan T. Ash, Subhi Goel, Akshay Krishnamurthy, and Cyril Zhang. 2023. Transformers learn shortcuts to automata. In _International Conference on Learning Representations_.
* Maass et al. [1991] W. Maass, G. Schnitger, and E. D. Sontag. 1991. On the computational power of sigmoid versus boolean threshold circuits. In _Proceedings 32nd Annual Symposium of Foundations of Computer Science_, pages 767-776.
* Volume 48_, ICML'16, page 1614-1623.
* Merrill [2019] William Merrill. 2019. Sequential neural networks as automata. In _Proceedings of the Workshop on Deep Learning and Formal Languages: Building Bridges_, 

pages 1-13, Florence. Association for Computational Linguistics.
* Merrill and Sabharwal (2023a) William Merrill and Ashish Sabharwal. 2023a. The expressive power of transformers with chain of thought. arXiv preprint arXiv:2310.07923. Cited by: SS4.
* W. Merrill and A. Sabharwal (2023b)A logic for expressing log-precision transformers. arXiv preprint arXiv:2210.02671. Cited by: SS5.
* W. Merrill and A. Sabharwal (2023c)The parallelism tradeoff: limitations of log-precision transformers. Transactions of the Association for Computational Linguistics11, pp. 531-545. External Links: Document, Link Cited by: SS5.
* W. Merrill, A. Sabharwal, and N. A. Smith (2022)Saturated transformers are constant-depth threshold circuits. Transactions of the Association for Computational Linguistics10, pp. 843-856. Cited by: SS5.
* W. Merrill and N. Tsilivis (2022)Extracting finite automata from RNNs using state merging. arXiv preprint arXiv:2201.12451. Cited by: SS5.
* W. Merrill, G. Weiss, Y. Goldberg, R. Schwartz, D. A. Smith, and E. Yahav (2020)A formal hierarchy of RNN architectures. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, Online, pp. 443-459. External Links: Document, Link Cited by: SS5.
* M. Lee Minsky (1954)Neural nets and the brain model problem. Ph.D. Thesis, Princeton University. Cited by: SS5.
* F. Nowak, A. Svate, A. Butoi, and R. Cotterell (2024)On the representational capacity of neural language models with chain-of-thought reasoning. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Thailand, Association for Computational Linguistics, Cited by: SS5.
* F. Nowak, A. Svate, L. Du, and R. Cotterell (2023)On the representational capacity of recurrent neural language models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, Singapore, Singapore, pp. 7011-7034. External Links: Document, Link Cited by: SS5.
* C. Olsson, N. Elhage, N. Nanda, N. Joseph, N. Das Sarma, T. Henighan, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conperly, D. Drain, Deep Ganguli, Z. Hatfield-Dodds, Danny Hernandez, S. Johnston, A. Jones, J. Kernion, L. Lovitt, K. Ndousse, D. Amodei, T. Brown, J. Clark, J. Kaplan, S. McCandlish, and C. Olah (2022)In-context learning and induction heads. Transformer Circuits Thread. Cited by: SS5.
* J. Perez, P. Barcelo, and J. Marinkovic (2021)Attention is Turing-complete.. Journal of Machine Learning Research22 (75), pp. 1-35. Cited by: SS5.
* H. Schwenk (2007)Continuous space language models. Computer Speech & Language21 (3), pp. 492-518. Cited by: SS5.
* C. E. Shannon (1948)A mathematical theory of communication. The Bell System Technical Journal27 (3), pp. 379-423. Cited by: SS5.
* H. T. Siegelmann and E. D. Sontag (1992)On the computational power of neural nets. In Proceedings of the Fifth Annual Workshop on Computational Learning Theory, COLT '92, New York, NY, USA, pp. 440-449. External Links: ISBN 978-1-4503-4 Cited by: SS5.
* P. Smolensky, M. C. Mozer, and D. E. Rumelhart (1996)Mathematical perspectives on neural networks. Psychology Press. Cited by: SS5.
* L. Strobl, W. Merrill, G. Weiss, D. Chiang, and D. Angluin (2023)Transformers as recognizers of formal languages: a survey on expressivity. arXiv preprint arXiv:2311.00208. Cited by: SS5.
* A. Svate, R. Shing Moon Chan, and R. Cotterell (2024)A theoretical result on the inductive bias of RNN language models. arXiv preprint arXiv:2402.15814. Cited by: SS5.
* A. Svate and R. Cotterell (2023)Recurrent neural language models as probabilistic finite-state automata. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, Singapore, pp. 8069-8086. Cited by: SS5.
* A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin (2017)Attention is all you need. In Advances in Neural Information Processing Systems, Vol. 30, pp. . Cited by: SS5.
* G. Weiss, Y. Goldberg, and E. Yahav (2021)Thinking like transformers. In Proceedings of the 38th International Conference on Machine Learning, Proceedings of Machine Learning Research, pp. 11080-11090. Cited by: SS5.
* K. Wen, Y. Li, B. Liu, and A. Risteski (2023)Transformers are uninterpretable with myopic methods: a case study with bounded Dyck grammars. In Thirty-seventh Conference on Neural Information Processing Systems, Cited by: SS5.
* S. Michael Xie, A. Raghunathan, P. Liang, and T. Ma (2022)An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations, Cited by: SS5.
* Z. Yang, Z. Dai, R. Salakhutdinov, and W. W. Cohen (2018)Breaking the softmax bottleneck: a high-rank rnn language model. In International Conference on Learning Representations 

Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. 2021. Self-attention networks can process bounded hierarchical languages. In _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_, pages 3770-3785, Online. Association for Computational Linguistics.

 

Modelling Assumptions

As encouraged by Strobl et al. (2023), we provide in Tab. 2 a short summary of the assumptions behind the specific transformer architecture we consider in this work. This aims to facilitate the placement of the results in the context of existing work and to make the results more accessible to the reader.

## Appendix B Proofs: Hard Attention

This section provides detailed proofs of all theorems about the representational capacity of hard-attention transformer LMs stated in the main part of the paper.

### Computing Logical AND with an MLP

**Definition B.1**.: _A \({\rm ReLU}\)-activated **multi-layer-perceptron** (MLP) \({\rm MLP}\colon\mathbb{R}^{N}\to\mathbb{R}^{M}\) is a function defined as the composition of functions \(\mathbf{f}_{1},\cdots,\mathbf{f}_{L}\)_

\[{\rm MLP}\left(\mathbf{x}\right)\stackrel{{\text{\tiny def}}}{{=}} \left(\mathbf{f}_{L}\circ\mathbf{f}_{L-1}\circ\cdots\circ\mathbf{f}_{1}\right)\left( \mathbf{x}\right) \tag{20}\]

_where each \(\mathbf{f}_{\ell}\) for \(\ell\in[L]\) is defined as_

\[\mathbf{f}_{\ell}\left(\mathbf{x}\right) \stackrel{{\text{\tiny def}}}{{=}}{\rm ReLU}\left( \mathbf{W}_{\ell}\mathbf{x}+\mathbf{b}_{\ell}\right)\quad\ell\in[L-1] \tag{21a}\] \[\mathbf{f}_{L}\left(\mathbf{x}\right) \stackrel{{\text{\tiny def}}}{{=}}\mathbf{W}_{L} \mathbf{x}+\mathbf{b}_{L} \tag{21b}\]

_where \(\mathbf{W}_{\ell}\in\mathbb{R}^{N_{\ell}\times M_{\ell}}\) is a weight matrix with dimensions \(N_{\ell}\) and \(M_{\ell}\) specific to layer \(\ell\), \(\mathbf{b}_{\ell}\in\mathbb{R}^{M_{\ell}}\) is a bias vector. We refer to MLPs by the number of hidden layers, e.g., a one-layer-MLP is an MLP with one hidden layer._

In our construction, simulating a \(n\)-gram LM with a transformer LM requires a component of the transformer to perform the logical AND operation between specific entries of binary vectors \(\mathbf{x}\in\mathbb{B}^{D}\). The following lemma shows how this can be performed by an MLP with appropriately set parameters.

**Lemma B.1**.: _Consider \(m\) indices \(i_{1},\ldots,i_{m}\in[D]\) and vectors \(\mathbf{x},\mathbf{v}\in\mathbb{B}^{D}\) such that_

\[v_{i}=\mathbbm{1}\left\{i\in\left\{i_{1},\ldots,i_{m}\right\}\right\}, \tag{22}\]

_i.e., with entries \(1\) at indices \(i_{1},\ldots,i_{m}\). Then, it holds for the MLP \({\rm MLP}\left(\mathbf{x}\right)\stackrel{{\text{\tiny def}}}{{=}}{ \rm ReLU}\left(\mathbf{v}^{\top}\mathbf{x}-(m-1)\right)\) that_

\[{\rm MLP}\left(\mathbf{x}\right)=1\mbox{ if and only if $x_{i_{k}}=1$ for all $ k=1,\ldots,m$.} \tag{23}\]

_In other words,_

\[{\rm MLP}\left(\mathbf{x}\right)=x_{i_{1}}\wedge\cdots\wedge x_{i_{m}}. \tag{24}\]

Proof.: By the definition of \(\mathbf{v}\), \(\mathbf{v}^{\top}\mathbf{x}\leqslant m\) for all \(\mathbf{x}\in\mathbb{B}^{D}\). Furthermore, \(\mathbf{v}^{\top}\mathbf{x}=m\) if and only if \(x_{i_{k}}=1\) for all \(k=1,\ldots,m\). The \({\rm ReLU}\) function maps all other values to \(0\) and does not change the output value \(1\) in this case.

 



### Proofs of the Hard Attention Case

This subsection contains all proofs of the representational capacity of hard attention transformer LMs. We tackle three cases: the simulation with \(n-1\) heads and a single layer, with \(n-1\) layers and a single head, and lastly, the simulation with a single head and a single layer. All sections first outline the intuition behind the proofs and then provide the details in the form of finer-grained lemmata.

#### b.2.1 Simulation with _n_\(-1\) Heads: The Intuition

We now outline the intuition behind the construction of a hard attention transformer LM simulating an _n_-gram LM, as first presented in Fig. 1.9 To ease the exposition, we start with the final step of the construction: Assuming we have identified the appropriate history \({\mathbf{y}}^{t-1}_{t-n+1}\) after combining the head values using the head combining function \(\mathcal{H}\), we show how \(p_{\mathcal{T}}\) can encode the conditional probability distribution \(p\left(y_{t}~{}|~{}{\mathbf{y}}_{t-n+1:t-1}\right)\). The intuition of this step is simple: Knowing what the individual \(p\left(y_{t}~{}|~{}{\mathbf{y}}^{t-1}_{t-n+1}\right)\) for \(y_{t}\in\overline{\Sigma}\) are, we can simply put their logits into a vector and combine the constructed vectors for all possible histories into the output matrix \(\mathbf{E}\):10,11

Footnote 9: For simplicity, we disregard the role of residual connections in the following outline. Residual connections are, however, considered in the full proof later.

Footnote 10: To be able to take the log of \(0\) probabilities, we work over the set of _extended_ reals \(\overline{\mathbb{R}}=\mathbb{R}\cup\{-\infty,\infty\}\).

Footnote 11: Throughout the paper, we implicitly index the matrices directly with symbols and histories. We assume that the symbols and histories are ordered in some way and that the matrices are ordered accordingly.

\[E_{y,{\mathbf{y}}^{t-1}_{t-n+1}}\overset{\text{def}}{=}\log p\left( y_{t}~{}|~{}{\mathbf{y}}^{t-1}_{t-n+1}\right) \tag{25}\]

In the following, we write \(\mathsf{enc}\left({\mathbf{y}}_{<t}\right)\) as shorthand notation for \(\mathsf{enc}\left({\mathbf{y}}_{<t}\right)\overset{\text{def}}{=}F\left({\bf x}^{L }_{t-1}\right)\) (i.e., the representation which is linearly transformed by \(\mathbf{E}\) to compute \(p\left(y_{t}~{}|~{}{\mathbf{y}}_{<t}\right)\) after normalization) where \(\mathbf{X}^{L}=\mathcal{T}\left(\mathcal{R}\right)\left({\mathbf{y}}_{<t}\right)\). If we one-hot encode the identified history with \(\mathcal{T}\) as

\[\mathsf{enc}\left({\mathbf{y}}_{<t}\right)\overset{\text{def}}{=}[{\mathbf{y}}^{t-1}_ {t-n+1}] \tag{26}\]

we can then, using the formulation of the transformer LM from Definition 2.11, use the \(\mathsf{enc}\left({\mathbf{y}}_{<t}\right)\) to look up the appropriate column in \(\mathbf{E}\) containing the logits of the conditional probabilities given the identified history for all possible \(y_{t}\in\overline{\Sigma}\), i.e., \(\left(\mathbf{E}~{}\mathsf{enc}\left({\mathbf{y}}_{<t}\right)\right)_{y}=\log p \left(y~{}|~{}{\mathbf{y}}^{t-1}_{t-n+1}\right)\).

We now consider the preceding step of the simulation: Identifying the history given that the \(n-1\) heads identified the symbols \(y_{1},\ldots,y_{n-1}\) in the positions they attended to. If we concatenate the values of the \(n-1\) heads into a vector \(\mathbf{v}\), this vector of size \(\left(n-1\right)|\overline{\Sigma}|\) will contain the **multi-hot** representation of the history of interest:

\[\mathbf{v}=\begin{pmatrix}[\![y_{1}]\\ \vdots\\ [\![y_{n-1}]\!]\end{pmatrix} \tag{27}\]

and \(\mathbf{v}_{i|\overline{\Sigma}|+j}=1\) if and only if \(m\left(y_{i}\right)=j\) for a bijection \(m\colon\overline{\Sigma}\to\big{[}|\overline{\Sigma}|\big{]}\) that determines the indices of the one-hot representations of the symbols. We would then like to transform this vector into a vector \(\mathbf{u}\in\mathbb{R}^{|\Sigma|^{n-1}}\) such that

\[u_{i}=1\iff i=s\left(y_{1},\ldots,y_{n-1}\right) \tag{28}\]

for a bijection \(s\colon\underbrace{\overline{\Sigma}\times\cdots\times\overline{\Sigma}}_{ n-1\text{ times}}\to\big{[}|\overline{\Sigma}|^{n-1}\big{]}\). This can be equivalently written as

\[u_{i}=1\iff v_{j|\overline{\Sigma}|+m\left(y_{j}\right)}=1\text{ for all $ j=1,\ldots,n-1$} \tag{29}\]

where \(i=s\left(y_{1},\ldots,y_{n-1}\right)\). This is an instance of performing the logical AND operation, which can be implemented by an MLP as described in Lemma B.1. This MLP will form the transformation \(\mathcal{H}\) combining the information obtained from all the heads of the transformer.

This brings us to the final part of the proof: Identifying the symbols at the previous _n_\(-1\) positions by the _n_\(-1\) transformer heads. To show how this can be done, let us consider the parameters we can still set to define a transformer

* The position-augmented symbol representations \({\bf r}\). Inspired by concurrent work from Merrill and Sabharwal (2023a), we use the representations of the form \[{\bf r}\left(y_{t},t\right)=\begin{pmatrix}\llbracket y_{t}\rrbracket\\ \sqrt{\frac{1}{t}}\\ \sqrt{1-\frac{1}{t}}\\ \sqrt{\frac{1}{t+1}}\\ \sqrt{1-\frac{1}{t+1}}\\ \vdots\\ \sqrt{\frac{1}{t+n-1}}\\ \sqrt{1-\frac{1}{t+n-1}}\end{pmatrix}\in\mathbb{R}^{2n}\] (30) This results in vectors \({\bf r}\left(y_{t},t\right)\) of size \(|\overline{\Sigma}|+2+2\left(n-1\right)\). Note that such a symbol representation function can be implemented by concatenating or adding symbol- (\(\llbracket y_{t}\rrbracket\)) and position-(\(\sqrt{\frac{1}{t}},\ldots,\sqrt{1-\frac{1}{t+n-1}}\)) specific components, which is in line with most practical implementations of the transformer architecture.
* The attention scoring function \(f\). We will use the standard dot-product scoring function \[f\left({\bf q},{\bf k}\right)\stackrel{{\text{def}}}{{=}} \langle{\bf q},{\bf k}\rangle.\] (31) \(f\) will, together with the positional encodings, allow us to easily single out the relevant positions in the string.
* The parameters of each of the attention heads, that is, the transformations \(Q\), \(K\), and \(V\). Each of those will take the form of a linear transformation of the symbol (and positional) representations. We describe them and their roles in more detail below.

The parameters of all the heads will be identical, with the only difference being a single parameter that depends on the "index" of the head, \(h\). In the following, we describe the construction of a single head. At any time step \(t\) (i.e., when modeling the conditional distribution \(p\left(y_{t}~{}|~{}\mathbf{y}_{<t}\right)\)), the head \(h\) will attend to the symbol at position \(t-h\), \(y_{t-h}\). In Fig. 1, for example, Head \(3\) attends to the position \(t-3\), which is denoted by the stronger arrow to that position. We now describe the individual transformations \(Q_{h}\), \(K_{h}\), \(V_{h}\), and \(O_{h}\) of the head \(h\). All of them will be _affine_ transformations. Since we are considering only the first layer of the transformer, we can think of the inputs to the layer as the original symbol representations together with their position encodings (rather than some contextual representations at higher levels). As mentioned, the head \(h\) will be responsible for identifying the symbol at position \(t-h\). Therefore, we want it to put all its attention to this position. In other words, given the query \({\bf q}_{t-1}\), we want the attention function in Eq. (31) to be uniquely maximized by the key of the symbol at position \(t-h\). Notice that, therefore, the key does not have to depend on the identity of the symbol at position \(t-h\)--only the positional information matters. Let us then consider the following query and key transformations for head \(h\):

\[Q_{h}\colon{\bf r}\left(y_{t},t\right) \mapsto\begin{pmatrix}\sqrt{\frac{1}{t}}\\ \sqrt{1-\frac{1}{t}}\end{pmatrix} \tag{32}\] \[K_{h}\colon{\bf r}\left(y_{t},t\right) \mapsto\begin{pmatrix}\sqrt{\frac{1}{t+k}}\\ \sqrt{1-\frac{1}{t+h}}\end{pmatrix}. \tag{33}\]

Given such a query and such keys, the scoring function computes

\[f\left({\bf q}_{t},{\bf k}_{j}\right)=\bigg{\langle}\begin{pmatrix}\sqrt{\frac {1}{t}}\\ \sqrt{1-\frac{1}{t}}\end{pmatrix},\begin{pmatrix}\sqrt{\frac{1}{j+h}}\\ \sqrt{1-\frac{1}{j+h}}\end{pmatrix}\bigg{\rangle}.\] (34) 

Eq. (34) is an inner product between two unit vectors, and is therefore maximized if and only if they are the same, that is, if \(j=t-h\). This is exactly the position that we want the head \(h\) to attend to.12 Intuitively, both transformations keep only the positional information. The query transformation "injects" the knowledge of which position should maximize the attention score, while the key transformation simply "exposes" the positional information about the symbol. The constants \(1\) and \(-1\) and the index of the position ensure that the inner product simply computes the difference between the position of the symbol and the position of interest.

Footnote 12: Note that while the choice of the positional encodings in this construction is uncommon in practice, the popular sinusoidal positional encodings [57] have also been linked to the ability of the transformer-based models to attend to specific positions of interest based on linear transformations of the positional encodings [68].

This leaves us with the question of how to use the position of the symbol of interest \((t-h)\) to extract the one-hot encoding of the symbol at that position. Due to the information contained in the symbol representations \({\bf r}\left(y_{j}\right)\), this is trivial:

\[V\colon{\bf r}\left(y_{t},t\right)\mapsto\llbracket y_{j}\rrbracket. \tag{35}\]

With this, the identity of the symbol is carried forward through the attention mechanism. Notice that the only head-dependal transformation is the query transformation--it depends on the index of the head, determining the position of interest, meaning that every head defines a different query transformation, while the keys and values transformations are the same among all heads. This concludes the outline of the proof.

#### b.2.2 Simulation with _n_\(-1\)**Heads: Proofs**

This subsection formally proves the construction intuited in Appendix B.2.1 by proving a sequence of lemmata that formalize each of the steps described in the intuition. Specifically,

1. Lemma B.2 shows how the one-hot encodings of individual symbols in the history can be combined into the one-hot encoding of the history.
2. Lemma B.3 shows that the scoring function is maximized at the position of interest.
3. Lemma B.4 shows how the one-hot encodings of the symbols in the history can be identified by the hard attention mechanism.
4. The proof of Theorem 3.1 shows how the construction of the one-hot encoding of the current history allows us to define the appropriate next-symbol conditional distribution of the _n_-gram LMs.

**Lemma B.2**.: _Let \(\mathcal{T}\) be a transformer with \(H=n-1\) heads. Let \({\bf z}_{h}=\llbracket y_{t-h}\rrbracket\) be the output of the \(h^{th}\) head at time \(t\). Then, there exists a function \(\mathcal{H}\) implemented by a single-layer MLP such that_

\[\mathcal{H}\left({\bf z}_{1},\ldots,{\bf z}_{n-1}\right)=\llbracket{\mathbf{y}}_{ t-n+1}^{t-1}\rrbracket. \tag{36}\]

Proof.: Let \({\mathbf{y}}=y_{1}\ldots y_{n-1}\in\underline{\Sigma}^{n-1}\) and let \(i\) be the index in \(\big{[}|\Sigma|^{n-1}\big{]}\) that corresponds to \({\mathbf{y}}\). Furthermore, let \(i_{1},\dots,i_{n-1}\) be the indices corresponding to the symbols \(y_{1},\dots,y_{n-1}\) in \({\bf z}_{1},\dots,{\bf z}_{n-1}\). Then, we have

\[\llbracket{\mathbf{y}}_{t-n+1}^{t-1}\rrbracket_{i}=1~{}\Longleftrightarrow~{}z_ {1,i_{1}}=1~{}\wedge~{}\cdots\wedge z_{n-1,i_{n-1}}=1 \tag{37}\]

Eq. (36) is an instance of the logical AND operation on the indices encoding the individual histories, which can be implemented by an MLP as shown in Lemma B.1. 

The next lemma presents a useful equality about the standard dot-product attention scoring function: For unit vectors, the attention score is maximized if and only if the vectors are identical. We will use this fact in our construction to attend to particular positions in the string.

**Lemma B.3**.: _Given a fixed \(t\in\mathbb{N}\), \(f\), define_

\[g\left(j\right)\stackrel{{\text{\tiny def}}}{{=}}\left<\begin{pmatrix} \sqrt{\frac{1}{t-1}}\\ \sqrt{1-\frac{1}{t-1}}\end{pmatrix},\begin{pmatrix}\sqrt{\frac{1}{j+h}}\\ \sqrt{1-\frac{1}{j+h}}\end{pmatrix}\right>\] (38) 

_for \(j\in[t-1]\). Then, \(g\) is maximized at \(j=t-1-h\):_

\[\operatorname*{argmax}_{j\in[t-1]}\left(\bigg{<}\begin{pmatrix}\sqrt{\frac{1}{ t-1}}\\ \sqrt{1-\frac{1}{t-1}}\end{pmatrix},\begin{pmatrix}\sqrt{\frac{1}{j+h}}\\ \sqrt{1-\frac{1}{j+h}}\end{pmatrix}\right>\right)=t-1-h. \tag{39}\]

Proof.: The two arguments to the inner product in Eq. (39) are unit vectors. Inner products of unit vectors are at most \(1\), with the maximum achieved only if the two vectors are identical. This means that the function in Eq. (39) is maximized if and only if

\[\begin{pmatrix}\sqrt{\frac{1}{t-1}}\\ \sqrt{1-\frac{1}{t-1}}\end{pmatrix} =\begin{pmatrix}\sqrt{\frac{1}{j+h}}\\ \sqrt{1-\frac{1}{j+h}}\end{pmatrix}\Longleftrightarrow \tag{40a}\] \[\sqrt{\frac{1}{t-1}} =\sqrt{\frac{1}{j+h}}\iff\] (40b) \[\frac{1}{t-1} =\frac{1}{j+h}\iff\] (40c) \[j =t-1-h. \tag{40d}\]

The following lemma presents the core of the proof of Theorem 3.1, exhibiting the construction of a transformer head that can single out and one-hot encode a symbol at a specific position in the input string. The lemma relies on a simple pre-processing of the input string, where the string is prepended (padded) with \(n-1\) beginning of string symbols \(y_{1}=\ldots=y_{n-1}\stackrel{{\text{def}}}{{=}}\textsc{bos}\), which is common practice in language modeling literature, especially when talking about \(n\)-gram LMs. We will denote \(\underline{\Sigma}\stackrel{{\text{def}}}{{=}}\Sigma\cup\{ \textsc{bos}\}\). This enables a cleaner presentation of the concrete construction of the attention mechanism.

**The idea of the proof.** Lemma B.4 contains a number of technical definitions of the parameters of a transformer layer (cf. Definition 2.6. Together, they describe a single head of a transformer layer (which will contain \(H=n-1\) such heads) that is able to extract the one-hot encoding of a particular symbol in the history. The layer of \(n-1\) heads will then be able to extract the \(n-1\) symbols, as required by Lemma B.2. We now describe a single head more formally. Define the following position-augmented symbol representation function of the transformer head \(h\):

\[\mathbf{r}\left(y,t\right)=\begin{pmatrix}&\llbracket y\rrbracket\\ &\boldsymbol{0}_{\lfloor\underline{\Sigma}\rceil}\\ &\sqrt{\frac{1}{t}}\\ &\sqrt{1-\frac{1}{t}}\\ &\sqrt{\frac{1}{t+1}}\\ &\sqrt{1-\frac{1}{t+1}}\\ &\vdots\\ &\sqrt{\frac{1}{t+n-1}}\\ &\sqrt{1-\frac{1}{t+n-1}}\\ \end{pmatrix}\in\mathbb{R}^{2|\underline{\Sigma}|+2n}. \tag{41}\]

Here, \(\llbracket\cdot\rrbracket\in\{0,1\}^{|\underline{\Sigma}|}\) one-hot encodes symbols over \(\underline{\Sigma}\). This means that the entire static representations contain multiple components:

* two \(|\underline{\Sigma}|\)-dimensional slots for _symbol_ representations and
* \(n\) 2-dimensional slots for head-specific _positional_ representations 

We further define

\[f\left({\bf q},{\bf k}\right) \stackrel{{\text{\tiny def}}}{{=}}\langle{\bf q},{\bf k}\rangle, \tag{42}\]

\[Q\left({\bf x}\right) \stackrel{{\text{\tiny def}}}{{=}}{\bf Qx}, {\bf Q} \in\mathbb{R}^{2\times(2|\Sigma|+2n)}, \tag{43a}\] \[K\left({\bf x}\right) \stackrel{{\text{\tiny def}}}{{=}}{\bf Kx}, {\bf K} \in\mathbb{R}^{2\times(2|\Sigma|+2n)},\] (43b) \[V\left({\bf x}\right) \stackrel{{\text{\tiny def}}}{{=}}{\bf Vx}, {\bf V} \in\mathbb{R}^{(2|\Sigma||+2n)\times(2|\Sigma|+2n)},\] (43c) \[O\left({\bf x}\right) \stackrel{{\text{\tiny def}}}{{=}}{\bf Ox}, {\bf O} \in\mathbb{R}^{(2|\Sigma|+2n)\times(2|\Sigma|+2n)}, \tag{43d}\]

\[{\bf Q}_{:,2|\underline{\Sigma}|+1:2|\underline{\Sigma}|+2} \stackrel{{\text{\tiny def}}}{{=}}{\bf I}_{2} \tag{44a}\] \[{\bf K}_{:,2|\underline{\Sigma}|+2h+1:2|\underline{\Sigma}|+2h+2} \stackrel{{\text{\tiny def}}}{{=}}{\bf I}_{2},\] (44b) \[{\bf V}_{|\underline{\Sigma}|+1:2|\underline{\Sigma}|;1:| \underline{\Sigma}|} \stackrel{{\text{\tiny def}}}{{=}}{\bf I}_{|\underline{ \Sigma}|}\] (44c) \[{\bf O}_{1:\underline{\Sigma}|,1:\underline{\Sigma}|} \stackrel{{\text{\tiny def}}}{{=}}-{\bf I}_{|\underline{ \Sigma}|} \tag{44d}\]

We can visualize these matrices as

\[{\bf Q} = \tag{45a}\] \[{\bf K} =\] (45b) \[{\bf V} =\] (45c) \[{\bf O} = \tag{45d}\]

where \({\bf 0}_{N}\) is a \(N\)-dimensional vector of zeros, \({\bf I}_{N}\) is the \(N\)-dimensional identity matrix and the unspecified elements of \({\bf Q},{\bf K},{\bf V},{\bf O}\) are \(0\).

**Lemma B.4**.: _Let \(\Sigma\) be an alphabet and \({\bm y}\in\Sigma^{*}\). For any \(t=1,\ldots,|{\bm y}|\), the \(h^{th}\) transformer head (\(h\in[n-1]\)) defined with the parameters specified in Eq. (41) to Eq. (44d) outputs_

\[{\bf z}_{t-1}=\begin{pmatrix}{\bf 0}_{|\Sigma|}\\ \llbracket y_{t-1-h}\rrbracket\\ {\bf 0}_{2n}\end{pmatrix}. \tag{46}\]

_In particular, this means that the output \({\bf z}_{t-1}\) at time step \(t-1\) contains the one-hot encoding of the symbol at position \(t-h-1\).13_ 

Proof.: Fix \(t\in[|{\mathbf{y}}|]\). We compute the representation of \({\mathbf{y}}_{<t}\) computed by the head. First, observe that the matrix defining the query function \(Q\) projects onto the first two components of the positional encoding from the representation of \(y_{t-1}\):

\[Q\left({\bf r}\left(y_{t-1}\right)\right)={\bf Q}~{}{\bf r}\left(y_{t-1}\right)= \begin{pmatrix}\sqrt{\frac{1}{t-1}}\\ \sqrt{1-\frac{1}{t-1}}\end{pmatrix}. \tag{47}\]

Similarly, the matrix defining the key transformation projects onto the \(h^{\text{th}}\) positional encoding slot, i.e., the dimensions \(2|\underline{\Sigma}|+2h+1\) and \(2|\underline{\Sigma}|+2h+2\):

\[K\left({\bf r}\left(y_{j}\right)\right)={\bf K}~{}{\bf r}\left(y_{j}\right)= \begin{pmatrix}\sqrt{\frac{1}{j+h}}\\ \sqrt{1-\frac{1}{j+h}}\end{pmatrix} \tag{48}\]

for \(j=-n+2,\ldots,t-1\).

This results in the scoring function

\[f\left({\bf q}_{t-1},{\bf k}_{j}\right)\stackrel{{\mbox{\tiny def }}}{{=}}\langle{\bf q}_{t-1},{\bf k}_{j}\rangle=\bigg{\langle}\begin{pmatrix} \sqrt{\frac{1}{t-1}}\\ \sqrt{1-\frac{1}{t-1}}\end{pmatrix}\!,\begin{pmatrix}\sqrt{\frac{1}{j+h}}\\ \sqrt{1-\frac{1}{j+h}}\end{pmatrix}\!\!\bigg{\rangle}. \tag{49}\]

In particular, as shown in Lemma B.3, \(f\) is maximized for

\[j=t-1-h. \tag{50}\]

This means that

\[\operatorname*{hardmax}\left(f\left({\bf q}_{t-1},{\bf k}_{1}\right),f\left({ \bf q}_{t-1},{\bf k}_{2}\right),\ldots,f\left({\bf q}_{t-1},{\bf k}_{t-1} \right)\right)_{j}=\mathbbm{1}\left\{j=t-1-h\right\}. \tag{51}\]

The definition of \({\bf V}\) further means that

\[V\left({\bf r}\left(y_{j}\right)\right)={\bf V}~{}{\bf r}\left(y_{j}\right)= \begin{pmatrix}{\bf 0}_{|\underline{\Sigma}|}\\ {\llbracket y_{j}\rrbracket}\\ {\bf 0}_{2n}\end{pmatrix}\!, \tag{52}\]

giving us, by Eq. (9a),

\[{\bf a}_{t-1}={\bf v}_{t-1-h}+{\bf x}_{t-1}=\begin{pmatrix}\llbracket y_{t-1} \rrbracket\\ \llbracket y_{t-1-h}\rrbracket\\ {\bf 0}_{2n}\end{pmatrix}\!. \tag{53}\]

The definition of \(O\) then gives us

\[{\bf z}_{t-1} = O\left({\bf a}_{t-1}\right)+{\bf a}_{t-1}\] (54a) \[= {\bf OA}_{t-1}+{\bf a}_{t-1}\] (54b) \[= \begin{pmatrix}-{\bf I}_{|\underline{\Sigma}|}&\\ \phantom{-}\hskip-6.0pt\begin{array}{c}\\ \phantom{-}\hskip-6.

 

Lemmas B.2 and B.4 show that we can define a transformer that one-hot encodes the history of interest \({\mathbf{y}}_{t-n+1}^{t-1}\). We now show how to define an output matrix **E** to define a weakly equivalent transformer LM. Concretely, we define \(\textbf{E}\in\mathbb{R}^{|\overline{\Sigma}|\times|\overline{z}|^{m-1}}\) with

\[E_{y,{\mathbf{y}}_{t-n+1}^{t-1}}\stackrel{{\text{def}}}{{=}}\log p \left(y_{t}~{}|~{}{\mathbf{y}}_{t-n+1}^{t-1}\right). \tag{55}\]

**Theorem 3.1**.: _For any n-gram LM, there exists a weakly equivalent single-layer hard attention transformer LM with n \(-1\) heads._

Proof.: Let \(\mathcal{T}\) be a transformer LM with \(n-1\) heads defined with the parameters specified in Eq. (41) to Eq. (44d) (\(h\in[\![n-1]\)). Let \({\mathbf{y}}\in\{\textsc{bos}\}^{n-1}\Sigma^{*}\) with \(|{\mathbf{y}}|=T\) be a string and \(\mathbf{X}^{L}=\left(\mathbf{x}_{-n+2}^{L\top};\ldots;\mathbf{x}_{T}^{L\top} \right)=\mathcal{T}\left({\mathbf{y}}\right)\). We derive:

\[p_{\mathcal{T}}\left({\mathbf{y}}\right) =p_{\mathcal{T}}\left(\textsc{eos}~{}|~{}{\mathbf{y}}\right)\prod_{ t=1}^{T}p_{\mathcal{T}}\left(y_{t}~{}|~{}{\mathbf{y}}_{<t}\right)\] (56a, Autoregressive LM.) \[=\operatorname{softmax}\bigl{(}\textbf{E}\ F\left(\mathbf{x}_{ T}^{L}\right)\bigr{)}_{\textsc{eos}}\prod_{t=1}^{T}\operatorname{softmax}\bigl{(}\textbf{ E}\ F\left(\mathbf{x}_{t-1}^{L}\right)\bigr{)}_{y_{t}}\] (56b, Definition 2.11.) \[=\operatorname{softmax}\bigl{(}\textbf{E}\left[{\mathbf{y}}_{T-n+2: T}\rrbracket\right)_{\textsc{eos}}\prod_{t=1}^{T}\operatorname{softmax}\bigl{(}\textbf{E} \left[{\mathbf{y}}_{t-n+1}^{t-1}\rrbracket\right)_{y_{t}}\] (56c, Lemma 2.2 ) \[=\frac{\exp\left(\textbf{E}\left[{\mathbf{y}}_{T-n+2:T}\rrbracket \right)_{\textsc{eos}}}{\sum_{y\in\overline{\Sigma}}\exp\left(\textbf{E} \left[{\mathbf{y}}_{T-n+2:T}\rrbracket\right)\right)}\prod_{t=1}^{T}\frac{\exp \left(\textbf{E}\left[{\mathbf{y}}_{t-n+1:t}^{t-1}\rrbracket\right)_{y_{t}}}{\sum _{y\in\overline{\Sigma}}\exp\left(\textbf{E}\left[{\mathbf{y}}_{T-n+1:t-1} \rrbracket\right)_{y}}\right.\] (56d, Definition of softmax.) \[=\frac{\exp\left(\log p\left(\textsc{cos}~{}|~{}{\mathbf{y}}_{t-n+2}^{T} \right)\right)}{\sum_{y\in\overline{\Sigma}}\exp\left(\log p\left(y~{}|~{}{ \mathbf{y}}_{t-n+2}^{T}\right)\right)}\prod_{t=1}^{T}\frac{\exp\left(\log p\left( y~{}|~{}{\mathbf{y}}_{t-n+1}^{t-1}\right)\right)}{\sum_{y\in\overline{\Sigma}}\exp \left(\log p\left(y~{}|~{}{\mathbf{y}}_{t-n+1}^{t-1}\right)\right)}\] (56e, Eq. (55).) \[=p\left(\textsc{eos}~{}|~{}{\mathbf{y}}_{T-n+2:T}\right)\prod_{t=1}^{T}p \left(y_{t}~{}|~{}{\mathbf{y}}_{t-n+1}^{t-1}\right)=p\left({\mathbf{y}}\right)\] (56f, The \[n\]-gram LM is autoregressive.)

#### b.2.3 Simulation with \(n-1\) Layers: The Intuition

This section presents the construction of a transformer LM with a _single_ head but \(n-1\) layers that is weakly equivalent to a given \(n\)-gram LM. We again outline the intuition first before giving the technical details below as part of Lemma B.5. The construction we present resembles the one from the proof of Theorem 3.1. The main difference is that, instead of using \(n-1\) attention heads to identify the history, we instead use \(n-1\) transformer _layers_. Intuitively, each of the \(n-1\) layers iteratively adds one of the \(n-1\) symbols needed to identify the history, resulting in the identification of the full history after the \((n-1)^{\textrm{st}}\) layer. Once the history is identified, the \((n-1)^{\textrm{st}}\) layer can compute the conditional distribution over the next symbol as in the proof of Theorem 3.1. This can be illustrated as the following sequence of 

transformations:14

Footnote 14: We leave out the positional encodings for clarity.

\[{\bf X}^{1}\colon\begin{pmatrix}[y_{1}]&[y_{2}]&[y_{3}]&\cdots&[ y_{t}]&\cdots&[yr]\\ {\bf 0}_{|\Sigma|}&{\bf 0}_{|\Sigma|}&{\bf 0}_{|\Sigma|}&\cdots&{\bf 0}_{|\Sigma|}& \cdots&{\bf 0}_{|\Sigma|}\\ {\bf 0}_{|\Sigma|}&{\bf 0}_{|\Sigma|}&{\bf 0}_{|\Sigma|}&\cdots&{\bf 0}_{|\Sigma|}& \cdots&{\bf 0}_{|\Sigma|}\\ \vdots&\vdots&\vdots&\ddots&\vdots&\ddots&\vdots\\ {\bf 0}_{|\Sigma|}&{\bf 0}_{|\Sigma|}&{\bf 0}_{|\Sigma|}&\cdots&{\bf 0}_{|\Sigma|}& \cdots&{\bf 0}_{|\Sigma|}\end{pmatrix} \tag{57a}\] \[{\bf X}^{2}\colon\begin{pmatrix}[y_{1}]&[y_{2}]&[y_{3}]&\cdots&[ y_{t}]&\cdots&[ yr]\\ {\bf 0}_{|\Sigma|}&[y_{1}]&[y_{2}]&\cdots&[y_{t-1}]&\cdots&[y_{T-1}]\\ {\bf 0}_{|\Sigma|}&{\bf 0}_{|\Sigma|}&{\bf 0}_{|\Sigma|}&\cdots&{\bf 0}_{| \Sigma|}&\cdots&{\bf 0}_{|\Sigma|}\\ \vdots&\vdots&\vdots&\ddots&\vdots&\ddots&\vdots\\ {\bf 0}_{|\Sigma|}&{\bf 0}_{|\Sigma|}&{\bf 0}_{|\Sigma|}&\cdots&\cdots&{\bf 0}_{| \Sigma|}\end{pmatrix}\] (57b) \[\vdots\] \[{\bf X}^{n-1}\colon\begin{pmatrix}[y_{1}]&[y_{2}]&[y_{3}]& \cdots&[y_{t}]&\cdots&[y_{T}]\\ {\bf 0}_{|\Sigma|}&[y_{1}]&[y_{2}]&\cdots&[y_{t-1}]&\cdots&[y_{T-1}]\\ {\bf 0}_{|\Sigma|}&{\bf 0}_{|\Sigma|}&[y_{1}]&\cdots&[y_{t-2}]&\cdots&[y_{ T-1}]\\ \vdots&\vdots&\vdots&\ddots&\vdots&\ddots&\vdots\\ {\bf 0}_{|\Sigma|}&{\bf 0}_{|\Sigma|}&{\bf 0}_{|\Sigma|}&\cdots&[y_{t-n+1}]& \cdots&[y_{T-n+2}]\end{pmatrix} \tag{57c}\]

Such representations can be constructed by starting with the initial symbol encoding (\({\bf X}^{1}\)) and then _passing over_ the information from the \(t^{\text{th}}\) symbol to the contextual representations of the \((t+1)\,^{\text{st}}\) symbol in \({\bf X}^{2}\) by adding a _shifted_ the contextual representation of the \(t^{\text{th}}\) symbol. This is performed \(n-1\) times, resulting in contextual representations of the form \({\bf X}^{n-1}\) where the \(t^{\text{th}}\) contextual representation contains the (ordered) information about the preceding \(n-1\) symbols, i.e., the required history. Intuitively, the \(\ell^{\text{th}}\) transformation of the contextual representation \({\bf x}_{t}^{\ell-1}\) should be of the form

\[{\bf x}_{t}^{\ell}=\underbrace{\sum_{t}^{\ell-1}}_{\text{The previous $\ell-1$ symbols.}}+\underbrace{\downarrow{\bf x}_{t-1}^{\ell-1}}_{\text{ The symbol $\ell$ symbols back shifted one ``cell'' downward.}} \tag{58}\]

The first term is simply the residual connection of the transformer layer. The second term is the shifted representation of the symbol \(y_{t-\ell}\), which can be performed by a simple linear transformation in the value transformation \(V\).

#### b.2.4 Simulation with \(n-1\) Layers: Proofs

We now make the intuition presented in the previous section more formal. Concretely, we only investigate how to identify the correct \(n-1\) symbols in the history with the \(n-1\) layers. We then rely on Lemma B.2 and the derivation from the proof of Theorem 3.1 again to convert this information into a weakly equivalent transformer LM. Let \(\ell\in[n-1]\). Define the following parameters of the attention head of the \(t^{\text{th}}\) transformer layer:

\[{\bf r}\left(y,t\right)\stackrel{{\mbox{\tiny def}}}{{=}} \begin{pmatrix}&&&&&\\ &{\bf 0}_{|\Sigma|}&&&&\\ &\vdots&&&&&\\ &{\bf 0}_{|\Sigma|}&&&\\ &\sqrt{\frac{1}{t}}&&&\\ &\sqrt{1-\frac{1}{t}}&&&\\ &\sqrt{\frac{1}{t+1}}&&& 

\[f\left({\bf q},{\bf k}\right) \stackrel{{\text{def}}}{{=}}\langle{\bf q},{\bf k}\rangle, \tag{60}\]

\[Q\left({\bf x}\right) \stackrel{{\text{def}}}{{=}}{\bf Qx},\quad{\bf Q}\in \mathbb{R}^{2\times((n-1)|\Sigma|+4)} \tag{61a}\] \[K\left({\bf x}\right) \stackrel{{\text{def}}}{{=}}{\bf Kx},\quad{\bf K}\in \mathbb{R}^{2\times((n-1)|\Sigma|+4)}\] (61b) \[V\left({\bf x}\right) \stackrel{{\text{def}}}{{=}}{\bf Vx},\quad{\bf V}\in \mathbb{R}^{((n-1)|\Sigma|+4)\times((n-1)|\Sigma|+4)},\] (61c) \[O\left({\bf x}\right) \stackrel{{\text{def}}}{{=}}{\bf Ox},\quad{\bf O}\in \mathbb{R}^{((n-1)|\Sigma|+4)\times((n-1)|\Sigma|+4)}, \tag{61d}\]

\[{\bf Q}_{:,(n-1)|\Sigma|+1:(n-1)|\Sigma|+2} \stackrel{{\text{def}}}{{=}}{\bf I}_{2}, \tag{62a}\] \[{\bf K}_{:,(n-1)|\Sigma|+3:(n-1)|\Sigma|+4} \stackrel{{\text{def}}}{{=}}{\bf I}_{2},\] (62b) \[{\bf V}_{1+\ell|\Sigma|:1+(\ell+1)|\Sigma|,1+(\ell-1)|\Sigma|:1+ \ell|\Xi|} \stackrel{{\text{def}}}{{=}}{\bf I}_{|\Sigma|}, \tag{62c}\]

where the unspecified elements of \({\bf Q},{\bf K},\) and \({\bf V}\) are \(0.\) Schematically, \({\bf V}\) looks as follows:

\[{\bf V}= \begin{pmatrix}(\ell-1)\underset{\text{symbol slots}}{\underbrace{( \ell-1)|\Sigma|}}&\ell\overset{\text{ob symbol slot}}{\underbrace{|\Sigma|}}&( n-1-\ell)\underset{\text{symbol slots}}{\underbrace{(n-1-\ell)|\Sigma|}}&4\\ \end{pmatrix}\begin{matrix}4&&\text{position slots}\\ \end{matrix} \tag{63}\] \[{\bf V}= \begin{pmatrix}&&&{\bf I}_{|\Sigma|}&&&&\end{matrix}\]

That is, \({\bf I}_{|\Sigma|}\) occupies the "cell" \(\ell\) cells down and \(\ell-1\) right of the top-left corner. Moreover, the matrix \({\bf O}\) is a matrix of all zeros, meaning that \(O\left({\bf x}\right)={\bf 0}\) for all \({\bf x}.\) Again, notice that the position-augmented symbol representation function \({\bf r}\) can be implemented by concatenating or summing a symbol- and a position-specific component.

**Lemma B.5**.: _With the parameters defined above, it holds that_

\[{\bf x}_{t-1}^{\ell}=\begin{pmatrix}\llbracket y_{t-1}\rrbracket\\ \llbracket y_{t-1-1}\rrbracket\\ \llbracket y_{t-1-2}\rrbracket\\ \vdots\\ \llbracket y_{t-1-\ell}\rrbracket\\ {\bf 0}_{|\underline{\Sigma|}}\\ \hline\vdots\\ {\bf 0}_{|\underline{\Sigma|}}\\ \sqrt{\frac{1}{t}}\\ \sqrt{1-\frac{1}{t}}\\ \sqrt{\frac{1}{t+1}}\\ \sqrt{1-\frac{1}{t+1}}\\ \end{pmatrix} \tag{64}\]

Proof.: We prove the lemma by induction. As in the proof of Lemma B.4, we pad the input string with \(n-1\)bos symbols to resolve the case when \(t-1<n-1.\)

Base Case.For \(\ell=1,\) the claim follows from the definition of the symbol representation function \({\bf r}.\) 

Inductive Step.For \(\ell>1\), we assume that the claim holds for \(\ell-1\). We then prove that it holds for \(\ell\) as well. By the construction of the keys and values matrices, as in the proof of Lemma B.4, it holds that the attention head puts all its attention for query \({\bf q}_{t-1}^{\ell}\) on the key \({\bf k}_{t-2}^{\ell}\). This means that

\[{\bf a}_{t-1}^{\ell}={\bf x}_{t-1}^{\ell-1}+{\bf v}_{t-2}^{\ell}. \tag{65}\]

By the induction hypothesis, we have that



Inserting this into Eq. (65), we get that \({\bf a}_{t-1}^{\ell}\) satisfies the required equality:

\[{\bf a}_{t-1}^{\ell}={\bf x}_{t-1}^{\ell-1}+{\bf v}_{t-2}^{\ell} \tag{69a}\] \[=\left(\begin{array}{c}\llbracket y_{t-1}\rrbracket\\ \vdots\\ \llbracket y_{t-1-(\ell-1)}\rrbracket\\ {\bf 0}_{|\Sigma|}\\ \vdots\\ {\bf 0}_{|\Sigma|}\\ \sqrt{\frac{1}{t}}\\ \sqrt{1-\frac{1}{t}}\\ \sqrt{\frac{1}{t+1}}\\ \sqrt{1-\frac{1}{t+1}}\end{array}\right)+\left(\begin{array}{c}\mathbf{0}_{| \Sigma|}\\ \vdots\\ \mathbf{0}_{|\Sigma|}\\ \mathbf{0}_{|\Sigma|}\\ \sqrt{\frac{1}{t}}\\ \sqrt{1-\frac{1}{t}}\\ \sqrt{\frac{1}{t+1}}\\ \sqrt{1-\frac{1}{t+1}}\end{array}\right)=\left(\begin{array}{c}\llbracket y _{t-1}\rrbracket\\ \vdots\\ \llbracket y_{t-1-(\ell-1)}\rrbracket\\ \llbracket y_{t-2-(\ell-1)}\rrbracket\\ \vdots\\ \mathbf{0}_{|\Sigma|}\\ \sqrt{\frac{1}{t}}\\ \sqrt{1-\frac{1}{t}}\\ \sqrt{\frac{1}{t+1}}\\ \sqrt{1-\frac{1}{t+1}}\end{array}\right)=\left(\begin{array}{c}\llbracket y _{t-1}\rrbracket\\ \vdots\\ \llbracket y_{t-1-(\ell-1)}\rrbracket\\ \llbracket y_{t-2-(\ell-1)}\rrbracket\\ \vdots\\ \mathbf{0}_{|\Sigma|}\\ \sqrt{\frac{1}{t}}\\ \sqrt{1-\frac{1}{t}}\\ \sqrt{\frac{1}{t+1}}\\ \sqrt{1-\frac{1}{t+1}}\end{array}\right) \tag{69b}\]

Since the computation of \({\bf x}_{t-1}^{\ell}\overset{\text{def}}{=}{\bf z}_{t-1}^{\ell}=O\left({\bf a }_{t-1}^{\ell}\right)+{\bf a}_{t-1}^{\ell}={\bf a}_{t-1}^{\ell}\) with the definition of \(O\) as the zero function gives us \({\bf z}_{t-1}^{\ell}={\bf a}_{t-1}^{\ell}\), we get the required equality, which finishes the proof. 

This allows us to prove the following theorem.

**Theorem 3.2**.: _For any \(n\)-gram LM, there exists a weakly equivalent \((n-1)\)-layer hard attention transformer LM with a single head._

Proof.: Lemma B.5 shows that the \(n-1\)-layer transformer can identify the history of interest. Applying Lemma B.2 and the same derivation as in the proof of Theorem 3.1 shows that we can construct a weakly equivalent hard attention transformer. 

#### b.2.5 Simulation with a Single Layer and a Single Head: Intuition

While the construction presented here is considerably less intuitive than that of Theorem 3.1, the steps of the proof remain the same--they include the identification of the individual symbols and their positions in the history, combining them into the one-hot encoding of the entire history, and using that to compute the correct next-symbol conditional distributions. This proof focuses on encoding the entire history of interest \({\mathbf{y}}_{t-n+1}^{t-1}\) into a single vector in a way that can be decoded to index the conditional probability distribution as in Definition 2.11. This can then be used to index the appropriate conditional probability distributions as in the proof of Theorem 3.1.

Again, we outline the intuition first. Fix \(t\leqslant|{\mathbf{y}}|\). We decompose the computation of the representation \(\mathsf{enc}\left({\mathbf{y}}_{<t}\right)\) constructed by a single-layer-single-head transformer network as follows.

1. Attending exactly to the history of interest with a single head and a single layer. This can be done by assigning the same attention _score_ with the scoring function to all positions within the history \({\mathbf{y}}_{t-n+1}^{t}\) and a lower score to all other positions. Keeping the definitions of \(Q\) and \(V\) from Theorem 3.1, we can achieve that by defining \[f\left({\bf q}_{t-1},{\bf k}_{j}\right)=-\mathrm{ReLU}(\langle{\bf q}_{t-1},{ \bf k}_{j}\rangle)\] (70) which assigns positions in the history score \(0\) and others negative scores. This is illustrated in Fig. 3. Concretely, the scoring function \(f\) together with hard attention results in attention weights of the form \[s_{j}=\frac{1}{n-1}\mathbbm{1}\left\{j\geqslant t-n+1\right\}.\] (71) 

2. Storing the order of the symbols in the history. We will define the position-augmented representations as position-scaled one-hot encodings of the input symbols. In particular, define \[{\bf r}\left(y_{t-1},j\right)\stackrel{{\text{def}}}{{=}}\begin{pmatrix} 10^{-j}\cdot\llbracket y_{t-1}\rrbracket\\ t\\ 1\end{pmatrix}\in\mathbb{R}^{|\Sigma|+2}\] (72) This effectively stores the information about both the position of the current symbol (with the magnitude) as well as the identity of the symbol \(y_{j}\) (with the index of the non-zero entry). Unlike in the multi-head or multi-layer case, note that in this case, the function \({\bf r}\) is not a concatenation (or addition) of two separate components (one for symbols and one for positions).

Ignoring residual connections, Eq. (71) then implies that

\[{\bf a}_{t-1}=\sum_{j=t-n+1}^{t-1}\frac{1}{n-1}10^{-j}\llbracket y_{j} \rrbracket. \tag{73}\]

For simplicity, we now write \({\bf a}\stackrel{{\text{def}}}{{=}}{\bf a}_{t-1}\). The individual entries of \({\bf a}\) will correspond to symbols in \(\underline{\Sigma}\). The _digits_ of these symbols will encode the positions of the symbols in the history. Concretely, \({\bf a}\) will have a non-zero digit in the \(i^{\text{th}}\) position if and only if the symbol \(y\) appears in the string \(\boldsymbol{y}_{<t}\) at position \(i\) for \(i\in\llbracket t-1\rrbracket\). For example, in a \(5\)-gram LM over the alphabet \(\Sigma=\{a,b\}\), the contextual representation \({\bf a}\) for the history \(\boldsymbol{y}_{t-4:t-1}=abaa\) will be

\[{\bf a}=\frac{1}{4}\begin{pmatrix}10^{t-1}+10^{t-3}+10^{t-4}\\ 10^{t-2}\end{pmatrix}=\frac{1}{4}10^{-t}\begin{pmatrix}0.1011\\ 0.0100\end{pmatrix}. \tag{74}\]

Such representations therefore uniquely encode the history of interest.

Decoding the representations of the history.The representations \({\bf a}\), therefore, contain the information about the history of interest compactly represented in a \(|\Sigma|\)-dimensional vector \({\bf a}\). We now want to construct a function that transforms the constructed vector \({\bf a}\) into a one-hot encoding of the history. To make \({\bf a}\) invariant with respect to \(t\), we first scale it by \(\frac{n-1}{10^{t-t}}\) and define

\[{\bf a}^{\prime}=\frac{n-1}{10^{n-t}}{\bf a}.\] (75) 

Then

\[a_{y}^{\prime}=\sum_{i=1}^{n-1}\mathbbm{1}\left\{y=y_{i+t-n}\right\}10^{-i}. \tag{76}\]

In words, the \(|\Sigma|\) entries of \(\mathbf{a}^{\prime}\) are of the form \(a_{y}^{\prime}=0.d_{1}\ldots d_{n-1}\), where \(d_{i}=1\) if and only if \(y\) appears in the history \(\mathbf{y}_{<t}\) at position \(t-n+i\).

We now focus on a specific symbol \(y\in\underline{\Sigma}\) with the entry \(a_{y}^{\prime}\) in the vector \(\mathbf{a}^{\prime}\) and decode it into a vector that can be used to construct a one-hot encoding of the relevant history \(\mathbf{y}_{t-n+1}^{t-1}\) with \(F\). The "decoding function" will take the form of a \(n-1\)-layer \(\mathrm{ReLU}\)-activated MLP. Intuitively, each of the \(n-1\) layers will contain \(|\underline{\Sigma}|\) neurons, where each of them will compute the values \(d_{i}\) for a particular \(y\in\Sigma\). Now, \(d_{1}\) can be computed as

\[d_{1}=\mathbbm{1}\left\{10^{1}\cdot a_{y}-1+10^{-n}>0\right\}. \tag{77}\]

Since \(\mathbbm{1}\left\{\cdot\right\}\) is not a continuous function, and, unlike in Appendix B.1, the arguments can now take arbitrary real values, it cannot be implemented using a composition of \(\mathrm{ReLU}\) functions. We can, however, simulate the discontinuous function with a linear combination of two \(\mathrm{ReLU}\) functions that together define the same function as \(\mathbbm{1}\left\{\cdot\right\}\) on a subset of \(\mathbb{R}\) relevant for our purposes. Notice that, as long as \(d_{1}=1\), we have that \(10^{1}\cdot a_{y}-1+10^{-n}\geq 10^{-n}\) while we have \(10^{1}\cdot a_{y}-1+10^{-n+1}=10^{1}\cdot 0-1+10^{-n}<0\) if \(d_{1}=0\). This means that our approximation of \(\mathbbm{1}\left\{\cdot\right\}\) only has to map values greater or equal to \(10^{-n}\) to \(1\), rather than all positive values. This allows us to continuously transition from \(0\) to \(1\) as the input to the \(\mathrm{ReLU}\) function increases from \(0\) to \(10^{-n}\). Such a piecewise linear approximation of \(\mathbbm{1}\left\{\cdot\right\}\) can be easily implemented by a linear combination of \(\mathrm{ReLU}\) functions, i.e., with an MLP. See Fig. 4 for an illustration of the approximation.

\(d_{2}\) can then be computed as

\[d_{2}=\mathbbm{1}\left\{10^{2}\cdot a_{y}-10^{1}d_{1}-1+10^{-n}>0\right\} \tag{78}\]

and, in general,

\[d_{i}=\mathbbm{1}\left\{10^{i}\cdot a_{y}-\sum_{j=1}^{i-1}10^{i-j}d_{j}-1+10 ^{-n}>0\right\}. \tag{79}\]

The computation of \(d_{i}\) requires \(i\) layers (since \(d_{j}\) for \(j<i\) have to be computed first), meaning that \(n-1\) layers are required in total. Altogether, these layers compute the values \(d_{i}\) for a single \(y\in\Sigma\). Replicating this computation for every \(y\in\Sigma\) and concatenating the results gives us the desired contextual representation \(\mathbf{z}\).

With this construction, it holds for every \(y\in\Sigma\) that \(d_{i}=1\) if and only if the symbol \(y\) appears in the history \(\mathbf{y}_{t-n+1}^{t-1}\) at position \(t-n+i\). This, therefore, gives us enough information to reconstruct the multi-hot encoding of the history of interest. As in Theorem 3.1, this can then be converted into a one-hot encoding using another \(\mathrm{ReLU}\) layer to implement the logical AND operation. This intuition is formalized in the following section.

#### b.2.6 Simulation with a Single Layer and a Single Head: Proofs

We define the following parameters of the transformer head.15 

* Static encodings \[{\bf r}\left(y_{t-1},j\right)\stackrel{{\text{def}}}{{=}}\begin{pmatrix} 10^{-j}\cdot\llbracket y_{t-1}\rrbracket\\ 1\\ t\end{pmatrix}\in\mathbb{R}^{|\Sigma|+3}\] (80)
* Query, key, value, and output transformations \[Q\left({\bf x}\right)\stackrel{{\text{def}}}{{=}}{ \mathbf{Q}}{\bf x}+{\mathbf{b}}_{Q}, {\mathbf{Q}}\in\mathbb{R}^{2\times(|\Sigma|+2)},{\mathbf{b}}_{ Q}\in\mathbb{R}^{2},\] (81a) \[K\left({\bf x}\right)\stackrel{{\text{def}}}{{=}}{ \mathbf{K}}{\bf x}, {\mathbf{K}}\in\mathbb{R}^{2\times(|\Sigma|+2)},\] (81b) \[V\left({\bf x}\right)\stackrel{{\text{def}}}{{=}}{ \mathbf{V}}{\bf x}, {\mathbf{V}}\in\mathbb{R}^{|\Sigma|\times(|\Sigma|+2)},\] (81c) \[O\left({\bf x}\right)\stackrel{{\text{def}}}{{=}}{ \mathbf{I}}_{|\Sigma|}{\bf x},\] (81d)

\[{\mathbf{Q}}_{:,|\underline{\Sigma}|+1;|\underline{\Sigma}|+2} \stackrel{{\text{def}}}{{=}}\begin{pmatrix}0&1\\ -1&0\end{pmatrix},\] (82a) \[{\mathbf{b}}_{Q} \stackrel{{\text{def}}}{{=}}\begin{pmatrix}-\left(n -2\right)\\ 0\end{pmatrix}\] (82b) \[{\mathbf{K}}_{:,|\underline{\Sigma}|+1;|\underline{\Sigma}|+2} \stackrel{{\text{def}}}{{=}}{\mathbf{I}}_{2}\] (82c) \[{\mathbf{V}}_{:,1:|\underline{\Sigma}|} \stackrel{{\text{def}}}{{=}}{\mathbf{I}}_{|\underline{ \Sigma}|}\] (82d) where the unspecified elements of \({\mathbf{Q}},{\mathbf{K}},{\mathbf{V}},{\mathbf{O}}\) are \(0\).

A large part of the proof of correctness will rely on identifying the digits of the contextual representations of strings. We will rely heavily on the following definition.

**Definition B.2**.: _Let \(x\in\mathbb{Q}\cap\left(10^{-(N+1)},10^{-1}\right]\) be a rational-valued number with at most \(N\) digits in its decimal representation. We define \(d_{i}\left(x\right)\) as the \(i^{th}\) digit of \(x\) for \(i\in[N]\). We also group these \(N\) values into the vector \({\mathbf{d}}\left(x\right)\):_

\[{\mathbf{d}}\left(x\right)\stackrel{{\text{def}}}{{=}}\begin{pmatrix} d_{1}\left(x\right)\\ \vdots\\ d_{N}\left(x\right)\end{pmatrix}. \tag{83}\]

**Lemma B.6**.: _A transformer with the parameters and functions defined in Eq. (80)-Eq. (82d) computes for string \({\boldsymbol{y}}\in\Sigma^{*}\)_

\[{\mathbf{a}}_{t-1}=\sum_{j=t-n+1}^{t-1}\frac{1}{n-1}10^{-j}\llbracket y_{j} \rrbracket. \tag{84}\]

Proof.: By construction, we have

\[{\mathbf{q}}_{t-1} =\begin{pmatrix}t-1-\left(n-2\right)\\ -1\end{pmatrix}\] (85a) \[{\mathbf{k}}_{j} =\begin{pmatrix}1\\ j\end{pmatrix}\] (85b) \[{\mathbf{v}}_{j} =10^{-j}\cdot\llbracket y_{t-1}\rrbracket,\] (85c) 

thus

\[f\left({\mathbf{q}_{t-1},\mathbf{k}_{j}}\right) =-\mathrm{ReLU}(\langle\mathbf{q}_{t-1},\mathbf{k}_{j}\rangle)\] (86a) \[=-\mathrm{ReLU}\bigg{(}\bigg{<}\binom{t-1-\left(n-2\right)}{-1}, \binom{1}{j}\bigg{>}\bigg{)}\] (86b) \[=-\mathrm{ReLU}(t-1-n+2-j)\] (86c) \[=-\mathrm{ReLU}(t-n+1-j)\] (86d) \[=\begin{cases}0&\textbf{if $j\geq t-n+1$}\\ <0&\textbf{otherwise}\end{cases}\] (86e) meaning that \[s_{j}=\frac{1}{n-1}\mathbbm{1}\left\{j\geq t-n+1\right\}.\] (87) This means that \[\mathbf{a}_{t-1}=\sum_{j=t-n+1}^{t-1}\frac{1}{n-1}\mathbf{v}_{j}=\sum_{j=t-n+1 }^{t-1}\frac{1}{n-1}10^{-j}\llbracket y_{j}\rrbracket\] (88) which is what we needed to prove. \(\blacksquare\)

**Lemma B.7**.: _Define_

\[\mathbf{a}^{\prime}\stackrel{{\text{def}}}{{=}}\frac{n-1}{10^{ n-t}}\;\mathbf{a}_{t-1} \tag{89}\]

_with \(\mathbf{a}_{t-1}\) from Lemma b.6. Indexing the \(|\Sigma|\) elements of \(\mathbf{a}^{\prime}\) directly with \(y\in\underline{\Sigma}\), it holds that_

\[d_{i}\left(a_{y}^{\prime}\right)=1\iff y_{t-n+i}=y \tag{90}\]

_for all \(i\in\llbracket N\rrbracket\) and \(d_{i}\left(x\right)=0\) for all \(i>N\)._

Proof.: By Lemma B.6, a contains entries of the form

\[a_{y} =\frac{1}{n-1}\sum_{j=t-n+1}^{t-1}\mathbbm{1}\left\{y=y_{j} \right\}10^{-j}\] (91a) \[=\frac{1}{n-1}\sum_{j^{\prime}=1}^{t-1-\left(t-n\right)}\mathbbm{ 1}\left\{y=y_{j^{\prime}+t-n}\right\}10^{-\left(j^{\prime}+t-n\right)}\] (91b, Change of variables.) \[=\frac{1}{n-1}\sum_{j=1}^{n-1}\mathbbm{1}\left\{y=y_{j+t-n} \right\}10^{n-t-j}\] (91c, Change of variables.) \[=\frac{10^{n-t}}{n-1}\sum_{j=1}^{n-1}\mathbbm{1}\left\{y=y_{j+t-n }\right\}10^{-j}\] (91d, Distributivity.)

for \(y\in\underline{\Sigma}\). Then

\[a_{y}^{\prime} \stackrel{{\text{def}}}{{=}}\frac{n-1}{10^{n-t}}a_{y} \tag{92a}\] \[=\frac{n-1}{10^{n-t}}\frac{10^{n-t}}{n-1}\sum_{j=1}^{n-1} \mathbbm{1}\left\{y=y_{j+t-n}\right\}10^{-j}\] (92b) \[=\sum_{j=1}^{n-1}\mathbbm{1}\left\{y=y_{j+t-n}\right\}10^{-j}. \tag{92c}\]

This implies that \(d_{i}\left(a_{y}^{\prime}\right)=1\iff y_{t-n+i}=y\) for \(i\in\llbracket N\rrbracket\) and that \(d_{i}\left(x\right)=0\) for \(i>N\), which is what we wanted to prove 



**Lemma B.8**.: _Let \(x\in{\mathbb{Q}}\,\cap\left(10^{-\left(N+1\right)},10^{-1}\right]\) with \(d_{i}\left(x\right)\in\left\{0,1\right\}\) for \(i\in[N]\). Then, \(d_{i}\left(x\right)\) satisfy the equality_

\[d_{i}\left(x\right)=\mathbbm{1}\left\{10^{i}\cdot x-\sum_{j=1}^{ i-1}10^{i-j}\;d_{j}\left(x\right)-1+10^{-\left(N+1\right)}>0\right\}. \tag{93}\]

_for all \(i\in[N]\)._

Proof.: Let \(x\in{\mathbb{Q}}\,\cap\left(10^{-\left(N+1\right)},10^{-1}\right]\) with \(d_{i}\left(x\right)\in\left\{0,1\right\}\) for \(i\in[N]\). Then, by definition of \(d_{i}\left(x\right)\), we have that

\[x=\sum_{j=1}^{N}d_{j}\left(x\right)10^{-j}. \tag{94}\]

This means that

\[10^{i}x =10^{i}\sum_{j=1}^{N}d_{j}\left(x\right)10^{-j} \tag{95a}\] \[=\sum_{j=1}^{N}d_{j}\left(x\right)10^{i-j}\] (95b) \[=\sum_{j=1}^{i-1}d_{j}\left(x\right)10^{i-j}+d_{i}\left(x\right)+ \sum_{j=i+1}^{N}d_{j}\left(x\right)10^{i-j},\] (95c) implying that \[10^{i}\cdot x-\sum_{j=1}^{i-1}10^{i-j}\;d_{j}\left(x\right)=d_{i}\left(x \right)+\sum_{j=i+1}^{N}d_{j}\left(x\right)10^{i-j}. \tag{96}\]

Suppose now that \(d_{i}\left(x\right)=1\). Then

\[10^{i}\cdot x-\sum_{j=1}^{i-1}10^{i-j}\;d_{j}\left(x\right)-1+10 ^{-\left(N+1\right)} =d_{i}\left(x\right)+\sum_{j=i+1}^{N}d_{j}\left(x\right)10^{i-j}-1 +10^{-\left(N+1\right)} \tag{97a}\] \[=1+\sum_{j=i+1}^{N}d_{j}\left(x\right)10^{i-j}-1+10^{-\left(N+1 \right)}\] (97b) \[=\sum_{j=i+1}^{N}d_{j}\left(x\right)10^{i-j}+10^{-\left(N+1 \right)}>0,\] (97c) meaning that \[\mathbbm{1}\left\{10^{i}\cdot x-\sum_{j=1}^{i-1}10^{i-j}\;d_{j} \left(x\right)-1+10^{-\left(N+1\right)}\right\}=1=d_{i}\left(x\right). \tag{98}\]

Suppose, on the contrary, that \(d_{i}\left(x\right)=0\). Then

\[10^{i}\cdot x-\sum_{j=1}^{i-1}10^{i-j}~{}d_{j}\left(x\right)-1+10^{ -\left(N+1\right)} =d_{i}\left(x\right)+\sum_{j=i+1}^{N}d_{j}\left(x\right)10^{i-j}-1 +10^{-\left(N+1\right)}\] (99a) \[=0+\sum_{j=i+1}^{N}d_{j}\left(x\right)10^{i-j}-1+10^{-\left(N+1 \right)}\] (99b) \[=\sum_{j=i+1}^{N}d_{j}\left(x\right)10^{i-j}+10^{-\left(N+1 \right)}-1\] (99c) \[=\sum_{j^{\prime}=1}^{N-\left(i+1\right)}d_{j^{\prime}+i+1}\left( x\right)10^{i-j^{\prime}-i-1}+10^{-\left(N+1\right)}-1\] (99d) \[=\underbrace{\sum_{j^{\prime}=1}^{N-\left(i+1\right)}d_{j^{\prime }+i+1}\left(x\right)10^{-j^{\prime}-1}+10^{-\left(N+1\right)}}_{<1}-1<0\] (99e) meaning that \[\mathbbm{1}\left\{10^{i}\cdot x-\sum_{j=1}^{i-1}10^{i-j}~{}d_{j}\left(x \right)-1+10^{-\left(N+1\right)}\right\}=0=d_{i}\left(x\right)\text{.}\] (100) This finishes the proof. 

**Lemma B.9**.: _Let \(x\in\mathbb{Q}\cap\left(10^{-\left(N+1\right)},10^{-1}\right]\) with \(d_{i}\left(x\right)\in\left\{0,1\right\}\) for \(i\in\left[N\right]\). Given \(d_{j}\left(x\right)\) for \(j<i\), \(d_{i}\left(x\right)\) can be computed by a single layer MLP._

Proof.: By Lemma B.8, we can use the knowledge of \(d_{j}\left(x\right)\) for \(j<i\), \(d_{i}\left(x\right)\) to implement the function

\[d_{i}\left(x\right)=\mathbbm{1}\left\{10^{i}\cdot x-\sum_{j=1}^{i-1}10^{i-j}~ {}d_{j}\left(x\right)-1+10^{-\left(N+1\right)}>0\right\}. \tag{101}\]

with an MLP. For any \(i\in\left[N\right]\), the inner function

\[\begin{split}\left(\begin{matrix}d_{1}\left(x\right)\\ \vdots\\ d_{i-1}\left(x\right)\\ x\end{matrix}\right)&\mapsto 10^{i}\cdot x-\sum_{j=1}^{i-1}10^{ i-j}~{}d_{j}\left(x\right)-1+10^{-\left(N+1\right)}\\ x&\hskip 256.490134pt=\left\langle\underbrace{\begin{pmatrix}-10^{ i-1}\\ \vdots\\ -10^{1}\\ 10^{i}\end{pmatrix}}_{\mathbf{W}^{\top}},\left(\begin{matrix}d_{1}\left(x\right)\\ \vdots\\ d_{i-1}\left(x\right)\\ x\end{matrix}\right)\right\rangle\underbrace{-1+10^{-\left(N+1\right)}}_{ \mathbf{b}}\end{split} \tag{102b}\]

is an affine transformation. The indicator function in Eq. (101), however, is discontinuous and can thus not be implemented by a composition of continuous \(\mathrm{ReLU}\) MLPs. Here, we take advantage of the fact that

\[10^{i}\cdot x-\sum_{j=1}^{i-1}10^{i-j}~{}d_{j}\left(x\right)-1+10^{-\left(N+1 \right)}\in\underbrace{\left(-\infty,0\right]\cup\left[10^{-\left(N+1\right)}, \infty\right)}_{\mathcal{I}}.\] (103) 

The MLP

\[{\rm MLP}_{\mathcal{I}}\left(z\right) =10^{N+1}\left({\rm ReLU}(z)-{\rm ReLU}\Big{(}z-10^{-(N+1)}\Big{)}\right) \tag{104a}\] \[=\begin{pmatrix}10^{N+1}&-10^{N+1}\end{pmatrix}{\rm ReLU}\bigg{(} \binom{1}{1}z+\binom{0}{-10^{-(N+1)}}\bigg{)} \tag{104b}\]

matches \(\mathbbm{1}\left\{\cdot>0\right\}\) on \(\mathcal{I}\), as we show next. See Fig. 4 for an illustration of the approximation. First, assume that \(z\leqslant 0\). Then

\[{\rm MLP}_{\mathcal{I}}\left(z\right)=10^{N+1}\left({\rm ReLU}( z)-{\rm ReLU}\Big{(}z-10^{-(N+1)}\Big{)}\right)=10^{N+1}\left(0-0\right)=0= \mathbbm{1}\left\{0>0\right\}. \tag{105}\]

On the contrary, assuming \(z\geqslant 10^{-(N+1)}\), we have that

\[{\rm MLP}_{\mathcal{I}}\left(z\right) =10^{N+1}\left({\rm ReLU}(z)-{\rm ReLU}\Big{(}z-10^{-(N+1)} \Big{)}\right) \tag{106a}\] \[=10^{N+1}\left(z-\Big{(}z-10^{-(N+1)}\Big{)}\right)\] (106b) \[=10^{N+1}\left(z-z+10^{-(N+1)}\right)\] (106c) \[=10^{N+1}\left(10^{-(N+1)}\right)\] (106d) \[=1=\mathbbm{1}\left\{z>0\right\}. \tag{106e}\]

We can therefore construct the MLP \({\rm MLP}\) computing Eq. (101) on \(\mathcal{I}\) by a composition of Eq. (102b) (computing \(z\) in Eq. (104b)) and the MLP \({\rm MLP}_{\mathcal{I}}\) from Eq. (104b). 

**Lemma B.10**.: _Let \(N\in\mathbb{N}\). There exists an MLP \({\rm MLP}\) such that_

\[{\rm MLP}\left(x\right)=\mathbf{d}\left(x\right) \tag{107}\]

_for all \(x\in\mathbb{Q}\cap\big{(}10^{-(N+1)},10^{-1}\big{)}\) with \(d_{i}\left(x\right)\in\{0,1\}\) for \(i\in[N]\)._

Proof.: At a very high level, we will construct an \(N\)-layer MLP performing the transformations

\[x\mapsto\begin{pmatrix}x\\ x\\ \vdots\\ x\end{pmatrix}\mapsto\begin{pmatrix}d_{1}\left(x\right)\\ d_{2}\left(x\right)\\ \vdots\\ x\end{pmatrix}\mapsto\cdots\mapsto\begin{pmatrix}d_{1}\left(x\right)\\ d_{2}\left(x\right)\\ \vdots\\ d_{N}\left(x\right)\end{pmatrix}=\mathbf{d}\left(x\right). \tag{108}\]

By Lemma B.9, all individual transformations can be performed exactly by a single-layer MLP. The composition of the \(N\) layers results in the vector \(\mathbf{d}\left(x\right)\).

The transformation \(x\mapsto\begin{pmatrix}x\\ x\\ \vdots\\ x\end{pmatrix}\) is a simple linear transformation. We now construct the \(\ell^{\text{th}}\) layer \(\mathbf{f}_{\ell}\) of

the MLP with parameters \(\mathbf{W}_{\ell}\) and \(\mathbf{b}_{\ell}\) (cf. Definition B.1), assuming that it has the input 

layer \({\mathbf{f}_{\ell}}\) has to satisfy

\[{\mathbf{f}}\left(\begin{pmatrix}d_{1}\,(x)\\ d_{2}\,(x)\\ \vdots\\ d_{\ell-1}\,(x)\\ x\\ x\\ \vdots\\ x\end{pmatrix}\right)=\begin{pmatrix}d_{1}\,(x)\\ d_{2}\,(x)\\ \vdots\\ d_{\ell-1}\,(x)\\ d_{\ell}\,(x)\\ x\\ \vdots\\ x\end{pmatrix}, \tag{109}\]

i.e., it must copy all \(n-1\) entries apart from the \(\ell^{\text{th}}\) one. Thus, \(\mathbf{W}_{k,k^{\prime}}=\mathbbm{1}\left\{k=k^{\prime}\right\}\) for \(k\neq\ell\) and \(\mathbf{b}_{k}=0\) for all \(k\neq\ell\) (here, we write \(\mathbf{W}\) and \(\mathbf{b}\) for \(\mathbf{W}_{\ell}\) and \(\mathbf{b}_{\ell}\) to avoid clutter). To define the remaining \(\ell^{\text{th}}\) row, we use Lemma B.9. It tells us that defining

\[\mathbf{W}_{\ell,:} \stackrel{{\mbox{\tiny def}}}{{=}}\begin{pmatrix}10^ {\ell-1}&\cdots&10^{1}&10^{\ell}&0&\cdots&0\\ 10^{\ell-1}&\cdots&10^{1}&10^{\ell}&0&\cdots&0\end{pmatrix}\] (110a) \[b_{\ell} \stackrel{{\mbox{\tiny def}}}{{=}}\begin{pmatrix}-1+ 10^{-n}\\ -1+10^{-n}-10^{-n}\end{pmatrix}=\begin{pmatrix}-1+10^{-n}\\ -1\end{pmatrix}\] (110b) will result in the \[\ell^{\text{th}}\] row of \[{\mathbf{f}_{\ell}}\] computing exactly \[d_{\ell}\,(x)\] after being multiplied by the matrix \[\mathbf{W}^{\prime}\stackrel{{\mbox{\tiny def}}}{{=}}\begin{pmatrix} 10^{n}&10^{n}\end{pmatrix}.\] (111) This is what Eq. ( 109 ) requires. The parameters \[\mathbf{W}_{\ell,:}\] and \[b_{\ell}\] represent the parameters of the affine transformation in Lemma B.9. Note that the matrix \[\mathbf{W}^{\prime}\] is not part of the original definition of the MLP. However, it can easily be absorbed into the matrix \[\mathbf{W}_{\ell+1}\] in the actual implementation (at the cost of duplicating the size of the hidden state). We keep it here to make the presentation more intuitive. Since any layer \[{\mathbf{f}_{\ell}}\] can be defined like this, and the final MLP is their composition, this finishes the proof. 

**Lemma B.11**.: _Given \(\mathbf{a}_{t-1}\) from Lemma B.6, it holds that_

\[\left\lVert\mathbf{a}_{t-1}\right\rVert_{1}=\frac{10^{n-1-t}}{ n-1}\frac{1-\frac{1}{10^{n-1}}}{1-\frac{1}{10}}. \tag{112}\]

Proof.: By Lemma B.7, we have that

\[\mathbf{a}\stackrel{{\mbox{\tiny def}}}{{=}}\mathbf{ a}_{t-1}=\sum_{j=t-n+1}^{t-1}\frac{1}{n-1}10^{-j}\llbracket y_{ j}\rrbracket \tag{113}\]

We compute:

\[\left\lVert\mathbf{a}\right\rVert_{1} =\frac{1}{n-1}\sum_{j=t-n+1}^{t-1}10^{-j}\] (114a) \[=\frac{1}{n-1}\sum_{j^{\prime}=0}^{t-1-(t-n+1)}10^{-(j^{\prime}+ t-n+1)}\] (114b) \[=\frac{1}{n-1}10^{-(t-n+1)}\sum_{j=0}^{n-2}\frac{1}{10^{j}}\] (114c) \[=\frac{10^{-(t-n+1)}}{n-1}\frac{1-\frac{1}{10^{n-1}}}{1-\frac{1} {10}}\] (114d) \[=\frac{10^{n-1-t}}{n-1}\frac{1-\frac{1}{10^{n-1}}}{1-\frac{1}{10 }}.\] (114e) 



**Corollary B.1**.: _Given \({\bf a}\stackrel{{\mbox{\tiny def}}}{{=}}{\bf a}_{t-1}\) from Lemma B.6 and \({\bf a^{\prime}}\stackrel{{\mbox{\tiny def}}}{{=}}{\bf a}^{\prime}_ {t-1}\) from Lemma B.7, it holds that_

\[\frac{1}{\left\|{\bf a}\right\|_{1}}\cdot 10\cdot\frac{1-\frac{1}{10^{n-1}}}{1- \frac{1}{10}}\cdot{\bf a}={\bf a}^{\prime}. \tag{115}\]

Proof.: By Lemma B.11, we can write

\[\left\|{\bf a}\right\|_{1}=\frac{10^{n-1-t}}{n-1}\frac{1-\frac{1}{10^{n-1}}}{1 -\frac{1}{10}}=\frac{10^{n-t}}{n-1}\cdot Z \tag{116}\]

and

\[\frac{1}{\left\|{\bf a}\right\|_{1}}=\frac{n-1}{10^{n-t}}\cdot\frac{1}{Z} \tag{117}\]

where \(Z\stackrel{{\mbox{\tiny def}}}{{=}}\frac{1}{10}\cdot\frac{1}{10 ^{n-1}}{1-\frac{1}{10}}\) is a constant independent of \(t\). Then

\[\frac{1}{\left\|{\bf a}\right\|_{1}}\cdot\frac{1}{10}\cdot\frac{1 -\frac{1}{10^{n-1}}}{1-\frac{1}{10}}\cdot{\bf a} =\frac{n-1}{10^{n-t}}\cdot\frac{1}{Z}\cdot\underbrace{\frac{1}{1 0}\cdot\frac{1-\frac{1}{10^{n-1}}}{1-\frac{1}{10}}}_{Z}\cdot{\bf a} \tag{118a}\] \[=\frac{n-1}{10^{n-t}}\cdot\frac{1}{Z}\cdot Z\cdot{\bf a}\] (118b) \[=\frac{n-1}{10^{n-t}}\cdot{\bf a}\] (118c) \[={\bf a}^{\prime} \tag{118d}\]

**Lemma B.12**.: _Let \(\Sigma\) be an alphabet. Given the transformer parameters and functions defined in Eq. (80)-Eq. (82d), there exists an MLP \(F\) (whose inputs are \(\left\|\cdot\right\|_{1}\)-normalized) that maps the contextual representations \({\bf a}\) of \(\boldsymbol{y}_{<t}\) into a one-hot encoding of \(\boldsymbol{y}^{t-1}_{t-n+1}\) for any string \(\boldsymbol{y}\in\Sigma^{*}\) and any \(t\in[\left|\boldsymbol{y}\right|]\)._

Proof.: By Lemma B.7, we have that

\[{\bf a}=\sum_{j=t-n+1}^{t-1}\frac{1}{n-1}10^{-j}[y_{j}] \tag{119}\]

and that \(d_{i}\left(a^{\prime}_{y}\right)=1\iff y_{t-n+i}=y\) for \({\bf a}^{\prime}\stackrel{{\mbox{\tiny def}}}{{=}}\frac{n-1}{10^{ n-t}}\)\({\bf a}\) (where \(a^{\prime}_{i}=0\) for \(i>n-1\), from the same lemma). We can express the entries of the one-hot encoding of the history \(\boldsymbol{y}^{t-1}_{t-n+1}\), \([\![\boldsymbol{y}^{t-1}_{t-n+1}]\!]\), as

\[[\![\boldsymbol{y}^{t-1}_{t-n+1}]\!]_{y_{t-n+1}\ldots y_{t-1}}=1\iff d_{1} \left(a^{\prime}_{y_{t-n+1}}\right)\wedge\cdots\wedge d_{n-1}\left(a^{\prime}_ {y_{t-1}}\right). \tag{120}\]

By Lemma B.10, the vectors \({\bf d}\left(a^{\prime}_{y_{t-n+1}}\right),\dots,{\bf d}\left(a^{\prime}_{y_{ t-1}}\right)\) can be computed by an \(n-1\)-layer MLP. Each of these vectors is of size \(n-1\) and, among others, contains the values \(d_{1}\left(a^{\prime}_{y_{t-n+1}}\right),\dots,d_{n-1}\left(a^{\prime}_{y_{t-1 }}\right)\). Since the entries \([\![\boldsymbol{y}^{t-1}_{t-n+1}]\!]_{y_{t-n+1}\ldots y_{t-1}}\) can be expressed as the results of the logical AND operation, their computation can be performed by a single-layer \(\mathrm{ReLU}\) MLP as per Lemma B.1. The MLP \(F\) can therefore be constructed as a composition of three functions:

1. The scaling \({\bf a}\mapsto\frac{1}{\left\|{\bf a}\right\|_{1}}\cdot\frac{1}{10}\cdot\frac{ \cdot\frac{1}{10^{n-1}}}{1-\frac{1}{10}}\cdot{\bf a}\), which results in \({\bf a}^{\prime}\) by Corollary B.1.
2. The concatenation of the \(\left|\underline{\Sigma}\right|\)\(n-1\)-layer MLPs computing \({\bf d}\left(a^{\prime}_{y}\right)\) for all \(y\in\underline{\Sigma}\). This results in \((n-1)\left|\underline{\Sigma}\right|\) binary values altogether.



3. The MLP performing the AND operation between the entries of \({\bf d}\left(a_{y}^{\prime}\right)\).

This finishes the proof. 

**Theorem 3.3**.: _For any \(n\)-gram LM, there exists a weakly equivalent single-layer hard attention transformer LM with a single head._

Proof.: To show that there exists a weakly equivalent single-layer-single-head transformer LM to any \(n\)-gram LM, we combine the lemmata in this section. Let \(\mathcal{T}\) be a transformer LM defined with the parameters and functions defined in Eq.80-Eq.82d. By Lemma B.6, the representations \({\bf a}_{t-1}=\sum_{j=t-n+1}^{t-1}\frac{1}{n-1}10^{-j}[\![y_{j}]\!]\) computed by \(\mathcal{T}\) contain information about the symbols and their positions in the history \(\mathbf{y}_{t-n+1}^{t-1}\). By Lemma B.12 then, \({\bf a}\) can be mapped to the one-hot encoding of the history with a \(n-1\)-layer MLP \(F\). This one-hot encoding can then be used to index (the logits of) the probabilities stored in the output matrix \({\bf E}\) defining a weakly equivalent transformer. 

## Appendix C Sparse Attention

We now prove a lemma analogous to Lemma B.4. It shows that a sparse attention transformer head can isolate a particular symbol in the string. First, define the following position-augmented symbol representation function of the transformer head \(h\):

\[{\bf r}\left(y,t\right)\stackrel{{\text{\tiny def}}}{{=}}\begin{pmatrix} \llbracket y\rrbracket\\ \mathbf{0}_{|\Sigma|}\\ 1\\ t\end{pmatrix}\in\{0,1\}^{2|\Sigma|+2} \tag{121}\]

and the scoring function

\[f\left({\bf q},{\bf k}\right)\stackrel{{\text{\tiny def}}}{{=}} -\big{|}{\bf q}^{\top}{\bf k}\big{|}. \tag{122}\]

Here, the position-augmented symbol representation function \({\bf r}\) can again be implemented by concatenating or summing a symbol- and a position-specific component. Lastly, we define the transformation matrices

\[Q\left({\bf x}\right)\stackrel{{\text{\tiny df}}}{{=}}{\bf Qx}+{ \bf b}_{Q}, {\bf Q} \in\mathbb{R}^{2\times(2|\Sigma|+2)},{\bf b}_{Q}\in\mathbb{R}^{2}, \tag{123a}\] \[K\left({\bf x}\right)\stackrel{{\text{\tiny df}}}{{= }}{\bf Kx}, {\bf K} \in\mathbb{R}^{2\times(2|\Sigma|+2)},\] (123b) \[V\left({\bf x}\right)\stackrel{{\text{\tiny df}}}{{= }}{\bf Vx}, {\bf V} \in\mathbb{R}^{(2|\Sigma|+2)\times(2|\Sigma|+2)},\] (123c) \[O\left({\bf x}\right)\stackrel{{\text{\tiny df}}}{{= }}{\bf Ox}, {\bf O} \in\mathbb{R}^{(2|\Sigma|+2)\times(2|\Sigma|+2)}, \tag{123d}\]

\[{\bf Q}_{;2|\Sigma|+1:2|\Sigma|+2} \stackrel{{\text{\tiny df}}}{{=}}{\bf I}_{2} \tag{124a}\] \[{\bf b}_{Q} \stackrel{{\text{\tiny df}}}{{=}}{\scriptsize\begin{pmatrix} 0\\ -h\end{pmatrix}}\] (124b) \[{\bf K}_{;2|\Sigma|+1:2|\Sigma|+2} \stackrel{{\text{\tiny df}}}{{=}}{\scriptsize\begin{pmatrix} 0&1\\ -1&0\end{pmatrix}},\] (124c) \[{\bf V}_{|\Sigma|+1:2|\Sigma|,1:|\Sigma|} \stackrel{{\text{\tiny df}}}{{=}}{\bf I}_{|\Sigma|}\] (124d) \[{\bf O}_{1:|\Sigma|,1:|\Sigma|} \stackrel{{\text{\tiny df}}}{{=}}{\bf-I}_{|\Sigma|} \tag{124e}\]

where the unspecified elements of \({\bf Q},{\bf K},{\bf V},{\bf O}\) are \(0\).

**Lemma C.1**.: _Let \(\Sigma\) be an alphabet and \(\mathbf{y}\in\Sigma^{*}\). For any \(t\in[|\mathbf{y}|]\), a transformer head defined with the parameters above outputs_

\[{\bf z}_{t-1}=\begin{pmatrix}\mathbf{0}_{|\Sigma|}\\ \llbracket y_{t-1-h}\rrbracket\\ \mathbf{0}_{2}\end{pmatrix}. \tag{125}\]

_In particular, this means that the output \({\bf z}_{t-1}\) at time step \(t-1\) contains the one-hot encoding of the symbol at position \(t-h-1\)._ 

Proof.: The construction is largely identical to the one in Theorem 3.1, with one crucial difference: It relies on simpler, but _unbounded_ positional encodings and a less-standard, but still easily implementable attention scoring function in the form of an MLP.

It is easy to see that the query and value transformations result in: equationparentequation

\[{\bf q}_{t-1} =Q\left({\bf r}\left(y_{t},t\right)\right)=\binom{1}{t-1-h} \tag{126a}\] \[{\bf k}_{j} =K\left({\bf r}\left(y_{j},j\right)\right)=\binom{-j}{1}\text{.} \tag{126b}\]

Thus, we get that equationparentequation

\[f\left({\bf q}_{t-1},{\bf k}_{j}\right) =-|{\bf q}^{\top}_{t-1}{\bf k}_{j}| \tag{127a}\] \[=-\Bigg{|}\binom{1}{t-1-h}^{\top}\binom{-j}{1}\Bigg{|}\] (127b) \[=-|j-(t-h-1)|\text{.} \tag{127c}\]

\(f\) clearly has a unique maximum at \(j^{*}=t-1-h\). Moreover, by construction, for any \(j\neq j^{*}\). This is a crucial property of the scoring function and one that allows sparsemax to _uniquely_ attend to \(j^{*}\); by Lemma C.2, it holds that

\[\operatorname{sparsemax}\left(f\left({\bf q}_{t-1},{\bf k}_{1}\right),\ldots ,f\left({\bf q}_{t-1},{\bf k}_{t-1}\right)\right)=\mathrm{hardmax}\left(f \left({\bf q}_{t-1},{\bf k}_{1}\right),\ldots,f\left({\bf q}_{t-1},{\bf k}_{ t-1}\right)\right)\text{,} \tag{128}\]

meaning that

\[\operatorname{sparsemax}\left(f\left({\bf q}_{t-1},{\bf k}_{1}\right),\ldots ,f\left({\bf q}_{t-1},{\bf k}_{t-1}\right)\right)_{j}=\mathbbm{1}\left\{j=t-1- h\right\} \tag{129}\]

as in the proof of Lemma B.4. This is exactly the same as Eq. (51). Since \(O\) and \(V\) result in the same vectors as in Lemma B.4, the remainder of the proof is the same as in Lemma B.4. 

**Lemma C.2**.: _Let \({\bf x}\in\mathbb{R}^{D}\). If \(\underset{d=1}{\overset{D}{\max}}x_{d}\geq\underset{d\not\!\mathrm{argmax} \left({\bf x}\right)}{\max}\quad x_{d}+1\), then \(\operatorname{sparsemax}\left({\bf x}\right)=\mathrm{hardmax}\left({\bf x}\right)\)._

Proof.: Let \({\bf x}\in\mathbb{R}^{D}\) and let \(x_{(1)}\geq x_{(2)}\geq\ldots\geq x_{(D)}\) be the non-increasing entries of \({\bf x}\). Due to the additive invariance of the softmax (Martins and Astudillo (2016, Proposition 2)), we can assume that \(x_{(1)}=0\) and \(x_{(2)}\leq-1\). By Martins and Astudillo (2016, Proposition 1),

\[\operatorname{sparsemax}\left({\bf x}\right)_{d}=\max\left(0,x_{d}-\tau\left( {\bf x}\right)\right)\text{,} \tag{130}\]

where

\[\tau\left({\bf x}\right)\stackrel{{\mbox{\tiny def}}}{{=}}\frac{ \sum_{j=1}^{k\left({\bf x}\right)}x_{(j)}-1}{k\left(x\right)} \tag{131}\]

and

\[k\left({\bf x}\right)\stackrel{{\mbox{\tiny def}}}{{=}}\max \left(k\in\left[D\right]\mid 1+kx_{(k)}>\sum_{j=1}^{k}x_{(j)}\right)\text{.} \tag{132}\]

It suffices to show that \(k\left({\bf x}\right)=1\). For \(k=1\), we get

\[1+1\cdot x_{(1)}=1+1\cdot 0=1=1\geq 0=x_{(1)}\text{.}\] (133) 

For \(k=2\), we get

\[1+2\cdot x_{(2)} =1+x_{(2)}+x_{(2)}\] (134a) \[\leqslant 1-1+x_{(2)}\] (134b) \[=x_{(2)}\] (134c) \[=x_{(1)}+x_{(2)}\] (134d) \[=\sum_{j=1}^{2}x_{(j)},\] (134e) meaning that the condition from Eq. ( 132 ) is not fulfilled for \[k=2\] . This implies that \[k\left(\mathbf{x}\right)=1\] and \[\tau\left(\mathbf{x}\right)=x_{(1)}-1\] . Thus, we get that \[\operatorname*{sparsemax}\left(\mathbf{x}\right)_{(1)}=\max\left(0,x_{(1)}-x _{(1)}+1\right)=1\] (135) and \[\operatorname*{sparsemax}\left(\mathbf{x}\right)_{(d)}=\max\left(0,x_{(d)}-x _{(1)}+1\right)=\max\left(0,\underbrace{x_{(1)}}_{<-1}-0+1\right)=0\] (136) for \[d>1\] . 

This allows us to prove the main theorem for sparse-attention transformer LMs.

**Theorem 4.1**.: _For any n-gram LM, there exists a weakly equivalent single-layer sparse attention transformer LM with \(n-1\) heads._

Proof.: Lemma C.1 shows how individual heads of the transformer can identify the symbols in the position of interest. \(n-1\) of them can identify the entire history. The proof then follows the same reasoning as that of Theorem 3.1. 

Adapting the same proof strategy to Theorem 3.2 would naturally result in an analogous result for \(n-1\) layers and a single head.

Notice that Lemma C.1 requires different and less standard positional encodings, which are, crucially, unbounded. Constructing a sparse attention transformer with bounded positional encodings seems more difficult; the contextual representations would in that case either converge or be non-unique with \(t\to\infty\) and since the sparsemax always contracts (Martins and Astudillo, 2016, Proposition 2), attending to individual positions would be difficult. While the positional encodings and the scoring function used in the proof of Lemma C.1 are somewhat less standard than those used in Lemma B.4, similar positional encodings and the same scoring function have been used in theoretical analyses before and even in practical implementations (Perez et al., 2021).