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

Authors: alex  havrilla
wenjing  liao
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 13:56:29

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.

Main paper: 
 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. Schmidt-Hieber [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.In Theorem 1, we consider the global empirical risk minimizer Tn from n i.i.d. training data {(x i , f (x i ))} n i=1 , given by Tn = arg min T∈TEQUATIONunder a properly chosen transformer network architecture T . We prove that, the generalization error of Tn satisfiesEQUATIONwhere Q denotes the distribution of x, and Õ 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:EQUATIONwith sufficient data. Our approximation theory in Theorem 2 predicts the following neural scaling law between the approximation error and the network size N :Squared Approximation Error := inf T∈T ∥T -f ∥ 2 L ∞ (M) ≲ N -α N , where α N = 2β d (4)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 length1 .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. 
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.Definition 1 (Transformer Neural Network). We define a transformer neural network T as a composition of functions of the formT(x) = D • B L T • ... • B 1 • (PE + E(x)) (5)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 takeE = E ′ • U where U ∈ R l×D and E ′ ∈ R d embd ×1applied 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 asT (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 ,and have ∥T θ ∥ L ∞ (R D ) ≤ R, ∥θ∥ ∞ ≤ κwhere ∥T θ ∥ L ∞ (R D ) ≤ R bounds the output of T and ∥θ∥ ∞ bounds the weight magnitude of T θ .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τ = inf{r > 0 : ∃x ̸ = y ∈ M, v ∈ R D such that r = ∥x -v∥ = ∥y -v∥ = inf z∈M ∥z -v∥}.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.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 Tn 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 byEQUATIONIf 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 parametersL T = O log(d) , L FFN = O log(n) , w FFN = O 1 , l = O dn d 2β+d d embd = O 1 , m = O dn d 2β+d , κ = O d 2 n 2d 2β+din 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 Tn given by (1) satisfiesEQUATIONwhere Õ 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.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 parametersL T = O log(d) , L FFN = O log(ϵ -1 ) , w FFN = O 1 , l = O dϵ -d β d embd = O 1 , m = O dϵ -d β , κ = O d 2 ϵ -2d β ,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∥T θ -f ∥ L ∞ (M) ≤ ϵ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 Mn=1 . This will give us the expression f(x) = C M n=1 f n (x)1 Un (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:EQUATIONwherefn = f n • ϕ -1 n : [0, 1] d → R and ϕ n : M → [0, 1] d isa projection onto the local tangent space. We then construct transformers to approximate the fn , ϕ n , 1 Un 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 fn on d-dimensional coordinates. L T = O log(d) , L FFN = O 1 , w FFN = O 1 , l = O dϵ -d β d embd = O 1 , m = O dϵ -d β , κ = O d 2 ϵ -2d β ,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 functionT θ ∈ T (L T , L FFN , w FFN , l, d embd , m, R, κ) such that ∥T θ -f ∥ L ∞ ([0,1] d ) ≤ ϵ.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 t1 , h t2 , 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] : A transformer 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] .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.• 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). 
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 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. 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:α D = 2 2 + d = 2 1 d 2 1 d + 1 = α N α N + 1 , α N = α D 1 -α DSee Table 1 for ID-free estimations of empirically observed exponents in the literature [Sharma and Kaplan, 2022, Hoffmann et al., 2022] .Chinchilla αN = 0.076 α D = 0.070 αN = 0.34 α D = 0.25 αD = 0.095 α N = 0.106 αD = 0.28 α N = 0.33 Table 1 : ID-free estimation of scaling exponents for GPT-2 and Chincilla. 
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] . 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. 
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.  
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. 
 
 Appendix: First we extract bounds on scaling laws in the case of regression squared error. We have the bound from the proof of Theorem 1M f (x) -Tn (x) 2 dQ(x) ≤ Õ ϵ 2 + Dd 2 ϵ -d β nwhere ϵ is the approximation error such thatinf T∈T ∥f -T∥ L ∞ (M) < ϵ Let the model size N of a transformer T ∈ T be N = L T (d 2 embd (3m + L FFN )) = log(d)(25(3ϵ -d β + log(ϵ -1 ))) = Õ ϵ -d β .Write the squared generalization error asL sq (N, n) = M f (x) -Tn (x) 2 dQ(x).ThenL sq (N, n) ≤ Õ ϵ 2 + Dd 2 ϵ -d β n = Õ N -2β d + N nIn the model scaling regime, when data is plentiful, we haveN 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α N = 2β dFor 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 haveN -2β d ≍ N n ⇐⇒ n ≍ N 1+ 2β d ≍ N 2β+d d ⇐⇒ N ≍ n d 2β+dwhere ≍ denotes in the same order. Substituting this into the error bound givesL sq (N, n) ≤ Õ 2 N n = Õ n d 2β+d -1 = Õ n -2β 2β+d ,which gives rise toα D = 2β 2β + d .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 R(f ) := P (y ̸ = f (x)). The Bayes decision rule has the form f * (x) = 1 {η(x) ≥ 1/2} where 1 denotes the indicator function and η(x) := P (y = 1|x) (8) 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 ∥η θ -η∥ L ∞ (M) ≤ ϵ (9) 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 E(f θ ) = R(f θ ) -R(f * ) = P (y ̸ = f θ (x)) -P (y ̸ = f * (x)) = P (y = f * (x)) -P (y = f θ (x)) = E x [1 {f * (x) = 1} η(x) + 1 {f * (x) = 0} (1 -η(x)) -1 {f θ (x) = 1} η(x) -1 {f θ (x) = 0} (1 -η(x))] = E x [|2η(x) -1|1{f θ (x) ̸ = f * (x)}] .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,P x (0 < |η(x) -1/2| ≤ t) ≤ c M t γ .(10) 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:EQUATIONPlugging the regression error in (9) to the excess risk bound in ( 11) with δ = ϵ, we obtainE(f θ ) ≤ 2c M ϵ 1+γ .When the margin assumption satisfies γ = 1, the excess risk bound becomesE(f θ ) ≤ 2c M ϵ 2 ,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.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:EQUATION)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.H will represent the embedding matrix of a transformer and h i will represent the ith token (column) of the embedding matrix. h j i will denote the jth component of the vector h i . ∥ • ∥ p , p > 1, will denote the pth norm of vectors, and ∥ • ∥ p,q will denote the component-wise matrix norm. For a matrixA ∈ 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 ith standard basis vector where the ith 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.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 formB(H) = FFN(MHA(H) + H) + MHA(H) + H (13)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).A Q,K,V (H) = V Hσ((KH) T QH) (14)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 ith column h i of H asEQUATION. This allows us to interpret the ith 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 asEQUATIONwhere 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 thatEQUATION)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 formFFN(h) = W LFFN σ(W LFFN-1 ...σ(W 1 h + b 1 )... + b LFFN-1 ) + b LFFN where W 2 , ..., W LFFN-1 ∈ R d embd ×d embd , W 1 ∈ R wFFN×d embd , W LFFN ∈ R d embd ×wFFN , b 1 , ..., b LFFN ∈ R d embdand the activation is ReLU. Note, each feed-forward layer is applied tokenwise to an embedding matrix H.We can also define a class FF N of feed-forward networks: 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). 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.MHA(m) = MHA θ | MHA θ is a multi-headedLastly, 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.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.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 someN ≥ 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 ith coordinate of x and n respectively. Then we construct a partition of unity (PoU) viaEQUATIONwith the trapezoid function ψ : R d → R asψ(x) =    1 |x| ≤ 1 1 -|x| 1 ≤ |x| ≤ 2 0 |x| ≥ 2 .We can now write our approximation for f asf (x) = n ϕ n (x)f n where f n = f (x n ) withx n being the center point of U n . This yields the following approximation error:∥f -f ∥ L ∞ ([0,1] d ) = sup x∈[0,1] d f (x) -f (x) = sup x∈[0,1] d f (x) - n∈{1,...,N } d f n ϕ n (x) = sup x∈[0,1] d n∈{1,...,N } d f (x)ϕ n (x) - n∈{1,...,N } d f n ϕ n (x) = sup x∈[0,1] d n∈{1,...,N } d |f (x) -f n |ϕ n (x) = sup x∈[0,1] d |x i -n i N -1 |≤ 1 N -1 |f (x) -f n | ≤ sup x∈[0,1] d |x i -n i N -1 |≤ 1 N -1 H f ∥x -x n ∥ β 2 ≤ sup x∈[0,1] d 2 d max |x i -n i N -1 |≤ 1 N -1 H f ∥x -x n ∥ β 2 ≤ sup x∈[0,1] d 2 d H f d β (N -1) β = 2 d d β H f (N -1) β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 N d+N d forming input x ′ ∈ R 1×d+N d+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+N d+N d d) of the form      x 1 ... x d 0 N d ... 0 N d d 0 ... 0 0 N d ... 0 N d d 0 ... 0 0 N d ... 0 N d 0 ... 0 0 N d ... 0 N d d 0 ... 0 0 N d ... 0 N d d     where d embd = 5. To this we add a fixed positional encoding PE ∈ R d embd ×(d+N d+N d d) , producing the first hidden embedding matrixH 1 =    x 1 ... x d 0 N d ... 0 N d d 0 ... 0 0 N d ... 0 N d d I 1 ... ... ... ... I d+dN +N d d 1 ... ... ... ... 1   where I i ∈ R 2 represents an interaction term determining when each token embedding will interact with another in the attention mechanism. We can setEQUATION)) 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 ψ(3N (x ij-1 N -1 )) Now we pre-compute the terms ψ(3N (x ij-1 N -1 )) from which we can construct our partition of unity. We defines 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 ij-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+N d .Define B 1 = B(N d, 7) i.e. B 1 has a multi-headed attention layer with N d attention heads and a feed-forward layer with depth 7. We will construct each of the ijth blocks using the Interaction Lemma 3. We may choose data kernelsQ data ij = 0 0 0 0 1 0 0 0 0 1 ∈ R 2×d embd , K data ij = 1 0 0 0 0 0 0 0 0 -j N -1 + 1 ∈ R 2×d embdso that for the token embedding h ij = h d+i * d+j , we haveA ij (h ij ) = l k=1 σ(⟨Q ij h ij , K ij h k ⟩)V ij h k = l k=1 σ(⟨Q data ij h ij , K data ij h k ⟩ + ⟨Q I ij h ij , K I ij h k ⟩ -C)e 2 = l k=1 σ(h 1 k - j -1 N -1 + 1 + ⟨Q I ij h ij , K I ij h k ⟩ -C)e 2 = σ(x i - j -1 N -1 + 1)e 2 = (x i - j -1 N -1 + 1)e 2where we can choose token h ij to only interact with token h i . Otherwise we have A ij (h t ) = 0 forh t ̸ = h ij . Further, the weights of θ Aij of A ij are bounded such that ∥θ Aij ∥ ∞ = 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 byH 1.5 =    x 1 ... x d 0 N d ... ... 0 N d d 0 ... 0 x 1 + c ... x d -1 + c 0 N d d I 1 ... ... ... ... ... I d+N d+N d d 1 ... ... ... ... ... 1    .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. DefineW 1 =      0 1 0 0 -1 + M 0 -1 0 0 M 0 0 1 0 0 0 0 0 1 0 0 0 0 0 1      , b 1 = 0 W 2 =      1 0 0 0 -M 0 1 0 0 -M 0 0 1 0 0 0 0 0 1 0 0 0 0 0 1      , b 2 = 0.Here W 7 =     1 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0     , b 7 = 0which zeros out all except the first two components. This produces the next embedding matrix H 2 :H 2 =    x 1 ... x d s 1,1 ... s d,N 0 N d d 0 ... 0 0 N d ... ... 0 N d d I 1 ... ... ... ... ... I d+N d+N d d 1 ... ... ... ... ... 1    .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 ith partial product of 2 k-1 terms. Formally, 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,2i-1 • p n,k-1,2i . Storage and assembly will be done in the last dp n,k,i = p n,k-1,2i-1 • p n,k-1,2i 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+ N d + 1, ..., d + N d + 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+N d . 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+N d+d d p=1 (n p -1)N d-p +i . Concretely, for a fixed attention head A n,i define the data kernelsQ data n,i = 0 0 0 0 1 0 0 0 0 0 K data n,i =1 0 0 0 0 0 0 0 0 0 p n,k,i = p n,k-1,2i-1 • p n,k-1,2i with p n,1,i = s i,n i for n ∈ {1, ..., N } d , 1 ≤ k ≤ log 2 (d), 1 ≤ i ≤ d 2 k .via the Interaction Lemma 3 we may pick A n,i so thatA n,i h n,i = l k=1 σ(⟨Q n,i h n,i , K n,i h k ⟩)V n,i h k = l k=1 σ(h 5 n,i h 1 k + Q I n,i I n,i • K I n,i I k -C)e 1 = σ(1 • s i,n i )e 1 = s i,n i e 1where 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 haveA n,i h k = 0 when k ̸ = d + N d + 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 whereH 3 =    x 1 ... x d s 1,1 ...P 1 ∈ R N d d is P 1 = p {1} d ,1,1 p {1} d ,2,2 ... p {1} d ,1,d p {2}×{1} d-1 ,1,1 ... p {N } d ,1,d = [s 1,1 s 2,1 ... s d,1 s 1,2 ... s d,N ]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 asA n,i for n ∈ {1, ..., N } d , 1 ≤ i ≤ d 2 . Define data kernels Q data n,i = 1 0 0 0 0 0 0 0 0 0 K data n,i = 1 0 0 0 0 0 0 0 0 0where we choose the interaction terms so that the token containing p n,1,2i-1 only interacts with the token containing p n,1,2i 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,2i-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 . WriteA n,i h n,i = l k=1 σ(⟨Q n,i h n,i , K n,i h k ⟩)V n,i h k = σ(p n,1,2i-1 p n,1,2i )e 2 = σ(p n,2,i )e 2 = p n,2,i e 2 .Adding the result onto H 3 gives the intermediate embedding matrix Note: * denotes a non-essential entry which can take any value. We then design our FFN network as W 1 x whereH 3.5 =    x 1 ... x d sW 1 =      -1 1 0 0 0 0 -1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0      , b 1 = 0which 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 givesH 4 =    * ... P 2 0 d+N d+N d d ... 0 I 1 ... I d+N d+N d d 1 ... 1   whereP 2 = P 1.5 ∈ R N d d .More generally define the kth block implementing the kth level of the parallel product asB k+2 = B(N d d 2 k , FF N (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,2i-1 only interacts with the token containing p n,k-1,2i to produce p n,k,i :A n,k,i h n,i = l t=1 σ(⟨Q n,i h n,i , K n,i h t ⟩)V n,i h t = σ(p n,k-1,2i-1 p n,k-1,2i )e 2 = p n,k,i e 2 .In the FFN set W 1 as before, yieldingH k+2 =    * ... P k 0 d+N d+N d d ... 0 I 1 ... I d+N d+N d d 1 ... 1   whereP k = p {1} d ,k,1 * ... p {N } d ,k,1 * ... . At k = log(d) we have H log(d)+2 =    * ... P log(d) 0 d+N d+N d d ... 0 I 1 ... I d+N d+N d d 1 ... 1   where componentP d+N d+ 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 kernelsQ n = 1 0 0 0 0 0 0 0 0 1 K n = 0 0 0 0 f n 0 0 0 0 R where ∥f ∥ L ∞ ([0,1] d ) ≤R and the token h d+N d+ d p=1 (n p -1)N p +1 = h n,1 only interacts with itself so thatA n h n,1 = l k=1 σ(h 1 n,1 • f n + R + Q I n I n,1 • K I n I k -C)e 2 = σ(f n ϕ n (x) + R)e 2 = (f n ϕ n (x) + R)e 2where the last equality comes from|f n ϕ n (x)| ≤ |f n | ≤ R. We implement FFN to subtract R via W 1 =      0 0 0 0 0 0 1 0 0 -R 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0      so that H log(d)+4 =    * ... ... 0 d+N d ... P log(d)+4 I 1 ... I d+N d+N d d 1 ... 1   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 asQ =      0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 R      K =      0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1      V =      0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0     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 givesA 1 h t = l k=1 σ(⟨Qh t , Kh k ⟩)V h k = l k=1 σ(h 2 k + R)e 1 = (lR + n∈{1,...,N } d f n ϕ n )e 1 .We can subtract lR with the FFN asW 1 =      1 0 0 0 -lR 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0      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.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.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 Un . This will give us the expressionEQUATIONwhere 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:f (x) = C M n=1 fn • ϕ n (x)1 Un (x)where fn : R d → R and ϕ n : M → R d is a projection. This step is critical: we now must only approximate a low-dimensional function fn and a simple projection ϕ n instead of the highdimensional function f n . After this, we must simply approximate each fn , ϕ n , 1 Un using transformer neural networks.Step . . , 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 Mn=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 cn (M) . Via Lemma 8, as long as r ≤ τ /4, each ϕ n is a diffeomorphism between U n and a subset in T cn (M). Here the number of charts C M is a constant depending on the complexity of the manifold M. We can bound it as For the open covering {U n } C M n=1 , there exists a smooth partition of unity Tu, 2011] . Each ρ n is supported only in its corresponding chart U n . We additionally include an indicator function 1 Un for each neighborhood U n , resulting the decompositionC M ≲ SA(M){ρ n } C M n=1 such that f (x) = C M n=1 ρ n (x)f (x) for x ∈ M [f (x) = C M n=1 f (x)ρ n (x) fn(x)1 Un (x).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 cn (M) ⊆ [0, 1] d . We can writef n (x) = (f n • ϕ -1 n ) • ϕ n (x) = fn • ϕ n (x)where fn : R d → R and again ϕ n : M ⊆ R D → R d is an orthogonal projection. Then we can write f asf (x) = C M n=1 fn • ϕ n (x)1 Un (x).(20)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 Un . We will approximate each component with transformer neural networks. Specifically, let T fn approximate fn , T ϕn approximate ϕ n , and T 1 Un approximate 1 Un , 1 ≤ n ≤ C M . Then we will have the overall approximationf (x) ≈ T f (x) = C M n=1 T fn • T ϕn (x)T 1 Un (x).This yields the error decomposition∥f -T f ∥ L ∞ (M) = C M n=1 ∥ fn • ϕ n 1 Un -T fn • T ϕn T 1 Un ∥ L ∞ (M) ≤ C M n=1 ∥ fn • ϕ n T 1 Un -T fn • T ϕn T 1 Un ∥ L ∞ (M) + ∥ fn • ϕ n 1 Un -fn • ϕ n T 1 Un ∥ L ∞ (M) ≤ C M n=1 ∥ fn • ϕ n -T fn • T ϕn ∥ L ∞ (Un) + E n,3 ≤ C M n=1 ∥ fn • ϕ n -T fn • ϕ n ∥ L ∞ (Un) + ∥T fn • ϕ n -T fn • T ϕn ∥ L ∞ (Un) + E n,3 = C M n=1 E n,1 + E n,2 + E n,3where the error terms E n,1 , E n,2 , E n,3 denoteEQUATION) respectively. It remains to construct transformer approximations for each term.Step 2.1: Approximating fn and Bounding E n,1 We first handle the E n,1 error in (21). Fix1 ≤ n ≤ C M . We can write ∥ fn • ϕ n -T fn • ϕ n ∥ L ∞ (Un) ≤ ∥ fn -T fn ∥ L ∞ ([0,1] d )so it suffices to approximate fn : [0, 1] d → R. Now choose the desired accuracy δ n,1 > 0. Via Lemma 1 we can construct an approximation such that 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):∥ fn -T fn ∥ L ∞ ([0,1] d ) < δ n,E n,2 = ∥T fn • ϕ n -T fn • T ϕn ∥ L ∞ (Un)We will implement ϕ n using a transformer block B ∈ B(dD, 5). We can identify the tangent space at the point c n asT cn (M) = span(v n,1 , ..., v n,d )where the vectors v n,i ∈ R D , 1 ≤ i ≤ d, form an orthonormal basis of T cn (M). Then we can write the projection map ϕ n asϕ n (x) = s n (V T n (x -c n ) + u n ) 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 formH =    x 1 ... x D 0 d 0 ... ... 0 I 1 ... ... I l 1 ... ... 1    ∈ R d embd ×lWe 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 kernelsQ data i = 0 0 0 0 1 0 0 0 0 1 K data i = v j n,i 0 0 0 -v j n,i c j n 0 0 0 0 2Mvia the Interaction Lemma 3 we can define A i,j so that token h D+i interacts with token h j so thatA i,j (h D+i ) = σ(⟨Q data i h D+i , K data i h j ⟩) = σ(v j n,i x j -v j n,i c j n + 2M ) = σ(v j n,i (x j -c j n ) + 2M ) = (v j n,i (x j -c j n ) + 2M )e 1and 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 ) + 2M ≥ 0 for all n, i, j. Then the output of the multi-headed attention MHA on h D+i isMHA(h D+i ) = 1≤k≤d 1≤j≤D A k,j h D+i = D j=1 A i,j h D+i = D j=1 v j n,i (x j -c j n )e 1 + 2M e 1 = ⟨v n,i , x -c n ⟩e 1 + 2DM e 1This yields the intermediate embedding matrixH + MHA(H) =    x 1 ... x D ⟨v n,1 , x -c n ⟩ + 2DM ... ⟨v n,d , x -c n ⟩ + 2DM 0 ... ... ... ... 0 I 1 ... ... ... ... I l 1 ... ... ... ... 1   We can then design a five layer FFN layer to subtract 2DM and then add u n and multiply s n only from the embedding tokens D + 1, ..., D + d. The first layer will simply subsract 2DM and then add u n and multiply s n from each token. This can be implemented by the matrixW 1 =      s n 0 0 0 -2DM s n + u n s n 0 0 0 0 0 0 0 1 0 0 0 0 0 1 0 0 0 0 0 1     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.H ′ = B(H) =    x 1 ... x D s n (⟨v n,1 , x -c n ⟩ + u n ) ... s n (⟨v n,d , x -c n ⟩ + u n ) 0 ... ... ... ... 0 I 1 ... ... ... ... I l 1 ... ... ... ... 1   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 Un (x) and Bounding E n,3 Again fix 1 ≤ n ≤ C M . By construction we can write1 Un (x) = 1 ∥x -c n ∥ 2 2 < r 2 0 ∥x -c n ∥ 2 2 ≥ r 2Suppose we are given an input embedding matrix of the formH 1 =    x 1 ... x D 0 D 0 ... ... 0 I 1 ... ... I l 1 ... ... 1    ∈ R d embd ×2DWe 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 H 2 = B 1 (H 1 ) =    x 1 ... x D x 1 -c 1 n ... x D -Q data i = 1 0 0 0 0 0 0 0 0 0 K data i = 1 0 0 0 0 0 0 0 0 0Then via the Interaction Lemma 3 we can construct A i so that the token h D+i interacts with itself such thatA i (h D+i ) = σ((x i -c i n ) 2 )e 2 = (x i -c i n) 2 e 2 and 0 otherwise. The output of the resulting multi-headed attention givesH ′ 2 =    x 1 ... x D x 1 -c 1 n ... x D -c D n 0 ... 0 (x 1 -c 1 n ) 2 ... (x D -c D n ) 2 I 1 ... ... ... ... I l 1 ... ... ... ... 1   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 2D . This gives the outputH 3 = B 2 (H 2 ) =    x 1 ... x D (x 1 -c 1 n ) 2 ... (x D -c D n ) 2 0 ... ... ... 0 I 1 ... ... ... ... I l 1 ... ... ... ... 1   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 kernelsQ data i = 0 0 0 0 1 0 0 0 0 0 K data i =1 0 0 0 0 0 0 0 0 0 via the Interaction Lemma 3 we construct A i so that h D+1 interacts only with h D+1+i such thatA i (h D+1 ) = σ(⟨Q data i h D+1 , K data i h D+1+i ⟩)e 1 = σ((x i+1 -c i+1 n ) 2 )e 1 = (x i+1 -c i+1 n ) 2 e 1The output of the multi-headed attention is thenMHA(h D+1 ) = D-1 i=1 A i (h D+1 ) = D-1 i=1 (x i+1 -c i+1 n )e 1and 0 otherwise. This gives the intermediate outputH ′ 3 =    x 1 ... x D D i=1 (x i -c i n ) 2 ... (x D -c D n ) 2 0 ... ... ... ... 0 I 1 ... ... ... ... I l 1 ... ... ... ... 1   where in particular h 1 D+1 = ∥x -c n ∥ 2 2 . It remains to approximate the indicator1 r 2 (s) = 1 s < r 2 0 s ≥ r 22.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 of1 Un is E n,3 = ∥ fn • ϕ n 1 Un -fn • ϕ n T 1 Un ∥ L ∞ (M). We know on the boundary of U n that fn • ϕ 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 indicator1r 2 ,∆ (s) =      1 s ≤ r 2 -∆ 1 -1 ∆ s + r 2 -∆ ∆ r 2 -∆ < s < r 2 0 s ≥ r 2for some ∆ > 0. Clearly for s ≤ r 2 -∆ and s ≥ r 2 we have 1r 2 ,∆ (s) = 1 r 2 (s). So we argue fn is small for x ∈ M in the rangeK 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 :E n,3 ≤ c r ∆for some constant c. So it suffices to make ∆ small. It then remains to implement 1r 2 ,∆ . We will do so using a feed-forward layer FFN ∈ FFN (O(log( 1 ∆ ))). Note that we can write the target exactly as1r 2 ,∆ (s) = σ 1 -σ 1 ∆ (s -(r 2 -∆))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 layersO(log( 1 ∆ )) as 1r 2 ,∆ (s) = σ   1 -σ    1 ∆ 1 log( 1 ∆ ) • σ... • σ    1 ∆ 1 log( 1 ∆ ) (s -(r 2 -∆)))...         which prevents the weights( 1 ∆ ) 1 log( 1 ∆ )from blowing up. We can then approximate each 1 Un 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 asf (x) = C M n=1 fn • ϕ n (x)1 Un (x).We argued we can successfully approximate 1. fn up to δ n,1 accuracy in L ∞ via a transformer neural network with O(dδ where each c fn = fn • ϕ n (x) and c Un = 1 Un (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 kernelsQ data i = 1 0 0 0 0 -1 0 0 0 M K data i = 1 0 0 0 0 0 0 0 0 1with interaction kernels chosen so that h D+i only interacts with h D+C M +i under A i . ThenA i (h D+i ) = σ(c fi c Ui -c fi + M )e 1 = (c fi c Ui -c fi + M )e 1We 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.EQUATIONIt simply remains to control the errors. Recall we have the error bound∥f -T f ∥ L ∞ (M) ≤ C M n=1 E n,1 + E n,2 + E n,3 ≤ C M n=1 δ n,1 + δ n,3 = C M δ n,1 + C M δ n,3Then for a target accuracy ϵ we chooseδ n,1 = ϵ 2C M and δ n,3 = ϵ 2C M so that ∥f -T f ∥ L ∞ (M) ≤ C M δ n,1 + C M δ n,3 = C M ϵ 2C M + C M ϵ 2C M = ϵ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)).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.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 Tn 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.EQUATION)The first (bias) term in ( 24) can be controlled via Theorem 2:E 1 n n i=1 ( Tn (x i ) -f (x i )) 2 = E inf T∈T 1 n n i=1 (T(x i ) -f (x i )) 2 ≤ inf T∈T E 1 n n i=1 (T(x i ) -f (x i )) 2 = inf T∈T M T(x) -f (x) 2 dQ ≤ inf T∈T M ∥T -f ∥ 2 L ∞ (M)dQ < ϵ 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 boundEQUATIONwhere N ( δ 4R , 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 N (δ, T , ∥ • ∥ ∞ ) ≤ 2 L 2 T +1 L FFN M 3L T d 18L 2 T embd w 18L 2 T LFFN FFN κ 6L 2 T LFFN m L 2 T l L 2 T δ 4d 2 embd w 2 FFN D(m+LFFN)L T . Lemma 2 is proved in Appendix F.2.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 obtainE M ( Tn (x) -f (x)) 2 dQ -2E 1 n n i=1 ( Tn (x i ) -f (x i )) 2 ≤ inf δ>0 104R 2 3n log N δ 4R , T , ∥ • ∥ ∞ + 4 + 1 2R δ ≤ 104R 2 3n log N 1 4Rn , T , ∥ • ∥ ∞ + 4 + 1 2R 1 nwhere we set δ = 1 n . This yieldsN δ 4Rn , T , ∥ • ∥ ∞ ≤ log 2 L 2 T +1 L FFN M 3L T d 18L 2 T embd w 18L 2 T LFFN FFN κ 6L 2 T LFFN m L 2 T l L 2 T n 4d 2 embd w 2 FFN D(m+LFFN)L T ≤ 4d 2 embd w 2 FFN D(m + L FFN )L T (5L 2 T L FFN log(2M L FFN d embd κmln)) ≤ 20 log(2M L FFN d embd w FFN κmln)Dd 2 embd w 2 FFN mL 3 T L 2 FFN . Recall for target accuracy ϵ > 0 we can choose L T = log(d), L FFN = log(ϵ -1 ), w FFN ≤ 4C M , d embd = 5, l ≤ C M dϵ -d β , m ≤ C M dϵ -d β , κ ≤ (M C M dϵ -d β ) 2 , l ≤ C M dϵ -d β. This simplifies the above toN ( δ 4R , T , ∥ • ∥ ∞ ) ≤ 900 log(d) 2 log(60d 3 ϵ -4 d β n) C 3 M Ddϵ -d β ≤ Õ Dd 2 ϵ -d βwhere Õ hides log terms and constants. Then we can bound the variance term asE M ( Tn (x) -f (x)) 2 dQ -2E 1 n n i=1 ( Tn (x i ) -f (x i )) 2 ≤ Õ Dd 2 ϵ -d β nPutting together the bounds on the bias and variance yields a bound on the empirical risk:E M ( Tn (x) -f (x)) 2 dQ ≤ Õ ϵ 2 + Dd 2 ϵ -d β nand 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:E M ( Tn (x) -f (x)) 2 dQ -2E 1 n n i=1 ( Tn (x i ) -f (x i )) 2 ≤ Õ Dd 2 n -2β 2β+das desired.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|T(x) -T ′ (x)| = |D • B L • ... • B 1 • (PE + E(x)) -D ′ • B ′ L • ... • B ′ 1 • (PE + E ′ (x))| ≤ ∥B L • ... • B 1 • (PE + E(x)) -B ′ L • ... • B ′ 1 • (PE + E ′ (x))∥ ∞ .To handle this term, we write∥B L • ... • B 1 • (PE + E(x)) -B ′ L • ... • B ′ 1 • (PE + E ′ (x))∥ ∞ ≤ ∥B L • ... • B 1 • (PE + E(x)) -B L • ... • B 1 • (PE + E ′ (x))∥ ∞ + ∥B L • ... • B 1 • (PE + E ′ (x)) -B ′ L • ... • B ′ 1 • (PE + E ′ (x))∥ ∞ .We will first bound the second term.Consider two multi-headed attention layers MHA 1 , MHA 2 ∈ MHA(m) with attention headsA i j , i ∈ {1, 2}, 1 ≤ j ≤ m. For H ∈ [0, M ] d embd ×l we compute ∥MHA 1 (H) -MHA 2 (H)∥ ∞ = ∥ m j=1 A 1 j (H) - m j=1 A 2 j (H)∥ ∞ = ∥ m j=1 A 1 j (H) -A 2 j (H) ∥ ∞ ≤ m j=1 ∥A 1 j (H) -A 2 j (H)∥ ∞ 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 ∥A 1 (H) -A 2 (H)∥ ∞ ≤ max 1≤j≤l ∥A 1 (h j ) -A 2 (h j )∥ ∞ = max 1≤j≤l ∥ l k=1 σ(⟨Q 1 h j , K 1 h k ⟩)V 1 h k - l k=1 σ(⟨Q 2 h j , K 2 h k ⟩)V 2 h k ∥ ∞ ≤ max 1≤j≤l l k=1 ∥σ(⟨Q 1 h j , K 1 h k ⟩)V 1 h k -σ(⟨Q 2 h j , K 2 h k ⟩)V 2 h k ∥ ∞ ≤ max 1≤j≤l l k=1 ∥σ(⟨Q 1 h j , K 1 h k ⟩)V 1 h k -σ(⟨Q 2 h j , K 2 h k ⟩)V 1 h k ∥ ∞ + ∥σ(⟨Q 2 h j , K 2 h k ⟩)V 1 h k -σ(⟨Q 2 h j , K 2 h k ⟩)V 2 h k ∥ ∞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∥σ(⟨Q 1 h j , K 1 h k ⟩)V 1 h k -σ(⟨Q 2 h j , K 2 h k ⟩)V 1 h k ∥ ∞ =|⟨Q 1 h j , K 1 h k ⟩ -⟨Q 2 h j , K 2 h k ⟩|∥V 1 h k ∥ ∞ ≤ |⟨Q 1 h j , K 1 h k ⟩ -⟨Q 1 h j , K 2 h k ⟩| + |⟨Q 1 h j , K 2 h k ⟩ -⟨Q 2 h j , K 2 h k ⟩| ∥V 1 h k ∥ ∞ ≤ |⟨Q 1 h j , K 1 h k -K 2 h k ⟩| + |⟨Q 1 h j -Q 2 h j , K 2 h k ⟩| ∥V 1 h k ∥ ∞ ≤ |∥Q 1 h j ∥ 2 ∥K 1 h k -K 2 h k ∥ 2 + ∥Q 1 h j -Q 2 h j ∥ 2 ∥K 2 h k ∥ 2 ∥V 1 h k ∥ ∞ ≤ |∥Q 1 ∥ 1 ∥h j ∥ ∞ ∥K 1 -K 2 ∥ 1 ∥h k ∥ ∞ + ∥Q 1 -Q 2 ∥ 1 ∥h j ∥ ∞ ∥K 2 ∥ 1 ∥h k ∥ ∞ ∥V 1 h k ∥ ∞ ≤ 2κ 2 d 4 embd M 2 η ∥V 1 h k ∥ ∞ ≤2κ 3 d 6 embd M 3 η 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∥σ(⟨Q 2 h j , K 2 h k ⟩)V 1 h k -σ(⟨Q 2 h j , K 2 h k ⟩)V 2 h k ∥ ∞ =∥V 1 -V 2 ∥ 1 ∥⟨Q 2 h j , K 2 h k ⟩h k ∥ ∞ ≤d 2 embd κη∥⟨Q 2 h j , K 2 h k ⟩h k ∥ ∞ ≤d 2 embd κη∥Q 2 h j ∥ 2 ∥K 2 h k ∥ 2 ∥h k ∥ ∞ ≤d 6 embd κ 3 M 3 η. Thus we can bound the attention difference ∥A 1 (H) -A 2 (H)∥ ∞ as ∥A 1 (H) -A 2 (H)∥ ∞ ≤ max 1≤j≤l l k=1 ∥⟨Q 1 h j , K 1 h k ⟩V 1 h k -⟨Q 2 h j , K 2 h k ⟩V 1 h k ∥ ∞ + ∥⟨Q 2 h j , K 2 h k ⟩V 1 h k -⟨Q 2 h j , K 2 h k ⟩V 2 h k ∥ ∞ ≤ max 1≤j≤l l k=1 3κ 3 d 6 embd M 3 η = 3κ 3 d 6 embd M 3 lη.We can also now bound the multi-headed difference∥MHA 1 (H) -MHA 2 (H)∥ ∞ as ∥MHA 1 (H) -MHA 2 (H)∥ ∞ ≤ m j=1 ∥A 1 j (H) -A 2 j (H)∥ ∞ ≤ 3κ 3 d 6 embd M 3 η = 3κ 3 d 6 embd M 3 mlη.A bound on the first residual layer in the transformer block immediately follows via∥(H + MHA 1 (H)) -(H + MHA 2 (H))∥ ∞ = ∥MHA 1 (H) -MHA 2 (H)∥ ∞ .To bound the difference in the output of the FFN layer in the transformer block we have forf 1 , f 2 ∈ FFN(L FFN ) and input H ∈ [0, M ′ ] d embd ×l we have ∥f 1 (H) -f 2 (H)∥ ∞ ≤ L FFN (d embd M ′ + 2)(κd embd ) LFFN-1 η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 , ∥B 1 (H) -B 2 (H)∥ ∞ = ∥ H + MHA 1 (H) + FFN 1 (H + MHA 1 (H)) -H + MHA 2 (H) + FFN 2 (H + MHA 2 (H)) ∥ ∞ = ∥MHA 1 (H) + FFN 1 (H + MHA 1 (H)) -MHA 2 (H) + FFN 2 (H + MHA 2 (H))∥ ∞ ≤ ∥MHA 1 (H) -MHA 2 (H)∥ ∞ + ∥FFN 1 (H + MHA 1 (H)) -FFN 2 (H + MHA 2 (H))∥ ∞ ≤ 3κ 3 d 6 embd M 3 mlη + ∥FFN 1 (H + MHA 1 (H)) -FFN 2 (H + MHA 2 (H))∥ ∞where we can bound the first term via analysis above. For the second term write∥FFN 1 (H + MHA 1 (H)) -FFN 2 (H + MHA 2 (H))∥ ∞ ≤∥FFN 1 (H + MHA 1 (H)) -FFN 1 (H + MHA 2 (H))∥ ∞ + ∥FFN 2 (H + MHA 2 (H)) -FFN 1 (H + MHA 2 (H))∥ ∞ .To handle the first term we must bound he Lipschitz constant of FFN1 . Let H, H ′ ∈ R d embd ×l . Write ∥FFN 1 (H) -FFN 1 (H ′ )∥ ∞ = max 1≤j≤l ∥FFN 1 (h j ) -FFN 1 (h ′ j )∥ ∞ = max 1≤j≤l ∥W LFFN σ(W LFFN-1 ...σ(W 1 h j + b 1 )... + b LFFN-1 ) + b LFFN -W LFFN σ(W LFFN-1 ...σ(W 1 h ′ j + b 1 )... + b LFFN-1 ) -b LFFN ∥ ∞ ≤ max 1≤j≤l ∥W LFFN ∥ 1 ∥W LFFN-1 ...σ(W 1 h j + b 1 )... + b LFFN-1 -W LFFN-1 ...σ(W 1 h ′ j + b 1 )... + b LFFN-1 ∥ ∞ ≤ max 1≤j≤l LFFN i=1 ∥W i ∥ 1 ∥h j -h ′ j ∥ ∞ ≤ w 2LFFN FFN κ LFFN ∥H -H ′ ∥ ∞where we again use that ∥W x∥ ∞ ≤ ∥W ∥ 1 ∥x∥ ∞ and that σ(•) is 1-Lipschitz. Applying this to the first term from above we get∥FFN 1 (H + MHA 1 (H)) -FFN 1 (H + MHA 2 (H))∥ ∞ ≤ w 2LFFN FFN κ LFFN ∥MHA 1 (H) -MHA 2 (H)∥ ∞ ≤ w 2LFFN FFN κ LFFN 3κ 3 d 6 embd M 3 mlη = 3κ 3+LFFN w 2LFFN FFN d 6 embd M 3 mlη. We can bound ∥FFN 2 (H + MHA 2 (H)) -FFN 1 (H + MHA 2 (H))∥ ∞ as ∥FFN 2 (H + MHA 2 (H)) -FFN 1 (H + MHA 2 (H))∥ ∞ ≤∥FFN 2 (H + MHA 2 (H)) -FFN 1 (H + MHA 2 (H))∥ ∞ ≤ L FFN (w FFN M ′ + 2)(κw FFN ) LFFN-1 η where ∥H + MHA 2 (H)∥ ∞ ≤ M ′ .Putting all the estimates together, we have∥B 1 (H) -B 2 (H)∥ ∞ ≤3κ 3 d 6 embd M 3 mlη + ∥FFN 1 (H + MHA 1 (H)) -FFN 2 (H + MHA 2 (H))∥ ∞ ≤3κ 3 d 6 embd M 3 mlη + 3κ 3+LFFN w 2LFFN FFN d 6 embd M 3 mlη + L FFN (w FFN M ′ + 2)(κw FFN ) LFFN-1 η ≤ 4κ 3+LFFN w 2LFFN FFN d 6 embd M 3 ml + L FFN (w FFN M ′ + 2)(κw FFN )LFFN-1 η and it remains to control M ′ in terms of M . We know ∥H + MHA(H)∥ ∞ ≤ M ′ and so we compute∥H + MHA(H)∥ ∞ ≤ M + ∥MHA(H)∥ ∞ ≤ M + m j=1 max 1≤i≤l l k=1 ∥⟨Q j h i , K j h k ⟩V j h k ∥ ∞ ≤ M + mld 6 embd κ 3 M ≤ 2d 6 embd κ 3 mlM.Thus, we can complete our estimate on∥B 1 (H) -B 2 (H)∥ ∞ as ∥B 1 (H) -B 2 (H)∥ ∞ ≤ 4κ 3+LFFN w 2LFFN FFN d 6 embd M 3 ml + L FFN (w FFN (2d 6 embd κ 3 mlM ) + 2)(κw FFN ) LFFN-1 η. Now consider the multi-block difference ∥B L T • ... • B 1 (H) -B ′ L T • ... • B ′ 1 (H)∥ ∞ ≤ ∥B L T • ... • B 1 (H) -B L T • ... • B ′ 1 (H)∥ ∞ + ∥B L T • ... • B ′ 1 (H) -B ′ L T • ... • B ′ 1 (H)∥ ∞ . To bound the second difference we must bound the output ∥B(H)∥ ∞ , B ∈ B(m, L FFN ). Suppose ∥H∥ ∞ ≤ M . Write ∥B(H)∥ ∞ = ∥FFN(H + MHA(H)) + MHA(H) + H∥ ∞ ≤ ∥H∥ ∞ + ∥MHA(H)∥ ∞ + ∥FFN(H + MHA(H))∥ ∞ ≤ 2d 6 embd κ 3 mlM + ∥FFN(H + MHA(H))∥ ∞ so we must bound the output of the FFN layer. We have forH ′ ∈ R d embd ×l with ∥H ′ ∥ ∞ ≤ M , we have ∥FFN(H ′ )∥ ∞ ≤ (κd 2 embd ) LFFN-1 M + κ(κd 2 embd ) LFFN-2 . Since ∥H + M HA(H)∥ ∞ ≤ 2d 6 embd κ 3 mlM , we have ∥B(H)∥ ∞ ≤ 2d 6 embd κ 3 mlM + (κw 2 FFN ) LFFN-1 2d 6 embd κ 3 mlM + κ(κw 2 FFN ) LFFN-2 ≤ 4d 4 embd w 2LFFN FFN κ LFFN+1 mlM. Iterating this gives ∥B L T -1 • ... • B 1 (H)∥ ∞ ≤ (4d 4 embd w 2LFFN FFN κ LFFN+1 ml) L T -1 M.Then we can bound the above second term as∥B L T • ... • B ′ 1 (H) -B ′ L T • ... • B ′ 1 (H)∥ ∞ ≤ 4κ 3+LFFN d 6+2LFFN embd ((4d 2LFFN+4 embd κ LFFN+1 ml) L T -1 M ) 3 ml + L FFN (d embd (2d 6 embd κ 3 ml((4d 2LFFN+4 embd κ LFFN+1 ml) L T -1 M )) + 2)(κd embd ) LFFN-1 η ≤4 3L T -2 L FFN d 6L T LFFN+12L T -10LFFN-6 embd κ 3L T LFFN+3L T -2LFFN M 3 m L T l L T η ≤4 3L T L FFN M 3 d 18L T LFFN embd w 18L T LFFN FFN κ 6L T LFFN m L T l L T η.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 ∥B(H) -B(H ′ )∥ ∞ = ∥FFN(H + MHA(H)) -FFN(H ′ + MHA(H ′ )) + MHA(H) -MHA(H ′ )∥ ∞ ≤ ∥FFN(H + MHA(H)) -FFN(H ′ + MHA(H ′ ))∥ ∞ + ∥MHA(H) -MHA(H ′ )∥ ∞ ≤ ∥FFN(H + MHA(H)) -FFN(H ′ + MHA(H ′ ))∥ ∞ + ∥MHA(H) -MHA(H ′ )∥ ∞ ≤ w 2LFFN FFN κ LFFN ∥MHA(H) -MHA(H ′ ) + H -H ′ ∥ ∞ + ∥MHA(H) -MHA(H ′ )∥ ∞ ≤ w 2LFFN FFN κ LFFN ∥H -H ′ ∥ ∞ + (w 2LFFN FFN κ LFFN + 1)∥MHA(H) -MHA(H ′ )∥ ∞ and it suffices to control the Lipschitz constant of MHA ∈ MHA(m). For the same H, H ′ compute ∥MHA(H) -MHA(H ′ )∥ ∞ ≤ max 1≤i≤l m j=1 l k=1 ∥⟨Q j h i , K j h k ⟩V j h k -⟨Q j h ′ i , K j h ′ k ⟩V j h ′ k ∥ ∞ .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 j h i , K j h k ⟩V j h k -⟨Q j h ′ i , K j h ′ k ⟩V j h ′ k ∥ ∞ ≤∥⟨Q j h i , K j h k ⟩V j h k -⟨Q j h ′ i , K j h ′ k ⟩V j h k ∥ ∞ + ∥⟨Q j h ′ i , K j h ′ k ⟩V j h ′ k -⟨Q j h ′ i , K j h ′ k ⟩V j h k ∥ ∞ ≤ ⟨Q j h i , K j h k ⟩ -⟨Q j h ′ i , K j h ′ k ⟩ ∥V j h k ∥ ∞ + ⟨Q j h ′ i , K j h ′ k ⟩ -⟨Q j h ′ i , K j h ′ k ⟩ ∥V j h ′ k -V j h k ∥ ∞ ≤ ⟨Q j h i , K j h k ⟩ -⟨Q j h ′ i , K j h ′ k ⟩ d 2 embd κM + ⟨Q j h ′ i , K j h ′ k ⟩ -⟨Q j h ′ i , K j h ′ k ⟩ d 2 embd κ∥h k -h ′ k ∥ ∞ ≤ ⟨Q j h i , K j h k ⟩ -⟨Q j h ′ i , K j h ′ k ⟩ d 2 embd κM + 2d 4 embd κ 2 M 2 d 2 embd κ∥h k -h ′ k ∥ ∞ ≤4d 6 embd κ 3 M 2 ∥h k -h ′ k ∥ ∞ .Applying this gives∥MHA(H) -MHA(H ′ )∥ ∞ ≤ 4d 6 embd κ 3 M 2 ml∥H -H ′ ∥ ∞ . Then ∥B(H) -B(H ′ )∥ ∞ ≤ w 2LFFN FFN κ LFFN ∥H -H ′ ∥ ∞ + (w 2LFFN FFN κ LFFN + 1)∥MHA(H) -MHA(H ′ )∥ ∞ ≤ w 2LFFN FFN κ LFFN ∥H -H ′ ∥ ∞ + (w 2LFFN FFN κ LFFN + 1)4d 6 embd κ 3 M 2 ml∥H -H ′ ∥ ∞ ≤ 8w 2LFFN FFN κ LFFN d 6 embd κ 3 M 2 ml∥H -H ′ ∥ ∞ .With this we can bound the first term of the multi-block difference∥B L T • ... • B 1 (H) -B L T • ... • B ′ 1 (H)∥ ∞ ≤ 8d 2LFFN+6 embd κ LFFN d 6 embd κ 3 ∥B L T -1 • ... • B 1 (H)∥ 2 ∞ ml * ∥B L T -1 • ... • B 1 (H) -B ′ L T -1 • ... • B ′ 1 (H)∥ ∞ ≤ 8d 2LFFN+6 embd κ LFFN d 6 embd κ 3 ((4d 2LFFN+4 embd κ LFFN+1 ml) L T -1 M ) 2 ml * ∥B L T -1 • ... • B 1 (H) -B ′ L T -1 • ... • B ′ 1 (H)∥ ∞ ≤ 2 7 d 5L T LFFN embd κ 2L T LFFN d 6 embd κ 3 M 2 m L T l L T * ∥B L T -1 • ... • B 1 (H) -B ′ L T -1 • ... • B ′ 1 (H)∥ ∞ .Putting together the estimates on both terms gives ∥B L T • ... • B 1 (H) -B ′ L T • ... • B ′ 1 (H)∥ ∞ ≤ 2 7 d 5L T LFFN embd κ 2L T LFFN d 6 embd κ 3 M 2 m L T l L T * ∥B L T -1 • ... • B 1 (H) -B ′ L T -1 • ... • B ′ 1 (H)∥ ∞ + 4 3L T L FFN M 3 d 18L T LFFN embd κ 6L T LFFN m L T l L T η ≤ 2 7L T L FFN M 3 d 18L T LFFN embd κ 6L T LFFN m L T l L T * ∥B L T -1 • ... • B 1 (H) -B ′ L T -1 • ... • B ′ 1 (H)∥ ∞ which∥B L T • ... • B 1 (H) -B ′ L T • ... • B ′ 1 (H)∥ ∞ ≤ 2 7L T L FFN M 3 d 18L T embd w 18L T LFFN FFN κ 6L T LFFN m L T l L T L T η ≤ 2 7L 2 T L FFN M 3L T d 18L 2 T embd w 18L 2 T LFFN FFN κ 6L 2 T LFFN m L 2 T l L 2 T η.It simply remains to deal with the encoding layer H = PE + E(x). Both E and P E are fixed among architectures and further∥PE + E(x)∥ ∞ = ∥x∥ ∞ + 1 ≤ M + 1. Thus ∥B ′ L T • ... • B ′ 1 • (PE + E ′ (x))∥ ∞ ≤ (4d 2LFFN+4 embd κ LFFN+1 ml) L T M.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∥T(x) -T ′ (x)∥ ∞ = ∥D • B L T • ... • B 1 • (PE + E(x)) -D ′ • B ′ L T • ... • B ′ 1 • (PE + E ′ (x))∥ ∞ ≤ 2 L 2 T +1 L FFN M 3L T d 18L 2 T LFFN embd κ 6L 2 T LFFN m L 2 T l L 2 T η = C T η. 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|θ| = |θ E | + |θ D | + L T i=1 |θ Bi | = d embd + d embd D + L T |θ B | = d embd + d embd D + L T (|θ MHA | + |θ FFN |) = d embd + d embd D + L T (3d 2 embd m + L FFN w 2 FFN ) ≤ 4d 2 embd w 2 FFN D(m + L FFN )L T .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 byN (T , δ, ∥ • ∥ ∞ ) ≤ κ η 4d 2 embd w 2 FFN D(m+LFFN)L T = C T κ δ 4d 2 embd w 2 FFN D(m+LFFN)L T = 2 L 2 T +1 L FFN M 3L T d 18L 2 T embd w 18L 2 T LFFN FFN κ 6L 2 T LFFN m L 2 T l L 2 T δ 4d 2 embd w 2 FFN D(m+LFFN)L T .G ) = σ(⟨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 t1 ) = B(h t1 , h t2 )e i and otherwise A(h t ) = 0 for t ̸ = t 1 . Further we have∥θ A ∥ ∞ = O(d 4embd κ 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 asQ =      Q B 0 0 0 0 Q I 0 0 0 0 0 0 1      , K =      K B 0 0 0 0 K I 0 0 0 0 0 0 -C      , V = e i e T d embdwhere 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 It 2 is a projection onto the unit interaction vector I t2 (the interaction term for h t2 ) and Q I is a dilation and rotation of I t1 onto I t2 i.e. Q I I t1 = CI t2 . We must now compute A(H). For an arbitrary 1 ≤ t ≤ l we can compute the action of A on h t asA(h t ) = l k=1 σ(⟨Qh t , Kh k ⟩)V h k = l k=1 σ(⟨Q B h t , K B h k ⟩ + ⟨Q I I t , K I I k ⟩ -C)e iNow we must case on whether the input token t = t 1 .Case 1: t = t 1 Write A(h t1 ) = l k=1 σ(⟨Q B h t1 , K B h k ⟩ + ⟨Q I I t1 , K I I k ⟩ -C)e iTo 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 t1 , K I I t2 ⟩ = ⟨CI t2 , I t2 ⟩ since by construction Q I I t1 = I t2 and K I is a projection onto I t2 so that K I I t2 = I t2 . This further simplifies as ⟨CI t2 , I t2 ⟩ = C since I t2 is unit length. Thusσ(⟨Q B h t1 , K B h k ⟩ + ⟨Q I I t1 , K I I k ⟩ -C) = σ(⟨Q B h t1 , K B h k ⟩ + C -C) = σ(⟨Q B h t1 , K B h k ⟩)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⟨Q I I t1 , K I I k ⟩ ≤ ∥Q I I t1 ∥ 2 ∥K I I k ∥ 2 = C∥P It 2 I k ∥ 2 < Cwhere we use Cauchy-Schwarz at the first inequality and note that ∥P It 2 I k ∥ 2 < 1 since k ̸ = t 2 for the second inequality. Then, for large enough C, we haveσ(⟨Q B h t1 , K B h k ⟩ + ⟨Q I I t1 , K I I k ⟩ -C) ≤ σ(⟨Q B h t1 , K B h k ⟩ + C∥P It 2 I k ∥ 2 -C) = 0 so we simply must choose C so that ⟨Q B h t1 , K B h k ⟩ + C∥P It 2 I k ∥ 2 -C < 0. Compute ⟨Q B h t1 , K B h k ⟩ + C∥P It 2 I k ∥ 2 -C < 0 ⇐⇒ C > ⟨Q B h t1 , K B h k ⟩ 1 -∥P It 2 I k ∥ 2We can bound⟨Q B h t1 , K B h k ⟩ ≤ ∥Q B h t1 ∥ 2 ∥K B h k ∥ 2 ≤ ∥Q B ∥ 1,1 ∥h t1 ∥ ∞ ∥K B ∥ 1,1 ∥h t1 ∥ ∞ ≤ ∥Q B ∥ ∞,∞ d 2 embd ∥K B ∥ ∞,∞ d 2 embd M 2 ≤ d 4 embd κ 2 M 2 so it remains to control 1 1-∥P I t 2 I k ∥2 when k ̸ = t 2 by upper bounding ∥P It 2 I k ∥ 2 . Compute ∥P It 2 I k ∥ 2 ≤ max 1≤t1̸ =t2≤l ∥P It 1 I t2 ∥ 2 ≤ ∥P I0 I 1 ∥ 2 = ⟨I 0 , I 1 ⟩ = cos( π 2l ) ≤ 1 - 1 2 π 2 2! • 2 2 l 2where we note the projection is maximized by any adjacent interaction terms I t , I t+1 . Then⟨Q B h t1 , K B h k ⟩ 1 -∥P It 2 I k ∥ 2 ≤ d 4 embd κ 2 M 2 π 2 22!•2 2 l 2 = O(d 4 embd κ 2 M 2 l 2 )which gives a bound on C. So for large enough C we can concludeA(h t1 ) = σ(⟨Q B h t1 , K B h k ⟩)It remains to consider the case t ̸ = t 1 .Case 2: t ̸ = t 1 Now we consider A applied to tokens other than h t1 . We haveA(h t ) = l k=1 σ(⟨Q B h t , K B h k ⟩ + ⟨Q I I t , K I I k ⟩ -C)e i 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.: 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⟨Q I I t , K I I t2 ⟩ = ∥Q I I t ∥ 2 ∥K I I t2 ∥ 2 cos(θ t,t2 ) = C cos(θ t,t2 )where θ t,t2 is the angle between Q I I t and K I I k . Crucially, because t ̸ = t 1 , Q I I t ̸ = CI t2 and so θ t,t2 > 0 and cos(θ t,t2 ) < 1. Then we can say⟨Q B h t , K B h t2 ⟩ + ⟨Q I I t , K I I t2 ⟩ -C < 0 ⇐⇒ ⟨Q B h t , K B h t2 ⟩ + C cos(θ t,t2 ) -C < 0 ⇐⇒ C > ⟨Q B h t , K B h t2 ⟩ 1 -cos(θ t,t2 )We upper bound the numerator by O(d 4 embd κ 2 M 2 l 2 ). Note that the Q I rotates by clockwise t1-t2 l π 2 radians. So Q I I t = C(cos( (t+t1-t2)π 2l), sin( (t+t1-t2)π 2l)). Thus we can upper bound cos(θ t,t2 ) ascos(θ t,t2 ) = ⟨I t+t1-t2 , I t2 ⟩ ≤ ⟨I 0 , I 1 ⟩ = cos( π 2l ) ≤ 1 - 1 2 π 2 2! • 2 2 l 2 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 thatFFN(h t ) =              h t t ∈ {1, ..., k}     0 I 1 t I 2 t 1     otherwise 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 W 1 = I d embd +      0 0 Cv 1 Cv 2 0 0 0 Cv 1 Cv 2 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0      , b 1 = 0 W 2 = I d embd +      0 0 -Cv 1 -Cv 2 0 0 0 -Cv 1 -Cv 2 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0      , b 2 = 0.where I d embd is the identity matrix of size d embd × d embd . Thenz = σ(W 1 h t + b 1 ) =       σ(h 1 t + CI t • v) σ(h 2 t + CI t • v) σ(h d embd -2 t ) σ(h d embd -1 t ) σ(h d embd t )       =      σ(h 1 t + CI t • v) σ(h 2 t + CI t • v) I 1 t I 2 t 1      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 + CI t • v) = h i t + CI t • v. Applying the second layer yields W 2 z + b 2 = h t as desired. Further we can compute a bound on C > 0 as |CI t • v| > ∥H∥ ∞ ⇐⇒ C > ∥H∥∞ |It•v| . We compute |I t • v| = |⟨I t , R π 2 I k + I k+1 ∥I k + I k+1 ∥ 2 ⟩| ≥ 1 2 |⟨I k , R π 2 I k ⟩ + ⟨I k , R π 2 I k+1 ⟩| = 1 2 |⟨I k , R π 2 I k+1 ⟩| = | 1 2 ⟨I 0 , R π 2 I 1 ⟩| = 1 2 sin π 2l ≥ 1 4 π 2lwhere the last inequality holds for l ≥ 2.Thus ∥H∥∞ |It•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 formH =    x 1 ... x D 0 D 0 ... ... 0 I 1 ... ... I l 1 ... ... 1    ∈ R d embd ×lwhere l = 2D such that ∥H∥ ∞,∞ ≤ M 2 . Then there exists a transformer block B ∈ B(D, 3) suchB(H) =    x 1 ... x D x 1 -c 1 ... x D -c D 0 ... ... ... ... 0 I 1 ... ... ... ... I l 1 ... ... ... ... 1    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 kernelsQ data i = 1 0 0 0 0 0 0 0 0 1 K data i = 0 0 0 0 1 0 0 0 0 -c i + M .Then, via the Interaction Lemma 3, we can construct A i so that token h D+i interacts with h i such thatA i (h D+i ) = σ(⟨Q data i h D+i , K data i h i ⟩)e 1 = σ(x i -c i + M )e 1 = (x i -c i + M )e 1and A i (h t ) = 0 when t ̸ = D + i. Then the residual multi-headed attention yields MHA(H) + H =    x 1 ... x D x 1 -c 1 + M ... x D -c D + M 0 ... .FFN replace (H) =      -h 1 1 + h 2 1 ... -h 1 l + h 2 l 0 ... 0 0 ... 0 0 ... 0 0 ... 0     In this case we say FFN replace replaces the the first row of H with the second row.Proof of Lemma 6. SetW =     -1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0    Applying W token-wise gives the desired result.). Let m 1 , m 2 , l 1 , l 2 , L FFN1 , L FFN2 , w FFN1 , w FFN2 , d embd ∈ N. Fix transformer blocks B ∈ B(m 1 , L FFN1 , w FFN1 ) with input H 1 ∈ R d embd ×l1 and B 2 ∈ B(m 2 , L FFN2 , w FFN2 ) with input H 2 ∈ R d embd ×l2. 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 existsB 3 ∈ B(m 1 + m 2 , max(L FFN1 , L FFN2 ) + 2, w FFN1 + w FFN2 ) which takes as input H 3 ∈ R d embd ×(l1+l2) such that B 3 (H 3 ) = B 3 [H ′ 1 H ′ 2 ] = [B 1 (H 1 ) B 2 (H 2 )] ∈ R d embd ×(l1+l2) where H ′ i ∈ R d embd ×liis the same as H i except for the interaction terms where we haveI ′ 3,t = (cos( t l1+l2 π 2 ), sin( t l1+l2 π 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,i1 , 1 ≤ i 1 ≤ m 1 ,denote the attention heads of MHA 2 and A 2,i2 , 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 l1+l2 π 2 ), sin( t l1+l2 π 2 )). For l 1 + 1 ≤ t ≤ l 1 + l 2 define h 3,t = h 3,t-l1 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 MHA3 . When 1 ≤ j ≤ m 1 write Q 3,j = Q ′ 1,j where Q ′ 1,jis the same Q 1,j except for a new interaction kernel Q ′I 1,j . Let h 1,sj , s j ∈ {1, ..., l 1 }, be the source token in H 1 interacting with Q I 1,j and h 1,tj ,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,sj onto I 1,tj ). Then define Q ′I 1,j as a rotation onto the new target interaction kernel I 3,tj . 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,tj . Then define the new interaction kernel as a projection onto the new target interaction term I 3,tj . Simply set V 3,j = V 1,j . Now we can check A 3,j (H 3 ) 1,...,l1 = A 1,j (H 1 ) up to the interaction terms. Compute A 3,j (h 3,sj ) = l1+l2 k=1 σ(⟨Q 3,j h 3,sj , K 3,j h 3,k ⟩)V 3,j h 3,k = σ(⟨Q 3,j h 3,sj , K 3,j h 3,tj ⟩)V 3,j h 3,tj = σ(⟨Q 1,j h 1,sj , K 1,j h 1,tj ⟩)V 1,j h 1,tj = A 1,j (h 1,sj ) 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 MHA 3 (H 3 ) = [B 1 (H 1 ) B 2 (H 2 )] 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 kth 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 ( I3,p+I3,p+1 ∥I3,p+I3,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 asW 3,k = W ′ 1,k 0 0 W ′ 2,k ∈ R (wFFN1+wFFN2)×(wFFN1+wFFN2) , b 3,k = b 1,k b 2,k ∈ R wFFN1+wFFN2We must also construct an embedding network W 3,0 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.Code is available at https://github.com/Dahoas/transformer_manifolds_learning 


Decision: This paper develops a theoretical analysis of scaling laws for transforms for intrinsically low-dimensional data and validates the theory via empirical observation of training LLMS on text data. The topic is salient and will be of interest to the community. The reviewers found the work novel and highly non-trivial, opening a window into the differences between transformers and simpler architectures. I recommend acceptance.

Reviews: review: paper_summary: paper_summary: This 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.

The 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.

Finally, 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.
strength_weakness: strength_weakness: The 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:

Despite its novel theoretical contributions, this paper has several significant weaknesses that warrant discussion:

1. **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.

2. **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.

3. **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.

These 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.
questions: questions: 1. 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?

2.  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.
limitations: limitations: Some limitations I discussed above are already in the paper but many are not.
paper_summary: paper_summary: This paper makes a series of contributions:

- 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)
- New Covering Number for Transformers (although I believe this is not presented in the main text)
- Predicting Transformer Empirical Scaling Laws
- Some investigations into how transformer hyperparameters affect the estimated intrinsic dimensionality of text data

 
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.
strength_weakness: strength_weakness: Note: 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.

- “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.
- [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.

I have many questions (below) that may become weaknesses, but I don't wish to penalize the authors if I've misunderstood their work.
questions: questions: - 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?
- 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.
- 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?
- 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 
- Line 71: Is the word “layers” missing? i.e. “requiring only O(log(d)) layers independent of the desired accuracy…”
- Figure 3: Why are there 5 GPT2 model sizes? HuggingFace offers only 4. 
- Figure 3: Why are only 4 Pythia models used? I believe that there are 8.
- 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.
limitations: limitations: N/A
paper_summary: paper_summary: This 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:
1. The paper aims to establish mathematical theories to explain and predict scaling laws observed in transformer models like large language models (LLMs).
2. It presents both statistical estimation and approximation theories for transformers when the input data lies on a low-dimensional manifold.
3. 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.
4. 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.
5. The authors test their theoretical predictions empirically by training LLMs on natural language datasets and find close agreement with observed scaling laws.
6. 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
strength_weakness: strength_weakness: 1. 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.
2. Limited scope: The paper focuses primarily on regression tasks, while many practical applications of transformers (like language modeling) involve more complex objectives.
3. 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.
4. 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.
5. 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.
questions: questions: 1. 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.
2. In Theorem 1 and 2, the parameters of the network should satisfy a specific magnitude assumption, is this assumption significant and practical in applications?
3. 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?
4. Why did you assume the Lipschitz regularity of language modeling objective equals 1? Is there any explanation for it? 
5. The results shown in figure3 seem not good so what’s the significant reason?
6. The order of subscripts seems incorrect in Table 1.
limitations: limitations: 1. The assumptions in the theorem seem not practical.
2. The construction in the proof of estimation theory is hard and not practical in empirical inference.
3. The estimation of the intrinsic dimension is not clear.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: 1. 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.
questions: questions: 1. What is the number of hidden neurons in FFNs?
2. 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$?
limitations: limitations: 1. As the author discussed, this work doesn't consider the training dynamics of the learning problem.
2. 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$?

Comments: title: Rebuttal cont.
Rebuttal cont.
Rebuttal cont.
Rebuttal cont.
Rebuttal cont.
Rebuttal cont.
Response to Authors' Rebuttal (Part 2)
Response to Authors' Rebuttal (Part 1)
Rebuttal cont.
Rebuttal cont.
Rebuttal cont.
comment: We thank the reviewer improving their score and for the helpful discussion!
>  Each component (or sub-sequence of components) could be taken to be time-series measurements or a sequence of pixels [...]

I see. That is an interesting interpretation. I did not get this picture from the paper. Apologies if I missed it but if I didn't miss it, I think it would be good to highlight it.

> our characterization of theory as "predicting" scaling laws [...]

I think "estimating" would be better to use or at least greatly highlight that the prediction made is better but still far from reality so intrinsic dimension is only a part of a puzzle. I guess the bottom line is it's clear that deep learning can only work if the data are "low-dimension" in nature so it's natural that the same would apply to transformer. It is not surprising that things are a straight line on a log-log curve since even PAC bounds would look like a straight line on the log-log curve so how well the slope matches is the only thing that matters for the prediction. In any case, I think a large portion of my questions and concerns have been addressed and I think the theoretical construction is quite interesting and worthy of follow-up so I am raising my score to 6. Thank you for the discussion.
Thanks the authors for their response. My questions are partially addressed and I will maintain my score.
- How do you incorporate sequential structure?

A vector $x \in \mathbb{R}^D$ is very general. Each component (or sub-sequence of components) could be taken to be time-series measurements or a sequence of pixels (as is done in ViT). We then use a position encoding to embed this sequential structure.

- It is ok that a theoretical framework is not accurate but the claim you are making in the paper is "predicting" and "explaining". These are tall orders and I don't think the empirical results sufficiently support your claims. By systematically off I just mean that they are already diverging somewhat significantly for SQL and Tiny story whereas the deviation in Openwebtext is much more reasonable. Scaling law is only interesting precisely because one can extrapolate from it and I don't think we can extrapolate from the theoretical prediction. 

If you strongly object to our characterization of theory as "predicting" scaling laws we would be willing to use a softer characterization (perhaps "estimating"?) in an updated version of the paper.
> Data with an explicit sequential structure is included in our formulation,

How do you incorporate sequential structure? Isn't $x\in \mathbb{R}^D$ is just a single vector.

> we are more accurate for smaller models [...]

It is ok that a theoretical framework is not accurate but the claim you are making in the paper is "predicting" and "explaining". These are tall orders and I don't think the empirical results sufficiently support your claims. By systematically off I just mean that they are already diverging somewhat significantly for SQL and Tiny story whereas the deviation in Openwebtext is much more reasonable. Scaling law is only interesting precisely because one can extrapolate from it and I don't think we can extrapolate from the theoretical prediction. 

> In Figure 4 we demonstrate the estimated ID is stable across a number of model hyperparameter [...]

I agree that in Figure 4 the measurement seems somewhat stable but they are all measured by transformers (so are your references I believe), but my point is that how can one be sure that it is actually the intrinsic dimension of the data? If I use an RNN to measure the intrinsic dimension would I still get the same quantities? I think what this is actually measuring is some notion of simplicity as measured by *transformers trained by SGD*. In any case, this is somewhat of a philosophical question so I won't include it in my decision.
We thank the reviewer for their response.

- Thank you for the clarification. So if I understand correctly, there is no sequential aspect to the data under this model? Wouldn't this make it even less realistic?

Data with an explicit sequential structure is included in our formulation, making our formulation more general. Additionally, we remark transformers are often applied in practice to data which is not inherently sequential (such as images).

- I am focusing on the final loss value, which I believe is comparable since the downstream performances are largely a function of the final loss according to the prevailing scaling hypothesis. What I am concerned about is that the prediction is systematically off, so it is hard for it to convince me that the theory predicts what happens in reality.

To reiterate our previous point, we are simply saying that comparing the raw loss differential between 125M parameter models to the raw loss differential of 34B - 70B parameter models is not a fair comparison. The loss differential for 125M models will be much larger, so one should not judge the accuracy of our theory in this way.

We admit that our predictions at this larger scale of model (34B - 70B) may not be as relatively precise as they are for smaller models. However, we are more accurate for smaller models ($\pm 0.05$ loss differential), and believe that this accuracy meaningfully demonstrates the practical correctness of our theory (when compared to other theoretical predictions). 

Additionally, you say you are concerned our predictions are systemically off. Do you mind being more specific? Do you believe there is a significant source of error even in the smaller model regime?

- But if you train different models I assume you'd get different measurements of ID, so it's clearly not just noise. But this is a minor point due to lack of better alternative, so I would not include this in my final decision.

In Figure 4 we demonstrate the estimated ID is stable across a number of model hyperparameters (model depth, embedding dimension, size, and context length) with the ID changing by at most $\pm 5$ (with the ambient dimension being as large as **1536**). Changing our scaling predictions to use another of the estimated IDs would result in at most $\pm 0.03$ additional error.
Thank you for the reply.

> The pre-embeded input $x \in \mathbb{R}^D$ can be viewed as a sequence components [...]

Thank you for the clarification. So if I understand correctly, there is no sequential aspect to the data under this model? Wouldn't this make it even less realistic?

> This does not allow for a direct comparison [...]

I am focusing on the final loss value, which I believe is comparable since the downstream performances are largely a function of the final loss according to the prevailing scaling hypothesis. What I am concerned about is that the prediction is systematically off, so it is hard for it to convince me that the theory *predicts* what happens in reality.

>  it is trained on the distribution of data we aim to estimate the ID of [...]

But if you train different models I assume you'd get different measurements of ID, so it's clearly not just noise. But this is a minor point due to lack of better alternative, so I would not include this in my final decision.
Thank you for reading our rebuttal and taking the time to respond. We address some of your points below.

- After reading your response, I wonder if there might be some miscommunication here. Sharma and Kaplan state that to estimate ID: "we use the activations from the last token in each sequence to measure the ID, though the ID does not vary significantly across token positions (see figure 10)", whereas your lines 219-220 state " we will estimate the intrinsic dimension of the input data by estimating the intrinsic dimension of token embeddings." I interpreted your sentence to mean that you took either the embedding matrix or the unembedding matrix (which may be the same, if you are tying/sharing the matrices). After rereading the paragraph on lines 219 to 230, I think you are using the word "embedding" when Sharma and Kaplan (as well as I) use "activations".

    Could you please clarify what exactly you are doing? Are you using activations or the embedding vectors that comprise the (un)embedding matrix? If you are using activations, this seems far more sensible to me.
    
Ah sorry about the confusion. We are using activations from the final transformer layer (pre logit transformation). Figure 9 measures the ID resulting from instead using activations from earlier transformer layers.

- I fear we might have miscommunicated here. I wasn't questioning plotting in log-log scale. Rather, as I understand, you appear to have log transformed your data and then plotted the log-transformed data linearly, rather than plotting your data and transforming the axes logarithmically. In matplotlib pseudo-code, you appear to have plotted:

    log_n = np.log10(n)
    log_sge = np.log10(sge)
    plt.plot(log_n, log_sge)

    rather than plotting

    plt.plot(n, sgd)
    plt.xscale("log")
    plt.yscale("log")

    I feel like the latter is far more common, hence why I suggested it. I did label this point as a nit, so if you feel strongly that the former is preferable and tell me so, that's fine.
    
Thanks for clarifying! We agree adjusting the plot scale (as you suggest) is cleaner and will update our paper accordingly.

- This doesn't actually tell me what MLE ID is or why PR is inappropriate here (certainly it's linear, but there's missing next step). However [...] based on this information, MLE ID seems like a more reasonable choice (since I lack knowledge that would favor or disfavor this choice).

A non-linear measure of ID is generally preferred since it can capture non-linear structure that a linear measure of ID will miss. See [1,2] for a demonstration why real-world data has a low non-linear ID but high linear ID. As a simple example, consider a three-dimensional (unit) helix. Inspecting the singular values of the helix would suggest ID = 3 when in reality the helix has ID = 1. MLE is simply a standard choice of non-linear ID measure.

- Ask: If I can make one last request of the authors, I think a "Future Directions" section would be a nice addition to help readers decide which next research problems are worth pursuing.

Thank you for the suggestion. We would be happy to include this in a longer version of the paper.

References:

[1] Nonlinear Dimensionality Reduction by Locally Linear Embedding, https://www.science.org/doi/full/10.1126/science.290.5500.2323?casa_token=yfncfe5drp0AAAAA%3AeCd9RAaQNuvRbWlvXsOnGOlpX0BtboQ4U3k-eQYu0oztePn9TZOXPGoPktxvph4GfvA-bmIKMBSO

[2] A Global Geometric Framework for Nonlinear Dimensionality Reduction, https://www.science.org/doi/full/10.1126/science.290.5500.2319?casa_token=Lo9xDQgvfxAAAAAA%3Ap2bdmZCL2lej6ljrkru8QxldbkbpEjDfSWW3w3ix0D8-kjSw-rI68bZWBb_PczxfiifAZT8vdqc2
- I apologize for the typo in my review where I meant to type 0.02 instead of 0.2, but the point I believe still stands. Let's take the Bigcode-SQL plot in Figure 2. By eyeballing, I'd say the final prediction is roughly -0.1 (empirical fit) and -0.05 (predicted). This translates to roughly .046 difference in validation loss. This is a huge difference. For a comparison, let's look at llama2 [2]. In figure 5, going from 34B to 70B, the perplexity decreased by about less than 0.1, roughly, that's a 0.028 difference in validation loss but represents a huge change in the model performance. As such, I still think the prediction made by the proposed framework is pretty far away from reality. If I misinterpreted the plot in any way, please let me know.



In Figure 2 we are plotting the data scaling exponent $\alpha_D$ (i.e. how does loss decrease as samples increase) for relatively small models (~125M parameters). However, in your example you are comparing these **data scaling** results to Llama's **model scaling** results going from 34B to 70B. This does not allow for a direct comparison. For models in such a large regime the loss differential will naturally be much smaller. It may be the competing loss terms are uniformly small enough that our theoretical framework, as currently applied, is less useful in this setting. However, in the smaller (but equally valid) regime of 125M models, where loss can range from 2.5 - 4.0, we consider an absolute error of $\pm 0.05$ to be accurate and compelling evidence that our theory is able to produce predictions reflecting deep learning phenomena in practice.

We additionally remark that although our theory is not as precisely accurate as empirical measurements, it is more accurate than any other theoretical predictions we are aware of. This is because we importantly leverage the intrinsic dimension of the training distribution which is known to have a strong impact on model performance [5]. We reiterate that our experiments are not meant to suggest replacing current empirical predictions with theoretical ones but rather to validate the correctness of our theory which is the **core** contribution of this paper. Further, the experiments demonstrate for text a recurring observation made in other modalities: intrinsic dimension of data strongly impacts the learning process. We are able to rigorously quantify this impact with our theory.

- I understand that it [our esimtate of ID] is a standard practice but it doesn't make it a reasonable one. You are making a claim about an intrinsic property of the data-generating process yet the actual computation relies on an arbitrarily chosen model. Real data is most likely not generated by a transformer so your measure of intrinsic dimension cannot be correct, so it is not clear to me what I can take away from empirical validation that is based on an erroneous foundation. Note that this is independent of the theoretical framework.

We note the model used to measure ID is **not** arbitrarily chosen: it is trained on the distribution of data we aim to estimate the ID of. In general, the practice of training a model to learn a good representation in order to measure quantities of interest in data is extremely common (see e.g. [3,4,5]). Although the approximation offered may have some noise, often it is still useful in making predictions. Empirically, this has been demonstrated many times [4,5,6,7]. For these reasons, we (as well as many others) believe this method of measuring ID to be an entirely reasonable choice.

Reference

[1] Fantastic Generalization Measures are Nowhere to be Found. Gastpar et al.

[2] llama 2: Open Foundation and Fine-Tuned Chat Model.

[3] DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining, https://arxiv.org/abs/2305.10429

[4]  Scaling Laws from the Data Manifold Dimension, https://jmlr.org/papers/v23/20-1111.html

[5] The Intrinsic Dimension of Images and Its Impact on Learning, https://arxiv.org/abs/2104.08894

[6] The Shape of Learning: Anisotropy and Intrinsic Dimensions in Transformer-Based Models, https://arxiv.org/abs/2311.05928

[7] An Intrinsic Dimension Perspective of Transformers for Sequential Modeling, https://openreview.net/forum?id=0UzYWLzPBjA
Thank you for taking the time to read our rebuttal and respond. We address some of your points below:

- Indeed we use embedding in practice too but in practice, we have a sequence of tokens/inputs and here a single input is being embedded into a sequence of latent tokens. In other words, the modeling assumption is that the whole sequence is a linear projection of a low-dimensional input. I don't see an immediate connection between this setting and real data. Could you explain this more?

The pre-embeded input $x \in \mathbb{R}^D$ can be viewed as a sequence components. This is similar to what vision transformers do: take a vector $x \in \mathbb{R}^D$ as input and embed it as a sequence of pixel tokens. This is the same as our formulation. We would like to clarify that we do **not** assume that the whole sequence is a linear projection of a low-dimensional input. Rather, we apply a linear transformation to **each component** of the input i.e. tokenwise (as is usually done with transformers).


- Respectfully, I disagree that a good explanation of generalization in deep learning can be independent of the optimization procedure, see [1] for a more detailed discussion. Nonetheless, I think the jury is still out on this question so I will not include this in my final decision. Regarding scaling law, I will elaborate on this in the next point.

We agree that an understanding of optimization is essential for a complete understanding of generalization. However, other components (e.g. approximation and statistics) are equally necessary and useful for explaining certain **components** of the generalization error: namely those coming from the bias and variance. This is exactly what our work (and more generally the entre field of non-parametric statistics) does. These error components **directly result** in the model and data scaling laws observed empirically when other sources of error are small. Further, our bound is mini-max optimal (up to logarithmic factors). Using these results, which importantly depend on the intrinsic dimension of the data, we are able to theoretically predict empirical scaling laws more accurately than any other (theoretical works) we are aware of. We would again like to note that this theory is the **main contribution** of our work and therefore should be considered in its evaluation.

In regards to [1], we agree that it is important to consider the algorithm used in generalization error analysis. As far we as we know, the convergence landscape of a multi-layer transformer network is a widely open-question in of itself. However, it does not negate the importance of the theory we contribute.
> Participation ratio is a linear measure of intrinsic dimension. Maximum Likelihood Estimation (MLE) [4] measures the intrisic dimension for low-dimensional nonlinear models by applying the principle of maximum likelihood to the distances between close neighbors.

This doesn't actually tell me what MLE ID is or why PR is inappropriate here (certainly it's linear, but there's missing next step). However...

> MLE has been commonly used in prior works [1,2] to estimate the intrinsic dimension of natural data and assess impact on downstream model performance. 

... based on this information, MLE ID seems like a more reasonable choice (since I lack knowledge that would favor or disfavor this choice).

> We used only four models because we were already able to extract an empirical scaling law over two orders of model size. However, if the reviewer feels including the full suite would strengthen our experimental results we would be happy to plot more intermediate model sizes.

I strongly suspect that others will ask you this question, so I recommend doing it, but I now better understand that it probably won't matter as much.

> our ablations in Figure 4 are meant to demonstrate the stability of our measurements of intrinsic dimension under various reaslistic model hyperparameters - not to suggest any asymptotic behavior

Thank you for clarifying this point. I suggest perhaps lightly rephrasing "Overall, we find the estimated ID is fairly stable across each factor." to integrate the clarification you shared with me here.

Based on the authors' rebuttal, I will increase my score to a 6 and decrease my confidence to a 1. To justify why:

- I don't feel competent to assess the novelty or significance of the theoretical contributions (results or proof techniques), which is really the heart of this paper
- the authors addressed my more empirical concerns

**Ask**: If I can make one last request of the authors, I think a "Future Directions" section would be a nice addition to help readers decide which next research problems are worth pursuing.
I thank the authors for their rebuttal. Responding sequentially:

> Again, we would like to emphasize that the main contributions of this paper are **theoretical**.

I understand. Sadly, I'm not well equipped to evaluate the maths, which is why my confidence is so low. I did not bid on this paper, but I will try my best regardless.

> A Lipschitz assumption with $\beta = 1$ implies the variation of the output for the target function is proportional to the variation of the input which is also widely used in literature. In general we are unaware of numerical methods for measuring the Holder index of the target function from samples.

Without knowledge of any problem with this choice or any alternatives, this seems reasonable.

> Estimating the intrinsic dimension via model embeddings is a relatively established practice [1]. Further, in the case of text we have no other choice as existing ID estimation algorithms are designed for continuous data. 

After reading your response, I wonder if there might be some miscommunication here.

Sharma and Kaplan state that to estimate ID: "we use the activations from the last token in each
sequence to measure the ID, though the ID does not vary significantly across token positions
(see figure 10)", whereas your lines 219-220 state " we will estimate the intrinsic dimension of the input data by estimating the intrinsic dimension of token embeddings." I interpreted your sentence to mean that you took either the embedding matrix or the unembedding matrix (which may be the same, if you are tying/sharing the matrices).  After rereading the paragraph on lines 219 to 230, I think you are using the word "embedding" when Sharma and Kaplan (as well as I) use "activations".

Could you please clarify what exactly you are doing? Are you using activations or the embedding vectors that comprise the (un)embedding matrix?

If you are using activations, this seems far more sensible to me.

> [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.

I fear we might have miscommunicated here. I wasn't questioning plotting in log-log scale. Rather, as I understand, you appear to have log transformed your data and then plotted the log-transformed data linearly, rather than plotting your data and transforming the axes logarithmically. In matplotlib pseudo-code, you appear to have plotted:

```
log_n = np.log10(n)
log_sge = np.log10(sge)
plt.plot(log_n, log_sge)
```

rather than plotting

```
plt.plot(n, sgd)
plt.xscale("log")
plt.yscale("log")
```

I feel like the latter is far more common, hence why I suggested it. I did label this point as a nit, so if you feel strongly that the former is preferable and tell me so, that's fine.
Thank you for the detailed response. Some of my questions have been addressed, but many are still unresolved.

> We note that most common real-world tasks including language modeling [...]

Indeed we use embedding in practice too but in practice, we have a sequence of tokens/inputs and here a *single* input is being embedded into a sequence of latent tokens. In other words, the modeling assumption is that the whole sequence is a linear projection of a low-dimensional input. I don't see an immediate connection between this setting and real data. Could you explain this more?

>  Note that our statistical learning theory does not make any claim about how T is achieved [...]

Respectfully, I disagree that a good explanation of generalization in deep learning can be independent of the optimization procedure, see [1] for a more detailed discussion. Nonetheless, I think the jury is still out on this question so I will not include this in my final decision. Regarding scaling law, I will elaborate on this in the next point.

> Our margin of error in Figure 2 when predicting [...]

I apologize for the typo in my review where I meant to type 0.02 instead of 0.2, but the point I believe still stands. Let's take the Bigcode-SQL plot in Figure 2. By eyeballing, I'd say the final prediction is roughly -0.1 (empirical fit) and -0.05 (predicted). This translates to roughly $\exp(-0.05) - \exp(-0.1) = 0.046$ difference in validation loss. This is a huge difference. For a comparison, let's look at llama2 [2]. In figure 5, going from 34B to 70B, the *perplexity* decreased by about less than 0.1, roughly, that's a 0.028 difference in validation loss but represents a huge change in the model performance. As such, I still think the prediction made by the proposed framework is pretty far away from reality. If I misinterpreted the plot in any way, please let me know.

> Approximation of the intrinsic dimension of data via model embeddings is a relatively established practice [...]

I understand that it is a standard practice but it doesn't make it a reasonable one. You are making a claim about an intrinsic property of the data-generating process yet the actual computation relies on an arbitrarily chosen model. Real data is most likely not generated by a transformer so your measure of intrinsic dimension cannot be correct, so it is not clear to me what I can take away from empirical validation that is based on an erroneous foundation. Note that this is independent of the theoretical framework.


**Reference**

[1] Fantastic Generalization Measures are Nowhere to be Found. Gastpar et al.

[2] llama 2: Open Foundation and Fine-Tuned Chat Model.
**Questions:**

1. 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.

We refer the reviewer our response to the above bullet point. Maximum Likelihood Estimatio (MLE) gives rise to a non-linear measure of intrinsic dimension by applying the principle of maximum likelihood to the distances between close neighbors [6]. It has been commonly used in prior works [2,3] to estimate the intrinsic dimension of natural data and assess impact on downstream model performance. Our theoretical predictions of scaling laws for transformers using these estimates of ID are closer to empirical predictions in comparison with existing work [2]. 



2. In Theorem 1 and 2, the parameters of the network should satisfy a specific magnitude assumption, is this assumption significant and practical in applications?

In practice neural networks are often trained with implicit regularization such as weight-normalization to prevent the weight parameters from blowing up. This makes it easy for such networks to satisfy our assumption on the magnitude of their weights. Note: Our magnitude upper bound is fairly large ($O(dn^{\frac{2}{2\beta +d}}M)$), allowing for plenty of flexibility during model training. 

3. 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?

We find that transformer scaling laws are stable in certain ranges of data and model size as shown in Figures 2 and 3. It is well known that, empirically, scaling laws can break down when the data and model size further increases [7], probably due to optimization error or computational limit, which is not the focus of our paper.

4. Why did you assume the Lipschitz regularity of language modeling objective equals 1? Is there any explanation for it?

In general we expect the target function to have some regularity as otherwise it would be impossible to estimate [1]. Assuming the output varies according to the variation of the input, Holder continuity is the proper and mostly widely considered assumption in nonparametric estimation. A Lipschitz assumption with $\beta = 1$ implies the variation of the output for the target function is proportional to the variation of the input which is widely used in literature. In general we are unaware of any methods for measuring the Holder exponent of the target function from samples.

5. The results shown in figure3 seem not good so what’s the significant reason?

The results in Figure 3 have $\pm 0.03$ i.e. nearly as accurate as Figure 2 with $\pm 0.02$ error. We regard both as fairly accurate predictions of the scaling exponents.

6. The order of subscripts seems incorrect in Table 1.

Thank you for flagging this! The values for $\alpha_N$ and $\alpha_D$ should be flipped.

**Limitations:**

1. The construction in the proof of estimation theory is hard and not practical in empirical inference.

Our construction in the approximation theory shows the universal approximation ability of transformers for Holder functions. These results are significant because **they allow us to quantitatively and precisely control the approximation error $\epsilon$ as a function of the transformer architecture size**. Importantly, this construction to control the approximation does not say anything about the learned parameters of the empirical risk minimizer or the parameters learned during the optimizaion process. We view this as an advantage of our theory, as we do not need to know anything about the empirical risk minimizer's learned parameters to control its generalization error.


2. The estimation of the intrinsic dimension is not clear.

Re-iterating, MLE gives rise to a non-linear measure of intrinsic dimension by applying the principle of maximum likelihood to the distances between close neighbors [6]. It has been commonly used in prior work [2,3] to estimate the intrinsic dimension of natural data and assess impact on downstream model performance.

**References:**

[1]  A Distribution-Free Theory of Nonparametric Regression, https://link.springer.com/book/10.1007/b97848

[2]  Scaling Laws from the Data Manifold Dimension, https://jmlr.org/papers/v23/20-1111.html

[3] The Intrinsic Dimension of Images and Its Impact on Learning, https://arxiv.org/abs/2104.08894

[4] The Shape of Learning: Anisotropy and Intrinsic Dimensions in Transformer-Based Models, https://arxiv.org/abs/2311.05928

[5] An Intrinsic Dimension Perspective of Transformers for Sequential Modeling, https://openreview.net/forum?id=0UzYWLzPBjA

[6] Maximum Likelihood Estimation of Intrinsic Dimension, https://www.stat.berkeley.edu/~bickel/mldim.pdf

[7] Scaling Laws for Neural Language Models, https://arxiv.org/abs/2001.08361
4. 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

Participation ratio is a linear measure of intrinsic dimension. Maximum Likelihood Estimation (MLE) [4] measures the intrisic dimension for low-dimensional nonlinear models by applying the principle of maximum likelihood to the distances between close neighbors. MLE has been commonly used in prior works [1,2] to estimate the intrinsic dimension of natural data and assess impact on downstream model performance. 

4. Figure 3: Why are there 5 GPT2 model sizes? HuggingFace offers only 4.

For GPT2 we reported the results published in [5]. 

5. Figure 3: Why are only 4 Pythia models used? I believe that there are 8.

We used only four models because we were already able to extract an empirical scaling law over two  orders of model size. However, if the reviewer feels including the full suite would strengthen our experimental results we would be happy to plot more intermediate model sizes.

6. 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.

We believe that, in science, the role of theory is not just to predict unobserved events but to provide **rigorous explanation** for observed phenomena. We would describe the latter situation as the theory **correctly predicting** the observed phenomena. Concretely, in our case, we refer to the exponents $\alpha_D, \alpha_N$ as being *predicted* by our theory as a function of input quantities $(\beta, d)$. As you point out, we make three previously unreported predictions in Figure 2. These are **new** in the sense that scaling laws on these datasets have not been previously reported in the literature. Additionally, to our knowledge, scaling behavior for the Pythia suite has also not been previously analyzed. Practically speaking, it would be very difficult to produce a new dataset/model substantially different from what is currently used in industry as, from the scaling perspective, this would require the collection of a tens of trillions token dataset or the training of a hundred-billion parameter model. Our goal, as a theory paper, is instead to rigorously explain existing phenomena and validate with novel, modestly-sized experiments on existing data. Further, these experiments are fully reproducible, as all code and data is open-source.

**References:**

[1]  Scaling Laws from the Data Manifold Dimension, https://jmlr.org/papers/v23/20-1111.html

[2] The Intrinsic Dimension of Images and Its Impact on Learning, https://arxiv.org/abs/2104.08894

[3] A Distribution-Free Theory of Nonparametric Regression, https://link.springer.com/book/10.1007/b97848

[4] Maximum Likelihood Estimation of Intrinsic Dimension, https://www.stat.berkeley.edu/~bickel/mldim.pdf

[5] Scaling Laws for Neural Language Models, https://arxiv.org/abs/2001.08361
**Questions:**

1. 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?

Approximation of the intrinsic dimension of data via model embeddings is a relatively established practice [4,5,6]. In the case of text we have no other choice as existing ID estimation algorithms are designed for continuous data. However, this inapplicability should not be taken to mean that textual data does not have low-dimensional structure. While using embeddings to measure this low-dimensional structure may introduce some noise, we expect a good model (such as an LLM) to have learned good enough representations to preserve the majority of this structure. We do ablations at the end of the paper to examine how sensitive this estimation is to various hyperparameters and find it is relatively stable. While embeddings from any layer could be chosen, we choose the final layer to stay consistent with prior work [4,5,6].

2. 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.

Via our approximation theory, the number of parameters in a transformer is 

\begin{align*}
    O(L_T\cdot(m\cdot d_{embd}+L_{FFN}\cdot d_{embd})) &= O(\log(d)(\log(\epsilon^{-1})+d\epsilon^{-\frac{d}{\beta}})) \\
    &= O(\log(d)d\epsilon^{-\frac{d}{\beta}})
\end{align*}

The number of parameters in a FFN with the same target $\epsilon$ is
\begin{align*}
    O(Depth \cdot Width) = O(\log(\epsilon^{-1})
    \log(d)d\epsilon^{-\frac{d}{\beta}})
\end{align*} where the depth depends logarithmically on $\epsilon^{-1}$ [7]. So the transformer does indeed only require a factor of $\log(\epsilon^{-1})$ fewer parameters than the FFN [7]. Additionally, most of the parameters in our proposed transformer architecture class come from the attention heads. Computationally speaking, this is desirable as the attention heads can be more easily parallelized than sequential feed-forward layers.

**References:**

[1] Scalable Diffusion Models with Transformers, https://arxiv.org/abs/2212.09748

[2] A Distribution-Free Theory of Nonparametric Regression, https://link.springer.com/book/10.1007/b97848

[3] Scaling Laws for Neural Language Models, https://arxiv.org/abs/2001.08361

[4]  Scaling Laws from the Data Manifold Dimension, https://jmlr.org/papers/v23/20-1111.html

[5] The Shape of Learning: Anisotropy and Intrinsic Dimensions in Transformer-Based Models, https://arxiv.org/abs/2311.05928

[6] An Intrinsic Dimension Perspective of Transformers for Sequential Modeling, https://openreview.net/forum?id=0UzYWLzPBjA

[7] Error bounds for approximations with deep ReLU networks, https://arxiv.org/abs/1610.01145