# Memorization Capacity of Multi-Head Attention in Transformers

Sadegh Mahdavi1,2, Renjie Liao1,2, Christos Thrampoulidis1
1University of British Columbia
2Vector Institute for AI
<EMAIL_ADDRESS>

###### Abstract

Transformers have become the go-to architecture for language and vision tasks,
yet their theoretical properties, especially memorization capacity, remain
elusive. This paper investigates the memorization abilities of multi-head
attention mechanisms, examining how many example sequences they can memorize,
as a function of the number of heads and sequence length. Motivated by
experimental findings on vision transformers, we introduce novel assumptions
about the linear independence of input data, distinct from the commonly used
general-position assumption. Under these assumptions, we demonstrate that an
attention layer with $H$ heads, dimension $d$, and context size $n<d,$
featuring $\Theta(Hd^{2})$ parameters, can memorize $\Omega(Hn)$ examples. Our
analysis sheds light on how different attention heads handle various example
sequences, aided by the softmax operator’s saturation property. We validate
our findings through experiments on synthetic data.

##  1 Introduction

Transformers have emerged as the de-facto architecture for achieving state-of-
the-art performance across different domains such as natural language
processing and vision tasks, as demonstrated by recent breakthroughs, e.g.
(Dehghani et al., 2023; Touvron et al., 2023; OpenAI, 2023). These models are
characterized by their immense scales, often comprising billions of parameters
to achieve their high performances. Given their sizes and complexities, a
natural question arises: How effectively can they memorize the training data?
This query is important both from a privacy perspective (Carlini et al.,
2020), and as a stepping stone towards quantifying the model’s ability to
generalize to new data (Zhang et al., 2017). Additionally, studying how the
memorization of transformers emerges can offer insights into the distinct
roles played by their different building blocks.

Concretely, the memorization capacity of a model $f_{\bm{\mathcal{W}}}$ refers
to the minimal set of parameters ${\bm{\mathcal{W}}}$ such that
$f_{\bm{\mathcal{W}}}(\bm{x})=\bm{y}$ holds for all input-output pairs
$(\bm{x},\bm{y})$ in the training dataset.111The definition of memorization in
this context is distinct from concepts such as ”associative memory” Demircigil
et al. (2017); Ramsauer et al. (2020) and ”neural memory” Sukhbaatar et al.
(2015); Geva et al. (2020) which involve associating an input pattern with a
given query and have been recently studied for transformers. Historically,
this question has been extensively analyzed for fully connected networks
(FCNs), dating back to at least the 1980s when Baum (1988) proved the first
set of results for _single-hidden-layer_ Threshold networks with binary
labels. The best known bound on the memorization capacity of single-hidden-
layer ReLU FCNs has only recently been established by Bubeck et al. (2020),
following numerous recent studies, e.g. (Zhang et al., 2017; Yun et al.,
2019). Recent studies have also extensively considered extensions to deeper
models, e.g. (Vershynin, 2020; Park et al., 2021; Vardi et al., 2022; Yun et
al., 2019). However, as of now, the best known results consist of order-wise
bounds, involving unknown constants and logarithmic factors. The majority of
these works on memorization of FCNs are established for continuous inputs
assumed in “General Position,” a relatively mild assumption, as it holds true
under scenarios like small random perturbations of input examples.

In contrast to FCNs, understanding the memorization capacity of transformers
is still in its infancy. A common Transformer architecture consists of two
main components: (i) a Multi-head Attention (MHA) module and (ii) a ReLU
fully-connected network. The MHA module creates convex combinations of its
input elements by computing softmax similarities between input
representations. This operation, unique to transformers when compared to FCNs,
is leveraged in practice for its ability to selectively attend to information
from various representations and effectively process long sequences (Vaswani
et al., 2017). Yet, its theoretical underpinnings remain largely unexplored.
This paper aims to contribute towards closing this gap by investigating the
memorization capacity of the MHA module. To isolate its role from the ReLU
FCN, we take a step toward analyzing Attention-only models.

Concretely, we study a single-layer MHA module with $H$ heads that takes as
inputs a $n\times d$ key/context matrix comprising representations of $n$
tokens, and a $d\times 1$ query vector. We are particularly interested in
understanding how the number of attention heads, and the context size of the
input affect the memorization in this architecture. Our contributions are as
follows:

$\bullet$ We introduce a new set of input-data assumptions that are more
relaxed compared to General Position assumptions. Our first assumption
requires the set of all query vectors to have Kruskal rank at least $n$, while
our second assumption needs context vectors of each example to be linearly
independent. These are motivated and validated by experiments on real-world
models, such as ViT, following a single layer, whether it be Attention-only or
a complete Transformer, that is either randomly initialized or pretrained. In
contrast, General-Position does _not_ hold in either of these settings.

$\bullet$ For this set of practically relevant assumptions, we prove that a
single-layer MHA module with $H$ heads, embedding dimension $d$, key/query
dimension $d_{h}$, value dimension $d_{v}$, context size $n<d$, output
dimension ${d_{\text{out}}}\leq d_{v}$, equipped with a total of
$O(Hd(d_{h}+d_{v}))$ trainable parameters, is capable of memorizing
$\Omega\left(H\min(n,d_{h})\right)$ input examples. Specifically when
${d_{\text{out}}}=d_{h}=d$ and $n=\Theta(d)$, this finding recovers the
optimal order of real-valued numbers that can be memorized with the given
number of parameters.

$\bullet$ Our proof illuminates how different attention heads play distinct
roles in memorizing different sets of examples, and how this role allocation
is facilitated by the softmax. To show memory capacity increases linearly with
$H$, our idea is to make each head individually “responsible” for memorizing
$\min(n,d_{h})-1$ examples. By leveraging the saturation property of softmax
and introducing specific key-query weight adjustments, we design desired
softmax logits for individual heads while ensuring minimal interference with
other examples.

$\bullet$ In addition to experiments validating our assumptions on the input
sequences, we design and discuss synthetic experiments corroborating our
findings: namely, the impact of head number $H$, context length $n$ and
key/query dimension $d_{h}$ on memorization, as well as, the occurrence of the
saturation property of the softmax operator.

##  2 Related Work

Theoretical understanding of Transformers. Driven by the success of
Transformers, there have been several attempts at understanding their
underlying mechanisms, e.g., (Li et al., 2023b; Luo et al., 2022; von Oswald
et al., 2022; Dong et al., 2021; Edelman et al., 2022; Jelassi et al., 2022;
Li et al., 2023a; Sahiner et al., 2022). Closer to us, (Yun et al., 2020a; b)
study the expressivity of Transformers from a universal-approximation
perspective, by constructing networks with an exponential number of layers and
parameters in input dimensions (context size and dimension) to cover the
entire input space. In contrast, our aim is to find a small network to
memorize a _finite_ number of samples. Bhojanapalli et al. (2020) studies the
expressivity of MHA on a single input sequence, yielding guidelines for
setting the head dimension $d_{h}$ at least as high as the context size $n$ to
preserve expreissivity. To the best of our knowledge, our work is the first to
explicitly address memorization of MHA with more than one input data points.
Only very recently, Kim et al. (2023) has extended previous results for ReLU
networks (Vardi et al., 2022) to Transformers with ReLU-activation feed-
forward layer and a _single_ Attention head, showing that a sublinear number
of parameters suffice for memorization. Here, in order to better understand
the distinct role of the Attention mechanism, we study the memorization of
one-layer MHA. Also, unlike those previous works, we focus on a single
prediction per input. These make our results not directly comparable to
theirs. Also, our assumptions on the input data and proof techniques are
different.
Memory capacity of FCNs. Contrary to Transformers, the memory capacity of FCNs
with ReLU/Threshold/Sigmoid activation has been extensively analyzed. Baum
(1988) provided the first memory bound, proving under General Position that a
single hidden-layer Threshold can memorize $T$ binary labels with $O(T)$
parameters. The best known result on memorization under ReLU activation and
real-valued inputs was obtained only very recently by Bubeck et al. (2020).
These bounds are tight for the respective class of assumptions since
memorization for a dense class of input data (e.g., General Position) with
piece-wise analytic definable functions requires number of parameters at least
linear on the number of examples (Sontag, 1997). Building on these
foundational works, recent work have extended these results to different
input-data assumptions, smooth activations, and multi-layer networks
(Vershynin, 2020; Yun et al., 2019; Rajput et al., 2021; Zhang et al., 2021;
Bombari et al., 2022). Additionally, recent works yield bounds on ReLU
networks with $\tilde{O}(\sqrt{T})$ parameters by utilizing
$\tilde{\Omega}(\sqrt{T})$-deep networks (Vardi et al., 2022; Park et al.,
2021). Following the tradition of early studies on FCNs that explored
memorization in simplified one-hidden layer networks, we delve into the memory
capacity of a single-layer MHA mechanism. Our objective is to isolate the MHA
from the ReLU MLP and analyze its unique properties. We anticipate that our
discoveries will inspire additional research and extensions for more advanced
Transformer models.

##  3 Problem Setup

Notation. $[n]$ denotes the set of positive integers $\\{1,2,\ldots,n\\}$.
Boldface upper-case \ lower-case letters, such as $\bm{A}$ \ $\bm{v}$, denote
matrices \ column-vectors. $[\bm{A};\bm{B}]$ \
$\begin{bmatrix}\bm{A},\bm{B}\end{bmatrix}$ denote row-wise \ column-wise
concatenation. $\bm{x}_{[n]}=\\{\bm{x}_{1},\bm{x}_{2},\dots,\bm{x}_{n}\\}$
denotes the set of $[n]$ indexed vectors.

We study a single-layer _multi-head Attention_ (MHA) mechanism. An MHA layer
$\mathcal{A}$ with $H$ heads consists of three sets of matrices
${\bm{W}_{K}}_{h},{\bm{W}_{Q}}_{h}\in\mathbb{R}^{{d}\times{d_{\text{h}}}},{\bm{W}_{V}}_{h}\in\mathbb{R}^{{d}\times{d_{v}}}$
for each head $h\in[H]$, one weight matrix
${\bm{W}_{O}}\in\mathbb{R}^{H{d_{v}}\times{d}}$ to combine the outputs of
different heads, and finally, a read-out weight matrix
${\bm{W}_{D}}\in\mathbb{R}^{d\times{d_{\text{out}}}}$ to get the final model’s
output. We denote the entire set of parameters as
${\bm{\mathcal{W}}}=\\{\\{{\bm{W}_{Q}}_{h},{\bm{W}_{K}}_{h},{\bm{W}_{V}}_{h}\\}_{h=1}^{H},{\bm{W}_{O}},{\bm{W}_{D}}\\}$.
An input example to the MHA model consists of (i) a key matrix
$\bm{E}=\begin{bmatrix}{\bm{e}}_{1}^{\top};{\bm{e}}_{2}^{\top};\dots;{\bm{e}}_{n}^{\top}\end{bmatrix}\in\mathbb{R}^{n\times{d}}$
containing $n$ tokens, and, (ii) a query vector ${\bm{e}}\in\mathbb{R}^{{d}}$.
The output of the model as a function of the input example and of the set of
parameters is given according to the following computational mechanism:
$\displaystyle\bm{\alpha}_{h}$
$\displaystyle:=\bm{E}{\bm{W}_{K}}_{h}{\bm{W}_{Q}}_{h}^{\top}{\bm{e}}$
$\displaystyle(\bm{\alpha}_{h}\in\mathbb{R}^{n})$ (1)
$\displaystyle\bm{\theta}_{h}$
$\displaystyle:=\operatorname{Softmax}\left(\bm{\alpha}_{h}\right)$
$\displaystyle(\bm{\theta}_{h}\in\mathbb{R}^{n})$ (2)
$\displaystyle\bm{z}_{h}$ $\displaystyle:=\bm{E}^{\top}\bm{\theta}_{h}$
$\displaystyle(\bm{z}_{h}\in\mathbb{R}^{{d}})$ (3) $\displaystyle\bm{p}_{h}$
$\displaystyle:={\bm{W}_{V}}_{h}^{\top}\bm{z}_{h}\quad$
$\displaystyle(\bm{p}_{h}\in\mathbb{R}^{{d_{v}}})$ (4) $\displaystyle\bm{o}$
$\displaystyle:={\bm{W}_{O}}^{\top}\begin{bmatrix}\bm{p}_{1};\bm{p}_{2};\dots;\bm{p}_{H}\end{bmatrix}$
$\displaystyle(\bm{o}\in\mathbb{R}^{{d}})$ (5) $\displaystyle\hat{\bm{y}}$
$\displaystyle:={\bm{W}_{D}}^{\top}\bm{o}$
$\displaystyle(\bm{y}\in\mathbb{R}^{d_{\text{out}}}).$ (6)

Eqs. (1)–(4) hold for all heads $h\in[H]$, $\bm{o}$ is the output of the MHA
layer, and, $\hat{\bm{y}}$ is the final output of the model. In Transformer
terminology, the key matrix $\bm{E}$ is also called the _context_. Moreover,
when ${\bm{e}}\in{\bm{e}}_{[n]}$, the MHA layer is called _Self-Attention_ ,
and when the query vector is not among key vectors, the MHA layer is a
particular form of _Cross-Attention_. Here, we impose _no_ restrictions on the
specific type of MHA as long as certain assumptions on the input data are
satisfied (see Sec. 4).

We consider a training set $\mathcal{T}$ consisting of $T$ input examples. To
each example $t\in[T]$, consisting of $n$ context tokens and one query token,
corresponds a label $\bm{y}^{(t)}$. Thus, we have
$\mathcal{T}:=\\{(\bm{E}^{(t)},{\bm{e}}^{(t)},\bm{y}^{(t)})\\}_{t=1}^{T}$. Our
goal is to find a parameter set ${\bm{\mathcal{W}}}$ such that the MHA model
is able to memorize any label set given the key matrices and query vectors,
i.e., $\bm{y}^{(t)}=f_{\bm{\mathcal{W}}}(\bm{E}^{(t)},{\bm{e}}^{(t)})$ for all
$t\in[T]$. The maximum number $T$ of examples for which we are able to achieve
this serves as a lower bound on the memorization capacity of the model.

To illustrate the role of an input example distinct components (key matrix,
query vector, label), consider image classification using (say) Vision
Transformer (ViT) (Dosovitskiy et al., 2021). Two types of tokens are given as
input to this model: A single “[CLS]” token, and several “image patch” tokens.
The image patches contain information about the input image, while the “[CLS]”
token is a special token in that the model’s class output is read from the
prediction of this token. For each image $t\in[T]$, the key matrix
$\bm{E}^{(t)}$ is the matrix containing image patches and the “[CLS]” token,
the query vector ${\bm{e}}^{(t)}$ corresponds to the “[CLS]” token only, and
$\bm{y}^{(t)}$ is the label of the image (here, the label is a scalar; thus,
$d_{\text{out}}=1$.) Note that, the type of MHA is here Self-Attention since
the “[CLS]” token is inside $\bm{E}^{(t)}$ as well, otherwise, the context
would have contained only the image patches.

Often, a positional encoding is also added to the embedding of each token
before feeding them to the Transformer. Namely, for all $t\in[T]$:
${\bm{e}}^{(t)}:=\bm{x}^{(t)}+\operatorname{pos}_{0},$ and
${\bm{e}}_{i}^{(t)}:=\bm{x}_{i}^{(t)}+\operatorname{pos}_{i}$, for all
$i\in[n],$ where $\bm{x}_{i}^{(t)}$ and $\bm{x}^{(t)}$ are the raw embedding
of input tokens, and $\operatorname{pos}_{i}\in\mathbb{R}^{d}$ are fixed and
unique positional encoding vectors for each position index. In the case of
Self-Attention, the query vector takes only one positional encoding since it
is already present among ${\bm{e}}_{[n]}$. We will show in Sec. 4.4 that the
use of positional encoding motivates our Assumption 2.

##  4 Main Results

Two types of assumptions are typical in prior works on memorization: (1) norm-
based (Vershynin, 2020; Park et al., 2021; Kim et al., 2023), and, (2) linear-
independence (Baum, 1988; Bubeck et al., 2020) assumptions. The former assumes
well-separated and/or bounded-norm input data, while the latter only requires
some notion of linear independence among the data points. While norm-based
assumptions are well-suited for discrete data such as language, linear
independence is usually better suited in continuous data such as images, or
continuous distributions, e.g. Gaussian and Uniform.

In this paper, we focus on the second type of assumptions. Thus, it is useful
to recall the notions of Kruskal Rank and General Position.

######  Definition 1 (Kruskal Rank).

Kruskal Rank of a set of vectors is the largest number $i$ such that every
subset of size $i$ is linearly independent.

######  Definition 2 (General Position).

A set of vectors
$\\{\bm{x}_{1},\bm{x}_{2},\dots,\bm{x}_{m}\\}\subset\mathbb{R}^{d}$ is in
General Position if it has the maximal possible Kruskal Rank of $d$.

In contrast to previous works on the memorization capacity of FCNs, we depart
from the commonly used General Position assumptions and tailor our assumptions
to transformers. Specifically, we make the following assumptions.

######  Assumption 1.

The set of all query vectors
$\\{{\bm{e}}^{(t)}|{\bm{e}}^{(t)}\in\mathbb{R}^{d}\\}_{t=1}^{T}$ has Kruskal
Rank at least $n$.

######  Assumption 2.

For each example $t\in[T]$, the context matrix
$\bm{E}^{(t)}\in\mathbb{R}^{n\times d}$ has rank $n$.

These assumptions find their motivation in experiments conducted on real data,
which reveal scenarios where these assumptions hold, while the General
Position assumption consistently fails (see Sec. 5.1).

We are now ready to state our main result below.

######  Theorem 1.

Consider a multi-head attention layer $\mathcal{A}$ with $H$ heads, embedding
dimensions $d$, ${d_{v}}\geq{d_{\text{out}}}\geq 1$, and $d_{h}\geq 1$. Let
$\mathcal{T}=\left\\{\left(\bm{E}^{(t)},{\bm{e}}^{(t)},\bm{y}^{(t)}\right)\right\\}_{t=1}^{T}$
be a training set with context size $n<d$. Define $r:=\min(n,d_{h})$. If
Assumptions 1 and 2 hold, and $T\leq H\left(r-1\right)+1$, then, there exists
a set of parameters ${\bm{\mathcal{W}}}$ such that $\mathcal{A}$ can memorize
$\mathcal{T}$.

Theorem 1 shows that an MHA with $2Hdd_{h}+Hdd_{v}+d{d_{\text{out}}}$
parameters can memorize at least $H\left(\min(n,d_{h})-1\right)+1$ examples
under Assumptions 2 and 1. To gain further insights on the implications of
this theorem consider three concrete settings: (1) Suppose $d_{h}=d$ and
scalar model output $d_{v}={d_{\text{out}}}=1$, then a MHA with
$\Theta(Hd^{2})$ can memorize $\Omega(Hn)$ examples. (2) Suppose
$d_{h}=n<d,d_{v}={d_{\text{out}}}=1$ then a MHA with $\Theta(Hdn)$ can
memorize $\Omega(Hn)$ examples. (3) Suppose $d_{h}=d_{v}=d$ and
$d$-dimensional model output ${d_{\text{out}}}=d$, then a MHA with
$\Theta(Hd^{2})$ can memorize $\Theta(Hn)$ examples with $d$-dimensional label
each. The third scenario is particularly interesting. To see this further,
consider the typical setting where $n$ and $d$ are of the same order
($n=\Theta(d)$), then the established bound becomes optimal. This optimality
is reflected in the fact that $\Omega(Hnd)=\Omega(Hd^{2})$ real-valued outputs
are memorized, and a trivial upper bound for the memorization of an MHA is the
number of its parameters, which is $\Theta(Hd^{2})$222$d$-dimensional outputs
is common in prior works, e.g., (Yun et al., 2020b; a), with the distinction
that they consider sequence-to-sequence outputs..

Additionally, Theorem 1 suggests there is _no_ gain in memorization capacity
provided $d_{h}>n$. This agrees with using $d_{h}<d$ in practice. We also
confirm this in synthetic experiments in Section 5.2. A similar conclusion is
also drawn by Bhojanapalli et al. (2020), where authors show by having
$d_{h}<n$, Attention loses expressivity in terms of convex combinations
$\bm{\theta}_{h}$ the Attention can represent. However, their finding alone
does not allow direct conclusions about the impact on memorization. Our
theorem directly connects the trade-off between $n$ and $d_{h}$ to
memorization.

The assumption that $n<d$ is general practice. For instance, in prominent
models such as Vision Transformer (Dosovitskiy et al., 2021), BERT (Devlin et
al., 2019), and GPT-3 (Brown et al., 2020), the $n/d$ ratio is $197/768$,
$512/768$, and $2048/12288$, respectively. Furthermore, in Sections 4.4 and
5.1, we demonstrate that Assumption 2 typically holds in practice due to
positional encoding, and Assumption 1 is valid after applying a single
Attention layer. Unlike previous works on FCNs that regard General Position as
an assumption (Baum, 1988; Bubeck et al., 2020), our approach innovatively
incorporates Assumption 1, which imposes a weaker constraint than General
Position since $n<d$. In fact, our experiments reveal that Assumption 1 is
satisfied in real data, even though the more commonly used General Position
assumption is often not met.

###  4.1 Proof of Theorem 1

Our proof consists of two steps: (1) We focus on the intermediate
representations originating from $\bm{z}_{h}$ in Equation (3). We demonstrate
with a suitable set of attention weights
$\\{({\bm{W}_{K}}_{h},{\bm{W}_{Q}}_{h})\\}_{h=1}^{H}$, we can achieve linear
independence across different data points, resulting in a rich representation.
(2) We prove the existence of
$\\{\\{{\bm{W}_{V}}_{h}\\}_{h=1}^{H},{\bm{W}_{O}},{\bm{W}_{D}}\\}$ for
achieving memorization by taking advantage of the rich representation from the
previous step.

Let us define matrix $\bm{Z}$, which is a function input data and
${\bm{W}_{K}}_{h},{\bm{W}_{Q}}_{h},h\in[H]$ as follows:

$\displaystyle\bm{Z}$
$\displaystyle:=\begin{bmatrix}\bm{z}^{(1)^{\top}};\bm{z}^{(2)^{\top}};\dots;\bm{z}^{(T)^{\top}}\end{bmatrix}\in\mathbb{R}^{T\times
dH},~{}\text{where}~{}\bm{z}^{(t)}:=\begin{bmatrix}\bm{z}^{{(t)}^{\top}}_{1},\bm{z}^{{(t)}^{\top}}_{2},\dots,\bm{z}^{{(t)}^{\top}}_{H}\end{bmatrix}^{\top}.$
(7)

Each vector ${\bm{z}}^{(t)}$ is a concatenation of $H$ $d$-dimensional vectors
each representing a convex combination of context vectors of the $t$-th
example parameterized in terms of key/query vectors and
$\bm{W}_{Kh},\bm{W}_{Qh}$. The matrix $\bm{Z}$ gathers all these $dH$
dimensional representations, one for each example in a matrix form. Then,
based on Eq. (1) to Eq. (6), the predicted labels $\hat{\bm{y}}^{(t)}$ are
given by

$\displaystyle\hat{\bm{y}}^{(t)}:=\underbrace{\bm{Z}}_{\text{Step
1}}\underbrace{\operatorname{diag}({\bm{W}_{V}}_{1},{\bm{W}_{V}}_{2},\dots,{\bm{W}_{V}}_{H}){\bm{W}_{O}}{\bm{W}_{D}}}_{\text{Step
2}}\in\mathbb{R},\quad\forall t\in[T].$ (8)

In the first step, using an inductive proof technique, we show that $\bm{Z}$
has a high rank, and use this information in the second step to solve a linear
system of equations and achieve memorization.
_Step 1. The Rank of $\bm{Z}$._ Our key technical result is proving the
following lower bound on $\operatorname{rank}(\bm{Z})$.

######  Proposition 1.

Under the conditions of Theorem 1, there exists parameters
$\\{({\bm{W}_{K}}_{h},{\bm{W}_{Q}}_{h})\\}_{h=1}^{H}$ such that
$\operatorname{rank}(\bm{Z})\geq\min\left\\{H(r-1)+1,T\right\\}.$

We provide a proof sketch of this proposition here and defer the full proof to
Appendix B. The high-level idea is to inductively add heads and use each new
head for a distinct set of $r-1$ examples, which in turn translates to an
increase on the rank of $\bm{Z}$. The new head exploits a new distinct set of
examples without affecting the rank of the previous induction step. To
accomplish this, we leverage the saturation property of the softmax and the
linear independence assumptions. For simplicity, we only consider here the
case $d_{h}\geq d>n$, for which $r=n$.

In order to get $H(n-1)+1$ rank, each time a head is added, the rank of
$\bm{Z}$ should be increased by $n-1$. For each head $h$, we only focus on
$h(n-1)+1$ examples, therefore, the subsequent step can be viewed as adding a
new head with tunable weights, and a set of $n-1$ fresh data points. Namely,
the induction step of $\bm{Z}$ can be described in view of the following block
matrix:

$\bm{Z}=\left[\begin{array}[]{ccc|c}&&&\\\
&\bm{Z}_{a}^{\prime}&&\bm{Z}_{a}^{\prime\prime}\\\ &&&\\\
\hline\cr&\bm{Z}_{b}^{\prime}&&\bm{Z}_{b}^{\prime\prime}\\\
\makebox[0.0pt][l]{$\smash{\underbrace{\phantom{\begin{matrix}&&&\end{matrix}}}_{\text{$H^{\prime}d$}}}$}&&&\makebox[0.0pt][l]{$\smash{\underbrace{\phantom{\begin{matrix}~{}~{}~{}~{}~{}~{}\end{matrix}}}_{\text{$d$}}}$}{}{}{}{}{}{}\end{array}\right]\stackrel{{\scriptstyle\text{row-
operations}}}{{\longrightarrow}}\left[\begin{array}[]{ccc|c}&&&\\\
&\bm{Z}_{a}^{\prime}&&\bm{Z}_{a}^{\prime\prime}\\\ &&&\\\
\hline\cr&\textbf{0}&&\bm{Z}_{b}^{\prime\prime}+\mathbf{A}\\\
\makebox[0.0pt][l]{$\smash{\underbrace{\phantom{\begin{matrix}&&&\end{matrix}}}_{\text{$H^{\prime}d$}}}$}&&&\makebox[0.0pt][l]{$\smash{\underbrace{\phantom{\begin{matrix}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}\end{matrix}}}_{\text{$d$}}}$}{}{}{}{}{}{}{}{}{}{}{}{}{}\end{array}\right]\begin{array}[]{@{}
l @{}}\left.\begin{array}[]{@{}c@{}}\hbox{}\hfil\\\ \hbox{}\hfil\\\
\hbox{}\hfil\end{array}\right\\}{\scriptstyle T^{\prime}:=T-n+1}\\\
\left.\begin{array}[]{@{}c@{}}\hbox{}\hfil\\\
\hbox{}\hfil\end{array}\right\\}{\scriptstyle n-1}\end{array}$

$\bm{Z}^{\prime}_{a}\in\mathbb{R}^{T^{\prime}\times H^{\prime}d}$ denotes the
matrix derived from the induction hypothesis,
$\bm{Z}^{\prime}_{b}\in\mathbb{R}^{(n-1)\times H^{\prime}d}$ contains the
embeddings of new examples corresponding to the $H^{\prime}\geq 1$ heads of
the previous induction steps,
$\bm{Z}^{\prime\prime}_{a}\in\mathbb{R}^{T^{\prime}\times d}$ contains the
embeddings of the previous examples corresponding to the new head
$H=H^{\prime}+1$, and finally,
$\bm{Z}^{\prime\prime}_{b}\in\mathbb{R}^{(n-1)\times d}$ contains the
embeddings of the new examples corresponding to the new head $H$. The upper-
triangular block matrix on the right-hand side is obtained from $\bm{Z}$ only
by row operations (see the appendix for details); thus, (i) its rank is the
same as the rank of $\bm{Z}$, and, (ii) the matrix
$\mathbf{A}:=\mathbf{A}(\bm{Z}_{a}^{\prime\prime})$ is a function of the block
$\bm{Z}_{a}^{\prime\prime}$.

In view of the upper-triangular block formulation above, our objective becomes
clear.

Goal: Adjust the weights ${\bm{W}_{K}}_{H}$, ${\bm{W}_{Q}}_{H}$ of the new
head $H$, to increase the rank of $\bm{Z}^{\prime\prime}_{b}+\mathbf{A}$ while
dealing with the challenge that such tuning can also impact
$\bm{Z}_{a}^{\prime\prime}$, subsequently affecting the rank of
$\mathbf{A}=\mathbf{A}(\bm{Z}_{a}^{\prime\prime})$. We will demonstrate a
method for precisely tuning these weights to control the rank of
$\bm{Z}_{b}^{\prime\prime}$ without altering the rank of
$\bm{Z}_{a}^{\prime\prime}$.

Approach: We approach this in two sub-steps, providing an overview of each
below. Recall
$\bm{\theta}_{H}^{(t)}(\mathbf{W}):=\operatorname{Softmax}\left({\mathbf{E}^{(t)}\mathbf{W}\mathbf{e}^{(t)}}\right)$
is the vector of softmax coefficients, associated with the $t$-th input and
the new head $H$, parameterized by
$\mathbf{W}:={\bm{W}_{K}}_{H}{\bm{W}_{Q}}_{H}^{\top}$. Also note that
$\bm{Z}_{a}^{\prime\prime}$ and $\bm{Z}_{b}^{\prime\prime}$ are parameterized
by $\bm{\theta}_{H}^{(t)}(\mathbf{W})$. Consequently, we will use
${\bm{Z}_{a}^{\prime\prime}}(\mathbf{W})$ and
${\bm{Z}_{b}^{\prime\prime}}(\mathbf{W})$ to explicitly denote
parameterization with a specific weight matrix.

$\bullet$ Sub-step 1.A. Tune the new head’s weights and memorize the new set
of examples. Given an arbitrary, but fixed matrix
$\bm{M}\in\mathbb{R}^{(n-1)\times d}$, whose precise value will be chosen
later in the second substep, we demonstrate that it is possible to design
weights $\mathbf{W}^{*}:=\mathbf{W}^{*}(\bm{M})$ to achieve a target rank for
$\bm{Z}_{b}^{\prime\prime}$, specifically to achieve
$\operatorname{rank}(\bm{Z}_{b}^{\prime\prime}(\mathbf{W}^{*})+\bm{M})\geq
n-1.$ The construction is based on an $n$-step induction, making use of
Assumption 2 that the rows of $\mathbf{E}^{(t)}$ span an $n$-dimensional
space.

$\bullet$ Sub-step 1.B. Suppress the impact of the new head’s weights on the
previous examples. To select the matrix $\bm{M}$, we establish the following
procedure for fixing $\bm{Z}_{a}^{\prime\prime}$. We first establish (see
Lemma 3) the existence of a matrix $\mathbf{W}^{+}$, satisfying two
conditions: (i) ${\mathbf{E}^{(t)}\mathbf{W}^{+}\mathbf{e}^{(t)}}=0$ for all
$t=T^{\prime}+1,\ldots,T$, and (ii)
${\mathbf{e}_{i}^{(t)}\mathbf{W}^{+}\mathbf{e}^{(t)}}\neq{\mathbf{e}_{j}^{(t)}\mathbf{W}^{+}\mathbf{e}^{(t)}}$
for all $i\neq j\in[n]$ and $t\in[T^{\prime}]$. Given $\mathbf{W}^{+}$, define
$\bm{\theta}_{+}^{(t)}:=\mathbf{1}[\arg\max_{i\in[n]}\mathbf{e}_{i}^{(t)}\mathbf{W}^{+}\mathbf{e}^{(t)}]$
for $t\in[T^{\prime}]$. The crucial insight here is that owing to the
saturating property of the softmax (see Lemma 4) and property (ii) above, for
_any_ weight matrix $\mathbf{W}$ (which may depend on $\mathbf{W}^{+}$), we
have
$\bm{\theta}_{+}^{(t)}=\lim_{c\rightarrow\infty}\bm{\theta}_{H}^{(t)}\left(\mathbf{W}+c\mathbf{W}^{+}\right)$
for all $t\in[T^{\prime}]$. Consequently, for _any_ $\mathbf{W}$, the matrix
$\lim_{c\rightarrow\infty}\bm{Z}_{a}^{\prime\prime}(\mathbf{W}+c\mathbf{W}^{+})=[{\bm{E}^{(1)}}^{\top}\,\bm{\theta}_{+}^{(1)}\cdots{\bm{E}^{(T^{\prime})}}^{\top}\,\bm{\theta}_{+}^{(T^{\prime})}]^{\top}=:\bm{Z}_{a,+}$
is independent of the choice of $\mathbf{W}$.

The proof is finalized by integrating the two sub-steps. Specifically, we
utilize $\bm{Z}_{a,+}$ from sub-step 1.B and select
$\bm{M}=\mathbf{A}(\bm{Z}_{a,+})$. Next, we apply sub-step 1.A to identify the
weights $\mathbf{W}^{*}=\mathbf{W}^{*}(\mathbf{W}^{+})$ in a manner that
ensures
$\operatorname{rank}(\bm{Z}_{b}^{\prime\prime}(\mathbf{W}^{*})+\mathbf{A}(\bm{Z}_{a,+}))\geq
n-1$. Finally, we set the overall weights of the new head as
$\mathbf{W}=\mathbf{W}^{*}+c\mathbf{W}^{+}$. As the parameter $c$ approaches
infinity, the softmax coefficients for the first $T^{\prime}$ examples will be
exclusively determined by $\mathbf{W}^{+}$ (not $\mathbf{W}^{*}$), as observed
in Sub-step 1.B. Furthermore, thanks to property (i) of Sub-step 1.B, we have
$\bm{Z}_{b}^{\prime\prime}(\mathbf{W}^{*}+c\mathbf{W}^{+})=\bm{Z}_{b}^{\prime\prime}(\mathbf{W}^{*})$
for all values of $c$. Collectively, these steps complete the proof of the
proposition.

In the Appendix B, we provide all the details. This involves showing that a
finite value of scaling $c$ above suffices; thus, there exist _finite_ weights
with the desired property. Moreover, for the scenario $d_{h}\leq d$, we
demonstrate that the matrices $\mathbf{W}^{*},\mathbf{W}^{+}$ can be chosen to
have rank at most $d_{h}$.

_Step 2. Solving system of equations._ We can find the remaining weights to
memorize at least $\operatorname{rank}(\bm{Z})$ examples. WLOG, assume that
$T=\operatorname{rank}(\bm{Z})$, otherwise we can ignore some of the data
points and keep only $\operatorname{rank}(\bm{Z})$ data points that make
$\bm{Z}$ full-rank. We find weights
$\\{\\{{\bm{W}_{V}}_{h}\\}_{h=1}^{H},{\bm{W}_{O}},{\bm{W}_{D}}\\}$ to memorize
all $T$ examples as follows:

$\displaystyle\bm{Z}\bm{W}_{V}=\left[\begin{array}[]{c|c}\bm{y}^{{(1)}^{\top}}&\\\
\vdots&\bm{0}_{T\times(d_{v}-d_{out})}\\\
\bm{y}^{{(T)}^{\top}}&\end{array}\right],~{}~{}\begin{bmatrix}{\bm{W}_{V}}_{1}\\\
{\bm{W}_{V}}_{2}\\\ \vdots\\\
{\bm{W}_{V}}_{H}\end{bmatrix}:=\bm{W}_{V},~{}~{}{\bm{W}_{O}}:=\begin{bmatrix}\bm{I}_{{d_{v}}\times
d}\\\ \bm{I}_{{d_{v}}\times d}\\\ \vdots\\\ \bm{I}_{{d_{v}}\times
d}\end{bmatrix},~{}~{}{\bm{W}_{D}}:=\bm{I}_{d\times{d_{\text{out}}}},$

where $\bm{W}_{V}$ is a solution to the first expression (note that such
$\bm{W}_{V}$ exists since $\bm{Z}$ has full row rank). Substituting in Eq.
(8), this gives $\hat{\bm{y}}^{(t)}=\bm{y}^{(t)}$ for all $t\in[T]$.
Therefore, we have proved that the memorization capacity of one layer is at
least $\operatorname{rank}(\bm{Z})$. But, since we also proved a lower bound
of $\operatorname{rank}(\bm{Z})\geq H(r-1)+1$, we conclude that the
memorization capacity is at least $H(r-1)+1$.

######  Remark 1.

In the setting of Theorem 1, if Assumption 1 is violated (i.e., the Kruskal
Rank of query vectors is less than $n$), define $Q<n$ to be the Kruskal Rank
of query vectors
$\\{{\bm{e}}^{(t)}|{\bm{e}}^{(t)}\in\mathbb{R}^{d}\\}_{t=1}^{T}$. Then, a
weaker lower bound on the memorization still holds: $T\leq
H(\min(Q,d_{h})-1)+1$.

###  4.2 Tightness of the lower bound on the rank

Our proof of Theorem 1 relies on lower bounding $\operatorname{rank}(\bm{Z})$
in Proposition 1. One might naturally wonder how tight this bound is. The
proposition below proves that the bound is indeed tight.

######  Proposition 2.

Under the conditions of Theorem 1, if all the contexts are shared, i.e.,
$\bm{E}^{(t)}:=\bm{E}$ for all $t\in[T]$, then for any
$\\{({\bm{W}_{K}}_{h},{\bm{W}_{Q}}_{h})\\}_{h=1}^{H}$, we have
$\operatorname{rank}(\bm{Z})\leq H(n-1)+1$.

We defer the proof to Appendix E. The proposition shows that the lower bound
of Proposition 1 is tight when the context is shared across data points and
$d_{h}\geq n$. Furthermore, the setting of Proposition 2 better highlights the
dependency factor of $r-1$ rather than $r$. To illustrate this, consider this
setting for $r=n=1$, i.e., each context contains a single token and is shared
across all examples. Then, all $\bm{\theta}^{(i)}_{h}$ become scalars, whose
values are equal to one. Consequently, the vector $\bm{z}^{(t)}_{h}$ copies
the first (and only) context token, which is shared across all examples.
Hence, no matter how large $H$ or $d$ are, the output of MHA for all examples
is identical, and the memory capacity of the attention is exactly
$T=H(r-1)+1=1$ (i.e., a single example).

###  4.3 Comparison with two-layer ReLU networks

Given Assumptions 1 and 2, we compare the memorization capacity of MHA with
one-hidden layer FCN networks of the same size. An FCN does not distinguish
between key/query vectors. Hence, we flatten out the concatenations of
$\bm{E}^{(t)}$ and $\mathbf{e}^{(t)}$, resulting in input
$\bm{x}_{\text{FCN}}^{(t)}:=[{\bm{e}}_{1}^{(t)};\dots;{\bm{e}}_{n}^{(t)};{\bm{e}}^{(t)}]\in\mathbb{R}^{(n+1)d},t\in[T]$.
The proposition below upper bounds the memorization power of a ReLU FCN.

######  Proposition 3.

Suppose a two-layer ReLU network with input dimension $(n+1)d$, hidden
dimension $m$, and output dimension ${d_{\text{out}}}$. Then, the memorization
power of this network for the class of datasets
$\mathcal{T}=\\{(\bm{E}^{(t)},{\bm{e}}^{(t)},\bm{y}^{(t)})\\}_{t=1}^{T}$ with
context size $n<d$ and satisfying Assumptions 1 and 2 is at most
$T\leq(n+1)m/{d_{\text{out}}}+(m+1)$ examples.

We defer the proof to Appendix F. A ReLU FCN with
$\Theta(ndm+m{d_{\text{out}}})$ parameters can memorize at most
$O(nm/{d_{\text{out}}}+m)$ examples under Assumptions 1 and 2. This setting is
particularly interesting with $m=\Theta(H)$ and ${d_{\text{out}}}=1$, where a
two-layer ReLU network with $\Theta(Hnd)$ can memorize at most $O(Hn)$
examples. This _upper_ bound on memorization of a ReLU network together with
the _lower_ bound of $\Omega(Hn)$ as established in Theorem 1 for MHA with the
same order of parameters, shows that the latter is at least as powerful, in
terms of memorization, as the former (when data satisfy Assumptions 1 and 2).
This also aligns with practical Transformer implementations, where a similar
order of parameters is allocated to both Attention and ReLU FCN (e.g.,
(Touvron et al., 2023).

###  4.4 Analysis of the validity of assumptions

We confirm our assumptions are valid either on the embedding layer or after
just one Attention layer. Detailed experimental verification of this full-rank
assumption is provided in Section 5.1.

First, we motivate Assumption 2 by considering the common practice of
incorporating positional encoding into token representations before inputting
data into a Transformer. A widely used positional encoding method utilizes
sinusoidal positional encodings (Vaswani et al., 2017). Consequently, the
positional encoding matrix is inherently full rank. In view of the positional
encoding equations (see Sec. 3), this characteristic extends to the context
matrix $\bm{E}^{(t)}$, establishing the validity of Asm. 2.

Second, we provide an explanation for why Assumption 1 generally becomes valid
after just one Attention layer. Consider the case of image classification
using ViT, where query vectors ${\bm{e}}^{{(t)}}$ are all composed of the
“[CLS]” token combined with positional encoding at position zero. Here,
Assumption 1 does not initially hold. However, we demonstrate below that it
becomes valid after a single Self-Attention layer with a skip connection. The
key insight lies in the behavior of the “[CLS]” token: as it traverses the
attention layer, it effectively mixes information from the surrounding
context, which is distinct for each image. Therefore, after a single (even
randomly initialized and fixed) Self-Attention layer, we establish that
Assumptions 2 and 1 are satisfied. This allows us to directly apply Theorem 1
to the second layer of Attention to achieve memorization. This line of
argument is conceptually similar to the approach employed by Vershynin (2020),
who first designs layers with random weights to enrich the representation of
data points, and then proves the existence of a learned subsequent perception
layer for memorizing the labels in a FCN.

Now, we provide the details on how one Attention-layer with skip connection
validates Assumption 1. To begin with, we recall how a Self-Attention layer
with a skip connection works. Consider a single example with context matrix
$\bm{E}$ and query token ${\bm{e}}$. Given a set of attention weights
${\bm{\mathcal{W}}}$, the Self-Attention layer with skip connection transforms
the tokens according to the following equations:

$\displaystyle{\bm{e}}^{\prime}:={\bm{e}}+f_{\bm{\mathcal{W}}}\big{(}\bm{E},{\bm{e}}\big{)},\quad{\bm{e}}^{\prime}_{i}:={\bm{e}}_{i}+f_{\bm{\mathcal{W}}}\big{(}\bm{E},{\bm{e}}_{i}\big{)}\quad\text{for
all }i\in[n],\quad\bm{E}^{\prime}$
$\displaystyle:=[{\bm{e}}_{1}^{\prime^{\top}};{\bm{e}}_{2}^{\prime^{\top}};\dots;{\bm{e}}_{n}^{\prime^{\top}}],$

where ${\bm{e}}^{\prime}$, and $\bm{E}^{\prime}$ correspond to the output of
query, and context vectors, respectively. Note that the query matrix is part
of the context matrix due to Self-Attention. Let us assume the query vector is
in the first row of $\bm{E}$, i.e. ${\bm{e}}_{1}={\bm{e}}$, and consequently
${\bm{e}}^{\prime}={\bm{e}}^{\prime}_{1}$.

######  Remark 2.

Theorem 1 also holds true for Attention layers with skip connection (see
Appendix C).

We are now ready to state a proposition on the mixing power of this type of
Attention with skip connection. To do this, we introduce the following relaxed
assumption.

######  Assumption 3.

For training set
$\mathcal{T}=\left\\{\left(\bm{E}^{(t)},{\bm{e}}^{(t)},\bm{y}^{(t)}\right)\right\\}_{t=1}^{T}$,
define
$\tilde{{\bm{e}}}^{(t)}:={\bm{e}}^{(t)}+\frac{1}{n}\sum_{i\in[n]}{\bm{e}}_{i}^{(t)}$
for all $t\in[T]$. Then the set
$\mathcal{S}:=\left\\{\tilde{{\bm{e}}}^{(1)},\tilde{{\bm{e}}}^{(2)},\dots,\tilde{{\bm{e}}}^{(T)}\right\\}$
has Kruskal Rank at least $n$.

######  Proposition 4.

Define Self-Attention layer $\mathcal{A}_{0}$ with single head, skip
connection, ${d_{v}}=d$, and weights
${\bm{W}_{K}}={\bm{W}_{Q}}=\bm{0},{\bm{W}_{V}}={\bm{W}_{O}}=\bm{I}$. Suppose
training set
$\mathcal{T}=\left\\{\left(\bm{E}^{(t)},{\bm{e}}^{(t)},\bm{y}^{(t)}\right)\right\\}_{t=1}^{T}$
satisfies Assumptions 2 and 3. Let
$\bm{E}^{\prime{(t)}},{\bm{e}}^{\prime{(t)}}$ be the output of
$\mathcal{A}_{0}$ for each example $t\in[T]$ and define new dataset
$\mathcal{T}^{\prime}:=\left\\{\left(\bm{E}^{\prime{(t)}},{\bm{e}}^{\prime{(t)}},\bm{y}^{(t)}\right)\right\\}_{t=1}^{T}$.
Then, $\mathcal{T}^{\prime}$ satisfies both Assumptions 2 and 1.

Assumption 3 assumes that the input data as a whole is sufficiently
distinctive across different examples. In particular, this is a more relaxed
assumption compared to Assumption 2. For instance, in the case of images, it
corresponds to the average of all image patches (plus a constant vector) being
sufficiently distinctive across different images. Proposition 4 establishes
that under the given assumptions (specifically, Asm. 3 in conjunction with
Asm. 2), the application of a trivial Self-Attention layer results in the
fulfillment of Assumption 1 at the output of the Attention layer, all while
preserving Assumption 2. This implies that the “[CLS]” token is capable of
performing token mixing even with a basic Attention mechanism. Consequently,
Theorem 1 can be directly applied to a second layer of Attention, facilitating
effective memorization of the examples.

##  5 Experiments

###  5.1 Testing assumptions on Vision Transformer (ViT)

As mentioned in Section 4.4, we show that our assumptions on the input tokens
hold after applying a single Attention layer. To test this, we mainly focus on
image classification and Vision Transformer (ViT) (Dosovitskiy et al., 2021;
Wolf et al., 2020). We consider the following models: 1) Embedding Layer: A
randomly initialized embedding layer. This layer contains only positional
encoding and a linear embedding layer (i.e., the input tokens
$\bm{E}^{(t)},{\bm{e}}^{(t)}$). 2) Random Attention: A randomly initialized
Self-Attention layer with similar architecture to a ViT layer, without the FCN
of a Transformer. 3) Random ViT: A randomly initialized base ViT. We only look
at the output of the first layer and discard the subsequent layers. 4) Trained
ViT: Similar to the random ViT, with the difference that we take the weights
from the ViT pre-trained on ImageNet (Deng et al., 2009).

We evaluate the mentioned models on 2000 images sampled from ImageNet. To
empirically test Assumption 2, we verify whether the context vectors are all
linearly independent for each example. On the other hand, testing Assumption 1
is computationally difficult since computing Kruskal Rank is NP-Hard.
Therefore, it can only be approximately tested in polynomial time. To address
this, we randomly sample $n$ query vectors from the training set $\mathcal{T}$
and construct a matrix of size $n\times d$. We then test if the rank of the
matrix is equal to $n$. After repeating this test 5000 times, we report that
the assumption holds if the rank check holds for at least $99\%$ of the test
instances. The results of our empirical tests are summarized in Table 1.
Specifically, we find that only Assumption 1 is violated in the Embedding
layer, whereas both assumptions hold true for the remaining models. Notably,
the commonly made assumption of General Position fails in all cases, and the
Kruskal Rank is typically much less than $d$, as shown by Figure 1. We also
empirically test Assumption 3 and find that it already holds for the Embedding
layer. Following up on Section 4.4, this may explain why Assumption 1 already
holds after applying any type of Attention. 333Our code is available at
https://github.com/smahdavi4/attention-memorization

Figure 1: Testing Kruskal Rank of query tokens on the output of one layer
Random Attention on ImageNet. The Kruskal Rank is only slightly larger than
$n$ (Assumption 1), and much smaller than $d$ (General Position).

Model | Gen. Pos. | Asm. 1 | Asm. 2
---|---|---|---
Embedding | $\times$ | $\times$ | $\checkmark$
Rand Attention | $\times$ | $\checkmark$ | $\checkmark$
Rand ViT | $\times$ | $\checkmark$ | $\checkmark$
Trained ViT | $\times$ | $\checkmark$ | $\checkmark$

Table 1: Testing assumptions on ImageNet and ViT: while only one assumption is
violated in the embedding layer, both hold true after a single layer of
Attention or Transformer. The General Position assumption is violated in all
cases.

###  5.2 Empirical validation of memorization improvement factors

Theorem 1 provides three conclusions: (1) When fixing dimension $d$,
increasing the number of heads $H$ improves memorization. (2) When further
fixing the number of heads $H$, increasing the context size $n$ improves
memorization. (3) When fixing $d,n$, increasing $d_{h}$ only helps up to
$d_{h}<n$, and there is no memorization gain beyond that. By conducting
synthetic experiments, we verify the conclusions drawn by the theorem.
Throughout, we fix dimension $d=64$. Moreover, to follow Assumptions 2 and 1,
we generate random query inputs, and random shared context inputs (see
Appendix I.1 for experiments with non-shared context) according to a uniform
distribution, as follows:

$\displaystyle\bm{x}_{1},\bm{x}_{2},\dots,\bm{x}_{n}\overset{\text{{iid}}}{\sim}\mathcal{U}(0,1)^{64},\quad\bm{x}^{(t)}\overset{\text{{iid}}}{\sim}\mathcal{U}(0,1)^{64},\quad\bm{x}_{i}^{(t)}$
$\displaystyle:=\bm{x}_{i}\quad\text{for all }i\in[n],~{}t\in[T],$

where $\bm{x}_{i}^{(t)},\bm{x}^{(t)}$ are the raw inputs. Each input is then
given to the model consisting of an embedding layer, positional encoding, and
an Attention layer. To test the first conclusion, we fix $n=32,d_{h}=d$ and
increase $H$. To test the second conclusion, we fix $H=20,d_{h}=d$ and
increase context size $n$. In addition to classification tasks, we also
include regression tasks in the appendix.

The results of our experiments are presented in Figure 2 for classification
task with various dataset sizes $T$ and uniform labels
$y\in\mathcal{U}([100])$ from $100$ classes (see Figure 5 in the appendix for
the regression task). We measure memorization by computing average accuracy
across examples. Observe that memorization increases monotonically with $n$
and linearly with $H$. To quantify the linearity with $H$, we measure the
$R^{2}$ score of the best linear fit of the non-saturated part in Figure 2(a)
and observe $R^{2}\geq 0.98$ in all cases, confirming a high degree of
linearity. Also observe in Figure 2(c) that there appears no benefit on
memorization fraction beyond $d_{h}>n$.

(a)

(b)

(c)

Figure 2: Testing memorization as a function of number of heads (a), context
size (b), and head size (c) for classification under Assumptions 2 and 1.
Examples generated synthetically with $d=64$ and shared context (see
Proposition 2). Memorization increases linearly with $H$, monotonically with
$n$, and monotonically with $d_{h}$ as long as $d_{h}\leq n$.

##  6 Conclusion and Discussion

In this paper, we proved a lower bound on the memorization capacity of
attention layers under a rather mild linear independence assumption. Our proof
technique relies on establishing a tight bound on the rank of a matrix of
intermediate representations in the attention layer. We discussed implications
of our lower bound which we also confirmed by additional synthetic
experiments. Our work opens up several promising avenues for future research.
First, it would be valuable to extend our theoretical results beyond one layer
of attention networks and to sequence-to-sequence learning scenarios. Second,
while we studied some specific cases of interest, establishing general upper
bounds on the memorization capacity of MHA remains an important open question.
Finally, refining our data assumptions to narrow the gap between practical
memorization abilities and theoretical guarantees could shed more light on the
memorization capacity of Attention networks (see Appendix I). Ultimately, we
hope that gaining a better understanding of the memorization abilities of
Transformers could facilitate more efficient implementations and better use of
data for training more effective, generalizable, and privacy-preserving
models.

#### Acknowledgments

This work was funded, in part, by NSERC DG Grants (No. RGPIN-2022-04636 and
No. RGPIN-2021-03677), the Vector Institute for AI, Canada CIFAR AI Chair, and
Oracle Cloud credits. Resources used in preparing this research were provided,
in part, by the Province of Ontario, the Government of Canada through CIFAR,
and companies sponsoring the Vector Institute (www.vectorinstitute.ai),
Advanced Research Computing at the University of British Columbia, the Digital
Research Alliance of Canada (alliancecan.ca), and the Oracle for Research
program.

## References

  * Baum (1988) Eric B Baum.  On the capabilities of multilayer perceptrons.  _Journal of Complexity_ , 4(3):193–215, 1988\.  ISSN 0885-064X.  doi: https://doi.org/10.1016/0885-064X(88)90020-9.
  * Bhojanapalli et al. (2020) Srinadh Bhojanapalli, Chulhee Yun, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar.  Low-rank bottleneck in multi-head attention models.  In Hal Daumé III and Aarti Singh (eds.), _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 864–873. PMLR, 13–18 Jul 2020.
  * Bombari et al. (2022) Simone Bombari, Mohammad Hossein Amani, and Marco Mondelli.  Memorization and optimization in deep neural networks with minimum over-parameterization.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), _Advances in Neural Information Processing Systems_ , volume 35, pp. 7628–7640. Curran Associates, Inc., 2022.  URL https://proceedings.neurips.cc/paper_files/paper/2022/file/323746f0ae2fbd8b6f500dc2d5c5f898-Paper-Conference.pdf.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei.  Language models are few-shot learners.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 1877–1901. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper_files/paper/2020/file/1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf.
  * Bubeck et al. (2020) Sébastien Bubeck, Ronen Eldan, Yin Tat Lee, and Dan Mikulincer.  Network size and size of the weights in memorization with two-layers neural networks.  In _Neural Information Processing Systems_ , 2020.  URL https://api.semanticscholar.org/CorpusID:227275586.
  * Carlini et al. (2020) Nicholas Carlini, Florian Tramèr, Eric Wallace, Matthew Jagielski, Ariel Herbert-Voss, Katherine Lee, Adam Roberts, Tom B. Brown, Dawn Xiaodong Song, Úlfar Erlingsson, Alina Oprea, and Colin Raffel.  Extracting training data from large language models.  In _USENIX Security Symposium_ , 2020.
  * Dehghani et al. (2023) Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenatton, Lucas Beyer, Michael Tschannen, Anurag Arnab, Xiao Wang, Carlos Riquelme Ruiz, Matthias Minderer, Joan Puigcerver, Utku Evci, Manoj Kumar, Sjoerd Van Steenkiste, Gamaleldin Fathy Elsayed, Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jasmijn Bastings, Mark Collier, Alexey A. Gritsenko, Vighnesh Birodkar, Cristina Nader Vasconcelos, Yi Tay, Thomas Mensink, Alexander Kolesnikov, Filip Pavetic, Dustin Tran, Thomas Kipf, Mario Lucic, Xiaohua Zhai, Daniel Keysers, Jeremiah J. Harmsen, and Neil Houlsby.  Scaling vision transformers to 22 billion parameters.  In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pp. 7480–7512. PMLR, 23–29 Jul 2023.  URL https://proceedings.mlr.press/v202/dehghani23a.html.
  * Demircigil et al. (2017) Mete Demircigil, Judith Heusel, Matthias Löwe, Sven Upgang, and Franck Vermet.  On a model of associative memory with huge storage capacity.  _Journal of Statistical Physics_ , 168:288–299, 2017.
  * Deng et al. (2009) Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , pp. 248–255. Ieee, 2009.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: pre-training of deep bidirectional transformers for language understanding.  In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers)_ , pp. 4171–4186. Association for Computational Linguistics, 2019\.  doi: 10.18653/v1/n19-1423.  URL https://doi.org/10.18653/v1/n19-1423.
  * Dong et al. (2021) Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas.  Attention is not all you need: pure attention loses rank doubly exponentially with depth.  In Marina Meila and Tong Zhang (eds.), _Proceedings of the 38th International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , pp. 2793–2803. PMLR, 18–24 Jul 2021.  URL https://proceedings.mlr.press/v139/dong21a.html.
  * Dosovitskiy et al. (2021) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  In _9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021_. OpenReview.net, 2021.  URL https://openreview.net/forum?id=YicbFdNTTy.
  * Edelman et al. (2022) Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang.  Inductive biases and variable creation in self-attention mechanisms.  In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pp. 5793–5831. PMLR, 17–23 Jul 2022.  URL https://proceedings.mlr.press/v162/edelman22a.html.
  * (14) Wikimedia Foundation.  Wikimedia downloads.  URL https://dumps.wikimedia.org.
  * Garg et al. (2022) Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant.  What can transformers learn in-context? a case study of simple function classes.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), _Advances in Neural Information Processing Systems_ , volume 35, pp. 30583–30598. Curran Associates, Inc., 2022.  URL https://proceedings.neurips.cc/paper_files/paper/2022/file/c529dba08a146ea8d6cf715ae8930cbe-Paper-Conference.pdf.
  * Geva et al. (2020) Mor Geva, R. Schuster, Jonathan Berant, and Omer Levy.  Transformer feed-forward layers are key-value memories.  _ArXiv_ , abs/2012.14913, 2020.  URL https://api.semanticscholar.org/CorpusID:229923720.
  * Horn & Johnson (1990) Roger A. Horn and Charles R. Johnson.  _Matrix Analysis_.  Cambridge University Press, 1990.  ISBN 0521386322.  URL http://www.amazon.com/Matrix-Analysis-Roger-Horn/dp/0521386322%3FSubscriptionId%3D192BW6DQ43CK9FN0ZGG2%26tag%3Dws%26linkCode%3Dxm2%26camp%3D2025%26creative%3D165953%26creativeASIN%3D0521386322.
  * Jelassi et al. (2022) Samy Jelassi, Michael E. Sander, and Yuan-Fang Li.  Vision transformers provably learn spatial structure.  _ArXiv_ , abs/2210.09221, 2022.
  * Kim et al. (2023) Junghwan Kim, Michelle Kim, and Barzan Mozafari.  Provable memorization capacity of transformers.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=8JCg5xJCTPR.
  * Li et al. (2023a) Hongkang Li, M. Wang, Sijia Liu, and Pin-Yu Chen.  A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity.  _ArXiv_ , abs/2302.06015, 2023a.
  * Li et al. (2023b) Yingcong Li, M Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak.  Transformers as algorithms: Generalization and implicit model selection in in-context learning.  _arXiv preprint arXiv:2301.07067_ , 2023b.
  * Luo et al. (2022) Shengjie Luo, Shanda Li, Shuxin Zheng, Tie-Yan Liu, Liwei Wang, and Di He.  Your transformer may not be as powerful as you expect.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), _Advances in Neural Information Processing Systems_ , volume 35, pp. 4301–4315. Curran Associates, Inc., 2022.  URL https://proceedings.neurips.cc/paper_files/paper/2022/file/1ba5f64159d67775a251cf9ce386a2b9-Paper-Conference.pdf.
  * OpenAI (2023) OpenAI.  Gpt-4 technical report.  _ArXiv_ , abs/2303.08774, 2023.  URL https://api.semanticscholar.org/CorpusID:257532815.
  * Park et al. (2021) Sejun Park, Jaeho Lee, Chulhee Yun, and Jinwoo Shin.  Provable memorization via deep neural networks using sub-linear parameters.  In Mikhail Belkin and Samory Kpotufe (eds.), _Proceedings of Thirty Fourth Conference on Learning Theory_ , volume 134 of _Proceedings of Machine Learning Research_ , pp. 3627–3661. PMLR, 15–19 Aug 2021.  URL https://proceedings.mlr.press/v134/park21a.html.
  * Radford et al. (2019) Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever.  Language models are unsupervised multitask learners.  2019\.
  * Rajput et al. (2021) Shashank Rajput, Kartik Sreenivasan, Dimitris Papailiopoulos, and Amin Karbasi.  An exponential improvement on the memorization capacity of deep threshold networks.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), _Advances in Neural Information Processing Systems_ , volume 34, pp. 12674–12685. Curran Associates, Inc., 2021.
  * Ramsauer et al. (2020) Hubert Ramsauer, Bernhard Schafl, Johannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gruber, Markus Holzleitner, Milena Pavlovi’c, Geir Kjetil Ferkingstad Sandve, Victor Greiff, David P. Kreil, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter.  Hopfield networks is all you need.  _ArXiv_ , abs/2008.02217, 2020.  URL https://api.semanticscholar.org/CorpusID:220968978.
  * Sahiner et al. (2022) Arda Sahiner, Tolga Ergen, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pilanci.  Unraveling attention via convex duality: Analysis and interpretations of vision transformers.  In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pp. 19050–19088. PMLR, 17–23 Jul 2022.  URL https://proceedings.mlr.press/v162/sahiner22a.html.
  * Sontag (1997) Eduardo D. Sontag.  Shattering All Sets of ‘k’ Points in “General Position” Requires (k — 1)/2 Parameters.  _Neural Computation_ , 9(2):337–348, 02 1997\.  ISSN 0899-7667.  doi: 10.1162/neco.1997.9.2.337.  URL https://doi.org/10.1162/neco.1997.9.2.337.
  * Sukhbaatar et al. (2015) Sainbayar Sukhbaatar, Jason Weston, Rob Fergus, et al.  End-to-end memory networks.  _Advances in neural information processing systems_ , 28, 2015.
  * Touvron et al. (2023) Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al.  Llama: Open and efficient foundation language models.  _arXiv preprint arXiv:2302.13971_ , 2023.
  * Vardi et al. (2022) Gal Vardi, Gilad Yehudai, and Ohad Shamir.  On the optimal memorization power of reLU neural networks.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=MkTPtnjeYTV.
  * Vaswani et al. (2017) Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In _NIPS_ , 2017.
  * Vershynin (2020) Roman Vershynin.  Memory capacity of neural networks with threshold and rectified linear unit activations.  _SIAM J. Math. Data Sci._ , 2:1004–1033, 2020.  URL https://api.semanticscholar.org/CorpusID:226352137.
  * von Oswald et al. (2022) Johannes von Oswald, Eyvind Niklasson, E. Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov.  Transformers learn in-context by gradient descent.  _ArXiv_ , abs/2212.07677, 2022.
  * Wolf et al. (2020) Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush.  Huggingface’s transformers: State-of-the-art natural language processing, 2020.
  * Yun et al. (2019) Chulhee Yun, Suvrit Sra, and Ali Jadbabaie.  Small relu networks are powerful memorizers: a tight analysis of memorization capacity.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.
  * Yun et al. (2020a) Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar.  Are transformers universal approximators of sequence-to-sequence functions?  In _International Conference on Learning Representations_ , 2020a.  URL https://openreview.net/forum?id=ByxRM0Ntvr.
  * Yun et al. (2020b) Chulhee Yun, Yin-Wen Chang, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar.  O(n) connections are expressive enough: Universal approximability of sparse transformers.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 13783–13794. Curran Associates, Inc., 2020b.  URL https://proceedings.neurips.cc/paper_files/paper/2020/file/9ed27554c893b5bad850a422c3538c15-Paper.pdf.
  * Zhang et al. (2017) Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals.  Understanding deep learning requires rethinking generalization.  In _International Conference on Learning Representations_ , 2017.  URL https://openreview.net/forum?id=Sy8gdB9xx.
  * Zhang et al. (2021) Jiawei Zhang, Yushun Zhang, Mingyi Hong, Ruoyu Sun, and Zhi-Quan Luo.  When expressivity meets trainability: Fewer than $n$ neurons can work.  _Advances in Neural Information Processing Systems_ , 34:9167–9180, 2021.

Symbol | Description
---|---
$\mathcal{A}$ | A single-layer multi-head Attention (MHA) mechanism
$\mathcal{T}$ | Training set
$H$ | Number of heads in the MHA layer
$T$ | Training set size
$n$ | Context length
${\bm{W}_{K}}_{h},{\bm{W}_{Q}}_{h}\in\mathbb{R}^{{d}\times{d_{\text{h}}}},{\bm{W}_{V}}_{h}\in\mathbb{R}^{{d}\times{d_{v}}}$ | Weight matrices for each head $h\in[H]$
${\bm{W}_{O}}\in\mathbb{R}^{H{d_{v}}\times{d}}$ | Weight matrix to combine the outputs of different heads
${\bm{W}_{D}}\in\mathbb{R}^{d}$ | Final read-out weight of the model
${\bm{\mathcal{W}}}$ | The entire set of parameters of the MHA layer
$\bm{E}\in\mathbb{R}^{n\times{d}}$ | Key/Context matrix containing $n$ key/context vectors
${\bm{e}}\in\mathbb{R}^{{d}}$ | Query vector
$\bm{y}\in\mathbb{R}$ | Ground truth label
$\bm{\alpha}_{h},\bm{\theta}_{h},\bm{z}_{h},\bm{p}_{h}$ | Intermediate variables in the MHA layer
$\bm{o}$ | Output of the MHA layer
$\hat{\bm{y}}\in\mathbb{R}$ | Final output of the model
Table 2: Important notations used in the paper.

##  Appendix A Experiment details

In the synthetic experiment, for regression tasks, we generate uniform labels
$y^{(t)}\sim\mathcal{U}(0,1)$, and for classification tasks, we draw uniform
labels from $100$ classes $y^{(t)}\in\mathcal{U}([100])$. For classification,
we measure memorization by computing the average accuracy across examples in a
dataset. For regression, we use the mean squared error (MSE) loss to compare
memorization abilities. We train each task for at least $500$ epochs and at
least $50,000$ optimization steps (whichever is larger) with a batch size of
$256$, using an Adam optimizer with a learning rate of $0.001$, and a
scheduler with linear warmup and Cosine decay. For the experiments reported in
Figure 2(c) and Figure 3 since the optimization is more challenging, we train
for a total of $300,000$ optimization steps.

While our experimental results in Section 5.2 conclusively show that
increasing the number of heads $H$ or context size $n$ improves the model’s
memorization capability, we caution the reader that the experiments are done
purely through gradient descent optimization, which has limitations when it
comes to studying memorization. First, it is not guaranteed that if a
memorizing solution exists, GD would be able to find it. Second, optimization
requires hyperparameter tuning and a large number of gradient updates to
achieve memorization, especially in under-parameterized regimes and large-
context sizes (Garg et al., 2022). We mitigate these issues by running small-
scale experiments with a sufficiently high number of optimization steps and
adequate learning rate scheduling. Running the experiments reported in our
paper takes approximately 20 GPU days on an NVIDIA V100 GPU. However, this
procedure becomes highly resource-exhaustive and prohibitive in larger-scale
experiments. For instance, we find out that having large $d_{h}$ (beyond $n$)
can generally improve optimization speed, and lead to better memorization.
However, with sufficient large steps, this gap closes, and gains in
memorization vanish in the $d_{h}>n$ regime (as observed in Figure 2(c)).

##  Appendix B Proof of Proposition 1

For convenience let us define $r=\min(n,d_{h})$. WLOG let us assume that
$T\geq H(r-1)+1$. This is because for the case of $T<H(r-1)+1$ we can add
extra dummy data points that satisfy Assumptions 1 and 2 and make
$T=H(r-1)+1$, then prove $\bm{Z}$ has full row-rank, and finally remove the
extra rows of $\bm{Z}$ (i.e. dummy data points) to prove the statement.

We prove the lemma by induction on $H$. We will demonstrate that by adding
each head, we can increase the rank of $\bm{Z}$ by $r-1$. The key to proving
the induction is the following two claims:

######  Claim 1 (Fit $r$ examples in one head).

Assume that $r$ examples
$\left\\{\left(\bm{E}^{(t)},{\bm{e}}^{(t)}\right)\right\\}_{t=1}^{r}$ satisfy
Assumptions 1 and 2. Fix any head index $h\in[H]$. For any arbitrary set of
convex combination vectors $\bm{\theta}^{(t)}_{h,*}~{}\forall t\in[r]$, there
exists an Attention head with key/query weights
${\bm{W}_{K}}_{h},{\bm{W}_{Q}}_{h}$ which achieves such coefficients for the
set of $r$ data points.

######  Claim 2 (Fit $r-1$ examples in one head, leave the rest intact).

Assume that $T^{\prime}\geq r$ examples
$\left\\{\left(\bm{E}^{(t)},{\bm{e}}^{(t)}\right)\right\\}_{t=1}^{T^{\prime}}$
satisfy Assumptions 1 and 2, and define
$\mathcal{R}=\\{T^{\prime}-r+2,\dots,T^{\prime}-1,T^{\prime}\\}$ as the set of
last $r-1$ examples’ indices and $\mathcal{R}^{c}=[T]-\mathcal{R}$ as its
complement for convenience. Fix any head index $h\in[H]$. Then, there exist
$\bm{\theta}_{+}^{(t)},t\in\mathcal{R}^{c}$ such that for any arbitrary set of
desired convex combination vectors $\bm{\theta}_{*}^{(t)},t\in\mathcal{R}$ and
any $\delta>0$, there exists key/query weights
${\bm{W}_{K}}_{h},{\bm{W}_{Q}}_{h}$ such that both
$\bm{\theta}^{(t)}_{h}=\bm{\theta}_{*}^{(t)},\forall t\in\mathcal{R}$ and
$\left\lVert\bm{\theta}^{(t)}_{h}-\bm{\theta}_{+}^{(t)}\right\rVert_{1}\leq\delta,t\in\mathcal{R}^{c}$.

Claim 1 states that we can design an attention head that produces any
arbitrary $\bm{\theta}_{h}$ coefficients for $r$ examples. Furthermore, Claim
2 reinforces Claim 1 by demonstrating the possibility of tuning coefficients
$\bm{\theta}_{h}$ of $r-1$ examples while keeping the rest of the examples
approximately fixed with any precision $\delta$. However, this comes at the
cost of fitting one less example into a single head. The proofs of both claims
will be deferred until the end of the proposition. We are now ready to prove
the proposition by induction on $H$.

_Step 1. The base of induction $H=1$._ First, we prove the base of induction
for $H=1$. Let us only consider the first $r$ examples (i.e., $T^{\prime}=r$).
Notice that the matrix $\bm{Z}$ has $r$ rows and $d$ columns in this case.
Moreover, the $t$-th row of $\bm{Z}$ has the form

$\bm{Z}_{t}=\bm{E}^{{(t)}^{\top}}\bm{\theta}_{1}^{(t)}=\begin{bmatrix}{\bm{e}}_{1}^{(t)},{\bm{e}}_{2}^{(t)},\dots,{\bm{e}}_{n}^{(t)}\end{bmatrix}\bm{\theta}_{1}^{(t)},$

and is a convex combination of the context tokens for the example $t$.
According to Assumption 2, the span of context vectors ${\bm{e}}_{[n]}^{(t)}$
has dimension $n$. Therefore, we can inductively find
$\bm{\theta}_{1}^{(1)},\bm{\theta}_{1}^{(2)},\dots,\bm{\theta}_{1}^{(r)}$ such
that each row $\bm{Z}_{t}$ is not in the span of the previous $t-1$ rows. We
can continue this process for all $t\leq r$ since the span of the previous
rows has at most $t-1$ dimensions, but the convex hull of
${\bm{e}}_{[n]}^{(t)}$ has at least $n$ linearly independent elements ($r\leq
n$), so it is not covered by the previous $t-1$ rows. Therefore, there exists
$\\{\bm{\theta}_{1}^{(t)}\\}_{t=1}^{r}$ such that the resulting $\bm{Z}$
matrix has rank at least $r$. Then, using the result of Claim 1, there exists
Attention weights ${\bm{W}_{K}}_{1},{\bm{W}_{Q}}_{1}$ that produce the desired
set of convex combinations $\\{\bm{\theta}_{1}^{(t)}\\}_{t=1}^{r}$, and
subsequently the first $r$ rows of the desired feature matrix $\bm{Z}$.

_Step 2. Induction Step._ Next, assume that for $H=H^{\prime}$, there exists
heads with Attention weights
$\\{({\bm{W}_{K}}_{h},{\bm{W}_{Q}}_{h})\\}_{h=1}^{H^{\prime}}$ such that
matrix $\bm{Z}$ has rank $H^{\prime}(r-1)+1$. We prove for $H=H^{\prime}+1$
there exists Attention weights
$\\{({\bm{W}_{K}}_{h},{\bm{W}_{Q}}_{h})\\}_{h=1}^{H^{\prime}+1}$ such that the
resulting matrix $\bm{Z}$ attains rank $(H^{\prime}+1)(r-1)+1$. Let us refer
to the matrices derived from the induction hypothesis by a prime superscript
(i.e., $\bm{Z}^{\prime}$, and
$\\{({\bm{W}_{K}}_{h}^{\prime},{\bm{W}_{Q}}_{h}^{\prime})\\}_{h=1}^{H^{\prime}}$).
Now, for the case of $H=H^{\prime}+1$, we set the weights of the first
$H^{\prime}$ Attention heads to be equal to that of $H^{\prime}$. Namely,

${\bm{W}_{K}}_{h}:={\bm{W}_{K}}^{\prime}_{h},\quad{\bm{W}_{Q}}_{h}:={\bm{W}_{Q}}^{\prime}_{h}$

for all $h\in[H^{\prime}]$. Then, the feature matrix $\bm{Z}$ can be written
as the following block form:

$\bm{Z}=\left[\begin{array}[]{ccc|c}&&&\\\
&\bm{Z}^{\prime}&&\bm{Z}^{\prime\prime}\\\
\makebox[0.0pt][l]{$\smash{\underbrace{\phantom{\begin{matrix}&&&\end{matrix}}}_{\text{$H^{\prime}d$}}}$}&&&\makebox[0.0pt][l]{$\smash{\underbrace{\phantom{\begin{matrix}~{}~{}~{}~{}~{}~{}\end{matrix}}}_{\text{$d$}}}$}{}{}{}{}{}{}\\\
\end{array}\right],$

where $\bm{Z}^{\prime}\in\mathbb{R}^{T\times dH^{\prime}}$ is the matrix
derived from the hypothesis and $\bm{Z}^{\prime\prime}\in\mathbb{R}^{T\times
d}$ is the block corresponding to the new head $H^{\prime}+1$, whose
parameters we have not yet determined. Now, we know
$\operatorname{rank}(\bm{Z}^{\prime})\geq H^{\prime}(r-1)+1$. Define
$k:=\operatorname{rank}(\bm{Z}^{\prime})$. WLOG, assume that the first $k$
rows of the matrix $\bm{Z}^{\prime}$ are linearly independent (otherwise, we
can rearrange the rows of $\bm{Z}$, i.e. the examples, to satisfy this
property). Now, take the last $r-1$ rows that are linearly dependent on the
first $k$ rows. We defer the case for when less than $r-1$ such rows exist
(i.e., $k+r-1>T$) to the later stages of the proof.

Let us partition the matrix $\bm{Z}$ into two blocks: the first $T-r+1$ rows
and the last $r-1$ rows. Namely,

$\displaystyle\bm{Z}=\left[\begin{array}[]{c}\bm{Z}_{a}\\\
\hline\cr\bm{Z}_{b}\\\ \end{array}\right]=\left[\begin{array}[]{ccc|c}&&&\\\
&\bm{Z}_{a}^{\prime}&&\bm{Z}_{a}^{\prime\prime}\\\ &&&\\\
\hline\cr&\bm{Z}_{b}^{\prime}&&\bm{Z}_{b}^{\prime\prime}\\\
\end{array}\right].$ (15)

Observe that the rows of $\bm{Z}_{b}^{\prime}$ are all linear combinations of
rows of $\bm{Z}_{a}^{\prime}$. So, we can zero out the $\bm{Z}_{b}^{\prime}$
block by adding a linear combination of rows of $\bm{Z}_{a}^{\prime}$. Such
operations add a combination of rows of $\bm{Z}_{a}^{\prime\prime}$ to the
$\bm{Z}_{b}^{\prime\prime}$ as well. Notice that the rank of $\bm{Z}$ stays
the same since such operations do not change rank. So let us call
$\bm{Z}_{op}$ the resulting matrix:

$\bm{Z}_{op}=\left[\begin{array}[]{ccc|c}&&&\\\
&\bm{Z}_{a}^{\prime}&&\bm{Z}_{a}^{\prime\prime}\\\ &&&\\\
\hline\cr&\bm{0}&&\bm{Z}_{b}^{\prime\prime}+\bm{A}\\\ \end{array}\right],$

where

$\bm{A}\
:=\bm{A}(\bm{Z}^{\prime}_{a},\bm{Z}^{\prime}_{b},\bm{Z}^{\prime\prime}_{a}):=\bm{C}(\bm{Z}^{\prime}_{a},\bm{Z}^{\prime}_{b})\bm{Z}^{\prime\prime}_{a}\in\mathbb{R}^{(r-1)\times
d}$

contains a linear combination of rows of $\bm{Z}_{a}^{\prime\prime}$ (whose
coefficients are in $\bm{C}(\bm{Z}^{\prime}_{a},\bm{Z}^{\prime}_{b})$ and are
a function of $\bm{Z}^{\prime}_{a},\bm{Z}^{\prime}_{b}$). Since

$\operatorname{rank}(\bm{Z})=\operatorname{rank}(\bm{Z}_{op})\stackrel{{\scriptstyle\text{Lemma
\ref{lemma:rank_block}}}}{{\geq}}\operatorname{rank}(\bm{Z}_{b}^{\prime\prime}+\bm{A})+\operatorname{rank}(\bm{Z}_{a}^{\prime})\geq\operatorname{rank}(\bm{Z}_{b}^{\prime\prime}+\bm{A})+H^{\prime}(r-1)+1,$

we have to find suitable Attention weights such that
$\operatorname{rank}(\bm{Z}_{b}^{\prime\prime}+\bm{A})\geq r-1$ to achieve a
$(H^{\prime}+1)(r-1)+1$ rank for $\bm{Z}$ and complete the proof.

Now, consider Claim 2 for some fixed convex combination vectors
$\\{\bm{\theta}_{+}^{(t)}\\}_{t\in\mathcal{R}^{c}}$, and any desired
$\\{\bm{\theta}_{H^{\prime}+1}^{(t)}\\}_{t\in\mathcal{R}}$. For any
$\delta>0$, the claim statement gives key/query weights, and subsequently a
resulting intermediate matrix $\bm{Z}_{\delta}$. Let us define $\bm{Z}_{+}$ as
the _limiting matrix_ $\lim_{\delta\rightarrow 0^{+}}\bm{Z}_{\delta}$, and
similarly define its sub-blocks corresponding to this head by the same sign,
i.e., $\bm{Z}^{\prime\prime}_{a,+},\bm{Z}^{\prime\prime}_{b,+},\bm{A}_{+}$,
and similarly
$\bm{Z}^{\prime\prime}_{a,\delta},\bm{Z}^{\prime\prime}_{b,\delta},\bm{A}_{\delta}$.
Notice that $\bm{A}_{+}$ is only a function of
$\bm{Z}^{\prime}_{a},\bm{Z}^{\prime}_{b},\bm{Z}^{\prime\prime}_{a,+}$. As a
result, different choices of
$\\{\bm{\theta}_{H^{\prime}+1}^{(t)}\\}_{t\in\mathcal{R}}$ produce the same
$\bm{A}_{+}$ since $\\{\bm{\theta}_{+}^{(t)}\\}_{t\in\mathcal{R}^{c}}$ remains
fixed under Claim 2. Therefore, we first find
$\bm{\theta}^{(t)}_{H^{\prime}+1}$ such that
$\operatorname{rank}(\bm{Z}^{\prime\prime}_{b,+}+\bm{A}_{+})\geq r-1$, then we
find small enough $\delta$ that properly approximates the limiting matrices.
Similar to the base of induction (Step 1), we inductively build the individual
rows of $\bm{Z}^{\prime\prime}_{b,+}$, but the difference is that the original
matrix $\bm{Z}^{\prime\prime}_{b,+}$ is added to a fixed matrix $\bm{A}_{+}$.
Hence, at each inductive step $t$, the following modified constraint must be
satisfied:

$\left[\bm{Z}^{\prime\prime}_{b,+}+\bm{A}_{+}\right]_{t}\notin\operatorname{span}\left\\{\left[\bm{Z}^{\prime\prime}_{b,+}+\bm{A}_{+}\right]_{1},\left[\bm{Z}^{\prime\prime}_{b,+}+\bm{A}_{+}\right]_{2},\dots,\left[\bm{Z}^{\prime\prime}_{b,+}+\bm{A}_{+}\right]_{t-1}\right\\},$

which can be satisfied by a stronger constraint:

$\left[\bm{Z}^{\prime\prime}_{b,+}\right]_{t}\notin\operatorname{span}\left\\{\left[\bm{A}_{+}\right]_{t},\left[\bm{Z}^{\prime\prime}_{b,+}+\bm{A}_{+}\right]_{1},\left[\bm{Z}^{\prime\prime}_{b,+}+\bm{A}_{+}\right]_{2},\dots,\left[\bm{Z}^{\prime\prime}_{b,+}+\bm{A}_{+}\right]_{t-1}\right\\}.$

Hence, since the input context tokens are in an $n$ dimensional space, the set
of feasible valid rows $[\bm{Z}^{\prime\prime}_{b,+}]_{t}$ lies in an $n$
dimensional space as well, which allows us to continue for $n-1,(n\geq r)$
rows (and not $n$, due to presence of $\left[\bm{A}_{+}\right]_{t}$), find
suitable $\\{\bm{\theta}^{{(t)}}_{H^{\prime}+1}\\}_{t\in\mathcal{R}}$, and
build a full-rank matrix $\bm{Z}^{\prime\prime}_{b,+}$ with
$\operatorname{rank}(\bm{Z}^{\prime\prime}_{b,+}+\bm{A}_{+})\geq r-1$. Notice
that no well-defined key/weights generate $\bm{Z}_{+}$ as the weights need to
be infinite (see proof of Claim 2). However, we only need a small enough
$\delta$ such that

$\displaystyle\operatorname{rank}(\bm{Z}^{\prime\prime}_{b,\delta}+\bm{A}_{\delta})\geq\operatorname{rank}(\bm{Z}^{\prime\prime}_{b,+}+\bm{A}_{+})\geq
r-1$ (16)

for the weights to be well-defined. In Appendix B.2 we prove the existence of
such $\delta$. Therefore, there exists
${\bm{W}_{K}}_{H^{\prime}+1},{\bm{W}_{Q}}_{H^{\prime}+1}$ that produces
$\operatorname{rank}(\bm{Z}_{b}^{\prime\prime}+\bm{A})\geq r-1$, and the proof
is complete.

_The case of less than $r-1$ linearly dependent rows of $\bm{Z}^{\prime}$._ In
this case, we have $k+r-1>T$. Let us define $m:=T-k<r-1$. Since we have
assumed $T\geq(H^{\prime}+1)(r-1)+1$, this means that that $\bm{Z}^{\prime}$
already has a rank at least $(H^{\prime}+1)(r-1)+1-m$. So, the head index
$H^{\prime}+1$ needs to increase the rank of $\bm{Z}^{\prime}$ by $m$. Since
$m\leq r-1$, all the above steps also apply to $m$ and the rest of the proof
procedure is the same.

###  B.1 Proof of Claims 1 and 2

First, note each $\bm{\theta}^{(t)}_{h,*}$ is derived from applying Softmax
nonlinearity on the logits $\bm{\alpha}^{(t)}_{h,*}$. So, we need to solve a
system of equations for $\bm{\alpha}^{(t)}_{h,*}$ for each $t\in[T]$444The set
of logits that produce a convex combination is not necessarily unique.
However, uniqueness is not a requirement in this case, as we only require the
existence of some set of logits that produces a convex combination.. Since our
focus is on a _single_ head, we omit the subscript $h$ for simplicity.
Furthermore, we define $\bm{W}:={\bm{W}_{K}}{\bm{W}_{Q}}^{\top}$. Now, for
data indices $t_{1}<t_{2}<\dots<t_{m}\in[T]$ where $m\leq n$, we are
interested in solving the following system of equations:

$\displaystyle\bm{E}^{(t_{i})}\bm{W}{\bm{e}}^{(t_{i})}=\bm{\alpha}^{(t_{i})}_{*}\quad\forall
i\in[m].$

Since each $\bm{E}^{(t_{i})}\in\mathbb{R}^{n\times d}$ is full-rank with rank
$n$, it has a pseudo right inverse
$\bm{E}^{{(t_{i})}^{\dagger}}\in\mathbb{R}^{d\times n}$ such that
$\bm{E}^{{(t_{i})}}\bm{E}^{(t_{i})^{\dagger}}=\bm{I}$ . So, we solve the
following system of equations

$\displaystyle\bm{W}{\bm{e}}^{(t_{i})}=\bm{E}^{{(t_{i})}^{\dagger}}\bm{\alpha}^{(t_{i})}_{*}\quad\forall
i\in[m],$ $\displaystyle\Rightarrow$
$\displaystyle\bm{E}^{(t_{i})}\bm{W}{\bm{e}}^{(t_{i})}=\bm{\alpha}^{(t_{i})}_{*}\quad\forall
i\in[m].$

Now, let us define
$\bm{\beta}^{(t_{i})}_{*}:=\bm{E}^{{(t_{i})}^{\dagger}}\bm{\alpha}^{(t_{i})}_{*}$.
Then, we get the following linear system of equations to solve:

$\bm{W}\begin{bmatrix}{\bm{e}}^{(t_{1})},&{\bm{e}}^{(t_{2})},&\dots,&{\bm{e}}^{(t_{m})}\\\
\end{bmatrix}=\begin{bmatrix}\bm{\beta}^{(t_{1})}_{*},&\bm{\beta}^{(t_{2})}_{*},&\dots,&\bm{\beta}^{(t_{m})}_{*}\\\
\end{bmatrix}.$

Due to Assumption 1 the query tokens have Kruskal Rank at least $n$, and
$m\leq r\leq n$, so the query token matrix in the left-hand-side above has a
pseudo left inverse. Thus, $\bm{W}$ has at least one solution with rank at
most $m$. Assuming the solution is $\bm{W}^{*}\in\mathbb{R}^{d\times d}$,
there exists ${\bm{W}_{K}}_{h},{\bm{W}_{Q}}_{h}\in\mathbb{R}^{d\times d_{h}}$
such that ${\bm{W}_{K}}_{h}{\bm{W}_{Q}}_{h}^{\top}=\bm{W}^{*}$ for the weights
of the Attention head. So, we can prove Claim 1 by setting $m=r$ and
$(t_{1},t_{2},\dots t_{r})=(1,2,\dots,r)$.

Next, to prove Claim 2, we prove that if $m<r\leq n$,
$(t_{1},t_{2},\dots,t_{m})=(T-m+1,T-m+2,\dots,T)$, then we can not only
achieve any desired logits for the examples $(T-m+1,T-m+2,\dots,T)$, but also
keep the $\\{\bm{\theta}^{{(t)}}\\}_{t=1}^{T-m}$ approximately unchanged. So
far, we have found a $\bm{W}^{*}$ that achieves the set of desired logits.
Now, there exists a rank-one linear transformation
$\bm{W}^{+}\in\mathbb{R}^{d\times d}$ such that

$\displaystyle\bm{E}^{(t)}\bm{W}^{+}{\bm{e}}^{(t)}$
$\displaystyle=\bm{0}\quad\text{for all }T-m<t\leq T,$ (17)
$\displaystyle\left[\bm{E}^{(t)}\bm{W}^{+}{\bm{e}}^{(t)}\right]_{i}$
$\displaystyle\neq\left[\bm{E}^{(t)}\bm{W}^{+}{\bm{e}}^{(t)}\right]_{j}\quad\text{for
all }1\leq t\leq T-m,~{}i,j\in[n],i\neq j.$ (18)

This is due to Lemma 3.

So, we set $\bm{W}:=\bm{W}^{*}+c\bm{W}^{+}$, and choose a large enough $c$.
This way, the desired logits for $T-m<t\leq T$ stay the same since
$c\bm{W}^{+}$ has no effect on them, while for the rest of logits,
$c\bm{W}^{+}$ dominates $\bm{W}^{*}$ and saturates the Softmax operator. Lemma
4 proves the existence of such $c$, so, we can prove Claim 2 by finding
${\bm{W}_{K}}_{h},{\bm{W}_{Q}}_{h}\in\mathbb{R}^{d\times d_{h}}$ such that
${\bm{W}_{K}}_{h}{\bm{W}_{Q}}_{h}^{\top}=\bm{W}^{*}+c\bm{W}^{+}$. Note that
this construction is feasible because
$\operatorname{rank}(\bm{W}^{*}+c\bm{W}^{+})\leq m+1\leq r\leq d_{h}$.

###  B.2 Proof of existence of $\delta$ for Eq. (16)

Let us call the $\bm{\theta}^{(t)}_{h}$ derived from Claim 2 as
$\bm{\theta}_{\delta}^{(t)}$. So, for any $\delta>0$ we have:

$\displaystyle\left\lVert\bm{\theta}_{\delta}^{(t)}-\bm{\theta}_{+}^{(t)}\right\rVert_{1}$
$\displaystyle<\delta,\text{ for all }t\in\mathcal{R}^{c}$
$\displaystyle\bm{\theta}_{\delta}^{(t)}=\bm{\theta}_{+}^{(t)}$
$\displaystyle=\bm{\theta}_{*}^{(t)},\text{ for all }t\in\mathcal{R}$
$\displaystyle\Rightarrow\bm{Z}^{\prime\prime}_{b,\delta}$
$\displaystyle=\bm{Z}^{\prime\prime}_{b,+}.$

To prove the rank inequality argument, we use a corollary of the Weyl’s
singular value theorem (Lemma 5) and prove:

$\displaystyle\Delta_{\delta}:=\left(\bm{Z}^{\prime\prime}_{b,\delta}+\bm{A}_{\delta}\left)-\right(\bm{Z}^{\prime\prime}_{b,+}+\bm{A}_{+}\right),\quad\sigma_{\max}\left(\Delta_{\delta}\right)<\sigma_{\min}(\bm{Z}^{\prime\prime}_{b,+}+\bm{A}_{+}),$
(19)

which in turn leads to
$\operatorname{rank}(\bm{Z}^{\prime\prime}_{b,\delta}+\bm{A}_{\delta})\geq\operatorname{rank}(\bm{Z}^{\prime\prime}_{b,+}+\bm{A}_{+})$.
First, let us assume $\bm{Z}^{\prime\prime}_{b,+}+\bm{A}_{+}$ is not a zero
matrix (otherwise we must have $r-1=0$, in which case any $\delta>0$ would
satisfy the inequality). We now bound the spectral norm of the perturbation
matrix $\Delta_{\delta}$, i.e., $\sigma_{\max}(\Delta_{\delta})$.

$\displaystyle\sigma_{\max}(\Delta_{\delta})$
$\displaystyle=\left\lVert\Delta_{\delta}\right\rVert_{2}$
$\displaystyle=\left\lVert\bm{A}_{\delta}-\bm{A}_{+}\right\rVert_{2}$
($\bm{Z}^{\prime\prime}_{b,+}=\bm{Z}^{\prime\prime}_{b,\delta}$)
$\displaystyle=\left\lVert\bm{C}(\bm{Z}^{\prime}_{a},\bm{Z}^{\prime}_{b})\bm{Z}^{\prime\prime}_{a,\delta}-\bm{C}(\bm{Z}^{\prime}_{a},\bm{Z}^{\prime}_{b})\bm{Z}^{\prime\prime}_{a,+}\right\rVert_{2}$
$\displaystyle=\left\lVert\bm{C}(\bm{Z}^{\prime}_{a},\bm{Z}^{\prime}_{b})\left(\bm{Z}^{\prime\prime}_{a,\delta}-\bm{Z}^{\prime\prime}_{a,+}\right)\right\rVert_{2}$
$\displaystyle\leq\left\lVert\bm{C}(\bm{Z}^{\prime}_{a},\bm{Z}^{\prime}_{b})\right\rVert_{2}\left\lVert\bm{Z}^{\prime\prime}_{a,\delta}-\bm{Z}^{\prime\prime}_{a,+}\right\rVert_{2}$
$\displaystyle\leq\left\lVert\bm{C}(\bm{Z}^{\prime}_{a},\bm{Z}^{\prime}_{b})\right\rVert_{2}\left\lVert\bm{Z}^{\prime\prime}_{a,\delta}-\bm{Z}^{\prime\prime}_{a,+}\right\rVert_{F},$

where $\bm{C}(\bm{Z}^{\prime}_{a},\bm{Z}^{\prime}_{b})$ is a constant matrix
containing coefficients derived from zeroing out $\bm{Z}^{\prime}_{b}$ using
the rows of $\bm{Z}^{\prime}_{a}$. To bound the Frobenius norm of
$\bm{Z}^{\prime\prime}_{a,\delta}-\bm{Z}^{\prime\prime}_{a,+}$, notice that
$t$-th row of this matrix is:

$\displaystyle\left\lVert\left[\bm{Z}^{\prime\prime}_{a,\delta}-\bm{Z}^{\prime\prime}_{a,+}\right]_{t}\right\rVert_{2}$
$\displaystyle=\left\lVert{\bm{E}^{(t)}}^{\top}\bm{\theta}_{\delta}^{(t+T^{\prime})}-{\bm{E}^{(t)}}^{\top}\bm{\theta}_{+}^{(t)}\right\rVert_{2}$
$\displaystyle=\left\lVert{\bm{E}^{(t)}}^{\top}\left(\bm{\theta}_{\delta}^{(t)}-\bm{\theta}_{+}^{(t)}\right)\right\rVert_{2}$
$\displaystyle\leq\left\lVert{\bm{E}^{(t)}}^{\top}\right\rVert_{1,2}\left\lVert\left(\bm{\theta}_{\delta}^{(t)}-\bm{\theta}_{+}^{(t)}\right)\right\rVert_{1}$
$\displaystyle\leq e_{\max}\delta,$

where $e_{\max}$ is the maximum token Euclidian norm of all tokens in the
dataset. Therefore:

$\left\lVert\bm{Z}^{\prime\prime}_{a,\delta}-\bm{Z}^{\prime\prime}_{a,+}\right\rVert_{F}\leq\sqrt{T^{\prime}}e_{\max}\delta.$

Hence, choosing

$\delta\leq\frac{\sigma_{\min}(\bm{Z}^{\prime\prime}_{b,+}+\bm{A}_{+})}{\sqrt{T^{\prime}}e_{\max}\left\lVert\bm{C}(\bm{Z}^{\prime}_{a},\bm{Z}^{\prime}_{b})\right\rVert_{2}}$

gives the desired inequality in Eq. (19).

##  Appendix C Extension of Theorem 1 to Attention with skip connection

Theorem 1 holds true for when we have Attention with skip connection. Recall
that skip connection keeps all the computational mechanisms of MHA the same
except for Eq. (5), where the updated equation is:

$\displaystyle\bm{o}$
$\displaystyle:={\bm{W}_{O}}^{\top}\begin{bmatrix}\bm{p}_{1};\bm{p}_{2};\dots;\bm{p}_{H}\end{bmatrix}+{\bm{e}}$
$\displaystyle(\bm{o}\in\mathbb{R}^{{d}})$

The only modification that we have to make, is to modify the labels. Let us
consider the setting of Theorem 1, with dataset
$\mathcal{T}=\left\\{\left(\bm{E}^{(t)},{\bm{e}}^{(t)},\bm{y}^{(t)}\right)\right\\}_{t=1}^{T}$
and Attention weights ${\bm{\mathcal{W}}}$. Now, if we add the skip connection
to the Attention layer with the same set of weights, the new
$\hat{\bm{y}}^{(t)}$ will be:

$\hat{\bm{y}}^{(t)}={\bm{W}_{D}}^{\top}({\bm{e}}^{(t)}+{\bm{o}}^{(t)})=\bm{I}_{{d_{\text{out}}}\times
d}({\bm{e}}^{(t)}+{\bm{o}}^{(t)})=\bm{y}^{(t)}+\bm{I}_{{d_{\text{out}}}\times
d}{\bm{e}}^{{(t)}}.$

Therefore, we get labels $\bm{y}^{(t)}+\bm{I}_{{d_{\text{out}}}\times
d}{\bm{e}}^{{(t)}}$, which is the true label shifted by the first
${d_{\text{out}}}$ dimensions of the query vector. So, we need to define a new
dataset

$\mathcal{T}^{\prime}=\left\\{\left(\bm{E}^{(t)},{\bm{e}}^{(t)},\bm{y}^{(t)}-\bm{I}_{{d_{\text{out}}}\times
d}{\bm{e}}^{{(t)}}\right)\right\\}_{t=1}^{T},$

and find Attention weights ${\bm{\mathcal{W}}}^{\prime}$ for
$\mathcal{T}^{\prime}$ according to Theorem 1. Then, we can use the set of
Attention weights ${\bm{\mathcal{W}}}^{\prime}$ to get the correct label
$\bm{y}^{(t)}+\bm{I}_{{d_{\text{out}}}\times
d}{\bm{e}}^{{(t)}}-\bm{I}_{{d_{\text{out}}}\times d}{\bm{e}}^{{(t)}}$ for an
Attention layer with skip connection and dataset $\mathcal{T}$.

##  Appendix D Additional Lemmas

######  Lemma 1.

For any real-valued block matrix ${\bm{X}}$ of the form

${\bm{X}}=\left[\begin{array}[]{cc}\bm{A}&\bm{C}\\\
\bm{0}&\bm{B}\end{array}\right]$

we have
$\operatorname{rank}({\bm{X}})\geq\operatorname{rank}(\bm{A})+\operatorname{rank}(\bm{B})$.

######  Lemma 2.

Assume vectors ${\bm{e}}_{1},{\bm{e}}_{2},\dots,{\bm{e}}_{n}\in\mathbb{R}^{d}$
are linearly independent. Then,

${\bm{e}}^{\prime}_{i}:={\bm{e}}_{i}+\frac{1}{n}({\bm{e}}_{1}+{\bm{e}}_{2}+\dots+{\bm{e}}_{n}),$

are also linearly independent for all $i\in[n]$.

###### Proof.

Due to the linear independent assumption, the only solution to

$a_{1}{\bm{e}}_{1}+a_{2}{\bm{e}}_{2}+\dots+a_{n}{\bm{e}}_{n}=\bm{0}\quad\text{{s.t.}}\quad
a_{1},a_{2},\dots,a_{n}\in\mathbb{R}$

is $a_{1}=a_{2}=\dots=a_{n}=0$. Now, let us examine the same system of
equations for ${\bm{e}}^{\prime}_{i}$. Namely,

$b_{1}{\bm{e}}^{\prime}_{1}+b_{2}{\bm{e}}^{\prime}_{2}+\dots+b_{n}{\bm{e}}^{\prime}_{n}=\bm{0}\quad\text{{s.t.}}\quad
b_{1},b_{2},\dots,b_{n}\in\mathbb{R}.$

We prove zero is the only solution to this system to conclude the proof.

$\displaystyle
b_{1}{\bm{e}}^{\prime}_{1}+b_{2}{\bm{e}}^{\prime}_{2}+\dots+b_{n}{\bm{e}}^{\prime}_{n}$
$\displaystyle=$ $\displaystyle
b_{1}{\bm{e}}_{1}+b_{2}{\bm{e}}_{2}+\dots+b_{n}{\bm{e}}_{n}+\underbrace{\frac{1}{n}(b_{1}+b_{2}+\dots+b_{n})}_{\bar{b}}({\bm{e}}_{1}+{\bm{e}}_{2}+\dots+{\bm{e}}_{n})$
$\displaystyle=$
$\displaystyle(b_{1}+\bar{b}){\bm{e}}_{1}+(b_{2}+\bar{b}){\bm{e}}_{2}+\dots+(b_{n}+\bar{b}){\bm{e}}_{n}=\bm{0}.$

According to the linear independence assumption, we must have

$b_{i}+\bar{b}=0\quad\text{for all }i\in[n].$

By summing all the terms up, we conclude that $b_{i}=0$ for all $i\in[n]$ is
the only solution and this concludes the proof. ∎

######  Lemma 3.

Consider a set of full-rank matrices
$\bm{E}^{(1)},\bm{E}^{(2)},\dots,\bm{E}^{(T)}\in\mathbb{R}^{n\times d}$ where
$n<d$. Moreover, consider a set of vectors
${\bm{e}}^{(1)},{\bm{e}}^{(2)},\dots,{\bm{e}}^{(T)}\in\mathbb{R}^{d}$ with a
Kruskal Rank at least $n$. Then, for any $m\leq n-1$ there exists a rank-one
matrix $\bm{W}\in\mathbb{R}^{d\times d}$ such that

$\displaystyle\bm{E}^{(t)}\bm{W}{\bm{e}}^{(t)}$
$\displaystyle=\bm{0}\quad\text{for all }1\leq t\leq m,$ (20)
$\displaystyle\left[\bm{E}^{(t)}\bm{W}{\bm{e}}^{(t)}\right]_{i}$
$\displaystyle\neq\left[\bm{E}^{(t)}\bm{W}{\bm{e}}^{(t)}\right]_{j}\quad\text{for
all }m<t\leq T,~{}i,j\in[n],i\neq j.$ (21)

###### Proof.

Let us define a basis for $\mathbb{R}^{d}$ as
${\bm{b}}_{1},{\bm{b}}_{2},{\bm{b}}_{3},\dots,{\bm{b}}_{d}\in\mathbb{R}^{d}$
such that ${\bm{b}}_{i}={\bm{e}}^{(i)}$ for all $i\in[m]$, and
${\bm{b}}_{m+1},\dots\bm{b}_{m+2},\dots,{\bm{b}}_{d}$ are an arbitrary
extension of ${\bm{e}}_{[m]}$ to a basis. Then, a linear transformation
$\bm{W}$ is uniquely determined by defining
${\bm{c}}_{1},{\bm{c}}_{2},\dots\bm{c}_{d}\in\mathbb{R}^{d}$ and setting

$\bm{W}{\bm{b}}_{i}={\bm{c}}_{i}\quad\text{for all }i\in[d].$

First, to satisfy Eq. (20), we set ${\bm{c}}_{i}:=\bm{0}$ for all $i\in[m]$,
which gives $\bm{W}{\bm{b}}_{i}=\bm{W}{\bm{e}}^{(i)}={\bm{c}}_{i}=\bm{0}$ for
all $i\in[m]$. Let us decompose every vector ${\bm{e}}^{(t)}$ into the
mentioned basis, i.e.,

${\bm{e}}^{(t)}=a_{1}^{(t)}{\bm{b}}_{1}+a_{2}^{(t)}{\bm{b}}_{2}+\dots+a_{d}^{(t)}{\bm{b}}_{d}\quad\text{for
all }m+1\leq t\leq T.$

Since $\\{{\bm{e}}^{(t)}\\}_{t=1}^{T}$ have a Kruskal Rank at least $n$, then
for any $t\in[T]$ at least one of
$a_{m+1}^{(t)},a_{m+2}^{(t)},\dots,a_{d}^{(t)}$ must be nonzero, otherwise,
${\bm{e}}^{(t)}$ can be written as a linear combination of the first $m$
vectors. Moreover, there exists fixed
$u_{m+1},u_{m+2},\dots,u_{d}\in\mathbb{R}$ such that

$a_{m+1}^{(t)}u_{m+1}+a_{m+2}^{(t)}u_{m+2}\dots+a_{d}^{(t)}u_{d}\neq
0\quad\text{ for all $m+1\leq t\leq d$},$

for instance, we can choose $u_{i}$ iid from a Standard Gaussian Distribution
and satisfy this property almost surely. Now, let us define

$\displaystyle{\bm{c}}$
$\displaystyle\sim\mathcal{N}\left(\bm{0},\bm{I}\right)$ (22)
$\displaystyle{\bm{c}}_{i}$ $\displaystyle:=u_{i}{\bm{c}}\quad\text{ for all
$m+1\leq i\leq d$.}$ (23)

Therefore,

$\displaystyle\bm{W}{\bm{e}}^{(t)}$
$\displaystyle=a_{1}^{(t)}\bm{W}{\bm{b}}_{1}+a_{2}^{(t)}\bm{W}{\bm{b}}_{2}+\dots+a_{d}^{(t)}\bm{W}{\bm{b}}_{d}$
(24)
$\displaystyle=a_{1}^{(t)}{\bm{c}}_{1}+a_{2}^{(t)}{\bm{c}}_{2}+\dots+a_{d}^{(t)}{\bm{c}}_{d}$
(25)
$\displaystyle=a_{m+1}^{(t)}{\bm{c}}_{m+1}+a_{m+2}^{(t)}{\bm{c}}_{m+2}\dots+a_{d}^{(t)}{\bm{c}}_{d}\quad$
(26)
$\displaystyle=\left(a_{m+1}^{(t)}u_{m+1}+a_{m+2}^{(t)}u_{m+2}\dots+a_{d}^{(t)}u_{d}\right){\bm{c}}\quad\text{for
all }m+1\leq t\leq T,$ (27)

with one nonzero coefficient. Now, since ${\bm{c}}$ is standard gaussian, then

$\bm{W}{\bm{e}}^{(t)}\sim\mathcal{N}\left(\bm{0},\left(a_{m+1}^{(t)}u_{m+1}+a_{m+2}^{(t)}u_{m+2}\dots+a_{d}^{(t)}u_{d}\right)^{2}\bm{I}\right).$

Moreover, $\bm{E}^{(t)}$ is full-rank, so none of its rows are equal to each
other. As a result,

$\left({\bm{e}}_{i}^{(t)}-{\bm{e}}_{j}^{(t)}\right)^{\top}\bm{W}{\bm{e}}^{(t)}$

follows a Gaussian distribution for all $i,j\in[n],~{}i\neq j$. Since
$\\{0\\}$ has zero Lebesgue measure, the probability of having two equal
elements is zero. So, Eq. 21 is satisfied almost surely. Note that this
construction produces a rank-one $\bm{W}$ since the rank of
$[{\bm{c}}_{1},{\bm{c}}_{2},\dots,{\bm{c}}_{d}]$ is at most one. Therefore,
there exists a rank-one matrix $\bm{W}$ that satisfies Eq. (21) and this
completes the proof.

∎

######  Lemma 4.

Consider an input matrix $\bm{E}\in\mathbb{R}^{n\times d}$, an input vector
${\bm{e}}\in\mathbb{R}^{d}$, and a weight matrix
$\bm{W}^{+}\in\mathbb{R}^{d\times d}$ such that

$\displaystyle\quad\left[\bm{E}\bm{W}^{+}{\bm{e}}\right]_{i}\neq\left[\bm{E}\bm{W}^{+}{\bm{e}}\right]_{j}\quad\text{for
all }i,j\in[n],i\neq j.$ (28)

For any $\bm{W}^{*}\in\mathbb{R}^{d\times d}$ define
$\bm{W}_{c}:=\bm{W}^{*}+c\bm{W}^{+}$ where $c\in\mathbb{R}$ is a scalar. Then,

$\bm{\theta}_{+}:=\lim_{c\rightarrow\infty}\bm{\theta}_{c}:=\lim_{c\rightarrow\infty}\operatorname{Softmax}\left(\left[\bm{E}\bm{W}_{c}{\bm{e}}\right]\right),$

is only a function of $\bm{W}^{+}$. Consequently, $\bm{\theta}_{c}$ converges
to $\bm{\theta}_{+}$ in $L^{1}$ sense, i.e., for any $\delta>0$, there exists
$c_{\delta}$ such that

$\left\lVert\bm{\theta}_{+}-\bm{\theta}_{c_{\delta}}\right\rVert_{1}<\delta.$

###### Proof.

Let us define

$\bm{\alpha}^{+}:=\bm{E}\bm{W}^{+}{\bm{e}},\quad\text{and}\quad\bm{\alpha}^{*}:=\bm{E}\bm{W}^{*}{\bm{e}}.$

Now, $i$th element of $\bm{\theta}$ is given by:

$\displaystyle[\bm{\theta}_{+}]_{i}$
$\displaystyle=\lim_{c\rightarrow\infty}[\bm{\theta}_{c}]_{i}=\lim_{c\rightarrow\infty}\frac{\exp\left(\bm{\alpha}^{*}_{i}+c\bm{\alpha}^{+}_{i}\right)}{\sum_{j\in[n]}\exp\left(\bm{\alpha}^{*}_{j}+c\bm{\alpha}^{+}_{j}\right)}$
$\displaystyle=\lim_{c\rightarrow\infty}\frac{\exp\left(\bm{\alpha}^{*}_{i}\right)}{\exp\left(\bm{\alpha}^{*}_{i}\right)+\sum_{j\neq
i}\exp\left(\bm{\alpha}^{*}_{j}+c(\bm{\alpha}^{+}_{j}-\bm{\alpha}^{+}_{i})\right)}$
$\displaystyle=\bm{1}\left[i=\arg\max\\{\bm{\alpha}^{+}_{1},\bm{\alpha}^{+}_{2},\dots,\bm{\alpha}^{+}_{n}\\}\right],$

where the last equality is due to $\bm{\alpha}^{+}_{i}$ being distinct.
Therefore, $[\bm{\theta}_{+}]_{i}$ is a one-hot vector determined by
$\bm{W}^{+}$. Consequently, since each individual entry of $\bm{\theta}_{c}$
converges to $\bm{\theta}_{+}$, we get convergence in $L_{1}$ and this
completes the proof. ∎

######  Lemma 5.

Let $\bm{A},\bm{B}\in\mathbb{R}^{n\times m}$ be two matrices satisfying

$\sigma_{\max}(\bm{B}-\bm{A})<\sigma_{\min}(\bm{A}),$

where $\sigma_{\min},\sigma_{\max}$ are the smallest and largest non-zero
singular values. Then, we have

$\operatorname{rank}(\bm{B})\geq\operatorname{rank}(\bm{A}).$

###### Proof.

From Weyl’s singular value theorem (Horn & Johnson, 1990), we know

$\left|\sigma_{k}(\bm{A}+\Delta)-\sigma_{k}(\bm{A})\right|\leq\sigma_{\max}(\Delta),$

for any perturbation matrix $\Delta$, and singular value index
$k\in[\min(n,m)]$. Setting $\Delta=\bm{B}-\bm{A}$ and using the lemma’s
assumption we have:

$\left|\sigma_{k}(\bm{B})-\sigma_{k}(\bm{A})\right|\leq\sigma_{\max}(\bm{B}-\bm{A})<\sigma_{\min}(\bm{A})\leq\sigma_{k}(\bm{A})$

for all $k$ such that $\sigma_{k}(\bm{A})>0$. Hence, for any non-zero singular
value $k$, $\bm{B}$ has a corresponding non-zero singular value and
$\operatorname{rank}(\bm{B})\geq\operatorname{rank}(\bm{A})$. ∎

##  Appendix E Proof of Proposition 2

First, note that when the contexts are shared, the matrix $\bm{Z}$ has the
following decomposition:

$\bm{Z}^{\top}=\left(\bm{I}_{H}\otimes\bm{E}^{T}\right)\bm{\Theta}^{\top},$

where $\otimes$ is the Kronecker product and
$\bm{\Theta}\in\mathbb{R}^{T\times Hn}$ is defined similar to $\bm{Z}$ (i.e.,
the block in $t$-th row and $h$-th column is
$\bm{\theta}^{{(t)}^{\top}}_{h}$). Notice that
$\bm{\Theta}\in\mathbb{R}^{T\times Hn}$, therefore, the rank of $\bm{Z}$ can
be at most $Hn$. However, we prove here that the column rank of $\bm{\Theta}$
is at most $H(n-1)+1$ due to the structure of $\bm{\theta}^{(t)}_{h}$ being
convex combination vectors. With a slight abuse of notation, let us write
$\bm{\Theta}$ as:

$\bm{\Theta}=\begin{bmatrix}\bm{\Theta}_{1},&\bm{\Theta}_{2},&\bm{\Theta}_{3}&\dots,&\bm{\Theta}_{H}\end{bmatrix},$

where each $\bm{\Theta}_{h}\in\mathbb{R}^{T\times n}$ corresponds block
belonging to head $h\in[H]$. Observe that $\bm{\Theta}_{h}$ is a right
stochastic matrix (i.e., rows sum up to one) for all $h\in[H]$. Thus, each row
of $\bm{\Theta}_{h}-\bm{\Theta}_{1}$ sums up to zero and the matrix has at
least one zero eigenvalue. Therefore,

$\displaystyle\operatorname{rank}(\bm{Z})$
$\displaystyle\leq\operatorname{rank}(\bm{\Theta})=\operatorname{rank}\left(\begin{bmatrix}\bm{\Theta}_{1},&\bm{\Theta}_{2},&\bm{\Theta}_{3},&\dots,&\bm{\Theta}_{H}\end{bmatrix}\right)$
$\displaystyle=\operatorname{rank}\left(\begin{bmatrix}\bm{\Theta}_{1},&\bm{\Theta}_{2}-\bm{\Theta}_{1},&\bm{\Theta}_{3}-\bm{\Theta}_{1},&\dots,&\bm{\Theta}_{H}-\bm{\Theta}_{1}\end{bmatrix}\right)$
$\displaystyle\leq\operatorname{rank}(\bm{\Theta}_{1})+\operatorname{rank}(\bm{\Theta}_{2}-\bm{\Theta}_{1})+\dots+\operatorname{rank}(\bm{\Theta}_{H}-\bm{\Theta}_{1})$
$\displaystyle\leq n+(n-1)+(n-1)+\dots+(n-1)=H(n-1)+1\,.$

Hence $\operatorname{rank}(\bm{Z})$ is upper bounded by $H(n-1)+1$ and this
concludes the proof.

##  Appendix F Proof of Proposition 3

We give a proof by contradiction, so let us assume
$T>(n+1)m/{d_{\text{out}}}+m+1$. Consider ReLU network $f$ as

$f_{d_{1},d_{2},{d_{\text{out}}}}(\bm{x};\bm{W}_{1},{\bm{b}}_{1},\bm{W}_{2},{\bm{b}}_{2}):=\bm{W}_{2}^{\top}\operatorname{ReLU}\left(\bm{W}_{1}^{\top}\bm{x}+{\bm{b}}_{1}\right)+{\bm{b}}_{2}.$

where $\bm{W}_{1}\in\mathbb{R}^{d_{2}\times
d_{1}},{\bm{b}}_{1}\in\mathbb{R}^{d_{2}},\bm{W}_{2}\in\mathbb{R}^{{d_{\text{out}}}\times
d_{2}}$, and ${\bm{b}}_{2}\in\mathbb{R}^{{d_{\text{out}}}}$ are the network
parameters. First, we consider $d_{1}=n,d_{2}=m$. This network possesses a
total of $(n+1)m+(m+1){d_{\text{out}}}$ parameters. Since a trivial upper
bound on the number of real-valued labels a network can memorize is at most as
high as the number of its parameters, for any

$T>\frac{(n+1)m+(m+1){d_{\text{out}}}}{{d_{\text{out}}}}=\frac{(n+1)m}{{d_{\text{out}}}}+m+1$

and inputs $\bm{x}^{(1)},\bm{x}^{(2)},\dots,\bm{x}^{(T)}\in\mathbb{R}^{n}$,
there exists labels
$\bm{y}^{(1)},\bm{y}^{(2)},\dots,\bm{y}^{(T)}\in\mathbb{R}^{{d_{\text{out}}}}$
such that no $f_{n,m,{d_{\text{out}}}}$ can memorize this dataset. In other
words, for any $\bm{W}_{1},\bm{W}_{2},{\bm{b}}_{1},{\bm{b}}_{2}$, there exists
index $i\in[T]$ such that

$f_{n,m,{d_{\text{out}}}}(\bm{x}_{i};\bm{W}_{1},{\bm{b}}_{1},\bm{W}_{2},{\bm{b}}_{2})\neq\bm{y}_{i}.$

Now, let us consider any arbitrary
$\bm{x}^{(1)},\bm{x}^{(2)},\dots,\bm{x}^{(T)}\in\mathbb{R}^{n}$ in General
Position (i.e., maximal Kruskal rank of $n$), and their corresponding set of
labels
$\bm{y}^{(1)},\bm{y}^{(2)},\dots,\bm{y}^{(T)}\in\mathbb{R}^{{d_{\text{out}}}}$
for which no $f_{n,m,{d_{\text{out}}}}$ can memorize the labels. Consider the
following extended input dataset:

$\displaystyle{\bm{c}}_{\text{FCN}}:=[{\bm{s}}_{1};{\bm{s}}_{2};\dots;{\bm{s}}_{n};\bm{0}_{(d-n)}]\in\mathbb{R}^{(n+1)d-n},~{}~{}\bm{x}_{\text{FCN}}^{(t)}:=[{\bm{c}}_{\text{FCN}};\bm{x}^{(t)}]\in\mathbb{R}^{(n+1)d}\quad\text{
for all }t\in[T],$ (29)

where ${\bm{s}}_{i}$ are standard basis vectors in $\mathbb{R}^{d}$. Let us
call
$\mathcal{T}_{\text{FCN}}:=\\{(\bm{x}_{\text{FCN}}^{(t)},\bm{y}^{(t)})\\}_{t\in[T]}$.
First, we claim that no network $f_{(n+1)d,m,{d_{\text{out}}}}$ can memorize
$\mathcal{T}_{\text{FCN}}$. To show this, we create a mapping between any
network of $f_{(n+1)d,m,{d_{\text{out}}}}$ to networks of
$f_{n,m,{d_{\text{out}}}}$. For any
$\bm{W}^{\prime}_{1}\in\mathbb{R}^{m\times(n+1)d},{\bm{b}}^{\prime}_{1}\in\mathbb{R}^{m},\bm{W}^{\prime}_{2}\in\mathbb{R}^{{d_{\text{out}}}\times
m}$, and ${\bm{b}}^{\prime}_{2}\in\mathbb{R}^{{d_{\text{out}}}}$ define:

$\displaystyle\bm{W}_{1}$
$\displaystyle:=\bm{W}^{\prime}_{1}[:,(n+1)d-n+1:(n+1)d]$
$\displaystyle(\bm{W}_{1}\in\mathbb{R}^{m\times n})$
$\displaystyle{\bm{b}}_{1}$
$\displaystyle:={\bm{b}}^{\prime}_{1}-\bm{W}^{\prime}_{1}[:,1:(n+1)d-n]^{\top}{\bm{c}}_{\text{FCN}}$
$\displaystyle({\bm{b}}_{1}\in\mathbb{R}^{m})$ $\displaystyle\bm{W}_{2}$
$\displaystyle:=\bm{W}^{\prime}_{2}$
$\displaystyle(\bm{W}_{2}\in\mathbb{R}^{{d_{\text{out}}}\times m})$
$\displaystyle{\bm{b}}_{2}$ $\displaystyle:={\bm{b}}^{\prime}_{2},$
$\displaystyle({\bm{b}}_{2}\in\mathbb{R}^{{d_{\text{out}}}})$

where $\bm{W}^{\prime}_{1}[:,a:b]$ is the sub-block of matrix
$\bm{W}^{\prime}_{1}$ selected by columns $a$ to $b$. With the above
assignment, we get:

$f_{(n+1)d,m,{d_{\text{out}}}}(\bm{x}_{\text{FCN}}^{(t)};\bm{W}^{\prime}_{1},{\bm{b}}^{\prime}_{1},\bm{W}^{\prime}_{2},{\bm{b}}^{\prime}_{2})=f_{n,m,{d_{\text{out}}}}(\bm{x}^{(t)};\bm{W}_{1},{\bm{b}}_{1},\bm{W}_{2},{\bm{b}}_{2}),$

for all $t\in[T]$. Therefore, since no ReLU network $f_{n,m,{d_{\text{out}}}}$
is able to memorize $\\{(\bm{x}^{(t)},\bm{y}^{(t)})\\}_{t\in[T]}$, no ReLU
network $f_{(n+1)d,m,{d_{\text{out}}}}$ is able to memorize
$\mathcal{T}_{\text{FCN}}$ either.

So far, we have found a dataset $\mathcal{T}_{\text{FCN}}$ that a two-layer
ReLU network with $m$ hidden neurons can not memorize. Next, we prove that
$\mathcal{T}_{\text{FCN}}$ can be derived from a valid dataset
$\mathcal{T}=\left\\{\left(\bm{E}^{(t)},{\bm{e}}^{(t)},\bm{y}^{(t)}\right)\right\\}_{t=1}^{T}$,
which satisfies both Assumptions 1 and 2. To observe this, consider the
following inputs:

$\displaystyle\bm{E}^{(t)}$
$\displaystyle:=[{\bm{s}}_{1}^{\top};{\bm{s}}_{2}^{\top};\dots,{\bm{s}}_{n}^{\top}]$
$\displaystyle{\bm{e}}^{(t)}$ $\displaystyle:=[\bm{0}_{(d-n)};\bm{x}^{(t)}],$

and labels $\bm{y}^{(t)}$ the same as $\mathcal{T}_{\text{FCN}}$. Then, we get

$\bm{x}_{\text{FCN}}^{(t)}=[{\bm{e}}_{1}^{(t)};\dots;{\bm{e}}_{n}^{(t)};{\bm{e}}^{(t)}].$

So, what is left is to prove that $\mathcal{T}$ satisfies Assumptions 1 and 2.
First, $\\{{\bm{e}}^{(t)}\\}_{t\in[T]}$ has the same Kruskal rank as
$\\{\bm{x}^{(t)}\\}_{t\in[T]}$, which is $n$, so Assumption 1 is satisfied.
Second, each $\bm{E}^{(t)}$ is full-rank since each row $i\in[n]$ is the
$i$-th standard basis for all $t\in[T]$.

Therefore, we have found a dataset $\mathcal{T}$ with size
$T>(n+1)m/{d_{\text{out}}}+m+1$ that no ReLU network with $m$ hidden
dimensions can memorize, and this completes the proof.

##  Appendix G Proof of Proposition 4

First, we derive the output tokens
$\bm{E}^{\prime^{(t)}},{\bm{e}}^{\prime^{(t)}}$ as a function of input tokens.
We omit the head subscript since we have only one head. Since
${\bm{W}_{K}}={\bm{W}_{Q}}=\bm{0}$, we have $\bm{\alpha}^{(t)}=\bm{0}$, and
consequently $\bm{\theta}^{(t)}=\frac{1}{n}\bm{1}$. Therefore, the attention
head only takes the average of each token, i.e.,

$\bar{{\bm{e}}}:=\frac{1}{n}\left({\bm{e}}_{1}^{(t)}+{\bm{e}}_{2}^{(t)}+\dots+{\bm{e}}_{n}^{(t)}\right).$

Now, since the Attention head also has a skip connection, the output of each
token becomes:

$\displaystyle{\bm{e}}^{\prime^{(t)}}$
$\displaystyle={\bm{e}}^{(t)}+\bm{o}^{(t)}={\bm{e}}^{(t)}+\bar{{\bm{e}}},$
(30) $\displaystyle{\bm{e}}_{i}^{\prime^{(t)}}$
$\displaystyle:={\bm{e}}_{i}^{(t)}+\bm{o}_{i}^{(t)}={\bm{e}}_{i}^{(t)}+\bar{{\bm{e}}}\quad\text{for
all }i\in[n]$ (31)

First, since we have ${\bm{e}}^{\prime^{(t)}}=\tilde{\bm{e}}^{(t)}$ in
Assumption 3, this means that the set
$\\{{\bm{e}}^{\prime^{(t)}}\\}_{t=1}^{T}$ has a Kruskal Rank at least $n$, and
Assumption 1 holds for $\mathcal{T}^{\prime}$. Next, since the $\bm{E}^{(t)}$
are full-rank, adding the average of tokens to $\bm{E}^{(t)}$ preserves the
rank (Lemma 2), so $\bm{E}^{\prime{(t)}}$ are also full-rank and Assumption 2
is satisfied for $\mathcal{T}^{\prime}$.

##  Appendix H Softmax Saturation and Optimization-Based Memorization

In the proof of Theorem 1, we took advantage of the saturation properties of
Softmax. We generate a dataset of size $T$ with a shared context (as in
Section 5.2) with $d=64$, $n=32$, and discrete labels
$y^{(t)}\in\mathcal{U}([100])$ drawn iid from a uniform distribution. We test
with two different head sizes $H=4$ and $H=8$. For each head size, we first
empirically find the maximal number of data points the model can perfectly
memorize using gradient-based optimization. We find for $H=4$, this number is
around $T=550$, and for $H=8$ around $T=1300$. We train both models for
$300,000$ optimization steps. We call a head _saturated_ for an example if the
query token has one softmax coefficient $>0.99$ to some context token in that
head, i.e., head $h\in[H]$ is saturated for example $t\in[T]$ if
$\bm{\theta}^{(t)}_{h}$ has at least one entry $>0.99$.

(a)

(b)

(c)

(d)

(e)

(f)

Figure 3: Testing the saturation property of Softmax on synthetic data. On the
top row, we have $H=8$, and on the bottom row $H=4$. (a) and (d) show whether
softmax in each head is saturated for the first $20$ examples of the dataset,
results are qualitatively similar for the rest of the examples. (b) and (e)
show the histogram of the number of saturated heads for each example across
the dataset ($H\times T$ in total). (c) and (f) show the histogram of the
maximum softmax coefficient for each head and example across all the datasets
($H\times T$ in total). All the figures suggest that with perfect
memorization, the majority of heads become saturated across the examples, with
almost all examples having at least one non-saturated head.

The results are depicted in Figure 3. Interestingly, a significant number of
heads are saturated for the majority of examples, while almost all examples
have at least one corresponding non-saturated head. This is in line with the
intuition we provided in the proof of Theorem 1, where our proof relied on
designating a set of examples to one head and saturating that head for the
rest of the examples.

##  Appendix I Upper bound on the rank of matrix Z with stronger assumptions

In Proposition 2 we provided an upper bound on the rank of $\bm{Z}$, which
matches the lower bound of Proposition 1. One might wonder whether under
stronger assumptions (and perhaps far from reality), a full-rank $\bm{Z}$ with
rank $Hd$ is achievable or not. Here, we prove that even the "General
Position" on _all_ tokens assumption is not enough, and there exists a very
simple assumption that limits the rank of $\bm{Z}$ to $H(d-1)+1$.

######  Proposition 5.

Assume that we have a multi-head attention layer $\mathcal{A}$ with $H$ heads,
embedding dimensions ${d}={d_{\text{h}}}=d$, and ${d_{v}}=1$. Let
$\mathcal{T}=\left\\{\left(\bm{E}^{(t)},{\bm{e}}^{(t)},\bm{y}^{(t)}\right)\right\\}_{t=1}^{T}$
be a training set with context size $n<d$. If all context vectors satisfy
$\bm{1}^{\top}{\bm{e}}_{i}^{(t)}=C$ for all $t\in[T],~{}i\in[n]$ and some
universal scalar $C\in\mathbb{R}$, then for any set of weights
$\\{({\bm{W}_{K}}_{h},{\bm{W}_{Q}}_{h})\\}_{h=1}^{H}$, we have

$\operatorname{rank}(\bm{Z})\leq H(d-1)+1.$

Proposition 5 gives an upper bound on the rank of $\bm{Z}$ when individual
elements of each context token sum up to a constant $C$. Note that all the
context vectors can be in general position and still satisfy this condition.
Moreover, this upper bound puts _no assumption_ on the query vectors. Hence
even with the strongest assumptions, obtaining a full-rank matrix $\bm{Z}$ is
out of reach. However, whether stronger assumptions can maintain dependence on
$d$ instead of $n$ (i.e., give $H(d-1)+1$ memorization lower bound instead of
$H(r-1)+1$) is still an open question for future work.

###### Proof of Proposition 5.

First, note that since $\bm{1}^{\top}{\bm{e}}_{i}^{(t)}=C$, and
$\bm{\theta}_{i}^{(t)}$ are convex combinations we have:

$\bm{1}^{\top}\bm{z}_{h}^{(t)}=\bm{1}^{\top}\bm{E}{{(t)}^{\top}}\bm{\theta}_{h}^{(t)}=C\bm{1}^{\top}\bm{\theta}_{h}^{(t)}=C\quad\text{for
all }t\in[T].$

With some abuse of notation, let us write the feature matrix
$\bm{Z}\in\mathbb{R}^{T\times Hd}$ as follows:

$\bm{Z}=\begin{bmatrix}\bm{Z}_{1},&\bm{Z}_{2},&\dots,&\bm{Z}_{H}\\\
\end{bmatrix},$

where each $\bm{Z}_{h}\in\mathbb{R}^{T\times d}$ is a sub-block of $\bm{Z}$
corresponding to the $h$-th head for all $h\in[H]$. Then, since each row of
$\bm{Z}_{h}$ (i.e., $\bm{z}_{h}^{(t)}$) sums up to $C$, with an argument
similar to the proof of Proposition 2 we get:

$\displaystyle\operatorname{rank}(\bm{Z})$
$\displaystyle=\operatorname{rank}\left(\begin{bmatrix}\bm{Z}_{1},&\bm{Z}_{2},&\bm{Z}_{3},&\dots,&\bm{Z}_{H}\end{bmatrix}\right)$
$\displaystyle=\operatorname{rank}\left(\begin{bmatrix}\bm{Z}_{1},&\bm{Z}_{2}-\bm{Z}_{1},&\bm{Z}_{3}-\bm{Z}_{1},&\dots,&\bm{Z}_{H}-\bm{Z}_{1}\end{bmatrix}\right)$
$\displaystyle\leq\operatorname{rank}(\bm{Z}_{1})+\operatorname{rank}(\bm{Z}_{2}-\bm{Z}_{1})+\dots+\operatorname{rank}(\bm{Z}_{H}-\bm{Z}_{1})$
$\displaystyle\leq d+(d-1)+(d-1)+\dots+(d-1)$ $\displaystyle=H(d-1)+1.$

∎

###  I.1 Synthetic experiments under general position assumption

We also run synthetic experiments under the general position assumption on
_all_ input tokens. We consider a similar setting to Section 5.2, but we also
generate context vectors from a uniform distribution as well, i.e.

$\bm{x}_{1}^{(t)},\bm{x}_{2}^{(t)},\dots,\bm{x}_{n}^{(t)},\bm{x}^{(t)}\overset{\text{{iid}}}{\sim}\mathcal{U}(0,1)^{64}\quad\text{for
all }t\in[T].$

This way, all tokens (i.e., context and query) inside, and across all examples
would satisfy the general position assumption. Preliminary results in Figure 4
suggest that the linearity in $H$ still remains, while the linearity in $n$ is
weakened. We conjecture that a bound with complexity $O(Hd)$ would hold, which
could be interesting for future works. Yet, as we showed in Section 5.1, from
a practical point of view, the General Position assumption might be too strong
to be realistic.

(a)

(b)

(c)

(d)

Figure 4: Testing memorization under general position assumption for both
context and query vectors. While the linearity in $H$ still remains, the
linearity in $n$ no longer strictly holds, suggesting $H(d-1)+1$ to be a
better alternative under general position assumption.

##  Appendix J Testing Assumptions on Natural Language Tasks

We further test Assumptions 2 and 1 on BERT Devlin et al. (2019), and GPT2
Radford et al. (2019) on the task of pre-training on Wikipedia Foundation .
Following the setting introduced in Section 5.1, we examine the
embedding/first layer of trained/untrained GPT2/BERT.

  * •

BERT: We follow a setting similar to the pre-training of BERT. For each
example, an input of the format _“[CLS] Sentence1 [SEP] Sentence2”_ is formed,
where _Sentence 1_ and _Sentence 2_ are two pieces of text from the Wikipedia
corpus. Here _Sentence 2_ is the sentence after _Sentence 1_ with probability
$0.5$ and randomly chosen from the corpus with probability $0.5$. The “[CLS]”
token’s output predicts whether _Sentence 2_ follows the first sentence or
not. We consider the “[CLS]” token as the query token ${\bm{e}}^{(t)}$, and
the whole input as the context $\bm{E}^{(t)}$ for all $t\in[T]$. The
positional encoding is sinusoidal positional encoding.

  * •

GPT2: This model is trained autoregressively to predict the next token in each
piece of text. We consider last-token prediction from the context. That is, we
set the query ${\bm{e}}^{(t)}$ to be the last token of an example, and the
context $\bm{E}^{(t)}$ to be the whole input example. In contrast to BERT, The
positional encoding in this model is learned instead of fixed.

Table 3: Testing assumptions on Wikipedia for BERT and GPT2. L0 refers to the embedding layer, and L1 refers to the output of the first layer. | GPT2 | BERT
---|---|---
Model | Random | Trained | Random | Trained
| L0 | L1 | L0 | L1 | L0 | L1 | L0 | L1
General Position | $\times$ | $\times$ | $\times$ | $\times$ | $\times$ | $\times$ | $\times$ | $\times$
Assumption 1 | $\times$ | $\checkmark$ | $\times$ | $\checkmark$ | $\times$ | $\times$ | $\times$ | $\times$
Assumption 2 | $\checkmark$ | $\checkmark$ | $\checkmark$ | $\times$ | $\checkmark$ | $\checkmark$ | $\checkmark$ | $\checkmark$

The results of testing Assumptions are shown in Table 3. While linear
independence assumptions are not well-suited for discrete data, we observe
that most assumptions still hold. For Assumption 2, the sinusoidal positional
encodings in BERT preserve the linear independence for all cases. In GPT2, we
observe that the learned positional encodings in pre-trained GPT2 have a lower
rank and are no longer linearly independent. For Assumption 1, the GPT2 model
already satisfies it after one layer. In the case of BERT, the $n/d$ ratio is
$512/1024$, and while the Kruskal Rank of query vectors is below 512
(violation of Assumption 1), we observe that the Kruskal Rank is high enough
(291 for the case of trained and 349 for the case of random BERT). Hence, the
result of Remark 1 still yields an acceptable lower bound.

(a)

(b)

Figure 5: Testing memorization on regression task with a similar setting as
Figure 2, except for real-valued labels. The same results and conclusions in
terms of monotonicity with $H$ and $n$ hold true for regression. The variance
in regression plots is magnified in some cases due to the log scale of the
plot.