Beyond Scaling Laws: Understanding Transformer Performance with Associative Memory

## Abstract

Abstract Increasing the size of a Transformer model does not always lead to enhanced performance. This phenomenon cannot be explained by the empirical scaling laws. Furthermore, improved generalization ability occurs as the model memorizes the training samples.
We present a theoretical framework that sheds light on the memorization process and performance dynamics of transformer-based language models. We model the behavior of Transformers with associative memories using Hopfield networks, such that each transformer block effectively conducts an approximate nearest-neighbor search. Based on this, we design an energy function analogous to that in the modern continuous Hopfield network which provides an insightful explanation for the attention mechanism. Using the majorization-minimization technique, we construct a global energy function that captures the layered architecture of the Transformer. Under specific conditions, we show that the minimum achievable cross-entropy loss is bounded from below by a constant approximately equal to 1. We substantiate our theoretical results by conducting experiments with GPT-2 on various data sizes, as well as training vanilla Transformers on a dataset of 2M tokens.

## 1 Introduction

Transformer-based neural networks have exhibited powerful capabilities in accomplishing a myriad of tasks such as text generation, editing, and question-answering.
These models are rooted in the Transformer architecture (Vaswani et al., ) which employs the self-attention mechanisms to capture the context in which words appear, resulting in superior ability to handle long-range dependencies and improved training efficiency.
In many cases, models with more parameters result in better performance measured by perplexity (Kaplan et al., ), as well as in the accuracies of end tasks (Khandelwal et al., ; Rae et al., ; Chowdhery et al., ). As a result, larger and larger models are being developed in the industry. Recent models (Smith et al., ) can reach up to 530 billion parameters, trained on hundreds of billions of tokens with more than 10K GPUs.

Nevertheless, it is not always the case that bigger models result in better performance. For example, the 2B model MiniCPM (Hu et al., ) exhibits comparable capabilities to larger language models, such as Llama2-7B (Touvron et al., ), Mistral-7B (Jiang et al., ), Gemma-7B (Banks and Warkentin, ), and Llama-13B (Touvron et al., ). Moreover, as computational resources for training larger models increase, the size of available high-quality data may not keep pace. It has been documented that the generalization abilities of a range of models increase with the number of parameters and decrease when the number of training samples increases (Belkin et al., ; Nakkiran et al., ; d’Ascoli et al., ), indicating that generalization occurs beyond the memorization of training samples in over-parameterized neural networks (Power et al., ).
Therefore, it is crucial to understand the convergence dynamics of training loss during memorization, both in relation to the model size and the dataset at hand. There has been an increasing interest in the empirical scaling laws under constraints on the training dataset size (Muennighoff et al., ).
Let $N$ denote the number of the parameters of the model and $D$ denote the size of the dataset.
Extensive experiments have led to the conclusion of the following empirical scaling laws (Kaplan et al., ) in terms of the model performance measured by the test cross-entropy loss $L$ on held-out data:

$$ $\displaystyle L(N,D)$ $\displaystyle=\left(\left(\frac{N_{c}}{N}\right)^{\frac{\alpha_{N}}{\alpha_{D} }}+\frac{D_{c}}{D}\right)^{\alpha_{D}},$ (1) $$

where $N_{c},D_{c},\alpha_{N},\alpha_{D}$ are constants.
Unfortunately, this scaling law does not explain why in many cases smaller models perform better.

In this paper, we focus on the theoretical aspects of the dependencies between the achievable performance, indicated by the pre-training loss, for transformer-based models, and the model and data sizes during memorization. It has been observed that a family of large language models tends to rely on knowledge memorized during training (Hsia et al., ), and the larger the models, the more they tend to encode the training data and organize the memory according to the similarity of textual context (Carlini et al., ; Tirumala et al., ). Therefore, we model the behavior of the Transformer layers with associative memory, which associates an input with a stored pattern, and inference aims to retrieve the related memories. A model for associative memory, known as the Hopfield network, was originally developed to retrieve stored binary-valued patterns based on part of the content (Amari, ; Hopfield, ). Recently, the classical Hopfield network has been generalized to data with continuous values. The generalized model, the Modern Continuous Hopfield Network (MCHN), has been shown to exhibit equivalence to the attention mechanism (Ramsauer et al., ). With MCHN, the network can store well-separated data points with a size exponential in the embedding dimension. However, the MCHN only explains an individual Transformer layer and relies heavily on regularization.

Transformer-based models consist of a stack of homogeneous layers. The attention and feed-forward layers contribute to the majority of the parameters and are also the key components of the attention mechanism. It has been shown that the feed-forward layers can be interpreted as persistent memory vectors, and the weights of the attention and feed-forward sub-layers can be merged into a new all-attention layer without compromising the model’s performance (Sukhbaatar et al., ). This augmentation enables a unified view that integrates the feed-forward layers into the attention layers, thereby significantly facilitating our analysis. Furthermore, the layered structure of the transformer networks induces a sequential optimization, reminiscent of the majorization-minimization (MM) technique (Ortega and Rheinboldt, ; Sun et al., ), which has been extensively utilized across domains such as signal processing and machine learning. Using the MM framework, we construct a global energy function tailored for the layered structure of the transformer network.

Our model provides a theoretical framework for analyzing the performance of transformer-based language models as they memorize training samples. Large language models only manifest capabilities for certain downstream tasks once the training loss reaches a specific threshold (Du et al., ). In practice, the training of large language models is terminated when the loss curves plateau. On the one hand, the validation loss offers valuable insights for budgetary considerations; it has been observed that even after training on up to 2T tokens, some models have yet to exhibit signs of saturation (Touvron et al., ). On the other hand, early stopping can potentially compromise the model’s generalization capabilities (Murty et al., ). On the other hand, implementing early stopping can potentially compromise the generalization capabilities of the models (Murty et al., ).
We conduct a series of experiments utilizing GPT-2 across a spectrum of data sizes, concurrently training vanilla Transformer models on a dataset comprising 2M tokens. The experimental outcomes substantiate our theoretical results.
We believe this work offers valuable theoretical perspectives on the optimal cross-entropy loss, which can inform and enhance decision-making regarding model training.

## 2 Related Work

#### Scaling laws

As discussed in the introduction, we have seen consistent empirical evidence that the performance of models increases as both the size of the models and the volume of training data scale up (Kaplan et al., ; Khandelwal et al., ; Rae et al., ; Chowdhery et al., ).
Intensive experiments have also been conducted to explore neural scaling laws under various conditions, including constraints on computational budget (Hoffmann et al., ), data (Muennighoff et al., ), and instances of over-training (Gadre et al., ).
In these analyses, a decomposition of the expected risk is utilized, leading to the following fit:

$$ $\hat{L}(N,D)=E+\frac{A}{N^{\alpha}}+\frac{B}{D^{\beta}}.$ (2) $$

For Chinchilla models, the fitted parameters are (Hoffmann et al., )

$$ $\alpha=0.34,\quad\beta=0.28,\quad E=1.61,\quad A=406.4,\quad B=410.7.$ $$

A line of research concerns the generalization of over-parameterized neural networks (Belkin et al., ; Nakkiran et al., ; Power et al., ).
Recent experiments show that over-trained transformers exhibits inverted U-shaped scaling behavior (Murty et al., ), which cannot be explained by the empirical scaling laws.

#### Energy-based models

Energy-based models, motivated by statistical physics, have become a fundamental modeling tool in various fields of machine learning over the past few decades (LeCun et al., ). The central idea is to model the neural network through a parameterized probability density function $p_{\theta}(x)$ for $x\in\mathbb{R}^{n}$ and to express the distribution in terms of a learnable energy function $E_{\theta}(x):\mathbb{R}^{n}\mapsto\mathbb{R}$ whose parameters correspond to the model’s parameters as

$$ $p_{\theta}(x)=\frac{\exp(-E_{\theta}(x))}{Z_{\theta}}.$ $$

Here, $Z_{\theta}=\int_{x}\exp(-E_{\theta}(x))$ is the normalizing constant known as the partition function.

#### Hopfield models

Classical Hopfield networks (Amari, ; Hopfield, ) were introduced as paradigmatic examples of associative memory. The network’s update dynamics define an energy function, whose fixed points correspond to the stored memories. An important indicator is the number of patterns that the model can memorize, known as the network’s storage capacity. Modifications to the energy function (Krotov and Hopfield, ; Demircigil et al., ) result in higher storage capacities (see Table 1 in Appendix A). The original model operates on binary variables. The modern continuous Hopfield network (MCHN) (Ramsauer et al., ) generalizes the Hopfield model to the continuous domain, making it an appealing tool for understanding the attention mechanism in Transformers, which also take vector embeddings in the real domain as inputs. Given an input (e.g., a prompt), the Hopfield layer retrieves a memory by converging to a local minimum of the energy landscape, and the update rule has a nice correspondence to the query-key-value mechanism in attention. Krotov () proposes a Hierarchical Associative Memory (HAM) model that enables the description of the neural network with a global energy function, as opposed to energy functions for individual layers.

## 3 Model

We consider tokenized training samples $\mathcal{D}=\{s^{1},s^{2},\ldots,s^{d}\},$ where each element is a sequence of tokens whose length is bounded by a number $T_{\mathrm{max}\in\mathbb{N}}.$ Let $\widetilde{\mathcal{D}}=\{\tilde{s}^{1},\tilde{s}^{2},\ldots,\tilde{s}^{d^{
\prime}}\}$ be the set of held-out validation samples. The preprocessing of the dataset is described in Appendix D.
The size $D\in\mathbb{N}$ of the dataset is proportional to the number of samples $d\in\mathbb{N},$ such that $D\approx T_{\mathrm{max}}d.$
Let $d_{\mathrm{emb}}\in\mathbb{N}$ be the embedding dimension of the tokens, so each input sequence has $n=T_{\mathrm{max}}d_{\mathrm{emb}}$ dimensions.
Let $N\in\mathbb{N}$ be the total number of parameters, most of which are attributed to the attention layers and the feed-forward layers. Then $N\approx Ald_{\mathrm{emb}}^{2}=\frac{Ald_{\mathrm{emb}}}{T_{\mathrm{max}}}n$ for some constant $A$ (see Appendix D).
We use a generic distance metric $d(\cdot,\cdot)$ in the Euclidean space.

### 3.1 Associative memories

We consider models trained with a causal language modeling objective.
Given an input sequence of tokens $s_{\leq t}=(s_{1},s_{2},\ldots,s_{t}),$ the $l$-layer Transformer outputs a distribution
over the next token $s_{t+1}.$ The transformer models are trained by maximizing the log-probability of the correct token $s_{t+1}$ given $s_{\leq t}.$ After training, the model develops the ability to generate desired content, such as predicting the correct
continuations. Thus, the sequences can be viewed as patterns, as in the setting of associative memories, where stored patterns (e.g., sequences) can be retrieved using partial contents of the patterns. The tokenized patterns consist of a subset of the training samples $\mathcal{D}$ based on the following observation.

###### Observation 1

The models tend to memorize the patterns of the training data.

Empirical studies on large language models have shown that the larger the models are, the more they tend to memorize training data (Carlini et al., ; Tirumala et al., ). This memorization allows the models to learn important patterns, such as world knowledge (Hsia et al., ), individual words (Chang and Bergen, ), and linguistic structure (Chang and Bergen, ). Therefore, we make the following assumption regarding memorization.

###### Assumption 1

After saturation, the model memorizes the training samples $\mathcal{D}$ as patterns $\{\rho^{1},\rho^{2},\ldots,\rho^{d}\},$ where $\rho^{i}\in\mathbb{R}^{n}$ for $i=1,2,\ldots,d.$

To be economical with notations, we use $\mathcal{D}$ to directly address the patterns $\mathcal{D}=\{\rho^{1},\rho^{2},\ldots,\rho^{d}\}.$
By memorizing the samples, we mean that the patterns are stored within the model and can be retrieved when provided with an adequate prompt. Specifically, we follow the definitions in (Ramsauer et al., ) for stored and retrieved patterns.

###### Definition 1

For every pattern $\rho^{i},$ denote by $B_{i}:=\{x\in\mathbb{R}^{n}:d(x,a_{i})\leq r_{i}\}$ an $n$-ball such that $\rho^{i}\in B_{i}.$ The pattern $\rho^{i}$ is said to be stored if there exists a single fixed point $\rho^{i*}\in B_{i}$ to which all points $x\in B_{i}$ converge, and $B_{i}\cap B_{j}=\emptyset$ for $i\neq j.$ Such $B_{i}$ is said to be associated to the pattern $\rho^{i},$ and we denote $B_{i}\sim\rho^{i}.$ The pattern $\rho^{i}$ is said to be retrieved if the converged point is $\epsilon$-close to the fixed point $\rho^{i*}.$

We also assume that the small set of held-out test samples exhibits the same patterns as those in the training set. In practice, the test samples are randomly selected from the same dataset as the training samples, preserving the distribution.

###### Assumption 2

Every element in the set of patterns $\widetilde{\mathcal{D}}$ associated with the test data is within $B_{i}$ for some $i\in[d].$ Therefore, we assume $\widetilde{\mathcal{D}}\subset\mathcal{D}.$

###### Observation 1

###### Assumption 1

###### Definition 1

###### Assumption 2

### 3.2 Transformer blocks

Transformers (Vaswani et al., ) are made of a stack of homogeneous layers, where each consists of a multi-head attention sub-layer, a feed-forward sub-layer, an add-norm operation with a skip connection, and layer normalization.
As an example of a typical Transformer, the GPT-2 architecture is discussed in Appendix D. The multi-head attention and feed-forward (FF) layers account for most of the parameters in the model.

#### Attention mechanism

The attention mechanism arguably contributes most to the overall performance of the transformer models.
The attention mechanism takes three matrices $K\in\mathbb{R}^{d_{\mathrm{emb}}\times d_{k}},Q\in\mathbb{R}^{d_{\mathrm{emb}}
\times d_{k}},$ and $V\in\mathbb{R}^{d_{\mathrm{emb}}\times d_{v}}$ as inputs that can be interpreted as keys, queries, and values. Setting $d_{v}=d_{\mathrm{emb}}$ facilitates the inclusion of residual connections. In a single update, an attention matrix is obtained using the update rule
$\mathrm{Attention}(Q,K,V)=V\cdot\mathrm{softmax}(\frac{QK^{\mathsf{T}}}{\sqrt{
d_{k}}})$

#### Feed-forward layers

The FF layers take up almost two third of the model’s parameters.
It has been shown that the FF layers operate essentially as key-value memories (Geva et al., ) such that
$\mathrm{FF}(x)=f(x\cdot K^{\mathsf{T}})\cdot V,$
where $K,V$ are parameter matrices and $f$ is a non-linear activation function such as ReLU.
In fact, the FF layers can be merged into the attention without degrading the Transformer’s performance (Sukhbaatar et al., ).
Thus, we have the following observation.

###### Observation 2

The attention layer and the feed-forward layer can be conceptually integrated into a unified transformer layer.

The attention layers and the FF layers contribute to the majority of the model’s parameters, such that the number of parameters $N$ is proportional to the square of the embedding dimension. The ratio depends on the number of layers and the hidden dimensions of the transformer blocks.
In the current work, we do not consider other modifications such as lateral connections, skip-layer connections, or other compressive modules such as (Xiong et al., ; Fei et al., ; Munkhdalai et al., ).

###### Observation 2

## 4 A New Energy Function

Figure: (a) $-\mathrm{LogSumExp}$
Refer to caption: extracted/5596353/energy_logsumexp.png

We first introduce a new energy function that does not rely on additional regularization terms. We then adapt this function to the layered transformer blocks using the majorization-minimization technique. For reference, related energy functions for Hopfield networks are listed in Table 1 within Appendix A. In particular, the energy function for the modern continuous Hopfield network (Ramsauer et al., ) is

$$ $\displaystyle E_{\mathrm{MCHN}}^{\beta}(x)$ $\displaystyle=-\mathrm{LogSumExp}(\beta,Mx)+\frac{1}{2}x^{T}x+\beta^{-1}\log d +\frac{\max_{i}\|x_{i}\|^{2}}{2},\quad\text{where}$ $\displaystyle\mathrm{LogSumExp}(\beta,Mx):=\beta^{-1}\log\left(\sum_{i=1}^{d} \exp(\beta\rho_{i})\right).$ $$

Notice that the negative LogSumExp function was adapted from (Demircigil et al., ). However, in the continuous domain, the negative LogSumExp function is not convex, which makes it a less suitable candidate for the energy function. The MCHN energy then adds regularization terms to create a convex energy function. These regularization terms involve both the max norm of the input and the number of patterns.

Instead of designing different regularization terms, we define a new energy function through an auxiliary function

$$ $g(x):=\min_{1\leq i\leq d}d(x,\rho_{i}),$ (3) $$

which functions as a nearest neighbor search over the the set of patterns $\mathcal{D}.$
So it holds that $g(x)\geq 0,$ with $g(x)=0$ if and only if $x\in\mathcal{D}=\{\rho_{1},\ldots,\rho_{d}\}.$
According to Assumption 1, the model has memorized all the patterns; thus, the inference corresponds to a search algorithm based on some distance $d(\cdot,\cdot).$ In the sequel, we use the squared Euclidean 2-norm $d(x,y)=\|x-y\|^{2}$ for simplicity, although other norms could be used in general.

We consider a new energy function $E(x)$ which also takes the form of LogSumExp. It is worth noting that the softmax function is the gradient of the LogSumExp function. By summing up the negative distance between $x$ and each stored pattern, the function assigns smaller values to points near the patterns. Our proposed energy function is

$$ $E(x)=-\log\left(\sum_{i=1}^{n}\exp(-d(x,\rho_{i}))\right).$ (4) $$

By replacing the dot product in the MCHN energy with a distance metric, $E(x)$ achieves the same goal without additional regularization. As shown in Figures 1(a) and 1(b), as an extension of (Demircigil et al., ), the negative LogSumExp is not convex in the real domain, so regularization terms are applied. Figures 1(d) and 1(c) show that the landscape of the proposed energy resembles that of the MCHN energy.
In (Ramsauer et al., ), it is shown that $E_{\mathrm{MCHN}}$ induces stationary points near the stored patterns. Here, the proposed function $E(x)$ serves as a smooth surrogate of the desired function $g(x)$ in Eq. (3), therefore also demonstrates the retrieval ability.

###### Proposition 1

Given $\mathcal{D}=\{\rho_{1}\ldots,\rho_{d}\},$ the proposed energy $E(x)$ satisfies

$$ $g(x)-\log d\leq E(x)\leq g(x).$ $$

The proof of Proposition 1 is due to Lemma 3 and is postponed to Appendix B.
Furthermore, we show that $E(x)$ is close to the MCHN energy.

###### Proposition 2

Let $\beta=2$ we have

$$ $|E(x)-(2E_{\mathrm{MCHN}}^{\beta=2}(x)-\log d)|\leq\max_{1\leq i\leq d}\|\rho_ {i}\|^{2}-\min_{1\leq i\leq d}\|\rho_{i}\|^{2}.$ $$

The proof is given in Appendix B. Fig. 1 provides visualizations of the two propositions using one-dimensional patterns.
The following result follows directly from the above inequalities.

###### Proposition 3

$$ $\min_{1\leq i\leq d}\|\rho_{i}\|^{2}-\max_{1\leq i\leq d}\|\rho_{i}\|^{2}\leq g (x)-2E_{\mathrm{MCHN}}^{\beta=2}(x)\leq\max_{1\leq i\leq d}\|\rho_{i}\|^{2}- \min_{1\leq i\leq d}\|\rho_{i}\|^{2}+\log d.$ $$

Since the proposed energy and the MCHN energy both approximate the search for the nearest pattern (desired stationary point), according to Theorem 4 in (Ramsauer et al., ), in each transformer layer, the probability density of the transformer layer, corresponding to the retrieval, is

$$ $p(x)=\frac{1}{Z}\exp(-E(x)\rvert_{\Omega}),\quad\Omega=\bigcup_{i=1}^{d}B_{i}$ $$

where $Z$ is the normalizing factor and $B_{i}$ is as defined in Definition 1. We assume that $B_{i}=\{x\in\mathbb{R}^{n}:d(x,\rho^{i})\leq r_{i}\}$ is centered at the $i$-th pattern.
We make the following assumption on the samples in order to study the memorization. In practice, this requires the training dataset to be of very good quality.

###### Assumption 3

The patterns in $\mathcal{D}$ are well-separated, i.e., $B_{i}\cap B_{j}=\emptyset,\forall 1\leq i<j\leq d.$

Under Assumption 3, the energy function, confined in $\Omega,$ can be replaced by the nearest neighbor search $g(x).$ So we can make the following claim.

###### Claim 1

Suppose the patterns are well-separated, then the probability density is

$$ $p(x)=\frac{1}{Z}\exp(-g(x)).$ (5) $$

###### Proposition 1

###### Proposition 2

###### Proposition 3

###### Assumption 3

###### Claim 1

### 4.1 The layered structure

Previous Hopfield models could only handle a single hidden layer, whereas Transformers often consist of a stack of homogeneous blocks of attention and FF layers. To model the multi-layered structure of Transformers, we employ a technique known as majorization-minimization (MM) (Ortega and Rheinboldt, ; Sun et al., ), which aims to accelerate optimization using surrogate convex functions. We argue that the layered structure serves the same purpose when the patterns memorized by all layers encompass the set of training samples.

We divide the set of samples into $\mathcal{D}=\cup_{i=1}^{l}\mathcal{D}_{i},$ where $\mathcal{D}_{i}=\{\rho^{i_{1}},\rho^{i_{2}},\ldots,\rho^{i_{d}}\}.$ Then,
the energy function for each layer can be written as

$$ $\displaystyle E_{t}(x)=\frac{1}{Z_{t}}\exp(-g_{t}(x)),\quad\text{where}\quad g _{t}(x):=\min_{1\leq j\leq d}d(x,\rho^{t_{j}}).$ $$

Denote by $x^{(0)}$ the embedding vector input into the first transformer layer and $x^{(t)}\in\mathbb{R}^{n}$ the output of the $t$-th layer for $t=1,2,\ldots,l.$ Let $E_{t}(x)$ be the energy function associated with the Hopfield model of the $t$-th layer, then the sequential structure of the transformer network is achieved by forwarding the output $x^{(t-1)}$ to the $t$-th layer as input, i.e.,

$$ $\displaystyle x^{(t)}$ $\displaystyle=\operatorname*{arg\,min}_{x\in\mathcal{X}_{t}}E_{t}(x),\quad t=1 ,\ldots,l$ (6) $\displaystyle\mathcal{X}_{t}$ $\displaystyle=\{x\in\mathbb{R}^{n}:d(x,x^{(t-1)})\leq\delta_{t}\}$ $$

where the retrieved fix point attractor in the $t$-th layer is $\delta_{t}$-close to $x^{(t-1)}$ in $d(\cdot,\cdot)$ for some $\delta_{t}>0.$
Such sequential optimization step is equivalent to the MM technique where every minimization step locally approximates the objective function. In particular, Eq. (6) corresponds to the surrogate function Eq. $(3)$ in (Sun et al., ).
Therefore, we define a global energy function

$$ $E_{\mathrm{global}}(x):=-\mathrm{LogSumExp}((-E_{1}(x),-E_{2}(x),\ldots,-E_{l} (x))).$ (7) $$

$E_{\mathrm{global}}(x)$ is continuous but not convex. As opposed to the HAM (Krotov, ), the global energy function is not a linear combination of the component energies.
According to Lemma 3, we have

$$ $\min_{1\leq i\leq l}E_{i}(x)-\log l\leq E_{\mathrm{global}}(x)<\min_{1\leq i \leq l}E_{i}(x).$ (8) $$

So

$$ $E_{t}(x)|_{x\in\mathcal{X}_{t}}\geq E_{\mathrm{global}}(x)|_{x\in\mathcal{X}_{ t}}+c_{t}$ $$

as in Eq. $(2)$ in (Sun et al., ). The probability density function corresponding to the layered transformer network can then be written as

$$ $p_{\theta}(x)=\frac{1}{Z_{\theta}}\exp(-E_{\mathrm{global}}(x)),\quad x\in\Omega$ (9) $$

where $\theta$ denotes the model’s parameters and $Z_{\theta}$ is the normalizing constant.

###### Remark 1

If the model is severely over-parameterized, the energy function can approximate the energy of the sample distribution well and is not confined to the form expressed in Eq. (9).

###### Remark 1

## 5 Cross-Entropy Loss

We now proceed to analyze the loss of Transformer networks. The cross-entropy loss, which measures the difference between predicted probabilities and actual labels, is commonly used for training Transformer models. The attention mechanism includes a softmax operation that outputs a probability distribution $p\in\Delta_{n}.$ In practice, the final softmax output is then fed into a task-specific layer for downstream tasks, such as predictions and classifications. Thus, we compare the last softmax output of the transformer blocks with the target distribution.

We demonstrate that the cross-entropy loss can be formulated using the log partition function of the model’s distribution, which reflects how the attention weights are allocated based on the patterns. Let us consider the cross-entropy loss on the test dataset $\widetilde{\mathcal{D}}.$
Generally, the cross-entropy loss is the negative log-likelihood computed over a mini-batch. Since we are considering un-batched validation samples, the loss is normalized by the size $d^{\prime}.$

According to Eq. (8), there exist a layer $t$ such that
$E_{\mathrm{global}}(x)$ is close to $E_{t}(x),$ i.e.,

$$ $E_{\mathrm{global}}(x)=E_{t}(x)-\log l+c(x),$ (10) $$

for some $c(x)\in C^{\infty}(\mathbb{R}^{n})$ such that $0\leq c(x)<\log l.$ We further assume that $c(x)=c\in[0,\log l)$ is constant.
Under Assumption 1, the target distribution, which encodes all the patterns in $\mathcal{D}$, is given by

$$ $p_{\mathcal{D}}(x)=\sum_{i=1}^{d}p_{i}\delta(x-\rho^{i}),\quad x\in\mathbb{R}^ {n}$ $$

where $\delta(\cdot)$ is the Dirac delta function such that $\delta(x)=0,\forall x\neq 0$
and $p_{i}=\Pr(x=\rho^{i})$ is the probability mass assigned to pattern $\rho_{i}$ for $i=1,2,\ldots,d.$ Suppose the data points are homogeneous, i.e., $p_{i}=\frac{1}{d},$ then $P_{\mathcal{D}}(x)=\frac{1}{d}\sum_{i=1}^{d}\delta(x-\rho^{i}),$ and the corresponding test samples $\widetilde{\mathcal{D}}$ induces the distribution

$$ $P_{\widetilde{\mathcal{D}}}(x)=\frac{1}{d^{\prime}}\sum_{i=1}^{d^{\prime}} \delta(x-\rho^{\sigma{(i)}}),\quad\sigma(\cdot)\in\mathrm{Sym}([d]).$ (11) $$

We have the following result regarding the cross-entropy loss.

###### Proposition 4

Let $L$ be the cross-entropy loss of the above model, then

$$ $1<L=\log Z_{t}+\frac{1}{Z_{t}}+\log l-c\leq\log Z_{t}+\frac{1}{Z_{t}}+\log l,$ $$

where $c\in[0,\log l).$

The proof is deferred to Appendix C.1.
Note that the empirically obtained loss function Eq. (2) for the Chinchilla model converges to $\hat{L}(N,D)=1.61$ as $N\rightarrow\infty$ and $D\rightarrow\infty,$ which substantiates our theory that $L(N,D)\approx\log Z_{t}+\frac{1}{Z_{t}}\geq 1,$ with minimum obtained when $Z_{t}=1.$

###### Remark 2

The cross-entropy can be written as

$$ $H(p_{\mathcal{D}},p_{\theta})=D_{\mathrm{KL}}(p_{\mathcal{D}}\|p_{\theta})+H(p _{\mathcal{D}}).$ $$

When the model is severely over-parameterized, the energy function can well approximate the energy of the sample distribution.
In this case, the minimal cross-entropy equals the entropy of the training samples.

Next, we take a closer look at the layer partition function. We have

$$ $\displaystyle Z_{t}$ $\displaystyle=\int_{x\in\Omega}\exp(-g_{t}(x))\ \mathrm{d}\mu=\int_{x\in\Omega }\exp(-\min_{i}d(x,\rho^{t_{i}}))\ \mathrm{d}\mu$ $\displaystyle=\sum_{i=1}^{d}\int_{x\in B_{t_{i}}}\exp(-\|x-\rho^{t_{i}}\|^{2}) \ \mathrm{d}x\stackrel{{\scriptstyle(c)}}{{=}}2\sum_{i=1}^{d}\pi^{\frac{n}{2}} \frac{\gamma(n,r_{i})}{\Gamma(\frac{n}{2})}$ (12) $$

where $(c)$ is due to Appendix C.2, $r_{i}$ is the radius of $B_{i}.$
According to Lemma 5, we have

$$ $e^{-r_{i}}V_{n}(r_{i})\leq 2\pi^{\frac{n}{2}}\frac{\gamma(n,r_{i})}{\Gamma( \frac{n}{2})}=\int_{x\in B_{i}}\exp(-\|x-\rho^{i}\|^{2})\ \mathrm{d}x\leq V_{n }(r_{i}),$ $$

where

$$ $V_{n}(r)=\frac{\pi^{\frac{n}{2}}r^{n}}{\Gamma(1+\frac{n}{2})}$ $$

is the hyper-volume of the $n$-dimensional ball of radius $r.$
Note that the volume of the unit ball $V_{n}(1)$ in higher dimensions decreases fast with respect to the increase in dimensionality.
The gamma function can be approximated using Stirling’s approximation for large values of its argument.
Applying Stirling’s approximation to the gamma function in the denominator of the volume formula for large $n$, we get:

$$ $\displaystyle V_{n}(1)$ $\displaystyle\approx\frac{\pi^{n/2}}{\sqrt{2\pi(n/2)}\left(\frac{n/2}{e}\right )^{n/2}}=\frac{1}{\sqrt{n\pi}}\left(\frac{2\pi e}{n}\right)^{\frac{n}{2}}$ $$

For $V_{n}(r)$ to have a volume of $O(1)$, the radius $r$ must be approximately $\sqrt{\frac{n}{2\pi e}}$ asymptotically.
Bringing $r=\sqrt{\frac{n}{2\pi e}}$ to Eq. (5), we get

$$ $\frac{d\cdot V_{n}(\sqrt{\frac{n}{2\pi e}})}{\exp(\sqrt{\frac{n}{2\pi e}})} \leq Z_{t}\leq d\cdot V_{n}(\sqrt{\frac{n}{2\pi e}}).$ (13) $$

In Section 3, we argue that

$$ $N\approx\frac{Ald_{\mathrm{emb}}}{T_{\mathrm{max}}}n,\quad D=T_{\mathrm{max}}d.$ $$

Therefore, for $Z_{t}$ to reach $Z_{t}=1,$ we need $N=O(D^{2}).$

In Table 2 in Appendix A, we compare the reported cross-entropy loss of various transformer-based models in the literature. Usually, a family of models ranging in a variety of sizes is reported, and we select the largest ones. We observe that similar cross-entropy loss is achieved across a wide range of architectural shapes (including depth, width, attention heads, FF dimensions, and context lengths). Nevertheless, the losses all satisfy $L>1.$

###### Remark 3

We remark that some models add auxiliary regularization terms such as the z-loss (Chowdhery et al., ; Yang et al., ) during their training. In these cases, the scaling laws should take into consideration the additional terms. Also, modifications to the transformer blocks, such as additional layer normalization may contribute to the lower bound of the cross-entropy.

###### Proposition 4

###### Remark 2

###### Remark 3

## 6 Empirical Results

We explore the hypothesis regarding the radius $r$ in Section 5 using a pre-trained GPT-2 medium model. Additionally, we train various GPT-2 small models and vanilla Transformer models to analyze their cross-entropy losses.

### 6.1 Empirical evaluation of the radius

Figure: Figure 2: Nearest neighbor distances of output activations using $25\%,50\%,75\%,$ and $100\%$ of the output. The mean and median of the distances are all around 20, which is approximately $2\sqrt{n/2\pi e}$ as hypothesized.
Refer to caption: extracted/5596353/histogram-20k.png

We evaluate the radius $r$ in $Z_{l}$ of a pre-trained GPT-2 medium model.
We use the 24-layer pre-trained GPT-2 model (Radford et al., )(^1^11available at [https://github.com/openai/gpt-2](https://github.com/openai/gpt-2)). The medium size model has 355M parameters.
The model is pre-trained with a causal language modeling objective (next sentence prediction) on a very large (40 GB) text corpus extracted from web pages. The hidden dimension $d_{\mathrm{emb}}=1024.$

We test the model on the OpenWebText (Gokaslan and Cohen, ) dataset which is a reproduction of the WebText dataset used for training the GPT-2 model. The dataset contains 9B tokens from 8,013,769 documents.
We randomly sample $80$K chunks of $256$ tokens from the dataset. These cover approximately $1\%$ of the documents and constitute approximately $0.2\%$ of the tokens used for training. For each sample chunk, we record the activation vector of the last layer for prediction of the next token. Since the GPT-2 model is causal, the activation depends on all activations of earlier sequences. As discussed above, each vector should be close to a stored pattern $\rho^{l_{i}}.$
We find the distance between each activation vector and its nearest neighbor in terms of the $L_{2}$ norm. In Fig. 2, we plot the histogram of the nearest neighbor distances for these output activations using $25\%,50\%,75\%,$ and $100\%$ of the output vectors. In all these cases, the mean and median are approximately 20, so that a typical $B_{i}$ of pattern $\rho^{i}$ has radius $10.$ This corroborates our theory in Eq. (13) in Section 5, according to which the radius is of order $\sqrt{\frac{n}{2\pi e}}=\sqrt{\frac{1024}{2\pi e}}=7.74.$ Note that the activation is only collected for $1\%$ of the documents, the estimated radius may be greater than actual one in the model.

### 6.2 Training GPT-2

Figure: Figure 3: Cross-entropy loss of GPT-2 small model trained on (left) $100\%,$ (middle) $1\%,$ and (right) $0.1\%$ of OpenWebText-9B dataset with a typical training time.
Refer to caption: extracted/5596353/GPT1-1.png

We train a 12-layer transformer language model with the GPT-2 small tokenizer and architecture ($l=12,d_{\mathrm{emb}}=768,T_{\mathrm{max}}=1024$) on the publicly available OpenWebText dataset (Gokaslan and Cohen, ), a reproduction of the WebText dataset used for training the original GPT-2 model, which contains 9B tokens from 8,013,769 documents

We train three models using different amount of data. We create subsets containing the first $1\%$ (90M) and $0.1\%$ (9M) of the OpenWebText data, and then train over a typical training time.
For all the settings, we split the dataset into a training set and a smaller test set. We use a batch-size of 12 on the training set and evaluate the cross-entropy loss on the test set. The training and test cross-entropy losses are reported in Fig. 3. The training with $0.1\%$ (9M) of the OpenWebText data clearly experiences over-fitting, and the training loss vanishes over iterations. This is because the training samples are not well-separated so the energy of the model reduces to the summation of some delta functions and no longer satisfies Eq. (5). When training with 90M tokens and 9B tokens, the training cross-entropy losses are both above 2. Since the model size is about the order $O(D^{2})$ when trained on 90M tokens, the model can achieve similar training loss and validation loss compared to the setting with 9B tokens.

### 6.3 Training Vanilla Transformers

Figure: (a) Vanilla Transformer with 6 layers
Refer to caption: extracted/5596353/vanilla6.png

We next train vanilla transformer models using a small amount of high-quality data. The of Question-Formation dataset, proposed by McCoy et al. (), consists of pairs of English sentences in declarative formation and their corresponding question formation. The dataset contains $D=2$M tokens. The sentences are context-free with a vocabulary size of 68 words, and the task is to convert declarative sentences into questions.

We follow the settings in (Murty et al., ) to train two vanilla Transformers ($d_{\mathrm{emb}}=512,T_{\mathrm{max}}=5000$) with $l=6$ layers and $l=10$ layers respectively. We use a batch size of 8 and tie the input and output matrices as done in (Press and Wolf, ). The training losses are shown in Fig. 4, where the losses stabilize at a value of around 1 as predicted in Proposition 4.

## 7 Conclusion

We model transformer-based networks with associative memory and study the cross-entropy loss with respect to model and data sizes. By proposing a new energy function in Eq. 5, which does not rely on additional regularization terms as is common in modern continuous Hopfield networks, we demonstrate that the proposed energy function corresponds to a nearest neighbor search across patterns memorized during training. We then construct a global energy function for the layered structure of the transformer models using the majorization-minimization technique.

In practice, we have observed that the majority of transformer models tend to achieve a cross-entropy loss of approximately 2.2. The optimal balance between model and data sizes, however, is often determined by the collective expertise of practitioners. Additionally, the performance of these models can be compromised by both early and delayed stopping.

We believe the current paper represents an important step towards understanding the convergence and generalization behaviors of large transformer models. It provides insights into the theoretically optimal cross-entropy loss, which can inform both budgetary planning and model termination strategies.

Acknowledgments

The author thanks Dr. Yongqi Xu for stimulating discussions and practical assistance with the experiments.

## Appendix A Deferred Tables

**Table 1: Table of selected related works for Hopfield network, enumerating their domain, energy function, and memory capacity. For all the works above, $n$ represents the dimension of the input vector. $W$ is the outer product of the patterns. $M$ is the matrix of patterns. $r$ is the order of polynomial $F(\cdot),$ $d$ is the number of patterns, and $c$ is a positive constant.**
| Reference | Domain | Energy | Capacity |
| --- | --- | --- | --- |
| Hopfield () | $\{-1,+1\}^{n}$ | $E(x)=-\frac{1}{2}x^{T}Wx-b^{T}x$ | $O(n)$ |
| Krotov and Hopfield () | $\{-1,+1\}^{n}$ | $E(x)=-\sum_{i=1}^{n}F(\rho_{i}^{T}x)$ | $\Theta(n^{r})$ |
| Demircigil et al. () | $\{-1,+1\}^{n}$ | $E(x)=-\mathrm{LogSumExp}(Mx)$ | $\Theta(2^{\frac{n}{2}})$ |
| Ramsauer et al. () | $\mathbb{R}^{n}$ | $E(x)=-\mathrm{LogSumExp}(\beta,Mx)+$ | $\Theta(c^{\frac{n-1}{4}})$ |
|  |  | E(x) = $\frac{1}{2}x^{T}x+\beta^{-1}\log d+\max_{i}\|x_{i}\|^{2}/2$ |  |

**Table 2: Large transformer-based language models and their reported cross-entropy loss.**
| Model | Model Size | Data Size | $L$ | Reference |
| --- | --- | --- | --- | --- |
| Transformer | $1.5$B | $22$B | $2.5$ | (Kaplan et al., ) |
| Chinchilla | $70$B | $1.4$T | $2.2$ | (Hoffmann et al., ) |
| PaLM 2 | $16$B | $100$B | $2.4$ | (Anil et al., ) |
| GPT-2 | $8.7$B | $178$B | $2.3$ | (Muennighoff et al., ) |
| MiniCPM | $2.4$B | $140$B | $2.4$ | (Hu et al., ) |

## Appendix B Some Properties of the Energy Functions

We introduce some useful properties of the LogSumExp function defined below. This is particularly useful because The softmax function, widely utilized in the Transformer models, is the gradient of the LogSumExp function.
As shown in (Grathwohl et al., ), the LogSumExp corresponds to the energy function of the a classifier.

$$ $\mathrm{LogSumExp}(x):=\log\sum_{i=1}^{n}e^{x_{i}},\quad x=(x_{1},\ldots,x_{n} )\in\mathbb{R}^{n}.$ $$

###### Lemma 1

$\mathrm{LogSumExp}(x)$ is convex.

Proof

$$ $\displaystyle t\mathrm{LogSumExp}(x)+(1-t)\mathrm{LogSumExp}(y)=\log(\sum_{i=1 }^{n}e^{x_{i}})^{t}(\sum_{i=1}^{n}e^{y_{i}})^{1-t}$ $\displaystyle\geq\log\sum_{i=1}^{n}e^{tx_{i}+(1-t)y_{i}}=\mathrm{LogSumExp}(tx +(1-t)y)\quad\forall t\in[0,1].$ $$

###### Lemma 2

Suppose $x=(x_{1},\ldots,x_{n})\in\mathbb{R}^{n},$ then we have

$$ $\max_{1\leq i\leq n}x_{i}<\mathrm{LogSumExp}(x)\leq\max_{1\leq i\leq n}x_{i}+ \log n.$ $$

Proof
Taking log on each side of the inequality

$$ $\exp(\max_{1\leq i\leq n}x_{i})<\sum_{i=1}^{n}\exp(x_{i})\leq\sum_{i=1}^{n} \exp(\max_{1\leq i\leq n}x_{i})$ $$

yields the results.

Consequently, we have the following smooth approximation for the min function.

###### Lemma 3

Suppose $x=(x_{1},\ldots,x_{n})\in\mathbb{R}^{n},$ then we have

$$ $\min_{1\leq i\leq n}x_{i}-\log n\leq-\mathrm{LogSumExp}(-x)<\min_{1\leq i\leq n }x_{i}.$ $$

###### Lemma 4

For $x=(x_{1},\ldots,x_{n}),y=(y_{1},\ldots,y_{n})\in\mathbb{R}^{n},$ we have

$$ $|\mathrm{LogSumExp}(x)-\mathrm{LogSumExp}(y)|\leq\|x-y\|_{\infty},$ $$

where $\|x\|_{\infty}:=\max_{1\leq i\leq n}|x_{i}|.$

Proof
Let

$$ $f(t):=\mathrm{LogSumExp}(tx+(1-t)y),\quad\forall t\in[0,1].$ $$

According to the mean value theorem, $\exists s\in(0,1)$ such that

$$ $\mathrm{LogSumExp}(x)-\mathrm{LogSumExp}(y)=f^{\prime}(s)=\frac{\sum_{i=1}^{n} \exp(sx_{i}+(1-s)y_{i})(x_{i}-y_{i})}{\sum_{i=1}^{n}\exp(sx_{i}+(1-s)y_{i})}.$ $$

So

$$ $\displaystyle|\mathrm{LogSumExp}(x)-\mathrm{LogSumExp}(y)|\leq\frac{\sum_{i=1} ^{n}\exp(sx_{i}+(1-s)y_{i})\|x-y\|_{\infty}}{\sum_{i=1}^{n}\exp(sx_{i}+(1-s)y_ {i})}=\|x-y\|_{\infty}.$ $$

###### Lemma 1

###### Lemma 2

###### Lemma 3

###### Lemma 4

### B.1 Proof of Proposition 2

Proof
Let $\xi=\max_{1\leq i\leq d}\|\rho_{i}\|,$ then we have

$$ $\displaystyle 2E_{\mathrm{MCHN}}^{\beta=2}(x)$ $\displaystyle=-\log\left(\sum_{i=1}^{d}\exp(2\rho_{i}^{\mathsf{T}}x)\right)+ \log d+\|x\|^{2}+\xi^{2}$ $\displaystyle=-\log\left(\sum_{i=1}^{d}\exp(2\rho_{i}^{\mathsf{T}}x)\right)- \log(\exp(-(\|x\|^{2}+\xi^{2})))+\log d.$ $$

So

$$ $\displaystyle 2E_{\mathrm{MCHN}}^{\beta=2}(x)-\log d$ $\displaystyle=-\log(\sum_{i=1}^{d}\exp(2\rho_{i}^{\mathsf{T}}x-\xi^{2}-\|x\|^{ 2}))$ $\displaystyle=-\log(\sum_{i=1}^{d}\exp(\|\rho_{i}\|^{2}-\xi^{2}-\|\rho_{i}-x\| ^{2})).$ $$

Therefore, due to Lemma 4, we have

$$ $\displaystyle|E(x)-(2E_{\mathrm{MCHN}}^{\beta=2}(x)-\log d)|$ $\displaystyle=|\mathrm{LogSumExp}(\|\rho_{i}\|^{2}-\xi^{2}-\|\rho_{i}-x\|^{2}) -\mathrm{LogSumExp}(-\|x-\rho_{i}\|^{2})|$ $\displaystyle\leq\max_{1\leq i\leq d}|\|\rho_{i}\|^{2}-\xi^{2}|=\max_{1\leq i \leq d}\|\rho_{i}\|^{2}-\min_{1\leq i\leq d}\|\rho_{i}\|^{2}.$ $$

## Appendix C Deferred Proofs from Section 5

### C.1 Proof of Proposition 4

Proof

$$ $\displaystyle L(N,D)$ $\displaystyle=H(p_{\widetilde{\mathcal{D}}},p_{\theta})=-\frac{1}{d^{\prime}} \sum_{x\in\widetilde{\mathcal{D}}}\log(p_{\theta}(x))=-\mathds{E}_{x\sim p_{ \widetilde{\mathcal{D}}}}[\log p_{\theta}(x)]$ $\displaystyle=\log Z_{\theta}\int_{x\in\Omega}P_{\widetilde{\mathcal{D}}}(x)\ \mathrm{d}\mu+\frac{1}{d^{\prime}}\int_{x\in\Omega}\ \sum_{i=1}^{d^{\prime}} \delta(x-\rho^{\sigma{(i)}})E_{\mathrm{global}}(x)\ \mathrm{d}\mu$ $\displaystyle=\log Z_{\theta}+\frac{1}{d^{\prime}}\sum_{\rho^{\sigma(i)}}E_{ \mathrm{global}}(x)$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{=}}\log Z_{\theta}+\frac{1}{Z_{t}}+ \log l\ \stackrel{{\scriptstyle(b)}}{{=}}\ \log Z_{t}+\frac{1}{Z_{t}}+\log l-c$ (14) $$

where $(a)$ is because $g(\rho^{\sigma(i)})=0,$ and $(b)$ is due to Eq. (10), where we have

$$ $\displaystyle Z_{\theta}$ $\displaystyle=\int_{x\in\Omega}\exp(-E_{\mathrm{global}}(x))\mathrm{d}x=\frac{ l}{e^{c}}\int_{x\in\Omega}\exp(-E_{t}(x))\mathrm{d}x,\quad\text{and}$ $\displaystyle\log Z_{\theta}$ $\displaystyle\approx\log l-c+\log\int\exp(-E_{t}(x))\mathrm{d}x=\log l-c+\log Z _{t}.$ $$

### C.2

Let $B_{t}(x)$ denote the $n$-ball with radius $t$ centered at $x,$ $A_{n-1}=\frac{2\pi^{n/2}}{\Gamma(\frac{n}{2})}$ be the hyper-volume of the $(n-1)$-dimensional unit sphere,

$$ $\gamma(n,r)=\int_{0}^{r}t^{n-1}e^{-t}\ \mathrm{d}t,\quad\Gamma(n,r)=\int_{r}^{ \infty}t^{n-1}e^{-t}\ \mathrm{d}t$ $$

be the (incomplete) gamma functions. Then

$$ $\displaystyle\int_{x\in B_{i}}\exp(-\|x-\rho^{i}\|^{2})\ \mathrm{d}x$ $\displaystyle=\int_{\|x-\rho^{i}\|<r_{i}}\exp(-\|x-\rho^{i}\|^{2})\ \mathrm{d}x$ $\displaystyle=\int_{\|y\|<r_{i}}\exp(-\|y\|^{2})\ \mathrm{d}y$ $\displaystyle=\int_{0}^{r_{i}}\int_{\partial B_{t}(0)}e^{-t^{2}}\ \mathrm{d} \mathcal{H}^{n-1}\mathrm{d}t$ $\displaystyle=\int_{0}^{r_{i}}e^{-t^{2}}\mathcal{H}^{n-1}(\partial B_{t}) \mathrm{d}t$ $\displaystyle=\int_{0}^{r^{i}}e^{-t^{2}}A_{n-1}t^{n-1}\ \mathrm{d}t$ $\displaystyle=\frac{2\pi^{\frac{n}{2}}}{\Gamma(\frac{n}{2})}\int_{0}^{r_{i}}t^ {n-1}e^{-t^{2}}\ \mathrm{d}t$ $\displaystyle=2\pi^{\frac{n}{2}}\frac{\gamma(n,r_{i})}{\Gamma(\frac{n}{2})}$ $$

###### Lemma 5

The incomplete gamma function $\gamma(n,r)$ satisfies

$$ $e^{-r}\frac{r^{n}}{n}\leq\gamma(n,r)\leq\frac{r^{n}}{n}$ $$

Proof
For $0\leq x\leq r,$ we have

$$ $x^{n-1}e^{-r}\leq x^{n-1}e^{-x}\leq x^{n-1}.$ $$

Integrating from 0 to $r$ on each side yields the result.

###### Theorem C.1 (Stirling’s approximation)

For any complex $z,$ the Stirling’s approximation gives that

$$ $\Gamma(z)=\sqrt{\frac{2\pi}{z}}\left(\frac{z}{e}\right)^{z}\left(1+O(\frac{1}{ z})\right).$ $$

For large $z,$

$$ $\Gamma(z+1)\approx\sqrt{2\pi z}\left(\frac{z}{e}\right)^{z}.$ $$

###### Lemma 5

###### Theorem C.1 (Stirling’s approximation)

## Appendix D Transformer Details: Using GPT-2 as an Example

The original GPT-2 model was trained on a 40GB large dataset called WebText that is made of data derived from outbound links from Reddit.
The model is trained on the next sentence prediction (NSP) task in a self-supervised manner. A pre-trained tokenizer can be applied to convert the text into tokens using a fixed vocabulary. A max token length $T_{\mathrm{max}}$ (e.g., $T_{\mathrm{max}}=1024$) is set, so during training, if the number of tokens is greater than $T_{\mathrm{max}},$ the documents will be truncated. The model is trained causally, which means that the prediction for the next token only depends on the inputs from earlier tokens. The model was trained with a global batch size of 512, and the test perplexity still improves if given more training time.

GPT-2 uses a byte-level version of Byte Pair Encoding (BPE), and the vocabulary size is $50,257$. The hidden dimension for the medium size model is $d_{\mathrm{emb}}=1024$. So the input sequence is of $T_{\mathrm{max}}d_{\mathrm{emb}}$ dimension. These sequences are passed through the model. For the medium size model, these include 24 transformer encoder blocks with 1024 hidden units and 16 self-attention heads (i.e., $l=24,d_{\mathrm{emb}}=1024,n_{h}=16$).
The number of parameters in the multi-head attention layer is $l\cdot n_{h}\cdot(3\cdot d_{\mathrm{emb}}\cdot d_{\mathrm{emb}}/n_{h})=3ld_{
\mathrm{emb}}^{2}=75,497,472.$ The number of parameters in the dense weights and layer normalization is $l(d_{\mathrm{emb}}^{2}+2d_{\mathrm{emb}})=ld_{\mathrm{emb}}^{2}+2ld_{\mathrm{
emb}}=25,214,976,$ and the number of parameters in the feed-forward weight matrices and bias is $l(2d_{\mathrm{emb}}\cdot d_{\mathrm{FF}}+d_{\mathrm{emb}}+d_{\mathrm{FF}})=6ld
_{\mathrm{emb}}^{2}+4ld_{\mathrm{emb}}=151,093,248,$ with $d_{\mathrm{FF}}=3072=3d_{\mathrm{emb}}.$
As we can observe, the multi-head attention and feed-forward layers account for most of the parameters in the model, and $N\approx Ald_{\mathrm{emb}}^{2}$ with some constant $A\approx 15$ in this case.

The loss used for the GPT-2 model is the log-probability of a dataset divided by the number of canonical units (e.g., a character, a byte, a word), which is equivalent to the cross-entropy loss. The cross-entropy loss is commonly used to measure the divergence between the predicted probabilities and the true labels.
For the NSP task, the model is trained to predict the next token in a sequence based on the context of the previous tokens. So the cross-entropy is taken between the predicted probabilities $\Pr_{\theta}(x_{i})$ of the token $x_{i}$ and the labels’ probabilities $\Pr_{D}(x_{i})$ for all tokens $x_{i}$ in the vocabulary, i.e.,

$$ $-\frac{1}{D}\sum_{i=1}^{D}\log(p_{\theta}(x_{i}))=-\sum_{x\sim p_{D}}p_{D}(x) \log(p_{\theta}(x)).$ $$

Another commonly used loss is the perplexity, which is equivalent to the exponentiated version of the cross-entropy.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: (Q1) Why bother making a distinctive definition of a validation set at all? If $\tilde{\mathcal{D}} \subset \mathcal{D}$ [L191] i.e., the validation set is some subset of the memorized training points, why consider it at all?

(Q2) The analysis of sections 4 and 5 seems to have nothing to do with the actual transformer architecture. 

(Q3) [L348-352] The attention softmax is irrelevant to the final cross-entropy, and including this statement here only confused me. Can the authors justify why we care about the softmax in the attention when considering the softmax in the cross-entropy loss?

**(More) Typos:**
- Eq (3) has incorrect indices of summation, if it is supposed to align with $E_{\text{MCHN}}$.
- L252: Typo, the definition of LogSumExp RHS uses $\rho^{i}$, when the function is defined in terms of $Mx$.

**Improvements**
(I1) Scaling laws (eq (1)) $N$ and $D$ are not defined. Is $N$ the same as the definition in line 153? Is $D$ the same as the definition in line 151? If so, please put these definitions closer to when they are first introduced

(I2) [L122] A bit non-standard notation for the partition function $Z_{\theta}$ -- generally this is expressed with integrating variable `dx` on the RHS

**Additional Comments**
(C1) The following statement is deceptive:

> As the attention layers and the FF layers contribute to the majority of the model’s parameters...

This is only true when we ignore the embedding and unembedding layers of a Transformer, which can be incredibly large with large vocab sizes. For smaller transformer models with large vocab sizes, it is possible that the embedding and unembedding matrices dominate the parameter count. 
Take for example `Model 1` in App F.1, where the number of parameters in the embedding+unembedding layers given the vocab size is ~96M parameters, whereas the reported total number of transformer params is 40M.
Please clarify, perhaps "The MHA and FF layers account for most of the parameters outside the embedding/unembedding layers" or "The fundamental operations of the Transformer are the MHA and FF layers", since parameter count is not the reason why you are focusing on these operations.

(C2) A gripe: the authors treat their version of the transformer as "standard", when it is arguably wrong or at least far from standard. This is summarized in the following quote from L419-420

> Also, modifications to the transformer blocks, such as additional layer normalization may contribute to the lower bound of the cross-entropy.

It is an open challenge to get high-performant transformers without normalization, and all modern transformers (including the models tested in this work!) use some form of normalization.
soundness: 2
strengths: (S1) **A novel theoretical analysis** of transformers. The paper describes a causal transformer as a sequence of energy functions (but these energy functions only model attention, see weaknesses) that serve as surrogates for optimizing a global energy function (eq. (6)) according to  the MM algorithm. This is an interesting perspective that could shed light on what the fundamental operation of a transformer could be.
summary: This paper introduces a theoretical framework for studying *causal* Transformer models and their scaling properties from the perspective of Associative Memories. The paper studies a "stacked" colelction of attractor energy functions to model a multi-layered Transformer where the argmin of one layer's energy is passed to the next layer. The next layer will further optimize its energy function in the local region around the output of the first layer, and so on. The authors validate their theoretical claims on vanilla transformers and GPT models.
weaknesses: (W1) **The equation for $E_{MCHN}^\beta(x)$ [L250] is incorrect (or at the very least undefined)**, since $M$ seems to be written as a linear operator over a vector of dimension $n=T_{\text{max}}d_{\text{emb}}$, which is not the attention operation of transformers because the `LogSumExp` in this equation is the summation index over the memorized examples instead of the sequence dimension... Could the authors please clarify this equation? If it is incorrect, the following propositions in the paper no longer hold.

(W2) **The equation for attention [L203] is incorrect**. $Q$, $K$, and $V$ should each have a sequence dimension, whereas no sequence dimension is mentioned in this equation. Is this intentional or a typo? Given my other reservations about the soundness of the analysis, I am inclined to believe that the authors misunderstand the attention operation of transformers.

(W3) **Assumption 1 is a BIG assumption**. Assumption 1 assumes that *every sample in the training set is memorized verbatim* as a vector of length $n$, which is not a reasonable assumption for an LLM or even any Associative Memory trained on abundant data. Can the authors discuss the implications of this assumptions not holding in real-world scenarios? Or provide some empirical evidence/theoretical justification for why their results might still be relevant despite this assumptions?

(W4) **Eq (5) shows optimal $x^{(t)}$ being passed to the next layer $t+1$**. However, the update rule of a vanilla Transformer is not guaranteed to minimize the energy $E_{t}(x)$ and I wonder how the authors could have made the assumption that each step returns the `argmin` of the energy function for a particular layer. Vanilla transformers (using the config specified in App. F.1) are allowed to learn any value matrix in the attention, which could lead to the defined energy *increasing* after each update and not satisfy the assumption of Eq. (5)

(W5) **Unconvincing (and arguably unrelated) empirical results**. This is a fundamental weakness of the paper. The empirical results study vanilla Transformer architectures: architectures that include layer normalization and feed-forward networks, do not use the proposed metric in $g(x)$ from Eq. (2), whose update rule is not guaranteed to minimize the energy function on L250, and whose attention operation is not reasonably modeled by Eq. (4). I don't know how to reconcile this weakness. Given the fundamental incompatibility with the experimental results and the actual Transformer architecture, I am surprised that the theoretical results "align closely" [L449] with the empirical results and would ask the authors to dig deeper into why this could be the case.

(W5) **Overloaded notations make exposition unclear**. The mathematical notation $d$ is a little confusing. It represents:

1. The number of samples in the training set or validation set [L148-149]
2. The embedding dimension of each token [L152]
3. The metric [L156]

**Summary**

As mentioned in the Strengths section above, I find the perspective of MM as a description of what the Transformer architecture is doing fascinating. I am willing to reconsider my score if the authors can clarify where I have fundamentally misunderstood the paper. I cannot recommend this paper in its current state for acceptance.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: Using the euclidean distance inside an Hopfield networks: This has already been done, and it has shown good performance in simple image retrieval tasks [1]. Hence, interestingly your approach could be seen as using universal Hopfield networks to model the energy function of the transformer model.    


[1] Millidge, Beren, et al. "Universal hopfield networks: A general framework for single-shot associative memory models." International Conference on Machine Learning. PMLR, 2022
soundness: 4
strengths: The analysis in well made, the problem tackled is an important one in the literature, and the work is well structured and clear. Furthermore, the connection with the nearest neighbour search is (to my knowledge) novel. The experimental evaluation is also well performed.
summary: The paper investigates the scaling laws of transformer models, under the lenses of energy-based models, specifically, Hopfield networks. The authors expand on the classic work of Ramsauer on MCHNs, and show how hierarchical models can be described using a simpler energy function, based on the euclidean distance instead of the dot product, that implements a kind of nearest neighbor. Such an energy function is not local anymore (that is, it is not the sum of the energies of the single layers).
weaknesses: The section on cross entropy loss is a little less clear than the rest of the work: Maybe you could add another ‘proposition’ or a statement that summarizes the results of the second part of the section?

Scaling laws in the related works: N and D are never introduced. 

“*It has been shown that the feed-forward layers can be interpreted as persistent memory vectors, and the weights of the attention and feed-forward sub-layers can be merged into a new all-attention layer without compromising the model’s performance*” This part is not well explained in the introduction in my opinion: could the authors dedicate to this a couple more sentences? Or, remove it and leave it as it is later in the paper?   

The provided energy function is not ‘local’, as it is in other energy based models such as standard Hopfield networks, predictive coding, etc. This is quite interesting. However, I feel there is little discussion on the consequences of this. Why is this important? What does it tell us that is different from the previous literature?

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: 1. Is it possible to collect the activation outputs up to 5% of the documents? If it's not possible due to limited compute, I understand. However, I think collecting activation outputs from 1% of the documents might be too limited. 

2. I believe you forgot to include the citation for RMSNorm, see the paragraph which describes the hypersphere volume on page 8 of the manuscript and where you mentioned LayerNorm.

3. In the description for the cross entropy loss $L$ in Proposition 4, you should perhaps replace that equation with eq. 13 in appendix D, since there is a mention of $c$ which is not found in the equation in the main text. 

4. There is an inconsistency in labeling the equations in the main text. Please organize them better since the mathematical descriptions provided in this paper are great. For example, I understand that $N \approx \frac{A l d_{max}}{T_{max}} n$ and $D = T_{max} n$ is mentioned in the prior sections of the manuscript. It is still certainly an important equation to be labeled/numbered.
soundness: 4
strengths: 1. The paper highlights the characteristics (i.e., pros and cons) of the Modern Hopfield Network's energy function given high dimensionality and continuous values, while providing an alternative energy function which maintains convexity. This energy function, i.e., logsumexp of a euclidean distance between a target $x$ and a stored pattern $\rho^{(i)}$, describes the dynamic of each Transformer block as an associative memory retrieval system.

2. Using the above (layer-wise) energy function, the paper also introduces a global energy function which accounts for the stacking of homogenous blocks in Transformer model. This particular global energy function enables the description of probability density function of the entire Transformer model. Using this description, the paper works out the optimal condition for the partition function $Z$ relating back to the scaling law from [Hoffmann et al (2022)](https://proceedings.neurips.cc/paper_files/paper/2022/file/c1e2faff6f588870935f114ebe04a3e5-Paper-Conference.pdf).

3. The paper also provides fantastic mathematical descriptions of memorization using their proposed energy function.
summary: Nowadays, it is common to scale up a Transformer model to achieve great performance across a variety of tasks. However, increasing the size of the model does not necessarily lead to performance improvement. Meanwhile, an interesting phenomenon occurs--as the model's ability to memorize the training samples so does its generalization ability. Consequently, this work provides a theoretical framework which elucidates the memorization process in the Transformer model via the lens of Modern Associative Memory or Hopfield models. The contributions include a novel analysis of memorization in each Transformer block as a function of modern Hopfield energy, and mathematical illustrations of the dependency of the size of the model and dataset respectively that is needed for the model to attain optimal performance. Such illustrations highlight the optimal cross entropy loss is bounded by a constant.
weaknesses: 1. The authors might not be aware of this paper due to its title. But the proposed energy function (eq. 3) in the paper was recently introduced/used in [Saha et al. (2023)](https://proceedings.mlr.press/v202/saha23a/saha23a.pdf), with the exception that they use $\beta$ in their formulation and descriptions of the function's properties. However, I do not think this fact compromises the novelty of the work.

2. Much of the current analyses of storage capacity in Hopfield models relies heavily on uncorrelated patterns. Consequently, the claim of exponential capacity does not necessarily hold true for natural (or correlated) data. Similarly, a weakness of this work is assumption 3, since the work focuses on text data. The authors can instead make a claim that the Transformer blocks are storing the latent data, obtained from passing the text data to an embedding layer ---which should be much more uncorrelated--- to reinforce their narrative.

3. The writing for the experimentation is unclear, especially to unfamiliar readers. For example, I think the authors should elaborate more on the findings of the mean and median nearest neighbor distances, and why such results corroborate eq. 12. Finally, I'm not sure what section 6.3 is trying to demonstrate. How was $\sigma$ chosen? To find $D^*$, you would have to train multiple models across a variety of data sizes $D^{(k)}$ for $k=1...M$ sizes. Perhaps, there should be an elaboration on the selection of the size of different models and each chosen data sizes in relation to $N \approx \frac{A l d_{max}}{T_{max}} n$ and $D = T_{max} n$. The inclusion of these details should fortify the reasonings behind these experiments.

---

confidence: 4
contribution: 3
initial_rating: 3
presentation: 3
questions: A slight nitpick, but on line 060 the authors state that the Hopfield network has only been recently generalized to continuous values. However, the continuous value Hopfield network was introduced only shortly after the 1982 paper. See Hopfield, J. "Neurons with graded response have collective computational properties like those of two-state neurons", 1984. The Dense Associative Memory (Modern Hopfield Network), introduced by Krotov and Hopfield 2016, was generalized to continuous values by Ramsauer et al. 2020, so perhaps this is what the authors intended? Reading on, this is discussed more in Section 2, subsection on Hopfield models, so my apologies for citing literature the authors are aware of! A slight rephrasing of line 060 may still be prudent.

The discussion on Hopfield models in Section 2 is excellent, but towards the end (e.g. works on the Energy Transformer and U-Hop networks) are seemingly only tangentially related to the main text. These could be left in for completion, but are they necessary for a reader's understanding of the main text? Perhaps with more explicit ties in later sections these would make more sense for a reader.

Section 3 lays out the model architecture. Section 3.1 discusses the associative memory formulation and how sequences of tokens may be interpreted as patterns in associative memories. Since each prediction will have access to a different number of previous tokens (e.g. predicting token 5 will have access to 4 previous tokens, predicting token 6 will have access to the previous 5, and so on) does this mean the associative memory here must be capable of handling retrieval for patterns of different lengths? Or are all patterns to be padded to a uniform length, such as in a traditional transformer, in which was the associative memory would have a much lower proportion of initial contents for subsequences that are early in the task (e.g. predicting the 5th token compared to predicting the 500th)? Perhaps I am misunderstanding the associative memory structure or pattern structure here? The initial paragraph of Section 3 seems to imply the latter, that all patterns are equal to the maximum token length. Is this correct? If so, how does this affect the retrieval of patterns that have only a small number of tokens available?

Definition 1 is a nice formalization of storage and retrieval in (continuous) associative memories. The condition that the intersection of any two balls must be null seems superfluous, however. If the intersection of B_i and B_j exists that would imply that the same point (i.e. any point in the intersection) would converge to both rho^i and rho^j, which is impossible. Unless this formalization of the associative memory somehow allows for multiple points of convergence? Is the condition of null intersection required for this definition, or could it be removed / discusses as a consequence of the requirement of convergence?

In reference to Assumption 2, are we free to assume that the held-out samples are likely to be stored in a similar fashion to the training set? In typical associative memories even samples from the same distribution as the training set are not stored, instead patterns from the training set have their basin of attraction (here, associated ball) expand to cover the hyper-volume of the validation distribution --- does this intuition not carry over to this architecture for some reason? If not, why not? 

The main difference between the MCHN energy function and the proposed energy function appears to be the removal of the regularization terms. Instead, the negative distances are fed into the LogSumExp function, rather than the pattern matrix times the probe vector. To my understanding, this implies that the proposed energy function requires all patterns in the training set to be kept explicitly in memory such that the calculation in Equation 3 can be performed. Later, in Equation 4, the energy function near the patterns of the training set is replaced by the nearest neighbor search which again requires knowledge of all training patterns. Is this correct? In typical associative memories we try to avoid storing the training set patterns directly, as this would defeat the purpose of a model that can learn to retrieve items. However, this paper seems to use the nearest-neighbor-search explicitly which is relatively counterintuitive in the context of associative memories as a whole. The formalization of the search as an energy function is interesting, but I am not yet convinced this leads to a useful associative memory --- if we have the training set explicitly stored, why not conduct a nearest-neighbor search over this data instead of using the energy function? 

Assumption 3 references back to the condition that the all balls around training set patterns are well separated. Is this not repeated from the discussion in Section 3? The authors also discuss briefly what this assumption means in practice, but is this property common in real datasets? Or is this a property that is nice to reason about, but in practice can be violated somewhat? It feels that the requirement in Assumption 3 limits the proposed energy function to only very nicely behaved data --- is that correct?

In Section 6.1, from line 465 onwards, the experimental radius of each ball is determined to be 10. Equation 12 would predict this radius to be 7.74. This seems like a fairly large discrepancy; some 30% larger than expected. Is this large a discrepancy expected in practice? The final line of the Section offers some respite, but perhaps increasing the number of documents to verify that the presented radius is larger that the actual value would be useful.
soundness: 4
strengths: This paper offers a theoretical analysis of the optimal loss for transformer-architecture models, which is corroborated in practical experiments. A novel energy function for the Dense Associative Memory (which the paper refers to as the Modern Continuous Hopfield Network) is proposed, which gives interesting behaviors in comparison to the original energy function from Ramsauer et al. 2020. The work is robustly proven with mathematical analysis that is rooted in both statistical mechanics and prior work on associative memories, and is hence of great quality. The clarity of the paper is reasonable, with the caveat that much of the theoretical work is based in mathematics that is reasonably dense, and is often deferred to the appendix. The theoretical findings appear very significant in the context of transformer-architecture models, although the experimental results may not be as strong in this regard.
summary: This paper looks at the performance of transformer-architecture models, particularly looking at the performance of such models with respect to the number of parameters. Recently proposed scaling laws may not reflect the true effect of parameter count on performance, and some smaller transformers outperform larger ones. The authors employ associative memories to model the behavior of transformers, specifically modelling each transformer unit with a single Hopfield network. Individual Hopfield energies are analyzed in a fashion similar to the Majorization Minimization technique. This allows an energy function to be associated with a transformer block, from which the paper forms a "global energy function" over the entire model. The global energy function relates the transformer size, the dataset size, and the optimal performance of the model. In particular, the authors find that the optimal loss is bounded by a constant term. Experiments are conducted that corroborate these results using a variety of transformer architectures on several datasets. Although some of these experiments were bounded by the computational resources (the authors note the work could be considered preliminary investigations) the results are still seemingly in agreement with the theoretical predictions made in previous Sections.
weaknesses: Several decisions / assumptions of the paper are made without discussion that would be useful for the reader. For example, it seems that all training set patterns to be stored explicitly in memory for the proposed energy function to be implemented, which is contrary to the traditional understanding of an associative memory. Almost all of the mathematical proofs are relegated to appendices, which is somewhat disappointing considering how vital these are to the main body of research.

Meta-review: The authors present a mostly theoretical analysis of how performance in Transformers scales with model size, based on an approximation of Transformers as a sequence of energy-based models. The reviewers agreed the analysis was interesting and novel, but two of the reviewers took issue with the empirical results in the paper, having concerns that the analysis did not support the conclusions of the theoretical framework. In response to this, the authors simply moved the empirical results to the appendix instead of improving the presentation or providing new results. Given the paper is not a theory paper in the sense of proving some novel theorem, but is meant to be a useful approximation to real networks, some empirical justification for the simplifying assumptions is necessary. Therefore I recommend against acceptance.

Decision: Reject