arXiv CS open access

Neural Quadratic Forms: A Unified Minimal Model for Sudden Learning and Scaling Laws

2026 · arxiv_cs

arXiv CS · Papers · 2026

deep-learning machine-learning neural-networks

Conceptio Open Archive

Analyze & explore related research in Conceptio

Connect this record directly into your research workflow or AI agent. Explore citations across open repositories, query full text, and export structured proof bundles.

Explore in Search →Connect via MCP / API ↗

Neural Quadratic Forms: A Unified Minimal Model for Sudden Learning and Scaling Laws Liu Ziyin1, * , Yizhou Xu2,∗ , Tomaso Poggio1 , Isaac Chuang1 1 Massachusetts Institute of Technology 2 École Polytechnique Fédérale de Lausanne

August 14, 2026

arXiv:2608.13335v1 [cs.LG] 13 Aug 2026

Abstract Neural networks trained by gradient descent on a smooth cost function can nevertheless learn in steps: the cost holds on long plateaus and then drops abruptly. Meanwhile, training losses instead follow smooth power laws. Variants of both behaviors occur in architectures with very different microscopic structures, which is the signature of a few relevant collective variables. We show that a symmetry fixes what those variables are: a network layer is a sum over interchangeable units, so relabeling the units leaves it unchanged; given smoothness and the condition that a unit’s gradient vanish at the origin, symmetry then enforces a universal leading form for the expansion about the near-zero weights present at the start of training, the quadratic Tr[W W ⊺ A(x)], in which every architectural detail is confined to a single “structure matrix” A(x) that we compute for each architecture. Perceptrons, attention layers, mixtures of experts, and convolutions become one model at different A. Its training dynamics then close on the “order parameter” M = W W ⊺ and, whenever the data matrices share an eigenbasis, reduce to a Lotka–Volterra equation whose modes switch on one after another. The smaller the initial weights, the further apart the switch-on times, and the plateaus appear as a singular limit of a smooth flow; when many modes are unresolved the same events merge into a power law in training time whose exponent the theory predicts. We confirm both numerically across training methods and architectures.

Introduction

Neural networks exhibit two apparently incompatible forms of regularity. Along a single training trajectory, the loss can remain nearly constant for long intervals and then fall abruptly as new features are acquired. Across model size, dataset size, or training compute, by contrast, losses often follow smooth power laws over many orders of magnitude, with weak sensitivity to architectural detail [39, 7, 49]. These cross-scale laws are accurate enough to guide large-scale design decisions: compute-optimal scaling can favor a smaller model trained on substantially more data over a much larger, undertrained one [36]. Variants of both behaviors occur across architectures with very different microscopic structure, including multilayer perceptrons, convolutional networks, mixtures of experts, and transformers [62, 46, 75, 43, 37, 40]. Sharp microscopic events together with macroscopic laws that are indifferent to microscopic detail are the familiar signature of a small number of relevant collective variables, and they invite the construction of a minimal model. To a physicist, these behaviors are not unique to neural networks. Smooth deterministic dynamics can produce sharp, reproducible transitions when different modes cross instability thresholds one at a time, as at the onset of convection in a fluid. Long plateaus followed by rapid changes resemble the induction periods of autocatalytic reactions. Moreover, a transition may remain smooth for every finite value of a parameter but become sharp in a limiting regime, just as phase transitions emerge only in the thermodynamic limit. Finally, the weak dependence on microscopic details suggests universality: the large-scale behavior is governed by a small number of collective variables. These physical phenomena can be understood by finding the collective variables and the equation they obey. Such a search naturally begins with a symmetry, and neural architectures supply an unusually clean one. Neural * These two authors contributed equally.

Figure 1: Permutation symmetry of the units allows us to characterize the model with structure matrices A. (a) Perceptrons, attention layers, mixtures of experts and convolutions are all sums over interchangeable components, so a permutation leaves the represented function unchanged. (b) Training begins with all weights near zero, so each module may be Taylor expanded about that point. Given smoothness and the condition that a component’s gradient vanish whenever that component’s own weights do, symmetry then enforces the same leading quadratic form in every case (Theorem 1), and the architecture survives only in the structure matrix A(x). Here x is a single input to the module. networks are canonically assembled from repeated, exchangeable components: hidden units in a perceptron, heads in an attention layer, experts in a mixture, channels in a convolution. Relabeling those components leaves the represented function unchanged, so a module of width d is invariant under the symmetric group Sd acting on its components, in the same way that a collection of identical particles is invariant under exchange. The analogy is more than verbal: the training of two-layer networks has been mapped directly onto the mean-field dynamics of interacting particles [52], permutation symmetry has been proposed as the operational definition of a neuron [77], and the breaking and restoration of parameter symmetries has been argued to organize the order in which features are acquired [79, 80]. Symmetry of this kind is the standard point of departure for a Landau construction [44, 32, 21]: identify the symmetry group, identify a small parameter that holds the system near a distinguished point of configuration space, expand about that point, and retain only the lowest-order terms the symmetry permits. Here the group is Sd , the distinguished point is the origin of parameter space, and the small parameter is the initialization scale ϵ. Microscopic detail should then survive only in the phenomenological coefficients of the expansion, which is the sense in which universality across architectures is to be expected and which is the logic underlying the statistical-mechanical tradition in learning theory [63, 25, 74, 8, 20]. Carrying out this program produces a single expression, summarized in Figure 1. We prove (Theorem 1) that any module fx (w1 , . . . , wd ) that is Sd -symmetric in its components [Figure 1(a)], is three times continuously differentiable, and satisfies a gradient condition, stated precisely below, obeys fx (W ) − fx (0) = µ⊺ g(x) + Tr[W W ⊺ A(x)] + O (∥W ∥3 ) .

(1)

Here x is a single input, W = (w1 , . . . , wd ) ∈ Rp×d collects the d components, each carrying p parameters, µ = ∑di=1 wi , and the symmetric matrix A(x) ∈ Rp×p is fixed by the architecture and by x alone; we call Eq. (1) the neural quadratic form (NQF) and A(x) the structure matrix. It turns out that the term µ⊺ g(x) vanishes for most of the practical architectures, and so the NQF can often be written with only the second-moment term. Everything architecture-specific now sits in A(x), as a sparsity pattern and a set of couplings [Figure 1(b)]. Attention is the striking case: only its value and readout blocks appear at this order, the query and key matrices entering first at quartic order. What controls the model is likewise not the individual components but the order parameter M = W W ⊺ , whose size is set by p. Quadratic parameterizations are not themselves new as proxies for nonconvex learning [33, 5, 18, 12, 48, 23]. Matrix sensing, phase retrieval, diagonal networks and quadratic networks have each been analyzed on their own terms, to explain implicit low-rank bias, spectral initialization, or feature growth; exact analyses of deep linear networks 2

display the same mode-wise plateaus under gradient descent [62, 3]; and recent work argues that quadratic models remain quantitatively accurate even for large language models [53]. What Eq. (1) adds is that these are not a family of analogous models but one model, and that the coupling A(x) is computed from the architecture [Figure 1(a,b)]. Seventeen models previously studied as separate solvable proxies, including those above, are instances of Eq. (1) at particular A. A practical consequence is that A(x) can be evaluated for a proposed layer before that layer is ever trained. Several distinct regularities travel under the name of a neural scaling law, and they are worth separating before any theory is compared. Power laws are reported in the loss against optimization time within a single run; in the trained loss against model or dataset size; and along a compute-optimal frontier trading the two. What follows is a law of the first kind, for excess loss against rescaled training time, controlled by the same spectral structure that enters theories of the other two. How the trained loss depends on model size, dataset size and compute, and how those resources should be traded against one another, is a separate question we do not take up. Within that scope, the closest theory is that of Bahri et al. [7], who separate a variance-limited regime, where the loss falls as the inverse of model or dataset size, from a resolution-limited one whose exponent is set by the data manifold or the tangent-kernel spectrum. Theirs is a theory of trained loss against scale; ours is of loss against time. Kernel treatments reach power-law learning curves from spectral structure by a related route [15, 17, 49, 14], but they linearize about initialization and freeze the representation [38, 22], suppressing the feature growth that produces abrupt transitions [61, 10]; sudden learning is then treated separately, as saddle-to-saddle transitions for particular targets [1, 54, 56, 11], or as the loss of stability of a trivial symmetric solution [71]. Our question spans the two: can the spectrum behind a macroscopic law emerge from within a feature-learning dynamics, and can that same dynamics produce the abrupt acquisition of individual modes? Answering both questions with one dynamics is tractable only because that dynamics is reducible to a few collective variables. That reduction is our first result. Under stochastic gradient descent, the dynamics of any NQF close on the pair (M, µ) = (∑i wi wi⊺ , ∑i wi ), no matter how many trainable parameters the model contains (Theorem 2). The entire parameter trajectory is therefore represented by these low-dimensional moments; and within the NQF description a module of width d admits a compressed equivalent of width kV + 1, set by the rank of the data, which reproduces its predictions on the training data at every step under a rescaling (Theorem 3). This measures the low-dimensional training dynamics [34, 45, 4]. The second result is that the closed dynamics is in many cases solvable. In the common eigenbasis of the data, the flow of M becomes the generalized Lotka–Volterra equation of population ecology [35, 51, 16], with the eigenvalues zk of M in the role of species abundances, and we solve it in four regimes (Theorems 4–7), each exhibiting a combination of sudden learning and neural scaling law phenomena. This work is organized as follows. Section 2 derives the normal form and Section 3 computes A(x) for standard architectures. Section 4 analyzes the dynamics, and Section 5 solves the dynamics for special cases. Section 6 derives the predictions for sudden learning and neural scaling laws. Section 7 tests them numerically, and Section 9 discusses the limitations and the extension to higher orders.

Neural Quadratic Forms

This sections presents the central result of this paper. We show that permutation symmetry alone gives us a universal norm quadratic expansion, which is the result of Theorem 1. In Section 3 we will show how each neural architecture is different. Notation. Throughout, x denotes a single input to the module: one training example drawn from a set X . The module itself is written fx , a scalar-valued function of the parameters at fixed x; the subscript records that x is held fixed while the parameters vary. Its d components carry p parameters each, collected as w1 , . . . , wd ∈ Rp , so fx ∶ Rp×d → R. Vector-valued outputs are treated in Appendix B.2 and change nothing essential. All notations used in our paper are listed in Table 2. The starting point of the theory is that there is a universally shared mathematical structure across different types of layers and architectural modules (see Section 3). Definition 1 (Permutation symmetry). fx ∶ Rp×d → R satisfies the permutation symmetry (Sd - symmetry) if fx (w1 , . . . , wd ) = fx (wσ(1) , . . . , wσ(d) ) for any permutation σ. Leveraging the terminology from MLPs, each wi can be seen as the weights of a neuron of a hidden layer. It is thus natural to call wi a “neuron,” and we will stick to this terminology. In fact, permutation symmetries have been 3

suggested as the “right” way to precisely define a neuron in deep learning [77]. To be more precise, one could also call wi the “coordinate” of the i-th neuron. Thus, in our work, any subset of the weights that is the minimal unit of permutation symmetry will be called a “neuron.” In this terminology, the weights of a self-attention head are also a neuron. For a simplified presentation of the theory, we assume one further benign condition that is obeyed by almost any neural architecture in use (see Section 3). Property 1. Zero gradient at zero (ZGZ): for any neuron i and any wj≠i , ∇wi fx (w1 , . . . , wd )∣wi =0 = 0. The following theorem shows that for models with permutation symmetry, their Taylor expansions take a highly universal form. In some sense, this result can be seen as a variant of the fundamental theorem of symmetric polynomials [69]. Theorem 1 (Neural Quadratic Forms). Let fx ∶ Rp×d → R be a three times continuously differentiable model with respect to its neurons W = (w1 , . . . , wd ), where wi ∈ Rp . Assume fx satisfies the permutation symmetry and the ZGZ conditions. Then, d

fx (W ) = fx (0) + ∑ Tr[wi wi⊺ A(x)] + O(∥W ∥3 ),

(2)

i=1

where A(x) ∈ Rp×p is a symmetric matrix dependent only on x, and fx (0) is the model evaluated at W = 0. Equivalently, d

lim ϕ(fx (ϕ−1/2 W ) − fx (0)) = ∑ Tr[wi wi⊺ A(x)].

ϕ→∞

(3)

i=1

Remark. We note that the ZGZ condition is nonessential. Removing it leads to essentially the same result but with more complicated notations. We present these results in Appendix B.1. Moreover, in practice, the ZGZ assumption can be replaced by the stronger (but equally common) assumption of having a per-neuron Z2 symmetry, meaning fx (. . . , wi , . . . ) = fx (. . . , −wi , . . . ) for all i. This means the 1st- and 3rd-order tensors in the Taylor expansion are zero, naturally elevating the error bound from O(∥W ∥3 ) to O(∥W ∥4 ). Therefore, we will refer to any model in the form of Eq.(2) as a neural quadratic form (NQF). For example, fx (W ) = c0 + Tr[W W ⊺ A(x)]

(4)

will be a generic NQF with d as its width and A as its “architecture” or “structure.” We will from now on refer to A as the structure matrix. In Appendix B.2 we extend the results to multidimensional outputs. Note also the dimension of M . Both M and A(x) are p × p, where p is the number of parameters of one neuron: the width d having already been summed away. Since the neurons are exchangeable, they may be read as identical particles with a p-dimensional state space, and ∑i wi and M = ∑i wi wi⊺ as the first two moments of their empirical distribution, M playing the role of a density matrix. The next section shows how different forms of A distinguish different architectures. Note that f depends on the neurons only through their second moment M = W W ⊺ , a point we take up in Section 4.

NQF for Different Architectures

In fact, a point often unclear to practitioners is that the existence of permutation symmetries is so universal that essentially any architecture module that has a notion of “width” automatically has the permutation symmetry property. This makes Theorem 1 applicable to almost any neural module one encounters in practice. Thus, in the NQF perspective, it is the structure matrices that determine the learning dynamics of neural networks. See Table 1 for common models that are reducible to an NQF. We now discuss a few examples in more detail. Two-Layer MLP. Consider a two-layer Multi-Layer Perceptron (MLP) with a scalar output fx ∶ Rd×(k+1) → R without bias terms. Let the parameters associated with the i-th hidden neuron be denoted as a single vector wi = [u⊺i , vi ]⊺ , where ui ∈ Rk is the input weight vector and vi ∈ R is the readout (output) weight. The model is given by: d

fx (w1 , . . . , wd ) = ∑ vi ϕ(u⊺i x) i=1

where x ∈ Rk is the input and ϕ ∶ R → R is the activation function. 4

(5)

Table 1: Six representative models that are instances of the neural quadratic form (1), together with A(x) that distinguishes them. There are two ways to read the table: many generic architectures reduce approximately to an NQF, and many previously separate models are special cases of the NQF. The complete list appears in Table 3. Structure matrix A(x) ⎡ ⎤ ′ ⎢0k×k x⎥ ⎥ A(x) = ϕ 2(0) ⎢⎢ ⊺ ⎥ ⎢ x ⎥ 0 ⎣ ⎦ ⎡ ⎤ ′ ⎢ 0 ∑p xp ⎥⎥ m×m A(X) = ϕ 2(0) ⎢⎢ ⎥ ⎢∑p x⊺p 0 ⎥⎦ ⎣ ⎡ ⎤ ⎢02dk D ⎥ 0 0 ⎢ ⎥ ⎥ 1⎢ ⎥ A(X) = 2 ⎢ 0 0 x ⊗ I avg d v⎥ ⎢ ⎢ ⎥ ⊺ ⎢ 0 ⎥ xavg ⊗ Idv 0 ⎣ ⎦ ⎡ ⎤ ⊺ ′ ⎢ ⎥ 0 x G(x) ⎥, A(x) = ψ 2(0) ⎢⎢ ⎥ ⎢G(x) x⊺ ⎥ 0 ⎣ ⎦ G ∶= ∇Θ E(x; 0)

Model

NQF notation

Refs.

Two-layer MLP

fx (W ) = ∑di=1 vi ϕ(u⊺i x)

Single-layer CNN

fx (W ) = d P ∑i=1 vi ∑p=1 ϕ(ki⊺ xp )

Multi-head attention

fx (W ) = d ∑i=1 vi⊺ WiV X ⋅ softmax(⋅)

Mixture of experts

fx (W ) = d ∑i=1 ψ(ri⊺ x) E(x; Θi )

Phase retrieval

ŷa = a⊺ W W ⊺ a

Aa = aa⊺

[18, 47]

Diagonal linear network

fx (u) = ∑k xk u2k /4

Ax = 14 Diag(x)

[12, 59]

[29, 61] [42, 24]

[67]

[64, 57]

Proposition 1. If ϕ is three times continuously differentiable and satisfies ϕ(0) = 0, then the model (5) satisfies the assumptions in Theorem 1 with ϕ′ (0) 0k×k x A(x) = [ ⊺ ]. (6) x 0 2 Self-Attention. Consider a Multi-Head Attention (MHA) model computing a scalar output from an input query token x ∈ RD and a context matrix X ∈ RD×N (where N is the sequence length). Let the parameters associated with the i-th attention head be grouped into a single parameter vector wi = vec(WiQ , WiK , WiV , vi ), where WiQ , WiK ∈ Rdk ×D are the query and key matrices, WiV ∈ Rdv ×D is the value matrix, and vi ∈ Rdv is the readout vector for the scalar output. Assuming no bias terms, the model is formulated as: d

fx (w1 , . . . , wd ) = ∑ vi⊺ WiV X ⋅ softmax ( i=1

X ⊺ (WiK )⊺ WiQ x √ ) dk

(7)

where softmax is applied over the N dimension. Proposition 2. The MHA model (7) satisfies the assumption in Theorem 1. The matrix A(X) takes the following block form: ⎡0 0 ⎤ 0 0 ⎢ ⎥ ⎥ 1 ⎢⎢0 0 0 0 ⎥ ⎥ A(X) = ⎢ (8) ⎢ 0 xavg ⊗ Idv ⎥⎥ 2 ⎢0 0 ⎢0 0 x⊺avg ⊗ Idv ⎥ 0 ⎣ ⎦ where xavg = N1 X1N ∈ RD , 1N ∈ RN is a vector of all ones, and ⊗ denotes the Kronecker product. This means that to leading order, what determines the behavior of a self-attention head are the value and output matrices, and the self-attention behaves like an averaging operator. In fact, the query and key weights appear only at quartic order.

Query-Key-Only Head. Now, consider a variant of MHA where the readout vector and value matrix are fixed and merged into a constant vector c ∈ RD . While this is rarely the case in practice, it is actually a common theoretical model of the learning dynamics of attention (e.g., [6, 50], or in which the readout and query-key matrices are trained separately [70]). The model evaluates a scalar output as: d

fx (w1 , . . . , wd ) = ∑ c⊺ X ⋅ softmax ( i=1

X ⊺ (WiK )⊺ WiQ x √ ) dk

(9)

Let the parameters for the i-th head be the query and key matrices wi = [vec(WiQ )⊺ , vec(WiK )⊺ ]⊺ ∈ R2dk D . Proposition 3. Model (9) satisfies all assumptions in Theorem 1. Furthermore, A(X) takes the block form: 1 0 A(X) = √ [ ⊺ ux ⊗ Idk 2 dk

xu⊺ ⊗ Idk ] 0

(10)

where u = ΣX c ∈ RD , and ΣX = N1 XX ⊺ − ( N1 X1N )( N1 X1N )⊺ ∈ RD×D is the sample covariance matrix. Single-head attention logits. Finally, we can also consider the single-head attention model, which can, interestingly, also be regarded as an independent NQF because the query and key weight matrices contain the permutation symmetry as a subgroup. Proposition 4. By treating the i-th rows of W Q and W K as wi ∈ R2D , the single-head attention model fx (w1 , . . . , wdk ) = c⊺ X ⋅ softmax (

X ⊺ (W K )⊺ W Q x √ ) dk

(11)

satisfies all assumptions in Theorem 1. Furthermore, A(X) ∈ R2D×2D takes the block form: 1 0 A(X) = √ [ ⊺ ux 2 dk

xu⊺ ] 0

(12)

where u = ΣX c ∈ RD , and ΣX is the sample covariance matrix. Mixture of Experts. Consider a Mixture of Experts (MoE) model computing a scalar output fx from an input x ∈ Rk . Let the parameters associated with the i-th expert branch be grouped into a single vector wi = [ri⊺ , vec(Θi )⊺ ]⊺ , where ri ∈ Rk is the routing (gating) weight vector, and Θi represents the internal parameters of the expert network E(x; Θi ). We employ an independent gating mechanism [57]: d

fx (w1 , . . . , wd ) = ∑ ψ(ri⊺ x)E(x; Θi )

(13)

i=1

where ψ ∶ R → R is a smooth gating activation function (e.g., Tanh or GeLU) satisfying ψ(0) = 0. Each expert E(x; Θi ) is a smooth neural network satisfying E(x; 0) = 0. Proposition 5. Assuming the gating function ψ and the expert networks E are at least three times continuously differentiable, the MoE model (13) satisfies the assumption in Theorem 1 with A(x) =

1 0 [ ′ ⊺ 2 ψ (0)∇Θi E(x; 0)x

ψ ′ (0)x∇Θi E(x; 0)⊺ ]. 0

(14)

Convolutional Neural Networks. Consider a single-layer Convolutional Neural Network (CNN) computing a scalar output fx from an input feature map X. Let the parameters associated with the i-th convolutional channel be grouped into a single vector wi = [ki⊺ , vi ]⊺ , where ki ∈ Rm represents the flattened weights of the convolutional filter, and vi ∈ R is the readout weight applied after pooling. Assuming the network uses global sum-pooling over the P spatial patches of the input (denoted as xp ∈ Rm for p = 1, . . . , P ), the model is formulated as: d

fx (w1 , . . . , wd ) = ∑ vi ∑ ϕ(ki⊺ xp ), i=1

where ϕ ∶ R → R is the activation function. 6

p=1

(15)

Proposition 6. If ϕ is three times continuously differentiable and satisfies ϕ(0) = 0, then the CNN (15) satisfies the assumption in Theorem 1 with ϕ′ (0) 0m×m ∑P p=1 xp A(X) = [ P ⊺ ]. (16) 0 2 ∑p=1 xp

Notions of Layer and Representation. While the form Tr[W AW ⊺ ] is universal, the matrix A always takes a sparse and strongly off-diagonal form. In particular, for any neural layer above, A=

1 0 C(x) [ ⊺ ], 0 2 C (x)

(17)

where C ∈ Rd1 ×d2 is an arbitrary rectangular matrix with suitable dimensions. This means that one can also divide W Z into corresponding blocks: W = [ 1 ], such that Z2 Z Z⊺ M = W W ⊺ = [ 1 1⊺ Z2 Z1

Z1 Z2⊺ ] Z2 Z2⊺

(18)

This implies that the output of the model only depends on Z2 Z1⊺ : fx = Tr[Z2 Z1⊺ C]. In the case of a vector C = x, Z1 is also a vector and so fx = Z2 Z1⊺ x, a two-layer linear network. Therefore, one can abstractly think of Z1⊺ as the first-layer weight and Z2 as the second-layer weight. The representation also becomes definable. The first-layer latent representation of C(x) is thus h(x) = Z1⊺ C(x). (19) For a structured model, such as above, it is possible to talk about a latent representation and define the notion of layers. Whenever the matrix A can be written in this form, we refer to the model as a “feedforward” model.

Learning Dynamics of the NQF

Sections 2 and 3 concern a model when expanded, and in this section, we study the learning dynamics of these models. A key question this section answers is how many numbers are needed to follow training. The answer is that the number does not grow with the model: whatever the width, the trajectory of an NQF is a flow on the pair (M, µ). This is stated as Theorem 2. One consequence follows immediately: learning is highly redundant, in that a model can be compressed with no change to its trajectory at all (Theorem 3). We work throughout with the most general form of NQF without the ZGZ condition (See Appendix B.1). If we denote M ∶= ∑di=1 wi wi⊺ and µ ∶= ∑di=1 wi , we can write a general NQF as fx (W ) = fx (0) + g(x)⊺ µ + µ⊺ B(x)µ + Tr[M A(x)].

(20)

Without loss of generality, A and B matrices can always be regarded as symmetric matrices. Theorem 2 (Master Theorem for NQF). Let A(x), B(x) be symmetric. Under SGD ∆W = −η ∑ ∇W L(fx (W ))

(21)

x∈B

with identical data sampling, two neural quadratic models have Ma (t) = Mb (t), µa (t) = µb (t) at any time t if at initialization: Ma (0) = Mb (0), µa (0) = µb (0). (22) The learning dynamics are completely determined by M , µ: ∆µ = −η(d ⋅ v + Hµ)

(23)

∆M = −η(vµ⊺ + µv ⊺ + HM + M H) + η 2 (d ⋅ vv ⊺ + vµ⊺ H + Hµv ⊺ + HM H)

(24)

∂L where d is the number of neurons, v ∶= ∑x∈B ℓ′x (g(x) + 2B(x)µ) and H ∶= ∑x∈B 2ℓ′x A(x). Here we denote ℓ′x ∶= ∂f x and B denotes a minibatch.

This theorem can be extended to a form that describes a broad class of optimization methods including vanilla SGD, SGD with weight decay, and Polyak Momentum. With weight decay, a simple decay −γM term appears in the equation, where γ is the weight decay strength. See Appendix B.3. A deep neural network can often be viewed as a composition of multiple layers, each of which can be approximated by an NQF model. We theoretically and empirically study multilayer NQFs in Appendix B.4. Compressibility of NQF. A key implication of Theorem 2 is that the learning process of neural networks is highly redundant. We will say that two neural networks fx (θ) and gx (θ′ ) have the same learning dynamics if for any x, fx (θt ) = gx (θt′ ).

(25)

Namely, the identity is defined on the functional side, and it is possible for two models to have the same dynamics even if they have different parameters. Actually, there exist infinitely many parameter configurations whose learning dynamics are identical. This also implies that the learning process of a very large neural network is identical to and can be fully captured by a much smaller network – a phenomenon that has been termed the “dynamical lottery ticket hypothesis” [28, 69]. Theorem 3. Let an NQF have d neurons, defined by functions g(x), B(x), A(x), learning rate η, and initial parameter statistics µ(0) and M (0). Let kV ∶= dim (span ⋃x∈X ({g(x)} ∪ Col(A(x)) ∪ Col(B(x)))) ≤ p represent the dimension of the joint subspace spanned by the vectors g(x) and all column vectors of B(x) and A(x) across all x ∈ X , where X denotes the training set. If d > kV + 1, then there exists a smaller NQF with d′ = kV + 1 neurons, ′ characterized by g̃(x) = g(x), B̃(x) = B(x), Ã(x) = dd A(x) and learning rate η̃ = dd′ η, such that: 1. fd′ = fd for all inputs x ∈ X ; 2. fd′ = fd for all training steps t ≥ 0 under SGD with identical data sampling. In fact, independent of the original width d, there exists a finite-size NQF whose learning dynamics is identical to that of the original. This partially explains the commonly observed low dimensionality of the training dynamics [34, 45, 4]. Moreover, when the data is low-rank, this theorem states that the dimension of the learning dynamics is at most that of the data, a direct explanation of the folklore belief that structures in the data make learning simple. In Appendix B.5 we further characterize the compression error.

Exactly Solvable Cases of Learning Dynamics

The learning dynamics of NQF do not have a general solution. However, with special initialization or data distributions, solutions of the learning dynamics are obtainable. Here, we focus on training NQF with the gradient flow algorithm. We consider the empirical Mean Squared Error (MSE) loss over a dataset of m samples {xµ , yµ }m µ=1 : L(W ) =

1 m 2 ⊺ ∑ (Tr[W W A(xµ )] − yµ ) m µ=1

(26)

where W = [w1 , . . . , wd ] ∈ Rp×d is the parameter matrix. Let ∆µ (t) = Tr[W (t)W (t)⊺ A(xµ )] − yµ denote the residual for the µ-th sample. The gradient flow Ẇ = −∇W L(W ) is given by: Ẇ (t) = −

4 m ∑ ∆µ (t)A(xµ )W (t) = −H(t)W (t) m µ=1

(27)

m 4 where H(t) ∶= m ∑µ=1 ∆µ (t)A(xµ ) ∈ Rp×p . Notice that the network’s output is determined by the positive semi-definite matrix M = W W ⊺ ∈ Rp×p . By differentiating M (t), we observe that the dynamics naturally close in the space of M :

Ṁ (t) = Ẇ W ⊺ + W Ẇ ⊺ = −H(t)M (t) − M (t)H(t)

(28)

This kind of equation appears frequently in fluid dynamics and polymer physics, and this dynamical equation can be seen as a zero-dimensional homogeneous flow, and M can be identified as the Reynolds stress [60]. The nonlinear matrix differential equation (28) is generally impossible to solve. However, for some special cases it has exact solutions. 8

We will assume the following commutativity condition – this is a common assumption for solving the learning dynamics of exactly solvable models, e.g., [62, 33]. Assumption 1. The dataset consists of m mutually commuting symmetric data matrices: A(xµ ) = P Λµ P ⊺ , where P is an orthogonal matrix and Λµ = diag(λµ,1 , . . . , λµ,p ). Under Assumption 1, we can project the dynamics into the common eigenbasis defined by P . Let M̃ (t) = P ⊺ M (t)P . The residual can be rewritten as: p

∆µ (t) = Tr[M (t)A(xµ )] − yµ = Tr[M̃ (t)Λµ ] − yµ = ∑ M̃kk (t)λµ,k − yµ .

(29)

k=1

Equation (29) reveals that the network’s output and the loss gradient only depend on the diagonal elements of M̃ (t). We define zk (t) ∶= M̃kk (t) ≥ 0 as the k-th feature. Proposition 7. Under Assumption 1, the off-diagonal elements are slaved to the diagonal elements and do not influence the predictions or the loss. Their evolution is solved by: t

M̃ij (t) = M̃ij (0) exp (− ∫

(H̃ii (τ ) + H̃jj (τ )) dτ ) .

(30)

Substituting ∆µ (t) = ∑pj=1 λµ,j zj (t) − yµ into (28), we obtain the evolution of the k-th feature: 8 m ∑ ∆µ (t)λµ,k zk (t). m µ=1

(31)

p ⎛ m ⎞ ⎞ 8 ⎛m ∑ yµ λµ,k − ∑ ∑ λµ,k λµ,j zj (t) zk (t). m ⎝µ=1 ⎝ ⎠ ⎠ j=1 µ=1

(32)

żk (t) = − Expanding ∆µ (t) yields: żk (t) =

(32) is the Generalized Lotka-Volterra (GLV) equation from population ecology. In general, it still does not have analytic solutions1 . For the following sections, we will consider several special cases where the dynamics are actually solvable. The connection to the GLV allows some direct interpretation of zk as the abundance of a species and Ck,j ∶= ∑µ λµ,k λµ,j as the competition (when positive) and cooperation (when negative) between species [35]. Proportional Samples. Theorem 4. Suppose Assumption 1 holds, and further assume that the data matrices can be expressed as A(xµ ) = cµ A m m 4 4 for a symmetric matrix A and constants cµ . Define α ∶= m ∑µ=1 c2µ and β ∶= m ∑µ=1 cµ yµ . Then, the trajectory of W (t) is given by: W (t) = P exp(−ξ(t)Λ)P ⊺ W (0) (33) where exp(−ξ(t)Λ) = diag (exp(−λ1 ξ(t)), . . . , exp(−λp ξ(t))), and the auxiliary scalar variable ξ(t) satisfies the following implicit integral equation: ξ(t)

t=∫

ds α ∑pj=1 λj zj (0) exp(−2λj s) − β

(34)

Remark. The exact solution in [73] is a special case of Theorem 4, where Λ only has two eigenvalues and (34) can be integrated explicitly. Orthogonal Samples. Theorem 5. Suppose Assumption 1 holds, and further assume that A(xµ )A(xν ) = 0 for all µ ≠ ν. Then, the trajectory of W (t) is given by W (t) = P exp(−Σ(t))P ⊺ W (0), (35) m 4 where Σ(t) = m ∑µ=1 ξµ (t)Λµ is a diagonal matrix, and the auxiliary scalar variable ξµ (t) for each µ ∈ {1, . . . , m} evolves independently according to the following implicit integral equation: ξµ (t)

t=∫

. p 8 λµ,k s) − yµ ∑k=1 λµ,k zk (0) exp (− m

1 It has implicit solutions as shown in [33].

(36)

Orthogonal Features. Theorem 6. Suppose Assumption 1 holds, and further assume that: ∑m µ=1 λµ,k λµ,j = 0 for all k ≠ j. Define rk ∶=

8 m ∑ λµ,k yµ , m µ=1

Ckk ∶=

8 m 2 ∑λ . m µ=1 µ,k

(37)

Then the trajectory of W (t) is given by: W (t) = P D(t)P ⊺ W (0),

(38) √

where D(t) = diag(d1 (t), . . . , dp (t)) is a diagonal matrix with dk (t) = otherwise. The exact solution for zk (t) is: zk (t) =

zk (t)/zk (0) if zk (0) ≠ 0 and dk (t) = 0

rk zk (0) Ckk zk (0) + (rk − Ckk zk (0)) exp(−rk t)

(39)

k (0) otherwise. if rk ≠ 0 and zk (t) = 1+Czkk zk (0)t

Isotropic Samples. The previous solutions rely on the commutativity assumption (Assumption 1). In this section, we explore a different solvable regime where the data matrices do not commute, but instead satisfy an isotropic condition. Assumption 2. Assume the dataset consists of m ≥ p(p + 1)/2 symmetric data matrices {A(xµ )}m µ=1 and their second moment tensor satisfies: 1 m c (40) ∑ Aij (xµ )Akl (xµ ) = (δik δjl + δil δjk ) m µ=1 2 where c > 0 is a constant, and δ is the Kronecker delta. m 4 Theorem 7. Suppose Assumption 2 holds. Let Y = m ∑µ=1 yµ A(xµ ) ∈ Rp×p . Then, the trajectory of M (t) = ⊺ W (t)W (t) has a closed-form solution: −1

M (t) = exp(Y t)M (0) [I + 8cΦ(t)M (0)]

exp(Y t),

(41)

where I is the identity matrix, and Φ(t) = ∫0 exp(2Y τ )dτ . If Y is non-singular, Φ(t) = 21 Y −1 (exp(2Y t) − I). Furthermore, if the initial state M (0) and Y are simultaneously diagonalizable, then, the eigenvalues of M (t), denoted as zk (t), evolve according to zk (t) =

γk zk (0) 4czk (0) + (γk − 4czk (0)) exp(−2γk t)

(42)

zk (0) otherwise, where γk is the k-th eigenvalue of Y . if γk ≠ 0 and zk (t) = 1+8cz k (0)t

In Appendix B.7 we futher propose a general solution which unifies Theorem 6 and the simultaneously diagonalizable case of Theorem 7, as well as the exact solutions of linear networks in [62]. Neural Tangent Kernel and Feature Learning. While the main focus of our work is sudden learning at a small initialization, the NQF models can also be used to study lazy training and feature learning. We discuss this point in Appendix B.6.

Sudden Learning and Neural Scaling Laws

When learning of single features or data points is step-function-like, they can be naturally composed to give an arbitrary learning curve (as long as it is monotonically decreasing). For example, in this section, we show that when the data correlations follow a power-law structure, the learning curves naturally become power laws with predictable exponents. This leads to a precise prediction of the exponents of neural scaling laws. In this section, we only provide an informal statement. Please refer to Appendix A.17 for formal statements and the proofs. 10

Feature-wise Descent Under the conditions of Theorem 6 (or Theorem 7 with M (0) and Y simultaneously diagonalizable), define ζk = rk (in Theorem 6) or ζk = γk (in Theorem 7), which can be understood as the effective growth rate for the feature zk . As the initialization scale ϵ → 0, the characteristic time for zk to be activated is t∗k ∼

1 1 ln . ζk ϵ

(43)

Consequently, the gaps between different feature activations diverge in the limit of small initialization, which exhibits saddle-to-saddle dynamics: lim ∣t∗k′ − t∗k ∣ = ∞ (44) ϵ→0

if ζk′ ≠ ζk . 2 γ2 rk (under Theorem 6) or Vk ∶= ck (under Theorem 7), which can be understood as the Moreover, define Vk ∶= Ckk effective target strength. Assume that ζk and Vk exhibit the following power-law decays: ζk ∝ k −α2 , Vk ∝ k −α1 .

(45)

Then, under the infinite-width limit p → ∞ and the small initialization limit ϵ → 0, the limiting excess loss has a power-law decay α −1

E(τ ) = Θ (τ

1 − α

(46)

8 Sample-wise Descent Under the conditions of Theorem 5, suppose ζµ ∶= m maxk λµ,k yµ > 0, which can be understood as the effective growth rate of the µ-th sample. When the initialization scale ϵ → 0, the empirical MSE converges to a saddle-to-saddle trajectory 1 1 m 2 ∗ lim+ L (τ ln ) = (47) ∑ y ⋅ I(τ < τµ ), ϵ→0 ϵ m µ=1 µ

where I(⋅) denotes the indicator function and τµ∗ ∶= ζ1µ . Moreover, suppose that the ζµ and yµ exhibit the following power-law decay: ζµ ∝ µ−γ2 , yµ ∝ µ−γ1

(48)

for γ1 > 1/2 and γ2 > 0. Then under the infinite sample limit m → ∞ and the small initialization limit ϵ → 0, the limiting MSE follows a power-law decay: L(τ ) = Θ (τ

−

2γ1 −1 γ2

(49)

The dynamics in Theorem 4 can be viewed as being mathematically equivalent to the solution in Theorem 5 restricted to m = 1, and thus there is only one plateau.

Experiments

All the experiments in this section are described in Appendix C. The code that reproduces every figure is available at https://github.com/xu-yz19/Neural-Quadratic-Forms. NQF approximation To empirically validate the NQF approximation (Theorem 1), we train a two-layer MLP (Proposition 1), a single-layer CNN (Proposition 6), a query-key-only attention head (Proposition 3) and a multi-head attention model (Proposition 2) in Figure 2 with different optimizers (gradient descent, Polyak momentum β = 0.9, and Adam). All models use Gaussian initialization with scale σ. The input is chosen to be Gaussian, and the label is generated by a low-rank NQF. The NQF tracks the original almost exactly across every architecture and optimizer. At large initialization, the NQF visibly departs from the original, showing that the NQF approximation is only valid in the small-initialization regime. In Appendix C, we further provide additional experiments concerning the NQF approximation for the teacher-student setting and real datasets.

Figure 2: NQF approximation across initialization, optimizer, and architecture. Two initialization scales are juxtaposed—small init (σ = 0.01, left) and large init (σ = 0.2, right). Within each block, rows are optimizers and columns are architectures. Each module is trained together with its NQF approximation on an NQF teacher with singular values s = (1, 12 , 14 , 18 ). Curves show the recovered k-th singular value ŝk (t)/sk versus training step; colored solid = original module, black dashed = NQF. Feature-wise saddle-to-saddle dynamics To validate the feature-wise descent in Section 6, we choose {A(xµ )}m µ=1 to be diagonal matrices satisfying the condition of Theorem 6 and train the model using full-batch GD from a small initialization. As shown in Figure 3, the numerical trajectories match the predictions of Theorem 6. Furthermore, the excess loss exhibits sequential plateaus at the predicted characteristic timescales t∗k . On the right side of Figure 3, we further choose m = 1000, rk , Vk to follow power laws and confirm that the slope predicted by Section 6 matches the curve obtained by GD. Sample-wise saddle-to-saddle dynamics To validate the sample-wise descent in Section 6, we construct a synthetic dataset consisting of m mutually orthogonal data matrices. The student model is trained via full-batch GD from a small initialization scale. As illustrated in Figure 4, the numerical trajectories obtained via GD match the analytical curves in Theorem 5. Specifically, the left panel shows that the MSE for each individual sample remains stagnant at its initial plateau before undergoing a sudden, sharp decay toward zero near t∗µ . Consequently, as shown in the right panel, the total MSE loss transitions through m sequential plateaus. Power law on MLP

We consider a Fourier MLP 1 P f (x) = √ ∑ Wk tanh(sk ak ψk (x)), P k=1

(50)

where the trainable parameters are wk = [ak , Wk ]⊺ . Here, ψk (x) represents orthonormal Fourier features ψk (x) = √ 2 cos(2πkx) sampled on a uniform 1D grid xµ ∈ {0, 1/M, . . . , (M − 1)/M }. To induce the power-law structures required by Section 6, we choose sk = k −θ and construct a target function y(x) = ∑P k=1 bk ψk (x) with decaying coefficients bk = k −β . In Appendix C we will show that the NQF of this Fourier MLP satisfies the conditions in Section 6 (feature-wise descent). We train the model using full-batch GD. We use the Gaussian initialization with a small scale. We compare this architecture against two baselines: (1) its NQF approximation (Proposition 1); and (2) a standard MLP trained on the raw input x. The results are shown in Figure 5, which shows an empirical power-law loss and matches the prediction in Section 6. Meanwhile, the MLP trained on the raw input does not exhibit a power-law loss, which demonstrates the necessity of the structure required by Section 6 to attain the power laws. The small systematic steepening of the empirical power law relative to theory is the finite-initialization effect of Appendix C, which vanishes as ϵ → 0.

Figure 3: Sudden learning in a solvable NQF, the phenomenon the theory of this paper is built to explain. The experiment is under the feature-wise descent setting of Section 6. Solid lines represent the theory in Theorem 6 and markers k ) where represent the results obtained by GD. Vertical dashed lines mark the characteristic timescales t∗k = r1k ln ( Crkk ϵ ⊺ each feature reaches half-saturation. Left: Evolution of the eigenvalues of W W . Middle: The excess MSE loss over time, which exhibits sequential plateaus. Each sharp drop in the loss aligns with t∗k . Right: Under the feature-wise descent setting of Section 6, the blue dashed line indicates the theoretical slope and confirms that our predicted slope matches the loss trajectory before the finite-size cutoff.

Related Works

Quadratic and factorized model classes. Quadratic parameterizations have long served as solvable proxies for nonconvex learning, which includes low-rank matrix factorization [33, 5], phase retrieval [18], diagonal linear networks [12], quadratic networks [48, 23], and so on. These models have been used to explain implicit low-rank bias, spectral initialization, and feature learning. A recent work [53] suggests that quadratic models can even accurately approximate LLMs. NQFs differ from these model-specific analyses by deriving the quadratic features A(x) from the architecture: matrix sensing, diagonal networks, MLPs, CNNs, attention heads, and MoE branches become instances of the same local normal form. Learning dynamics and feature learning. Exact analyses of deep linear networks show rapid transitions between mode-wise plateaus under gradient descent [62]. The neural tangent kernel and lazy-training literature describe the opposite regime, in which a network is well approximated by its linearization around initialization and training is governed by a nearly fixed kernel [38, 22]. Feature-learning work has emphasized that realistic networks change their representation during training [61, 10]. NQFs sit between these perspectives: the model is still a local expansion near small initialization, but the retained quadratic term makes the kernel and representation evolve through the finitedimensional order parameter M , giving closed nonlinear dynamics for feature growth and competition. Sudden learning and scaling laws. Empirical neural scaling laws show that loss often follows power laws in data, model size, or compute [39]. Several theoretical works explain such laws through spectral structure: high-eigenvalue modes are learned earlier, and power-law eigenspectra or task-model alignments induce power-law learning curves [15, 17, 49, 7, 23]. Sudden learning has also been formalized in SGD analyses where target structure is acquired sequentially through saddle-to-saddle transitions [1]. The NQF perspective connects these two viewpoints: individual features or samples undergo sharp activation events at initialization-dependent times, while a power-law distribution of feature alignments, variances, or sample difficulties aggregates these transitions into a smooth scaling law [54, 56, 11].

Discussion

We have derived a normal form for the training of a smooth nonlinear neural networks from one structural assumption: that the repeated components of a layer are exchangeable. This property is present in essentially every architecture that has a notion of width. Within that scope, the resulting neural quadratic form is architecture independent, and what distinguishes one architecture from another is the sparsity pattern and the strength of its structure matrix A. The learning 13

Figure 4: The experiment is under the setting of the sample-wise descent in Section 6. Solid lines represent the theory of Theorem 5 obtained by numerically solving Eq. (36), and markers represent the empirical results obtained by fullbatch GD. Vertical dashed lines mark the characteristic timescales t∗µ = ζµ−1 ln (1/ϵ) according to (47). Left: Evolution of the per-sample loss ∆µ (t)2 = (Tr[W W ⊺ A(xµ )] − yµ )2 , showing independent and sequential convergence. Right: The total empirical MSE loss over time, which exhibits m sequential plateaus. Each discrete drop in the total loss matches the respective timescale t∗µ .

Figure 5: Curves show the normalized MSE L(t)/L(0); dotted lines are the slopes predicted in Section 6. Left: For θ = 1, the Fourier MLP (blue) descends as the predicted power law t−(2β−1)/(θ+β) (see Appendix C) and is tracked by its NQF approximation (black); a standard tanh MLP trained on the raw input x (red), which does not exhibit the power law. Right: The power-law exponent of the Fourier MLP (solid lines) matches the theory (dashed lines) for θ ∈ {0, 0.5, 1}. dynamics of any NQF close on the finite-dimensional order parameters M and µ, they reduce in solvable regimes to generalized Lotka–Volterra flows, and they generate saddle-to-saddle transitions whose composition produces neural scaling laws with predictable exponents. Read as a Landau construction, the calculation says something specific about universality in deep learning. The dictionary is worth setting down carefully, because two of its entries are not the ones a first reading suggests. The symmetry group is Sd ; the expansion is truncated at the lowest symmetry-allowed nonconstant order; and the architecture enters only through the coupling A(x), in the same way that microscopic detail enters a Landau functional only through phenomenological coefficients. That much is the standard construction, and it is what makes universality across architectures the expected outcome. The remaining two entries need more care. The first is the order parameter: M = W W ⊺ . It is Sd -invariant by construction, and invariant as well under the larger redundancy W ↦ W O, O ∈ O(d), that the quadratic form does not see; it therefore does not transform in a nontrivial representation of the symmetry group, as the order parameter of a textbook Landau construction does. What it shares with the Edwards–Anderson overlap of spin-glass theory, and with the overlaps wi⊺ wj that serve as the order parameters of online learning in committee machines [25, 63], is that it is an invariant which nonetheless vanishes in one phase and not in the other: zk = 0 is the unlearned state and zk > 0 the learned one. The symmetry broken at ignition is the freedom of W itself, which the flow on M has quotiented out,

and M is the invariant that registers the breaking. The second is the control parameter ϵ, which is not the small parameter of the expansion. The initialization scale ϵ does not decide whether a feature is learned, since every mode with rk > 0 will be learned at every ϵ > 0; what tunes the transition is rk , or rk − γ once weight decay is included, whose sign decides the stability of zk = 0 in the way that a reduced temperature decides the stability of the disordered state – a point we will explore in a future work. What ϵ controls instead is the sharpness of the crossover, and its natural counterpart is system size: the gaps between successive ignitions grow as ln(1/ϵ) while the transitions themselves do not, so ln(1/ϵ) plays the role of N , and ϵ → 0 is the analogue of the thermodynamic limit. Sudden learning is then the ignition of individual modes of M , and the ϵ → 0 limit is a genuine singular limit in which a smooth flow acquires sharp transitions; they are sharp for the same reason that a phase transition is sharp only in an infinite system. Limitations. plainly.

Being a minimal model of training, the model certainly has many limitations, which are worth stating

1. Small initialization. The NQF is the leading nonconstant term of a local expansion, and it is accurate only when the initialization scale is small compared with the feature length scale set by the data. Figure 2 shows the approximation failing at large initialization, which is the behavior the theory predicts for itself. Appendix B.6 partially relaxes this by treating the lazy and feature-learning regimes within the same framework. 2. Smoothness. Theorem 1 assumes three-times differentiability. Rectified activations at the origin, bias terms, normalization layers, and modules that apply a nonlinearity after aggregating their components each violate one of these and require additional treatment. 3. Other plateaus. Our experiments show plateaus that the NQF does not account for, which we attribute to cubic and higher terms. Extending to neural cubic and quartic forms (Appendix B.1) is the natural next step, as is a systematic treatment of multilayer NQFs, begun in Appendix B.4. 4. Power-law spectra are an assumption, not a prediction. Deriving A(x) from the architecture fixes the operator whose spectrum controls the scaling exponent. It does not, however, predict that the spectrum has a power-law tail. That tail is a hypothesis of the scaling theorem of Section 6, and our Fourier-feature experiment constructs it deliberately. Whether it arises generically, for realistic data and architectures, is the question that would have to be settled in future works. 5. Adaptive optimizers. The linear-update corollary of Appendix B.3 covers gradient descent, weight decay, and Polyak momentum, but not Adam, which is nonetheless tracked accurately by the NQF in our experiments (Figure 2). Outlook. The correspondences to physics suggest that the technical machinery developed for interacting populations, hydrodynamic instabilities, and polymer flow can be brought to bear on learning dynamics, and conversely that trained networks provide an unusually well-instrumented laboratory for these classical models: every mode amplitude is directly observable, the coupling matrix is known exactly, and the control parameter can be tuned over many decades. Whether the mapping survives beyond quadratic order, and whether the ecological picture of competition between features predicts anything about generalization rather than optimization, are the questions we regard as most worth pursuing.

acknowledgments We gratefully acknowledge the generous support of NTT Research, Inc. I.L.C. acknowledges support in part from the Institute for Artificial Intelligence and Fundamental Interactions (IAIFI) through NSF Grant No. PHY-2019786. This work was also supported by the Center for Brains, Minds and Machines (CBMM), funded by NSF STC award No. CCF-1231216.

References [1] Emmanuel Abbe, Enric Boix Adserà, and Theodor Misiakiewicz. Sgd learning on neural networks: Leap complexity and saddle-to-saddle dynamics. In Proceedings of the Thirty Sixth Conference on Learning Theory, volume 195 of Proceedings of Machine Learning Research, pages 2552–2623. PMLR, 2023. [2] Hisham Abou-Kandil, Gerhard Freiling, Vlad Ionescu, and Gerhard Jank. Matrix Riccati equations in control and systems theory. Birkhäuser, 2012. [3] Madhu S. Advani, Andrew M. Saxe, and Haim Sompolinsky. High-dimensional dynamics of generalization error in neural networks. Neural Networks, 132:428–446, 2020. [4] Armen Aghajanyan, Sonal Gupta, and Luke Zettlemoyer. Intrinsic dimensionality explains the effectiveness of language model fine-tuning. In Proceedings of the 59th annual meeting of the association for computational linguistics and the 11th international joint conference on natural language processing (volume 1: long papers), pages 7319–7328, 2021. [5] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In Advances in Neural Information Processing Systems, volume 32, 2019. [6] Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. Advances in neural information processing systems, 36:48314–48362, 2023. [7] Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jaehoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. Proceedings of the National Academy of Sciences, 121(27):e2311878121, 2024. [8] Yasaman Bahri, Jonathan Kadmon, Jeffrey Pennington, Sam S Schoenholz, Jascha Sohl-Dickstein, and Surya Ganguli. Statistical mechanics of deep learning. Annual Review of Condensed Matter Physics, 11:501–528, 2020. [9] Adrien Bardes, Jean Ponce, and Yann LeCun. VICReg: Variance-invariance-covariance regularization for selfsupervised learning. In International Conference on Learning Representations, 2022. [10] Daniel Beaglehole, Adityanarayanan Radhakrishnan, Parthe Pandit, and Mikhail Belkin. Mechanism of feature learning in convolutional neural networks. arXiv preprint arXiv:2309.00570, 2023. [11] Gérard Ben Arous, Murat A. Erdogdu, Nuri Mert Vural, and Denny Wu. Learning quadratic neural networks in high dimensions: SGD dynamics and scaling laws. In Advances in Neural Information Processing Systems, volume 38, 2025. [12] Raphaël Berthier. Incremental learning in diagonal linear networks. Journal of Machine Learning Research, 24(171):1–26, 2023. [13] Fabrizio Boncoraglio, Vittorio Erba, Emanuele Troiani, Yizhou Xu, Florent Krzakala, and Lenka Zdeborová. Single-head attention in high dimensions: A theory of generalization, weights spectra, and scaling laws. In Proceedings of the 43rd International Conference on Machine Learning, 2026. [14] Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. A dynamical model of neural scaling laws. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of PMLR, 2024. [15] Blake Bordelon, Abdulkadir Canatar, and Cengiz Pehlevan. Spectrum dependent learning curves in kernel regression and wide neural networks. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 1024–1034. PMLR, 2020. [16] Guy Bunin. Ecological communities with Lotka–Volterra dynamics. Phys. Rev. E, 95:042414, 2017. [17] Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks. Nature Communications, 12(2914), 2021.

[18] Emmanuel J. Candès, Xiaodong Li, and Mahdi Soltanolkotabi. Phase retrieval via wirtinger flow: Theory and algorithms. IEEE Transactions on Information Theory, 61(4):1985–2007, 2015. [19] Yuan Cao and Quanquan Gu. Tight sample complexity of learning one-hidden-layer convolutional neural networks. Advances in Neural Information Processing Systems, 32, 2019. [20] Giuseppe Carleo, Ignacio Cirac, Kyle Cranmer, Laurent Daudet, Maria Schuld, Naftali Tishby, Leslie VogtMaranto, and Lenka Zdeborová. Machine learning and the physical sciences. Rev. Mod. Phys., 91:045002, 2019. [21] Paul M. Chaikin and Tom C. Lubensky. Principles of Condensed Matter Physics. Cambridge University Press, Cambridge, 1995. [22] Lénaı̈c Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In Advances in Neural Information Processing Systems, volume 32, 2019. [23] Leonardo Defilippis, Yizhou Xu, Julius Girardin, Vittorio Erba, Emanuele Troiani, Lenka Zdeborová, Bruno Loureiro, and Florent Krzakala. Scaling laws and spectra of shallow neural networks in the feature learning regime. In The Fourteenth International Conference on Learning Representations, 2026. [24] Simon Du, Jason Lee, Yuandong Tian, Aarti Singh, and Barnabas Poczos. Gradient descent learns one-hiddenlayer cnn: Don’t be afraid of spurious local minima. In International Conference on Machine Learning, pages 1339–1348. PMLR, 2018. [25] Andreas Engel and Christian Van den Broeck. Statistical Mechanics of Learning. Cambridge University Press, Cambridge, 2001. [26] Vittorio Erba, Emanuele Troiani, Luca Biggio, Antoine Maillard, and Lenka Zdeborová. Bilinear sequence regression: A model for learning from long sequences of high-dimensional tokens. Physical Review X, 15(2):021092, 2025. [27] Mathieu Even, Scott Pesme, Suriya Gunasekar, and Nicolas Flammarion. (S)GD over diagonal linear networks: Implicit bias, large stepsizes and edge of stability. In Advances in Neural Information Processing Systems, volume 36, 2023. [28] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference on Learning Representations, 2019. [29] Kenji Fukumizu. A regularity condition of the information matrix of a multilayer perceptron network. Neural Networks, 9(5):871–879, 1996. [30] Rong Ge, Jason D Lee, and Tengyu Ma. Matrix completion has no spurious local minimum. Advances in neural information processing systems, 29, 2016. [31] Yoav Goldberg and Omer Levy. word2vec explained: Deriving mikolov et al.’s negative-sampling wordembedding method, 2014. [32] Nigel Goldenfeld. Lectures on Phase Transitions and the Renormalization Group. Addison-Wesley, Reading, MA, 1992. [33] Suriya Gunasekar, Blake Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nathan Srebro. Implicit regularization in matrix factorization. In Advances in Neural Information Processing Systems, volume 30, pages 6151–6159. Curran Associates, Inc., 2017. [34] Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace. arXiv preprint arXiv:1812.04754, 2018. [35] Josef Hofbauer and Karl Sigmund. Evolutionary games and population dynamics. Cambridge university press, 1998.

[36] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models. In Advances in Neural Information Processing Systems, volume 35, 2022. [37] Minyoung Huh, Brian Cheung, Tongzhou Wang, and Phillip Isola. Position: The platonic representation hypothesis. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 20617–20642. PMLR, 2024. [38] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems, volume 31, 2018. [39] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020. [40] Prakhar Kaushik, Shravan Chaudhari, Ankit Vaidya, Rama Chellappa, and Alan Yuille. The universal weight subspace hypothesis, 2025. [41] Yehuda Koren, Robert Bell, and Chris Volinsky. Matrix factorization techniques for recommender systems. Computer, 42(8):30–37, 2009. [42] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25, 2012. [43] Daniel Kunin, Giovanni Luca Marchetti, Feng Chen, Dhruva Karkada, James Simon, Michael Deweese, Surya Ganguli, and Nina Miolane. Alternating gradient flows: A theory of feature learning in two-layer neural networks. Advances in Neural Information Processing Systems, 38:4377–4424, 2025. [44] Lev Davidovich Landau and Evgenii Mikhailovich Lifshitz. Statistical Physics: Volume 5, volume 5. Elsevier, 2013. [45] Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the intrinsic dimension of objective landscapes. In International Conference on Learning Representations, 2018. [46] Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning. Advances in Neural Information Processing Systems, 35:34651–34663, 2022. [47] Antoine Maillard, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Phase retrieval in high dimensions: Statistical and computational phase transitions. Advances in Neural Information Processing Systems, 33:11071– 11082, 2020. [48] Antoine Maillard, Emanuele Troiani, Simon Martin, Lenka Zdeborová, and Florent Krzakala. Bayes-optimal learning of an extensive-width neural network from quadratically many samples. Advances in Neural Information Processing Systems, 37:82085–82132, 2024. [49] Alexander Maloney, Daniel A. Roberts, and James Sully. A solvable model of neural scaling laws. arXiv preprint arXiv:2210.16859, 2022. [50] Rodrigo Maulen Soto, Pierre Marion, and Claire Boyer. Attention-based clustering. Advances in Neural Information Processing Systems, 38:66455–66506, 2025. [51] Robert M. May. Will a large complex system be stable? Nature, 238:413–414, 1972. [52] Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on learning theory, pages 2388–2464. PMLR, 2019. [53] Alexandru Meterez, Pranav Ajit Nair, Depen Morwani, Cengiz Pehlevan, Sham Kakade, and Alex Damian. A defense of the quadratic model, 2026. 18

[54] Eric Michaud, Ziming Liu, Uzay Girit, and Max Tegmark. The quantization model of neural scaling. Advances in Neural Information Processing Systems, 36:28699–28722, 2023. [55] Tomas Mikolov, Ilya Sutskever, Kai Chen, Greg S. Corrado, and Jeff Dean. Distributed representations of words and phrases and their compositionality. In Advances in Neural Information Processing Systems, volume 26, 2013. [56] Yoonsoo Nam, Nayara Fonseca, Seok H Lee, Chris Mingard, and Ard A Louis. An exactly solvable model for emergence and scaling laws in the multitask sparse parity problem. Advances in Neural Information Processing Systems, 37:39632–39693, 2024. [57] Huy Nguyen, Nhat Ho, and Alessandro Rinaldo. Sigmoid gating is more sample efficient than softmax gating in mixture of experts. Advances in Neural Information Processing Systems, 37:118357–118388, 2024. [58] Xu Pan, Aaron Philip, Ziqian Xie, and Odelia Schwartz. Dissecting query-key interaction in vision transformers. In Advances in Neural Information Processing Systems, volume 37, 2024. [59] Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity. Advances in Neural Information Processing Systems, 34:29218–29230, 2021. [60] Stephen B. Pope. Turbulent Flows. Cambridge University Press, Cambridge, 2000. [61] Adityanarayanan Radhakrishnan, Daniel Beaglehole, Parthe Pandit, and Mikhail Belkin. Mechanism for feature learning in neural networks and backpropagation-free machine learning models. Science, 383(6690):1461–1467, 2024. [62] Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In International Conference on Learning Representations, 2014. [63] H. S. Seung, H. Sompolinsky, and N. Tishby. Statistical mechanics of learning from examples. Phys. Rev. A, 45:6056–6091, 1992. [64] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations, 2017. [65] Nathan Srebro, Jason Rennie, and Tommi Jaakkola. Maximum-margin matrix factorization. Advances in neural information processing systems, 17, 2004. [66] Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruction. Advances in Neural Information Processing Systems, 34:23831–23843, 2021. [67] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. [68] Julius Von Kügelgen, Yash Sharma, Luigi Gresele, Wieland Brendel, Bernhard Schölkopf, Michel Besserve, and Francesco Locatello. Self-supervised learning with data augmentations provably isolates content from style. Advances in neural information processing systems, 34:16451–16467, 2021. [69] Hong-Yi Wang, Di Luo, Tomaso Poggio, Isaac L Chuang, and Liu Ziyin. A universal compression theory for lottery ticket hypothesis and neural scaling laws. In The Fourteenth International Conference on Learning Representations, 2026. [70] Zixuan Wang, Eshaan Nichani, Alberto Bietti, Alex Damian, Daniel Hsu, Jason D. Lee, and Denny Wu. Learning compositional functions with transformers from easy-to-hard data. In Proceedings of Thirty Eighth Conference on Learning Theory, volume 291 of Proceedings of Machine Learning Research, pages 5632–5711. PMLR, 2025. [71] Tailin Wu, Ian Fischer, Isaac L. Chuang, and Max Tegmark. Learnability for the information bottleneck. Entropy, 21(10):924, 2019. 19

[72] Yizhou Xu, Antoine Maillard, Lenka Zdeborová, and Florent Krzakala. Fundamental limits of matrix sensing: Exact asymptotics, universality, and applications. In Proceedings of Thirty Eighth Conference on Learning Theory, volume 291 of Proceedings of Machine Learning Research, pages 5757–5823. PMLR, 2025. [73] Yizhou Xu and Liu Ziyin. Three mechanisms of feature learning in a linear network. In International Conference on Learning Representations, 2025. [74] Lenka Zdeborová and Florent Krzakala. Statistical physics of inference: Thresholds and algorithms. Advances in Physics, 65(5):453–552, 2016. [75] Yedi Zhang, Andrew Saxe, and Peter E. Latham. Saddle-to-saddle dynamics explains a simplicity bias across neural network architectures. In The Fourteenth International Conference on Learning Representations, 2026. [76] Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin. Quadratic models for understanding catapult dynamics of neural networks. In International Conference on Learning Representations, 2024. [77] Liu Ziyin. Symmetry induces structure and constraint of learning. In Forty-first International Conference on Machine Learning, 2024. [78] Liu Ziyin, Ekdeep Singh Lubana, Masahito Ueda, and Hidenori Tanaka. What shapes the loss landscape of self-supervised learning? In International Conference on Learning Representations, 2023. [79] Liu Ziyin, Mingze Wang, Hongchao Li, and Lei Wu. Parameter symmetry and noise equilibrium of stochastic gradient descent. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. [80] Liu Ziyin, Yizhou Xu, Tomaso Poggio, and Isaac Chuang. Parameter symmetry potentially unifies deep learning theory, 2025.

Contents 1

Introduction

Neural Quadratic Forms

NQF for Different Architectures

Learning Dynamics of the NQF

Exactly Solvable Cases of Learning Dynamics

Sudden Learning and Neural Scaling Laws

Experiments

Related Works

Discussion

A Theory A.1 List of Notations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.2 List of NQF Examples . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.3 Proof of Theorem 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.4 Proof of Proposition 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.5 Proof of Proposition 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.6 Proof of Proposition 3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.7 Proof of Proposition 4 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.8 Proof of Proposition 5 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.9 Proof of Proposition 6 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.10 Proof of Theorem 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.11 Proof of Theorem 3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.12 Proof of Proposition 7 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.13 Proof of Theorem 4 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.14 Proof of Theorem 5 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.15 Proof of Theorem 6 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.16 Proof of Theorem 7 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . A.17 Formal Statements of Section 6 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

21 21 23 23 24 25 26 26 27 27 28 28 30 30 31 32 33 34

B Additional Theory B.1 Removing ZGZ Condition . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.2 Multidimensional Output . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.3 Learning Dynamics of Other Optimization Methods . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.4 Learning Dynamics of Multi-layer NQF . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.5 Compression Error . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.6 NTK and Feature Learning . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . B.7 A general exact solution . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

39 39 41 41 42 44 45 46

C Experimental Details and Additional Experiments

Theory

A.1

List of Notations

Table 2 lists all notations used throughout the paper.

Table 2: Notation used throughout the paper. “Defined” gives the section or equation in which the symbol is introduced. Symbol Meaning Architecture and the normal form d width: number of exchangeable components p parameter dimension of one component wi the i-th component (“neuron”) W all components, W = (w1 , . . . , wd ) fx (W ) module output on input x Sd symmetric group acting on the d components A(x) structure matrix B(x), g(x) extra couplings when ZGZ is dropped C(x) off-diagonal block of A in a feedforward module ϵ initialization scale (small parameter of the expansion) Order parameters and dynamics M order parameter W W ⊺ , positive semidefinite µ first-moment order parameter ∑i wi 4 H residual-weighted structure operator m ∑µ ∆µ A(xµ ) η learning rate γ weight-decay strength kV dimension of the joint data subspace K(x, x′ ) empirical neural tangent kernel Data and spectra m number of training samples xµ , y µ µ-th input and target ∆µ residual Tr[M A(xµ )] − yµ L sample-averaged MSE loss L˜ unnormalized squared error ∑µ ∆2µ E excess loss L(t) − L(∞) P, Λµ common eigenbasis and eigenvalue matrix of A(xµ ) λµ,k k-th eigenvalue of A(xµ ) zk k-th feature, the k-th eigenvalue of M rk growth rate of feature k Ckj feature interaction matrix 4 Y, γk y A(xµ ) and its k-th eigenvalue m ∑µ µ c isotropy constant of the data second moment Sudden learning and scaling laws ζk effective growth rate (rk , or γk if isotropic) a 1 under Thm. 6, 2 under Thm. 7 Vk effective target strength of feature k t∗k ignition time of feature k τ rescaled time t/ ln(1/ϵ) α1 , α2 decay exponents of Vk and ζk γ1 , γ 2 decay exponents of yµ and ζµ β0 decay exponent of the initialization profile β decay exponent of the target coefficients bk θ spectral exponent of the Fourier features αeff effective exponent at finite ϵ

Space

Defined

N N Rp Rp×d R — Rp×p sym. Rp×p , Rp Rd1 ×d2 R>0

Section 2 Section 2 Section 2 Section 2 Section 2 Definition 1 Eq. (2) Eq. (20) Section 3 Section 6

Rp×p Rp Rp×p R>0 R≥0 N R

Eq. (28) Section 4 Eq. (27) Thm. 2 Section 4 Thm. 3 App. B.6

N — R R≥0 R≥0 R≥0 Rp×p R R≥0 R Rp×p Rp×p , R R>0

Section 5 Section 5 Section 5 Section 5 Section 6 Section 6 Asm. 1 Asm. 1 Section 5 Eq. (37) Eq. (37) Thm. 7 Asm. 2

R>0 {1, 2} R≥0 R>0 R>0 R>0 R>0 R R R R

Section 6 Eq. (43) Section 6 Eq. (43) Section 6 Section 6 Section 6 App. A.17 Section 7 Section 7 App. C

Table 3: Complete list of models that are special cases of the neural quadratic form, together with their structure matrices A(x). A six-row excerpt appears as Table 1 in the main text. Two conventions vary across the rows. The neurons are the columns of W in most entries, so that M = W W ⊺ ; in the tied query–key and linear SSL rows they are the rows of W , so that M = W ⊺ W . Model PSD matrix sensing / PSD factorization

Structure matrix A(x) Aa = Sym(Sa ); if Sa = Sa⊺ , Aa = Sa .

NQF notation ŷa = ⟨Sa , W W ⊺ ⟩

Rectangular matrix sensing / matrix factorization

ŷa = ⟨Sa , U V ⊺ ⟩ = Tr(U ⊺ Sa V ) R̂ij = u⊺i vj

Matrix completion / recommender entry

swc = u⊺w vc

Word-context embedding score Quadratic network Phase retrieval Diagonal linear network, square parametrization

fW (x) = ∑r (wr⊺ x)2 ŷa = a⊺ W W ⊺ a fx (u) = ∑k xk u2k /4

fx (W ) = ∑di=1 vi ϕ(u⊺i x)

Two-layer MLP Single-layer CNN

fx (W ) = ⊺ ∑di=1 vi ∑P p=1 ϕ(ki xp )

Mixture of Experts (MoE)

fx (W ) = ∑di=1 ψ(ri⊺ x)E(x; Θi )

[41, 30] [55, 31]

fx (W ) = ∑di=1 vi⊺ WiV X . . .

Query-key-only attention

fx = ∑di=1 c⊺ X ⋅ softmax(. . . )

Tied query-key attention

√1 u⊺ W ⊺ W x dk

Bilinear sequence regression

x⊺s Bxt , with B = U V ⊺

Linear SSL augmentation

∥W ∆∥2 = Tr(W ⊺ W ∆∆⊺ )

[61, 76, 23] [18, 47] [12, 59]

⎡ 0 Diag(x)⎤ ⎢ ⎥ ⎥ Ax = 21 ⎢ ⎢Diag(x) ⎥ 0 ⎣ ⎦ ⎡ ⎤ x⎥ ϕ′ (0) ⎢0k×k ⎥ A(x) = 2 ⎢ ⊺ ⎢ x 0⎥ ⎣ ⎦ ⎤ ⎡ ∑P ϕ′ (0) ⎢ 0m×m p=1 xp ⎥ ⎥ A(X) = 2 ⎢ P ⊺ ⎥ ⎢∑ 0 ⎦ ⎣ p=1 xp A(x) =

Multi-head attention (MHA)

A.2

[33, 5, 72]

Ax = xx⊺ Aa = aa⊺ Ax = 14 Diag(x)

fx (u, v) = u⊺ Diag(x)v

Diagonal linear network, two-factor signed form

⎡0 Sa ⎤ ⎢ ⎥ ⎥ Aa = 12 ⎢ ⊺ ⎢Sa 0⎥ ⎣ ⎦ ⎡ ⊺⎤ ⎢ 0 e i ej ⎥ ⎥ Aij = 21 ⎢ ⎢ ⊺ ⎥ ⎢ej ei 0 ⎥ ⎣ ⎦ ⊺⎤ ⎡ 0 e e w ⎢ c⎥ ⎥ Awc = 12 ⎢ ⊺ ⎢e c e w 0 ⎥ ⎣ ⎦

⎡

0 ψ ′ (0) ⎢ ⎢ 2 ⎢G(x) x⊺

Refs. [66, 33, 72]

x G(x)⊺ ⎤ ⎥ ⎥, with ⎥ 0 ⎦

⎣ G(x) ∶= ∇Θi E(x; 0) ⎤ ⎡0 0 0 0 ⎥ ⎢ ⎥ ⎢ ⎥ ⎢ 0 0 0 0 1 ⎢ ⎥ A(X) = 2 ⎢ 0 xavg ⊗ Idv ⎥ ⎥ ⎢0 0 ⎥ ⎢ ⎥ ⎢0 0 x⊺ ⊗ I 0 ⎦ ⎣ d avg v ⊺ ⎤ ⎡ 0 xu ⊗ Idk ⎥ ⎢ 1 ⎥, A(X) = √ ⎢ ⊺ ⎥ 2 dk ⎢ux ⊗ Id 0 ⎦ ⎣ k where u = ΣX c. ⊺ ⊺ √1 Atied x,X,c = 2 d (xu + ux ) k ⎡ 0 xs x⊺t ⎤ ⎥ ⎢ ⎥ As,t = 21 ⎢ ⎢xt x⊺s 0 ⎥ ⎦ ⎣ A∆ = ∆∆⊺ , with ∆ = x − χ

[27, 59] [29, 61] [42, 24, 19] [64, 57]

[67]

[6, 58, 50]

[58, 13] [65, 33, 26] [68, 9, 78]

List of NQF Examples

Table 3 gives the full list of models that reduce to the neural quadratic form, extending Table 1.

A.3

Proof of Theorem 1

Because fx is three times continuously differentiable, we can write its multi-variable Taylor expansion around the origin W = (0, . . . , 0): d

fx (W ) = fx (0) + ∑(∇wi fx (0))⊺ wi + i=1 2

∂ fx where Hij = ∂w ∣ i ∂wj

1 d d ⊺ 3 ∑ ∑ w Hij wj + O(∥W ∥ ), 2 i=1 j=1 i

(51)

∈ Rp×p represents the blocks of the Hessian matrix.

W =0

By the ZGZ condition, we know that ∇wi fx (0) = 0 23

for all i.

(52)

Thus, the first-order term ∑(∇wi fx (0))⊺ wi vanishes. Consider the off-diagonal Hessian blocks Hij where i ≠ j. By definition: Hij =

∂ (∇wi fx (W )) ∣ . ∂wj W =0

(53)

By the assumption that ∇wi fx = 0 as long as wi = 0 irrespective of wj , taking the derivative with respect to wj yields Hij = 0

for all i ≠ j.

(54)

The cross-terms in the second-order expansion are eliminated. We are now left with only the diagonal terms of the second-order expansion: fx (W ) = fx (0) +

1 d ⊺ 3 ∑ w Hii wi + O(∥W ∥ ) 2 i=1 i

(55)

By the Permutation Symmetry condition, all neurons are interchangeable. This implies that the curvature of the function with respect to any single neuron i, evaluated at the symmetric origin W = 0, must be identical for all neurons. Therefore: H11 = H22 = ⋅ ⋅ ⋅ = Hdd . (56) Let us define a new matrix A(x) = 12 Hii . Substituting this into our expansion yields: d

i=1

fx (W ) = fx (0) + ∑ wi⊺ A(x)wi + O(∥W ∥3 ) = fx (0) + ∑ Tr[wi wi⊺ A(x)] + O(∥W ∥3 ).

(57)

This finishes the proof.

A.4

Proof of Proposition 1

We verify Theorem 1’s prerequisites sequentially. 1. Smoothness: Since the activation function ϕ ∈ C 3 , fx is at least three times continuously differentiable. 2. Permutation Symmetry: Let σ be an arbitrary permutation of the indices {1, . . . , d}. d

fx (wσ(1) , . . . , wσ(d) ) = ∑ vσ(i) ϕ(u⊺σ(i) x) = ∑ vi ϕ(u⊺i x) = fx (w1 , . . . , wd ).

(58)

i=1

3. ZGZ: The gradient with respect to the combined neuron parameter wi consists of ∇ui fx = vi ϕ′ (u⊺i x)x, ∇vi fx = ϕ(u⊺i x).

(59)

∇ui fx ∣u =0,v =0 = 0 ⋅ ϕ′ (0)x = 0, ∇vi fx ∣u =0,v =0 = ϕ(0) = 0.

(60)

Setting ui = wi = 0 gives i

Therefore, the two-layer MLP satisfies all assumptions in Theorem 1. Finally let us compute A(x) = 21 Hii , where Hii = ∇2wi fx ∣w =0 is the Hessian matrix with respect to the i-th i neuron’s parameters wi = [u⊺i , vi ]⊺ . The Hessian Hii can be written as: ∇2ui fx Hii = [ ∇vi ∇ui fx

R ∇ui ∇vi fx RRRR ] R ∇2vi fx RRRR R

(61)

ui =0,vi =0

Top-left block (∇2ui fx ): ∇ui (∇ui fx ) = ∇ui (vi ϕ′ (u⊺i x)x) = vi ϕ′′ (u⊺i x)xx⊺ Evaluating at wi = 0 (which sets vi = 0), this term vanishes: ∇2ui fx ∣0 = 0k×k . 24

(62)

Bottom-right block (∇2vi fx ): ∇vi (∇vi fx ) = ∇vi (ϕ(u⊺i x)) = 0

(63)

Off-diagonal blocks (∇vi ∇ui fx ): Taking the derivative of ∇ui fx with respect to vi : ∇vi (vi ϕ′ (u⊺i x)x) = ϕ′ (u⊺i x)x

(64)

Evaluating at wi = 0 (which sets ui = 0), we obtain ϕ (0)x. By the symmetry of second derivatives, the top-right block is its transpose ϕ′ (0)x⊺ . Assembling the blocks yields the Hessian evaluated at the origin: ′

0 Hii = [ ′ k×k ⊺ ϕ (0)x

ϕ′ (0)x ], 0

(65)

which finishes the proof.

A.5

Proof of Proposition 2

1. Smoothness: The softmax function is infinitely differentiable everywhere in its domain. Thus, fx is C ∞ . 2. Permutation Symmetry: Let σ be an arbitrary permutation of the head indices {1, . . . , d}. d

fx (wσ(1) , . . . , wσ(d) ) = ∑ Headσ(i) (x, X) = ∑ Headi (x, X) = fx (w1 , . . . , wd ) i=1

(66)

i=1

3. ZGZ: Let Ai ∈ RN denote the softmax attention weights. V Wi = 0, WiQ = 0, and WiK = 0:

Evaluating the gradients at wi = 0 implies vi = 0,

∇vi fx ∣w =0 = WiV XAi ∣W V =0 = 0 i

(67)

∇WiV fx ∣w =0 = vi A⊺i X ⊺ ∣v =0 = 0 i

(68)

For the query and key matrices, the gradients involve the derivative of the softmax function, which we denote as a Jacobian matrix JAi . By the chain rule: ∇W Q fx = (. . . JAi . . . ) × (vi⊺ WiV X), ∇WiK fx = (. . . JAi . . . ) × (vi⊺ WiV X).

(69)

Because evaluating at vi = WiV = 0, (vi⊺ WiV X) becomes zero. Consequently, regardless of the Jacobian of the softmax, the overall gradients for the query and key matrices vanish: ∇W Q fx ∣w =0 = 0, i

∇WiK fx ∣w =0 = 0 i

(70)

We conclude that ∇wi fx ∣wi =0 = 0. All assumptions in Theorem 1 are satisfied. Finally let us compute A(X) = 21 Hii , where Hii = ∇2wi fx ∣w =0 . The output of the i-th head is: i

hi (wi ) = vi⊺ WiV X ⋅ softmax (

X ⊺ (WiK )⊺ WiQ x √ ) dk

(71)

X ⊺ (W K )⊺ W Q x

i i √ ∈ RN . Notice that L(wi ) is a homogeneous polynomial of Let the pre-softmax logits be L(wi ) = dk degree 2 with respect to wi . The Taylor expansion of the softmax function around L = 0 is:

1 1 1N + JS (0)L + O(∥L∥2 ) = 1N + O(∥wi ∥2 ) N N Now, substitute this back into the head formulation hi (wi ), we get: softmax(L) =

hi (wi ) = vi⊺ WiV X (

1 1 1N + O(∥wi ∥2 )) = vi⊺ WiV X1N + O(∥wi ∥4 ). N N

(72)

(73)

The only second-order term in the expansion comes from vi⊺ WiV ( N1 X1N ) = vi⊺ WiV xavg , where xavg = N1 X1N . To express this as a bilinear quadratic form wi⊺ A(X)wi , we can vectorize the term as: vi⊺ WiV xavg = vi⊺ vec(Idv WiV xavg ) = vi⊺ (x⊺avg ⊗ Idv )vec(WiV )

(74)

Placing this in the 4 × 4 block matrix corresponding to Q, K, V, v yields the required matrix A(X) and finishes the proof. 25

A.6

Proof of Proposition 3

Similarly to Proposition 2, smoothness and permutation symmetry are naturally satisfied. Let L = √1d X ⊺ (WiK )⊺ WiQ x ∈ RN . Evaluating the gradient with respect to WiQ at the origin: k

∇W Q fx ∣w =0 ∝ i

∂L

∣ K =0 ∂WiQ Wi =0

(75)

By symmetry, ∇WiK fx ∣w =0 = 0. Then the ZGZ condition holds. i We compute the second-order term of the Taylor expansion of the i-th head around wi = 0. The expansion of the softmax function around L = 0 is: softmax(L) =

1 1N + JS (0)L + O(∣L∣2 ), N

(76)

where JS (0) = N1 IN − N12 1N 1⊺N ∈ RN ×N is the Jacobian of the softmax function at the origin. Substituting L, the head’s output becomes: hi (wi ) = c⊺ X

X ⊺ (WiK )⊺ WiQ x 1 √ 1N + c⊺ XJS (0) + O(∣wi ∣4 ) N dk

(77)

The first term is a constant term. The second-order term is: 1 E2 ∶= √ c⊺ XJS (0)X ⊺ (WiK )⊺ WiQ x dk

(78)

Notice that the term XJS (0)X ⊺ inside the expression is the sample covariance matrix: XJS (0)X ⊺ =

1 1 1 XX ⊺ − ( X1N ) ( 1⊺N X ⊺ ) = ΣX N N N

(79)

Define u = ΣX c ∈ RD . We can rewrite the scalar E2 as: 1 E2 = √ u⊺ (WiK )⊺ WiQ x dk

(80)

Using the Kronecker product property vec(ABC) = (C ⊺ ⊗ A)vec(B), we rewrite E2 : 1 1 E2 = √ (WiK u)⊺ (WiQ x) = √ vec(WiK )⊺ (ux⊺ ⊗ Idk )vec(WiQ ) dk dk

(81)

Thus: 1 0 A(X) = √ [ ⊺ 2 dk ux ⊗ Idk

xu⊺ ⊗ Idk ], 0

(82)

which finishes the proof.

A.7

Proof of Proposition 4

Smoothness: The function is C ∞ as it is a composition of linear transformations and the softmax function. k Permutation Symmetry: Notice that the pre-softmax projection depends on the product (W K )⊺ W Q = ∑di=1 wiK (wiQ )⊺ . Any permutation σ of the neuron indices {1, . . . , dk } leaves the matrix product unchanged. Thus the model is invariant to the permutation of its hidden projection dimensions. ZGZ: As derived in the last section, the relevant part of the model is: 1 dk 1 E(W ) = √ u⊺ (W K )⊺ W Q x = √ ∑(u⊺ wiK )(x⊺ wiQ ) dk dk i=1

(83)

where u = ΣX c. Taking the gradient with respect to the components of the i-th neuron wi = [wiQ , wiK ]⊺ : ∇wQ fx ∝ (u⊺ wiK )x

(84)

∇wiK fx ∝ (x⊺ wiQ )u

(85)

Evaluating these gradients at wi = 0 yields zero vectors regardless of the states of wj (j ≠ i). The ZGZ condition holds. Deriving A(X): We focus on the second-order term for the i-th neuron: 1 1 √ (u⊺ wiK )(x⊺ wiQ ) = √ (wiK )⊺ ux⊺ wiQ dk dk

(86)

We need to express this in the quadratic form wi⊺ A(X)wi , Which gives: 1 0 A(X) = √ [ ⊺ 2 dk ux

xu⊺ ] 0

(87)

and finishes the proof.

A.8

Proof of Proposition 5

Smoothness: Assuming ψ and E are at least three times continuously differentiable, their product and sum ensure that the model fx is C 3 with respect to W . Permutation Symmetry: Let σ be an arbitrary permutation of the expert indices {1, . . . , d}. Evaluating the model with permuted parameters yields: d

⊺ fx (wσ(1) , . . . , wσ(d) ) = ∑ ψ(rσ(i) x)E(x; Θσ(i) ) = ∑ ψ(ri⊺ x)E(x; Θi ) = fx (w1 , . . . , wd ). i=1

(88)

i=1

ZGZ: We evaluate the partial gradients with respect to the components of wi at wi = 0 (which implies ri = 0 and Θi = 0). Taking the partial derivative with respect to ri yields: ∇ri fx = ψ ′ (ri⊺ x)E(x; Θi )x

(89)

= ψ ′ (0) ⋅ 0 ⋅ x = 0

(90)

∇Θi fx = ψ(ri⊺ x)∇Θi E(x; Θi )

(91)

= ψ(0) ⋅ ∇Θi E(x; Θi ) = 0 ⋅ ∇Θi E(x; Θi ) = 0

(92)

Since E(x; 0) = 0, we obtain: ∇r i f x ∣

wi =0

Taking the partial derivative with respect to Θi gives:

Since ψ(0) = 0, we have: ∇Θi fx ∣

wi =0

Thus, the ZGZ condition holds. The computation of A(x) is similar to that in Proposition 1.

A.9

Proof of Proposition 6

The proof is identical to that of Proposition 1.

A.10

Proof of Theorem 2

Let us concatenate the weight vectors wi ∈ Rp into W = [w1 , w2 , . . . , wd ] ∈ Rp×d . Let 1 ∈ Rd be the vector of all ones. With this notation, we can rewrite the summary statistics as µ = ∑di=1 wi = W 1 and M = ∑di=1 wi wi⊺ = W W ⊺ . Accordingly, the NQF formulation (20) can be rewritten as: fx (W ) = fx (0) + g(x)⊺ W 1 + 1⊺ W ⊺ B(x)W 1 + Tr[W ⊺ A(x)W ].

(93)

Taking the derivative of fx (W ) with respect to W , we have: ∇W (g(x)⊺ W 1) = g(x)1⊺ , ∇W (1⊺ W ⊺ B(x)W 1) = 2B(x)W 11⊺ = 2B(x)µ1⊺ and ∇W Tr[W ⊺ A(x)W ] = 2A(x)W . Summing these terms yields the gradient: ∇W fx (W ) = (g(x) + 2B(x)µ)1⊺ + 2A(x)W.

(94)

By the chain rule, the gradient of the loss L with respect to W over a mini-batch B is: ∇W L = ∑ ℓ′x ∇W fx (W ).

(95)

x∈B

Substituting (94) and using the definitions of v and H: ∇W L = v1⊺ + HW.

(96)

∆W = −η∇W L = −η(v1⊺ + HWt ).

(97)

∆µ = ∆W 1 = −η(v1⊺ 1 + HWt 1) = −η(d ⋅ v + Hµt ).

(98)

Under SGD, the weight matrix is updated as:

We obtain the dynamics for µ:

The update ∆M is computed as: ∆M = (W + ∆W )(W + ∆W )⊺ − W W ⊺ = W ∆W ⊺ + ∆W W ⊺ + ∆W ∆W ⊺ .

(99)

We evaluate these three components separately: ∆W W ⊺ = −η(v1⊺ + HW )W ⊺ = −η(vµ⊺ + HM ). W ∆W ⊺ = (∆W W ⊺ )⊺ = −η(µv ⊺ + M H) since H is symmetric. For the second-order term: ∆W ∆W ⊺ = η 2 (v1⊺ + HW )(1v ⊺ + W ⊺ H) = η 2 (d ⋅ vv ⊺ + vµ⊺ H + Hµv ⊺ + HM H).

(100)

Summing all three components, we establish: ∆M = −η(vµ⊺ + µv ⊺ + HM + M H) + η 2 (d ⋅ vv ⊺ + vµ⊺t H + Hµv ⊺ + HM H).

(101)

Therefore, if two NQFs satisfy Ma (0) = Mb (0) and µa (0) = µb (0) at initialization, we have Ma (t) = Mb (t) and µa (t) = µb (t). This completes the proof.

A.11

Proof of Theorem 3

Let us define the subspace V ⊆ Rp spanned by the data and g(x), A(x), B(x): V = span ⋃ ({g(x)} ∪ Col(A(x)) ∪ Col(B(x))).

(102)

x∈X

Let kV = dim(V) ≤ p, and let V ∈ Rp×k be a matrix whose columns form an orthonormal basis for V, where we recall that p ∶= dim wi . By definition, for any x, we have V V ⊺ g(x) = g(x), V V ⊺ A(x)V V ⊺ = A(x), and V V ⊺ B(x)V V ⊺ = B(x). Thus, to match the initial outputs and dynamics, we construct the smaller model within V such that its projected statistics match the original model: µ̃(0) = V V ⊺ µ(0), M̃ (0) = 28

d V V ⊺ M (0)V V ⊺ . d′

(103)

We construct each weight vector of the smaller model as wi′ = V ( d1′ V ⊺ µ(0) + ui ), where u1 , . . . , ud′ ∈ RkV satisfies: d′ d′ 1 ⊺ d ⊺ ⊺ ⊺ (104) ∑ ui = 0, and ∑ ui ui = ′ (V M (0)V − V µ(0)µ(0) V ) =∶ CV , d d i=1 i=1 which leads to (103). CV ∈ Rk×k is inherently positive semi-definite (PSD) and can be factorized as CV = LL⊺ for some L ∈ Rk×k . To find u1 , . . . , ud′ , consider the all-ones vector 1kV +1 ∈ RkV +1 . Its orthogonal complement is a kV -dimensional subspace in RkV +1 . Let the rows of a matrix Z ∈ RkV ×(kV +1) form an orthonormal basis for this subspace, satisfying Z1k+1 = 0 and ZZ ⊺ = IkV . We can set U ∶= [u1 , . . . , uk+1 ] ∶= LZ. It follows that: kV +1

∑ ui = U 1k+1 = LZ1kV +1 = 0,

and

U U ⊺ = LZZ ⊺ L⊺ = LIkV L⊺ = CV .

(105)

i=1

Thus, d′ = kV + 1 neurons are sufficient to match the effective initial statistics. Assume that at any step t, the statistics of the smaller model remain within the subspace V and match the projected original statistics: d (106) µ̃t = V V ⊺ µt , M̃t = ′ V V ⊺ Mt V V ⊺ d Due to the subspace definitions (e.g., V V ⊺ g(x) = g(x)), the output of the smaller model is: f˜x = fx (0) + g(x)⊺ µ̃t + µ̃⊺t B(x)µ̃t + Tr[M̃t A(x)] d d′ = fx (0) + g(x)⊺ V V ⊺ µt + µ⊺t V V ⊺ B(x)V V ⊺ µt + Tr [( ′ V V ⊺ Mt V V ⊺ ) ( A(x))] d d = fx (0) + g(x)⊺ µt + µ⊺t B(x)µt + Tr[Mt A(x)] = fx .

(107)

Thus, as long as the statistics maintain this proportional relationship, the outputs are identical. This satisfies the equivalence at initialization (t = 0). Because f˜x = fx , the loss derivative with respect to the output is identical: ℓ̃′x = ℓ′x . According to Theorem 2, the gradient variables for the smaller model evaluate to: ṽ = ∑ ℓ̃′x (g̃(x) + 2B̃(x)µ̃t ) = ∑ ℓ′x (g(x) + 2B(x)µt ) = v, H̃ = ∑ 2ℓ̃′x Ã(x) = x∈B

(108)

x∈B

d′ d′ ′ ∑ 2ℓx A(x) = H, d x∈B d

(109)

where we use B̃(x)µ̃t = B(x)V V ⊺ µt = B(x)µt because B(x)V V ⊺ = B(x) by definition of V. since v is a linear combination of g(x) and the columns of B(x), we have v = V V ⊺ v. Similarly, the symmetric matrix H is formed by A(x), meaning H = V V ⊺ HV V ⊺ . Applying Theorem 2 to the smaller model, we compute the update ∆µ̃: ∆µ̃t = −η̃(d′ ⋅ ṽ + H̃ µ̃t ) d d′ = − ( ′ η) (d′ ⋅ V V ⊺ v + ( V V ⊺ HV V ⊺ ) V V ⊺ µt ) d d = −ηV V ⊺ (d ⋅ v + Hµt ) = V V ⊺ ∆µt .

(110)

For the second-moment statistic, we have: ∆M̃t = −η̃(ṽ µ̃⊺t + µ̃t ṽ ⊺ + H̃ M̃t + M̃t H̃) + η̃ 2 (d′ ⋅ ṽṽ ⊺ + ṽ µ̃⊺t H̃ + H̃ µ̃t ṽ ⊺ + H̃ M̃t H̃). By using (106), the first-order term becomes: d − ( ′ η) V V ⊺ (vµ⊺t + µt v ⊺ + HMt + Mt H)V V ⊺ . d 29

(111)

The second term becomes d 2 d′ ( ′ η) V V ⊺ (d ⋅ vv ⊺ + vµ⊺t H + Hµt v ⊺ + HMt H)V V ⊺ . d d Summing these up, we obtain: ∆M̃t =

d V V ⊺ ∆Mt V V ⊺ . d′

(112)

Thus, at step t + 1, we have µ̃t+1 = µ̃t + ∆µ̃t = V V ⊺ µt+1 and M̃t+1 = dd′ V V ⊺ Mt+1 V V ⊺ . By mathematical induction, these projection relations hold for all training steps t ≥ 0, and we conclude that fd′ = fd holds throughout training. This completes the proof.

A.12

Proof of Proposition 7

By (29) and M̃˙ = −H̃ M̃ − M̃ H̃, we have: d M̃ij = −(H̃ii + H̃jj )M̃ij . dt

(113)

Because H̃ii and H̃jj do not depend on the off-diagonal elements, it can be integrated directly to yield t

M̃ij (t) = M̃ij (0) exp (− ∫

A.13

(H̃ii (τ ) + H̃jj (τ )) dτ ) .

(114)

Proof of Theorem 4

Recall that the residual for the µ-th sample can be written as: ∆µ (t) = Tr[W (t)W (t)⊺ cµ A] − yµ = cµ Tr[M (t)A] − yµ and H(t) =

4 m ∑ (cµ Tr[M (t)A] − yµ )cµ A = (αTr[M (t)A] − β)A, m µ=1

(115)

(116)

m m 4 4 where we use α = m ∑µ=1 c2µ and β = m ∑µ=1 cµ yµ . According to (27), the gradient flow for W (t) is:

Ẇ (t) = −H(t)W (t) = − (αTr[M (t)A] − β) AW (t)

(117)

To solve this equation, we introduce an auxiliary scalar variable ξ(t) defined as the integrated effective residual: t

ξ(t) = ∫

(αTr[M (τ )A] − β) dτ.

(118)

˙ = αTr[M (t)A] − β. Projecting the parameter matrix into the shared eigenbasis by By definition, ξ(0) = 0 and ξ(t) ⊺ defining W̃ (t) = P W (t), we can rewrite the dynamics of W̃ (t) as: ˙ (t) = P ⊺ Ẇ (t) = −ξ(t)P ⊺ ˙ ˙ W̃ AW (t) = −ξ(t)Λ W̃ (t).

(119)

Since Λ is a diagonal matrix, this system of differential equations is decoupled row-by-row. Integrating with respect to time from 0 to t yields the solution in the projected space: W̃ (t) = exp(−ξ(t)Λ)W̃ (0).

(120)

Multiplying by P on the left, we obtain the closed-form solution for W : W (t) = P exp(−ξ(t)Λ)P ⊺ W (0)

(121)

Next, we evaluate M (t) = W (t)W (t)⊺ to close the system for ξ(t). Projecting M (t) into the eigenbasis yields: M̃ (t) = P ⊺ M (t)P = W̃ (t)W̃ (t)⊺ = exp(−ξ(t)Λ)M̃ (0) exp(−ξ(t)Λ), 30

(122)

Since exp(−ξ(t)Λ) is a diagonal matrix, the diagonal entries of M̃ (t) decouple from the off-diagonal entries and depend only on the diagonal elements of M̃ (0). They evolve as: zk (t) = M̃kk (t) = (exp(−ξ(t)Λ)M̃ (0) exp(−ξ(t)Λ))kk = zk (0) exp(−2λk ξ(t))

(123)

The trace term can then be simplified via Tr[M (t)A] = Tr[M̃ (t)Λ] = ∑pj=1 λj zj (t). Substituting the expression of zj (t) into the derivative of ξ(t) results in a self-contained ODE: p dξ = α ∑ λj zj (0) exp(−2λj ξ) − β. dt j=1

(124)

Separating variables and integrating from 0 to ξ(t) leads to the implicit equation: ξ(t)

. α ∑pj=1 λj zj (0) exp(−2λj s) − β

(125)

⎛4 m ⎞ ⊺ 4 m ∑ ∆µ (t)A(xµ ) = P ∑ ∆µ (t)Λµ P . m µ=1 ⎝ m µ=1 ⎠

(126)

t=∫

This completes the proof.

A.14

Proof of Theorem 5

Recall that ∆µ (t) = Tr[M (t)A(xµ )] − yµ and H(t) =

To solve the dynamics Ẇ (t) = −H(t)W (t), we introduce a set of m independent auxiliary scalar variables, ξµ (t), defined as the integrated effective residual for each sample: t

ξµ (t) = ∫

∆µ (τ )dτ.

(127)

By definition, ξµ (0) = 0 and ξ˙µ (t) = ∆µ (t). Projecting the parameter matrix into the shared eigenbasis by defining W̃ (t) = P ⊺ W (t), the dynamics decouple into: m ⎞ ˙ (t) = −P ⊺ H(t)P W̃ (t) = − ⎛ 4 W̃ ∑ ξ˙µ (t)Λµ W̃ (t). ⎝ m µ=1 ⎠

(128)

m 4 By defining the diagonal matrix Σ(t) = m ∑µ=1 ξµ (t)Λµ , we can integrate this ODE to obtain

W̃ (t) = exp(−Σ(t))W̃ (0).

(129)

Multiplying by P on the left yields the closed-form representation for W (t). Next, we track the evolution of M (t) = W (t)W (t)⊺ . Projecting M (t) into the eigenbasis yields: M̃ (t) = P ⊺ M (t)P = exp(−Σ(t))M̃ (0) exp(−Σ(t)).

(130)

Because Σ(t) is a diagonal matrix, the diagonal entries of M̃ (t) evolve independently of the off-diagonal elements. The feature magnitudes evolve as: zk (t) = M̃kk (t) = zk (0) exp (−2Σkk (t)) = zk (0) exp (−

8 m ∑ ξν (t)λν,k ) . m ν=1

(131)

Now we utilize the data orthogonality condition A(xµ )A(xν ) = 0. This implies Λµ Λν = 0, meaning that for any specific feature index k, at most one dataset matrix can have a non-zero eigenvalue. Therefore, if λµ,k ≠ 0 for a specific sample µ, then λν,k = 0 for all ν ≠ µ.

This orthogonality decouples the sum in the exponent. When substituting zk (t) back into the derivative of the residual ξ˙µ (t), we only need to consider the non-zero elements λµ,k : p

ξ˙µ (t) = Tr[M̃ (t)Λµ ] − yµ = ∑ λµ,k zk (t) − yµ .

(132)

k=1

Applying the orthogonality property to the exponent of zk (t), we get: p dξµ 8 = ∑ λµ,k zk (0) exp (− λµ,k ξµ (t)) − yµ . dt k=1 m

(133)

Thus, the O(pd)-dimensional gradient flow reduces to m one-dimensional ODEs. Separating variables and integrating from 0 to ξµ (t) provides the implicit solution for each µ independently: ξµ (t)

t=∫

. p 8 λµ,k s) − yµ ∑k=1 λµ,k zk (0) exp (− m

(134)

This completes the proof.

A.15

Proof of Theorem 6

We start with (32): ⎡ ⎤ ⎢⎛ 8 m ⎥ ⎞ p ⎛8 m ⎞ żk (t) = zk (t) ⎢⎢ ∑ λµ,k yµ − ∑ ∑ λµ,k λµ,j zj (t)⎥⎥ . ⎠ j=1 ⎝ m µ=1 ⎠ ⎢⎝ m µ=1 ⎥ ⎣ ⎦

(135)

m 8 Using the definitions of rk and Ckj ∶= m ∑µ=1 λµ,k λµ,j , the dynamics can be written as: p ⎞ ⎛ żk (t) = zk (t) rk − ∑ Ckj zj (t) . ⎠ ⎝ j=1

(136)

Applying the orthogonal features condition, we have Ckj = 0 for all k ≠ j. This condition decouples the feature interactions. The ODE simplifies to the classic logistic equation: żk (t) = zk (t)(rk − Ckk zk (t)).

(137)

We have zk (t) = 0 if zk (0) = 0. With the initial condition zk (0) > 0, integrating from 0 to t yields the analytic solution: rk zk (0) zk (t) = (138) Ckk zk (0) + (rk − Ckk zk (0)) exp(−rk t) if rk ≠ 0. If rk = 0, the solution is zk (t) = zk (0)/(1 + Ckk zk (0)t). m 4 Finally, we reconstruct the parameter matrix W (t). Recall that H̃(t) = m ∑µ=1 ∆µ (t)Λµ is a diagonal mat ˙ (t) = trix for all t. Let Σ(t) = ∫0 H̃(τ )dτ , which is also diagonal. The dynamics in the projected space is W̃ −H̃(t)W̃ (t), yielding W̃ (t) = exp(−Σ(t))W̃ (0). Since Σ(t) is diagonal, the k-th row of W̃ (t) is scaled by the factor ⊺ exp(−Σkk (t)). We also know from the diagonal elements of M̃ (t) √ = W̃ (t)W̃ (t) that zk (t) = zk (0) exp(−2Σkk (t)). For any feature with zk (0) > 0, we have exp(−Σkk (t)) = zk (t)/zk (0). For a feature with zk (0) = 0, the k-th row of W̃ (0) √is a zero vector, which ensures that the k-th row of W̃ (t) remains zero for all t ≥ 0. Thus, defining dk (t) = zk (t)/zk (0) for zk (0) ≠ 0 and dk (t) = 0 otherwise consistently captures the row-wise scaling: W̃ (t) = D(t)W̃ (0). Multiplying by P on the left, we arrive at the trajectory for W : W (t) = P D(t)P ⊺ W (0). This completes the proof.

(139)

A.16

Proof of Theorem 7

m 4 Recall that the gradient flow of the parameters is Ẇ (t) = −H(t)W (t), where H(t) = m ∑µ=1 ∆µ (t)A(xµ ). By expanding the residual ∆µ (t) = Tr[M (t)A(xµ )] − yµ and isolating the label-dependent term Y , we have: m

1 Tr[M (t)A(xµ )]A(xµ ) − Y. µ=1 m

H(t) = 4 ∑

(140)

We evaluate the first term element-wise: p p ⎛ ⎞ ⎛ m 1 ⎞ 1 m =4∑∑ 4 ∑ Tr[M (t)A(xµ )]A(xµ ) ∑ Aij (xµ )Akl (xµ ) Mkl (t). m ⎠ ⎝ µ=1 m ⎠ ⎝ µ=1 k=1 l=1

(141)

Applying Assumption 2, this expression simplifies to: p p c 4 ∑ ∑ (δik δjl + δil δjk )Mkl (t) = 2c (Mij (t) + Mji (t)) = 4cMij (t), 2 k=1 l=1

(142)

where we utilized M (t) = M (t)⊺ . Thus, we have H(t) = 4cM (t) − Y . Substituting H(t) into Ṁ (t) = −H(t)M (t) − M (t)H(t), we obtain a matrix Riccati differential equation [2]: Ṁ = −(4cM − Y )M − M (4cM − Y ) = Y M + M Y − 8cM 2 .

(143)

To solve this non-linear matrix equation, we introduce M (t) = V (t)U (t)−1 , where V (0) = M (0) and U (0) = I. Differentiating this expression yields: Ṁ = V̇ U −1 − V U −1 U̇ U −1 = V̇ U −1 − M U̇ U −1 .

(144)

By matching the terms with the Riccati equation Ṁ = Y V U −1 − M (−Y U + 8cV )U −1 , the dynamics decouple into a system of two linear ordinary differential equations: V̇ = Y V,

(145)

U̇ = −Y U + 8cV.

(146)

The solution to (145) is given by V (t) = exp(Y t)M (0). Substituting V (t) into (146) gives: U̇ + Y U = 8c exp(Y t)M (0).

(147)

d Multiplying both sides by the integrating factor exp(Y t) leads to dt [exp(Y t)U (t)] = 8c exp(2Y t)M (0). Integrating from 0 to t and using U (0) = I, we get: t

exp(Y t)U (t) − I = 8c (∫

exp(2Y τ )dτ ) M (0) = 8cΦ(t)M (0).

(148)

0 t

Solving for U (t), we find U (t) = exp(−Y t) [I + 8cΦ(t)M (0)], which is always invertible because Φ(t) ∶= ∫0 exp(2Y τ )dτ is positive definite for t > 0, and M (0) = W (0)W (0)⊺ is positive semi-definite. Finally, reconstructing M (t) = V (t)U (t)−1 , we arrive at: −1

M (t) = exp(Y t)M (0) [I + 8cΦ(t)M (0)]

exp(Y t).

(149)

Furthermore, suppose Y and M (0) are simultaneously diagonalizable. Then, there exists an orthogonal matrix P such that Y = P ΓP ⊺ and M (0) = P Z(0)P ⊺ , where Γ and Z(0) are diagonal matrices containing the eigenvalues γk and zk (0), respectively. Since all matrices in the closed-form solution share the same eigenbasis, M (t) remains diagonal in this basis: M (t) = P Z(t)P ⊺ . For the k-th eigen-mode, the matrix functions reduce to scalar functions. The integral evaluates to Φk (t) = exp(2γk t)−1 for γk ≠ 0. Substituting these into (149) gives: 2γk zk (t) =

exp(2γk t)zk (0) k t)−1 1 + 8c ( exp(2γ ) zk (0) 2γk

(150)

Rearranging the terms, we recover the standard generalized logistic growth curve: zk (t) =

γk zk (0) . 4czk (0) + (γk − 4czk (0)) exp(−2γk t)

(151)

For the case where γk = 0, the integral simplifies to Φk (t) = ∫0 1dτ = t. Substituting this yields: zk (t) =

zk (0) . 1 + 8czk (0)t

(152)

This completes the proof.

A.17

Formal Statements of Section 6

Feature-wise Descent

The first part of the formal statement is the following theorem.

Theorem 8. Under the conditions of Theorems 6 (or Theorem 7 with M (0) and Y simultaneous diagonalizable), let ζ1 > ζ2 > ⋅ ⋅ ⋅ > ζp > 0, where we define ζk = rk in Theorem 6 (ζk = γk in Theorem 7). Then, for any fixed constant ρ ∈ (0, 1), let t∗k denote the time that zk (t∗k ) = ρzk (∞) with zk (∞) = limt→∞ zk (t). Let the initialization scales satisfy zk (0) = Θ(ϵ) for all k ∈ {1, . . . , p}, then as ϵ → 0+ , t∗k ∼

1 1 ln , aζk ϵ

(153)

where a = 1 in Theorem 6 (a = 2 in Theorem 7). Consequently, the gaps between consecutive feature activations diverge in the limit of small initialization, which exhibits saddle-to-saddle dynamics: lim(t∗k+1 − t∗k ) = ∞. (154) ϵ→0

Proof. Let the initialization scale be zk (0) = ck ϵ for all k, where ck > 0 is a constant and ϵ → 0. By Theorems 6 and 7, zk evolves according to zk (t) =

ζk zk (0) Ck zk (0) + (ζk − Ck zk (0)) exp(−aζk t)

(155)

where ζk = rk , a = 1, Ck = Ckk for Theorem 6, and ζk = γk , a = 2, Ck = 4c for Theorem 7. The non-zero steady-state saturation value is zk (∞) = ζk /Ck . Setting zk (t∗k ) = ρzk (∞) = ρ Cζkk , we substitute the initial condition and solve for t∗k : ρ

ζk ζ k ck ϵ = Ck Ck ck ϵ + (ζk − Ck ck ϵ) exp(−aζk t∗k )

(156)

Rearranging to isolate the exponential term yields: exp(−aζk t∗k ) =

Ck ck ϵ 1−ρ ( ) ζk − Ck ck ϵ ρ

(157)

Taking the logarithm and dividing by −aζk , we obtain the characteristic time: t∗k =

1 1 1 Ck ck 1 − ρ ln ( ) − ln ( ) aζk ϵ aζk ζk − Ck ck ϵ ρ

(158)

As ϵ → 0, the second term is O(1), and the time scales as: t∗k =

1 1 1 1 ln ( ) + O(1) ∼ ln . aζk ϵ aζk ϵ

This completes the proof. 34

(159)

Recall that rk and γk are essentially the correlation between input features and label, and Ckk ≥ 0 is essentially the input variances. The second part of the formal statement is the following. Theorem 9. Suppose the conditions of Theorem 8 hold. Consider the infinite-width limit p → ∞2 and assume the initialization satisfies zk (0) = ϵck k −β where the constants ck are bounded away from zero: 0 < cmin ≤ ck ≤ cmax < ∞. r2

γ2

k Let Eϵ (t) ∶= L(t) − L(∞) denote the excess loss. Define Vk ∶= 8Ckkk (under Theorem 6) or Vk ∶= 16c (under Theorem 7). Assume that ζk and Vk exhibit the following power-law decays:

ζk = cζ k −α2 , Vk = cw k −α1

(160)

for positive constants cζ , cw , α2 > 0. Furthermore, assume that α1 > 1 3 and β ≥ α1 − α2 to ensure finite initial excess loss. t . Then, under the small initialization limit ϵ → 0+ , the excess loss converges Define the rescaled time τ ∶= ln(1/ϵ) E(τ ) ∶= limϵ→0+ Eϵ (τ ln(1/ϵ)) for almost every τ and E(τ ) has a power-law decay α −1

E(τ ) = Θ (τ

1 − α

(161)

Proof. We first evaluate the empirical loss L(t) in terms of the features zk (t). 2 p m 1 In Theorem 6, the loss is L(t) = m ∑µ=1 (∑k=1 λµ,k zk (t) − yµ ) . Expanding the square and applying the or8 thogonality condition ∑µ λµ,j λµ,l = 0 for j ≠ l, the cross-terms vanish. Using the definitions Ckk = m ∑µ λ2µ,k and 8 rk = m ∑µ λµ,k yµ , we obtain: 1 p rk 2 L(t) = ∑ Ckk (zk (t) − ) + Const. (162) 8 k=1 Ckk 2

Since zk (∞) = rk /Ckk , the excess loss is Eϵ (t) ∶= L(t) − L(∞) = 81 ∑pk=1 Ckk (zk (t) − zk (∞)) . 2 m 1 Similarly, in Theorem 7, expanding L(t) = m ∑µ=1 (Tr[M (t)A(xµ )] − yµ ) and utilizing Assumption 2 yields L(t) = cTr[M (t)2 ] − 12 Tr[M (t)Y ] + Const. Because M (t) and Y are simultaneously diagonalizable, this simplifies 2

to L(t) = c ∑pk=1 (zk (t) − γ4ck ) + Const. With zk (∞) = γk /4c, the excess loss is Eϵ (t) = c ∑pk=1 (zk (t) − zk (∞)) . Therefore, in both regimes the excess loss at any time t can be written as 2

∞

Eϵ (t) = ∑ Vk (1 − k=1

zk (t) ) . zk (∞)

(163)

after taking p → ∞. β ≥ α1 − α2 guarantees that (1 − zk (0)/zk (∞))2 → 1 because c′ ϵk −β ζk c′k ϵk −β cζ k −α2 c′k cζ ϵ α1 −α2 −β zk (0) = = k = k , zk (∞) Vk cw k −α1 cw

(164)

where c′k ∶= c8k in Theorem 6 and c′k ∶= c4k in Theorem 7. Further α1 > 1 guarantees that at t = 0 the excess loss ∞ ∑k=1 Vk is finite. t We introduce the rescaled time τ = ln(1/ϵ) , or equivalently t = τ ln(1/ϵ). The fraction of unlearned feature magnitude, denoted as fk,ϵ (τ ) = 1 − zk (τzkln(1/ϵ)) , can be evaluated by substituting (155): (∞) zk (t) ζk ck ϵk −β /Ck = = −β zk (∞) ζk ck ϵk /Ck + (ζk /Ck − ck ϵk −β ) exp(−aζk t) 1 + (

1 ζk

ck k−β Ck ϵ

(165)

− 1) exp(−aζk τ ln(1/ϵ))

ζk Let C̃k (ϵ) ∶= ck k−β − 1. For sufficiently small ϵ, C̃k (ϵ) > 0. The fraction becomes: Ck ϵ

fk,ϵ (τ ) = 1 −

1 C̃k (ϵ)ϵaζk τ = . 1 + C̃k (ϵ)ϵaζk τ 1 + C̃k (ϵ)ϵaζk τ

(166)

Notice that the leading order behavior of the term C̃k (ϵ)ϵaζk τ is governed by ϵaζk τ −1 . Let τk∗ = aζ1k . We analyze the pointwise limit of fk,ϵ (τ ) as ϵ → 0+ : 2 Note that under Theorem 7 this requires m → ∞ simultaneously to satisfy Assumption 2

2 3 Note that under Theorem 7, V = ζk , which enforces α = 2α . Thus, we require α > 0.5 in that regime. 1 2 2 k 16c

• If τ > τk∗ , then aζk τ − 1 > 0. As ϵ → 0+ , C̃k (ϵ)ϵaζk τ ∼ ϵaζk τ −1 → 0. Hence, fk,ϵ (τ ) → 0. • If τ < τk∗ , then aζk τ − 1 < 0. As ϵ → 0+ , C̃k (ϵ)ϵaζk τ ∼ ϵaζk τ −1 → ∞. Hence, fk,ϵ (τ ) → 1. Therefore, for almost every τ , fk,ϵ converges to an indicator function: limϵ→0+ fk,ϵ (τ ) = I(τ < τk∗ ). Notice that for both Theorem 6 and Theorem 7, the relationship Ck = Θ(ζk2 /Vk ) = Θ(k α1 −2α2 ) holds (with α1 = 2α2 specifically under Theorem 7). Substituting this scaling into C̃k (ϵ) yields: Const β+α2 −α1 k − 1. ϵ

C̃k (ϵ) ≥

(167)

Because we assumed β ≥ α1 − α2 , the exponent of k is non-negative, and thus there exists a sufficiently small ϵ > 0 such that C̃k (ϵ) > 0 uniformly for all k ≥ 1. This ensures that 0 ≤ fk,ϵ (τ ) ≤ 1 across all features. Consequently, the terms in (163) are bounded by Vk fk,ϵ (τ )2 ≤ Vk . Since ∑∞ k=1 Vk < ∞, we can apply the Lebesgue dominated convergence theorem to interchange the limit and the summation: 2

∞

E(τ ) ∶= lim+ Eϵ (τ ln(1/ϵ)) = ∑ Vk ( lim+ fk,ϵ (τ )) = ∑ Vk ⋅ I(τ < τk∗ ). ϵ→0

ϵ→0

k=1

(168)

k=1

The condition τ < τk∗ translates to τ < acζ k1−α2 , which can be inverted to define a threshold: 1

k > (acζ τ ) α2 ∶= k ∗ (τ ).

(169)

Let K(τ ) = ⌊k ∗ (τ )⌋ + 1 be the smallest integer greater than k ∗ (τ ). The excess loss equals the tail sum: ∞ −α ∑ cw k 1 .

E(τ ) =

(170)

k=K(τ )

Because cw x−α1 is monotonically decreasing, the tail sum is bounded by the Riemann integrals: ∞

∞

∫

K(τ )

cw x−α1 dx ≤

∞

−α ∑ cw k 1 ≤ ∫ k=K(τ )

K(τ )−1

cw x−α1 dx.

Evaluating the integrals yields: cw cw −(α −1) −(α −1) (K(τ )) 1 ≤ E(τ ) ≤ (K(τ ) − 1) 1 . α1 − 1 α1 − 1

(171)

(172)

As τ → ∞, the threshold scales as K(τ ) ∼ k ∗ (τ ) = (acζ τ )1/α2 . Substituting k ∗ (τ ) into the bounds gives the asymptotics: 1

E(τ ) = Θ ((τ α2 )

−(α1 −1)

α −1

) = Θ (τ

1 − α

(173)

This completes the proof. Sample-wise Descent

The first part of the formal statement is the following theorem.

Theorem 10. Under the conditions of Theorem 5, for each sample µ ∈ {1, . . . , m}, let λµ,max = maxk λµ,k and assume that yµ > 0. Define the effective growth rate for each sample as: ζµ ∶=

8 λµ,max yµ m

(174)

m 1 Assume that ζ1 > ζ2 > ⋅ ⋅ ⋅ > ζm > 0. Let L(t) = m ∑µ=1 ∆µ (t)2 be the empirical MSE and the initialization satisfies zk (0) = Θ(ϵ) for all k ∈ {1, . . . , p}. t Define the rescaled time τ ∶= ln(1/ϵ) and the characteristic timescale τµ∗ ∶= ζ1µ . Then, in the small initialization + limit ϵ → 0 , the empirical loss converges pointwise (for almost every τ ) to a saddle-to-saddle trajectory transitioning through m plateaus: 1 1 m 2 ∗ lim+ L (τ ln ) = (175) ∑ y ⋅ I(τ < τµ ), ϵ→0 ϵ m µ=1 µ

where I(⋅) denotes the indicator function. 36

8 Proof. According to Theorem 5, for a specific sample µ, we denote zk (t) ∶= zk (0) exp (−λ̃k ξµ (t)) and λ̃k ∶= m λµ,k . Let Kµ = {k ∣ λµ,k = λµ,max } be the set of indices for the maximum eigenvalue. We define

Z(t) ∶= ∑ λµ,max zk (t) = Z(0) exp (−λ̃max ξµ (t))

(176)

k∈Kµ

where Z(0) ∶= ∑k∈Kµ λµ,max zk (0) = Θ(ϵ). We first prove that ξµ (t) ≤ 0. By definition, ξµ (0) = 0. The initial residual is ∆µ (0) = ∑pk=1 λµ,k zk (0) − yµ = O(ϵ) − yµ . Since yµ > 0, for sufficiently small ϵ, ∆µ (0) < 0. Because ξ˙µ (t) = ∆µ (t), ξµ (t) is initially decreasing. Suppose there exists a first time t1 > 0 such that ξµ (t1 ) = 0. At this instant, zk (t1 ) = zk (0) for all k, which implies ξ˙µ (t1 ) = ∆µ (0) < 0. Since ξµ (t) < 0 for t ∈ (0, t1 ), reaching zero at t1 from below requires ξ˙µ (t1 ) ≥ 0, which contradicts ξ˙µ (t1 ) < 0. Because ξµ (t) ≤ 0, the evolution of any feature depends on the sign of its eigenvalue. Since λµ,max > 0, the exponent −λ̃max ξµ (t) ≥ 0, which guarantees Z(t) ≥ Z(0) = Cϵ for all t ≥ 0. For any sub-dominant positive feature (k ∉ Kµ , λµ,k > 0), we can express its dynamics in terms of Z(t) by eliminating ξµ (t): zk (t) = zk (0) (

Z(t) ) Z(0)

γk

(177)

where γk ∶= λµ,k /λµ,max ∈ (0, 1). Conversely, for any negative feature (λµ,k < 0), the exponent −λ̃k ξµ (t) ≤ 0. This implies that zk (t) ≤ zk (0) for all t ≥ 0. We define the residual ∆µ (t) ∶= Z(t) + R(Z, ϵ) − yµ , where R(Z, ϵ) ∶= R+ (Z, ϵ) + R− (ϵ) and R+ (Z, ϵ) =

∑

λµ,k ck ϵ (

λµ,k >0,k∉Kµ

Z γk ) , Aϵ

R− (ϵ) = ∑ λµ,k zk (t).

(178)

λµ,k <0

Before proceeding, we establish a global bound for Z(t). Since the dynamics follow a gradient flow, the total loss is non-increasing (L(t) ≤ L(0)), which guarantees that the residual ∆µ (t) is globally bounded, i.e., ∣∆µ (t)∣ ≤ C∆ for some constant C∆ > 0. Rearranging the residual equation yields: Z(t) + R+ (Z, ϵ) = ∆µ (t) + yµ − R− (ϵ). We have already established that zk (t) ≤ zk (0) for all features with λµ,k < 0. This implies that the negative contribution R− (ϵ) is bounded: −R− (ϵ) ≤ O(ϵ). Consequently, the right-hand side of the equation is bounded by a constant MZ ∶= C∆ + yµ + O(ϵ). Since both Z(t) and R+ (Z, ϵ) are sums of non-negative terms (as λµ,k > 0 and zk (t) ≥ 0), we must have Z(t) ≤ MZ for all t ≥ 0. Let δ ∶= minλµ,k >0 (1 − γk ) > 0. Since MZ ≥ Z ≥ Cϵ, ∣R+ (Z, ϵ)∣ ≤ C1 ϵδ . Since zk (t) ≤ ck ϵ for negative eigenvalues, ∣R− (ϵ)∣ ≤ C2 ϵ. Combining these, for any bounded trajectory of Z(t), we can define a uniform bound relative to yµ : κ(ϵ) ∶= supt ∣R(Z(t), ϵ)∣/yµ = O(ϵmin(δ,1) ), which vanishes as ϵ → 0+ . Now, summing (31) over k ∈ Kµ yields: Ż(t) = −

8 Z(t) + R(Z, ϵ) λµ,max ∆µ (t)Z(t) = ζµ Z(t) (1 − ) m yµ

(179)

8 λµ,max yµ . With the definition of κ(ϵ), for Z ∈ [0, yµ ] we have: where ζµ = m

ζµ Z (1 − κ(ϵ) −

Z Z ) ≤ Ż ≤ ζµ Z (1 + κ(ϵ) − ) yµ yµ

(180)

Let Zlower (t) and Zupper (t) denote the solutions to these bounding logistic equations with the initial condition Zlower (0) = Zupper (0) = Z(0) = Aϵ: yµ (1 + κ(ϵ))Aϵ Aϵ + (yµ (1 + κ(ϵ)) − Aϵ) exp(−ζµ (1 + κ(ϵ))t) yµ (1 − κ(ϵ))Aϵ Zlower (t) = Aϵ + (yµ (1 − κ(ϵ)) − Aϵ) exp(−ζµ (1 − κ(ϵ))t)

Zupper (t) =

(181) (182)

We now apply the timescale change t = τ ln(1/ϵ) to evaluate the pointwise limit as ϵ → 0+ . For the upper bound, substituting the time yields the exponential term exp(−ζµ (1 + κ(ϵ))τ ln(1/ϵ)) = ϵζµ τ (1+κ(ϵ)) . Dividing the numerator and denominator by ϵ yields: yµ (1 + κ(ϵ))A 1 Zupper (τ ln ) = ϵ A + (yµ (1 + κ(ϵ)) − Cϵ)ϵζµ τ (1+κ(ϵ))−1

(183)

Let τµ∗ = 1/ζµ . Because κ(ϵ) → 0 as ϵ → 0+ , for any fixed τ ≠ τµ∗ , the sign of the exponent ζµ τ (1 + κ(ϵ)) − 1 is determined by ζµ τ − 1 for sufficiently small ϵ. If τ < τµ∗ , then the exponent is negative, meaning ϵζµ τ (1+κ(ϵ))−1 → ∞ and Zupper → 0. If τ > τµ∗ , then the exponent is positive, meaning the term vanishes and Zupper → yµ (1 + 0) = yµ . Applying the same time substitution to the lower bound gives the term ϵζµ (1−κ(ϵ))τ . Dividing by ϵ gives the denominator behavior governed by ϵζµ τ −1−ζµ τ κ(ϵ) . Since κ(ϵ) → 0, for any fixed τ ≠ τµ∗ , the sign of the exponent is determined by ζµ τ − 1 for sufficiently small ϵ. Thus, as with the upper bound, Zlower → 0 for τ < τµ∗ and Zlower → yµ (1 − κ(0)) = yµ for τ > τµ∗ . Since Zlower (t) ≤ Z(t) ≤ Zupper (t), the dominant feature subspace converges to a step function for τ ≠ τµ∗ : 1 lim Z (τ ln ) = yµ ⋅ I(τ > τµ∗ ) ϵ

ϵ→0+

(184)

Returning to the residual ∆µ (t) = Z(t) + R(Z, ϵ) − yµ . Because Z ≤ yµ and R(Z, ϵ) = O(ϵmin(δ,1) ) → 0 uniformly for bounded Z, we have limϵ→0+ ∆µ (t) = limϵ→0+ Z(t) − yµ . Squaring this residual yields: 1 2 lim+ ∆µ (τ ln ) = (0 − yµ )2 ⋅ I(τ < τµ∗ ) + (yµ − yµ )2 ⋅ I(τ > τµ∗ ) = yµ2 ⋅ I(τ < τµ∗ ) ϵ→0 ϵ

(185)

Summing over µ completes the proof. The second part of the formal statement is the following. Theorem 11. Suppose the conditions of Theorem 10 hold, and consider the infinite-sample limit m → ∞. Suppose that the ζµ and yµ exhibit the following power-law decay: ζµ = cζ µ−γ2 , yµ = cy µ−γ1

(186)

2 for some positive constants cζ , cy , γ2 > 0. Let the empirical MSE loss be defined as L̃(t) = ∑m µ=1 ∆µ (t) . We also 2 ∞ assume ∑∞ µ=1 (∑k=1 λµ,k zk (0)) < ∞ and γ1 > 1/2 such that the initial loss converges. t Define the rescaled time τ ∶= ln(1/ϵ) . The empirical loss converges L(τ ) ∶= limϵ→0+ L̃ (τ ln 1ϵ ) for almost every τ and the limiting loss follows a power-law decay:

L(τ ) = Θ (τ

−

2γ1 −1 γ2

(187)

Proof. According to Theorem 10, the residual for each sample converges for almost every τ : 1 2 lim+ ∆µ (τ ln ) = yµ2 ⋅ I(τ < τµ∗ ) ϵ→0 ϵ

(188)

where I(⋅) is the indicator function and τµ∗ = ζ1µ is the characteristic timescale. To interchange the limit and the infinite sum, we establish a uniform bound for the residuals. As shown in the proof of Theorem 10, for any trajectory of Z(t), we have Z(t) ≤ yµ (1 + κ(ϵ)) and ∣R(Z, ϵ)∣ ≤ yµ κ(ϵ), where κ(ϵ) → 0 as ϵ → 0+ . Therefore, the absolute residual is bounded by: ∣∆µ (t)∣ = ∣Z(t) + R(Z, ϵ) − yµ ∣ ≤ ∣Z(t)∣ + ∣R(Z, ϵ)∣ + yµ ≤ yµ (2 + 2κ(ϵ)).

(189)

For sufficiently small ϵ, we have κ(ϵ) ≤ 1/2, yielding a uniform bound ∆µ (t)2 ≤ 9yµ2 for all t ≥ 0. Under the 2 2 ∞ −2γ1 assumption γ1 > 1/2, the series ∑∞ < ∞. µ=1 9yµ = 9cy ∑µ=1 µ

Now we can apply the Lebesgue dominated convergence theorem: ∞ 1 2 ∞ L(τ ) = lim+ ∑ ∆µ (τ ln ) = ∑ yµ2 ⋅ I(τ < τµ∗ ), ϵ→0 µ=1 ϵ µ=1

(190)

We can then define a threshold on the index: 1

µ∗ (τ ) ∶= (cζ τ ) γ2 .

(191)

Let M (τ ) = ⌊µ∗ (τ )⌋ + 1 be the smallest integer greater than µ∗ (τ ). The limiting loss evaluates to the tail sum: ∞

L(τ ) =

c2y µ−2γ1 .

∑

(192)

µ=M (τ )

Since c2y x−2γ1 is decreasing, the sum is bounded by its corresponding Riemann integrals: ∞

∫

M (τ )

c2y x−2γ1 dx ≤ L(τ ) ≤ ∫

∞

c2y x−2γ1 dx

M (τ )−1

(193)

Evaluating the integrals yields explicit upper and lower bounds: c2y 2γ1 − 1

(M (τ ))

−(2γ1 −1)

≤ L(τ ) ≤

c2y 2γ1 − 1

(M (τ ) − 1)

−(2γ1 −1)

(194)

As τ → ∞, the threshold scales as M (τ ) ∼ µ∗ (τ ) = (cζ τ ) γ2 . Substituting this asymptotic scaling into both bounds provides the decay rate: −(2γ1 −1)

L(τ ) = Θ ((τ γ2 )

) = Θ (τ

−

2γ1 −1 γ2

(195)

This completes the proof.

Additional Theory

B.1

Removing ZGZ Condition

Theorem 12. Let fx ∶ Rp×d → R be a three times continuously differentiable model with respect to its neurons W = (w1 , . . . , wd ). Assume fx satisfies the permutation symmetry. Then, the model can be approximated around the origin as: d

i=1

fx (W ) = fx (0) + ∑ g(x)⊺ wi + ∑ wi⊺ A(x)wi + ∑ wi⊺ B(x)wj + O(∥W ∥3 ),

(196)

i≠j

where g(x) ∈ Rp , A(x) ∈ Rp×p is a symmetric matrix representing the self-interaction of each neuron and B(x) ∈ Rp×p is a symmetric matrix representing the cross-interaction between any pair of distinct neurons. Proof. We begin with the standard multi-variable Taylor expansion of fx around W = 0: d

⊺

fx (W ) = fx (0) + ∑ (∇wi fx (0)) wi + i=1

1 d d ⊺ 3 ∑ ∑ w Hij wj + O(∥W ∥ ), 2 i=1 j=1 i

(197)

∂ fx ∣ where Hij = ∂w are the p × p blocks of the Hessian matrix. i ∂wj W =0 By permutation symmetry, all neurons are interchangeable. This implies that the first derivative evaluated at the symmetric origin W = 0 must be identical for all neurons. Thus, there exists a vector g(x) such that:

∇wi fx (0) = g(x)

for all i.

(198)

Similarly, the permutation symmetry severely restricts the blocks of the Hessian matrix evaluated at W = 0. There can only be two types of blocks. 39

Diagonal blocks: Hii must be identical for all i. Let Hii = 2A(x). Off-diagonal blocks: Hij for all i ≠ j must be identical. Let Hij = 2B(x). Substituting these findings back into the Taylor expansion: d

fx (W ) = fx (0) + ∑ g(x)⊺ wi + i=1 d

1 1 d ⊺ ⊺ 3 ∑ wi (2A(x))wi + ∑ wi (2B(x))wj + O(∥W ∥ ) 2 i=1 2 i≠j d

(199)

= fx (0) + ∑ g(x) wi + ∑ wi⊺ A(x)wi + ∑ wi⊺ B(x)wj + O(∥W ∥3 ). i=1 i=1 i≠j ⊺

This concludes the proof. We can similarly expand the model to an arbitrary order. Theorem 13. Let fx ∶ Rp×d → R be a K + 1-times continuously differentiable model with respect to its neurons W = (w1 , . . . , wd ). Assume fx satisfies the permutation symmetry. Then, the K-th order Taylor expansion of the model around the origin W = 0 can be expressed as: K

fx (W ) = fx (0) + ∑ ∑

∑

k=1 λ⊢k i1 ,...,im distinct

1 2 m ] + O(∥W ∥K+1 ), Tλ (x)[wi⊗c , wi⊗c , . . . , wi⊗c 1 2 m

(200)

where λ = {c1 , c2 , . . . , cm } represents an integer partition of k, satisfying ∑m r=1 cr = k and cr ≥ 1. The innermost sum iterates over all m-tuples of mutually distinct indices (i1 , . . . , im ) drawn from {1, . . . , d} and Tλ (x) is a symmetric tensor of order k, determined by λ, representing the partial derivatives and combinatorial constants evaluated at the origin. Proof. We begin with the standard multi-variable Taylor expansion of fx around W = 0 up to order K: K

d 1 d K+1 k ). ∑ ⋅ ⋅ ⋅ ∑ (∇wj1 ...wjk fx (0)) [wj1 , . . . , wjk ] + O(∥W ∥ jk =1 k=1 k! j1 =1

fx (W ) = fx (0) + ∑

(201)

where ∇kwj1 ...wj fx (0) denotes the k-th order partial derivative tensor block evaluated at the origin. k For any fixed k, consider an arbitrary sequence of indices (j1 , . . . , jk ). This sequence can be grouped by the multiplicity of each distinct index it contains. Let there be m distinct indices i1 , . . . , im appearing with multiplicities c1 , . . . , cm respectively, such that ∑m r=1 cr = k and cr ≥ 1. This multiset of multiplicities defines an integer partition λ ⊢ k. By permutation symmetry, all neurons are interchangeable. At W = 0, any permutation of the input neurons leaves the function and its derivative tensors unchanged. Consequently, the derivative ∇kwj1 ...wj fx (0) depends on the k partition λ of its indices, and is identical for any specific choice of neurons i1 , . . . , im . Let Dλ (x) denote this derivative corresponding to the partition λ. We can rewrite the sum over the raw indices (j1 , . . . , jk ) by summing over all valid partitions λ ⊢ k, and then summing over all possible choices of m distinct indices (i1 , . . . , im ) from {1, . . . , d}: d

∑ ⋅ ⋅ ⋅ ∑ (∇wj1 ...wjk fx (0)) [wj1 , . . . , wjk ] = ∑

j1 =1

jk =1

∑

λ⊢k i1 ,...,im distinct

1 m ]. Cλ Dλ (x)[wi⊗c , . . . , wi⊗c 1 m

(202)

Here, Cλ is a combinatorial coefficient counting the number of ways to arrange the sequence (j1 , . . . , jk ) such that the distinct elements appear with the frequencies specified by λ. Substituting it back into (200), we can define Tλ (x) = Ck!λ Dλ (x), which yields: K

fx (W ) = fx (0) + ∑ ∑

∑

k=1 λ⊢k i1 ,...,im distinct

1 m ] + O(∥W ∥K+1 ). Tλ (x)[wi⊗c , . . . , wi⊗c 1 m

This concludes the proof. 40

(203)

B.2

Multidimensional Output

Corollary 1. Let Fx ∶ Rp×d → Rk be a three times continuously differentiable model with respect to its neurons W = (w1 , . . . , wd ), where wi ∈ Rp . Assume Fx satisfies the permutation symmetry4 and the ZGZ condition: the Jacobian matrix satisfies Jwi Fx (w1 , . . . , wd )∣w =0 = 0 ∈ Rk×p for any i. i Then, the model can be approximated around the origin as: d

Fx (W ) = Fx (0) + ∑ A(x)[wi , wi ] + O(∥W ∥3 ),

(204)

i=1

where A(x) ∈ Rk×p×p is a 3rd-order tensor dependent only on the input x. The operation A(x)[wi , wi ] yields a vector in Rk whose m-th component is wi⊺ A(m) (x)wi , with each A(m) (x) ∈ Rp×p being a symmetric matrix. (1)

(2)

⊺

(k)

(m)

Proof. Let Fx (W ) = [fx (W ), fx (W ), . . . , fx (W )] , where each component function fx three times continuously differentiable. (m) By permutation symmetry, each component fx satisfies

∶ Rp×d → R is

fx(m) (w1 , . . . , wd ) = fx(m) (wσ(1) , . . . , wσ(d) ).

(205)

Furthermore, the ZGZ condition implies that for every output dimension m: ∇wi fx(m) (W )∣w =0 = 0 ∈ Rp . i

(m)

Thus each component function fx

(206)

satisfies the conditions of Theorem 1, and we can apply it to each dimension m: d

fx(m) (W ) = fx(m) (0) + ∑ wi⊺ A(m) (x)wi + O(∥W ∥3 ),

(207)

i=1

where A(m) (x) = 21 Hii is the Hessian blocks of the m-th output dimension. To stack the equations for all k components, we define a 3rd-order tensor A(x) ∈ Rk×p×p such that its m-th slice along the first dimension represents A(m) (x). Consequently, the multi-dimensional expansion can be written as: (m)

Fx (W ) = Fx (0) + ∑ A(x)[wi , wi ] + O(∥W ∥3 ).

(208)

i=1

This concludes the proof.

B.3

Learning Dynamics of Other Optimization Methods

Corollary 2. Consider any learning algorithm whose update rule for the weight matrix W takes the form: ∆Wt = ∑ αk ∇W Lk + βk Wk

(209)

where αk , βk can be any scalar functions of the historical sufficient statistics {Mτ , µτ }τ ≤t . Then, the learning dynamics of the NQF are completely determined by {Mτ , µτ }τ ≤t . Proof. From Theorem 2, we know that the gradient of the loss with respect to the weight matrix W at any step k takes the form: ∇W Lk = vk 1⊺ + Hk Wk . (210) Given the update rule ∆Wt = ∑k≤t (αk ∇W Lk + βk Wk ), we substitute the gradient expression: ∆Wt = ∑ (αk vk 1⊺ + (αk Hk + βk I)Wk ). k≤t 4 F (w , . . . , w ) = F (w x x 1 d σ(1) , . . . , wσ(d) ) for any permutation σ.

(211)

We claim that for any step t, the weight matrix Wt can be written in the following form: Wt = Pt W0 + Qt 1⊺ ,

(212)

where Pt ∈ Rp×p is a transformation matrix and Qt ∈ Rp is a vector, dependent on the historical sequences {vk , Hk , αk , βk }k<t . We prove this by induction: For t = 0, W0 = I ⋅ W0 + 0 ⋅ 1⊺ . Thus, P0 = I and Q0 = 0, which holds. Assume the claim holds for all k ≤ t. For step t + 1, we have: Wt+1 = Wt + ∆Wt = Wt + ∑ (αk vk 1⊺ + (αk Hk + βk I)Wk ).

(213)

k≤t

Substitute Wk = Pk W0 + Qk 1⊺ into the equation: Wt+1 = (Pt W0 + Qt 1⊺ ) + ∑ [αk vk 1⊺ + (αk Hk + βk I)(Pk W0 + Qk 1⊺ )].

(214)

k≤t

By grouping the terms associated with W0 and 1⊺ , we get: Wt+1 = (Pt + ∑ (αk Hk + βk I)Pk ) W0 + (Qt + ∑ [αk vk + (αk Hk + βk I)Qk ]) 1⊺ . k≤t

(215)

k≤t

´¹¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¸¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¶

´¹¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹¸ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¹ ¶

∶=Pt+1

∶=Qt+1

This defines a recurrence relation for Pt+1 and Qt+1 that depends only on variables from steps k ≤ t. Now we show that µt and Mt only depend on Pt , Qt and the initial statistics. For the first moment µt : µt = Wt 1 = (Pt W0 + Qt 1⊺ )1 = Pt µ0 + d ⋅ Qt

(216)

Mt = Wt Wt⊺ = (Pt W0 + Qt 1⊺ )(Pt W0 + Qt 1⊺ )⊺ = Pt M0 Pt⊺ + Pt µ0 Q⊺t + Qt µt0 Pt⊺ + d ⋅ Qt Q⊺t .

(217)

For the second moment Mt :

Therefore, given the initial conditions M0 and µ0 , the subsequent values of Pt and Qt are completely determined, which in turn dictate the values of µt and Mt via the derived equations. This proves that for any linear-span learning algorithm, the dynamics of the NQF form a closed system over {Mτ , µτ }τ ≤t .

B.4

Learning Dynamics of Multi-layer NQF

Before stating the theorem, we first define an L-layer NQF. Definition 2 (Deep NQF). Let f (x) = h(L) (x) be an L-layer NQF with hidden representation h(l) (x) ∈ RDl : h(0) (x) = x

(218)

⊺ (l) (l) (l) (l) ⊺ (l) hk (x) = Ck + (gk ) µ(l) + (µ(l) ) Bk µ(l) + Tr [M (l) Ak ] ,

k = 1, . . . , Dl

(219)

where M (l) ∶= W (l) (W (l) )⊺ ∈ Rpl ×pl and µ(l) ∶= W (l) 1dl ∈ Rpl are the order parameters. The structure compo(l) (l) (l) (l) nents {Ck , gk , Bk , Ak } are functions that only depend on the previous layer’s output h(l−1) (x). Without loss of (l) (l) generality, Ak and Bk are symmetric matrices. Theorem 14 (Master Theorem for Deep NQF). Under SGD ∆W (l) = −η ∑ ∇W (l) L(f (x))

(220)

x∈B (l)

(l)

with identical data sampling, two models will have identical states (µa (t) = µb (t), Ma (t) = Mb (t)) for all layers l and time t if they share the same initial statistics: (l)

µ(l) a (0) = µb (0),

(l)

Ma(l) (0) = Mb (0) 42

∀l ∈ {1, . . . , L}.

(221)

The learning dynamics are closed, evolving as: ∆µ(l) = −η(dl ⋅ v (l) + H (l) µ(l) )

(222)

∆M (l) = −η(v (l) (µ(l) )⊺ + µ(l) (v (l) )⊺ + H (l) M (l) + M (l) H (l) ) + η 2 (dl ⋅ v (l) (v (l) )⊺ + v (l) (µ(l) )⊺ H (l) + H (l) µ(l) (v (l) )⊺ + H (l) M (l) H (l) ),

(223)

where dl is the width (number of neurons) of layer l. The vector v (l) ∈ Rpl and effective Hessian H (l) ∈ Rpl ×pl are defined as: Dl

∂L

v (l) ∶= ∑ ∑

(l)

(l) x∈B k=1 ∂hk (x) Dl

H (l) ∶= ∑ ∑ 2 x∈B k=1

(l)

(gk + 2Bk µ(l) )

∂L

(l) ∂hk (x)

(224)

(l)

(225)

Proof. By the chain rule, the gradient of the loss with respect to the weights W (l) is: Dl

∂L

∇W (l) L = ∑ ∑

(l) x∈B k=1 ∂hk (x)

(l)

∇W (l) hk (x).

(226)

Notice that µ(l) = W (l) 1dl , where 1dl ∈ Rdl is a vector of ones. The derivatives of the terms in hk are: (l)

(l)

∇W (l) ((gk )⊺ µ(l) ) = gk 1⊺dl

(227)

(l) (l) ∇W (l) ((µ(l) )⊺ Bk µ(l) ) = 2Bk µ(l) 1⊺dl (l) (l) ∇W (l) Tr [M (l) Ak ] = 2Ak W (l)

(228) (229)

Collecting these terms, the gradient factorizes into: ∇W (l) L = v (l) 1⊺dl + H (l) W (l) ,

(230)

where v (l) and H (l) are defined as in the theorem. The SGD update is therefore: ∆W (l) = −η(v (l) 1⊺dl + H (l) W (l) ).

(231)

For the mean vector µ(l) , we multiply the update by 1dl : ∆µ(l) = ∆W (l) 1dl = −η(v (l) (1⊺dl 1dl ) + H (l) W (l) 1dl ) = −η(dl ⋅ v (l) + H (l) µ(l) ).

(232)

(l)

For the second moment M (l) , we expand Mt+1 = (W (l) + ∆W (l) )(W (l) + ∆W (l) )⊺ : ∆M (l) = W (l) (∆W (l) )⊺ + ∆W (l) (W (l) )⊺ + ∆W (l) (∆W (l) )⊺ .

(233)

Substitute ∆W (l) and use the identity W (l) 1dl = µ(l) : W (l) (∆W (l) )⊺ = −η(W (l) 1dl (v (l) )⊺ + W (l) (W (l) )⊺ H (l) ) = −η(µ(l) (v (l) )⊺ + M (l) H (l) ).

(234)

The symmetric counterpart is −η(v (l) (µ(l) )⊺ + H (l) M (l) ). For the η 2 term: ∆W (l) (∆W (l) )⊺ = η 2 (v (l) 1⊺dl + H (l) W (l) )(1dl (v (l) )⊺ + (W (l) )⊺ H (l) ) = η 2 (v (l) (1⊺dl 1dl )(v (l) )⊺ + v (l) 1⊺dl (W (l) )⊺ H (l) + H (l) W (l) 1dl (v (l) )⊺ + H (l) W (l) (W (l) )⊺ H (l) ) = η 2 (dl ⋅ v (l) (v (l) )⊺ + v (l) (µ(l) )⊺ H (l) + H (l) µ(l) (v (l) )⊺ + H (l) M (l) H (l) ).

(235)

Summing these gives the update rule for ∆M (l) . Since the forward pass h(l) (x) depends only on {µ(l) , M (l) } and h(l−1) (x), all intermediate activations are func∂L also depend only on the sequence of order paramtions of the order parameters. Consequently, the gradients ∂h(l) (x) eters. Thus, the dynamics are closed independent of W (l) . 43

Theorem 3 can be generalized to multi-layer NQFs in the same way. Theorem 15 (Compressibility of Deep NQF). Let an L-layer General Deep NQF have layer widths d1 , . . . , dL , (l) (l) (l) defined by functions gk , Bk , Ak , learning rate η, and initial parameter statistics µ(l) (0) and M (l) (0) for each layer l ∈ {1, . . . , L}. Let V (l) ⊆ Rpl be the joint subspace spanned by the network components across all training samples x ∈ X and all training steps t ≥ 0 of the original model: Dl

V (l) ∶= span

(l)

(l−1)

⋃ (gk (ht

⋃

x∈X ,t≥0 k=1

(l)

(l−1)

) ∪ Col(Ak (ht

(l)

(l−1)

)) ∪ Col(Bk (ht

(236)

))).

Let kl ∶= dim(V (l) ) ≤ pl . If dl > kl + 1 for all l, then there exists a smaller Deep NQF with layer widths d′l = kl + 1, d′ (l) (l) characterized by Ãk = dll Ak , rescaled learning rate η̃ = dd′l η (layer-specific learning rates), such that: l 1. f˜x = fx for all inputs x ∈ X at initialization; (l)

(l)

2. µ̃t = Vl Vl⊺ µt and M̃t

= dd′l Vl Vl⊺ Mt Vl Vl⊺ for all layers l and all training steps t ≥ 0, where Vl ∈ Rpl ×kl is the (l)

orthogonal basis of V (l) . 3. The output and learning dynamics of the smaller model are identical to the original model: f˜x (t) = fx (t) for all t ≥ 0.

B.5

Compression Error

Corollary 3. For any original NQF with d neurons and input dimension p, and any target width d′ ≤ p, there exists a compressed NQF with d′ neurons such that, assuming the SGD update dynamics are locally Lipschitz continuous with constant L, the deviation in its sufficient statistics at any training step t is bounded by: ¿ p ⎛Á d′ À ∑ λ2 ⎞ exp(ηLt), (237) ∥ M̃t − Mt ∥ + ∥µ̃t − µt ∥2 ≤ Á d ⎝ k=d′ k ⎠ F where λ1 ≥ λ2 ≥ ⋅ ⋅ ⋅ ≥ λp ≥ 0 are the eigenvalues of the initial covariance matrix C ∶= dd′ (M (0) − d1 µ(0)µ(0)⊺ ). ′

Proof. Let the discrepancy in the sufficient statistics at training step t be denoted by the error states EM ∶= dd M̃t − Mt (t)

(t)

and Eµ ∶= µ̃t − µt . Define the total error norm as et ∶= ∥EM ∥F + ∥Eµ ∥2 . (0) To minimize the initial error, we first match the first moment by setting µ̃(0) = µ(0), which yields Eµ = 0. For ′ the second moment, following the construction in Theorem 3, finding U = [u1 , . . . , ud′ ] ∈ Rp×d requires solving: d′

∑ ui = 0,

and

U U ⊺ ≈ C ∶=

i=1

d 1 (M (0) − µ(0)µ(0)⊺ ) . ′ d d

(238)

Because the d′ vectors must sum to zero, they are confined to a (d′ − 1)-dimensional subspace. Consequently, the rank of the constructed covariance matrix U U ⊺ can be at most d′ − 1. Since we assume the target width d′ ≤ p, perfect reconstruction of the p × p full-rank matrix C is in general impossible. By the Eckart-Young-Mirsky theorem, the optimal rank-(d′ − 1) approximation C̃ that minimizes the Frobenius norm of the residual is given by the truncated eigendecomposition of C. Let C = ∑pk=1 λk qk qk⊺ be the eigendecompod′ −1 sition with eigenvalues λ1 ≥ ⋅ ⋅ ⋅ ≥ λp ≥ 0. We construct U such that U U ⊺ = C̃ ∶= ∑k=1 λk qk qk⊺ . The residual matrix is p ∆C = C − C̃ = ∑k=d′ λk qk qk⊺ . We initialize the compressed model such that its effective initial covariance is this optimal approximation. Thus, the initial matrix error is the residual: (0)

EM =

d′ M̃ (0) − M (0) = C̃ − C = −∆C. d

The initial total error is therefore the Frobenius norm of the truncated tail: ¿ Á p À ∑ λ2 . e0 = ∥∆C∥F + 0 = Á k=d′

(239)

(240)

For any step t, the dynamics of M and µ are governed by the update rules derived in Theorem 2. The increments ∆M and ∆µ are functions of the current state variables (Mt , µt ), the data sequence A(x), B(x), g(x), and the loss derivative ℓ′x . Because polynomials are locally Lipschitz continuous, and assuming the SGD trajectories and data samples remain bounded within the time horizon t, there exists a uniform Lipschitz constant L > 0 that bounds the divergence of the trajectories. Specifically, the difference in the updates between the uncompressed state (Mt , µt ) and the scaled ′ compressed state ( dd M̃t , µ̃t ) satisfies: ∥

d′ (t) ∆M̃t − ∆Mt ∥ + ∥∆µ̃t − ∆µt ∥2 ≤ ηL(∥EM ∥F + ∥Eµ(t) ∥2 ) = ηLet . d F

(241)

By the triangle inequality, the error at the next step is bounded by: et+1 ≤ et + ∥

d′ ∆M̃t − ∆Mt ∥ + ∥∆µ̃t − ∆µt ∥2 ≤ et + ηLet = (1 + ηL)et . d F

(242)

Applying this discrete recurrence relation recursively from step 0 to t, we obtain: et ≤ e0 (1 + ηL)t .

(243)

Using the standard exponential inequality (1 + x) ≤ exp(xt) for all x ≥ 0, we establish the upper bound for the error dynamics: ¿ p ⎛Á À ∑ λ2 ⎞ exp(ηLt). (244) et ≤ e0 exp(ηLt) = Á ⎝ k=d′ k ⎠ This completes the proof.

B.6

NTK and Feature Learning

Proposition 8. Under the same notations as Theorem 2, the empirical Neural Tangent Kernel (NTK) between two data points x and x′ , defined as Θ(x, x′ ) ∶= ⟨∇W fx (W ), ∇W fx′ (W )⟩, is determined by the summary statistics µ and M: Θ(x, x′ ) = d ⋅ u(x)⊺ u(x′ ) + 2u(x)⊺ A(x′ )µ + 2µ⊺ A(x)u(x′ ) + 4Tr(A(x)A(x′ )M ) (245) where u(x) ∶= g(x) + 2B(x)µ. Proof. By definition, the empirical NTK is the Frobenius inner product of the gradients with respect to the weight matrix W : Θ(x, x′ ) = Tr(∇W fx (W )⊺ ∇W fx′ (W )). (246) From the proof of Theorem 2, we know the sample-wise gradient can be written as: ∇W fx (W ) = (g(x) + 2B(x)µ)1⊺ + 2A(x)W = u(x)1⊺ + 2A(x)W.

(247)

Substituting this into the NTK definition yields: Θ(x, x′ ) = Tr((1u(x)⊺ + 2W ⊺ A(x))(u(x′ )1⊺ + 2A(x′ )W )).

(248)

We expand this expression to obtain: Θ(x, x′ ) = d ⋅ u(x)⊺ u(x′ ) + 2u(x)⊺ A(x′ )µ + 2µ⊺ A(x)u(x′ ) + 4Tr(A(x)A(x′ )M ).

(249)

This demonstrates that the empirical NTK at any point in the training trajectory can be evaluated without knowing the individual parameter states in W , as it relies solely on the sufficient statistics µ and M . Assume the network width is d. Let the NQF coefficients scale with d as g(x) = Θ(d−αg ), A(x) = Θ(d−αA ), and B(x) = Θ(d−αB ). To prevent the outputs and gradients from diverging at initialization, we require αA ≥ 0 and αB ≥ 1/2. Assume standard variance scaling at initialization (e.g., 1/d) such that the summary statistics scale as µ(0) = O(1) and M (0) = O(1). Based on Proposition 8, the initial magnitude of the NTK scales as Θ(x, x′ ) = O(d1−2αg + d1−2αB + d−2αA ). We assume that the initial NTK is Θ(x, x′ ) = Θ(1) by setting αg = 1/2. 45

Proposition 9 (NTK Evolution and Feature Learning). For the gradient flow, the learning dynamics have two regimes: 1. Lazy Training (Constant NTK): If αA > 0 and αB > 1, then limd→∞ Θ̇ = 0. The NTK remains invariant while the loss decreases by O(1). The NQF effectively behaves as a linear model. 2. Feature Learning (Evolving NTK): If αA = 0 or 12 ≤ αB ≤ 1, then Θ̇ = O(1). The NTK changes on the same timescale as the loss, allowing the model to learn data-dependent representations. Proof. By the chain rule under gradient flow Ẇ = −∇W L, the dynamics of the function output are f˙x = ⟨∇W fx , Ẇ ⟩ = − ∑x′ ∈B Θ(x, x′ )ℓ′x′ . From Proposition 8, the NTK is Θ = d ⋅ u⊺ u + 2u⊺ Aµ + 2µ⊺ Au + 4Tr(A2 M ), where u = g + 2Bµ. With αg = 1/2, the leading term d ⋅ g ⊺ g = Θ(1), ensuring Θ(x, x′ ) = Θ(1). Consequently, f˙x = O(1), guaranteeing that the loss changes by an O(1) amount in an O(1) time. We now evaluate the change in NTK during this O(1) learning process. The evolution of the sufficient statistics from Theorem 2 are: µ̇ = −(d ⋅ v + Hµ), Ṁ = −(vµ⊺ + µv ⊺ + HM + M H) (250) where v = ℓ′x u(x) and H = 2ℓ′x A(x). For magnitude analysis, we drop the constant scalar ℓ′x . Taking the time derivative of Θ yields: Θ̇ = 2d ⋅ u⊺ u̇ + 4u̇⊺ Aµ + 4u⊺ Aµ̇ + 4Tr(A2 Ṁ ) (251) Substituting u̇ = 2B µ̇ and µ̇ ≈ −d ⋅ u (since Hµ is dominated), the two leading-order terms dictating the macroscopic change Θ̇ are: 1. Term 1: d ⋅ u⊺ B µ̇ ≈ −d2 u⊺ Bu. Since αg = 1/2 and αB ≥ 1/2, we have u = g + 2Bµ = O(d−1/2 ). Thus, this term scales as d2 (d−1/2 )2 d−αB = O(d1−αB ). 2. Term 2: −d ⋅ u⊺ Au. This term scales as d(d−1/2 )2 d−αA = O(d−αA ). Other terms, such as Tr(A2 Ṁ ), scale as d−3αA or lower and are dominated. Therefore, the overall drift in NTK is bounded by Θ̇ = O(d1−αB + d−αA ). For the NTK to remain constant (while the loss decreases by O(1)), the derivative Θ̇ must vanish as d → ∞. This requires 1 − αB < 0 Ô⇒ αB > 1 and −αA < 0 Ô⇒ αA > 0. If these conditions are violated, the NTK shifts by an O(1) magnitude during the O(1) training time.

B.7

A general exact solution

The following theorem unifies a class of exact solutions in which the learning dynamics reduce to independent logistic equations. It contains Theorem 6 and the simultaneously diagonalizable special case of Theorem 7, as well as the exact solutions of linear networks in [62]. Theorem 16. Define the empirical covariance K(X) ∶=

4 m ∑ ⟨A(xµ ), X⟩F A(xµ ), m µ=1

Y ∶=

4 m ∑ yµ A(xµ ). m µ=1

(252)

Suppose that there exist pairwise orthogonal symmetric projectors {Πa }ra=1 and constants γa ∈ R and κa ≥ 0 such that Πa Πb = δab Πa , and for every a, b ∈ {1, . . . , r}, Y Πa + Πa Y = 2γa Πa ,

(253)

K(Πb )Πa + Πa K(Πb ) = 2κa δab Πa .

(254)

If M (0) = ∑ra=1 za (0)Πa with za (0) ≥ 0, then the solution can be written as: r

M (t) = ∑ za (t)Πa .

(255)

a=1

If γa ≠ 0, za (t) =

γa za (0) , κa za (0) + (γa − κa za (0)) e−2γa t 46

(256)

whereas, if γa = 0, za (t) =

za (0) . 1 + 2κa za (0)t

(257)

Proof. The dynamics of M (t) = W (t)W (t)⊺ can be written as Ṁ = Y M + M Y − K(M )M − M K(M ).

(258)

Consider a matrix in the subspace spanned by the projectors M = ∑ra=1 za Πa . The linear part of (258) is r

Y M + M Y = ∑ za (Y Πa + Πa Y ) = 2 ∑ γa za Πa , a=1

(259)

a=1

where we used (253). For the nonlinear part, we have r

a,b=1

a=1

K(M )M + M K(M ) = ∑ za zb [K(Πb )Πa + Πa K(Πb )] = 2 ∑ κa za zb δab Πa = 2 ∑ κa za2 Πa ,

(260)

where (254) was used in the second equality. Substitution of (259) and (260) into (258) gives r

Ṁ = 2 ∑ za (γa − κa za )Πa .

(261)

a=1

Thus each mode evolves independently according to ża (t) = 2za (t) (γa − κa za (t)) .

(262)

dza = 2t + constant, za (γa − κa za )

(263)

For γa ≠ 0, separation of variables gives ∫

from which (256) follows. When γa = 0, the equation reduces to ża = −2κa za2 ,

(264)

whose solution is (257). Corollary 4. Consider an NQF with structure matrices A(xµ ) =

1 0 [ ⊺ 2 Sµ

Sµ ], 0

Sµ ∈ Rnout ×nin .

(265)

U Writing W = [ ] with U ∈ Rnout ×d and V ∈ Rnin ×d . Define V C(Q) ∶=

1 m ∑ ⟨Sµ , Q⟩F Sµ , m µ=1

T ∶=

1 m ∑ yµ Sµ . m µ=1

(266)

Suppose that T possesses singular triplets {(τa , ua , va )}ra=1 satisfying T ⊺ ua = τa va ,

T va = τa ua ,

(267)

where {ua }ra=1 and {va }ra=1 are orthonormal. Assume additionally that the corresponding rank-one singular modes are eigenmatrices of C: C(ua va⊺ ) = ca ua va⊺ , ca ≥ 0. (268)

Let {ra }ra=1 ⊆ Rd be orthonormal directions, and assume the balanced and aligned initialization r

√ ρa (0) ua ra⊺ ,

U (0) = ∑

√

V (0) = ∑

ρa (0) va ra⊺ ,

(269)

ρa (t) va ra⊺ ,

(270)

a=1

where ρa (0) ≥ 0. Then the trajectory remains balanced and aligned: r

√

U (t) = ∑

ρa (t) ua ra⊺ ,

√

V (t) = ∑

a=1

where each singular mode evolves independently. If τa ≠ 0, the solution is ρa (t) =

τa ρa (0) , ca ρa (0) + (τa − ca ρa (0)) e−4τa t

(271)

ρa (0) . 1 + 4ca ρa (0)t

(272)

whereas, if τa = 0, ρa (t) = Proof. Introduce the symmetric block-matrix notation Q ]. 0

(273)

1 A(xµ ) = B(Sµ ). 2

(274)

4 m ∑ yµ A(xµ ) = 2B(T ). m µ=1

(275)

B(Q) ∶= [ Then

0 Q⊺

From (252), the target matrix Y is Y = For each singular mode, define 1 u ψa ∶= √ [ a ] , 2 va

Πa ∶= ψa ψa⊺ .

(276)

1 ⊺ (u ub + va⊺ vb ) = δab , 2 a

(277)

2 Tv Y ψa = √ [ ⊺ a ] = 2τa ψa . 2 T ua

(278)

Y Πa + Πa Y = 4τa Πa .

(279)

The orthonormality of the singular vectors implies ψa⊺ ψb = and hence Πa Πb = δab Πa . Using (267),

It follows that Thus (253) is satisfied with γa = 2τa . We next evaluate the empirical covariance on Πb . The off-diagonal block of Πb is 21 ub vb⊺ , and therefore K(Πb ) =

4 m ⊺ ⊺ ∑ ⟨A(xµ ), Πb ⟩F A(xµ ) = B (C(ub vb )) = cb B(ub vb ). m µ=1

(280)

For any a, b, 1 u v⊺ v B(ub vb⊺ )ψa = √ [ b ⊺b a ] = δab ψa . 2 vb ub ua

(281)

Figure 6: Left: Comparison between a two-layer MLP and its NQF approximation (Proposition 1) under a teacherstudent setting. Right: Comparison between a query-key-only attention model and its NQF approximation (Proposition 3). Models are trained via online SGD (learning rate 0.05, batch size 64) over 2, 000 iterations. The results are averaged across 5 independent runs, with shaded regions representing ±1 standard deviation. Consequently, K(Πb )Πa + Πa K(Πb ) = 2ca δab Πa .

(282)

Thus the condition (254) holds with κa = ca . Under the initialization (269), we have r

M (0) = W (0)W (0)⊺ = ∑ 2ρa (0)Πa .

(283)

a=1

Theorem 16 therefore applies with za (t) = 2ρa (t), γa = 2τa , κa = ca . Substitution into (262) yields ρ̇a = 4ρa (τa − ca ρa ).

(284)

(271) and (272) follow directly from (256) and (257). Finally, factorizing M (t) = ∑ 2ρa (t)Πa

(285)

using the fixed orthonormal directions {ra } gives (270). Remark (Recovery of the solution in [62]). Consider the standard two-layer linear network ŷ(x) = U V ⊺ x

(286)

m m 1 1 ⊺ ⊺ trained on samples {(xµ , yµ )}m µ=1 . Define Σx = m ∑µ=1 xµ xµ and T = m ∑µ=1 yµ xµ . The corresponding operator is

C(Q) = QΣx . Hence the condition (268) holds whenever Σx va = ca va

(287)

for the right singular vectors of T . In particular, for whitened inputs, Σx = I, we have ca = 1 for every mode. Then the solution is τa ρa (0) ρa (t) = . (288) ρa (0) + (τa − ρa (0)) e−4τa t Up to a rescaling of time, this is precisely the balanced solution of deep linear networks derived by [62].

Experimental Details and Additional Experiments

Figure 2 In Figure 2, all the models use output dimension 8 and the MSE loss. The two-layer MLP has an input dimension 20, hidden dimension 64 and the tanh activation without any biases. The single-layer CNN takes a singlechannel input of dimension 5 × 7, employing 64 kernels of dimension 4 × 5 (with stride 1 and no padding), followed by global sum-pooling. Both the single-head and the multi-head attention models have an input dimension 20, a hidden dimension of 8 and sequence length 16. The multi-head model further has four heads with dv = 20. For all the experiments, we do full batch learning with 600 samples. “Momentum” refers to GD with momentum 0.9. Learning rates differ in each setting. 49

Figure 7: Comparison of training dynamics between a two-layer MLP (Left), a one-layer CNN (Middle), a one-layer self-attention (Right) and their corresponding NQFs on the MNIST odd vs. even classification task. In all cases, the NQFs (in red) achieve similar performance as the original model (in blue).

Figure 8: Training dynamics of 4−layer MLPs and 2−layer attention models compared to their NQF approximations under different optimizers and initialization scales. The black line represents the original model. The orange line shows the 2-layer NQF approximations. The purple line represents an independently trained 1-layer NQF, initialized to match the original model at step zero. Solid lines represent the training loss and dashed lines represent the test loss. Under small initialization, the dynamics overlap, confirming the validity of deep NQF approximations. Additional Experiments on NQF Approximation On the left side of Figure 6, we train a two-layer bias-free MLP with a tanh activation function in the teacher-student framework. The teacher weights are standard Gaussian, and we train two student models: an MLP of the same structure and its corresponding NQF approximation. Both models are optimized using online SGD with a learning rate of 0.05 and a batch size of 64 over 2, 000 steps. Both the teacher and the student MLP have the input dimension 20, the hidden dimension 100 and the output dimension 1. Figure 6 shows that the trajectories of two student models match for the first 100 steps and then they become different. For the right side of Figure 6, we compare a query-key-only attention model (Proposition 3) with its NQF approximation and observe similarly that the NQF tracks the attention model accurately for the first 100 steps. The attention model computes a scalar output from a query token x ∈ RD and a context sequence X ∈ RD×N , with a fixed Gaussian readout vector c ∈ RD . We set D = 16, N = 10, dk = 8, and use H = 4 attention heads. Both the attention model and the NQF are optimized using online SGD with a learning rate of 0.02 and a batch size of 64 over 2, 000 steps. For both the MLP and the attention model, the teacher has standard Gaussian weights and student networks are initialized with identical Gaussian weights drawn from N (0, ϵ2 ) where ϵ = 10−3 . At each step, the input data is sampled i.i.d. from a Gaussian distribution. In Figure 7, we construct a two-layer bias-free MLP equipped with tanh activations, a one-layer CNN and a onelayer self-attention and their NQF counterparts, where the NQFs are initialized with identical small Gaussian weights (N (0, ϵ2 ) with ϵ = 0.01) as the original models. All the models are trained with Adam (full batch), MSE loss and learning rate 0.001. We construct a binary classification task: distinguishing odd from even digits on a subset of the 50

MNIST dataset with 4000 samples. As illustrated in Figure 7, the original models and their corresponding NQFs exhibit indistinguishable trajectories in both training loss and test accuracy. The two-layer MLP uses the hidden dimension 20, tanh activation and no biases. The one-layer CNN uses 50 channels with kernel size 15 × 15. The attention model uses a fixed readout vector, 10 heads and dk = 32. The input 28 × 28 image is regarded as 28 tokens with embedding dimension 28. In Figure 8, We formulate a standard teacher-student regression task. We examine two architectures: a 4-layer MLP and a 2-layer single-head attention. The inputs are drawn from a normal distribution, and the targets are generated by a corresponding teacher model with fixed Gaussian weights. According to Propositions 1 and 2, both models can be approximated by 2-layer NQFs. This is validated by Figure 8, where under small initialization, the training dynamics of both models are well approximated by 2-layer NQFs but not 1-layer NQFs. The MLP has the input dimension 8, the hidden dimension 16 and the output dimension 1. The initial weights are Gaussian with variance 0.05 under small initialization and 0.5 under large initialization. The 2-layer attention model has the input and hidden dimension 8, input length 10 and the output dimension 1. The initial weights are Gaussian with variance 0.1 under small initialization and 1 under large initialization. The training set has 200 samples and the test set has 1000 samples. When the initialization scale and the model are fixed, three models (original one, 2-layer NQF and 1-layer NQF) share the same initialization and the same learning rate. Figure 3 For the left and middle side of Figure 3, we choose m = 5, r = {1.0, 0.5, 0.25, 0.125, 0.0625} and Ckk = 1 for all k. The initialization scale is ϵ = 10−5 . {Aµ }m µ=1 are chosen to be diagonal matrices of size m × m. To satisfy the condition of Theorem 6, we randomly generate a m × m orthogonal matrix and use the µ−th column as the diagonal elements of Aµ . We train the model √ using full-batch GD with learning rate 0.01 to approximate the gradient flow and use the initialization W (0) = ϵI such that zk (0) = ϵ. Because W remains diagonal, zk (t) = [W (t)W ⊺ (t)]kk represents the eigenvalues of W W ⊺ . For the right side of Figure 3, we choose {Aµ }1000 µ=1 such that each matrix only has one non-zero element. Specifically, we first set rk = k −α2 , Vk = k −α1 √ for k = 1, 2, ⋯, 1000 and α1 = 2, α2 = 0.8. Then we can calculate Ckk = rk2 /8Vk . We choose Ak,kk = Ckk m/8 and other elements to be 0. The labels are chosen to be yk = mrk /(8Ak,kk ). The initialization is chosen to be zk (0) = 10−6 k −1.2 for k = 1, 2, ⋯, 1000. Figure 4 To validate Theorem 6, we construct a synthetic dataset utilizing block-diagonal data matrices. We consider a system with m = 4 samples and a parameter dimension of p = m × K. {A(xµ )}m µ=1 are configured as diagonal matrices with disjoint non-zero supports. Specifically, for any given sample µ ∈ {1, . . . , m}, only the diagonal entries within the index block [µK, µK + K − 1] are non-zero. This ensures that A(xµ )A(xν ) = 0 for all µ ≠ ν, satisfying the condition required by Theorem 5. We choose K = 3 and the three non-zero elements to be 1, 1.5, 2. In this case (36) remains irreducible and lacks a closed-form solution. Consequently, the theoretical trajectories presented in Figure 4 are obtained by numerically integrating the associated ODE for {ξµ (t)}4µ=1 using the standard Runge-Kutta method. Finite-initialization Effect On the right side of Figure 3 and Figure 5, one can notice that there is a slight difference between the empirical slope and the theoretical prediction. This is a finite-initialization effect. Under the feature-wise descent setting in Section 6, we can approximate the excess loss through ∞

E(k) ≈ ∫

Vx dx ∝ k −(α1 −1) ,

(289)

where k as a function of t is determined by the characteristic timescale t(k) ≈

1 ζk ln ( ) ζk Ck ϵ

(290)

E/dk ln E with ζk /Ck ∝ k α2 −α1 . We can define the effective slope as keff = − ddln , where the derivatives are d dk = − α1k−1 ln t/dk and d ln t α2 α2 − α1 1⎛ α1 − α2 ⎞ ⎟. = + = ⎜α2 − (291) ζ dk k k [ln ( Ckk ) + ln ( 1ϵ )] k ⎝ ln ( Cζkkϵ ) ⎠

Combining them, we obtain keff (k, ϵ) =

α1 − 1 . 2 α2 − α1 −α ζk ln( C ϵ ) k

(292)

For the experiments in Figure 3 (right) and Figure 5, we use α2 < α1 , and thus the effective slope is larger than the theoretical prediction for finite initialization (ϵ > 0). We can also write keff as a function of t. From (290) we have ζk (t)t = ln(1/ϵ) + ln (

ζk ) ≈ ln(1/ϵ) + α ln ζk (t) + C, Ck

(293)

2 and C is a constant independent of ϵ, t. At the leading order of t this gives where we denote α ∶= α1α−α 2

ζk (t) ⋅ t ≈ ln(1/ϵ) + α ln ln(1/ϵ) − α ln t + C.

(294)

Taking it into (292) we obtain keff (t) ≈

α1 − 1 α [1 + ], α2 ln(1/ϵ) + α ln ln(1/ϵ) + C − α ln t

(295)

which increases as t increases, as we observe in Figure 3 (right) and Figure 5. Figure 5 For Figure 5, we use β = 1.5, P = 64, M = 256 (thus 256 data points) and run full batch GD with learning rate 0.02 for 30000 steps. The standard MLP has 256 hidden units. All models are initialized with N (0, ϵ2 ) where ϵ = 10−4 . According to Proposition 1, (50) corresponds to an NQF with sk 0 Ak (x) = √ [ ψ k (x) 2 P

ψk (x) ] 0

(296)

for the k-th hidden unit. Now we verify that the conditions in Section 6 (feature-wise descent) are satisfied. Since A(x) is block-diagonal with elements {Ak (x)}P k=1 and Ak,µ Ak,ν = Ak,ν Ak,µ , we have Aµ Aν = Aν Aµ and thus Assumption 1 is satisfied. sk ψ (xµ ). Given the block-diagonal structure, the eigenvalues associated with the k-th unit for sample µ are λµ,k = ± 2√ P k Because ψk (x) are Fourier modes uniformly sampled on the grid, they are orthogonal, leading to ∑µ λµ,k λµ,j = 0 for k ≠ j, satisfying the orthogonal condition in Theorem 6. Finally, we map the experimental variables to that in Section 6: Ckk =

8 8 2 2 −2θ −(θ+β) . ∑ λµ,k ∝ sk = k , rk = ∑ λµ,k yµ ∝ sk bk = k m µ m µ r2

−2(θ+β)

Therefore, the rk decays with exponent α2 = θ + β. Vk is given by Vk = 8Ckkk ∝ k k−2θ Plugging α1 and α2 into Section 6, the excess loss E(t) decays asymptotically as: α −1

E(t) = Θ (t

1 − α

(297)

= k −2β , identifying α1 = 2β 5 .

2β−1

) = Θ (t− θ+β )

(298)

This provides the prediction for the scaling law in Figure 5. Power Law on One-hot Data We can also observe power laws from one-hot data. We choose each input sample xµ ∈ Rp (for µ ∈ {1, . . . , m}) to be xµ = λµ eµ , where eµ is the basis vector (with 1 at the µ-th coordinate and 0 α1 α1 elsewhere). We choose the scaling to be λµ ∝ µ−(α2 − 2 ) and the label to be yµ ∝ µ− 2 . Then we consider a quadratic network f (x) = x⊺ W W ⊺ x = Tr[W W ⊺ A(x)], (299) where A(x) = xx⊺ = λ2µ Eµµ , where Eµµ is the matrix with 1 at the (µ, µ) entry and 0 elsewhere. Then we can find that this is the identical setting of Figure 3 (right). 5 Theorem 9 assumes an initialization z (0) = ϵc k −β0 with β ≥ α − α , whereas we initialize all weights from the same N (0, ϵ2 ). This 0 1 2 k k does not lead to a significant difference because P = 64 is very small.

Neural Quadratic Forms: A Unified Minimal Model for Sudden Learning and Scaling Laws

Analyze & explore related research in Conceptio

Related documents