Optimal Memorization Capacity of Transformers

## Abstract

Abstract Recent research in the field of machine learning has increasingly focused on the memorization capacity of Transformers, but how efficient they are is not yet well understood.
We demonstrate that Transformers can memorize labels with O ~ ​ ( N ) ~ 𝑂 𝑁 \tilde{O}(\sqrt{N}) parameters in a next-token prediction setting for N 𝑁 N input sequences of length n 𝑛 n , which is proved to be optimal up to logarithmic factors.
This indicates that Transformers can efficiently perform memorization with little influence from the input length n 𝑛 n owing to the benefit of parameter sharing.
We also analyze the memorization capacity in the sequence-to-sequence setting, and find that O ~ ​ ( n ​ N ) ~ 𝑂 𝑛 𝑁 \tilde{O}(\sqrt{nN}) parameters are not only sufficient, but also necessary at least for Transformers with hardmax.
These results suggest that while self-attention mechanisms can efficiently identify input sequences, the feed-forward network becomes a bottleneck when associating a label to each token.

## 1 Introduction

In recent years, the Transformer architecture (Vaswani et al., ) has played a pivotal role in the field of machine learning, becoming indispensable for a variety of models in the community.
In addition to the original breakthroughs in natural language processing, such as the GPT series (Brown et al., ; Radford et al., , ), it has been observed that in numerous applications, higher accuracy can be achieved by replacing existing models with Transformers.
Specifically, models such as the Vision Transformer (Dosovitskiy et al., ) in image processing and the Diffusion Transformer (Peebles & Xie, ) in generative tasks have demonstrated exceptional performances in a wide variety of tasks.
These examples demonstrate how effective and versatile Transformers are for a diverse range of purposes.

Although the high performance of Transformers has led to their widespread use in practice, there are ongoing attempts to theoretically analyze what exactly contributes to their superior performance.
In particular, one important aspect of Transformers is their representational capabilities.
Previous studies have explored from a variety of angles why Transformers have high expressive capacity and can memorize vast amounts of data (Edelman et al., ; Gurevych et al., ; Takakura & Suzuki, ).
For example, it has been shown that Transformers are universal approximators (having the ability to approximate arbitrary functions) (Yun et al., ) or that a particular Transformer configuration can memorize a given set of data (Kim et al., ; Kajitsuka & Sato, ; Mahdavi et al., ; Madden et al., ).

Nevertheless, while various studies have suggested that Transformers are indeed capable of memorizing data, our understanding of how efficiently they can do so remains limited. Specifically, it is not yet fully clear how certain characteristics of Transformers, such as parameter sharing, influence the reduction of model parameters and overall efficiency with respect to their memorization capacity, the minimum size of networks required for memorizing any sequence of a given number of data.

There are several key advantages to investigating whether a Transformer can efficiently memorize data, such as the possibility of gaining a better understanding of Transformer’s strengths and providing useful insights for model design and selection.
In addition, knowledge of memorization efficiency can provide important information for evaluating generalization error (Belkin et al., ; Nakkiran et al., ).
Alternatively, if it turns out that Transformers do not offer a significant efficiency advantage over feed-forward networks, it may suggest that currently widely used Transformers may in fact be substitutable for feed-forward networks.

This paper investigates the efficiency of Transformers in achieving data memorization by analyzing the necessary and sufficient model complexity for this task.
To be more precise, we establish both upper and lower bounds on the number of parameters needed for memorization in the next-token prediction setting and demonstrate that they are of the same order up to logarithmic factors, thereby showing that Transformers can achieve data memorization with nearly optimal efficiency.

Furthermore, the upper bound on memorization capacity in the next-token prediction setting can be naturally extended to the sequence-to-sequence setting.
This upper bound is also proved to be optimal in the sequence-to-sequence setting, at least for Transformers with the hardmax function.

## 2 Related Work

### Memorization capacity

Research on memorization capacity began at least as late as the 1960s (Cover, ; Nilsson, ; Minsky & Papert, ).
Specifically, Nilsson () showed that one-hidden-layer neural networks with $N-1$ nodes is able to compute any label assignments for $N$ data points.
Later, (Baum, ) exhibited that $\lceil N/d\rceil$ neurons are sufficient for one-hidden-layer neural networks with threshold units to memorize any set of $N$ input-label pairs with the input dimension $d$, and Huang & Babri (); Zhang et al. () extended the results to more general activation functions.

The analysis of memorization capacity is closely linked to the concept of the Vapnik-Chervonenkis (VC) dimension.
While the memorization capacity of a model refers to the minimum size of the model required for memorizing any tuple of $N$ input-label pairs for some $N$, the VC dimension considers whether the model is capable of shattering, that is, memorizing any possible label assignments for some set of $N$ input points, which in turn provides a lower bound on the memorization capacity.
For example, Goldberg & Jerrum () estimated that the VC dimension of a feed-forward network with ReLU activation functions and $W$ parameters is at most $O(W^{2})$ by reducing the network to a boolean formula.
From this upper bound, it can be inferred that a feed-forward network with ReLU activation functions requires at least $\Omega(\sqrt{N})$ parameters to memorize arbitrary $N$ data points.
Bartlett et al. () further refined this analysis by examining the behavior of the network as a function of its parameters and analyzing it layer by layer, and demonstrated that the VC dimension of a ReLU network with width $W$ and depth $L$ is $O(WL\log W)$.

Remarkably, Park et al. () proposed a construction method under the assumption that the data points are separated by at least $\delta$, showing that a feed-forward network using sigmoid or ReLU activation functions with a sub-linear parameter order $O(N^{2/3}+\log\delta)$ can memorize $N$ data points.
Later, Vardi et al. () demonstrated that, under similar assumptions, a ReLU network with $O(\sqrt{N\log N})$ parameters suffices for memorizing arbitrary $N$ data points.
This result is optimal up to logarithmic factors, as it matches the lower bound $\Omega(\sqrt{N})$ implied by the VC dimension previously discussed.
Note that the assumption that data points are well separated is crucial to achieve sub-linear memorization capacity; in fact, it has been shown that at least $(N-1)/2$ parameters are required to memorize arbitrary $N$ distinct data points without such separation (Sontag, ).
Additionally, Siegel () proved that $\Omega(N)$ parameters are necessary for memorizing $N$ data points when the separation $\delta$ between data points is exponentially small with respect to $N$.

Memorization capacity is not only theoretically intriguing but also practically significant.
As the model size increases, classical learning theory predicts that the training error decreases while the generalization error follows a U-shaped curve.
However, recent observations of the double descent phenomenon (Belkin et al., ; Nakkiran et al., ) revealed that after achieving zero training loss, the generalization error begins to decrease again. Analyzing memorization capacity helps identify the critical model size at which this shift occurs, providing valuable insights into the dynamics of model performance.

**Table 1: Comparisons between our results and related work. The “1-layer” column in this table indicates which studies focus on single-layer Transformers. Additionally, the variable $\omega$ in the bounds presented by Madden et al. () represents the vocabulary size, or the number of distinct word vectors that appear in input sequences.**
| Paper | Setting | Input | $1$-layer | Upper bound | Lower bound |
| --- | --- | --- | --- | --- | --- |
| Kim et al. () | seq-to-seq | token-wise<br>$(r,\delta)$-separated | no | $\tilde{O}(n+\sqrt{nN})$ | - |
| Mahdavi et al. () | next-token like | linearly<br>independent | yes | $O(d^{2}N/n)$ | - |
| Kajitsuka & Sato () | seq-to-seq | token-wise<br>$(r,\delta)$-separated | yes | $O(d(2nN+d))$ | - |
| Madden et al. () | next-token | with positional<br>encoding | yes | $O(\omega N)$ | $\Omega(\omega N)$ |
| Ours | next-token | token-wise<br>$(r,\delta)$-separated | no | $\tilde{O}(\sqrt{N})$ | $\Omega(\sqrt{N})$ |
| \cdashline2-6 | seq-to-seq | token-wise<br>$(r,\delta)$-separated | no | $\tilde{O}(\sqrt{nN})$ | $\Omega(\sqrt{nN})$ |

### Expressivity of Transformers

One of the foundational studies on the representation power of Transformers is the work by Yun et al. (), who demonstrated that Transformers are universal approximators.
Their proof already incorporates the idea of constructing a contextual mapping from data points to contexts and linking these context ids to labels.
Kim et al. (), whose work is most closely related to our work, improved their contextual mapping approach and demonstrated that this mapping, constructed using $2n$ layers of self-attention for $N$ input sequences of length $n$, allows for memorization with $\tilde{O}(n+\sqrt{nN})$ parameters under the same assumption that data points are well separated as in Park et al. (); Vardi et al. ().
Later, Kajitsuka & Sato () showed that a single-layer, single-head Transformer already possesses memorization capacity under the same assumption, while self-attention with hardmax does not.
In contrast to the studies mentioned above, Mahdavi et al. () demonstrated that under the assumption that data points are linearly independent, a multi-head attention with $H$ heads and embedding dimension $d>n$ can memorize $\Omega(Hn)$ data points in a next-token prediction like setting.
Madden et al. () proved upper and lower bounds on the memorization capacity of one-layer Transformers with parameters of infinite precision in the next-token prediction setting.
Chen & Zou () investigated the behavior of Transformers with varying depths, and specifically demonstrated that a single-layer Transformer can achieve memorization if input sequences are sufficiently zero-padded.
However, they noted that their objective was not to explore efficient constructions.
The comparisons between our results and related work are summarized in Table [1](#S2.T1).
Note that all the papers listed here that investigate single-layer Transformers assume either infinite parameter precision or do not consider the bit-length required to represent parameters.

In addition to memorization capacity, there are studies highlighting other perspectives on Transformers, including their function approximation capacity (Gurevych et al., ; Takakura & Suzuki, ; Jiang & Li, ), and their ability to efficiently represent sparse functions (Edelman et al., ; Bhattamishra et al., ; Sanford et al., ; Trauger & Tewari, ; Wang et al., ).

## 3 Preliminaries

### 3.1 Notation

We denote vectors and matrices by bold lowercase and uppercase letters, respectively.
Given a vector ${\bm{v}}$, we denote its $i$-th element as $v_{i}$.
Given a matrix ${\bm{A}}$, we denote its $i$-th row as ${\bm{A}}_{i,:}$, its $j$-th column as ${\bm{A}}_{:,j}$ and the element at position $(i,j)$ as $A_{i,j}$.
For a natural number $m\in\mathbb{N}$, we use $[m]$ to denote the set $\{1,\dots,m\}$.
In the context of the self-attention mechanism, we use $\sigma_{S}$ to represent the column-wise softmax function.
Specifically, for a matrix ${\bm{A}}\in\mathbb{R}^{a\times b}$, $\sigma_{S}[{\bm{A}}]\in\mathbb{R}^{a\times b}$ is calculated by $\sigma_{S}[{\bm{A}}]_{i,j}:=\exp(A_{i,j})/\sum_{k=1}^{a}\exp(A_{k,j})$.
Likewise, we use $\sigma_{H}$ to denote the column-wise hardmax function.
Note that if there are multiple values in a column, its outputs are normalized so that they sum up to $1$.
Mathematically, for a matrix ${\bm{A}}\in\mathbb{R}$, $\sigma_{H}[{\bm{A}}]\in\mathbb{R}^{a\times b}$ is calculated as follows.

$$ $\displaystyle\sigma_{H}[{\bm{A}}]_{i,j}=\begin{cases}1/|I_{j}|&\text{if }A_{i,j}=\max_{k}A_{k,j},\\ 0&\text{otherwise},\end{cases}$ (1) $$

where $I_{j}:=\operatorname*{arg\,max}_{k}A_{k,j}:=\{k^{\prime}\in[a]\mid A_{k^{\prime},j}=\max_{k}A_{k,j}\}$ for any $j\in[b]$.
We use $\sigma_{R}$ to denote the ReLU activation function, that is, $\sigma_{R}[x]:=\max(0,x)$.
Unlike $\sigma_{S}$ and $\sigma_{H}$, $\sigma_{R}$ is always applied element-wise, regardless of whether the input is a vector or a matrix.
For any natural number $x\in\mathbb{N}$, $\operatorname{BIN}_{i:j}(x)\in\mathbb{N}$ represents the sequence of bits from the $i$-th bit to the $j$-th bit (counting from the left) of $x$, interpreted as a natural number.
For a vector ${\bm{v}}\in\mathbb{R}^{a}$, the $L^{2}$ norm of ${\bm{v}}$ is denoted by $\|{\bm{v}}\|_{2}:=\sum_{i=1}^{a}v_{i}^{2}$.
We use standard asymptotic notation.
Specifically, $f(n)=O(g(n))$ indicates that the function $f$ grows at most as fast as $g$ for sufficiently large $n$, and $f(n)=\tilde{O}(g(n))$ represents that $f$ grows at most as fast as $g$, up to logarithmic factors.
Likewise, $f(n)=\Omega(g(n))$ means that the function $f$ grows at least as fast as $g$ for sufficiently large $n$.
$f(n)\lesssim g(n)$ means that there exists a positive constant $c$ such that $f(n)\leq cg(n)$ holds.

In this paper, we basically use $n$ to denote the length of an input sequence, $N$ to denote the number of input sequences, $C$ to denote the number of classes, and $d$ to denote the dimensionality of each token.
Additionally, index $i$ is typically used to refer to the position of input sequences, while index $k$ is used to refer to the position of the token within an input sequence.

### 3.2 Transformer block

In this subsection, we introduce the architecture of Transformers (Vaswani et al., ).
We basically follow the notations by Kim et al. ().
Transformers are defined by stacking multiple Transformer blocks, each of which consists of a self-attention layer and a feed-forward layer.

Self-attention layer:
Given an input sequence ${\bm{Z}}\in\mathbb{R}^{m\times n}$, the output of a self-attention layer $\mathcal{F}^{(\mathrm{SA})}_{l}:\mathbb{R}^{m\times n}\to\mathbb{R}^{m\times n}$ at block $l\in[L]$ is calculated by

$$ $\displaystyle\mathcal{F}^{(\mathrm{SA})}_{l}\left({\bm{Z}}\right):={\bm{Z}}+\sum_{h=1}^{H}{\bm{W}}^{(O)}_{hl}{\bm{W}}^{(V)}_{hl}{\bm{Z}}\sigma_{S}\left[\left({\bm{W}}^{(K)}_{hl}{\bm{Z}}\right)^{\top}\left({\bm{W}}^{(Q)}_{hl}{\bm{Z}}\right)\right]\in\mathbb{R}^{m\times n},$ (2) $$

where ${\bm{W}}^{(V)}_{hl},\ {\bm{W}}^{(K)}_{hl},\ {\bm{W}}^{(Q)}_{hl}\in\mathbb{R}^{s\times m}$ and ${\bm{W}}^{(O)}_{hl}\in\mathbb{R}^{m\times s}$ are value, key, query and projection matrices at head $h\in[H]$ with head size $s$, respectively.

Feed-forward layer:
The output ${\bm{H}}\in\mathbb{R}^{m\times n}$ of the self-attention layer at block $l$ is then passed to the feed-forward layer, which performs the following token-wise operation:

$$ $\displaystyle\mathcal{F}^{(\mathrm{FF})}_{l}\left({\bm{H}}\right)_{:,k}:={\bm{H}}_{:,k}+{\bm{W}}^{(2)}_{l}\sigma_{R}\left[{\bm{W}}^{(1)}_{l}{\bm{H}}_{:,k}+{\bm{b}}^{(1)}_{l}\right]+{\bm{b}}^{(2)}_{l}\in\mathbb{R}^{m}\quad(k\in[n]),$ (3) $$

where ${\bm{W}}^{(1)}_{l}\in\mathbb{R}^{q\times m}$ and ${\bm{W}}^{(2)}_{l}\in\mathbb{R}^{m\times q}$ are weight matrices with hidden dimension $q$, and ${\bm{b}}^{(1)}_{l}\in\mathbb{R}^{q}$ and ${\bm{b}}^{(2)}_{l}\in\mathbb{R}^{m}$ are bias terms.

Using the self-attention layer and the feed-forward layer, the Transformer block $\mathcal{F}_{l}:\mathbb{R}^{m\times n}\to\mathbb{R}^{m\times n}$ at block $l\in[L]$ is defined as a composition of these two layers, that is, $\mathcal{F}_{l}:=\mathcal{F}^{(\mathrm{FF})}_{l}\circ\mathcal{F}^{(\mathrm{SA})}_{l}$, and the whole architecture of the Transformer $\mathcal{N}:\mathbb{R}^{d\times n}\to\mathbb{R}^{1\times n}$ is expressed by

$$ $\displaystyle\mathcal{N}:=\mathcal{E}_{\mathrm{out}}\circ\mathcal{F}_{L}\circ\cdots\circ\mathcal{F}_{1}\circ\mathcal{E}_{\mathrm{in}},$ (4) $$

with $d_{0}:=d$, and $\mathcal{E}_{\mathrm{in}}:\mathbb{R}^{d\times n}\to\mathbb{R}^{m\times n}$ and $\mathcal{E}_{\mathrm{out}}:\mathbb{R}^{m\times n}\to\mathbb{R}^{1\times n}$ are token-wise linear mappings.

In a Transformer, the width is determined by the combination of self-attention layers and feed-forward layers.
According to the definition proposed by Kim et al. (), the width of the Transformer model is defined as $\max(m,sH,q)$.
We define the depth of a Transformer by the number of blocks $L$.

###### Remark 3.1 .

The use of in/out token-wise linear mappings comes from the fact that Transformer blocks by definition have the same input and output dimensions.
The token-wise linear mappings can be removed at the cost of a linear dependence of the number of parameters required for memorization on the embedding dimension $d$.

###### Remark 3.1 .

### 3.3 Bit complexity

In this paper, we consider not only the number of parameters but also the number of bits required to represent the model.
Specifically, we adopt the definition of bit complexity proposed by Vardi et al. ().
According to this definition, the bit complexity of a parameter is defined as the number of bits needed to represent that parameter.
The bit complexity of a model is then defined as the maximum bit complexity among its individual parameters.

It is important to note that by multiplying the bit complexity of the model by the number of parameters, we can estimate the total number of bits required to represent the entire model.

## 4 Memorization Capacity of Transformers

In this section, we state the main theorems of this paper regarding the optimal memorization capacity of Transformers.
Section [4.1](#S4.SS1) defines the memorization capacity of Transformers and discuss the main challenge behind this concept.
In Sections [4.2](#S4.SS2) and [4.3](#S4.SS3), we provide upper and lower bounds on the number of parameters required for Transformers to achieve memorization in the next-token prediction setting and the seq-to-seq prediction setting, respectively.

### 4.1 Problem setting

The aim of this study is to analyze the memorization capacity of Transformers.
Informally, memorization capacity refers to the minimum size of a model that can memorize a specific number of arbitrary data points.
To be more precise, let $\mathcal{X}$ and $\mathcal{Y}$ be input space and output space, respectively.
Then, given $N$ input-label pairs $({\bm{X}}^{(1)},y^{(1)}),\dots,({\bm{X}}^{(N)},y^{(N)})\in\mathcal{X}\times\mathcal{Y}$, we are interested in the model complexity of a model $f:\mathcal{X}\to\mathcal{Y}$ such that $f({\bm{X}}^{(i)})=y^{(i)}$ holds for any $i\in[N]$.
In the case of Transformers, the input space $\mathcal{X}$ consists of input sequences made up of $n$ tokens, each of which is a $d$-dimensional vector.
Hence, we define the input space $\mathcal{X}$ as $\mathcal{X}:=\mathbb{R}^{d\times n}$.

Without any assumptions on the input data, it has been shown by Siegel (), that a linear order of parameters is required to memorize arbitrary $N$ data points.
To achieve a sub-linear memorization capacity, in this paper, we assume that the data points are well separated.
This concept is formalized as token-wise $(r,\delta)$-separatedness for Transformers (Kim et al., ; Kajitsuka & Sato, ).

###### Assumption 4.1 (Token-wise separation) .

Let $({\bm{X}}^{(1)},\dots,{\bm{X}}^{(N)})\in\mathbb{R}^{d\times n}$ be $N$ input sequences, each of which consists of $n$ word vectors with its dimension $d$.
Then, we say that $({\bm{X}}^{(1)},\dots,{\bm{X}}^{(N)})$ are token-wise $\mathbf{(r,\bm{\delta})}$-separated for some $r,\delta>0$ if the following two conditions are satisfied:

- 1.
for every $i\in[N]$ and $k\in[n]$, $\|{\bm{X}}^{(i)}_{:,k}\|_{2}\leq r$ holds.
- 2.
for every $i,j\in[N]$ and $k,l\in[n]$, either ${\bm{X}}^{(i)}_{:,k}={\bm{X}}^{(i)}_{:,l}$ or $\|{\bm{X}}^{(i)}_{:,k}-{\bm{X}}^{(i)}_{:,l}\|_{2}\geq\delta$ holds.

The notion of token-wise $(r,\delta)$-separatedness ensures that the word vectors appearing in the input sequences have an $L^{2}$ norm of at most $r$, and are separated by at least $\delta$ in $L^{2}$ norm from each other.

The main difficulty of memorization with Transformers, compared with feed-forward networks, lies in the fact that tokens with identical values do not necessarily correspond to the same label.
Instead, it is crucial to capture the context in which each token appears within the entire input sequence.
In Transformers, while feed-forward layers operate on individual tokens, self-attention layers are the only place that enables interactions between tokens within the input sequence.
Therefore, the central question we consider in this paper is:

$$ how efficiently can self-attention layers capture the context of tokens? $$

To explore this issue, we analyze both upper and lower bounds on the number of parameters required for memorization with Transformers in two settings: next-token prediction and sequence-to-sequence prediction.

###### Assumption 4.1 (Token-wise separation) .

### 4.2 Next-token prediction setting

#### 4.2.1 Upper bound

First, given $N$ input sequences of length $n$, consider the problem setting in which a Transformer memorizes labels corresponding to the $n$-th token of all input sequences.
We call this task next-token prediction setting.
In this problem setting, how many parameters does a Transformer architecture require?
Surprisingly, $\tilde{O}(\sqrt{N})$ is sufficient, that is, the input length $n$ has almost no effect on the number of parameters required for memorization, as the following theorem states.

In the next theorem, $\mathcal{F}^{(\mathrm{FF})}_{1}$ and $\mathcal{F}^{(\mathrm{FF})}_{2}$ represent feed-forward networks of arbitrary depth, unlike [eq. 3](#S3.E3), which is limited to two layers.
Note that deep feed-forward networks can also be implemented with standard Transformers, by setting the projection matrix of the self-attention layer in each block to zero.
Furthermore, the assumption of consistency on labels in [Theorem 4.1](#S4.Thmthm1) is a necessary requirement to perform memorization with a Transformer, due to its permutation equivariance.

###### Theorem 4.1 (Next-token prediction) .

Let $({\bm{X}}^{(1)},y^{(1)}),\dots,({\bm{X}}^{(N)},y^{(N)})\in\mathbb{R}^{d\times n}\times[C]$ be a sequence of input-label pairs such that

- 1.
$({\bm{X}}^{(1)},y^{(1)}),\dots,({\bm{X}}^{(N)},y^{(N)})$ are consistently labeled, in the sense that for any $i,j\in[N]$, we have $y^{(i)}=y^{(j)}$ if
$\displaystyle{\bm{X}}_{:,n}^{(i)}={\bm{X}}_{:,n}^{(j)}\quad\text{and}\quad{\bm{X}}^{(i)}={\bm{X}}^{(j)}\text{ up to permutations}.$
(5)
- 2.
$({\bm{X}}^{(1)},\dots,{\bm{X}}^{(N)})$ are token-wise $(r,\delta)$-separated for some $r\geq 1$ and $0<\delta\leq 1$.

Then, there exists a Transformer $\mathcal{N}:\mathbb{R}^{d\times n}\to\mathbb{R}^{n}$
with width $14$ and depth $\tilde{O}(\sqrt{N})$
that memorizes the dataset, that is,

$$ $\displaystyle\mathcal{N}\left({\bm{X}}^{(i)}\right)_{n}=\mathcal{E}_{\mathrm{out}}\circ\mathcal{F}^{(\mathrm{FF})}_{2}\circ\mathcal{F}^{(\mathrm{SA})}\circ\mathcal{F}^{(\mathrm{FF})}_{1}\circ\mathcal{E}_{\mathrm{in}}\left({\bm{X}}^{(i)}\right)_{n}=y^{(i)}$ (6) $$

holds for every $i\in[N]$, as long as $n,C,r\delta^{-1}=N^{O(1)}$ as $N\to\infty$.

###### Remark 4.1 (Deep sets) .

In fact, [Theorem 4.1](#S4.Thmthm1) can be extended to Deep Sets (Zaheer et al., ), which is a popular architecture to model a mapping from sets to labels.
For details on this result, see Appendix LABEL:sec:memorization_deep_sets.

###### Remark 4.2 (Embedding layer) .

A similar result holds for a Transformer with an embedding layer.
However, in this case, the presence of an embedding layer introduces a dependency on the size of the vocabulary, which may result in a non-optimal order of parameters in the worst-case scenario.

###### Remark 4.3 (Dependence on $d$ ) .

The Transformer architecture defined by [eq. 4](#S3.E4) includes token-wise linear mappings $\mathcal{E}_{\mathrm{in}}:\mathbb{R}^{d}\to\mathbb{R}^{m}$ and $\mathcal{E}_{\mathrm{out}}:\mathbb{R}^{m}\to\mathbb{R}^{d}$, leading to $\tilde{O}(d+\sqrt{N})$ parameters for a Transformer with depth $\tilde{O}(\sqrt{N})$ and width $14$.
As noted by Vardi et al. () and Kim et al. (), this dependence on the dimension $d$ is unavoidable to preserve the information of the input tokens.

[Theorem 4.1](#S4.Thmthm1) demonstrates that as long as the dimension $d$ is of the order $d=\tilde{O}(\sqrt{N})$, the Transformer with a single self-attention layer can memorize $N$ input sequences and their labels for next-token prediction with $\tilde{O}(\sqrt{N})$ parameters, showing negligible dependence on the input length $n$.
In contrast, to accomplish the same task with a feed-forward network, it is necessary to use $d\times n$ parameters to retain the information of the input sequence in $\mathbb{R}^{d\times n}$.
This illustrates a significant efficiency advantage of Transformers over feed-forward networks, thanks to parameter sharing.

###### Theorem 4.1 (Next-token prediction) .

###### Remark 4.1 (Deep sets) .

###### Remark 4.2 (Embedding layer) .

###### Remark 4.3 (Dependence on d 𝑑 d ) .

#### 4.2.2 Proof outline of Theorem 4.1

Here we provide an outline of the proof of [Theorem 4.1](#S4.Thmthm1).
See Appendix LABEL:sec:appendix_next_token_upper_bound for its full proof.

The proof strategy is to construct a contextual mapping as in Yun et al. (), Kim et al. () and Kajitsuka & Sato (), and then construct a mapping from the context id to the label.
Here, a contextual mapping is a function used to distinguish tokens in each input sequence with the following properties:

###### Definition 4.1 (Contextual mapping) .

Let ${\bm{X}}^{(1)},\dots,{\bm{X}}^{(N)}\in\mathbb{R}^{d\times n}$ be input sequences.
Then, a map $\mathcal{CM}:\mathbb{R}^{d\times n}\to\mathbb{R}^{n}$ is called an $\mathbf{(r,\bm{\delta})}$-contextual mapping if the following two conditions hold:

- 1.
For any $i\in[N]$ and $k\in[n]$, $\left|\mathcal{CM}({\bm{X}}^{(i)})_{k}\right|\leq r$ holds.
- 2.
For any $i,j\in[N]$ and $k,l\in[n]$ such that ${\bm{X}}^{(i)}_{:,k}\neq{\bm{X}}^{(j)}_{:,l}$ or ${\bm{X}}^{(i)}\neq{\bm{X}}^{(j)}$ up to permutations, $\left|\mathcal{CM}({\bm{X}}^{(i)})_{k}-\mathcal{CM}({\bm{X}}^{(j)})_{l}\right|>\delta$
holds.

In particular, $\mathcal{CM}({\bm{X}}^{(i)})_{k}$ is called a context id of the $k$-th token in ${\bm{X}}^{(i)}$.

Intuitively, the two conditions above ensure that the contextual mapping is injective from “distinct” data points to scalars.
If such a mapping can be constructed, then a mapping from context ids to labels can be realized using a feed-forward network with $\tilde{O}(\sqrt{N})$ parameters, as shown by Vardi et al. ().
In particular, if we can associate each distinct input sequence with a unique value, referred to as a sequence id, then the context id of, for example, the $k$-th token in ${\bm{X}}^{(i)}$ can be constructed from the sequence id of ${\bm{X}}^{(i)}$ and the token vector ${\bm{X}}^{(i)}_{:,k}$.
Therefore, the primary focus of our proof is on how to construct a mapping from each input sequence to its sequence id using a feed-forward network and a single self-attention layer.

From a high-level perspective, our goal is to construct a feed-forward network $\phi:\mathbb{R}^{d}\to\mathbb{R}$ with $\tilde{O}(\sqrt{N})$ parameters such that the sums

$$ $\displaystyle\sum_{k=1}^{n}\phi({\bm{X}}^{(1)}_{:,k}),\dots,\sum_{k=1}^{n}\phi({\bm{X}}^{(N)}_{:,k})$ (7) $$

are well-separated (^1^11For simplicity, here we assume that the input sequences ${\bm{X}}^{(1)},\dots,{\bm{X}}^{(N)}$ are distinct up to permutations.).
The sum $\sum_{k=1}^{n}\phi({\bm{X}}^{(i)}_{:,k})\ (i\in[N])$ is then used as the sequence id of ${\bm{X}}^{(i)}$.

Crucial observations for constructing $\phi$ with $\tilde{O}(\sqrt{N})$ parameters are as follows.

- 1.
To distinguish $N$ input sequences, it is sufficient to focus on at most $N$ distinct word vectors.
More precisely, given $N$ input sequences, there are at most $N$ distinct word vectors such that the input sequences can be identified by counting occurrences of these $N$ words $A=\{{\bm{v}}_{1},\dots,{\bm{v}}_{N}\}\subset\mathbb{R}^{d}$ (LABEL:lem:restriction_multiset).
- 2.
Although a feed-forward network requires $\Omega(\sqrt{N})$ parameters to memorize $N$ data points and their labels (Goldberg & Jerrum, ), a network that outputs zero for additional data points not among $N$ data points can be constructed without significantly affecting the order of the parameter count (LABEL:lem:extension_ff_memorization).
Together with the first observation, all we need is to construct a feed-forward network $\phi:\mathbb{R}^{d}\to\mathbb{R}$ such that
$\displaystyle\sum_{k=1,{\bm{X}}^{(1)}_{:,k}\in A}^{n}\phi({\bm{X}}^{(1)}_{:,k}),\dots,\sum_{k=1,{\bm{X}}^{(N)}_{:,k}\in A}^{n}\phi({\bm{X}}^{(N)}_{:,k})$
(8)
are well-separated.
- 3.
The final key observation is that rather than directly constructing $\phi$, we first consider the high-dimensional representation.
Specifically, given arbitrary bijection $g:A\to[N]$, we can map each input sequence ${\bm{X}}^{(i)}$ to a high-dimensional vector $\tilde{{\bm{X}}}^{(i)}\in\mathbb{R}^{N}$ as follows:
$\displaystyle\tilde{{\bm{X}}}^{(i)}:=\sum_{k=1,{\bm{X}}^{(i)}_{:,k}\in A}^{n}{\bm{e}}_{g({\bm{X}}^{(i)}_{:,k})},$
(9)
where ${\bm{e}}_{g({\bm{X}}^{(i)}_{:,k})}\in\mathbb{R}^{N}$ is a one-hot vector with $1$ in the $g({\bm{X}}^{(i)}_{:,k})$-th position.
From the first observation, $\tilde{{\bm{X}}}^{(1)},\dots,\tilde{{\bm{X}}}^{(N)}$ should be distinct, and we can compress this $N$-dimensional representation to a one-dimensional one, thus enabling the construction of $\phi$ (LABEL:lem:separating_function_multiset).

To ensure that a feed-forward network with $\tilde{O}(\sqrt{N})$ parameters can indeed implement the function $\phi$, we need to carefully analyze how separated the compressed versions of the high-dimensional representations $\tilde{{\bm{X}}}^{(1)},\dots,\tilde{{\bm{X}}}^{(N)}$ are.
Detailed proof of this implementation is provided in LABEL:lem:separation_of_multisets.

###### Definition 4.1 (Contextual mapping) .

#### 4.2.3 Lower bound

In this subsection, we evaluate the minimal model complexity required for memorization with Transformers in the next-token prediction setting to determine how close [Theorem 4.1](#S4.Thmthm1) is to optimal.

First, notice that the model obtained in [Theorem 4.1](#S4.Thmthm1) is optimal, in terms of bit counts.

###### Remark 4.4 (Optimality in terms of bit counts) .

As previously discussed in [Remark 4.3](#S4.Thmrem3), the Transformer model obtained in [Theorem 4.1](#S4.Thmthm1) has $\tilde{O}(\sqrt{N})$ parameters as long as $d=\tilde{O}(\sqrt{N})$.
On the other hand, the bit complexity of the model is $\tilde{O}(\log d+\sqrt{N})$ (see the formal statement in Appendix LABEL:sec:appendix_next_token_upper_bound).
Therefore, if $d=\tilde{O}(\sqrt{N})$, the total number of bits required to represent the model is $\tilde{O}(N)$.
Given that there are $2^{N}$ possible label assignments for $N$ distinct data points with binary labels, $\tilde{O}(N)$ bits are optimal up to logarithmic factors for this setting.

Having established the optimality in terms of bit counts, we now turn to evaluating how efficient the number of parameters of the Transformer model considered in [Theorem 4.1](#S4.Thmthm1) is.
The next theorem provides a lower bound on the number of parameters required for memorization in the next-token prediction setting.

###### Theorem 4.2 (Lower bound) .

Suppose a Transformer $\mathcal{N}$ can memorize any sequence of $N$ input-label pairs that are consistently labeled in the sense of [Theorem 4.1](#S4.Thmthm1).
Then, the Transformer $\mathcal{N}$ has at least $\Omega(\sqrt{N})$ parameters.

The proof of this theorem can be found in Appendix LABEL:sec:appendix_next_token_lower_bound.
This result indicates that the Transformer model described in [Theorem 4.1](#S4.Thmthm1) is also optimal in terms of the number of parameters.
Specifically, since memorization in the next-token prediction setting requires the ability to distinguish $N$ input sequences, this result provides the following crucial insight.

A Transformer with a single layer of self-attention already possesses necessary and sufficient expressive capacity to identify input sequences.

In fact, as indicated in the proof outline in Section [4.2.2](#S4.SS2.SSS2), we only employ the self-attention layer as an averaging operation in the model obtained by [Theorem 4.1](#S4.Thmthm1).
The observation that simple averaging provides sufficient representational power has been confirmed experimentally by Yu et al. () with their PoolFormer architecture.
In this paper, we provide theoretical support by demonstrating that a Transformer with just a simple averaging operation already has optimal memorization capacity.

###### Remark 4.4 (Optimality in terms of bit counts) .

###### Theorem 4.2 (Lower bound) .

### 4.3 Sequence-to-sequence prediction setting

Next, we consider the problem setting in which each token in an input sequence is assigned some label and a Transformer memorizes them all.
We call this task a sequence-to-sequence prediction setting, or seq-to-seq prediction for short.

It is readily apparent that the seq-to-seq prediction can be regarded as rearranging the input sequence so that each token is placed at the end of the sequence, and performing next-token prediction on $nN$ input sequences obtained in this way.
From this observation, we have the following corollary from [Theorem 4.1](#S4.Thmthm1).

###### Corollary 4.1 (Seq-to-seq prediction) .

Let $({\bm{X}}^{(1)},{\bm{y}}^{(1)}),\dots,({\bm{X}}^{(N)},{\bm{y}}^{(N)})\in\mathbb{R}^{d\times n}\times[C]^{n}$ be a sequence of input-label pairs such that

- 1.
$({\bm{X}}^{(1)},{\bm{y}}^{(1)}),\dots,({\bm{X}}^{(N)},{\bm{y}}^{(N)})$ are consistently labeled, in the sense that for any $i,j\in[N]$ and $k,l\in[n]$, we have $y^{(i)}_{k}=y^{(j)}_{l}$ if
$\displaystyle{\bm{X}}_{:,k}^{(i)}={\bm{X}}_{:,l}^{(j)}\quad\text{and}\quad{\bm{X}}^{(i)}={\bm{X}}^{(j)}\text{ up to permutations}.$
(10)
- 2.
$({\bm{X}}^{(1)},\dots,{\bm{X}}^{(N)})$ are token-wise $(r,\delta)$-separated for some $r\geq 1$ and $0<\delta\leq 1$.

Then, there exists a Transformer $\mathcal{N}:\mathbb{R}^{d\times n}\to\mathbb{R}^{n}$ with width $14$ and depth $\tilde{O}(\sqrt{nN})$
that memorizes the dataset, that is,

$$ $\displaystyle\mathcal{N}\left({\bm{X}}^{(i)}\right)_{k}=\mathcal{E}_{\mathrm{out}}\circ\mathcal{F}^{(\mathrm{FF})}_{2}\circ\mathcal{F}^{(\mathrm{SA})}\circ\mathcal{F}^{(\mathrm{FF})}_{1}\circ\mathcal{E}_{\mathrm{in}}\left({\bm{X}}^{(i)}\right)_{k}=y^{(i)}_{k}$ (11) $$

holds for every $i\in[N]$ and $k\in[n]$, as long as $C,r\delta^{-1}=(nN)^{O(1)}$ as $nN\to\infty$.

###### Remark 4.5 (Sparse Transformers) .

While [Corollary 4.1](#S4.Thmcor1) demonstrates that a Transformer with a single-layer self-attention can achieve memorization in the seq-to-seq prediction setting, it inevitably requires $O(n^{2})$ computational complexity due to the self-attention mechanism.
In line with recent efforts to improve the scalability of Transformers by making attention maps sparse (Zaheer et al., ; Yun et al., ), using two self-attention layers and appending an additional token to the input sequence allows us to achieve the same behavior with an $O(n)$ connections without affecting the order of parameter counts.
This idea of aggregating global information into the additional token has gained interest in recent studies (Darcet et al., ; Wang et al., ).

This corollary shows that at least $\tilde{O}(\sqrt{nN})$ parameters with bit complexity $\tilde{O}(\sqrt{nN})$ are enough to memorize $N$ input sequences of input length $n$.
The next question is: is this order optimal for the seq-to-seq prediction setting?
As in the case of next-token prediction setting ([Remark 4.4](#S4.Thmrem4)), we can leverage a similar argument to show that this is optimal, at least in terms of bit counts.

###### Remark 4.6 (Optimality in terms of bit counts) .

If $d=\tilde{O}(\sqrt{nN})$, the construction by [Corollary 4.1](#S4.Thmcor1) uses $\tilde{O}(\sqrt{nN})$ parameters with bit complexity $\tilde{O}(\sqrt{nN})$ to memorize $N$ input sequences of input length $n$, which amounts to $\tilde{O}(nN)$ bits.
If all word vectors in input sequences are different, there are $2^{nN}$ binary label patterns.
Therefore, to memorize such patterns, the number of states of the model must be at least $2^{nN}$, which means that $\log 2^{nN}=nN$ bits is required.

Unlike the next-token prediction setting, it is challenging to analyze the optimal lower bound on the number of parameters necessary to memorize $N$ input sequences with input length $n$ for the seq-to-seq prediction setting, mainly due to the presence of the softmax function (consider it is not yet well-understood even about the optimal memorization capacity of feed-forward networks with the sigmoid function).
However, we partially answer this question by considering a Transformer that uses not the softmax function, but instead the hardmax function, often viewd as an approximation of the softmax.

More rigorously, we introduce the following self-attention layer with the hardmax function, which we call the hard attention layer. For each block $l\in[L]$ and its input ${\bm{Z}}\in\mathbb{R}^{m\times n}$, the hard attention layer at block $l$ calculates

$$ $\displaystyle\mathcal{F}^{(\mathrm{HA})}_{l}\left({\bm{Z}}\right):={\bm{Z}}+\sum_{h=1}^{H}{\bm{W}}^{(O)}_{hl}{\bm{W}}^{(V)}_{hl}{\bm{Z}}\sigma_{H}\left[\left({\bm{W}}^{(K)}_{hl}{\bm{Z}}\right)^{\top}\left({\bm{W}}^{(Q)}_{hl}{\bm{Z}}\right)\right]\in\mathbb{R}^{m\times n},$ (12) $$

where $\sigma_{H}:\mathbb{R}^{n\times n}\to[0,1]^{n\times n}$ is the column-wise hardmax function, and ${\bm{W}}^{(V)}_{hl},\ {\bm{W}}^{(K)}_{hl},\ {\bm{W}}^{(Q)}_{hl}\in\mathbb{R}^{s\times m}$ and ${\bm{W}}^{(O)}_{hl}\in\mathbb{R}^{m\times s}$ are value, key, query and projection matrices at head $h\in[H]$ with head size $s$, respectively.
It is worth noting that a simple averaging operation can also be implemented using a hard attention layer by setting key and query matrices to zero.

With this definition, we demonstrate that the number of parameters by [Corollary 4.1](#S4.Thmcor1) is actually optimal up to logarithmic factors, at least for Transformers with the hardmax function.
To state the theorem, let $W$ be the number of parameters and $\bm{\theta}\in\mathbb{R}^{W}$ be a vector of all parameters of a Transformer.
We also denote by $\mathcal{N}_{\bm{\theta}}$ the Transformer to emphasize the presence of the parameter vector $\bm{\theta}$.

###### Theorem 4.3 (Lower bound) .

Let $\mathcal{N}_{\bm{\theta}}:\mathbb{R}^{d\times n}\to\mathbb{R}^{n}$ be a Transformer defined by [eq. 4](#S3.E4) with self-attention layers replaced with hard attention layers ([eq. 12](#S4.E12)).
In addition, suppose $\mathcal{N}_{\bm{\theta}}$ can shatter a set of $N$ input sequences ${\bm{X}}^{(1)},\dots,{\bm{X}}^{(N)}\in\mathbb{R}^{d\times n}$ with $X^{(i)}_{:,k}\neq X^{(j)}_{:,l}$ for any $i,j\in[N]$ and $k,l\in[n]\ (k\neq l)$, in the sense that for any label assignments ${\bm{y}}^{(1)},\dots,{\bm{y}}^{(N)}\in\{0,1\}^{n}$, there is a parameter vector $\bm{\theta}\in\mathbb{R}^{W}$ such that

$$ $\displaystyle\mathcal{N}_{\bm{\theta}}({\bm{X}}^{(i)})={\bm{y}}^{(i)}$ (13) $$

for any $i\in[N]$.
Then, the Transformer has at least $W=\Omega\left(\sqrt{\frac{nN}{\log(nN)}}\right)$ parameters.

The proof of LABEL:thm:seq_to_seq_lower_bound builds on the approach used by Bartlett et al. () to evaluate a lower bound on the VC dimension of feed-forward networks.
Specifically, considering a Transformer as a function in variable its parameter vector, we partition the parameter space of the Transformer in such a way that, within each cell of this partition, the function can be expressed as a polynomial in terms of its parameters, and then evaluate the number of cells and the properties of the polynomials within those cells.

The key novelty of the proof lies in the analysis of how parameter sharing and the hardmax function affect the memorization capacity of Transformers.
Parameter sharing in Transformers allows the model to effectively behave like a network with its width scaled by the number of tokens, without actually increasing the number of parameters.
However, the proof shows that merely increasing the width by a factor of $n$ does not lead to a fundamental improvement in the memorization capacity of the Transformer.
The full proof of [Theorem 4.3](#S4.Thmthm3) can be found in Appendix LABEL:sec:appendix_seq_to_seq_lower_bound.

[Theorem 4.3](#S4.Thmthm3) demonstrates that the number of parameters in the model from [Corollary 4.1](#S4.Thmcor1) is within logarithmic factors of the optimal lower bound.
In addition, it provides another crucial insight.
As shown in the next-token prediction setting, Transformers can identify $N$ input sequences with $\tilde{O}(\sqrt{N})$ parameters, which implies that they are capable of capturing the context of each token.
Therefore, in the seq-to-seq prediction setting, the primary bottleneck is not the contextual mapping of tokens, but rather the feed-forward layers’ capacity to map this token-level contextual information to labels.

We conclude this section by leaving an open problem.
Based on [Theorem 4.3](#S4.Thmthm3), for a Transformer to memorize $N$ sequences of length $n$ with $o(\sqrt{nN})$ parameters, it is necessary to exploit the unique characteristics of the softmax function, rather than using it as an approximation of hardmax.

###### Open Problem .

Does a Transformer using the softmax function require $\Omega(\sqrt{nN})$ parameters to memorize $N$ input-label pairs $({\bm{X}}^{(1)},{\bm{y}}^{(1)}),\dots,({\bm{X}}^{(N)},{\bm{y}}^{(N)})\in\mathbb{R}^{d\times n}\times[C]^{n}$?
Alternatively, is it possible to construct a Transformer with $o(\sqrt{nN})$ parameters that can shatter arbitrary $N$ token-wise $(r,\delta)$-separated input sequences in the seq-to-seq setting?

###### Corollary 4.1 (Seq-to-seq prediction) .

###### Remark 4.5 (Sparse Transformers) .

###### Remark 4.6 (Optimality in terms of bit counts) .

###### Theorem 4.3 (Lower bound) .

###### Open Problem .

## 5 Conclusions

In this paper, we showed that in the next-token prediction setting, a Transformer with $\tilde{O}(\sqrt{N})$ parameters can memorize $N$ input-label pairs, which we showed to be optimal up to logarithmic factors.
This result indicates that Transformers can perform next-token prediction with almost no impact from the length of the input sequence.
Notably, its proof indicates that even a single self-attention layer used as an averaging operation possesses sufficient expressive power to distinguish between input sequences efficiently.
Furthermore, we demonstrated that in the seq-to-seq prediction setting, $\tilde{O}(\sqrt{nN})$ parameters are also sufficient, and we proved that this is optimal up to logarithmic factors, at least for Transformers using hardmax.
This finding highlights that the main bottleneck in seq-to-seq prediction tasks lies in the feed-forward layers’ capacity to map each token to the corresponding label.

Given that a single layer of self-attention as an averaging operation suffices for distinguishing input sequences from a memorization perspective, our results suggest that the advantages of using self-attention might rather lie in the perspectives of optimization and generalization.
Another future direction is to prove the optimal memorization capacity for other equivariant models, as we believe that the proof techniques employed in our approach can be applied to them as well.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: See weaknesses above.
soundness: 3
strengths: - This paper is well-motivated and clearly presented, making it easy to follow.
- This paper is technically solid; the construction of a single-layer Transformer that can identify identical sequence ids from token-wise \((r,\delta)\)-separated sequences with near-optimal parameter order is novel. Additionally, the authors provide a parameter lower bound for the seq2seq case when the self-attention layer uses hard attention.
- Furthermore, the paper considers bit complexity, which is often overlooked in related work.
summary: The paper focuses on the memorization capacity of Transformers in next-token prediction and sequence-to-sequence settings, and shows the scaling in the number of input sequences and their lengths.
weaknesses: - Token-wise \((r,\delta)\)-separated sequences may not be an expressive enough model to fully capture the memorization capacity of Transformers. It appears that the authors construct a Transformer that captures contextual information by 'counting' the occurrences of tokens in a subset. However, Transformers typically capture contextual information from the mutual relationships among tokens via the attention mechanism. 
- In the proof of Theorem 4.1, the attention layer is set to uniform attention, functioning as a column-wise average operator. This over-simplification of the attention mechanism limits the ability of the work to reveal how attention contributes to memorization, as the attention mechanism should gather context-aware representations of tokens based on their relationships with others in the sequence. Therefore, how do the authors understand the role that the attention mechanism plays in the memorization capacity of the Transformer within their framework?
 - The proof for the parameter lower bound in seq2seq cases relies on hardmax activation, deviating from the real-world scenario where the attention layer adopts softmax activation. Is it possible to extend the proof to the case with softmax activation?

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: - If I understand correctly, neither of the setting in the paper correspond directly to the standard next-token prediction training paradigm in which the transformer needs to simultaneously predict the $i$th token given $X_{:,< i}$ for all $i$. Corollary 4.1 does not directly apply because of the varying sequence lengths, but it suggests a construction with $\tilde O(N)$ parameters+bit complexity. However, it's not clear if this is optimal since Theorem 4.3 is lower bound against the more general sequence-to-sequence task?

- Could you explain footnote 2 (line 1025)? I assume it's not possible to remove the $\log(r/\delta)$ dependence so is the goal to remove the $\log d$ or $\log C$ dependencies? If so, what is the idea?

- Is it possible to achieve the full tradeoff between bit-complexity and number of parameters? E.g. $N^{\alpha}$ parameters and $N^{1-\alpha}$ bit complexity for any $\alpha \in [1/2,1]$?
soundness: 3
strengths: - The proof sketch of Theorem 4.1 is clear and well motivated
- In the next-token setting, the paper proves matching upper and lower bounds, and in the sequence to sequence setting they match for hardmax transformers
- The authors provide clear comparisons with prior work throughout the paper
summary: This paper studies the number of parameters and the bit complexity required by a transformer to memorize a dataset of size $N$. In the context of next token prediction, the paper proves that for any dataset with well separated tokens, there exists a transformer with $\tilde O(\sqrt{N})$ parameters+bit-complexity which interpolates it, along with a matching lower bound that $\Omega(\sqrt{N})$ parameters are necessary. In the sequence to sequence setting, the paper proves a similar upper bound with $\tilde O(\sqrt{nN})$ parameters+bit complexity where $n$ is the sequence length, and they prove a matching lower bound for hardmax transformers.
weaknesses: - **bit complexity:** As the authors point out, the transformer needs at least $\Omega(N)$ bits, so to succeed with $\sqrt{N}$ parameters it needs a bit complexity of $\sqrt{N}$. However, this is very unrealistic as these models are trained in finite precisions, most often 16 bit precision. It is therefore unclear if these results say much about transformers ability to memorize in practice. This also affects the motivation, for example I would expected double descent (lines 102-107) to occur at $\Theta(N)$ parameters, not $\Theta(\sqrt{N})$ parameters.

- **experiments:** Related to the above point, the paper would benefit from some simple experiments that demonstrate whether such constructions are learnable or whether the empirical limit is $\Theta(N)$ or $\Theta(\sqrt{N})$ parameters. It could also be interesting to test this with various levels of precisions (FP16/32/64) to see whether higher precision really does let the transformer memorize with fewer parameters.

- minor typo on line 462: viewd -> viewed

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: Following the weaknesses part, I have the following questions.

1.  What is the significance of the logarithmic gap in the paper. For instance, do authors believe this gap can be closed with further analysis, or is it likely to be fundamental?

2.1 Is there any examples of real-world datasets that satisfy or approximately satisfy this assumption?
2.2 Can authors discuss how the results might change if this assumption is relaxed?
soundness: 3
strengths: 1. The results seem promising and strong, because both upper and lower bounds are provided.
2. The studied problem and method provide interesting insight into transformers.
summary: This paper studied the memorization capacity of transformers in two tasks, namely next-token prediction and sequence-to-sequence prediction. The paper showed that, by construction, a transformer with $\tilde{O}(\sqrt{N}log n)$ layers can memorize $N$ data points of length $n$. The paper further showed that a transformer with $\tilde{O}(\sqrt{nN})$ layers can memorize $N$ data points in the sequence-to-sequence prediction problem. A lower bound is provided to demonstrate the near-optimality of the upper bound.
weaknesses: 1. The title wording could be improved for accuracy, I would suggest 'near-optimal...' instead of 'optimal...' since there is a logarithm term mismatch between lower bound and upper bound.
2. The paper can have more discussion on the separatedness assumption. For example, how realistic is the assumption?

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. I don't completely follow the reasoning behind the statement: "the primary bottleneck is not the contextual mapping of tokens, but rather the feed-forward layers’ capacity to map this token-level contextual information to labels." Clarification on this point would be helpful.

2. If the feed-forward layer is identified as the bottleneck, is there an approach to mitigate this through architectural changes or modifications?

3. Are the data assumptions too restrictive? It feels surprising that the results hold if the attention layer primarily averages, especially when compared to findings in studies such as (a) and (b), despite the differences in setting.

4. Would it be possible to include experiments, even toy examples, to support some of the paper’s claims? Such experiments could significantly enhance the paper's impact.

I am open to increasing the score for the paper based on the authors' responses and feedback from other reviewers.

References:

a) Rajaraman, Nived, et al. "Transformers on Markov data: Constant depth suffices." arXiv preprint arXiv:2407.17686 (2024).  

b) Makkuva, Ashok Vardhan, et al. "Local to Global: Learning Dynamics and Effect of Initialization for Transformers." arXiv preprint arXiv:2406.03072 (2024).
soundness: 3
strengths: 1. The paper addresses the important problem of quantifying the memorization capacity of transformers.
2. The authors present theorems that cover a wide range of tasks for which transformers are used, making the analysis comprehensive.
summary: The paper examines the memorization capacity of Transformers. The authors claim that Transformers can memorize labels with $O(\sqrt{nN})$ parameters in next-token prediction for sequences. They argue that parameter sharing reduces input length sensitivity. In sequence-to-sequence tasks, $ O(\sqrt{nN})$ parameters are shown to be necessary and sufficient for hardmax transformers. While self-attention identifies input sequences effectively, the feed-forward network seems to be a limiting factor.
weaknesses: 1. The sketches of the proofs in the paper were somewhat difficult to follow without referring to the full proofs. Improving the presentation could help a wider audience better understand and appreciate the work.
2. While the paper is theoretical, and experiments may not be essential, including experiments to support the claims would strengthen the work.

Meta-review: This paper studies the number of parameters and the bit complexity required by a transformer to memorize a dataset of size. In the context of next token prediction, the paper proves that for any dataset with well separated tokens, there exists a transformer optimal 
 parameters+bit-complexity (both upper and lower bounds.) In the sequence to sequence setting, the paper proves a similar upper bound with, and they prove a matching lower bound for hardmax transformers. All reviewers believe this paper is technically strong, and the AC agrees and recommends acceptance.

Decision: Accept (Poster)