Title: Understanding Scaling Laws with Statistical and Approximation Theory for Transformer Neural Networks on Intrinsically Low-dimensional Data
Venue: NeurIPS 2024
Accepted: yes
Keywords: scaling laws

LLMs

approximation theory

statistical theory
Abstract: When training deep neural networks, a model's generalization error is often observed to follow a power scaling law dependent both on the model size and the data size. Perhaps the best known example of such scaling laws are for transformer-based large language models (**LLMs**), where networks with billions of parameters are trained on trillions of tokens of text. Yet, despite sustained widespread interest, a rigorous understanding of why transformer scaling laws exist is still missing. To answer this question, we establish novel statistical estimation and mathematical approximation theories for transformers when the input data are concentrated on a low-dimensional manifold. Our theory predicts a power law between the generalization error and both the training data size and the network size for transformers, where the power depends on the intrinsic dimension $d$ of the training data. Notably, the constructed model architecture is shallow, requiring only logarithmic depth in $d$. By leveraging low-dimensional data structures under a manifold hypothesis, we are able to explain transformer scaling laws in a way which respects the data geometry.  Moreover, we test our theory with empirical observation by training LLMs on natural language datasets. We find the observed empirical scaling laws closely agree with our theoretical predictions. Taken together, these results rigorously show the intrinsic dimension of data to be a crucial quantity affecting transformer scaling laws in both theory and practice.

Understanding Scaling Laws with Statistical and Approximation Theory for Transformer Neural Networks on Intrinsically Low-dimensional Data

Abstract

When training deep neural networks, a model's generalization error is often observed to follow a power scaling law dependent both on the model size and the data size. Perhaps the best known example of such scaling laws are for transformerbased large language models ( LLMs ), where networks with billions of parameters are trained on trillions of tokens of text. Yet, despite sustained widespread interest, a rigorous understanding of why transformer scaling laws exist is still missing. To answer this question, we establish novel statistical estimation and mathematical approximation theories for transformers when the input data are concentrated on a low-dimensional manifold. Our theory predicts a power law between the generalization error and both the training data size and the network size for transformers, where the power depends on the intrinsic dimension d of the training data. Notably, the constructed model architecture is shallow, requiring only logarithmic depth in d . By leveraging low-dimensional data structures under a manifold hypothesis, we are able to explain transformer scaling laws in a way which respects the data geometry. Moreover, we test our theory with empirical observation by training LLMs on natural language datasets. We find the observed empirical scaling laws closely agree with our theoretical predictions. Taken together, these results rigorously show the intrinsic dimension of data to be a crucial quantity affecting transformer scaling laws in both theory and practice.

1 Introduction

Deep learning has made remarkable breakthroughs in various real-world applications, such as natural language processing [Graves et al., 2013, Bahdanau et al., 2014, Liu et al., 2023, Vaswani et al., 2017], computer vision [Krizhevsky et al., 2012, Goodfellow et al., 2014, Song et al., 2020], healthcare [Miotto et al., 2018], and robotics [Gu et al., 2017]. A neural scaling law between the generalization error (or test loss) and several quantities, including the model size, the training data size, and the amount of compute, plays a key role in the performance of neural networks. Perhaps the best known example of such scaling laws are for transformer-based LLMs. Recent works in Hestness et al. [2017], Rosenfeld et al. [2019], Kaplan et al. [2020], Bahri et al. [2021] demonstrated a power law between the test loss and the network size, the training data size, and the amount of compute for transformer-based LLMs. Yet, despite sustained widespread interest, a rigorous understanding of why transformer scaling laws exist is still missing.

Understanding the theory behind neural scaling laws provides invaluable insights into practical applications of deep learning. A mathematical principal of neural scaling laws enables researchers and practitioners to describe and analyze the performance of neural networks with precision and

rigor. The neural scaling law between the generalization error and the network size can be partially explained via neural network representation theory [Yarotsky, 2016]. Further, the neural scaling law between the generalization error and the training data size n can be explained via statistical estimation theory. For feedforward neural networks [Schmidt-Hieber, 2020] and convolutional residual networks [Oono and Suzuki, 2019], a generalization error bound has been established for regression. SchmidtHieber [2020], Oono and Suzuki [2019] predicted Generalization Error ∼ n -c/D where n is the training data size, D is the data dimension and c is a constant. This predicted rate of convergence is extremely slow for high dimensional data when D is large, while the rate of convergence observed in real-world applications is significantly faster, which reveals a gap between theory and practice.

This gap can be bridged by exploiting low-dimensional structures of data. Real-world data sets often exhibit low-dimensional geometric structures due to rich local regularities, global symmetries, or repetitive patterns [Tenenbaum et al., 2000, Roweis and Saul, 2000]. According to Min et al. [2023, Figure 1], the intrinsic dimension of CIFAR-100, CelebA and ImageNet datasets are about 20 , 20 and 40 respectively. When the low-dimensional geometric structure of data is modeled by a manifold, the predicted scaling for regression, classification and distribution estimation becomes Generalization Error ∼ n -c/d , where n is the training data size, d is the intrinsic dimension of the data manifold, and c is a constant [Chen et al., 2022, Liu et al., 2021, Dahal et al., 2022, Nakada and Imaizumi, 2020]. In Sharma and Kaplan [2022], the neural scaling law between the test loss and the network size was predicted to be Test loss ∼ ( size ) -4 /d where d is the intrinsic dimension of data. While the theoretical studies focus on feedforward neural networks [Chen et al., 2022, Nakada and Imaizumi, 2020] and convolutional residual networks [Liu et al., 2021], a generalization to transformer-based neural networks [Vaswani et al., 2017] is of great interest but widely open.

This paper establishes mathematical approximation and statistical estimation theories to predict and justify the scaling law between the generalization error and the model/data size for transformer neural networks. We consider regression of a β -Hölder continuous function f : M → R where M is a d -dimensional compact Riemannian manifold isometrically embedded in R D . After embedding the input x ∈ M ⊂ R D to a proper sequence, we apply a transformer network on the embedded sequence to learn the function f . Our main results are on the statistical estimation and universal approximation theories of Hölder continuous functions on M by transformer neural networks.

Statistical Theory: In Theorem 1, we consider the global empirical risk minimizer ˆ T n from n i.i.d. training data { ( x i , f ( x i )) } n i =1 , given by

under a properly chosen transformer network architecture T . We prove that, the generalization error of ˆ T n satisfies

where Q denotes the distribution of x , and ˜ O hides constants and log n terms.

Approximation Theory: In Theorem 2, we construct a transformer network to universally approximate β -Hölder continuous functions on M with an arbitrarily given accuracy ε . Notably, the network is shallow, requiring only O ( log( d ) ) independent of the desired accuracy ϵ to approximate f locally. This highlights a major advantage of Transformers over feed-forward ReLU networks, which require O ( log( 1 ϵ ) ) layers to achieve the same accuracy.

In our proof, we embed the entries of x = [ x 1 , . . . , x D ] ∈ M into tokens such that the x i 's appear in a sequence. Our proof for the approximation theory explicitly constructs transformers to realize the interaction between different tokens efficiently via a crucial Interaction Lemma 3. This lemma allows us to flexibly implement many common operations including addition, multiplication, and parallelization, and so may of independent interest. In our proof for the statistical theory, we calculate the covering number of our transformer network class, which is also of independent interest.

Neural Scaling Laws and the Intrinsic Dimension: Our generalization error bound in (2) predicts the following neural scaling law between the generalization error and the data size n :

with sufficient data. Our approximation theory in Theorem 2 predicts the following neural scaling law between the approximation error and the network size N :

for a sufficiently large network class T . Our prediction of the power scaling law is consistent with our own empirical observations, and those in Kaplan et al. [2020] and Biderman et al. [2023]. More importantly, our theory quantifies the power α D , α N in terms of the intrinsic dimension of data.

Experimental Validation on LLMs: After establishing our theory we seek to validate it in practice by predicting empirical scaling laws for LLMs trained on natural language data. To test our predictions for the data scaling law, we pretrain a series of small (125 million parameter) LLMs on three datasets [Gokaslan et al., 2019, Eldan and Li, 2023, Kocetkov et al., 2022]. We find close agreement ( ± 0 . 02 ) between our predicted scaling exponent α D and the observed exponents ˆ α D . To evaluate our predictions for the model scaling exponent α N , we rely on publicly available scaling suites [Biderman et al., 2023, Radford et al., 2019] whose intrinsic data dimensions we can estimate. We find our predictions are still close but less accurate for α N . Finally, we carry out a series of ablations investigating factors impacting the estimated intrinsic data dimension d . For a fixed dataset, we find the estimated d is stable with respect to several factors including the model size, model embedding dimension, and context length 1 .

In summary, we make the following contributions:

A novel approximation theory for transformers approximating Hölder continuous functions on a d -dimensional manifold, requiring O (log( d )) depth independent of the accuracy ϵ .

A novel computation of the covering number of our transformer network class. This is used to establish generalization bounds exponentially depending on the intrinsic dimension d .

Empirical experiments demonstrating our theory predicts data scaling laws for LLMs as a function of the estimated intrinsic data dimension d .

An empirical study of several factors affecting the estimated intrinsic data dimension for transformers including model size, embedding dimension, layer depth, and context length.

We will present our main theory in Section 2, numerical validation of our theory and the prediction of neural scaling laws in Section 3. We will discuss related work in Section 4 and conclude our paper in Section 5. Our pre-training hyperparameters are given in Appendix A. The derivation of neural scaling laws is presented in Appendix B. Our notation is given in Appendix C, and proofs are presented in Appendix E and F.

2 Transformer Generalization and Approximation Theory

This paper establishes statistical estimation and mathematical approximation theory of transformers for the regression of Hölder functions on a low-dimensional manifold. We start by defining transformer neural networks.

2.1 Transformer Neural Networks

Definition 1 (Transformer Neural Network) . We define a transformer neural network T as a composition of functions of the form

which is parameterized by

L T : The number of transformer blocks B i in T .

m : The maximum number of attention heads per transformer block.

L FFN : The max depth of the feed-forward layers per block.

w FFN : The max width of the feed-forward layers per block.

d embd : The token embedding dimension.

D : The input dimension.

l : The number of hidden tokens.

κ : A bound on the magnitude of network parameters.

We define each component of the composition as

x ∈ R D is the input.

A linear embedding layer E : R D → R d embd × l . In this work we will always take E = E ′ ◦ U where U ∈ R l × D and E ′ ∈ R d embd × 1 applied columnwise is fixed. We call embedded output H = E ( x ) the first embedding matrix whose columns are referred to as tokens .

PE ∈ R d embd × l is a fixed matrix implementing the transformer positional encoding .

Transformer blocks B i : R d × l → R d × l for i ∈ { 1 , ..., L T } which are residual compositions of multi-headed attention ( MHA ) layers MHA and feed-forward layers FFN acting tokenwise.

A decoding layer D : R d embd × l → R which is fixed to outputting the first element of the last column.

We use the ReLU activation function σ ( x ) = max(0 , x ) in the network.

For a complete definition of the components of the transformer neural networks, we refer to Appendix D. The transformer network T may sometimes be written T θ which makes explicit the dependence on learnable weights θ . We can also define a class of transformer neural networks T of interest.

Definition 2 (Transformer Network Class) . We define a class of transformer networks as

T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) = { T θ | T θ is in (5) with at most m attention heads in each block,

L FFN layer feed-forward networks with hidden width w FFN , d embd token dimension , l hidden tokens ,

where ∥ T θ ∥ L ∞ ( R D ) ≤ R bounds the output of T and ∥ θ ∥ ∞ bounds the weight magnitude of T θ .

2.2 Assumptions

We consider a manifold M and the target function f : M→ R , with the following assumptions:

Assumption 1 (Manifold) . Let M be a compact Riemannian manifold with intrinsic dimension d isometrically embedded in R D . Because M is compact, there exists M > 0 such that ∥ x ∥ ∞ ≤ M for x ∈ M . Additionally, we assume M has positive reach τ > 0 .

Assumption 2 (Target function) . The target function f : M→ R is β -Hölder continuous on M , for some 0 < β ≤ 1 and Hölder constant H f > 0 , and in addition ∥ f ∥ L ∞ ( M ) ≤ R for some R > 0 .

In Assumption 1, the reach [Federer, 1959, Aamari et al., 2019] τ of M can be defined as

̸

Informally, reach is the smallest distance at which a projection onto the manifold is no longer unique. In practice this can be used to establish a bound on the number of charts covering the manifold.

2.3 Transformer Generalization Theory

Given n training samples { ( x i , f ( x i )) } n i =1 where { x i } n i =1 are i.i.d. samples of a distribution Q supported on M , we aim to learn an approximation ˆ T n to f by minimizing the empirical risk (1) over

a class of a transformer neural networks T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) . The corresponding generalization error is given by

If M and f satisfy Assumptions 1 and 2, we prove the following generalization error bound.

Theorem 1. Let M,τ,R,H f > 0 , 0 < β ≤ 1 , d, D ∈ N , M and f satisfy Assumption 1 and 2 respectively. Given n training samples { ( x i , f ( x i )) } n i =1 where { x i } n i =1 are i.i.d. samples of a distribution Q supported on M , if we use the transformer neural network class T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) with parameters

in the empirical risk minimization (1) , where O ( · ) hides terms in C M (the number of charts), D,H f , M , then the empirical risk minimizer ˆ T n given by (1) satisfies

where ˜ O hides logarithmic terms in n, d and linear terms in C M .

Theorem 1 is proved in Appendix F, via a bias-variance decomposition. The bias represents the approximation error of f by transformer neural networks, and the variance represents the stochastic error in the parameter estimation of transformer neural networks. To quantify the bias, we explicitly construct a transformer neural network to universally approximate β -Hölder continuous functions on M , to be detailed in Section 2.4. The variance is bounded by a novel calculation of the covering number of the transformer network class used in Theorem 1.

2.4 Transformer Approximation Theory

Theorem 2. Let M,τ,R,H f > 0 , 0 < β ≤ 1 , d, D ∈ N and M satisfy Assumption 1. For any ϵ ∈ (0 , 1) , there exists a transformer neural network class T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) with parameters

where O ( · ) hides terms in C M , D, H f , τ , such that for any target function f satisfying Assumption 2, if the network parameters θ are properly chosen, then the network yields a function T θ ∈ T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) with the approximation error

Theorem 2 is proved in Appendix E.2. In our proof, we decompose f ( x ) as a sum of terms over local neighborhoods U 1 , ..., U C M ⊆ M covering M . Approximations on overlapping neighborhoods containing x will then be combined via a partition of unity ( PoU ) { ρ n } C M n =1 which subordinates { U n } C M n =1 . This will give us the expression f ( x ) = ∑ C M n =1 f n ( x ) 1 U n ( x ) with f n = fρ n : M→ R . On each local neighborhood U n , we project the input x ∈ M ⊆ R D to the tangent coordinate in [0 , 1] d . This will give us the following local decomposition of the target function:

where ˜ f n = f n ◦ ϕ -1 n : [0 , 1] d → R and ϕ n : M→ [0 , 1] d is a projection onto the local tangent space. We then construct transformers to approximate the ˜ f n , ϕ n , 1 U n components in (7). A diagram of the constructed transformer network approximating f : M→ R is given in Figure 1. The following key lemma is used to efficiently approximate each low-dimensional function ˜ f n on d -dimensional coordinates.

H1

H2

H3

H4

Ho(log(d)

B1

→

da

B2

B

3

iu

f;(c)·1c;(z)

手(α)

CERD

LUCM

LUCM

1UCM

Legend

Φ1

n-iu

Chart

Determination

Embedding

Φ2

于2

f2-1U

Approximation

Pairing

ΦCM

fCM

fCm-1Uem

Lemma 1. Let H f , R > 0 , d ∈ N and 0 < β ≤ 1 . For any ϵ ∈ (0 , 1) , there exists a transformer neural network class T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) with parameters

where O ( · ) hides terms in H f , such that, for any β -Hölder continuous function f : [0 , 1] d → R , with Hölder constant no more than H f and ∥ f ∥ L ∞ ([0 , 1] d ) ≤ R , if the network parameters θ are properly chosen, this transformer network yields a function T θ ∈ T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) such that

Lemma 1 is proved in Appendix E.1. We develop a novel lemma Interaction Lemma 3, implementing a highly-sparse pairwise interaction between two arbitrary tokens h t 1 , h t 2 , as a crucial architectural building block allowing us to easily implement more complex functions, architecture serialization, and parallelization (7). This result highlights a distinct advantage of transformer function approximation over ReLU function approximation [Yarotsky, 2016]: Atransformer network only needs a constant O ( log( d ) ) number of layers to approximate f : [0 , 1] d → R independent of the desired accuracy ϵ . In contrast, the depth of ReLU feed-forward networks is in the order of log( ϵ -1 ) [Yarotsky, 2016] This is desirable from an empirical point of view, where wider networks instead of deeper ones tend to achieve superior performance [Kaplan et al., 2020, Lee et al., 2020].

3 Predicting Empirical Scaling Laws and Validation on LLMs

Our theory provides practical insights by predicting neural scaling laws for transformers, as given in (3) and (4), by explicitly quantifying the data scaling exponent α D and the model scaling exponent α N as a function of the intrinsic dimension ( ID ) d . If we assume the language modeling objective has Lipschitz regularity such that β = 1 in Assumption 2, then Theorem 1 predicts the scaling law between the squared generalization error and the data size n , as given in (3), with α D = 2 2+ d , and the model scaling law with exponent given by α N = 2 d . For a full derivation refer to Section B. We will observe how well our theory predicts these exponents both by pretraining small models from scratch and evaluating existing open-source model suites [Biderman et al., 2023].

In the following, we denote α D and α N as the scaling exponents predicted by our theory, where we numerically estimate the intrinsic dimension of data, denoted by d D . The empirical exponents are

OpenWebText (log) val loss vs. (log) number of train samples

Bigcode-SQL (log) val loss vs. (log) number of train samples

Tiny Stories (log) val loss vs. (log) number of train samples

0.70

Ground truth

Empirical fit: αp=0.112

0.15

Ground truth

Empirical fit: αo=0.126

0.30

Ground truth

Empirical fit: αo=0.142

Theory: αo=0.117 (ID ~ 15.1)

0.10

Theory: αo=0.110 (ID ~ 16.21)

Theory: αp=0.125 (ID ~ 14.04)

0.65

0.25

SSOl

SSO

SSO

e

0.60

0.05

IA (601)

0.20

(601)

(60l)

F 000

0.15

0.55

0.05

0.10

0.50

0.05

0.10

7.00

7.25

7.50

7.75

8.00

8.25

8.50

8.75

9.00

8.00

8.25

8.50

8.75

9.00

9.25

9.50

9.75

10.00

6.75

7.00

7.25

7.50

7.75

8.00

8.25

8.50

8.75

(log) Num samples

(log) Num samples

(log) Num samples

denoted by ˆ α D and ˆ α N . To obtain the data scaling exponent ˆ α D , we plot the test loss (comparable to squared error) versus the data size n in log-log scale, fit the log-log curve with a line, and then obtain ˆ α D from the magnitude of the slope. The model scaling exponent ˆ α N is obtained similarly.

Estimating the ID of Text To predict scaling exponents, we must first estimate the intrinsic dimension of our pretraining dataset D . While we can do this directly for image datasets [Russakovsky et al., 2014, Pope et al., 2021], we cannot do this directly for textual datasets. Instead, we will estimate the intrinsic dimension of the input data by estimating the intrinsic dimension of token embeddings. Specifically, we will represent each input token with its corresponding final-layer token embedding. Given a pretraining test set D test we embed a random l = 1024 length subsequence from each document D k ∈ D test . We then randomly sub-sample 32 final-layer tokens from the embedded subsequence and shuffle together all the embeddings. To estimate the ID of the embeddings we use the Maximum Likelihood Estimation ID algorithm [Levina and Bickel, 2004, Pedregosa et al., 2011] with K = 20 neighbors. We split the sampled embedding tokens into batches of 4096 and run the MLE estimator on each batch, averaging together for the final result. Unless otherwise specified, we embed each document D k ∈ D test using a 125 million parameter model M with L T = 12 layers and embedding dimension d embd = 768 . We first pretrain M on the full D train for 200 , 000 steps or until convergence.

Intrinsic dimension predicts the empirical data scaling exponent ˆ α D To validate our prediction of the dataset scaling exponent α D we pretrain a series of 125 million parameter GPT-style LLMs on three different datasets: OpenWebText [Gokaslan et al., 2019], the SQL portion of The Stack [Kocetkov et al., 2022], and Tiny Stories [Eldan and Li, 2023]. We train across three orders of dataset size to fit scaling laws. Detailed hyperparameters can be found in the Appendix A. We report the observed scaling laws in log scale in Figure 2. In addition, we plot our predicted test loss whose slope is given by α D = 2 2+ d D where d D is the our estimated intrinsic dimension.

Empirically, we find all three datasets produce nearly log-linear laws whose exponents lie between 0 . 1 < ˆ α D < 0 . 15 . Tiny Stories has the largest exponent, indicating the fastest rate of convergence, followed by the SQL dataset, followed by OpenWebText. This matches our rough intuition, since Tiny Stories is a synthetically generated dataset with less complexity than the other two datasets and thus easier to learn. The predicted exponents α D generally over-estimate ˆ α D but otherwise closely match up to ± 0 . 02 absolute error. Additionally, the predicted α D reflect the previouly mentioned differences in complexity of pretraining datasets. In particular, Tiny Stories has a smaller estimated ID than both OpenWebText and SQL, resulting in a larger predicted α D as desired.

Predicting empirical model scaling exponent ˆ α N with intrinsic dimension To validate our predictions of the model scaling exponent α N = 2 d D , we evaluate two model scaling suites: GPT2 [Radford et al., 2019] and Pythia [Biderman et al., 2023]. We refer to Kaplan et al. [2020] for GPT2's ˆ α N and estimate ˆ α N using OpenWebText as a proxy for GPT2's pretraining data. We compute α N for Pythia by evaluating each model on The Pile's test set [Gao et al., 2021]. We also estimate d D on the publicly available pretraining data to predict α N . The results are reported in Figure 3. Our predicted α N under-estimates the empirical ˆ α N . We conjecture this to be due to a number of factors including possible under-training of the largest models and the intrinsic entropy of the data distribution.

gpt2 model size vs. val loss

0.60

0.34

0.55

0.32

loss

0.50

SSOI

0.30

0.28

Val

0.45

Val

0

0

0.40

(601)

0.26

0.24

0.35

Ground truth

0.22

Ground truth

0.30

Empirical fit: αn = 0.073

Empirical fit:αn = 0.068

0.25

Theory: αn =0.132 (ID ~15.1)

0.20

Theory: αn =0.129 (ID ~ 15.56)

5.5

6.0

6.5

7.0

7.5

8.0

8.5

8.6

8.8

9.0

9.2

9.4

9.6

9.8

(log)Model size

(log)Model size

Estimated ID vs.Embedding dimension

20

Pythia

18

19

17

Estimated ID

Estimated ID

16

15

14

16

13 -

410M 1.4B

2.8B

6.9B

12.8B

384

768

1536

Model Size (billions of parameters)

Embedding dimension

OpenWebText EstimatedID vs.Layer depth

OpenWebTextEstimated IDvs.Contextlength

17

15.5

15.0 -

Estimated ID

16

14.5

15

14.0

1A

13.5

13.0

13

5

6

9

10

11

200

400

600

800

1000

Layer depth

Context length

Ablating the impact of model architecture on the estimated ID Practical application of our theory relies on a good estimate of the intrinsic dimension. However, there are many factors potentially biasing our estimate. Of particular interest is the embedding model's embedding dimension, depth, context length, and number of parameters. We ablate these factors in Figure 4, plotting estimated ID against each factor.

Overall, we find the estimated ID is fairly stable across each factor. As the number of parameters increases, the estimated ID of The Pile slightly increases from 15.56, via a 410 million parameter model, to 20.02 with a 12.8 billion parameter model. ID on OpenWebText behaves similarly when increasing the embedding dimension, increasing from 15.56 when d embd = 768 to 18.68 when d embd = 1536 . When fixing a single model, we find the ID across intermediate embedding layers is small initially but then increases and decreases again, stabilizing around the ID of the final layer. We observe that the ID appears to inversely correlate with sequence length, decreasing from 15 . 86 for very short sequences to 12 . 9 for sequences around 1024 tokens.

Predicting α N from α D (and vice versa) without estimating ID Above we estimated α D and α N by first estimating the intrinsic dimension d for a model's pretraining dataset. However, estimating d may not always be possible when pretraining data is not public. Alternatively, we can predict α D in terms of α N (and vice versa) without ever needing to estimate d :

See Table 1 for ID-free estimations of empirically observed exponents in the literature [Sharma and Kaplan, 2022, Hoffmann et al., 2022].

4 Related Work

The theoretical properties and advantages of transformers have been studied from many different perspectives [Jelassi et al., 2022, Zhang et al., 2022, Bai et al., 2023, Pérez et al., 2021, Sanford et al., 2024]. Most related to us are Yun et al. [2019], Edelman et al. [2022], Wei et al. [2022], Takakura and Suzuki [2023] in which transformers were studied from an approximation viewpoint. The work in Yun et al. [2019] proved that transformer models are universal approximators of continuous permutation equivariant sequence-to-sequence functions with compact support, while the network size suffers from the curse of dimensionality (the number of entries in the input sequence). Takakura and Suzuki [2023] studied the approximation and estimation ability of Transformers as seq-to-seq functions with infinite dimensional input, where anisotropic smoothness avoids the curse of dimensionality.

In the applications of Large Language Models (LLMs), empirical findings have demonstrated some correlation between the performance of transformers and the low-dimensional data structures [Razzhigaev et al., 2023, Min et al., 2023, Aghajanyan et al., 2020, Pandey, 2024]. Razzhigaev et al. [2023] investigated the intrinsic dimension of embeddings in transformer architectures, and suggested an encoder and decoder embedding property. Most similar to our work is Sharma and Kaplan [2022] which demonstrates an empirical connection between neural scaling laws and the intrinsic dimension of data. While they briefly discuss predictions for LLMs, their theory works best for predicting student-teacher model setups and image classification tasks which seem to enjoy more regularity (and a faster rate of convergence) than language modeling. Despite these empirical findings, we are not aware of any rigorous theoretical justification connecting the scaling laws of transformers with the intrinsic dimension of data. Our paper complements this line of research with statistical estimation and mathematical approximation theories which well-predict the behavior observed in practice.

5 Conclusion

Conclusion This paper establishes statistical and approximation theory results for transformers approximating Hölder continuous functions on low-dimensional data manifolds. The resulting bound on the generalization error suffers from exponential dependence only in the intrinsic dimension d . The constructed approximations of low-dimensional functions are shallow, requiring only O ( log( d ) ) layers independent of the desired accuracy. We demonstrate this theory is accurate in practice by predicting scaling laws in both model size and data size for LLMs trained on natural language datasets. We pay careful attention to the sensitivity of the estimated intrinsic data dimension, finding it is relatively stable with respect to several relevant hyperparameters.

Limitations and Broader Impact One important question unanswered by this work is how the intrinsic data dimension may affect the computational scaling exponent α C . Future work may investigate this direction. Additionally, our empirical experiments make the simplifying assumption that the underlying target function possesses Lipschitz regularity ( β = 1 ). Better estimates of the correct regularity would likely improve the accuracy of our predictions. More broadly, our work improves fundamental understanding of transformer-based LLMs and improves our ability to theoretically and safely predict future capabilities.

References

Eddie Aamari, Jisu Kim, Frédéric Chazal, Bertrand Michel, Alessandro Rinaldo, and Larry Wasserman. Estimating the reach of a manifold, 2019.

Armen Aghajanyan, Luke Zettlemoyer, and Sonal Gupta. Intrinsic dimensionality explains the effectiveness of language model fine-tuning. ArXiv , abs/2012.13255, 2020. URL https://api. semanticscholar.org/CorpusID:229371560 .

Jean-Yves Audibert and Alexandre B Tsybakov. Fast learning rates for plug-in classifiers. 2007.

Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0473 , 2014.

Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jaehoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. ArXiv , abs/2102.06701, 2021. URL https://api.semanticscholar.org/ CorpusID:231918701 .

Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection, 2023.

Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal. Pythia: A suite for analyzing large language models across training and scaling, 2023.

Minshuo Chen, Haoming Jiang, Wenjing Liao, and Tuo Zhao. Efficient approximation of deep relu networks for functions on low dimensional manifolds. Advances in neural information processing systems , 32, 2019.

Minshuo Chen, Haoming Jiang, Wenjing Liao, and Tuo Zhao. Nonparametric regression on lowdimensional manifolds using deep relu networks: Function approximation and statistical recovery. Information and Inference: A Journal of the IMA , 11(4):1203-1253, 2022.

John Conway and N. Sloane. Sphere Packings, Lattices and Groups , volume 290. 01 1988. ISBN 978-1-4757-2018-1. doi: 10.1007/978-1-4757-2016-7.

Biraj Dahal, Alexander Havrilla, Minshuo Chen, Tuo Zhao, and Wenjing Liao. On deep generative models for approximation and estimation of distributions on manifolds. Advances in Neural Information Processing Systems , 35:10615-10628, 2022.

Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning , pages 5793-5831. PMLR, 2022.

Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english?, 2023.

Herbert Federer. Curvature measures. Transactions of the American Mathematical Society , 93(3): 418-491, 1959.

Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling. CoRR , abs/2101.00027, 2021. URL https://arxiv.org/abs/2101.00027 .

Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus. http: //Skylion007.github.io/OpenWebTextCorpus , 2019.

Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems , 27, 2014.

Alex Graves, Abdel-rahman Mohamed, and Geoffrey Hinton. Speech recognition with deep recurrent neural networks. In 2013 IEEE international conference on acoustics, speech and signal processing , pages 6645-6649. Ieee, 2013.

Shixiang Gu, Ethan Holly, Timothy Lillicrap, and Sergey Levine. Deep reinforcement learning for robotic manipulation with asynchronous off-policy updates. In 2017 IEEE international conference on robotics and automation (ICRA) , pages 3389-3396. IEEE, 2017.

Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md Mostofa Ali Patwary, Yang Yang, and Yanqi Zhou. Deep learning scaling is predictable, empirically. arXiv preprint arXiv:1712.00409 , 2017.

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models, 2022.

Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems , 35:37822-37836, 2022.

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 , 2020.

Yongdai Kim, Ilsang Ohn, and Dongha Kim. Fast convergence rates of deep neural networks for classification, 2019. URL https://arxiv.org/abs/1812.03599 .

Denis Kocetkov, Raymond Li, Loubna Ben Allal, Jia Li, Chenghao Mou, Carlos Muñoz Ferrandis, Yacine Jernite, Margaret Mitchell, Sean Hughes, Thomas Wolf, Dzmitry Bahdanau, Leandro von Werra, and Harm de Vries. The stack: 3 tb of permissively licensed source code. Preprint , 2022.

Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems , 25, 2012.

Jaehoon Lee, Lechao Xiao, Samuel S Schoenholz, Yasaman Bahri, Roman Novak, Jascha SohlDickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent *. Journal of Statistical Mechanics: Theory and Experiment , 2020(12): 124002, December 2020. ISSN 1742-5468. doi: 10.1088/1742-5468/abc62b. URL http: //dx.doi.org/10.1088/1742-5468/abc62b .

Elizaveta Levina and Peter Bickel. Maximum likelihood estimation of intrinsic dimension. In L. Saul, Y. Weiss, and L. Bottou, editors, Advances in Neural Information Processing Systems , volume 17. MIT Press, 2004. URL https://proceedings.neurips.cc/paper_files/paper/2004/ file/74934548253bcab8490ebd74afed7031-Paper.pdf .

Hao Liu, Minshuo Chen, Tuo Zhao, and Wenjing Liao. Besov function approximation and binary classification on low-dimensional manifolds using convolutional residual networks. In International Conference on Machine Learning , pages 6770-6780. PMLR, 2021.

Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing. ACM Computing Surveys , 55(9):1-35, 2023.

Zeping Min, Qian Ge, and Zhong Li. An intrinsic dimension perspective of transformers for sequential modeling, 2023. URL https://openreview.net/forum?id=0UzYWLzPBjA .

Riccardo Miotto, Fei Wang, Shuang Wang, Xiaoqian Jiang, and Joel T Dudley. Deep learning for healthcare: review, opportunities and challenges. Briefings in bioinformatics , 19(6):1236-1246, 2018.

Ryumei Nakada and Masaaki Imaizumi. Adaptive approximation and generalization of deep neural network with intrinsic dimensionality. Journal of Machine Learning Research , 21(174):1-38, 2020.

Kenta Oono and Taiji Suzuki. Approximation and non-parametric estimation of resnet-type convolutional neural networks. In International conference on machine learning , pages 4922-4931. PMLR, 2019.

Rohan Pandey. gzip predicts data-dependent scaling laws, 2024. URL https://arxiv.org/abs/ 2405.16684 .

Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research , 12:2825-2830, 2011.

Phillip E. Pope, Chen Zhu, Ahmed Abdelkader, Micah Goldblum, and Tom Goldstein. The intrinsic dimension of images and its impact on learning. ArXiv , abs/2104.08894, 2021. URL https: //api.semanticscholar.org/CorpusID:233296562 .

Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing-complete. Journal of Machine Learning Research , 22(75):1-35, 2021. URL http://jmlr.org/papers/v22/20-302.html .

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019. URL https://api.semanticscholar.org/ CorpusID:160025533 .

Anton Razzhigaev, Matvey Mikhalchuk, Elizaveta Goncharova, Ivan Oseledets, Denis Dimitrov, and Andrey Kuznetsov. The shape of learning: Anisotropy and intrinsic dimensions in transformerbased models. ArXiv , abs/2311.05928, 2023. URL https://api.semanticscholar.org/ CorpusID:265128603 .

Jonathan S Rosenfeld, Amir Rosenfeld, Yonatan Belinkov, and Nir Shavit. A constructive prediction of the generalization error across scales. arXiv preprint arXiv:1909.12673 , 2019.

Sam T Roweis and Lawrence K Saul. Nonlinear dimensionality reduction by locally linear embedding. science , 290(5500):2323-2326, 2000.

Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael S. Bernstein, Alexander C. Berg, and Li FeiFei. Imagenet large scale visual recognition challenge. CoRR , abs/1409.0575, 2014. URL http://arxiv.org/abs/1409.0575 .

Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and logarithmic depth. ArXiv , abs/2402.09268, 2024. URL https://api.semanticscholar.org/ CorpusID:267657804 .

Johannes Schmidt-Hieber. Nonparametric regression using deep neural networks with ReLU activation function. The Annals of Statistics , 48(4):1875 - 1897, 2020. doi: 10.1214/19-AOS1875. URL https://doi.org/10.1214/19-AOS1875 .

Utkarsh Sharma and Jared Kaplan. Scaling laws from the data manifold dimension. J. Mach. Learn. Res. , 23(1), jan 2022. ISSN 1532-4435.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Shokichi Takakura and Taiji Suzuki. Approximation and estimation ability of transformers for sequence-to-sequence functions with infinite dimensional input. In International Conference on Machine Learning , 2023. URL https://api.semanticscholar.org/CorpusID:258967774 .

Joshua B Tenenbaum, Vin de Silva, and John C Langford. A global geometric framework for nonlinear dimensionality reduction. science , 290(5500):2319-2323, 2000.

Loring W. Tu. An Introduction to Manifolds . 01 2011. ISBN 978-1-4419-7399-3.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 12071-12083. Curran Associates, Inc., 2022. URL https://proceedings.neurips.cc/paper_files/paper/2022/file/ 4ebf1d74f53ece08512a23309d58df89-Paper-Conference.pdf .

Dmitry Yarotsky. Error bounds for approximations with deep relu networks. Neural networks : the official journal of the International Neural Network Society , 94:103-114, 2016. URL https://api.semanticscholar.org/CorpusID:426133 .

Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? ArXiv , abs/1912.10077, 2019. URL https://api.semanticscholar.org/CorpusID:209444410 .

Yufeng Zhang, Boyi Liu, Qi Cai, Lingxiao Wang, and Zhaoran Wang. An analysis of attention via the lens of exchangeability and latent variable models. arXiv preprint arXiv:2212.14852 , 2022.

A Pretraining Hyperparameters

Table 2 lists the default hyper-parameters we use for pretraining. All training was done on four RTX 6000s.

B Deriving Lanugage Model Scaling Laws from Statistical and Approximation Theory

B.1 Squared Regression Error

First we extract bounds on scaling laws in the case of regression squared error. We have the bound from the proof of Theorem 1

where ϵ is the approximation error such that

Let the model size N of a transformer T ∈ T be N = L T ( d 2 embd (3 m + L FFN )) = log( d )(25(3 ϵ -d β + log( ϵ -1 ))) = ˜ O ( ϵ -d β ) . Write the squared generalization error as

In the model scaling regime, when data is plentiful, we have N n << N -2 β d which implies the behavior of L sq ( N,n ) is dominated by N -2 β d . This gives us the model scaling exponent as

For the data scaling exponent α D we will choose N to balance both error terms. The will predict how data size and model size should scale together to achieve a minimal generalization error. We should have

Then

where ≍ denotes in the same order. Substituting this into the error bound gives

which gives rise to

The Bayes decision rule has the form

where 1 denotes the indicator function and

is the regression function of y on x . For binary classification, the goal is to estimate the probability function η ( x ) . If η is a β -Hölder function satisfying Assumption 2, our approximation theory in Theorem 2 gives rise to a transformer neural network η θ such that

for an arbitrary small ϵ > 0 .

We next consider the plug-in estimate of η θ [Audibert and Tsybakov, 2007]: f θ ( x ) = 1 { η θ ≥ 1 / 2 } . The excess risk of the plug-in estimate f θ is

̸

̸

̸

We next discuss how the regression error in (9) impacts the excess risk E ( f θ ) . For classification problems, the classification error depends on how well the conditional probability η ( x ) is estimated, and how many points are close to the decision boundary. Following Audibert and Tsybakov [2007], we assume the following margin condition: There exist c M > 0 and γ ≥ 0 , such that for any t > 0 ,

A smaller γ means more samples cluster around the decision boundary with a higher likelihood of falling into either class. This is not expected to be the case for natural data, where for most samples it is easy to tell to which class they belong [Kim et al., 2019, Figure 2]. As a result, we assume γ ≥ 1 .

Under the margin condition in (10), we follow Audibert and Tsybakov [2007] to decompose the excess risk E ( f θ ) such that for any δ > 0 :

̸

B.2 From Regression to Classification

After establishing statistical and approximation theory for transformers for regression, we now seek to apply our theory to classification. In language models, the next-token prediction task is a multi-class classification problem, where transformers are trained to predict the probability of the next word out of a large dictionary.

For simplicity, we consider binary classification here. Let ( x, y ) be a random couple taking values in M×{ 0 , 1 } with joint distribution P , where M is a manifold satisfying Assumption 1. Let P x be the marginal distribution of x . Here x stands for the input feature and y is the corresponding label. The classification goal is to predict the label y given the value of x . A decision rule is given by a function f : M→{ 0 , 1 } . The performance of a decision rule f is measured by the misclassification error

̸

̸

Plugging the regression error in (9) to the excess risk bound in (11) with δ = ϵ , we obtain

When the margin assumption satisfies γ = 1 , the excess risk bound becomes

which demonstrates a connection between the classification risk bound and the squared regression error for the η function in (8). This argument partially justifies the connection between empirical neural scaling laws in large language models and the squared regression error. We will leave a rigorous mathematical argument as future work.

B.3 Cross-Entropy Based Language Model

In practice, language models are trained and evaluated using cross-entropy loss. We consider the nexttoken prediction in language models. Per-sample (token), language models are trained to minimize the multi-class cross-entropy loss over a large number of classes/tokens:

where V is the vocabulary size of the model (number of classes), y ∗ i is the ground-truth label of x i , 1 y = y ∗ i is the indicator function for the event y = y ∗ i , and P ( y | x ) is the conditional probability of labels given x . For the next-token prediction, transformers are trained to estimate the probability function P ( y | x ) , and the next token is predicted to the word with the highest probability. The test loss is evaluated using cross-entropy on test data. We will leave the error bound on the cross-entropy loss as future work.

C Notation

H will represent the embedding matrix of a transformer and h i will represent the i th token (column) of the embedding matrix. h j i will denote the j th component of the vector h i . ∥·∥ p , p > 1 , will denote the p th norm of vectors, and ∥ · ∥ p,q will denote the component-wise matrix norm. For a matrix A ∈ R m × n , ∥ A ∥ p,q = ( ∑ n j =1 ( ∑ m i =1 | A ij | p ) q p ) 1 q . In particular, we denote ∥ A ∥ ∞ = ∥ A ∥ ∞ , ∞ and ∥ A ∥ 1 = ∥ A ∥ 1 , 1 for matrices. ∥ · ∥ L p ( U ) will denote the L p norm of a function on U ⊆ R D . Neural network blocks will always be bolded . Sometimes neural network blocks will also be sub-scripted with their corresponding vector of weights θ . We use ∥ θ ∥ ∞ to denote the largest magnitude of the weight parameters. However we will often omit θ and leave the dependence as implicit. In many places, we will also write θ NN to denote the weights for a neural network NN . e i will denote the i th standard basis vector where the i th component is 1 and all other entries are 0 . σ will always denote the ReLU activation. B ( x, r ) denotes the Euclidean ball of radius r centered at x ∈ R D . For a vector h t , h i : j t ∈ R j -i will denote the vector such that ( h i : j t ) k = h k t i.e. the contiguous components of h t from i, ..., j -1 .

D Details about Transformer Neural Networks

We have defined transformer neural networks in Definition 1. Some detailed definitions are given below.

Definition 3 (Transformer Block) . We define a transformer block B as a residual composition of the form

taking as input an embedding matrix H ∈ R d embd × l and MHA is a multi-headed attention layer and FFN is a feed-forward layer. To define the multi-headed attention layer we first define the attention mechanism.

Definition 4 (Attention) .

H2

FFN

+

H1.5

MHA

+

H1

where Q,K,V ∈ R d embd × d embd are referred to as the query, key, and value matrices and σ ( x ) = max(0 , x ) is the ReLU activation function. Often we will omit the dependence of A on Q,K,V . We note it will be convenient to write the action of A on the i th column h i of H as

. This allows us to interpret the i th column of the output of A as a linear combination of the values weighted by the interaction of h i th query and the h j th key. Multi-headed attention can then be defined as

where the output of each of j ∈ { 1 , ..., M } attention heads is concatenated and W O ∈ R d embd × md embd . Frequently we will simply take W O to be a sum so that

We also formally define the feed-forward layer FFN:

Definition 5 (Feed-forward Layer) . A feed-forward layer of depth L FFN and width w FFN is of the form

where W 2 , ..., W L FFN -1 ∈ R d embd × d embd , W 1 ∈ R w FFN × d embd , W L FFN ∈ R d embd × w FFN , b 1 , ..., b L FFN ∈ R d embd and the activation is ReLU. Note, each feed-forward layer is applied tokenwise to an embedding matrix H .

We can also define a class FFN of feed-forward networks:

Structured Token

C

Dynamic/Mutable

Dataterms

Static/lmmutable

Interactionterms

Constantterm

Definition 6 (Feed-forward Network Class) .

Sometimes we may omit the dependence of F on w FFN. In these cases we assume w FFN = d embd . We similarly define multi-headed attention and transformer block classes:

Definition 7 (Multi-headed Attention and Transformer Block Classes) .

MHA ( m ) = { MHA θ | MHA θ is a multi-headed attention network with weights θ with at most m attention heads. }

B ( m,L FFN , w FFN ) = { B θ | B θ is a transformer block with weights θ , a m -headed multi-headed attention layer , and a L FFN deep feed-forward layer with width w FFN } .

As with F , we may sometimes omit the dependence of B on w FFN. In these cases we take w FFN = d embd . We can parameterize the class of transformer neural networks T as Definition 2.

Lastly, our constructions will rely heavily on a particular structuring of the tokens/columns in the transformer's hidden states. The first two rows will contain mutable data used to compute the target function. The remaining rows will be immutable, containing positional/interaction vectors allowing us to uniquely identify each row and a constant term serving as a kind of 'scratchpad'. See Figure 6 for a diagram.

E Proof of Transformer Approximation Theory

We first prove the approximation theories in Section 2.4. We first prove Lemma 1 in Section E.1, and then prove Theorem 2 in Section E.2.

E.1 Proof of Lemma 1

Proof of Lemma 1. Fix f : [0 , 1] d → R and ϵ > 0 . We aim to approximate f efficiently with a transformer T ∈ T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) . We will proceed similarly to Yarotsky [2016] by approximating f locally with piecewise constant functions. First, we partition the domain uniformly into blocks indexed by n ∈ { 1 , ..., N } d for some N ≥ 1 with U n = { x ∈ [0 , 1] d : ∀ i ∈ { 1 , ..., d } , ∣ ∣ x i -n i N -1 ∣ ∣ ≤ 1 N -1 } . In this proof, we use x i and n i to denote the i th coordinate of x and n respectively. Then we construct a partition of unity ( PoU ) via

with the trapezoid function ψ : R d → R as

We can now write our approximation for f as

where f n = f ( x n ) with x n being the center point of U n . This yields the following approximation error:

Thus we can control this error simply by increasing N . It suffices to pick N = d ( 2 d H f ϵ ) 1 β +1 to obtain an ϵ error bound.

We next construct an approximation to ˆ f with a transformer neural network T. In fact we can represent ˆ f exactly as a transformer neural network.

At a high level, we will proceed by building up each ϕ n as an accumulated partial product p n in parallel over d transformer blocks. Each block will use at most d embd N d embd attention heads. Each head will be responsible for multiplying two partial products. Conveniently ϕ n can be exactly implemented with a two-layer FFN and applied column-wise to each term. The constant terms f n will then be multiplied in and summed at the final layer to compute the output.

Step 1: Embed the input Now fix input x ∈ R 1 × d . First we augment the input, via concatenation, with a sequence of 0s: 0 Nd + N d forming input x ′ ∈ R 1 × d + Nd + N d (note this is a linear operation). We can then construct a linear embedding layer [1 , 0 , 0 , 0 , 0] T = E ∈ R d embd × 1 resulting in the input embedding matrix H ∈ R d embd × ( d + Nd + N d d ) of the form

where d embd = 5 . To this we add a fixed positional encoding PE ∈ R d embd × ( d + Nd + N d d ) , producing the first hidden embedding matrix

where I i ∈ R 2 represents an interaction term determining when each token embedding will interact with another in the attention mechanism. We can set I i = (cos( i l π 2 ) , sin( i l π 2 )) where l is the number of hidden tokens. Note, this type of positional embedding is commonly known as the sinusoidal positional encoding and can be visualized as rotations of the unit vector e 1 around the first quadrant of the unit circle.

Step 2: Pre-compute ψ (3 N ( x i -j -1 N -1 )) Now we pre-compute the terms ψ (3 N ( x i -j -1 N -1 )) from which we can construct our partition of unity. We define s i,j = ψ (3( N -1)( x i -j -1 N -1 )) for 1 ≤ i ≤ d, 1 ≤ j ≤ N . We know ψ (3( N -1) · ) can be computed exactly with a two-layer FFN. So all we must do is prepare the input x i -j -1 N -1 . We will do this with one transformer block B 1 . The results will be stored in token columns h d +1 , ..., h d + Nd .

Define B 1 = B ( Nd, 7) i.e. B 1 has a multi-headed attention layer with Nd attention heads and a feed-forward layer with depth 7 . We will construct each of the ij th blocks using the Interaction Lemma 3. We may choose data kernels

so that for the token embedding h ij = h d + i ∗ d + j , we have

̸

where we can choose token h ij to only interact with token h i . Otherwise we have A ij ( h t ) = 0 for h t = h ij . Further, the weights of θ A ij of A ij are bounded such that ∥ θ A ij ∥ ∞ = O ( l 2 ) according to Lemma 3.

The added output of the multi-headed attention layer is then H 1 . 5 = MHA 1 ( H ) + H and is given by

It remains to subtract the positive term c = 1 and then apply ψ for the tokens d +1 , ..., d + dN while leaving the other tokens untouched. Define

Here W 1 shifts values in the second component of each token to the first and subtracts 1 . The second component is negated to be subtracted from H 1 . 5 . I t and the last component are preserved. A large positive number M > 2 ∥ x ∥ ∞ is added to prevent negative terms from getting erased by ReLU. W 2 removes M after the activation is applied. We can construct a two-layer network FFN 1 to apply ψ to the first component of each token h t . In tokens h d , ..., h d + Nd this produces the desired s ij terms.

Now we simply must ensure we zero-out changes to tokens outside the range d +1 , ..., d + Nd . Via Lemma 4 we can construct a two-layer feed-forward gating network FFN 2 such that FFN 2 ( h t ) = h t for t ≤ d and FFN 2 ( h t ) is zero except for the last three rows when t > d + dN . We can again invoke the same lemma to produce FFN 3 ∈ FFN (2) such that FFN 3 ( h t ) = h t when t > d and zero except for the last three rows otherwise. Applying FFN 2 and FFN 3 zeroes out tokens outside d +1 , ..., d + dN while preserving the rest. Finally, we define the last layer

which zeros out all except the first two components. This produces the next embedding matrix H

.

Step 3: Building up partial products Now we can start construction of the partial products for each patch ϕ n , n ∈ { 1 , ...., N } d . For a fixed patch n , we write p n,k,i to indicate the i th partial product of 2 k -1 terms. Formally, p n,k,i = p n,k -1 , 2 i -1 · p n,k -1 , 2 i with p n, 1 ,i = s i,n i where s i,n i = ψ (3( N -1)( x i -n i -1 N -1 )) and n ∈ { 1 , ..., N } d , 1 ≤ k ≤ log 2 ( d ) , 1 ≤ i ≤ d 2 k . See Figure 7 for a diagram of how each partial product is assembled. Architecturally, this will be done in log( d ) + 1 transformer blocks. The first block will copy the product terms into the last N d d tokens. The remaining log( d ) blocks compute the recursively assembled partial products by multiplying p n,k -1 , 2 i -1 · p n,k -1 , 2 i . Storage and assembly will be done in the last d + Nd +1 , ..., d + Nd + N d d tokens.

For the first transformer block we define B 2 = B ( N d d, 0) . We index each attention head as A n,i where n ∈ { 1 , ...., N } d and 1 ≤ i ≤ d . Fix a patch n and the corresponding product ∏ d i =1 s i,n i . These terms have been pre-computed and stored in tokens h d +1 , ..., h d + Nd . We will then use each attention head A n,i to copy s i,n i into its corresponding token h n,i = h d + Nd + d ∑ d p =1 ( n p -1) N d -p + i . Concretely, for a fixed attention head A n,i define the data kernels

2 :

(pn,3,1

(Pn,2,1

Pn,2,2

Pn,1,1

(Pn,1,2)

Pn,1,3

Pn,1,4)

S1,n1

S2,n2

S3,n3

S4,n4

via the Interaction Lemma 3 we may pick A n,i so that

̸

where the interaction terms zero-out all terms in the sum except when k = d + id + n i (which is the index of the token containing s i,n i ). Similarly to as in B 1 we have A n,i h k = 0 when k = d + Nd + ∑ d p =1 ( n p -1) N d -p + i . Set the feed-forward layer FFN = 0 . Then the next token embedding matrix can be written as

where P 1 ∈ R N d d is

where we recall p { 1 } d , 1 , 1 denotes the first term in the base level product for for the patch n = { 1 , ..., 1 } . Now define the next block B 3 ∈ B ( N d d 2 , 1) . Index the corresponding attention heads as A n,i for n ∈ { 1 , ..., N } d , 1 ≤ i ≤ d 2 . Define data kernels

where we choose the interaction terms so that the token containing p n, 1 , 2 i -1 only interacts with the token containing p n, 1 , 2 i so that they multiply to form p n, 2 ,i as shown in Figure 7. Note, in addition to computing the product, we must also subtract p n, 1 , 2 i -1 from each corresponding token so that the output of the residual from B 3 cancels with the existing terms in the embedding matrix H 3 . Write

Adding the result onto H 3 gives the intermediate embedding matrix

where P 1 . 5 ∈ R N d d is

Note: ∗ denotes a non-essential entry which can take any value. We then design our FFN network as W 1 x where

which when added to H 3 . 5 swaps the first component for the second component of each token and zeros out the second component. Then the next embedding matrix H 4 gives

where P 2 = P 1 . 5 ∈ R N d d . More generally define the k th block implementing the k th level of the parallel product as B k +2 = B ( N d d 2 k , FFN (1)) with attention heads A n,k,i where n ∈ { 1 , ..., N } d , 1 ≤ i ≤ d 2 k . We can construct each attention head A n,k,i so that the token containing p n,k -1 , 2 i -1 only interacts with the token containing p n,k -1 , 2 i to produce p n,k,i :

In the FFN set W 1 as before, yielding

where component P d + Nd + ∑ d p =1 ( n p -1) N p + i log( d )+2 = p n, log( d ) , 1 = ϕ n ( x ) as desired in (18). It simply now remains to multiply each ϕ n ( x ) by the corresponding local average f n and sum the result. This can be implemented with a single attention block B log( d )+3 ∈ B ( N d , 1) . For n ∈ { 1 , ..., N } d define A n with the data kernels

where

At k = log( d ) we have

where ∥ f ∥ L ∞ ([0 , 1] d ) ≤ R and the token h d + Nd + ∑ d p =1 ( n p -1) N p +1 = h n, 1 only interacts with itself so that

where the last equality comes from | f n ϕ n ( x ) | ≤ | f n | ≤ R . We implement FFN to subtract R via

where P n, 1 log( d )+4 = f n ϕ n for n ∈ { 1 , ..., N } d and is 0 elsewhere. To sum up the resulting terms we define the final transformer block B log( d )+4 ∈ B (1 , 0) containing a single attention head A 1 as

which when applied to each token simply sums up the second components of all tokens. The result when applied to h t where t = l gives

We can subtract lR with the FFN as

yielding the final result h 2 1 = ∑ n f n ϕ n . Applying the decoding head D returns the desired result. This shows we can construct the approximation ˆ f exactly with an O (log( d )) layer transformer.

E.2 Proof of Theorem 2

With Lemma 1, we can now move on to the proof of Theorem 2.

Proof of Theorem 2. Our goal is to approximate f : M → R in L ∞ ( M ) with a transformer neural network. Our construction proceeds similarly to the construction in Chen et al. [2019] with feedforward neural networks.

so that

First, we will decompose f ( x ) as a sum of terms over local neighborhoods U 1 , ..., U C M ⊆ M covering M . Approximations on overlapping neighborhoods containing x will then be combined via a partition of unity ( PoU ) { ρ n } C M n =1 which subordinates { U n } C M n =1 , while contributions from neighborhoods not containing x will be zeroed-out via an indicator 1 U n . This will give us the expression

where f n : M→ R agrees with f on U n . Next, on each local neighborhood we will project the input x ∈ M ⊆ R D to [0 , 1] d . This will give us the following local decomposition of the target function:

where ˜ f n : R d → R and ϕ n : M → R d is a projection. This step is critical: we now must only approximate a low-dimensional function ˜ f n and a simple projection ϕ n instead of the highdimensional function f n . After this, we must simply approximate each ˜ f n , ϕ n , 1 U n using transformer neural networks.

Step 1: Local Decomposition Denote the open Euclidean ball with center c and radius r in R D by B ( c, r ) . Given a manifold M satisfying Assumption 1, the collection of open balls { B ( c, r ) } c ∈M with some r ≤ τ/ 4 is an open cover of M . Since M is compact, there are a finite number of points, denoted by c n , for n = 1 , . . . , C M such that M ⊂ ∪ C M n =1 B ( c n , r ) . We can then decompose M into C M charts ( U n , ϕ n ) C M n =1 where U n = M∩ B ( c n , r ) and each ϕ n : U n → R d is an orthogonal projection of the local neighborhood U n onto the local tangent space T c n ( M ) . Via Lemma 8, as long as r ≤ τ/ 4 , each ϕ n is a diffeomorphism between U n and a subset in T c n ( M ) . Here the number of charts C M is a constant depending on the complexity of the manifold M . We can bound it as C M ≲ SA ( M ) r d d log( d ) where SA ( M ) is the surface area of M [Conway and Sloane, 1988].

For the open covering { U n } C M n =1 , there exists a smooth partition of unity { ρ n } C M n =1 such that f ( x ) = ∑ C M n =1 ρ n ( x ) f ( x ) for x ∈ M [Tu, 2011]. Each ρ n is supported only in its corresponding chart U n . We additionally include an indicator function 1 U n for each neighborhood U n , resulting the decomposition

We then aim to approximate each term in the sum on its corresponding local neighborhood U n and then sum the result. Because we are only locally approximating each ρ n , the indicator is critical in the approximation as it will be used to suppress contributions from U n such that the input x / ∈ U n .

Fix a local neighborhood U n . We will now project the input x ∈ M ⊆ R D to the local tangent space T c n ( M ) ⊆ [0 , 1] d . We can write

where ˜ f n : R d → R and again ϕ n : M⊆ R D → R d is an orthogonal projection. Then we can write f as

Step 2: Local Approximation In (20), we have rewritten f as a sum of functions which can be approximated locally on each U n and indicators 1 U n . We will approximate each component with transformer neural networks. Specifically, let T ˜ f n approximate ˜ f n , T ϕ n approximate ϕ n , and T 1 Un approximate 1 U n , 1 ≤ n ≤ C M . Then we will have the overall approximation

This yields the error decomposition

where the error terms E n, 1 , E n, 2 , E n, 3 denote

respectively. It remains to construct transformer approximations for each term.

Step 2.1: Approximating ˜ f n and Bounding E n, 1 We first handle the E n, 1 error in (21). Fix 1 ≤ n ≤ C M . We can write

so it suffices to approximate ˜ f n : [0 , 1] d → R . Now choose the desired accuracy δ n, 1 > 0 . Via Lemma 1 we can construct an approximation such that

This construction T ˜ f n has O (log( d )) transformer block layers, O (1) feed-forward layers, O ( dδ -d β n, 1 ) tokens, O (1) embedding dimension, O ( dδ -d β n, 1 ) attention heads per-layer, and weights with magnitude at most O ( dδ -2 d β n, 1 ) . Via the Parallelization Lemma 7 we can compute each approximation of ˜ f n in parallel using a transformer with O ( C M dδ -d β 1 ) attention heads and O ( C M d embd ) FFN width, where we set δ 1 = δ 1 , 1 = ... = δ C M , 1 , and O ( log( d ) ) layers.

Step 2.2: Approximating ϕ n Exactly Such That E n, 2 = 0 Fix 1 ≤ n ≤ C M . We must now control the E n, 2 error in (22):

We will implement ϕ n using a transformer block B ∈ B ( dD, 5) . We can identify the tangent space at the point c n as

where the vectors v n,i ∈ R D , 1 ≤ i ≤ d , form an orthonormal basis of T c n ( M ) . Then we can write the projection map ϕ n as

where s n ∈ (0 , 1] is a scaling factor, u n ∈ R d is a translation, and V n = [ v n, 1 ... v n,d ] ∈ R D × d so that ϕ n ( x ) ∈ [0 , 1] d . Suppose we have an input hidden embedding matrix of the form

We will implement the operation v j n,i · ( x j -c j n ) via an attention head A i,j , 1 ≤ i ≤ d , 1 ≤ j ≤ D . Define the data kernels

via the Interaction Lemma 3 we can define A i,j so that token h D + i interacts with token h j so that

and is 0 on other tokens h t = h D + i , where M bounds x ∈ M is such that v j n,i ( x j -c j n ) + 2 M ≥ 0 for all n, i, j . Then the output of the multi-headed attention MHA on h D + i is

̸

This yields the intermediate embedding matrix

We can then design a five layer FFN layer to subtract 2 DM and then add u n and multiply s n only from the embedding tokens D +1 , ..., D + d . The first layer will simply subsract 2 DM and then add u n and multiply s n from each token. This can be implemented by the matrix

Then, via the gating Lemma 4 we can construct a two-layer feed-forward network FFN gating such that, for a chosen 1 ≤ t u ≤ l , we have FFN gating ( h k ) = h k when k < t u and FFN gating ( h k ) is zero except for the last three rows for k ≥ t u . We can again invoke the lemma to construct another FFN layer so that for k ≥ t l = ⇒ FFN gating ( h k ) = h k and k < t l = ⇒ FFN gating ( h k ) is zero except for the last three rows. So, applying Lemma 4 twice, we can construct a four-layer feed-forward network FFN gating such that FFN gating ( h t ) = h t for D +1 ≤ D + d and is zero except for the last three rows otherwise. This yields the desired output.

Further, this approximation is exact , so that E n, 2 = 0 . We can then compute each ϕ n in parallel via the Parallelization Lemma 7 with a transformer having O ( C M dD ) attention heads and O ( C M d embd ) FFN width.

Step 2.3: Approximating 1 U n ( x ) and Bounding E n, 3 Again fix 1 ≤ n ≤ C M . By construction we can write

Suppose we are given an input embedding matrix of the form

We must start by computing ∥ x -c n ∥ 2 2 . Via the Addition Lemma 5 we can construct a transformer block B 1 ∈ B ( D, 3) such that

We now construct transfomer block B 2 ∈ B ( D, 2) squares each term in the sum. For 1 ≤ i ≤ D define the data kernels

Then via the Interaction Lemma 3 we can construct A i so that the token h D + i interacts with itself such that

and 0 otherwise. The output of the resulting multi-headed attention gives

Via Lemmas 6 and 4 we can construct FFN ∈ FFN (5) which replaces the first row with the second row for h D +1 , ..., h 2 D . This gives the output

and it remains only to take the sum. This can be done with B 3 ∈ B ( D -1 , 0) . For 1 ≤ i ≤ D -1 define the data kernels

via the Interaction Lemma 3 we construct A i so that h D +1 interacts only with h D +1+ i such that

The output of the multi-headed attention is then

and 0 otherwise. This gives the intermediate output

where in particular h 1 D +1 = ∥ x -c n ∥ 2 2 . It remains to approximate the indicator

2.3a: Approximating 1 r 2 ( s ) Recall our goal is to approximate the function f : M→ R in L ∞ . However, we cannot approximate 1 r 2 in L ∞ as it is a discontinuous function. To address this issue, recall the error term corresponding to the approximation of 1 U n is E n, 3 = ∥ ˜ f n ◦ ϕ n 1 U n -˜ f n ◦ ϕ n T 1 Un ∥ L ∞ ( M ) . We know on the boundary of U n that ˜ f n ◦ ϕ n = f n = fρ n must be 0 as ρ n is 0 and further fρ n is β -Hölder continuous as f is β -Hölder continuous and ρ n is smooth. This suggests the following approach. Define the approximating indicator

for some ∆ > 0 . Clearly for s ≤ r 2 -∆ and s ≥ r 2 we have ˆ 1 r 2 , ∆ ( s ) = 1 r 2 ( s ) . So we argue ˜ f n is small for x ∈ M in the range K n = { x ∈ M : r 2 -∆ < ∥ x -c n ∥ 2 2 < r 2 } . This is handled by Lemma 8 in Chen et al. [2022] which gives us the following bound on E n, 3 :

for some constant c . So it suffices to make ∆ small. It then remains to implement ˆ 1 r 2 , ∆ . We will do so using a feed-forward layer FFN ∈ FFN ( O (log( 1 ∆ ))) . Note that we can write the target exactly as

while is realized by a two-layer ReLU network. However, the weight 1 ∆ will scale unboundedly with the desired accuracy ϵ . So we must deepen the FFN to utilize a logarithmic number of layers O (log( 1 ∆ )) as

which prevents the weights ( 1 ∆ ) 1 log( 1 ∆ ) from blowing up. We can then approximate each 1 U n in parallel using a transformer with O ( C M D ) and three layers.

Step 3: Bringing it all together Recall our goal is to approximate f on M by writing f as

We argued we can successfully approximate

˜ f n up to δ n, 1 accuracy in L ∞ via a transformer neural network with O ( dδ -d β n, 1 ) width and O (log( d )) depth. All approximations can be computed in parallel via a transformer with O ( C M dδ -d β n, 1 ) attention heads and O ( log( d ) ) layers.

Each ϕ n exactly via a transformer neural network with O ( dD ) width and constant depth. All the ϕ n can be computed in parallel via a transformer with O ( C M dD ) heads and constant depth.

Each 1 U n up to δ n, 3 accuracy in L ∞ via a transformer with O ( D ) width and a feed-forward layer with O ( log( 1 ∆ ) ) depth. All 1 U n can be approximated in parallel via a transformer with O ( C M D ) attention heads and a constant number of transformer blocks.

Via the Parallelization Lemma 7 we can compute the approximations of 1 U n and ˜ f n ◦ ϕ n in parallel via a transformer with O ( C M ( dδ -d n, 1 + dD ) ) attention heads, O ( C M d embd ) FFN width, and O ( log( d ) )

transformer blocks. As a result will have the embedding matrix

where each c f n = ˜ f n ◦ ϕ n ( x ) and c U n = 1 U n ( x ) . We implement the final sum via two transformer blocks B 1 , B 2 . First define B 1 ∈ B ( C M , 1) with attention heads A i , 1 ≤ i ≤ C M with data kernels

with interaction kernels chosen so that h D + i only interacts with h D + C M + i under A i . Then

We can then design a two-layer FFN to subtract M from columns D +1 , ..., D + C M yielding the output embedding matrix

We can sum up the result with another transformer block B 2 ∈ B ( C M , 2) and storing the desired result in the first component of the last token t = l . Applying the decoding layer D yields the desired result.

It simply remains to control the errors. Recall we have the error bound

Then for a target accuracy ϵ we choose δ n, 1 = ϵ 2 C M and δ n, 3 = ϵ 2 C M so that

This completes the proof of Theorem 2.

The above result concludes our approximation theory results. To summarize, we first showed a (sufficiently smooth) function f : [0 , 1] d → R can be approximated up to ϵ accuracy in L ∞ with a transformer neural network with O ( dϵ -d ) width and O (log( d )) depth in Lemma 1. We then used this construction to show a function f : M→ R on a compact manifold M with intrinsic dimension d can be approximated up to ϵ accuracy in L 1 with a transformer neural network with width O (max( dD,dϵ -d )) and depth O (log( d )) .

F Proof of Transformer Generalization Theory

We next prove Theorem 1 in Appendix F.1. Much of this follows the same argument as in Chen et al. [2022] which focused on feedforward neural networks. The main novelty of this paper is a bound on the covering number of T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) given in Lemma 2, which is proved in Appendix F.2.

F.1 Proof of Theorem 1

Proof of Theorem 1. Suppose we have x 1 , ..., x n ∼ Q as i.i.d. training samples drawn from Q on the manifold M and their evaluations f ( x 1 ) , ..., f ( x n ) . Set the approximating function class as T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) . We define the transformer empirical risk minimizer ˆ T n in (1) as the transformer network which minimizes the empirical L 2 training loss. Our goal is to bound the squared generalization error in (6).

We first rewrite the squared generalization error by adding and subtracting (twice) the training objective where we note a factor of two is used to accelerate convergence of the statistical error.

The first (bias) term in (24) can be controlled via Theorem 2:

where the last line follows from Theorem 2. Note, we can pass the expectation inside the infimum via Jensen's inequality. Set d 2 T ( x ) = ( T ( x ) -f ( x )) 2 .

It remains to control the last two (variance) terms in (24). We will do so by controlling the covering number of the approximating function class T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) . Via Lemma 6 from Chen et al. [2022] we have the bound

where N ( δ 4 R , T , ∥ · ∥ ∞ ) denotes of the covering number of the network class T under the L ∞ norm. Our goal is now to bound this covering number of T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) . We do so with the following lemma:

Lemma 2. Consider a transformer neural network class T = T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) with input x ∈ R D satisfying ∥ x ∥ ∞ ≤ M .. Let δ > 0 . Then

With the covering number of T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) in hand we may now apply (25) to obtain

.

where we set δ = 1 n . This yields

Recall for target accuracy ϵ > 0 we can choose L T = log( d ) , L FFN = log( ϵ -1 ) , w FFN ≤ 4 C M , d embd = 5 , l ≤ C M dϵ -d β , m ≤ C M dϵ -d β , κ ≤ ( MC M dϵ -d β ) 2 , l ≤ C M dϵ -d β . This simplifies the above to

where ˜ O hides log terms and constants. Then we can bound the variance term as

Putting together the bounds on the bias and variance yields a bound on the empirical risk:

and it simply remains to pick ϵ to balance the error terms. We do this by choosing ϵ such that ϵ 2 = ϵ -d β n yielding ϵ = n -β 2 β + d . This gives our final bound on the expected empirical risk:

as desired.

F.2 Proof of Lemma 2 about the Transformer Covering Number

Proof of Lemma 2. The key is to bound the difference in ∥ · ∥ ∞ between two transformers T , T ′ ∈ T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) . Set η > 0 and choose T , T ′ so that ∥ θ -θ ′ ∥ ∞ < η . In other words, the weight parameters in T and T ′ differ at most by η . Fix x ∈ [0 , 1] D . Recall the decoder D : R d embd × l → R is fixed to output the first element of the first row. We compute

To handle this term, we write

We will first bound the second term.

Consider two multi-headed attention layers MHA 1 , MHA 2 ∈ MHA ( m ) with attention heads A i j , i ∈ { 1 , 2 } , 1 ≤ j ≤ m . For H ∈ [0 , M ] d embd × l we compute

Consider two attention heads A 1 , A 2 with weight parameters Q i , K i , V i , i ∈ { 1 , 2 } . For H ∈ [ -M,M ] d embd × l , we compute

where the last inequality comes from adding and subtracting the term ⟨ Q 2 h j , K 2 h k ⟩ V 1 h k . We bound the first term via

where we note ∥ Ax ∥ 2 ≤ ∥ Ax ∥ 1 ≤ ∥ A ∥ 1 ∥ x ∥ ∞ for A ∈ R d embd × d embd , x ∈ R d embd and ∥ A ∥ 1 ≤ d 2 embd ∥ A ∥ ∞ . This bounds the first term.

The second term can be bounded as

Thus we can bound the attention difference ∥ A 1 ( H ) -A 2 ( H ) ∥ ∞ as

We can also now bound the multi-headed difference ∥ MHA 1 ( H ) -MHA 2 ( H ) ∥ ∞ as

A bound on the first residual layer in the transformer block immediately follows via

To bound the difference in the output of the FFN layer in the transformer block we have for f 1 , f 2 ∈ FFN ( L FFN ) and input H ∈ [0 , M ′ ] d embd × l we have

which also bounds the second residual layer of the transformer block. Then to bound the difference of two transformer blocks B 1 , B 2 ∈ B ( m,L FFN ) write, for H ∈ [0 , M ] d embd × l ,

where we can bound the first term via analysis above. For the second term write

To handle the first term we must bound he Lipschitz constant of FFN 1 . Let H,H ′ ∈ R d embd × l . Write

where we again use that ∥ Wx ∥ ∞ ≤ ∥ W ∥ 1 ∥ x ∥ ∞ and that σ ( · ) is 1-Lipschitz. Applying this to the first term from above we get

We can bound ∥ FFN 2 ( H + MHA 2 ( H )) -FFN 1 ( H + MHA 2 ( H )) ∥ ∞ as

where ∥ H + MHA 2 ( H ) ∥ ∞ ≤ M ′ . Putting all the estimates together, we have

∥

B

1

(

H

)

-

B

2

(

H

)

∥

∞

and it remains to control M ′ in terms of M . We know ∥ H + MHA ( H ) ∥ ∞ ≤ M ′ and so we compute

Thus, we can complete our estimate on ∥ B 1 ( H ) -B 2 ( H ) ∥ ∞ as

Now consider the multi-block difference

To bound the second difference we must bound the output ∥ B ( H ) ∥ ∞ , B ∈ B ( m,L FFN ) . Suppose ∥ H ∥ ∞ ≤ M . Write

so we must bound the output of the FFN layer. We have for H ′ ∈ R d embd × l with ∥ H ′ ∥ ∞ ≤ M , we have

Since ∥ H + MHA ( H ) ∥ ∞ ≤ 2 d 6 embd κ 3 mlM , we have

Iterating this gives

Then we can bound the above second term as

To bound the first term ∥ B L T ◦ ... ◦ B 1 ( H ) -B L T ◦ ... ◦ B ′ 1 ( H ) ∥ ∞ we must control the Lipschitz constant of of B ∈ B ( m,L FFN ) . Let H,H ′ ∈ R d embd × l with ∥ H ∥ ∞ , ∥ H ′ ∥ ∞ ≤ M . Compute

and it suffices to control the Lipschitz constant of MHA ∈ MHA ( m ) . For the same H,H ′ compute

We can bound ∣ ∣ ⟨ Q j h i , K j h k ⟩ V j h k -⟨ Q j h ′ i , K j h ′ k ⟩ V j h ′ k ∣ ∣ as

∥⟨

≤∥⟨

Q

Q

∣

⟨

∣

Q

∣

∣

⟨

∣

∣

⟨

4

j

j

j

Q

Q

h

h

h

j

j

i

i

i

h

j

, K

, K

j

, K

i

h

i

j

, K

h

h

h

j

, K

k

k

k

h

j

V

h

⟩

⟩

j

h

V

⟩ - ⟨

k

j

6

h

κ

3

M

k

2

k

k

Q

⟩ - ⟨

⟩ - ⟨

∥

-⟨

-⟨

Q

j

Q

j

j

Q

Q

h

k

-

h

h

, K

h

j

h

j

h

h

′

i

′

i

′

i

′

i

′

i

, K

j

, K

j

, K

h

j

, K

′

k

j

∥

∞

.

h

j

h

′

∣

k

⟩

∣

′

k

′

k

∥

h

h

′

k

′

k

⟩

⟩

∣

∣

∣

∣

V

d

j

⟩

V

⟩

V

h

j

j

h

h

′

k

k

k

2

∥

∞

2

d

d

∥

∞

∥

∞

+

+

∣

∣

∥⟨

⟨

Q

+

∣

∣

⟨

+2

Q

j

h

Q

d

j

′

i

j

h

′

i

, K

, K

h

′

i

j

′

k

j

h

V

⟩

h

, K

4

′

k

j

⟩ - ⟨

j

h

κ

2

M

′

k

2

Q

⟩ - ⟨

d

h

j

′

k

h

Q

2

j

κ

′

i

h

∥

-⟨

Q

, K

j

′

i

h

j

h

, K

j

k

-

′

k

h

h

h

′

i

⟩

∣

′

k

′

k

∣

, K

∥

⟩

j

V

j

∣

∣

∥

d

∞

h

h

′

k

′

k

V

⟩

-

2

j

h

V

κ

∥

k

j

∥

h

h

k

∞

k

∥

-

∞

h

′

k

∥

∞

Applying this gives

Then

With this we can bound the first term of the multi-block difference

. Putting together the estimates on both terms gives

which reduces the bound on the L T layer transformer class to the bound on the bound on the L T -1 layer transformer class. Finally, via induction, we get

It simply remains to deal with the encoding layer H = PE + E ( x ) . Both E and PE are fixed among architectures and further ∥ PE + E ( x ) ∥ ∞ = ∥ x ∥ ∞ +1 ≤ M +1 . Thus

This allows us to complete the total bound on the transformer difference of T , T ′ ∈ T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) such that ∥ θ -θ ′ ∥ ∞ < η as

Now, we may compute N ( T , δ, ∥· ∥ ∞ ) . We must cover T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) with a δ -net such that, for all T ∈ T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) , we can find T ′ in the net such that ∥ T -T ′ ∥ ∞ < δ . Set η = δ C T . We construct the δ -net by uniformly discretizing the set of weights θ corresponding to T θ ∈ T with step size given by η . Then for any T ∈ T , we can find a T ′ in the grid such that ∥ θ T -θ T ′ ∥ ∞ < η . Then we have ∥ T -T ′ ∥ ∞ < C T η = C T δ C T = δ . We can compute the number of parameters in θ as

We can compute the number of steps per parameter as 2 κ η . Thus the covering number of T ( L T , L FFN , w FFN , l, d embd , m, R, κ ) is bounded by

G Building Blocks of Transformer Neural Networks

Lemma 3 (Interaction Lemma) . Set d embd = 5 and let κ, M > 0 . Fix l ∈ N , 1 ≤ t 1 , t 2 ≤ l and 1 ≤ i ≤ d embd . Let H ∈ R d embd × l be a structured transformer embedding matrix such that h d embd -2: d embd -1 t = I t and h d embd t = 1 . Further suppose ∥ H ∥ ∞ ≤ M for some M > 0 Let B : R d embd × R d embd → R be a kernel function on tokens which can be written in the form B ( h, h ′ ) = σ ( ⟨ Q B h, K B h ′ ⟩ ) for some h, h ′ ∈ R d embd , Q B , K B ∈ R ( d embd -3) × d embd with ∥ Q B ∥ ∞ , ∞ , ∥ K B ∥ ∞ , ∞ ≤ κ . Then we can construct an attention head A acting on H such that A ( h t 1 ) = B ( h t 1 , h t 2 ) e i and otherwise A ( h t ) = 0 for t = t 1 . Further we have ∥ θ A ∥ ∞ = O ( d 4 embd κ 2 M 2 l 2 ) . Note: we call the matrices Q B , K B 'data kernels'.

Proof of Lemma 3. Define the query, key, and value matrices as

where we call Q I , K I ∈ R 2 × 2 the interaction kernels , Q B , K B ∈ R d embd -3 × d embd the data kernels , and C > 0 simply a large positive number. We can choose Q I , K I so K I = P I t 2 is a projection onto the unit interaction vector I t 2 (the interaction term for h t 2 ) and Q I is a dilation and rotation of I t 1 onto I t 2 i.e. Q I I t 1 = C I t 2 . We must now compute A ( H ) . For an arbitrary 1 ≤ t ≤ l we can compute the action of A on h t as

Now we must case on whether the input token t = t 1 .

Case 1: t = t 1 Write

To handle the same we must again go by casework, this by casing on whether k = t 2 .

Case 1a): k = t 2 When k = t 2 we have ⟨ Q I I t 1 , K I I t 2 ⟩ = ⟨ C I t 2 , I t 2 ⟩ since by construction Q I I t 1 = I t 2 and K I is a projection onto I t 2 so that K I I t 2 = I t 2 . This further simplifies as ⟨ C I t 2 , I t 2 ⟩ = C since I t 2 is unit length. Thus

̸

, .

̸

Now we must consider the other terms in the sum when k = t 2 .

̸

̸

Case 1b): k = t 2 When k = t 2 we compute

where we use Cauchy-Schwarz at the first inequality and note that ∥ P I t 2 I k ∥ 2 < 1 since k = t 2 for the second inequality. Then, for large enough C , we have

̸

so we simply must choose C so that ⟨ Q B h t 1 , K B h k ⟩ + C ∥ P I t I k ∥ 2 -C < 0 . Compute

2 2

We can bound

̸

so it remains to control 1 1 -∥ P I t 2 I k ∥ 2 when k = t 2 by upper bounding ∥ P I t 2 I k ∥ 2 . Compute

̸

where we note the projection is maximized by any adjacent interaction terms I t , I t +1 . Then

which gives a bound on C . So for large enough C we can conclude

It remains to consider the case t = t 1 .

̸

̸

Case 2: t = t 1 Now we consider A applied to tokens other than h t 1 . We have

so we must argue each term in the sum is 0 by showing ⟨ Q B h t , K B h k ⟩ + ⟨ Q I I t , K I I k ⟩ -C < 0 . We again split into two cases.

̸

Case 2a): k = t 2 When k = t 2 we have ⟨ Q B h t , K B h k ⟩ + ⟨ Q I I t , K I I k ⟩ -C < 0 via the same argument as in case 1b .

̸

Case 2b): k = t 2 When k = t 2 we must instead more carefully bound the inner product. We can simplify the interaction terms as

̸

where θ t,t 2 is the angle between Q I I t and K I I k . Crucially, because t = t 1 , Q I I t = C I t 2 and so θ t,t 2 > 0 and cos( θ t,t 2 ) < 1 . Then we can say

̸

We upper bound the numerator by O ( d 4 embd κ 2 M 2 l 2 ) . Note that the Q I rotates by clockwise t 1 -t 2 l π 2 radians. So Q I I t = C (cos( ( t + t 1 -t 2 ) π 2 l ) , sin( ( t + t 1 -t 2 ) π 2 l )) . Thus we can upper bound cos( θ t,t 2 ) as

So as in case 1b we have C = O ( d 4 embd κ 2 M 2 l 2 ) . For sufficiently large C we have A ( h t ) = 0 as desired. This completes the proof of Lemma 3.

Lemma 4 (Gating FFNs) . Given an embedding matrix H ∈ R d embd × l which has all ones in the last row and whose tokens contain interaction terms and given a half-space of contiguous tokens h k , ..., h l or h 1 , ..., h k , we can design FFN ∈ FFN (2) such that

where    0 I 1 t I 2 t 1    is zero except for the last three coordinates. We call this FFN a gating feed-forward network . We additionally have ∥ θ FFN ∥ ∞ ≤ O ( l ∥ H ∥ ∞ ) .

Proof of Lemma 4. Without loss of generality assume we are given a contiguous prefix of tokens h 1 , ..., h k , k ≤ l . We can choose a vector v ∈ S 1 such that v · I t > 0 for t ∈ { 1 ..., k } and v · I t < 0 for t > k . We call such v the pivot vector . In particular we choose R π 2 ( I k + I k +1 ∥I k + I k +1 ∥ 2 ) where R π 2 is a quarter-circle rotation clockwise. For a large C > 0 , construct

where I d embd is the identity matrix of size d embd × d embd . Then

which is    0 I 1 t I 2 t 1    if I t · v < 0 . Applying the second layer to    0 I 1 t I 2 t 1    yields    0 I 1 t I 2 t 1    . Otherwise assume

I t · v > 0 . Then σ ( h i t + C I t · v ) = h i t + C I t · v . Applying the second layer yields

as desired. Further we can compute a bound on C > 0 as | CI t · v | > ∥ H ∥ ∞ ⇐⇒ C > ∥ H ∥ ∞ | I t · v | . We compute

where the last inequality holds for l ≥ 2 . Thus ∥ H ∥ ∞ | I t · v | ≤ 2 3 l ∥ H ∥ ∞ π < C .

Lemma 5 (Constant Addition) . Let M > 0 and c be a vector in R D such that ∥ c ∥ ∞ ≤ M 2 . Let H be an embedding matrix of the form

where l = 2 D such that ∥ H ∥ ∞ , ∞ ≤ M 2 . Then there exists a transformer block B ∈ B ( D, 3) such

with ∥ θ B ∥ ∞ = O ( lM ) . In this case we say B implements the addition of c to x .

Proof of Lemma 5. We begin by defining the attention heads A i , 1 ≤ i ≤ D . For each head A i we define the data kernels

Then, via the Interaction Lemma 3, we can construct A i so that token h D + i interacts with h i such that

as desired.

Lemma 6 (Replacing FFN) . Set d embd = 5 and l ∈ N . Let H ∈ R d embd × l be a structured intermediate embedding matrix. We can construct a feed-forward network FFN replace ∈ FFN (1) with the following output:

In this case we say FFN replace replaces the the first row of H with the second row.

Proof of Lemma 6. Set

̸

and A i ( h t ) = 0 when t = D + i . Then the residual multi-headed attention yields

Via the gating lemma 4 we can design a three layer FFN to subtract M from only h D +1 , ..., h 2 D . Thus

Applying W token-wise gives the desired result.

Lemma 7 (Transformer Parallelization) . Let m 1 , m 2 , l 1 , l 2 , L FFN 1 , L FFN 2 , w FFN 1 , w FFN 2 , d embd ∈ N . Fix transformer blocks B ∈ B ( m 1 , L FFN 1 , w FFN 1 ) with input H 1 ∈ R d embd × l 1 and B 2 ∈ B ( m 2 , L FFN 2 , w FFN 2 ) with input H 2 ∈ R d embd × l 2 . Further suppose both inputs are structured such that each token h i,t has an interaction term I t and constant term 1 . Further suppose all attention heads in both B 1 and B 2 are implemented using the Interaction Lemma and all FFN layers are either Gating FFN layers or are independent of the interaction terms. Then there exists B 3 ∈ B ( m 1 + m 2 , max( L FFN 1 , L FFN 2 ) + 2 , w FFN 1 + w FFN 2 ) which takes as input H 3 ∈ R d embd × ( l 1 + l 2 ) such that

where H ′ i ∈ R d embd × l i is the same as H i except for the interaction terms where we have I ′ 3 ,t = ( cos ( t l 1 + l 2 π 2 ) , sin ( t l 1 + l 2 π 2 )) for 1 ≤ t ≤ l 1 + l 2 . If B 3 satisfies this relationship we say B 3 parallelizes B 1 and B 2 .

Proof of Lemma 7. Let B 1 ( H 1 ) = FFN 1 ( MHA 1 ( H 1 )) + MHA 1 ( H 1 ) + H 1 and B 2 ( H 2 ) = FFN 2 ( MHA 1 ( H 2 )) + MHA 2 ( H 2 ) + H 2 . Let A 1 ,i 1 , 1 ≤ i 1 ≤ m 1 , denote the attention heads of MHA 2 and A 2 ,i 2 , 1 ≤ i 2 ≤ m 2 , denote the attention heads of MHA 2 . We construct B 3 as follows.

First we construct the new input embedding matrix H 3 from H 1 and H 2 . For 1 ≤ t ≤ l 1 set h 3 ,t = h 1 ,t except we define a new interaction term I 3 ,t = ( cos ( t l 1 + l 2 π 2 ) , sin ( t l 1 + l 2 π 2 )) . For

l 1 +1 ≤ t ≤ l 1 + l 2 define h 3 ,t = h 3 ,t -l 1 except where we again must change only the interaction term. The main difficulty in defining MHA 3 will then simply be updating the interaction kernels of the Interaction Lemma structured attention heads to respect the new interaction terms I t , 1 ≤ t ≤ l 1 + l 2 .

We now define the attention heads A 3 ,j in MHA 3 . When 1 ≤ j ≤ m 1 write Q 3 ,j = Q ′ 1 ,j where Q ′ 1 ,j is the same Q 1 ,j except for a new interaction kernel Q ′ I 1 ,j . Let h 1 ,s j , s j ∈ { 1 , ..., l 1 } , be the source token in H 1 interacting with Q I 1 ,j and h 1 ,t j , t j ∈ { 1 , ..., l 1 } , be the target token of the old interaction kernel Q I 1 ,j (recall Q I 1 ,j is defined as the rotation of I 1 ,s j onto I 1 ,t j ). Then define Q ′ I 1 ,j as a rotation onto the new target interaction kernel I 3 ,t j . Similarly define K 3 ,j = K ′ 1 ,j where K ′ 1 ,j = K 1 ,j up to a difference in interaction kernels. Let K I 1 ,j be the old interaction kernel which is a projection on to the target interaction term I 1 ,t j . Then define the new interaction kernel as a projection onto the new target interaction term I 3 ,t j . Simply set V 3 ,j = V 1 ,j . Now we can check A 3 ,j ( H 3 ) 1 ,...,l 1 = A 1 ,j ( H 1 ) up to the interaction terms. Compute

̸

where the second equality comes from the sparse construction of A 3 ,j , the third equality comes from the construction of Q 3 ,j , K 3 ,j , and the last equality comes from the definition of the sparsely interacting A 1 ,j . Otherwise A 3 ,j ( h 3 ,t ) = 0 for t = s j . Thus we can conclude A 3 ,j ( H 3 ) = A 1 ,j ( H 1 ) up to interaction terms. The case l 1 +1 ≤ j ≤ l 1 + l 2 proceeds analogously. Thus we conclude

up to the interaction terms. Now we must construct the parallelized feed-forward layer FFN 3 from FFN 1 and FFN 2 . Let W 1 ,k be the k th weight matrix of FFN 1 . By assumption we know W 1 ,k is either a gating layer or is independent of token interaction terms. If W 1 ,k is independent of interaction terms we can simply set W ′ 1 ,k = W 1 ,k . Otherwise if W 1 ,k is a gating layer on some half-space { 1 , ..., p } , 1 ≤ p ≤ l 1 , we must pick a new pivot vector v ′ = R π 2 ( I 3 ,p + I 3 ,p +1 ∥I 3 ,p + I 3 ,p +1 ∥ 2 ) . Set W ′ 1 ,k accordingly to the new v ′ . We may similarly choose W ′ 2 ,k as we did for W 1 ,k . Then we set the new layer W 3 ,k as

We must also construct an embedding network W 3 , 0

which duplicates the input token h 3 ,t ∈ R d embd for parallel processing. Finally, via lemma 4, we may construct a two-layer network FFN 1 ,gating such that FFN 1 ,gating ( h 3 ,t ) = h 3 ,t if 1 ≤ t ≤ l 1 and has zeroed out data rows otherwise. We similarly define FFN 2 ,gating for l 1 +1 ≤ t ≤ l 1 + l 2 . Finally we define the output layer W 3 , max( L FFN 1 ,L FFN 2 )+2 as

as desired. The case l 1 ≤ t ≤ l 1 + l 2 proceeds analagously.

H Other Lemmas

The following lemma [Chen et al., 2022, Lemma 2] is used for our construction of charts in the proof of Theorem 2.

Lemma 8 (Local Diffeomorphism) . Suppose Assumption 1 holds for manifold M and r ≤ τ/ 4 . Then the local neighborhood U n = B ( c n , r ) ∩ M is diffeomorphic to a subset of R d . In particular, the orthogonal projection P n onto the tangent space T c n ( M ) is a diffeomorphism.

NeurIPS Paper Checklist

[Yes] [No] [NA]

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We develop approximation and statistical theory predicting transformer scaling laws by mathematically proving approximation and statistical theory which we validate on several pretraining datasets and scaling suites.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss that our analysis does not extend to analyzing benchmark capabilities of LLMs, which are of much interest, and identify this direction as promising future work.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We clearly state all assumptions in Section 2 and give complete proofs in the appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We share all our hyperparameters, code, data, and models publicly for 100% reproducibility.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.

In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We share all our hyperparameters, code, data, and models publicly for 100% reproducibility.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We specify all hyper-parameters used in our experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: We do not report error bars.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We specificy the compute resources used to run our experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have reviewed the code of ethics and confirm we are consistent.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We include discussion of broader impacts in our conclusion.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: NA

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: Yes, we give proper citation with urls and follow all licensing.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We provide a full set of instructions in the released code.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: NA

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: NA

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work derived a generalization error bound of using transformer architecture to estimate beta-Holder continuous functions. The bounds depend on the intrinsic dimension of the data. The generalization error can be decomposed into an approximation error and a variance error. They further showed that TFs with finite number of blocks can approximate arbitrary beta-Holder continuous functions up to any precision.  Both their generalization and error bound and approximation error bound exhibit power-law decay, with exponents depending on the smoothness and intrinsic dimension. Empirical results on real datasets also validated their theoretical findings.\n\nSTRENGTHS:\nThe paper present novel approximation theory for transformers in terms of the intrinsic dimension.  Empirical observations are well-aligned with the theory.\n\nWEAKNESSES:\n1. The statement that TFs can approximate $f$ to any precision $\\epsilon$ with finite number of blocks is a bit inaccurate and misleading. Namely,  Theorem 2 assumes $L_{FFN}$ is of order $O(\\log(1/\\epsilon))$. This means in the standard TF architecture, we would have $O(\\log(1/\\epsilon))$ number of FFN-attention blocks instead of only a finite number.  It would be great if the author can clarify this point.\n\nQUESTIONS:\n1. What is the number of hidden neurons in FFNs?\n2. Does the result in (3) matches the minimax lower bound for estimating beta-holder class? What is the gain of using TFs instead of simple nonparametric methods to estimate $f$?\n\nLIMITATIONS:\n1. As the author discussed, this work doesn't consider the training dynamics of the learning problem.\n2. The exponents in (3), (4) highly depend on  the smoothness of $f$ (i.e., $\\beta$) and the intrinsic dimension of the data. However, the exact value of $\\beta$ is hard to predict in practice. Is there any practical way to estimate $\\beta$?", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis document appears to be a research paper on predicting scaling laws for transformer neural networks, particularly when applied to data with low intrinsic dimensionality. Here are some key points:\n1. The paper aims to establish mathematical theories to explain and predict scaling laws observed in transformer models like large language models (LLMs).\n2. It presents both statistical estimation and approximation theories for transformers when the input data lies on a low-dimensional manifold.\n3. The main results predict power law relationships between generalization error and both training data size and network size, where the power depends on the intrinsic dimension d of the training data.\n4. A key finding is that the transformer architecture constructed in their theory only requires logarithmic depth in d, which is an advantage over feedforward networks.\n5. The authors test their theoretical predictions empirically by training LLMs on natural language datasets and find close agreement with observed scaling laws.\n6. The paper argues that the intrinsic dimension of data is crucial in determining transformer scaling laws both theoretically and empirically. It provides formal definitions and theorems around transformer networks, generalization bounds, and approximation capabilities. The work aims to bridge gaps between theory and practice in understanding neural network scaling, leveraging the low-dimensional structure often present in real-world data\n\nSTRENGTHS:\n1. The paper provides a rigorous mathematical framework for understanding transformer scaling laws, which has been a significant open question in the field.\n2. The work attempts to reconcile theoretical predictions with empirical observations, particularly in the context of large language models.\n3. By incorporating the intrinsic dimension of data, the theory provides a more nuanced understanding of scaling laws that aligns better with real-world observations.\n4. The authors test their theoretical predictions on actual language models, providing evidence for the practical relevance of their work.\n\nWEAKNESSES:\n1. Simplified assumptions: The theory relies on assumptions about data lying on a low-dimensional manifold, which may not always hold in practice, especially for complex, high-dimensional data like natural language.\n2. Limited scope: The paper focuses primarily on regression tasks, while many practical applications of transformers (like language modeling) involve more complex objectives.\n3. Generalizability: It's unclear how well the theory generalizes to other types of transformer architectures or variations that differ from the specific formulation used in the paper.\n4. Computational complexity: The paper doesn't deeply address the computational aspects of their proposed methods, which could be a practical limitation for very large models.\n5. Data dimension estimation: The reliability and stability of estimating the intrinsic dimension of complex datasets (like text) remains a challenge and could impact the practical application of the theory.\n\nQUESTIONS:\n1. Can you state the process of estimating the intrinsic dimension of textual datasets which is important for your work? If the dimension depends on the pre-trained model, it seems the posterior estimate is not supering to be good and its practicability seems weak.\n2. In Theorem 1 and 2, the parameters of the network should satisfy a specific magnitude assumption, is this assumption significant and practical in applications?\n3. In figure2, only 5 points are presented in each subfigure and the x-axis’s scope is different, how does the empirical result appear in other regions?\n4. Why did you assume the Lipschitz regularity of language modeling objective equals 1? Is there any explanation for it? \n5. The results shown in figure3 seem not good so what’s the significant reason?\n6. The order of subscripts seems incorrect in Table 1.\n\nLIMITATIONS:\n1. The assumptions in the theorem seem not practical.\n2. The construction in the proof of estimation theory is hard and not practical in empirical inference.\n3. The estimation of the intrinsic dimension is not clear.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper makes a series of contributions:\n\n- Transformer Generalization Error: Loosely speaking, assuming a transformer is trained to approximate a Holder function in a regression setting, and assuming the data lives on the low dimensional manifold, then the generalization error of the transformer is upper bounded in a particular manner (Theorem 1)\n- New Covering Number for Transformers (although I believe this is not presented in the main text)\n- Predicting Transformer Empirical Scaling Laws\n- Some investigations into how transformer hyperparameters affect the estimated intrinsic dimensionality of text data\n\n \nI am not familiar with these more theoretical methods, and I have never taken a course on differential geometry. Consequently, my review will be quite limited. I focus more on the empirical methods since I am more familiar with these.\n\nSTRENGTHS:\nNote: I am not familiar with these more theoretical methods, and I have never taken a course on differential geometry. Consequently, my review will be quite limited. I focus more on the empirical methods since I am more familiar with these.\n\n- Having minor familiarity with Sharma and Kaplan, I think this work is a great extension towards language modeling.\n- The paper is clearly very well written and thorough\n- Figure 1 is visually nice (although I don’t have the background required to understand it)\n\nWEAKNESSES:\nNote: I am not familiar with these more theoretical methods, and I have never taken a course on differential geometry. Consequently, my review will be quite limited. I focus more on the empirical methods since I am more familiar with these.\n\n- “Overall, we find the estimated ID is fairly stable across each factor.” -> Looking at Figure 4, this feels very wrong to me. In 3 out of 4 subplots, the intrinsic dimensionality clearly seems to be changing with no asymptotic value in sight.\n- [nit] Figures 2 and 3: Please use a log scale. Don’t log transform the number of samples and then plot the log-transformed variable linearly.\n\nI have many questions (below) that may become weaknesses, but I don't wish to penalize the authors if I've misunderstood their work.\n\nQUESTIONS:\n- Looking at equation (3) and (4), the scaling exponents $\\alpha_D, \\alpha_N$ are partially determined by $\\beta$. Later, $\\beta$ is set to 1. How is this justified? Why would this be true empirically?\n- The paper critically concerns estimating the intrinsic dimensionality of the pretraining dataset D, but then claims this cannot be done directly for textual datasets. Why? No justification is given.\n- Following the above bullet, the paper then estimates the intrinsic dimension of the data by estimating the intrinsic dimension of the output token embedding matrix. Why is using this alternative quantity reasonable?\n- I’m not familiar with the “Maximum Likelihood Estimation ID algorithm”. What is this algorithm? Why was it chosen? How do the results depend on this choice? The ID estimation methods I'm more familiar with are quantities like participation ratio - how do other method \n- Line 71: Is the word “layers” missing? i.e. “requiring only O(log(d)) layers independent of the desired accuracy…”\n- Figure 3: Why are there 5 GPT2 model sizes? HuggingFace offers only 4. \n- Figure 3: Why are only 4 Pythia models used? I believe that there are 8.\n- What in this paper is a prediction? To me, the hallmark of a scientific prediction is a statement regarding what one should expect beyond already-existing data, followed by new experiments to confirm the new behavior. With the exception of Figure 2, I don't see much in the way of new predictions or new experiments to confirm predictions, and even Figure 2 is relatively weak in the sense of external verifiability.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the representational capabilities of transformers in regression tasks and their correlation with scaling laws. The authors present a novel analysis of the transformer's sample complexity on datasets with low intrinsic dimension $d$, or those residing on a $d$-dimensional manifold $\\mathcal{M}$. Their findings yield a bound reminiscent of standard non-parametric regression results, with the key distinction that it depends on the intrinsic dimension rather than the actual data dimension.\n\nThe analysis hinges on a specific transformer construction that performs approximation on subregions of $\\mathcal{M}$. Notably, when the model is sufficiently wide and deep, this construction can approximate the target function to $\\epsilon$ precision with a depth independent of $\\epsilon$. This characteristic contrasts favorably with feedforward models using ReLU activation, which typically require $O(\\epsilon^{-1})$ layers.\n\nFinally, the paper shows that if we estimate the intrinsic dimension of the data, the resulting estimate of the exponent can be informative of that predicted by the sample complexity from the theoretical result.\n\nSTRENGTHS:\nAs far as I know, the theoretical result of the paper is novel and highly non-trivial. To a certain degree, it sheds light on the fundamental difference between a transformer and simpler models like the feed-forward model. I foresee these results to be useful for future research into the statistical properties of transformers.\n\nWEAKNESSES:\nThe theoretical result is novel but I am not sure how insightful the bound of the construction is about what the model actually is doing and how the result is related to scaling law. There are several weaknesses in the paper which I will discuss below:\n\nDespite its novel theoretical contributions, this paper has several significant weaknesses that warrant discussion:\n\n1. **Limited applicability to real-world scenarios**: The theoretical model presented in the paper assumes a regression or supervised learning task, which diverges significantly from how transformer models are typically used in practice. The input representation ($x$ as a $D$-dimensional vector with the sequence being a linear transformation of $x$) also appears overly simplistic and may not adequately capture the complexity of real-world data and tasks. This disconnect raises questions about how well the theoretical results translate to practical applications.\n\n2. **Relevance of the proposed construction**: While the construction presented in the paper is mathematically interesting at a technical level, it's unclear whether it accurately represents the actual learning process of transformer models. The authors do not provide empirical evidence to support that their theoretical construction aligns with the internal representations or mechanisms developed by transformers during training. This lack of validation leaves a significant gap between theory and practice. I would be more than happy to be proven wrong with empirical evidence.\n\n3. **Tenuous connection to scaling laws**: The paper's attempt to link its theoretical results to empirical scaling laws appears weak. The predicted line for $\\alpha_D$ in Figure 2 shows a substantial divergence from actual scaling laws for most datasets, with OpenWebText being a notable exception. The claimed error margin of $\\pm 0.2$ is exceptionally large on a log-log scale, potentially overstating the accuracy of the predictions. Moreover, the fit for $\\alpha_N$ is even less convincing, further undermining the paper's claims about its relevance to scaling laws.\n\nThese weaknesses suggest that while the theoretical work presented is novel, its practical implications and connections to real-world transformer behavior may be limited. The paper would benefit from stronger empirical validation and a more thorough exploration of how its theoretical insights relate to the actual functioning of transformer models in practice.\n\nQUESTIONS:\n1. Why is the intrinsic dimension dependent on the model since they are supposed to be intrinsic to the data? How do I know the approximation of it is reasonable? For example, why is \"sub-sample final-layer tokens from the embedded subsequence and shuffle together all the embedding\" a reasonable thing to do?\n\n2.  How does the computational complexity of the proposed transformer construction compare to that of the feedforward model? If the transformer requires significantly more parameters to achieve the $\\epsilon$-independent depth, does this truly represent an advantage, given that other parameters in the bound still depend on $\\epsilon$? Can you provide a more comprehensive comparison that takes into account total computational resources, including the number of parameters and operations required? Without such a comparison, it's challenging to conclude whether this result definitively demonstrates the benefit of transformers over MLPs.\n\nLIMITATIONS:\nSome limitations I discussed above are already in the paper but many are not.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-19"}