PaperlayerView exact revision on arXiv

Source-linked AI summary

Neural Scaling Laws of Transformer Operator Network

Haoran Yan, Zhongjie Shi, Yuanzhe Xi, Peng Chen, Wenjing Liao

arXiv:2609.33533v1stat.MLcs.LG

TL;DR

Transformer-based operator learning has shown empirical scaling, but quantitative theories of its approximation and generalization errors have been missing. The paper develops bounds and a softmax-attention-aligned approximation construction, and reports power-law data scaling in experiments whose rate varies with input intrinsic dimension.

  • Problem

    Transformer operator-learning work had not quantified approximation error versus network size or analyzed generalization error, despite empirical indications of scaling with training data.

    Source: §1 · p. 2

  • Method

    The analysis constructs a two-level Softmax POU approximation, realizes it with a two-block Transformer, and derives approximation and generalization bounds.

    Source: §1 · pp. 2–3, §3.3 · p. 7

  • Results

    Burgers’ and KdV experiments exhibit the predicted quantitative data scaling law, with empirical decay exponents smaller for higher-dimensional input families.

    Source: §5 · p. 9, §6 · p. 10

  • Takeaways & Limitations

    The generalization scaling exponent reflects input intrinsic dimension, domain dimensions, and function and operator regularity; the output query count reduces the noise contribution.

    Source: §3.3 · p. 7

    The scaling laws are constructive upper bounds, not characterizations of optimal exponents or the performance of every trained architecture.

    Source: §6 · p. 10

Abstract

from arXiv · show

Transformers have emerged as powerful architectures for learning solution operators of physical systems. Empirically the prediction error has been observed to decrease when the data size and model size increase, suggesting neural scaling behavior. Yet a theoretical understanding of such scaling laws for transformer-based operator learning remains limited. In this work, we develop a theoretical framework for characterizing the approximation and generalization errors of transformer-based operator learning. Our analysis builds on a local-to-global approximation principle that is naturally aligned with the softmax attention mechanism and yields discretization-invariant output functions. On approximation theory, we derive a universal approximation error of transformer-based operator learning for H\"older-regular operators. On generalization theory, we establish a power scaling law between the prediction error and the training data size. The rate of convergence represented by the scaling exponent explicitly reflects the dimensions of the input and output domains, the regularity of the underlying functions and operators, and crucially, the intrinsic dimension of the input function class. By exploiting this intrinsic low-dimensional structure, our analysis yields a power-law generalization rate for operator learning, in contrast to the logarithmic-type power-law rates appearing in existing analyses of operator learning with feedforward neural networks. Numerical experiments validate the predicted power-law scaling and confirm that the convergence rate varies systematically with the intrinsic dimension of the input function class.

1 INTRODUCTION

The paper addresses the lack of quantitative approximation and generalization scaling results for transformer-based operator learning. It develops bounds and a two-level Softmax POU construction, yielding power scaling laws that account for function-space complexity and regularity.

1 INTRODUCTION

  • Existing transformer operator-learning analyses lack quantitative approximation scaling with network size and have not analyzed generalization error.The paper frames its question as how error scales with model size and data size.Source: §1 · p. 2
  • The framework assumes mild Hölder continuity and measures input-function-class complexity through its L2 covering number.The paper notes that this characterization fits applications where input functions depend on a small number of physical parameters.Source: §1 · p. 2
  • The paper derives a model-size approximation rate for Hӧlder-continuous operators using a constructed two-block softmax Transformer.The rate depends on operator and function regularity and input function-space complexity.Source: §1 · p. 2
  • For empirical risk minimization, the paper establishes a power-law generalization bound in training-data size.The stated squared-error rate is n^(-2/(2+s)).Source: §1 · p. 2
  • Accounting for low-dimensional input-function complexity yields power scaling, unlike the log-power bounds described for analyses that omit this structure.The bounds also quantify effects of regularity, discretization, model size, training-data size, and output noise.Source: §1 · p. 2
  • A two-level Softmax POU—over input functions and output locations—supports a local-to-global approximation realized by a two-block Transformer.The construction aligns with softmax attention and yields discretization-invariant outputs; point-wise feedforward layers perform exact linear feature assembly.Source: §1 · p. 3

2 OPERATOR LEARNING PROBLEM AND TRANSFORMER ARCHITECTURE

The setup learns a nonlinear operator from discretized input functions and pointwise output queries. The architecture encodes observations and a query as tokens, processes them with encoder blocks, and produces a scalar prediction.

2.1 PROBLEM SETUP

  • Training inputs are sampled independently, observed on a shared input grid, and paired with independently sampled output queries and noisy operator values.The output noise is independent, mean-zero, and sub-Gaussian.Source: §2.1 · p. 3

2.2 TRANSFORMER ARCHITECTURE

  • The Transformer sequence contains tokens for input locations and function values, plus one token for the output query.A token-wise affine embedding and fixed structural-positional encoding form the initial hidden representation; encoder blocks combine self-attention and point-wise feedforward layers before a linear readout.Source: §2.2 · p. 4

3 MAIN RESULTS

The analysis establishes approximation and generalization scaling laws for transformer operator learning under assumptions on operator stability and input/output regularity. The rates reflect function-space dimensions, regularity, and the input class’s intrinsic dimension.

3 MAIN RESULTS

  • The results establish quantitative approximation and generalization bounds for Transformer operator learning, yielding model-size and training-data scaling laws.The generalization result is for a clipped empirical risk minimizer.Source: §3 · p. 5, §3.3 · pp. 6–7

3.1 ASSUMPTIONS

  • The analysis assumes an operator stable under L2 input/output metrics, regular input and output functions, and an input class with a covering number controlled by intrinsic dimension d_U.The covering condition represents low-dimensional structure in the input function class.Source: §3.1 · p. 5

3.2 APPROXIMATION THEORY AND MODEL SCALING LAW

  • Approximation error scales as N_total^-1/s up to logarithmic factors, with the exponent reflecting intrinsic input dimension, domain dimensions, and function/operator regularity.The network uses two Transformer blocks; the result assumes the stated regularity and sequence-length compatibility conditions.Source: §3.2 · pp. 5–6

3.3 GENERALIZATION THEORY AND DATA SCALING LAW

  • The mean squared generalization error scales as n^-2/(2+s), up to logarithmic factors, with the exponent reflecting intrinsic dimension, domain dimensions, and function/operator regularity.The output query size n_y reduces the noise contribution through a bound depending on σ^2/n_y.Source: §3.3 · p. 7

4 PROOF SKETCH

The proof sketches establish approximation by first constructing a two-level Softmax POU oracle and realizing it with a two-block Transformer, then derive data scaling by balancing estimation variance and approximation bias.

4.1 PROOF SKETCH OF THEOREM 1.

  • Two Softmax POUs localize the input function over function-space anchors and the output query over output-domain anchors, yielding an oracle approximation.The oracle error is controlled through localization and input-discretization bounds that converge as the covers and discretization are refined.Source: §4.1 · pp. 7–8
  • Figure 2 illustrates how input-space localization and output-domain localization combine into the oracle approximation.The next proof step shows that a two-block Transformer realizes this construction.Source: Figure 2 · p. 8, §4.1 · p. 8
  • The Transformer realization pairs input and output anchors, assembles their joint localization logits, and uses Softmax attention to aggregate anchor values.The two attention layers implement the localizations and aggregation; the second feedforward layer is the identity and the readout extracts the scalar output.Source: §4.1 · p. 8

4.2 PROOF SKETCH OF THEOREM 2.

  • Theorem 2’s proof bounds estimation variance with a covering-number argument and controls bias using Theorem 1’s approximation result.Combining the bounds and balancing their terms yields the data scaling law.Source: §4.2 · pp. 8–9
  • Figure 3 displays empirical data-scaling curves for Burgers’ and KdV, with fitted power laws overlaid on test-MSE measurements.The caption specifies geometric-mean test MSE over five repeats and error bars of one sample standard deviation of log MSE.Source: Figure 3 · p. 9

5 EXPERIMENTS

Across Burgers’ and KdV, the higher-dimensional input families show smaller empirical data-decay exponents, consistent with the theorem’s predicted ordering.

5 EXPERIMENTS

  • The experiments test whether increasing input-class intrinsic dimension is associated with slower empirical data scaling for dissipative Burgers’ and dispersive KdV solution operators.Both tasks learn the final-time solution operator from the initial condition.Source: §5 · p. 9
  • For Burgers’, the input families have dU ∈ {2, 8}; for KdV, exact multi-soliton families have dU ∈ {4, 8}.Training uses independently sampled function counts from 32 to 1024, with the architecture and training protocol fixed within each problem.Source: §5 · p. 9
  • Burgers’ has p = 2.83 at dU = 2 and p = 2.51 at dU = 8, while KdV has p = 2.10 at dU = 4 and p = 1.47 at dU = 8.Both equations show smaller empirical decay exponents for the higher-dimensional input family; the exponents are not quantitative estimates of the theoretical upper-bound rate.Source: §5 · p. 9
  • Figure 3 shows power-law decay over n = 32–1024: Burgers’ has exponents 2.83 (dU = 2) and 2.51 (dU = 8), while KdV has 2.10 (dU = 4) and 1.47 (dU = 8).Source: §5 · p. 9

6 CONCLUSION

The paper concludes that its two-level Softmax POU analysis yields approximation and data scaling laws for Transformer-based operator learning, with experiments on Burgers’ and KdV showing dependence on input intrinsic dimension.

6 CONCLUSION

  • The analysis establishes approximation and generalization theories through a two-level Softmax POU construction and its explicit Transformer realization.The scaling laws capture operator regularity and the regularity and complexity of input and output function spaces.Source: §6 · p. 9
  • Burgers’ and KdV experiments exhibit the predicted quantitative data scaling law and its dependence on input intrinsic dimension.The paper limits this conclusion to the reported experimental settings.Source: §6 · p. 10
  • The scaling laws are constructive upper bounds, not characterizations of optimal exponents or every trained architecture’s performance.The authors identify sharper bounds and effects of architectural choices and optimization as future-work directions.Source: §6 · p. 10

A ADDITIONAL DEFINITIONS, TRANSFORMER ARCHITECTURE, AND EXTENSIONS

  • The appendices collect additional definitions, Transformer architecture details, and an extension of the approximation result to band-limited inputs.The listed definitions include Hölder norms, covering numbers, a uniform norm, and the Softmax map.Source: A ADDITIONAL DEFINITIONS, TRANSFORMER… · p. 13, §A.1 · p. 13

A.2 TRANSFORMER ARCHITECTURE

The appendix supplies formal Transformer architecture definitions, explains how band-limited inputs remove the discretization contribution to the scaling exponent, and outlines how the approximation proof combines its components.

A.2 TRANSFORMER ARCHITECTURE

  • The architecture appendix specifies the multi-head self-attention and point-wise feedforward maps used in the analysis.It also introduces the attention score matrix, column-wise Softmax attention matrix, and attention head.Source: §A.2 · p. 13

A.3 EXTENSION TO BAND-LIMITED INPUT FUNCTIONS

  • For band-limited inputs, a fixed number of samples recovers the required functional inner products exactly, so input discretization no longer contributes to the polynomial scaling exponent.The stated corollary assumes periodic inputs with fixed Fourier support; the input grid can be chosen independently of the target accuracy.Source: §A.3 · p. 14

B PROOF OF THEOREM 1: APPROXIMATION AND MODEL SCALING

  • The approximation proof develops the two-level Softmax POU oracle, establishes its Transformer realization, and combines them to prove the approximation theorem and derive model scaling.Source: B PROOF OF THEOREM 1: APPROXIMATION… · p. 14

B.1 ORACLE APPROXIMATION VIA TWO-LEVEL SOFTMAX POU

The oracle approximation uses Softmax partitions of unity over both the input-function space and output domain, replacing input-function inner products with fixed-grid quadrature. A two-block Transformer realizes the resulting approximation through pair-token construction and Softmax aggregation.

B.1 ORACLE APPROXIMATION VIA TWO-LEVEL SOFTMAX POU

  • Two Softmax partitions localize over the input-function space and output domain, and their combination defines the continuous two-level oracle approximation.The input-space inner products are replaced by deterministic quadrature because the Transformer observes input functions only on a fixed grid.Source: §B.1 · pp. 14–15
  • The total approximation error combines the input-space, output-domain, and quadrature errors; the required input-grid size depends on intrinsic dimension through a logarithmic factor in Mu.The stated grid-size dependence is described for the intrinsic dimension dU.Source: §B.1 · p. 16

B.2 TRANSFORMER REALIZATION OF THE ORACLE APPROXIMATION

  • A two-block Transformer realizes the discrete approximation with input samples, a query point, and a null column arranged in its sequence.The first block forms active columns indexed by pairs of input- and output-space anchors.Source: §B.2 · p. 16
  • Point-wise feedforward layers assemble joint logits from pair-token features, then a single-head attention layer aggregates the active tokens with Softmax weights.The construction targets features and anchor values for each input/output anchor pair.Source: §B.2 · p. 17, §3 · pp. 18–19

3. Parameter bounds. Define

  • The point-wise FFN layer has hidden width 2D and produces the assembled joint-logit representation for active columns.The corresponding construction assigns the null column separately.Source: §3 · p. 18

B.3 PROOF OF THEOREM 1

The proof constructs a Transformer approximation, bounds its implementation error and parameter complexity, then uses a covering-number bound with approximation error to establish the data-scaling result.

B.3 PROOF OF THEOREM 1

  • The construction preprocesses discretized input samples and a query point, builds anchor-pair features and logits, and aggregates them through a second attention block before scalar readout.The initial sequence includes a null column, and the active columns correspond to input/output anchor pairs.Source: §B.3 · pp. 19–20
  • The proof combines implementation and oracle approximation errors, then counts the dense architectural parameters of the constructed Transformer.The parameter count is expressed using D, P, input-grid size nx, and domain dimensions d1 and d2.Source: §B.3 · p. 21

C PROOF OF THEOREM 2: GENERALIZATION AND DATA SCALING

  • The generalization proof combines a Transformer-class covering-number bound for variance with Theorem 1's approximation bound for bias, then balances the terms to derive the data-scaling law.The covering-number result applies to the specified two-block Transformer class with fixed preprocessing and a corresponding clipped class.Source: C PROOF OF THEOREM 2: GENERALIZATION… · p. 24, §C.1 · p. 24

C.2 PROOF OF THEOREM 2

The proof of Theorem 2 combines bounds under its sampling and noise assumptions, while the supporting POU proofs control approximation through near- and far-center estimates in the input and output spaces.

C.2 PROOF OF THEOREM 2

  • Theorem 2's proof combines bounds for two terms to establish the generalization bound under the stated sampling and noise assumptions.The supplied proof then chooses a parameter to obtain the stated statistical rate, subject to its stated compatibility conditions.Source: §C.2 · pp. 25, 27

D.1.1 PROOF OF LEMMA 1

  • The input-space Softmax POU proof uses an L2-cover of the input-function class and bounds near-center error and Softmax mass on far centers.The proof combines these estimates for its resulting bound.Source: §D.1.1 · pp. 28–29

D.1.2 PROOF OF LEMMA 2

  • The output-domain Softmax POU proof uses an output-domain cover and bounds near-center error and Softmax mass on far centers.The proof combines these estimates for its resulting bound.Source: §D.1.2 · pp. 29–30

D.1.3 PROOF OF LEMMA 3

These proof sections develop constructions and bounds for transformer models, including quadrature-based components, feature and output errors, parameter perturbations, preprocessing, and a covering bound for the model class. They also point to supplementary Burgers’ and KdV experiment details.

D.1.3 PROOF OF LEMMA 3

  • A compound tensor-product Gauss–Legendre rule integrates the local Taylor polynomials exactly and uses a total of at most nx quadrature points.The Taylor remainder on each subcube is bounded by CT q^-α ∥h∥Cα(ΩU); zero-weight grid points are appended if needed.Source: §D.1.3 · p. 31

D.1.4 PROOF OF LEMMA 5

  • Lemma 5’s proof decomposes the prediction error into three terms, bounds them using prior lemmas and discrete logits, and combines the bounds uniformly over inputs.It concludes by taking the supremum over u ∈ U to prove (42).Source: §D.1.4 · p. 32

D.2.1 PROOF OF LEMMA 6

  • The preprocessing representation places discretized input and query information in designated coordinates, with remaining coordinates set to zero.The input occupies the first nx columns, the query point y the (nx+1)-st column, and later input coordinates are zero.Source: §D.2.1 · p. 33

D.2.2 PROOF OF LEMMA 7

  • The first MHA layer uses feature-extraction and auxiliary heads to target input, query, indicator, and positional features, with bounds on feature errors and parameters.The construction assigns (nx + 3)P feature-extraction heads, plus two auxiliary heads.Source: §D.2.2 · pp. 34, 36, 38

D.2.3 PROOF OF LEMMA 8

  • A point-wise ReLU FFN realizes the specified linear map, and feature bounds imply the inequalities used to prove the output bound.The proof uses the bound |δA| ≤ εfeat from Lemma 7.Source: §D.2.3 · pp. 39–40

D.2.4 PROOF OF LEMMA 9

  • The second attention layer assigns equal attention weights across target columns because its scores are independent of the target column, and the value projection satisfies MA2 ≤ 1.The scores are sj,j′ = kjqj′ = bΞj, independent of j′.Source: §D.2.4 · pp. 40–41

D.3 PROOF OF THE COVERING NUMBER BOUND

  • The covering-number proof compares networks with bounded parameters by separating parameter error from input error, beginning with their shared preprocessed representation.It assumes ∥Θ∥∞, ∥eΘ∥∞ ≤ Mmax and ∥Θ − eΘ∥∞ ≤ δ.Source: §D.3 · pp. 41–42
  • Because the preprocessing map is fixed, the input error (II) is zero, leaving only parameter error (I).Source: §D.3 · p. 42

U + DMmaxBU

  • The covering argument proceeds through two encoder blocks and a linear readout, then uses a parameter-cube cover; 1-Lipschitz clipping does not increase the covering number.Source: U + DMmaxBU · pp. 43–45

E ADDITIONAL EXPERIMENTAL DETAILS

E.1 COMMON EXPERIMENTAL SETUP

The experiments compare Burgers’ and KdV solution-operator learning across training-set sizes and input-function dimensions, using fixed architectures and held-out test evaluation. The two PDEs use distinct input families and reference-solution procedures.

E.1 COMMON EXPERIMENTAL SETUP

  • Both experiments use the same two-block encoder-only Transformer within each problem, with fixed architecture across sample sizes and intrinsic dimensions.Test MSE is averaged in physical units over 128 held-out functions and 1024 uniformly spaced spatial points; training uses noiseless targets and validation-based early stopping and checkpoint selection.Source: §E.1 · p. 45

E.2 VISCOUS BURGERS’ EQUATION

  • Burgers’ experiments use input families with dU = 2 or 8 and training sets of 32–1024 functions; reference solutions are recomputed using the periodic Cole–Hopf transformation.The reevaluation uses 4096 Fourier points, is validated by spatial refinement and an independent solver, and does not retrain the checkpoints.Source: §E.2 · p. 46

E.3 KDV EQUATION

  • KdV experiments use exact N-soliton solution families with dU = 4 or 8 and training sets of 32–1024 functions; test MSE is computed in physical units after amplitude rescaling.During training, input and output amplitudes are divided by 10, and predictions are multiplied by 10 before test evaluation.Source: §E.3 · p. 47

Related papers

Loading 2609.33533v1…