ArXiv: 2502.01032
🎯 Pitch
Neural networks can be exactly rewritten as polynomials that explain over 95% of their output variance, even for nonlinear architectures like ReLU MLPs. This closed-form conversion reveals a sharp phase transition during training where models suddenly shift from learning simple linear patterns to higher-order features, offering the first direct evidence of a ‘simplicity bias’ in network development.
1. Executive Summary
This paper derives closed-form least-squares optimal approximations of feedforward networks (MLPs and gated linear units) using polynomial functions of arbitrary degree, assuming Gaussian or Gaussian mixture input distributions. Applying this framework to a single-hidden-layer ReLU MLP trained on MNIST, the authors compute linear and quadratic approximants and trace their fraction of variance unexplained (FVU) across training time, revealing a phase transition where linear FVU rises sharply while quadratic FVU stays nearly constant—evidence for a distributional simplicity bias in which networks learn statistics of increasing polynomial degree. The quadratic approximants explain over 95% of the variance in network outputs at the end of training, and linear approximant features transfer causally to the original network, enabling SVD-based adversarial attacks that reduce the MLP to near-random accuracy after ablating only ten input dimensions. The derivation establishes that polynomial approximants can capture both in-distribution and out-of-distribution behavior with high fidelity, but only when the Gaussianity assumption holds and the network's function class remains within the reach of low-degree polynomials—hard problems outside this regime cannot be interpreted through this approach.
2. Context and Motivation
The Core Problem: Interpreting MLPs Without Training Auxiliary Models
The central problem this paper addresses is deceptively simple: given a pretrained MLP, how can we understand what it has learned without training a separate interpretability model on its activations? This matters because the dominant paradigm in mechanistic interpretability—dictionary learning via sparse autoencoders (SAEs)—solves a different problem than the one the authors want to solve. The paper draws a sharp distinction (Section 1) between learning features of the model versus learning features of the data. SAEs, by their construction, conflate the two: they are trained to reconstruct MLP activations under a sparsity constraint on some dataset, meaning their learned features are a function of both the model's computation and the specific data distribution used during SAE training. It is unclear, the authors argue, whether SAE features reflect something fundamental about the model's inductive biases, or whether they are simply picking up on statistical patterns in the training corpus that the model happens to be sensitive to.
The paper proposes an alternative: assume the input data is maximum-entropy subject to low-order moment constraints, which for known mean and covariance yields the Gaussian distribution . Under this assumption, the model can be approximated analytically without sampling any data at all—the approximation depends only on the model's weights and the assumed input statistics. This flips the interpretability relationship: instead of asking "what features of the data does the model respond to?", we ask "what features of the model itself are revealed when we assume the simplest possible input distribution consistent with first and second moments?" The distinction is subtle but important. An SAE might find a feature that fires on dog images because the training data contains many dogs; the polynomial approximant reveals what the model would do to any input with the same mean and covariance structure, regardless of whether that input resembles natural data. This is particularly relevant for understanding out-of-distribution behavior—what the model does when faced with inputs unlike anything in its training set—which SAEs, trained on in-distribution data, are poorly equipped to characterize.
Why This Problem Matters: Interpretability Without Distributional Assumptions
The practical significance of this framing flows from a growing concern in the interpretability community. SAEs have demonstrated impressive results in extracting monosemantic features from language model activations (Huben et al., 2023; Paulo et al., 2024), but they come with a fundamental limitation: they require choosing a dataset on which to train the dictionary. This choice is consequential. The features an SAE learns when trained on, say, Wikipedia text versus GitHub code versus Pile data will differ, because the activation distribution differs across these domains. If the goal is to find a "true" decomposition of the model's computation—features that the model uses regardless of the data distribution—then SAEs provide only a data-conditional answer.
The polynomial approximation approach addresses this by making the input distribution a principled, minimal-assumption choice rather than an empirical one. The Gaussian distribution is the unique distribution that maximizes entropy given known mean and covariance. This means it encodes no additional assumptions beyond those first and second moments. If a feature emerges from the polynomial decomposition under this maximum-entropy assumption, that feature is likely a genuine property of the model's weight matrices rather than an artifact of the training data distribution. The authors are careful not to claim this is always the right distribution—they acknowledge in Section 5 that Gaussianity "may not be suitable for all applications"—but they argue it is an independently motivated default that prevents the approximant from "overfitting to the specific properties of a training dataset."
On the theoretical side, this work connects to a broader question about what neural networks actually learn during training. The idea that networks exhibit a simplicity bias—learning simple functions before complex ones—has been proposed in various forms (Nakkiran et al., 2019; Refinetti et al., 2023; Belrose et al., 2024), but testing this hypothesis has been challenging because "simplicity" is difficult to operationalize. The polynomial approximation framework provides a natural metric: the degree of the polynomial needed to achieve a given . If a network's function is well-approximated by a linear function (), it is "simple" in this sense. If it requires quadratic or higher-order terms, it is "more complex." Tracing this across training time provides quantitative evidence for or against the simplicity bias hypothesis, and the authors exploit this in Section 4 to show a clear phase transition from linear to quadratic representations.
Prior Approaches and Where They Fall Short
The paper positions itself against several strands of prior work, each with specific limitations that the polynomial approximation framework addresses.
Sparse autoencoders for dictionary learning. The most prominent approach to understanding MLP representations is training SAEs to produce sparse overcomplete representations of MLP activations (Huben et al., 2023; Paulo et al., 2024). While effective at finding interpretable features, SAEs have the confound the paper identifies: it is unclear whether they recover model features or data features. Moreover, SAEs introduce a new model (the autoencoder) with its own inductive biases, hyperparameters, and training instability. Interpreting the SAE's features requires trusting that the SAE faithfully captures the original model's computation, which is itself an unverified assumption. The polynomial approximation approach is non-parametric in the sense that no new model is trained—the coefficients are derived analytically from the original model's weights and the assumed input distribution.
Bilinear layers and interaction matrices. The most directly related prior work is Pearce et al. (2024), who showed that bilinear layers (GLUs with identity activation) produce outputs that are quadratic forms in the input, and that the associated Hessian matrices (called "interaction matrices") can be eigendecomposed to reveal interpretable features. This approach is elegant but restrictive: it requires the base model to be architecturally bilinear, which is not a standard design choice. Most deployed models use ReLU, GELU, or SwiGLU activations, not identity-gated bilinear layers. The present paper generalizes this insight by showing that any pretrained MLP or GLU with standard activations can be converted into a polynomial (and hence analyzed via eigendecomposition) through closed-form least-squares approximation, without requiring the original model to be architecturally polynomial. This is a substantial practical advance: it means interpretability researchers can apply eigendecomposition-based analysis to off-the-shelf pretrained models rather than needing to train custom architectures from scratch.
Linear probes and logit lens. Techniques like the logit lens (Millidge & Black, 2022) interpret transformer components by projecting weight matrices into vocabulary space using the unembedding matrix. While effective for language models, these approaches typically analyze first-order (linear) relationships and do not extend naturally to higher-order interactions. The polynomial approximation framework provides a principled way to capture quadratic and higher-order structure, which is particularly important for MLPs whose expressivity comes from nonlinearities that linear probes fundamentally cannot capture.
Gradient-based polynomial fitting. One could, in principle, fit a polynomial approximant to a neural network by sampling from some input distribution, applying the network, and running stochastic gradient descent to minimize MSE. This is straightforward but has two major drawbacks the paper identifies (Section 3): it is computationally intensive (requiring many forward passes through the original network), and it does not "afford a deeper theoretical understanding of the network's inductive biases." The analytic derivation provides interpretable formulas that reveal why certain polynomial coefficients take the values they do, in terms of the original network's weights and the input distribution's parameters—something a black-box SGD fit cannot provide.
Distributional simplicity bias (DSB) studies. Prior work by Refinetti et al. (2023) and Belrose et al. (2024) has provided evidence that neural networks trained with SGD learn statistical moments in increasing order—first mean, then covariance, then higher-order interactions. However, testing this hypothesis has previously required training models with specific architectural constraints (e.g., separate linear and quadratic branches) or analyzing learning dynamics in simplified teacher-student setups. The polynomial approximation framework provides a post-hoc test: given any pretrained MLP, compute the of its best linear and quadratic approximants at each training checkpoint, and observe whether the linear drops before the quadratic drops (indicating the network is transitioning from linear to quadratic computation). This is more general and more directly applicable to real models than previous methods.
How This Paper Positions Itself
The paper frames its contribution through several deliberate positioning choices that distinguish it from existing work.
It is a derivation paper more than a method paper. While the practical outcome—polynomial approximants that can be visualized and attacked—is significant, the paper's primary intellectual contribution is the theoretical machinery that makes closed-form approximation possible. The Master Theorem (Theorem 3.1) is presented as the centerpiece: a reduction of multivariate Gaussian integrals involving arbitrary nonlinearities into linear combinations of univariate integrals. This is not an incremental improvement over prior work; it is a new mathematical result that enables an entire class of analyses that were previously intractable. The paper explicitly contrasts this with gradient-based fitting, emphasizing that "surprisingly, it is possible to derive analytic formulas for these polynomial approximants" (Section 3), suggesting the authors view the derivation itself as the contribution rather than its empirical consequences.
It targets model features, not data features. This is the paper's most explicit positioning move relative to the SAE literature. The introduction states: "it is not understood how much [SAEs] are learning features of the model versus features of the data." The proposed solution—maximum entropy inputs—is presented as a way to isolate the model's inductive biases from the data's statistical structure. This is not a claim that SAEs are bad or wrong; it is a claim that they answer a different question than the one this paper asks. The polynomial approximation approach asks: "What does this model compute, in the abstract, independent of any particular dataset?" SAEs ask: "What features does this model compute on this particular dataset?" Both are valid, but the distinction has implications for how we think about model understanding: if an SAE finds a feature that fires on cat images, is that because the model has a "cat detector" circuit, or because cat images happen to activate certain neurons that are actually computing something more abstract? The polynomial approach, by removing the data distribution from consideration, helps disambiguate these possibilities.
It situates itself within the simplicity bias research program. The quadratic approximant's high (>95%) and the phase transition in linear FVU are presented as evidence for the distributional simplicity bias hypothesis—networks learn simple (linear) functions first, then progressively incorporate higher-order (quadratic) structure. This connects an abstract mathematical derivation to a concrete empirical question about learning dynamics. The paper is not just providing a tool for interpretability; it is using that tool to test a scientific hypothesis about how neural networks learn. This dual purpose—interpretability method and scientific instrument—is central to how the paper positions its contribution.
It acknowledges sharp scope limitations. Unlike many papers that overclaim generality, this one is explicit about when the approach does and does not apply. The Gaussianity assumption is described as "independently motivated" but the authors note it "may not be suitable for all applications" (Section 5). The computational cost of quadratic approximants scales as for input dimension , making exact analytic computation intractable for modern architectures without special structure (Appendix B). The authors are transparent that "if the best quadratic approximation for an MLP is poor, there is little that can be done to interpret the model using our approach in isolation." These limitations are not hedged; they are stated as boundary conditions. This establishes credibility and helps readers understand exactly when the method is applicable (low-dimensional inputs, Gaussian-like distributions, networks whose functions are well-approximated by quadratics) and when it is not (high-dimensional transformers, non-Gaussian data, highly nonlinear networks).
It connects to adversarial robustness as an explanatory validation. The SVD-based adversarial attack in Section 4.4 is not presented as a security contribution but as evidence that the polynomial approximants "capture the out-of-distribution behavior of the network." The fact that ablating only 10 SVD components (out of 784 input dimensions) reduces the MLP to near-random accuracy, while leaving images human-intelligible (Figure 4), demonstrates that the linear approximant's principal directions are causally important for the original model—not just correlated, but mechanistically relevant. This causal transfer from approximant to original model is the strongest evidence the paper provides that the polynomial approximations are not merely descriptive but capture something fundamental about the network's computation. The positioning is careful: the adversarial attack is not the goal; it is a validation that the approximant has learned real features of the model rather than spurious correlations.
In summary, the paper addresses a specific gap—understanding MLP computations without training auxiliary interpretability models on specific datasets—by deriving closed-form polynomial approximants under maximum-entropy input assumptions. It positions this as both a practical tool for interpreting pretrained models and a scientific instrument for testing hypotheses about neural network learning dynamics, while being unusually explicit about the conditions under which the approach does and does not work.
3. Technical Approach
3.1 Reader orientation
This is primarily a theoretical derivation paper that shows how to convert a pretrained neural network's feedforward layers into an exactly equivalent polynomial function under a specific assumption about the input distribution. The system being built is a mathematical pipeline: given the weights of an MLP or GLU and the mean and covariance of the input distribution (assumed Gaussian), compute the coefficients of the polynomial (linear, quadratic, or higher-degree) that minimizes mean squared error between the polynomial's outputs and the original network's outputs. The problem this solves is interpretability without auxiliary models — rather than training a sparse autoencoder to explain what the network does on a particular dataset, derive analytically what the network computes in the abstract, revealing features of the model itself rather than features of the data.
3.2 Big-picture architecture (diagram in words)
The system has four conceptual components connected in a linear pipeline:
-
Input distribution specification — the user provides the mean vector
$\mu \in \mathbb{R}^d$and covariance matrix$\Sigma \in \mathbb{S}^d_+$of a Gaussian distribution$\mathcal{N}(\mu, \Sigma)$(or a Gaussian mixture with$k$components). This is the maximum-entropy distribution consistent with those moments, encoding no additional assumptions. -
Original network — a pretrained single-hidden-layer MLP
$f(x) = W_2 \phi(W_1 x + b_1) + b_2$or GLU$\text{GLU}_\phi(x, W, V, b, c) = \phi(xW + b) \odot (xV + c)$, where$\phi$is an elementwise nonlinearity (ReLU, GELU, etc.) and weights are frozen. -
Moment evaluation engine — analytically computes the Gaussian expectations
$\mathbb{E}[f(x)]$,$\text{Cov}[f(x), x]$, and higher-order cross-moments$\mathbb{E}[\phi(y_i) x_k x_l]$needed to set up the least-squares problem. This uses Stein's lemma (for ReLU/GELU) or Gauss-Hermite quadrature (for arbitrary$\phi$), plus the Master Theorem to reduce multivariate integrals to univariate ones. -
Least-squares solver — solves the ordinary least squares problem
$(\beta^*, \alpha^*) = \arg\min \mathbb{E}\|f(x) - (\beta^T z + \alpha)\|^2$where$z$is the polynomial feature vector (e.g.,$z = [x, \phi_2(x)]$for quadratics). Outputs the coefficients$\alpha, \beta$of the optimal polynomial approximant.
Information flows linearly: input distribution parameters $\rightarrow$ moment computation using network weights $\rightarrow$ assembly of normal equations $\rightarrow$ solution for polynomial coefficients. The key enabling piece is the Master Theorem, which makes step 3 computationally tractable by reducing high-dimensional Gaussian integrals to sums of univariate integrals.
3.3 Roadmap for the deep dive
- First, the linear approximant derivation (Section 3.1 in the paper) — the simplest case, establishing the overall approach and introducing Stein's lemma as the core tool for evaluating Gaussian expectations of ReLU/GELU.
- Second, the quadratic feature map and the reduction to OLS (Section 3.2) — how higher-degree polynomials are handled by lifting inputs into a higher-dimensional feature space, and why the covariance matrix of these features is the computational bottleneck.
- Third, the Master Theorem (Section 3.3) — the paper's central theoretical contribution: a combinatorial reduction of multivariate Gaussian integrals involving arbitrary nonlinearities into linear combinations of univariate integrals. This is what makes the entire approach computationally feasible.
- Fourth, Gaussian mixture extensions (Section 3.4) — how the law of total covariance extends the single-Gaussian derivation to mixture distributions (e.g., modeling class-conditional statistics for MNIST).
- Fifth, adaptation to gated linear units (Section 3.5) — how the same machinery extends to GLU variants with minimal modification, exploiting joint Gaussianity of the two linear branches.
- Sixth, practical implementation decisions (Appendix B) — how the quadratic case becomes computationally intractable at input dimension 784 without exploiting special structure, and the hybrid analytic-initialization-plus-SGD-finetuning strategy used in the experiments.
3.4 Detailed, sentence-based technical breakdown
The Overall Strategy: Reduce Network Approximation to Gaussian Integration
The paper's core idea is that approximating an MLP by a polynomial under a Gaussian input assumption reduces to evaluating a specific set of Gaussian integrals. Once these integrals are computed in closed form (or via quadrature), the polynomial coefficients follow directly from the standard ordinary least squares formulas. The key insight is that the nonlinearity $\phi$ in $f(x) = W_2 \phi(W_1 x + b_1) + b_2$ makes direct least-squares fitting appear intractable, but when $x$ is Gaussian, the pre-activations $y = W_1 x + b_1$ are also Gaussian (since affine transformations preserve Gaussianity), and the expectations $\mathbb{E}[\phi(y)]$, $\mathbb{E}[\phi'(y)]$, and related quantities factorize in ways that exploit the Gaussian structure. This means the problem decomposes into (a) computing univariate Gaussian expectations of $\phi$ and its derivatives, which are known analytically for ReLU and GELU, and (b) computing covariances between $\phi(y_i)$ and the input coordinates $x_j$, which Stein's lemma reduces to $\text{Cov}(y_i, x_j) \mathbb{E}[\phi'(y_i)]$ — again a univariate integral.
The approach contrasts with gradient-based polynomial fitting in two ways. First, it requires zero sampling: the expectations are computed analytically from the network weights and the assumed input moments, not estimated from Monte Carlo samples. Second, it produces exact least-squares coefficients (to the precision of the numerical integration) rather than SGD approximations, meaning there is no optimizer noise or convergence uncertainty.
The Linear Case: Ordinary Least Squares with Gaussian Inputs
The paper begins with the simplest nontrivial case: approximating an MLP with an affine function $g(x) = \beta^T x + \alpha$.
The optimization problem. The objective is standard ordinary least squares:
where $f(x)$ is the original MLP output, the expectation is taken over $x \sim \mathcal{N}(\mu, \Sigma)$, and the norm is the Euclidean ($\ell_2$) norm.
What it computes: This finds the linear-plus-bias function that minimizes average squared prediction error over the Gaussian input distribution. The solution balances two competing objectives: matching the network's average output (captured by $\alpha$) and matching its sensitivity to each input dimension (captured by $\beta$). Because the objective is convex in $(\beta, \alpha)$ and the expectation is over a known distribution, the solution is unique.
Why this form: The $\ell_2$ loss is the natural choice because it makes the optimization analytically tractable (it yields linear normal equations) and it corresponds to the maximum-likelihood estimator under Gaussian noise assumptions. Mean absolute error or other losses would not admit closed-form solutions via covariance matrices.
The closed-form solution. For OLS, the solution is known to be:
where $\text{Cov}[x] = \Sigma$ is the $d \times d$ input covariance matrix, $\text{Cov}[f(x), x]$ is the $d_{\text{out}} \times d$ cross-covariance between network outputs and inputs, $\mathbb{E}[f(x)]$ is the $d_{\text{out}}$-dimensional mean output, and $\mathbb{E}[x] = \mu$ is the $d$-dimensional input mean.
What it computes: $\beta$ is the matrix that best predicts the centered output from the centered input: each row tells us how much each input dimension contributes to the corresponding output dimension, after accounting for all cross-correlations via $\Sigma^{-1}$. $\alpha$ shifts the linear prediction to center it at the correct mean. Together, they define the affine function that is the best linear unbiased predictor of $f(x)$.
Why this form: The inverse covariance appears because OLS must disentangle correlated inputs — if two input dimensions are highly correlated, their individual $\beta$ coefficients will differ from what naive univariate regression would produce. The cross-covariance $\text{Cov}[f(x), x]$ captures the raw linear relationship between each output and each input. Multiplying by $\Sigma^{-1}$ adjusts for correlations among the inputs to produce the partial regression coefficients.
Computing $\mathbb{E}[f(x)]$. To evaluate the mean of $f(x) = W_2 \phi(W_1 x + b_1) + b_2$ under $x \sim \mathcal{N}(\mu, \Sigma)$, the paper exploits the fact that the pre-activations $y = W_1 x + b_1$ are Gaussian with mean $W_1 \mu + b_1$ and covariance $W_1 \Sigma W_1^T$. By linearity of expectation:
Each coordinate $\mathbb{E}[\phi(y_i)]$ is a univariate Gaussian integral: $\mathbb{E}_{z \sim \mathcal{N}(\mu_i, \sigma_i^2)}[\phi(z)]$ where $\mu_i$ is the $i$-th component of $W_1 \mu + b_1$ and $\sigma_i^2$ is the $(i,i)$-th diagonal entry of $W_1 \Sigma W_1^T$.
For ReLU: The paper derives (Appendix A.2.1) the closed form:
where $\Phi$ is the standard normal CDF and $\phi$ is the standard normal PDF. This formula has an intuitive interpretation: the first term $\mu \Phi(\mu/\sigma)$ is the mean multiplied by the probability that the ReLU is active; the second term $\sigma \phi(\mu/\sigma)$ is a correction for the truncation (the expected value of the positive part of a Gaussian truncated at zero is higher than the unconditional mean when $\mu$ is negative).
For GELU: Using the identity $\text{GELU}(x) = x \Phi(x)$, the paper derives (Appendix A.1.1):
The key difference from ReLU is the $\sqrt{1 + \sigma^2}$ in the denominator inside $\Phi$ and $\phi$, which reflects GELU's smoother transition (it's never exactly zero, so the effective variance is higher).
For arbitrary $\phi$: When no closed form exists, the paper recommends Gauss-Hermite quadrature — a numerical integration technique that evaluates the integrand at a small number of carefully chosen points and weights them to exactly integrate polynomials up to a specified degree against the Gaussian weight function.
Computing $\text{Cov}[f(x), x]$. The cross-covariance between network outputs and inputs is:
The first term $\mathbb{E}[f(x) x^T]$ is the challenging part. By linearity:
The core difficulty is evaluating $\mathbb{E}[\phi(y_i) x_j]$ for each hidden unit $i$ and input dimension $j$. This is where Stein's lemma enters.
Stein's lemma. For jointly Gaussian random variables $(X, Y)$ and a differentiable function $g$:
Applied to our case, $X = y_i$ (the pre-activation of neuron $i$), $Y = x_j$ (input coordinate $j$), and $g = \phi$:
What this lemma does: It replaces a joint expectation involving $\phi(y_i)$ and $x_j$ with a product of two terms: the covariance between $y_i$ and $x_j$ (which is simply the $(i,j)$ entry of $W_1 \Sigma$, computable directly from the weights and input covariance), and the expected derivative of $\phi$ evaluated at $y_i$ (a univariate Gaussian integral). This is a massive simplification: instead of computing $d \times d_\text{hidden}$ bivariate integrals, we compute $d_\text{hidden}$ univariate integrals of $\phi'$ and then multiply by pre-computed covariances.
Expected derivatives for ReLU and GELU. For ReLU, $\phi'(z) = \mathbf{1}[z > 0]$, so:
This is simply the probability that the pre-activation is positive — the fraction of the Gaussian mass on the active side of the ReLU.
For GELU, $\phi'(z) = \Phi(z) + z \phi(z)$ (Appendix A.1.2), yielding a more complex but still closed-form expression involving $\Phi$ and $\phi$ evaluated at $\mu_i / \sqrt{1 + \sigma_i^2}$.
Assembly. Putting it all together: for each hidden unit $i$ and input dimension $j$, compute $\text{Cov}(y_i, x_j) = (W_1 \Sigma)_{i,j}$, compute $\mathbb{E}[\phi'(y_i)]$ using the closed form for the chosen activation, and multiply. The resulting matrix $\mathbb{E}[\phi(y) x^T]$ then left-multiplies by $W_2$ to produce $\mathbb{E}[f(x) x^T]$. Subtract $\mathbb{E}[f(x)] \mathbb{E}[x]^T$ to obtain $\text{Cov}[f(x), x]$. Finally, solve $\beta = \Sigma^{-1} \text{Cov}[f(x), x]^T$ and $\alpha = \mathbb{E}[f(x)] - \beta^T \mu$.
Design rationale. The linear case establishes the template: (1) express the OLS solution in terms of expectations, (2) use Gaussian properties to factor the expectations into univariate integrals of $\phi$ and $\phi'$, (3) evaluate those integrals analytically or via quadrature, (4) assemble. Everything that follows for quadratic and higher-order approximants is an elaboration of this same template, with the Master Theorem generalizing Stein's lemma to handle higher-order cross-moments.
The Quadratic Case: Lifting to a Polynomial Feature Space
To fit a quadratic function $g(x) = x^T A x + \beta^T x + \alpha$ (where $A$ is a matrix of quadratic coefficients and $\beta$ is the linear coefficient vector), the paper reduces this to a linear regression problem in a higher-dimensional feature space.
The quadratic feature map. Define the feature vector $\phi_2(x)$ containing all pairwise products of input coordinates:
This vector has dimension $\frac{d(d+1)}{2}$ — all unique pairs including squares ($i=j$). For example, with $d=3$, $\phi_2(x) = [x_1^2, x_1 x_2, x_1 x_3, x_2^2, x_2 x_3, x_3^2]$.
Why this specific set: Any quadratic function can be uniquely expressed as a linear combination of these $\frac{d(d+1)}{2}$ monomials plus the $d$ linear terms and the constant. Using all pairs (not just upper-triangular) would introduce redundancy (since $x_i x_j = x_j x_i$), making the design matrix rank-deficient. The upper-triangular choice ensures a minimal, non-redundant basis.
Concatenation. Let $z = [x, \phi_2(x)]$ be the concatenation of the linear and quadratic features. The dimension of $z$ is $d + \frac{d(d+1)}{2} = \frac{d(d+3)}{2}$. Any quadratic function of $x$ can now be written as an affine function of $z$:
for some coefficient vector $\gamma$ (concatenating the linear and quadratic coefficients) and bias $\alpha$. This is the classic "kernel trick in reverse": instead of implicitly mapping to a high-dimensional space via a kernel, we explicitly construct the feature map and solve linear regression in that space.
The OLS problem in $z$-space. The problem becomes:
with solution:
What it computes: The same OLS solution as before, but now the "input" is the expanded feature vector $z$. This means we need three new quantities: $\mathbb{E}[z]$ (means of all linear and quadratic monomials), $\text{Cov}[z]$ (the covariance matrix among these monomials), and $\text{Cov}[f(x), z]$ (cross-covariance between network outputs and each monomial).
Why this form: The convexity and closed-form nature of OLS are preserved because the problem is still linear in the parameters $(\gamma, \alpha)$ — the nonlinearity is entirely pushed into the feature construction, not the optimization. This means the quadratic approximation is the globally optimal quadratic function under the $\ell_2$ loss, not a local optimum from gradient descent.
The computational bottleneck: $\text{Cov}[z]$. The covariance matrix of $z$ contains moments of the form $\mathbb{E}[x_i x_j x_k x_l]$ for various index combinations. For a standard Gaussian $x \sim \mathcal{N}(0, I)$, this matrix is diagonal with simple structure (entries are 1 for linear monomials, 2 for squared monomials), which is why Appendix B initializes quadratic approximants under this assumption and then finetunes. For general $\mathcal{N}(\mu, \Sigma)$, the paper invokes the noncentral Isserlis theorem: the expected product of $n$ jointly Gaussian variables equals a sum over all partitions of the variables into singletons (contributing their means) and pairs (contributing their covariances). Concretely:
For $n=4$ (fourth moments, needed for quadratic $\text{Cov}[z]$), this sum has $2^{n/2} \cdot n! / (n/2)!$ terms (15 for $n=4$), each computable from the known $\mu$ and $\Sigma$. However, $\text{Cov}[z]$ has $\frac{d(d+3)}{2} \times \frac{d(d+3)}{2}$ entries, which is $O(d^4)$ — for $d=784$ (MNIST flattened), that's approximately $9.47 \times 10^{10}$ entries, exceeding GPU memory.
Why this matters for practicality: The $O(d^4)$ memory and $O(d^6)$ solve time for the quadratic case mean that exact closed-form computation is only feasible for low-dimensional inputs (the paper does not report a specific tractability threshold, but $d \lesssim 100$ would be a reasonable guess given cubic solve complexity). For high-dimensional problems like MNIST (784 dimensions) or transformer residual streams (typically 512–4096 dimensions), the exact quadratic solution is intractable. This is why the paper uses the hybrid strategy in Appendix B: initialize with the exact solution for $\mathcal{N}(0, I)$ (which is cheap due to diagonal structure), then finetune with SGD on actual data samples — the initialization provides a good starting point, and the SGD corrections adjust for the non-identity covariance.
The second bottleneck: $\text{Cov}[f(x), z]$. The cross-covariance between network outputs and the quadratic features requires evaluating expectations of the form $\mathbb{E}[\phi(y_i) x_k x_l]$ for each hidden unit $i$ and each pair of input indices $(k, l)$ (including $k=l$ for squared terms). This is where the Master Theorem becomes essential: it reduces this three-variable expectation ($\phi(y_i)$, $x_k$, $x_l$) to a linear combination of univariate integrals of the form $\mathbb{E}[\phi(y_i) y_i^m]$ for $m = 0, 1, 2$.
The Master Theorem: Reducing Multivariate Gaussian Integrals to Univariate Ones
Theorem 3.1 is the paper's central theoretical contribution. It addresses a general problem: given a Gaussian random variable $X$, a set of jointly Gaussian variables $Y_1, \ldots, Y_n$, and a continuous function $g$, how can we evaluate $\mathbb{E}[g(X) \prod_{i=1}^n Y_i]$ without computing an $(n+1)$-dimensional integral?
The theorem statement. For jointly Gaussian $X, Y_1, \ldots, Y_n$ and continuous $g$:
where the coefficients $a_k$ are determined by the OLS regressions of each $Y_i$ on $X$ and the residual covariances among the $\epsilon_i = Y_i - \mathbb{E}[Y_i \mid X]$.
What it computes: The left-hand side is an $(n+1)$-dimensional Gaussian integral involving an arbitrary function $g$ and a product of $n$ Gaussian variables. The right-hand side is a sum of at most $n+1$ univariate integrals — each of the form $\mathbb{E}[g(X) X^k]$ for $k=0,\ldots,n$ — weighted by coefficients $a_k$ that depend only on the mean and covariance structure of $(X, Y_1, \ldots, Y_n)$, not on $g$.
Why this form matters: The reduction from multivariate to univariate is the key computational win. Univariate integrals of $\phi(y_i) y_i^k$ are tractable for ReLU (the paper derives a complete formula using the upper incomplete gamma function in Appendix A.3) and for GELU (using Owen's integral tables, Appendix A.1.3). Without this theorem, evaluating quadratic approximants would require $d_\text{hidden} \times \frac{d(d+1)}{2}$ three-dimensional numerical integrals, each potentially requiring thousands of quadrature points — completely intractable for practical networks.
Proof sketch and operational meaning. The proof proceeds in three conceptual steps, each with a clear operational interpretation:
- Conditional expectation decomposition. Each
$Y_i$is decomposed into its conditional expectation on$X$plus a residual:$Y_i = \mathbb{E}[Y_i \mid X] + \epsilon_i$. Because$(X, Y_i)$is bivariate Gaussian, the conditional expectation is linear:$\mathbb{E}[Y_i \mid X] = \alpha_i + \beta_i X$, where$\alpha_i$and$\beta_i$are the OLS intercept and slope from regressing$Y_i$on$X$. The residuals$\epsilon_i$are zero-mean, jointly Gaussian, and — crucially — independent of$X$(a property specific to Gaussian distributions). The product becomes:
-
Polynomial expansion. Expanding the product yields
$3^n$terms, each being a product where the$i$-th factor is drawn from$\{\alpha_i, \beta_i X, \epsilon_i\}$. Terms are grouped by how many$\beta_i X$factors they contain. A term with$k$such factors contributes to the$X^k$coefficient. The combinatorial calculation (Equation 14) expresses the coefficient$a_k$as a sum over all$\binom{n}{k}$ways to choose which$k$variables contribute$\beta_i X$factors, multiplied by the expected value of the residual product for the remaining$n-k$variables. -
Residual expectation via Isserlis' theorem. The expected value of a product of zero-mean Gaussian residuals
$\epsilon_i$is computed using Isserlis' theorem: it equals zero if the number of residuals is odd, and a sum over all perfect matchings of pairwise covariances if the number is even. The covariance between residuals$\epsilon_i$and$\epsilon_j$is:
which is the partial covariance of $Y_i$ and $Y_j$ given $X$.
Concrete instantiation for $n=2$ (the quadratic case). The paper provides the fully expanded formula (Equation 16):
where:
$\alpha_i, \beta_i$are the OLS coefficients from regressing$Y_i$on$X$:$\beta_i = \frac{\text{Cov}(X, Y_i)}{\text{Var}(X)}$and$\alpha_i = \mathbb{E}[Y_i] - \beta_i \mathbb{E}[X]$.$\text{Cov}(\epsilon_1, \epsilon_2)$is the residual covariance as defined above.- The three terms correspond to
$k=2$(both$Y$s contribute their$X$-dependent parts),$k=1$(one contributes its$X$-dependent part, the other contributes its residual), and$k=0$(both contribute residuals whose product expectation is their covariance).
What this means operationally for quadratic approximation: For each hidden neuron $i$ and each pair of input indices $(k, l)$, set $X = y_i$ (the pre-activation), $Y_1 = x_k$, $Y_2 = x_l$, and $g = \phi$. Then apply Equation 16. The required univariate integrals are $\mathbb{E}[\phi(y_i)]$ (already computed for the linear case), $\mathbb{E}[\phi(y_i) y_i]$, and $\mathbb{E}[\phi(y_i) y_i^2]$. For ReLU, these are evaluated using the generalized formula in Appendix A.3 (Equation 36):
where $\Gamma(s, x) = \int_x^\infty t^{s-1} e^{-t} dt$ is the upper incomplete gamma function, available in standard scientific computing libraries. The computation requires evaluating this sum for $n=0,1,2$ (to get the three needed moment integrals), reusing intermediate values for efficiency.
Design rationale: Why this particular decomposition? The Master Theorem works because of a special property of Gaussian distributions: conditional expectations are linear and residuals are independent of the conditioning variable. This means the $(n+1)$-dimensional integral factors into a sum of products of (a) deterministic coefficients from the $\alpha_i, \beta_i$ and the residual covariances, which depend only on the covariance structure, and (b) univariate integrals $\mathbb{E}[g(X) X^k]$ which depend on $g$ and the marginal distribution of $X$. This separation of concerns — structural coefficients from functional integrals — is what makes the theorem both general (any continuous $g$) and practical (the coefficients are computed once per network, the integrals once per activation function).
Gaussian Mixture Inputs: Extending Beyond a Single Gaussian
The derivation so far assumes $x \sim \mathcal{N}(\mu, \Sigma)$. But many real datasets — including MNIST — are better modeled as mixtures of Gaussians, with each mixture component corresponding to a class or cluster. The paper extends its framework to this case using the law of total covariance.
The law of total covariance. For random vectors $X$ and $Y$ whose joint distribution is a mixture over a discrete latent variable $Z \sim \text{Cat}(k)$ with $k$ components:
The total cross-covariance is the sum of (a) the expected within-component cross-covariance, and (b) the cross-covariance of the component-conditional means.
What it computes: The first term $\mathbb{E}[\Sigma_{XY \mid Z}]$ averages the cross-covariances that would be computed if we knew the component label — each component's cross-covariance is computed using the single-Gaussian formulas from the previous sections, with that component's $(\mu_z, \Sigma_z)$. The second term $\Sigma_{\mathbb{E}[X \mid Z], \mathbb{E}[Y \mid Z]}$ captures the cross-covariance arising from the fact that different components have different means — this is the "between-component" variation. For example, in MNIST, the mean of class "3" is different from the mean of class "7", and this mean variation contributes to the total cross-covariance even if within each class the relationship between $X$ and $Y$ is identical.
Why this form: The decomposition separates within-component structure (captured by the first term, computed analytically per component) from between-component structure (captured by the second term, computed from the component means and mixing weights). This means we can apply the single-Gaussian machinery to each mixture component independently, then combine results using only the component means and mixing probabilities.
Computational procedure:
-
For each mixture component
$z \in \{1, \ldots, k\}$with mean$\mu_z$, covariance$\Sigma_z$, and mixing weight$w_z$:- Compute
$\mathbb{E}[f(x) \mid Z=z]$using the single-Gaussian mean formula with$\mathcal{N}(\mu_z, \Sigma_z)$. - Compute
$\text{Cov}[f(x), x \mid Z=z]$using the single-Gaussian cross-covariance formula with$\mathcal{N}(\mu_z, \Sigma_z)$.
- Compute
-
Compute the expected within-component cross-covariance:
$\mathbb{E}[\Sigma_{XY \mid Z}] = \sum_{z=1}^k w_z \text{Cov}[f(x), x \mid Z=z]$. -
Compute the between-component contribution:
$\Sigma_{\mathbb{E}[X \mid Z], \mathbb{E}[Y \mid Z]} = \sum_{z=1}^k w_z (\mathbb{E}[f(x) \mid Z=z] - \mathbb{E}[f(x)]) (\mu_z - \mu)^T$. -
Sum to obtain
$\text{Cov}[f(x), x]$. Apply the OLS formula as before.
The same logic applies to $\text{Cov}[z]$: For the quadratic case, compute $\text{Cov}[z \mid Z=z]$ per component using the noncentral Isserlis theorem, then apply the law of total covariance to obtain the mixture covariance. This is where the computational burden compounds: each component requires computing a $\frac{d(d+3)}{2} \times \frac{d(d+3)}{2}$ matrix of fourth moments. For MNIST with 10 classes and $d=784$, this is the point where the paper declares exact computation "intractable" (Appendix B) and switches to the hybrid SGD approach.
Design rationale for the mixture extension: The authors could have simply fit a single Gaussian to the entire dataset and accepted the approximation error. The mixture extension is motivated by the fact that MNIST classes are well-separated in pixel space — the mean of "0" is very different from the mean of "1" — so a single Gaussian would be a poor model (it would have high variance in directions that separate classes, attributing to noise what is actually class structure). The mixture model captures this class structure while still being analytically tractable through the law of total covariance.
Gated Linear Units: Adapting the Framework to GLU Architectures
Modern transformers often use Gated Linear Units (GLUs) rather than standard MLPs. A GLU has the form:
where $\odot$ is elementwise (Hadamard) multiplication, $W$ and $V$ are weight matrices (potentially of different shapes, though typically the same), $b$ and $c$ are bias vectors, and $\phi$ is an elementwise activation (often identity for "bilinear" layers, Swish for SwiGLU, or ReLU for ReGLU).
Adaptation of the framework. The key observation is that when $x \sim \mathcal{N}(\mu, \Sigma)$, the two pre-activation vectors $y = xW + b$ and $z = xV + c$ are jointly Gaussian — each is an affine transformation of the same Gaussian vector, so their joint distribution is multivariate normal with means $\mathbb{E}[y] = \mu W + b$, $\mathbb{E}[z] = \mu V + c$, and cross-covariance $\text{Cov}(y_i, z_j) = (W^T \Sigma V)_{i,j}$.
Mean computation. Each output coordinate of the GLU is $\phi(y_i) z_i$. Its expectation is:
The covariance term can be evaluated via Stein's lemma (treating $\phi(y_i)$ as the function and $z_i$ as the Gaussian variable):
This is directly analogous to the MLP case: the first term captures how the gating modulates the linear pathway based on their covariance, and the second term is the product of the independent expectations.
Cross-covariance computation. For the linear approximant, we need $\text{Cov}[\text{GLU}_\phi(x), x]$. Each entry $(i, j)$ is $\mathbb{E}[\phi(y_i) z_i x_j] - \mathbb{E}[\phi(y_i) z_i] \mathbb{E}[x_j]$. The expectation $\mathbb{E}[\phi(y_i) z_i x_j]$ is a three-variable integral of the form covered by the Master Theorem with $n=2$, $X = y_i$, $Y_1 = z_i$, $Y_2 = x_j$, and $g = \phi$. The paper notes that "similar arguments apply, mutatis mutandis, to higher-order polynomial approximants" — meaning the same Master Theorem machinery handles quadratic and beyond for GLUs.
Why GLU is handled separately: At first glance, a GLU looks different from an MLP — it has a multiplicative interaction between two linear pathways rather than a single nonlinearity applied to a linear projection. But the mathematical structure is identical from the perspective of Gaussian integration: in both cases, the challenging term is $\mathbb{E}[\phi(\text{affine}_1(x)) \cdot \text{affine}_2(x) \cdot \text{affine}_3(x)]$, and the Master Theorem reduces it the same way regardless of whether $\text{affine}_2$ is a constant (as in the MLP $\mathbb{E}[\phi(y_i) x_j]$ case, where the "second affine" is simply $x_j$ with zero bias and identity weight) or a separate linear pathway (as in the GLU case).
Computational Complexity and the Practical Limits of the Exact Approach
The paper is notably honest about when its approach becomes intractable, and this honesty shapes the experimental design.
Scaling with degree $n$. For a degree-$n$ polynomial approximant, the feature vector $z$ has dimension $\binom{d+n}{n}$ (the number of monomials of degree $\leq n$ in $d$ variables). The covariance matrix $\text{Cov}[z]$ has $\binom{d+n}{n}^2$ entries, each requiring computation of a $2n$-order Gaussian moment via the noncentral Isserlis formula (which involves super-exponentially many terms in $n$). Solving the linear system takes $O(\binom{d+n}{n}^3)$ time. For $d=784$ and $n=2$, this is $O(10^{15})$ operations — clearly not feasible.
Why $n=2$ is the practical limit. The paper explicitly states: "In a world where intermediate activations can have dimensionality in the tens of thousands, $n=2$ is the highest that can feasibly be computed analytically" (Section 5, Limitations). This is not a theoretical bound — higher-degree approximants are mathematically well-defined — but a computational one driven by the cubic scaling of linear system solves with the expanded feature dimension.
The hybrid strategy (Appendix B). For the quadratic MNIST experiments, the paper uses a practical compromise: initialize with the exact solution for $\mathcal{N}(0, I)$ inputs (where $\text{Cov}[z]$ is diagonal and trivially invertible), then finetune with SGD on batches of samples from the actual Gaussian mixture distribution. Since the objective (quadratic MSE) is convex in the polynomial coefficients, SGD converges to the true optimum — the initialization just provides a warm start that avoids the need to materialize or invert the full $\text{Cov}[z]$ matrix. This is described as an engineering decision, not a theoretical limitation: "the resulting coefficients should be excellent approximations of the true least-squares values" because the objective is convex.
Cost of the linear approximant. For the linear case ($n=1$), the computation is fully tractable even for moderate $d$: $\text{Cov}[z] = \Sigma$ is $d \times d$, the solve is $O(d^3)$, and the Master Theorem reduces to Stein's lemma (only $n=1$). This is why the linear approximant experiments in the paper are exact, while the quadratic experiments use the hybrid approach.
Summary of Design Choices and Their Justifications
-
Maximum-entropy Gaussian assumption over empirical distributions: This prevents the polynomial approximant from overfitting to the specific properties of a training dataset. The Gaussian is the unique distribution that maximizes entropy given known mean and covariance, meaning it encodes the fewest additional assumptions. If a feature appears in the polynomial decomposition under this assumption, it is likely a genuine property of the network's weights, not a statistical artifact of the training data.
-
$\ell_2$loss for approximation: Makes the optimization convex with closed-form solution via the normal equations. This is essential for the analytic derivation — other loss functions (e.g.,$\ell_1$, cross-entropy on softmax outputs) would not admit closed-form Gaussian integrals. -
Stein's lemma for first-order cross-moments: Reduces bivariate integrals
$\mathbb{E}[\phi(y_i) x_j]$to univariate integrals$\mathbb{E}[\phi'(y_i)]$multiplied by known covariances. This is the key enabler for the linear case and would be impossible without the Gaussian assumption (Stein's lemma requires joint Gaussianity). -
Master Theorem for higher-order cross-moments: Generalizes Stein's lemma from
$n=1$to arbitrary$n$, using the fact that Gaussian conditional expectations are linear and residuals are independent. This is the paper's novel mathematical contribution and what makes quadratic (and theoretically higher) approximants computationally feasible. -
Noncentral Isserlis theorem for
$\text{Cov}[z]$: Handles the covariance among polynomial features, which involves moments up to order$2n$. The partition-based formula is computationally heavy but exact. -
Hybrid analytic-plus-SGD for high-dimensional quadratics: Acknowledges that
$O(d^6)$exact computation is intractable, and substitutes a convex SGD procedure that converges to the same optimum. The$\mathcal{N}(0, I)$initialization exploits the special diagonal structure in that case to provide a warm start. -
Gaussian mixture via law of total covariance rather than re-deriving per mixture: Reuses the single-Gaussian formulas for each component and combines using only component means and weights. This is cleaner than re-deriving the Master Theorem for mixture distributions directly.
4. Key Insights and Innovations
Innovation 1: The Maximum-Entropy Framing — Interpreting Models by Removing the Data Distribution
The most intellectually distinctive move in this paper is not the mathematical derivation itself, but the conceptual inversion it enables: instead of asking "what features does this model compute on this dataset?", ask "what features does this model compute when we assume the simplest possible input distribution consistent with low-order statistics?" This reframes the interpretability problem from one of characterizing model behavior on a specific data distribution to one of revealing the model's intrinsic inductive biases. The maximum-entropy principle provides the philosophical justification: the Gaussian $\mathcal{N}(\mu, \Sigma)$ is the unique distribution that maximizes entropy given known mean $\mu$ and covariance $\Sigma$. It encodes absolutely no additional assumptions. Under this distribution, any structure that emerges in the polynomial approximant must come from the model's weights — it cannot be attributed to statistical quirks of a particular dataset.
This is a fundamental departure from the dominant interpretability paradigm. Sparse autoencoders (Huben et al., 2023; Paulo et al., 2024) require choosing a dataset on which to train the dictionary. That choice is consequential: features learned on Wikipedia differ from features learned on code. The SAE user must trust that their chosen dataset is representative of what the model "really" computes, but there's no principled way to establish this. The polynomial approximant under maximum-entropy inputs answers a different question entirely — not "what does the model do on distribution D?", but "what does the model do in the abstract, independent of any particular D?" The authors explicitly make this distinction in the introduction when they note that SAEs conflate features of the model with features of the data, and that their approach aims to isolate the former.
The significance of this reframing extends beyond the specific polynomial technique developed in this paper. It suggests a general principle for interpretability research: when you want to understand the model rather than the data, strip away the data distribution as aggressively as possible. The Gaussian maximum-entropy assumption is one way to do this; future work might explore other minimal-assumption input specifications. The fact that quadratic approximants under this assumption explain over 95% of variance in MLP outputs on MNIST (Figure 2, end of training) provides preliminary evidence that this approach can capture real model structure despite — or perhaps because of — the deliberate simplification of the input distribution.
The distinction from Pearce et al. (2024) is instructive. That work showed that bilinear layers are quadratic forms in the input, and that their Hessians ("interaction matrices") are eigendecomposable into interpretable features. But this required building the model to be quadratic from the start — an architectural constraint. This paper shows that the same eigendecomposition-based analysis can be applied to any pretrained MLP by converting it into a polynomial approximant post-hoc. The intellectual move is from "build models whose structure is interpretable by construction" to "reveal structure in models that were not built to be interpretable." This is a significant practical and conceptual advance, because it means the vast library of existing pretrained models becomes amenable to this analysis without architectural modification.
Innovation 2: Polynomial Degree as a Dynamic Measure of Model Complexity During Training
The paper introduces a diagnostic concept that is both theoretically motivated and empirically operational: the degree of the best polynomial approximant as a measure of the complexity of the function a network has learned. This transforms the intuitive notion of "network complexity" (often measured by weight norms, sharpness, or spectral complexity) into something with direct functional meaning: a network is "simple" if a linear function explains nearly all its output variance, "somewhat complex" if quadratics are needed, and "highly complex" if higher-order terms are required.
This framing is novel because prior work on simplicity bias (Nakkiran et al., 2019; Refinetti et al., 2023; Belrose et al., 2024) had to rely on indirect proxies — spectral analysis of weight matrices, teacher-student setups with controlled target functions, or models with architecturally separated linear and nonlinear branches. These approaches either couldn't be applied post-hoc to arbitrary pretrained models, or required training custom architectures, or both. The polynomial approximation framework provides a post-hoc, architecture-agnostic diagnostic: take any pretrained MLP (or GLU), compute the $R^2$ of its best linear and quadratic approximants under the maximum-entropy input assumption, and compare. If the linear $R^2$ is high, the network is functionally linear; if it's low but the quadratic $R^2$ is high, the network is functionally quadratic; and so on.
The empirical result that makes this diagnostic compelling is the phase transition in linear FVU documented in Figure 2. Between 500 and 1000 training steps, the linear FVU rises sharply while the quadratic FVU remains nearly constant. The authors interpret this as a "quadratic phase" of training — the network, having initially learned a largely linear function, now incorporates second-order statistical structure. This is not a gradual drift but a qualitative change in the functional form, visible as a kink in the FVU curves. The fact that this transition appears at a specific moment in training, and that it coincides with the period where the network's accuracy is improving (the paper notes this in Section 4.2), provides concrete evidence for the distributional simplicity bias hypothesis: networks genuinely do learn statistics of increasing polynomial degree as training progresses.
The KL divergence plot (Figure 3) strengthens this interpretation by showing that the anomalous early decrease in FVU (present in the first few hundred steps of Figure 2) disappears when using KL divergence instead of FVU as the discrepancy measure. This suggests that the early decrease in FVU is a metric artifact, not a genuine simplification of the network — the KL divergence, which weights distributional differences differently, shows a monotonic complexity increase throughout training. This diagnostic move — using multiple discrepancy measures to distinguish genuine phenomena from metric artifacts — is a methodological contribution in its own right, even if the paper doesn't belabor it.
It's important to note what this diagnostic does and doesn't capture. It measures complexity relative to the maximum-entropy input distribution, not absolute functional complexity. A network that is highly nonlinear on natural images but linear on Gaussian noise would be classified as "simple" by this metric. This is not a bug but a feature: it means the diagnostic is specifically sensitive to structure that is inherent to the network's weight matrices, not structure learned from the training data distribution. The interpretation of the phase transition is therefore that the network's weight matrices themselves undergo a qualitative change during training — from encoding primarily linear relationships to encoding quadratic ones — independent of what the training data looks like.
Innovation 3: Causal Transfer from Polynomial Approximants to Original Networks — Interpreting by Attacking
The paper makes a methodological contribution that goes beyond descriptive approximation: it demonstrates that features extracted from the polynomial approximant are causally relevant to the original network's behavior. This is shown through the SVD-based adversarial attack in Section 4.4, but the deeper point is not about attack efficacy — it's about validation by intervention. The logic is: if the linear approximant accurately captures the directions in input space that the original network is sensitive to, then removing those directions (via orthogonal projection onto the nullspace of the top SVD components) should degrade the original network's performance. If it doesn't, the approximant is capturing spurious correlations, not genuine model features.
Figure 5 shows that this logic holds with striking precision. The original ReLU MLP's accuracy drops in "perfect lockstep" with the linear and quadratic approximants' accuracy as SVD components are ablated. Ablating just 4 out of 784 input dimensions brings the MLP below 50% accuracy; ablating all 10 SVD components (only ~1.3% of input dimensions) reduces it to near-random performance. Meanwhile, the projected images remain human-intelligible (Figure 4, bottom row) — this is a genuine adversarial effect, not just destroying the input. The causal transfer from approximant to original model is strong evidence that the approximant has captured something mechanistically real about how the network processes inputs.
This "interpretability by attack" methodology is conceptually significant independent of the specific polynomial technique. It addresses a persistent problem in interpretability: correlation does not imply mechanistic relevance. An SAE might find a feature that correlates with some behavior, but does the model actually use that feature to make decisions? Without causal intervention, it's impossible to know. The SVD projection attack provides a simple causal test: if the feature matters, removing it should change behavior. If it doesn't, the feature is epiphenomenal. The fact that the linear approximant's top principal directions pass this test — and pass it so decisively — validates the entire approximation framework.
There's a subtle but important distinction between this approach and standard adversarial attack methods. Typical adversarial attacks (FGSM, PGD) optimize perturbations to maximize loss for a given input. The SVD attack here is universal: it applies the same projection to all inputs, regardless of their class or content. This means the attack effectively identifies input directions that the model uses in general, not directions that happen to be fragile for specific examples. The fact that a universal, low-rank projection is so devastating suggests that the model's decision boundary depends heavily on a small number of linear features — precisely the features captured by the linear approximant's top SVD components. This is a finding about the model's inductive bias, not about adversarial vulnerability per se: the model has learned to rely on a low-dimensional linear subspace, making it simultaneously interpretable (via that same subspace) and fragile (to perturbations orthogonal to it).
Innovation 4: The Master Theorem as a Bridge Between Neural Network Analysis and Gaussian Integration Theory
While the paper's empirical contributions (simplicity bias evidence, adversarial transfer) are the most immediately impactful for practitioners, its most enduring intellectual contribution may be the Master Theorem (Theorem 3.1) — a combinatorial reduction that converts multivariate Gaussian integrals involving arbitrary continuous nonlinearities into linear combinations of univariate integrals. This is not an incremental improvement on existing techniques; it is a new mathematical result that enables an entire class of analyses that were previously computationally intractable.
To appreciate why this is fundamental rather than incremental, consider the alternative. Without the Master Theorem, computing the cross-covariance $\text{Cov}[f(x), z]$ for a quadratic approximant would require evaluating $d_\text{hidden} \times \frac{d(d+1)}{2}$ three-dimensional integrals numerically. Each such integral, if done via quadrature with $m$ points per dimension, would cost $O(m^3)$ — and $m$ would need to be large (hundreds) to achieve acceptable precision for $\phi$ applied to Gaussian variables. The total cost would be $O(d_\text{hidden} d^2 m^3)$, which is billions of quadrature evaluations even for tiny networks. The Master Theorem reduces this to $O(d_\text{hidden} \times n)$ univariate integrals (where $n$ is the polynomial degree), plus combinatorial coefficient computation that depends only on the covariance structure and is computed once per network.
The theorem works because it exploits a structural property of Gaussian distributions that is not true in general: conditional expectations are linear, and residuals are independent of the conditioning variable. Any distribution with this property would admit a similar reduction, but the Gaussian is unique among continuous distributions in satisfying it for all variables simultaneously. The proof's combinatorial argument — expanding a product of $n$ linear-plus-residual terms, grouping by how many $\beta_i X$ factors appear, and applying Isserlis' theorem to residual products — is elegant in its simplicity once stated, but far from obvious a priori. It's the kind of result that, after reading, makes you wonder why no one derived it before. The likely answer is that the neural network interpretability community and the Gaussian integration theory community have had little overlap until recently.
The theorem's generality is also notable. It applies to any continuous $g$, not just ReLU or GELU. This means it can handle activation functions for which no closed-form Gaussian integrals exist — the univariate integrals $\mathbb{E}[g(X) X^k]$ can be evaluated via Gauss-Hermite quadrature or other numerical methods, which are efficient in one dimension. The coefficients $a_k$ depend only on the covariance structure and are independent of $g$. This clean separation of structural coefficients from functional integrals mirrors the separation-of-concerns principle that makes the whole approach practical: compute the structural coefficients once for the network architecture and input distribution, then plug in different activation functions essentially for free.
One might ask: why doesn't the field already have this theorem? Integral tables for Gaussian expectations of specific functions (ReLU, GELU, etc.) have existed for decades (Owen, 1980). Stein's lemma is a standard tool in statistics. But systematically reducing a product of $n$ Gaussian variables times an arbitrary function to univariate form — and providing a combinatorial formula for the coefficients — is genuinely new. The paper positions this not as "we applied known math to a new domain" but as "we derived a new mathematical result that enables this domain." The distinction matters for how the contribution should be evaluated: it's a theoretical advance, not just an application.
Innovation 5: Negative Results as Boundary Specification — When the Method Works and When It Fails
The paper's explicit and detailed discussion of limitations (Section 5) is not merely honest — it is methodologically innovative in how it uses negative results and computational barriers to specify the precise conditions under which the approach is applicable. Rather than claiming generality and leaving readers to discover the failure modes themselves, the paper provides a clear specification: the method works when (a) the input distribution is approximately Gaussian or Gaussian mixture, (b) the network's function lies within the reach of low-degree polynomials (quadratic $R^2$ is high), and (c) the input dimension is small enough that the exact quadratic solution is computationally feasible (or a hybrid SGD approach is acceptable). Where these conditions fail — highly non-Gaussian data, deeply nonlinear networks, high-dimensional transformer activations — the paper explicitly says the approach "may not be suitable" or that "there is little that can be done."
This boundary specification is a form of diagnostic contribution: it tells future researchers not just that the method exists, but when to use it. The computational intractability of exact quadratic approximants for $d > \mathcal{O}(100)$ (implied by the $O(d^6)$ scaling) is presented not as a failure but as a natural limit: "in a world where intermediate activations can have dimensionality in the tens of thousands, $n=2$ is the highest that can feasibly be computed analytically." This is a useful constraint for the field — it means researchers working on large language models should not expect to directly apply exact quadratic approximations to transformer FFNs, but might apply linear approximations or use the hybrid SGD approach, or decompose the problem differently.
The ReST-EM-style failure for revision models in Appendix K (the companion work Belrose et al., 2024, not this paper) similarly serves as a boundary condition. The fact that additional sequential revisions hurt performance with a certain training methodology tells us something about when revision-based approaches are fragile — specifically, when on-policy data collection amplifies spurious correlations. Even though that result is from related work, the instinct to report negative results as illuminating boundaries rather than hiding them is characteristic of the paper's approach.
This style of contribution — negative results as boundary specification — is underappreciated in a field that often rewards only positive results. It is valuable because it prevents wasted effort: a researcher reading this paper knows immediately whether their problem setting is likely to benefit from polynomial approximation or whether they should look elsewhere. The authors even provide actionable alternatives for the failure cases: "we recommend initializing a gradient-based optimizer with our analytic polynomial approximant, then finetuning on samples from the desired target distribution" for non-Gaussian inputs, and they flag the need for tensor decomposition methods more sophisticated than eigendecomposition for overcomplete quadratic bases in high dimensions.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. All experiments use MNIST (LeCun et al., 1998), the 28×28 grayscale handwritten digit classification benchmark consisting of 60,000 training and 10,000 test images across 10 classes. The authors flatten each image into a 784-dimensional vector. MNIST is chosen explicitly because it is "known to be very well-modeled as a Gaussian mixture distribution" (Section 4.1), with the claim that "samples from a Gaussian mixture fit to MNIST are difficult for a human to distinguish from real samples" (citing Belrose et al., 2023, Figure 2). This makes it an ideal testbed for a method that assumes Gaussian mixture inputs.
-
Base model. A single-hidden-layer MLP with ReLU activation, trained from scratch on MNIST. The paper does not specify the hidden layer width in the main text; it is described only as "a single hidden layer MLP with ReLU activation" (Section 4.1). Training uses schedule-free AdamW (Defazio et al., 2024) with 1,000 warmup steps, batch size 64, and weight decay 0.1. Checkpoints are saved at log-spaced intervals across training. The choice of a shallow, narrow-domain model is deliberate: the paper is validating the derivation and the simplicity bias hypothesis, not competing on benchmark accuracy. The model is small enough that exact linear approximants are computationally tractable, and the quadratic approximants are feasible via the hybrid analytic-plus-SGD approach described in Appendix B.
-
Metrics. The primary metric is fraction of variance unexplained (FVU), defined as
$1 - R^2$where$R^2$is the coefficient of determination between the original network's outputs and the polynomial approximant's outputs, computed over the Gaussian or Gaussian mixture input distribution. For classification models, the outputs are the 10-dimensional logits before softmax. FVU measures how much of the network's output variance remains after subtracting the polynomial approximation — an FVU of 0 means perfect approximation, while an FVU of 1 means the approximant explains none of the variance. The paper also reports KL divergence (in nats) between the network's output distribution and the approximant's output distribution, evaluated on Gaussian mixture samples (Figure 3). For the adversarial attack experiments (Section 4.4), the metric is test set accuracy of the original MLP after applying input projections derived from the polynomial approximants. For qualitative feature visualization (Figure 1), the metric is visual interpretability of the top eigenvector or linear coefficient vector for a specific class. -
Baselines. The paper does not use baseline methods in the conventional sense (it is not proposing a new training algorithm or architecture that needs comparison to prior work). Instead, it compares linear approximants against quadratic approximants as measures of model complexity at each training checkpoint. The key comparison is between the FVU of the best linear fit and the FVU of the best quadratic fit at the same checkpoint — when linear FVU is low, the network is functionally linear; when linear FVU rises while quadratic FVU remains low, the network has entered a quadratic phase. For the adversarial attack transfer experiment (Figure 5), the implicit baseline is random chance accuracy (10% for 10-class MNIST), which the MLP approaches after ablating all 10 SVD components.
-
Generation budget / compute accounting. There is no "generation budget" in the sense of inference-time sampling, since this is a deterministic analytic method. Compute is measured in terms of training steps of the original MLP (x-axis of Figures 2 and 3), with checkpoints saved at log-spaced intervals from roughly step
$2^4$(16) to step$2^{18}$(262,144) in Figure 1. The computational cost of the polynomial approximation itself is characterized asymptotically:$O(d^6)$for the exact quadratic solution, which becomes intractable for$d=784$without the hybrid SGD approach. The paper does not report wall-clock time or FLOPs for computing approximants. -
Cross-validation / statistical protocol. The paper does not employ cross-validation or statistical significance testing. The FVU and KL divergence metrics are computed analytically (for the linear case) or via convex SGD optimization (for the quadratic case, Appendix B) over the entire Gaussian mixture distribution matching the MNIST training set statistics. There is no train/validation/test split for the polynomial approximation itself — the approximant is fit to the maximum-entropy distribution parameterized by the full dataset's class-conditional means and covariances, not to a finite sample. The adversarial attack evaluation uses the standard 10,000-example MNIST test set to measure accuracy after SVD projection, but this is a single evaluation without confidence intervals or multiple seeds. For the training dynamics analysis, a single training run is used — the paper does not report results over multiple random seeds to assess the reliability of the observed phase transition.
Main Quantitative Results
Training Dynamics: The Linear-to-Quadratic Phase Transition
The paper's central empirical finding is the behavior of the FVU of linear and quadratic approximants across training time, presented in Figure 2. At the start of training (step ~$2^4 = 16$), the linear FVU is approximately 0.025 (2.5% of variance unexplained), meaning a linear function explains ~97.5% of the network's output variance under the Gaussian mixture input distribution. This is consistent with the network being near-initialization, where random weight matrices produce approximately linear behavior.
Between 500 and 1,000 training steps, there is a sharp increase in linear FVU — from roughly 0.02 to roughly 0.08 by step $2^{10} = 1024$, continuing to rise to approximately 0.10–0.12 by step ~4,000 before reaching a plateau around 0.15–0.18 in the later stages of training (steps $2^{14} = 16,384$ and beyond). The authors describe this as a "clear phase transition where the FVU and KL divergence for the linear approximant start to increase sharply, and continue to do so until around 4K steps" (Section 4.2).
Critically, over the same 500–4,000 step period where linear FVU rises, the quadratic FVU remains nearly constant at approximately 0.01–0.02. This means a quadratic function continues to explain 98–99% of the output variance even as the linear explanation degrades. The paper interprets this as a "quadratic phase" of training: "the network learns to exploit second-order statistical information in the input" (Section 4.2). The quadratic FVU does eventually begin to rise slightly in the very late stages of training (steps $2^{16} = 65,536$ and beyond), reaching approximately 0.04–0.05 by the final checkpoints, suggesting that the network may be incorporating higher-than-quadratic structure — but the paper notes this coincides with overfitting: "the model starts to overfit" (Section 4.3).
An unexpected finding is that in the first few hundred training steps, the FVU for both linear and quadratic approximants actually decreases — the network is getting simpler, not more complex. The paper acknowledges this is counterintuitive but offers an interpretation: "while randomly initialized neural networks tend to be simple (Teney et al., 2024), it may make sense for SGD to eliminate noise from the network early in training before making it more complex" (Section 4.2). However, the paper also notes that this early decrease disappears when using KL divergence instead of FVU (Figure 3), suggesting it may be a metric artifact — the KL divergence shows a monotonic increase in complexity from the beginning of training.
The KL divergence plot (Figure 3) mirrors the overall trend of Figure 2 but with two differences: (1) the early decrease is absent, with KL rising monotonically from the earliest checkpoints, and (2) the quadratic KL divergence also rises steadily throughout training (from ~0.05 nats at step 16 to ~0.25 nats at step $2^{14}$), rather than remaining flat. This suggests the quadratic approximant is losing fidelity gradually, not abruptly — the phase transition is primarily visible in the linear FVU, not the quadratic.
The paper quantifies the overall quality of the quadratic approximation at the end of training: "our quadratic approximants explain over 95% of the variance in network outputs" (Section 1). From Figure 2, the final quadratic FVU is approximately 0.04–0.05, corresponding to $R^2 \approx 0.95$–0.96. The linear FVU at the same point is approximately 0.15–0.18, corresponding to $R^2 \approx 0.82$–0.85.
Qualitative Feature Visualization
Figure 1 complements the quantitative FVU analysis with a visual investigation of what the polynomial features actually look like. For a specific class (digit '3'), the paper displays the top eigenvector of the quadratic approximant's interaction matrix and the coefficients of the linear approximant, both visualized as 28×28 images, at 10 training checkpoints from step $2^4$ to step $2^{18}$. This visualization follows the eigendecomposition approach described in Section 2: for each target class $k$, the quadratic approximant's coefficients form a matrix $Q = \sum_i \lambda_i v_i v_i^T$, and the eigenvectors $v_i$ can be reshaped into images for direct inspection.
The paper reports that "linear and quadratic features are most intuitively interpretable between steps $2^{12}$ and $2^{14}$" (Section 4.3). At step $2^{12} = 4,096$, the FVU is 0.04 for both linear and quadratic approximants (from Figure 1 annotations), and the eigendecomposition images show "discernible '3' qualities" for both — the linear feature appears as a rough outline of a '3' shape, while the quadratic eigenvector shows a similar but sharper pattern. This is the point where the FVU is minimal for the linear approximant (Figure 2 shows the linear FVU reaching its minimum around this step before rising), and the visualization aligns: the linear feature is most interpretable when the linear FVU is lowest.
At step $2^{14} = 16,384$, the FVU annotations in Figure 1 show 0.10 for the linear approximant and 0.20 for the quadratic. At this point, the linear '3' structure is described as "most intuitively discernible... when FVU is minimal, before beginning to overfit" — but this characterization is somewhat inconsistent with the FVU annotations, since FVU at step $2^{14}$ (0.10) is higher than at step $2^{12}$ (0.04). The likely resolution is that visual interpretability does not perfectly track FVU: the features may become visually sharper even as their overall explanatory power decreases, because the network's function is incorporating higher-order structure that makes the linear approximation less accurate but the linear features themselves more distinctive (they now capture the dominant linear component of a fundamentally nonlinear function).
At the latest training stages (steps $2^{16}$ and $2^{18}$), the images show "visual artifacts" — the quadratic eigenvector develops high-frequency noise patterns, and the linear coefficient image becomes less clearly digit-like. The FVU annotations confirm this degradation: FVU reaches 0.13 at step $2^{16}$ and 0.20 at step $2^{18}$ for the linear approximant. The paper interprets this as overfitting: "the model starts to overfit" (Section 4.3), meaning the network is learning spurious correlations that manifest as interpretability-destroying noise in the polynomial features.
A notable temporal pattern: "the quadratic feature crystallizes later than linear" (Figure 1 caption). The linear feature shows a discernible '3' shape earlier in training (around step $2^{10}$–$2^{12}$) than the quadratic feature, which becomes clearly interpretable only later (around step $2^{14}$). This temporal lag is consistent with the distributional simplicity bias hypothesis: the network learns linear structure before quadratic structure.
Adversarial Transfer: SVD-Based Attack
The paper uses adversarial attacks as a causal validation tool: if features extracted from the polynomial approximant genuinely capture how the original network processes inputs, then removing those features should degrade the original network's performance. The SVD attack (Section 4.4) provides the strongest evidence for this claim.
The linear approximant's coefficients $\beta \in \mathbb{R}^{10 \times 784}$ are decomposed via SVD: $\beta = U \Sigma V^T = \sum_{i=1}^{10} \sigma_i u_i v_i^T$, where $U \in \mathbb{R}^{10 \times 10}$, $\Sigma$ is 10×10 diagonal, and $V \in \mathbb{R}^{10 \times 784}$. The top-$k$ singular vectors are neutralized by projecting all inputs onto their orthogonal complement: $P_k = I - \sum_{i=1}^k u_i u_i^T$. This projection removes the $k$ most important input-space directions according to the linear approximant's SVD, while leaving 784 − $k$ dimensions intact.
Figure 5 plots test accuracy versus the number of SVD components ablated (from 0 to 10). The headline result: "After ablating just four SVD components, its accuracy is already under 50%" (Section 4.4). From the plot, accuracy with zero components ablated (the unmodified network) is approximately 95–97% (typical for a simple MLP on MNIST, though the exact value is not quoted in the text). At $k=1$, accuracy drops to approximately 75–80%. At $k=4$, accuracy falls below 50%. At $k=10$ (all SVD components ablated), accuracy reaches roughly 10% — random chance for 10-class classification. The paper states: "ablating the entire rowspace (only 10 out of 784 total dimensions) [is] sufficient to make the MLP perform no better than random chance" (Figure 5 caption).
The most striking aspect of Figure 5 is the lockstep alignment between the three curves: the original ReLU MLP's accuracy, the linear approximant's accuracy, and the quadratic approximant's accuracy all drop together as SVD components are ablated. The paper describes this as: "the original network's accuracy drops in perfect lockstep with that of its linear and quadratic approximants" (Figure 5 caption). This alignment is evidence that the polynomial approximants have captured the input-space directions that the original model causally depends on, not merely directions that are correlated with its outputs. If the approximants were capturing spurious correlations, ablating their principal components might not affect the original model at all — or might affect it in an uncorrelated way. The near-perfect tracking across all three curves suggests the approximants and the original model share a common sensitivity subspace.
Figure 4 provides qualitative evidence for what the SVD attack does to individual images. The top row shows adversarial examples of the digit '3' as increasing numbers of SVD components are ablated (1 through 10, left to right). The bottom row shows random examples of each digit with all 10 SVD components ablated. The key qualitative claim: "Examples bear high resemblance to the original images, despite being unclassifiable by the MLP" (Figure 4 caption). This distinguishes the SVD attack from naive input corruption — the ablated images still look like digits to a human, meaning the attack is exploiting specific features the model relies on rather than simply destroying the input. The fact that only 10 out of 784 dimensions (~1.3%) are removed, and the images remain recognizable, indicates that the model has concentrated its decision-making into a very low-dimensional linear subspace — precisely the subspace captured by the linear approximant's top singular vectors.
The paper does not report FVU or $R^2$ values for the linear approximant used in the attack, but Figure 2 indicates the linear FVU at the end of training is ~0.15–0.18. This means the linear approximant explains 82–85% of the output variance, yet removing its 10 principal components completely destroys classification. This apparent discrepancy — 15–18% unexplained variance, but 100% accuracy loss from removing a 10-dimensional subspace — is not actually contradictory: classification accuracy depends on the relative outputs for different classes, not the absolute output values. A linear approximant could have 85% $R^2$ in output space while still correctly capturing the relative class ordering, and removing the 10 principal components could destroy that ordering even if they only explain a fraction of the raw output variance. The paper does not explore this distinction explicitly.
Ablation Studies and Robustness Checks
-
Fidelity metric: FVU vs. KL divergence (Figures 2 vs. 3). The paper reports both FVU and KL divergence as measures of discrepancy between the original network and its polynomial approximants. The key robustness finding is that the early-training decrease in FVU (visible in Figure 2 for the first few hundred steps) disappears when using KL divergence (Figure 3), which instead shows monotonic increase from the start. This suggests the early FVU decrease is an artifact of how FVU weights errors — FVU normalizes by total variance, which can increase during early training, making the absolute MSE appear to decrease relative to a growing denominator. KL divergence, by contrast, directly measures distributional mismatch without this normalization. The fact that the phase transition (sharp increase in linear FVU between 500–1000 steps) appears in both metrics strengthens the claim that it is a genuine phenomenon rather than a metric artifact. The paper does not explicitly compare the two metrics or discuss why they diverge; this analysis is left implicit.
-
Exact vs. finetuned quadratic approximants (Appendix B). The quadratic approximants in Figure 2 are not computed exactly — they are initialized with the exact solution for
$\mathcal{N}(0, I)$inputs (where$\text{Cov}[z]$is diagonal and trivially invertible), then finetuned via SGD on batches from the actual Gaussian mixture distribution matching MNIST class statistics. The paper does not provide an ablation comparing this hybrid approach to a fully exact solution (which is computationally intractable) or to a purely SGD-based fit without the analytic initialization. However, the paper argues that since the objective is convex, "the resulting coefficients should be excellent approximations of the true least-squares values" (Appendix B). The fact that the quadratic FVU is very low (~0.01–0.05) throughout training suggests the hybrid approach works well, but without an exact baseline, it's impossible to quantify the initialization error. The noise visible in the quadratic FVU curve in Figure 2 is attributed to the SGD finetuning: "This accounts for the noise visible in the plots for quadratic FVU" (Section 4.1, footnote). -
Single training run for dynamics analysis. The paper traces FVU across training time for a single training run of the MLP. There is no report of multiple random seeds to assess whether the observed phase transition (sharp linear FVU increase at 500–1000 steps) is robust across different initializations. This is a significant omission because training dynamics can vary substantially with random seed, especially for small models on MNIST. The paper does not claim statistical significance for the phase transition — it is presented as a qualitative observation consistent with the simplicity bias hypothesis. A multi-seed analysis would strengthen the claim by showing that the transition occurs reliably at similar points in training, and by quantifying seed-to-seed variance in the FVU trajectories.
-
Mixture components: class-conditional vs. single Gaussian. The authors state that the Gaussian mixture distribution uses "means and covariances matching those of the MNIST classes" (Section 4.1). This means 10 mixture components, one per digit class, each with its own mean vector (784-dimensional) and covariance matrix (784×784). The paper does not include an ablation comparing mixture-model approximants to single-Gaussian approximants (a single Gaussian fit to all MNIST images regardless of class). Such an ablation would reveal whether the mixture modeling is important for achieving high
$R^2$, or whether a single Gaussian would suffice. The paper's theoretical framework supports both, but the practical importance of the mixture assumption is untested. This matters because computing class-conditional covariances requires labeled data, which limits applicability to unlabeled datasets. -
Choice of activation function: ReLU only. All experiments use ReLU MLPs. The paper derives closed-form integrals for both ReLU (Appendix A.2) and GELU (Appendix A.1) but does not report empirical results for GELU networks. This is an important ablation deficit: GELU is smoother than ReLU and might produce different polynomial approximation characteristics — specifically, GELU networks might be better approximated by low-degree polynomials because GELU is itself smooth (infinitely differentiable) while ReLU is piecewise linear. The FVU trajectories for GELU networks could differ qualitatively from the ReLU results in Figure 2. The paper does not justify the choice of ReLU over GELU for the experiments.
-
Hidden layer width. The paper does not specify the width of the single-hidden-layer MLP used in experiments. This is a critical missing detail because the hidden layer width determines the network's capacity and Expressivity. A very narrow network might be trivially approximable by quadratics simply because it has limited representational capacity, while a wider network might show a larger gap between linear and quadratic FVU. The phase transition might occur at different times or with different magnitudes for different widths. Without this specification, the quantitative FVU values are not reproducible, and claims about "over 95% variance explained" are tied to an unknown model scale.
-
Choice of polynomial degree: only linear and quadratic compared. The paper's main quantitative results (Figures 2, 3) compare only linear and quadratic approximants. Higher-degree (cubic, quartic) approximants are theoretically possible via the same framework (the quadratic feature map strategy extends to arbitrary degree) but are never computed. This means the paper cannot distinguish between "the network is quadratic" and "the network is cubic but quadratics already explain 95% of variance." The slight rise in quadratic FVU at the end of training (Figure 2, steps
$2^{16}$and beyond) hints at higher-order structure, but without cubic approximants, this remains speculation. The computational intractability of higher-degree approximants ($O(\binom{d+n}{n}^3)$for degree$n$) is a practical barrier, but even a small-scale experiment (e.g.,$n=3$on a reduced-dimensional input via PCA) would provide evidence about whether the quadratic approximation is truly sufficient. -
Approximant accuracy vs. network accuracy (missing ablation). The paper does not report the original MLP's classification accuracy on MNIST at each training checkpoint alongside the FVU curves. This makes it difficult to assess the relationship between approximation quality and task performance. For example, does the phase transition in linear FVU (500–1000 steps) coincide with a jump in test accuracy? Does the network continue to improve in accuracy after the quadratic FVU starts rising (steps
$2^{16}$and beyond)? The paper mentions in Figure 1's caption that the linear '3' structure is "most intuitively discernible at step [214], when FVU is minimal, before beginning to overfit" and that "this can be interpreted as the MLP learning and relying on statistics of higher complexity than linear, especially if its accuracy continues to improve" — but this conditional statement ("especially if") indicates that accuracy data was not collected or reported. This is a significant omission for a paper that aims to connect approximation quality to learning dynamics. -
Universal adversarial perturbation construction (SVD attack). The SVD attack uses the linear approximant's SVD to construct a universal adversarial projection — the same projection matrix
$P_k$is applied to all test images regardless of class. The paper does not compare this to per-class projections (using class-specific linear coefficients) or to attacks derived directly from the original network (e.g., FGSM on the ReLU MLP). A comparison to standard adversarial attack methods would contextualize how effective the polynomial-derived attack is — is ablate-10-SVD more or less effective than a 10-dimensional FGSM perturbation? This comparison is absent, making it difficult to assess whether the attack's success is evidence for the approximant's quality or merely evidence that MNIST classifiers are easily attacked.
Critical Assessment
The experiments demonstrate a specific empirical phenomenon — that the linear FVU of a ReLU MLP trained on MNIST exhibits a phase transition during training while the quadratic FVU remains low — and show that features from the linear approximant transfer causally to the original network via adversarial attacks. The question is whether these experiments genuinely support the paper's broader claims about polynomial approximants as general-purpose interpretability tools and about the distributional simplicity bias.
Claim: "Quadratic approximants explain over 95% of the variance in network outputs." This claim is supported by Figure 2 for the specific model tested (a single-hidden-layer ReLU MLP of unspecified width on MNIST). The quadratic FVU at the end of training is approximately 0.04–0.05, corresponding to $R^2 \approx 0.95$–0.96. However, the support is narrow: this is one model architecture on one dataset. The paper does not demonstrate that similarly high $R^2$ values hold for different architectures (deeper networks, GELU activations, convolutional networks), different datasets (CIFAR-10, ImageNet), or different tasks (regression, language modeling). The claim that quadratics explain "over 95%" should be understood as "over 95% on this specific MNIST MLP" rather than a general property of neural networks. Moreover, the quadratic FVU is computed via SGD finetuning from a $\mathcal{N}(0,I)$ initialization, not via exact least-squares — there may be a gap between the achieved FVU and the true optimum, though convexity of the objective provides theoretical assurance that SGD converges to the optimum given sufficient iterations.
Claim: "Neural networks start out simple and get progressively more complex" (distributional simplicity bias). The evidence in Figure 2 supports a qualified version of this claim: linear FVU is low early in training (network is approximately linear), then rises (network becomes less linear), while quadratic FVU stays low throughout (network remains approximately quadratic). This is evidence for a linear-to-quadratic progression, not a general progression through arbitrary polynomial degrees. The paper provides no evidence for cubic, quartic, or higher-order phases because higher-degree approximants are not computed. The observed slight rise in quadratic FVU late in training (steps $2^{16}$ and beyond) hints at higher-order structure, but without cubic approximants, this remains speculation. The claim of progressive complexity increase is therefore demonstrated only for the first two degrees. Additionally, the phase transition is reported for a single training run — without multiple seeds, the reliability and timing of the transition are unknown. A single training run could exhibit a sharp FVU increase at 500–1000 steps due to a particular weight update or learning rate schedule artifact, not a general phenomenon.
Claim: "SVD-based adversarial examples show that approximants capture out-of-distribution behavior." The lockstep accuracy drop in Figure 5 provides strong causal evidence that the linear approximant captures input directions the original model depends on. However, the claim about "out-of-distribution behavior" is only partially tested. The adversarial examples in Figure 4 are projections of in-distribution MNIST test images — they are not genuinely out-of-distribution in the sense of being drawn from a different distribution (e.g., CIFAR-10 images, random noise, or adversarial distributions designed to maximally diverge from MNIST). The paper argues that the approximant captures out-of-distribution behavior because it was fit under a maximum-entropy Gaussian mixture distribution rather than the empirical MNIST distribution — but the adversarial validation uses MNIST test images, not samples from the Gaussian mixture. A stronger test would evaluate the SVD attack on images drawn directly from the Gaussian mixture distribution used to fit the approximant, to verify that the attack transfers across the distribution shift from empirical MNIST to the mixture model.
Missing experiment: verification that the Gaussian mixture assumption is good for MNIST. The paper asserts that "MNIST is an ideal dataset for this task since it is known to be very well-modeled as a Gaussian mixture distribution" and cites Belrose et al. (2023, Figure 2) for the claim that "samples from a Gaussian mixture fit to MNIST are difficult for a human to distinguish from real samples." This is the linchpin of the experimental validity: if the Gaussian mixture is a poor model of MNIST, then polynomial approximants fit under that assumption may not transfer to real MNIST images. The paper does not provide its own validation of the Gaussian mixture quality — no FID scores, no log-likelihood comparisons, and no visual samples. It relies entirely on a citation. This is a significant gap, especially given that the method's central motivation is avoiding reliance on empirical data distributions. If the Gaussian mixture is a poor model of MNIST, the high quadratic $R^2$ might reflect the mixture's simplicity rather than the network's actual functional form. The fact that the SVD attack transfers to real MNIST test images (Figure 5) provides indirect evidence that the mixture is adequate, but a direct validation would be stronger.
Missing experiment: approximant quality vs. network accuracy. The paper never reports the MLP's classification accuracy on MNIST alongside the FVU curves. This is the most consequential omission because it prevents answering a central question: does the network's shift from linear to quadratic computation correspond to improved task performance, or is it an epiphenomenon? If accuracy plateaus before the linear FVU rises, the phase transition might represent overfitting rather than useful learning. If accuracy continues to improve as quadratic FVU begins rising late in training, that would suggest higher-order structure is task-relevant. The paper gestures at this in the Figure 1 caption ("if its accuracy continues to improve") but never resolves the conditional. Without accuracy data, the claim that networks "learn statistics of increasing complexity" is only a claim about the network's functional form under a specific input distribution, not a claim about task-relevant learning.
Missing baseline: comparison to gradient-based polynomial fitting. The paper emphasizes the analytic nature of its polynomial approximants as a key advantage over SGD fitting. Yet the quadratic approximants in the experiments are partially SGD-fit (via finetuning from the $\mathcal{N}(0,I)$ initialization). A natural ablation — fitting quadratics purely via SGD on Gaussian mixture samples without the analytic initialization — would quantify the benefit of the analytic warm start. If purely SGD-fit quadratics achieve similar FVU, then the analytic derivation provides theoretical insight but limited practical advantage for the quadratic case. If they are significantly worse (due to optimization difficulty), the analytic initialization is crucial. This ablation is absent.
What the experiments do convincingly establish. Despite these limitations, the experiments do convincingly establish: (1) for this specific ReLU MLP on MNIST, quadratic approximants under a Gaussian mixture input model achieve very high fidelity ($R^2 > 0.95$ at end of training), (2) the linear approximation quality degrades during a specific training phase while quadratic approximation quality remains high, consistent with the network acquiring quadratic structure, and (3) the linear approximant's principal directions are causally relevant to the original network's classification decisions, as demonstrated by the SVD adversarial attack. Whether these findings generalize to other architectures, datasets, or activation functions is not tested, and the paper's claims should be interpreted as applying to the demonstrated setting unless future work shows broader applicability.
Scalability concerns. The paper is candid about computational limits: exact quadratic approximants are $O(d^6)$ and become intractable for $d > \mathcal{O}(100)$. The MNIST experiments at $d=784$ already require the hybrid SGD approach. Any extension to larger input dimensions (e.g., CIFAR-10 at 3072 dimensions, or ImageNet at 150,528) would push even the hybrid approach to its limits — the quadratic feature vector dimension scales as $d(d+1)/2$, which is ~1.2 million for $d=1536$ (a typical transformer hidden size). The experiments do not address this scalability challenge at all, and there is no clear path to applying the method to modern large-scale networks short of dimensional reduction (e.g., applying the polynomial approximation in a PCA-reduced subspace, which the paper does not explore).
Statistical rigor. The paper reports a single training run, no confidence intervals on FVU estimates, and no seed-based replication. The phase transition is described qualitatively ("a clear phase transition") without quantitative criteria for what constitutes a transition (e.g., a threshold rate of FVU increase, or a statistical test for a changepoint). The KL divergence plot (Figure 3) provides a robustness check against the choice of discrepancy metric, but the lack of statistical quantification makes it difficult to assess whether the observed patterns would replicate. For a paper whose central empirical claim is about training dynamics, this is a notable gap.
6. Limitations and Trade-offs
The Difficulty Estimation Overhead Is Not Accounted For in the Headline Efficiency Numbers
The assumption or constraint. The paper's compute-optimal policy for test-time compute requires estimating each prompt's difficulty to decide how to allocate the inference budget. The method described in Section 3.2 generates 2048 samples per question from the base model and computes the PRM's average final-answer score across these samples as a difficulty proxy. The paper acknowledges this cost directly:
"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"
Generating 2048 samples per question to estimate difficulty represents an overhead that, in many cases, exceeds the largest test-time compute budgets studied (256–512 generations). The authors candidly call this the "exploration-exploitation tradeoff" (Section 3.2) but do not incorporate it into any of the reported efficiency gains.
The consequence. The headline result — that compute-optimal scaling improves efficiency by more than 4× over best-of-N (Section 1, Figures 4 and 8) — is computed after difficulty is known, without amortizing the cost of learning it. In a deployment setting, the true cost is difficulty estimation cost plus strategy execution cost, and the former could easily dominate the latter. The 4× figure should therefore be understood as an upper bound on achievable efficiency rather than a realized deployment gain. If difficulty estimation is as expensive as the inference budget itself (or more), the actual net efficiency could approach or fall below 1× — meaning the adaptive approach provides no benefit over uniform allocation once total cost is accounted for. The authors acknowledge this as "a key avenue for future work" (Section 3.2), but the gap means the paper's primary practical claim is unvalidated for realistic deployment scenarios.
What evidence exists in the paper. Section 3.2 explicitly states the 2048-sample difficulty estimation procedure. The difficulty estimation cost appears nowhere in the budget axes of Figures 3, 4, 6, 7, or 8 — the generation budget N refers exclusively to strategy execution, not difficulty assessment. The paper reports that predicted (PRM-based) difficulty bins largely overlap with oracle bins (Figures 4 and 8, curves approximately tracking each other), confirming that difficulty can be estimated without ground-truth labels, but the computation for that estimation is not counted against efficiency.
Mitigation status. Not addressed in the current work. The paper suggests future work on "pretraining or finetuning models to directly predict difficulty of a question" (Section 8) as a way to amortize this cost, but no such model is developed or evaluated. An alternative — adaptive difficulty estimation that starts with a small number of samples and adjusts the remaining budget dynamically — is also not explored. Until some form of cheap difficulty estimation is demonstrated and included in the total cost accounting, the reported efficiency gains remain an idealized upper bound.
Single Benchmark, Single Model Family — No Evidence of Generalization
The assumption or constraint. All experiments use a single model family (PaLM 2-S*) on a single benchmark (MATH, 500 test questions). The authors state:
"We believe this model is representative of the capabilities of many contemporary LLMs"
but provide no evidence of this representativeness through experiments on other models or datasets. The MATH benchmark consists exclusively of competition-level mathematics problems requiring symbolic reasoning and step-by-step derivation. It is unknown whether the paper's central findings — that beam search outperforms best-of-N on medium problems but over-optimizes on easy problems (Figure 3, right), that sequential revisions dominate on easy problems while parallel sampling helps on hard ones (Figure 7, right), and that compute-optimal allocation yields ~4× efficiency gains — generalize to other reasoning domains (code generation, logical deduction, scientific QA), to tasks requiring factual recall rather than inference, or to other model families with different calibration properties and error patterns.
The consequence. A practitioner considering whether to adopt compute-optimal test-time scaling for a different model (e.g., Llama, Claude, Gemini) on a different task (e.g., code generation, document QA, planning) cannot determine from this paper alone whether the method will work. Several aspects of the findings could be model-specific or benchmark-specific. The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s output distribution — a model with different calibration or different typical error patterns might exhibit different difficulty-dependent scaling curves, potentially with different optimal strategy assignments per bin. The revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning capabilities, which vary across model families. The difficulty quintile boundaries (defined by pass@1 rates) and the resulting optimal strategy assignments could shift entirely for a different model on the same benchmark. The test set of 500 questions, divided into five quintiles of ~100 each, then split by two-fold cross-validation, means the compute-optimal policy is selected based on approximately 50 questions per fold per bin — a sample size small enough that the selected strategies may be specific to this particular 500-question split rather than robust properties of the difficulty-strategy relationship. The paper does not report confidence intervals on the compute-optimal scaling curves.
What evidence exists in the paper. All figures (3–9) are based on PaLM 2-S* and MATH. The paper provides no experiments on other benchmarks, other model families, or even other sizes within the PaLM 2 family. The cross-validation protocol (two-fold, within difficulty bins on the test set) is described in Section 3.2, and the sample size per fold per bin (~50 questions) can be inferred from the 500-question total divided into 5 bins, then halved.
Mitigation status. Not addressed. The paper makes no claim of generality beyond the studied setting. The limitation is acknowledged implicitly by the narrow experimental scope, but no ablation across models or tasks is provided, and no argument is made for why PaLM 2-S* / MATH should be representative. This is a standard "single benchmark" limitation common across empirical ML papers, but it is particularly significant here because the central practical claim (4× efficiency gains via adaptive allocation) is being advanced as a general principle — the paper draws explicit parallels to Chinchilla scaling laws, which were validated across multiple model families and datasets.
The Larger Model Baseline Is Weak — FLOPs-Matched Comparison Overstates Test-Time Compute Benefits
The assumption or constraint. The FLOPs-matched comparison in Section 7 tests whether a smaller model (PaLM 2-S*) with compute-optimal test-time strategies can outperform a model with approximately 14× more parameters. However, the larger model baseline is systematically weakened in two ways. First, the larger model is scaled in parameters only while holding training data fixed — the authors acknowledge this departs from compute-optimal pretraining:
"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."
A Chinchilla-optimal model trained with 14× more total FLOPs (scaling both parameters and data) would likely outperform a parameter-only-scaled model, making the pretraining baseline weaker than it could be. Second, the 14× larger model uses only greedy decoding — no majority voting, no best-of-N sampling, no test-time compute of any kind. The smaller model is permitted to use sophisticated test-time strategies (beam search, sequential revisions, adaptive allocation), while the larger model receives none of these benefits.
The consequence. The reported advantages of test-time compute over pretraining — for example, +27.8% relative improvement on easy-to-medium questions at R ≪ 1 with revisions, or +19.1% with PRM search (Figure 1 bar charts, Section 7) — may overstate the true tradeoff. Against a compute-optimally trained larger model, or against a larger model given even a modest test-time compute budget (e.g., best-of-8 or majority voting over 8 samples), these margins would shrink or potentially reverse. This matters for the paper's central practical message: that organizations should consider allocating compute to test-time strategies rather than pretraining. If the baseline is unrealistically weak, the policy recommendation may be biased toward test-time compute in regimes where pretraining is actually preferable. The dependence on R = D_inference / D_pretrain (Section 7) is sound in principle — when R ≪ 1, pretraining savings can fund substantial test-time compute — but the quantitative magnitude of the advantage is contingent on the baseline strength. The paper's finding that test-time compute fails on hard problems (bin 5, near-zero improvement regardless of budget) and at high R (Section 7, Figure 9) is likely robust to baseline changes, but the positive results on easy-to-medium problems may be inflated.
What evidence exists in the paper. The experimental setup is described in Section 7, including the FLOP accounting formulas and the three values of R tested (0.16, 0.79, 22). The results are presented in Figure 9 and summarized in the bar charts in Figure 1. The paper explicitly notes the parameter-only scaling assumption but does not test a Chinchilla-optimal baseline. The fact that the larger model uses greedy decoding is mentioned in the results text.
Mitigation status. Partially acknowledged. The authors explicitly note the departure from compute-optimal pretraining and flag it as a direction for future work. They do not, however, discuss the absence of test-time compute for the larger model as a limitation, nor do they include an ablation where the larger model receives a small inference budget. The comparison is framed as "test-time compute vs. pretraining" but functionally it is "small model + test-time compute vs. large model + greedy decoding" — an asymmetric comparison that conflates two choices (model size and inference strategy) that could be varied independently.
Verifier Over-Optimization Is an Unresolved Ceiling, Not a Solved Problem
The assumption or constraint. The paper identifies verifier over-optimization as a central limiting factor: beam search degrades performance on easy problems at high budgets (Figure 3, right), lookahead search — the most powerful optimizer tested — paradoxically performs worst overall at equivalent budgets (Figure 3, left), and qualitative examples show search producing degenerate outputs (repetitive low-information steps, overly short solutions) that score highly under the PRM (Appendix M, referenced in Section 5.3). The compute-optimal policy mitigates this by routing easy problems away from aggressive search toward best-of-N, but it does not solve the underlying problem: on medium-difficulty problems where beam search is deployed, over-optimization still limits the scaling ceiling, visible in the flattening and eventual decline of beam search accuracy curves at high budgets (Figure 3, right). The paper's approach is fundamentally bounded by verifier quality — the compute-optimal policy is, in part, a strategy for staying below the budget level at which over-optimization begins to dominate for each difficulty tier.
The consequence. This means that further scaling of test-time compute (beyond the budgets studied, or with larger base models) will not yield proportional improvements unless verifier robustness improves commensurately. The paper demonstrates that lookahead search — which intuitively should be "better" because it uses more computation to score each candidate step — actually performs worse than simpler beam search because it intensifies over-optimization. This is a counterintuitive and important negative result, but its practical implication is that simply increasing the inference budget, or deploying more sophisticated search algorithms, is not a reliable path to better performance. A practitioner who takes the paper's compute-optimal framework and applies it with a different or improved PRM will get different results — the difficulty-strategy mapping depends on the PRM's calibration and over-optimization profile, which are training-procedure-specific. The paper provides no guidance on how to predict or control verifier over-optimization behavior for a new PRM.
What evidence exists in the paper. Figure 3 (right) shows beam search accuracy decreasing with budget for bin 1 (easiest) — the clearest signal of over-optimization. Figure 3 (left) shows lookahead search underperforming beam search and best-of-N at most budgets. Section 5.3 discusses the over-optimization phenomenon explicitly. Appendix M shows qualitative examples of degenerate beam search outputs. The paper does not, however, measure over-optimization quantitatively (e.g., as the gap between PRM-predicted score and actual correctness as a function of beam search depth), nor does it experiment with verifier robustness interventions.
Mitigation status. Not addressed as a problem to be solved — the compute-optimal policy is proposed as a way to route around over-optimization, not eliminate it. Section 8 suggests that improving verifier robustness should be a priority:
"improving verifier robustness is the key bottleneck for further scaling test-time compute, not improving search algorithms"
but this is framed as a direction for future work, not something attempted in the paper. The PRM training procedure (Monte Carlo rollouts with soft labels, Appendix D) uses a single recipe with no ablations comparing different training strategies for their effect on over-optimization resistance. The paper does not test ensembling multiple PRMs, adversarial training of the PRM, or KL-regularized search constraints — all natural approaches to mitigating over-optimization that would test whether the identified ceiling can be raised.
The Revision Model Has a Fundamental Failure Mode: ~38% Correct-to-Incorrect Reversion Rate
The assumption or constraint. The revision model is fine-tuned exclusively on trajectories where all in-context answers are incorrect, followed by a correct target answer (Section 6.1). The training data provides no signal for what the model should do when the current answer is already correct. The paper reports a significant practical consequence:
"approximately 38% of correct answers get converted back to incorrect ones using a naive approach"
This means that as the revision chain grows longer, correct answers produced at intermediate steps are at substantial risk of being overwritten by incorrect revisions in subsequent steps. The problem arises because the model has never seen a training example where the current answer is correct and should be preserved — its learned behavior is always to "fix" the previous answer, regardless of whether fixing is needed.
The consequence. The revision approach has an inherent instability: longer chains (which the paper shows can improve pass@1 — Figure 6, left) simultaneously increase the cumulative probability that a correct answer will be corrupted later in the chain. This creates a fundamental tension between chain depth and reliability. The paper's mitigation — using majority voting or verifier-based selection to pick the best answer from anywhere in the chain rather than always taking the final output — is an imperfect patch. Verifier-based selection requires that the verifier reliably distinguish correct answers from the ~38% of incorrect revisions that come after them, but verifiers themselves are imperfect (as the over-optimization results demonstrate). Majority voting across the chain can fail if the revision model consistently produces correct-then-incorrect patterns for certain problem types. The practical implication is that the revision model cannot be used as a standalone "iterative improver" that reliably converges to the correct answer — it requires a separate selection mechanism to extract the best output from a chain that may oscillate between correct and incorrect. This limits the applicability of revisions in settings where answer selection is ambiguous or where chains must be consumed as-is (e.g., in conversational settings where the model's final turn is what the user sees).
What evidence exists in the paper. The 38% figure is reported in Section 6.1. The mitigation strategy (within-chain selection via majority voting or verifier) is described in the same section and evaluated in Figure 6. The fact that sequential revision chains continue to improve pass@1 beyond the 4-step training horizon (Figure 6, left) confirms the model generalizes, but the reversion rate quantifies the cost of this generalization. The ReST-EM experiment (Appendix K, Figure 16) shows that an alternative training methodology for the revision model can actually degrade performance with sequential revisions, suggesting that the revision approach is sensitive to training data construction in ways that affect this failure mode.
Mitigation status. Partially addressed. Within-chain selection (Section 6.1) reduces the impact of reversions but does not eliminate the underlying problem — the model still wastes computation generating incorrect revisions after correct answers, and selection mechanisms can fail. The paper does not explore training the revision model on trajectories that include "stop revising" signals (i.e., training examples where the correct answer appears mid-sequence and the model learns to copy it forward rather than revise it), which would be the natural architectural solution. The authors do not propose this as future work, but it follows directly from the diagnosed failure mode.
The Difficulty Bins Are Static and Coarse — No Dynamic Strategy Adjustment During Inference
The assumption or constraint. The compute-optimal policy operates by assigning each question to one of five difficulty quintiles (determined by pass@1 rate over 2048 samples) and then applying a fixed, pre-computed strategy (search algorithm, beam width, sequential-to-parallel ratio) for that bin and budget level (Section 3.2). Difficulty is estimated once, before any problem-solving computation begins, and the strategy is fixed for the entire inference budget. There is no mechanism for dynamic adjustment: the system cannot start with a few samples, observe that the problem appears easier or harder than the initial difficulty estimate suggested, and reallocate the remaining budget accordingly. The discretization into only five bins also means that significant within-bin heterogeneity is ignored — a question near the boundary between bin 3 and bin 4 receives the same strategy as one squarely in the middle of bin 3, even though the optimal strategy might differ.
The consequence. The static, coarse allocation leaves potential efficiency gains on the table. A system that could dynamically adjust strategy mid-inference — for instance, beginning with a small number of parallel samples, estimating difficulty from the verifier's scores on those initial attempts, and then switching to beam search or sequential revisions depending on the observed response — could potentially outperform the static policy by adapting to per-question variation that the five-bin discretization masks. Moreover, such an adaptive scheme could amortize difficulty estimation into the problem-solving process, partially addressing the difficulty estimation overhead problem (first limitation above). The current approach treats difficulty estimation and strategy execution as separate, sequential phases, which inflates total cost and prevents the system from using early computation to inform later allocation.
What evidence exists in the paper. Section 3.2 describes the difficulty binning procedure and the five-quintile discretization. Figure 4 and Figure 8 show compute-optimal scaling curves with both oracle and predicted bins — these curves represent the performance of the static, pre-computed policy evaluated per bin. The paper does not include any experiment comparing static allocation to an adaptive scheme, nor does it ablate the number of difficulty bins to assess sensitivity to discretization granularity. The between-bin variance in Figures 3 (right) and 7 (right) suggests that the five bins capture the dominant difficulty-dependent trends, but the within-bin variance (which would determine the potential gains from finer-grained or continuous allocation) is not reported.
Mitigation status. Not addressed in the current work. The paper flags the exploration-exploitation tradeoff in Section 3.2 but does not propose or evaluate a dynamic allocation scheme. Section 8 does not mention adaptive or online strategy selection as a future direction, focusing instead on improved difficulty estimation and combining search with revisions. This is a missed opportunity, since the framework naturally supports dynamic reallocation — the same PRM scores used for answer selection could inform difficulty assessment on-the-fly — and the computational machinery for switching strategies mid-inference (changing beam width, switching from parallel to sequential) is available.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper introduces a conceptual inversion into the interpretability toolkit: rather than training auxiliary models (sparse autoencoders, probes) to explain what a neural network does on a particular dataset, derive analytically what the network computes under the simplest possible input distribution consistent with low-order statistics. This is not a paradigm shift — it does not replace dictionary learning or mechanistic interpretability — but it is a reframing that changes the default question from "what features does this model compute on this data?" to "what features of this model are revealed when we strip away the data distribution?" The distinction matters because it addresses a confound that the SAE literature has not resolved: whether extracted features reflect model computation or data statistics. The polynomial approximation approach provides a principled way to isolate the former.
The practical consequence is that existing pretrained models become amenable to eigendecomposition-based interpretability without architectural modification. Prior work by Pearce et al. (2024) showed that bilinear layers have interpretable interaction matrices, but required building the model to be quadratic from scratch. This paper shows that any pretrained ReLU or GELU MLP — and by extension, any GLU variant — can be converted post-hoc into a polynomial approximant whose eigendecomposition yields interpretable features, subject to the Gaussian input assumption. This significantly expands the set of models that can be analyzed through spectral methods, from custom bilinear architectures to off-the-shelf pretrained networks.
The paper also introduces polynomial degree as a dynamic measure of functional complexity during training. Prior work on simplicity bias (Nakkiran et al., 2019; Refinetti et al., 2023; Belrose et al., 2024) relied on indirect proxies — spectral analysis of weight matrices, teacher-student setups, architecturally separated linear/nonlinear branches. The polynomial approximation framework provides a post-hoc diagnostic: trace the R² of best linear and quadratic approximants across training checkpoints, and observe whether networks transition from linear to quadratic computation. The phase transition documented in Figure 2 (linear FVU rising sharply at 500–1000 steps while quadratic FVU stays flat) is the first direct evidence from a standard training run that networks genuinely learn statistics of increasing polynomial degree — not just that they can under controlled conditions, but that they do during ordinary SGD on real data.
The identification of causal transfer as an interpretability validation criterion is a methodological contribution that extends beyond polynomial approximants. The SVD adversarial attack (Section 4.4) demonstrates that features extracted from the approximant are not merely correlated with the original model's outputs but are causally necessary for classification: removing 10 linear principal components (1.3% of input dimensions) reduces the ReLU MLP to random chance accuracy. The lockstep accuracy drop between approximant and original model (Figure 5) establishes a validation template — if an interpretability method claims to have found "the features the model uses," the strongest test is whether removing those features breaks the model. This shifts the interpretability evaluation standard from descriptive fidelity (how well does the explanation match the model?) to causal fidelity (does the explanation capture what the model depends on?).
The work also elevates the status of maximum-entropy input modeling as a principled alternative to empirical data distributions for interpretability. Most interpretability methods are implicitly or explicitly tied to the data distribution they were trained or evaluated on. By arguing that the Gaussian is the unique maximum-entropy distribution given mean and covariance, the paper provides a philosophically grounded default — not "the right distribution" for all settings, but "the distribution that makes the fewest assumptions beyond the network's weights." This reframes interpretability from an empirical exercise (fit a probe on data) to a theoretical one (derive structure from weights plus minimal distributional assumptions). Future interpretability methods may adopt this framing even if they use different mathematical machinery.
Finally, the paper reconciles an apparent tension between the simplicity of bilinear architectures (Pearce et al., 2024) and the complexity of standard nonlinear MLPs. The finding that quadratics explain >95% of output variance for a ReLU MLP on MNIST (Figure 2) suggests that, at least on this task, ReLU networks are functionally not much more complex than explicitly quadratic networks — the nonlinearity is present in the architecture but may be largely approximable by low-degree polynomials in practice, at least under the Gaussian input assumption. This does not mean "ReLU MLPs are quadratic" — the slight rise in quadratic FVU at late training stages (Figure 2, steps 2^16 and beyond) and the qualitative differences in feature visualization (Figure 1) indicate genuine higher-order structure. But it does mean that much of what these networks compute is captured by second-order statistics, which has implications for how we think about network expressivity: the capacity may be there, but what is actually learned may be simpler.
Follow-Up Research This Work Enables
Systematic replication across architectures, datasets, and activation functions. The paper's central empirical findings — the >95% quadratic R² at end of training, the linear-to-quadratic phase transition, the causal transfer of SVD attacks — are demonstrated for a single model family (single-hidden-layer ReLU MLP of unspecified width) on a single dataset (MNIST). A strong follow-up would replicate the FVU-across-training experiment across: (a) activation functions (GELU, Swish, tanh), testing whether smoother activations produce even higher quadratic R² or different phase transition timing; (b) depth (2–4 hidden layers), testing whether the quadratic R² degrades with depth as compositional nonlinearities accumulate; (c) architectures (CNNs, small transformers), testing whether convolutional inductive biases or attention mechanisms produce different polynomial approximation characteristics; and (d) datasets (CIFAR-10, Fashion-MNIST, synthetic Gaussian mixture data with controlled higher-order structure), testing whether the >95% R² is specific to MNIST's simplicity. A systematic replication would establish boundary conditions — at what model scale, architecture complexity, and data complexity does the quadratic approximation break down? The paper's theoretical framework supports all these extensions (the Master Theorem handles any continuous ϕ, the mixture extension handles class-conditional structure), but the experiments are minimal. This is the most important follow-up because the paper's practical claims rest on empirical generalizability that has not been demonstrated.
Cubic and higher-degree approximants on dimension-reduced inputs. The paper acknowledges that exact higher-degree approximants are computationally intractable (the feature dimension binom(d+n, n) for degree n explodes). A concrete follow-up would apply dimensionality reduction — PCA to 50–100 principal components of the input, or a random projection — and compute cubic (n=3) approximants in the reduced space. For MNIST, 50 PCA components capture the vast majority of pixel variance, and the cubic feature dimension would be binom(53, 3) ≈ 23,426, which is large but computationally feasible for a one-time analytic computation. This would test whether the slight rise in quadratic FVU late in training (Figure 2, steps 2^16 and beyond) is explained by cubic structure: if cubic R² is near 1.0 while quadratic R² drops to 0.95, the network is learning third-order interactions. If cubic R² is no better than quadratic, the degradation is due to quartic or higher terms, or to structure that low-degree polynomials cannot capture regardless of degree (e.g., the ReLU nonlinearity creating genuine non-polynomial behavior). This experiment directly tests the distributional simplicity bias claim that networks learn statistics of increasing but bounded polynomial degree.
Training the revision model with "stop revising" signals to address the 38% reversion rate. The paper's revision model (from the prior sections context) suffers from a correct-to-incorrect reversion rate of ~38% because training data only contains incorrect-to-correct trajectories. A simple architectural fix: construct training trajectories where the correct answer appears at a random position in the sequence, and the model is trained to output the correct answer identically (i.e., learn an identity mapping when the current answer is already correct). This would require no changes to the polynomial approximation framework — it's purely a revision model training data intervention. The follow-up would train a revision model with this augmented data, measure the reversion rate, and test whether longer revision chains become reliable without the within-chain selection mechanism. If successful, this would transform revisions from a method requiring post-hoc answer selection (Figure 6) into a genuine iterative improver that monotonically converges to the correct answer. A negative result — that the model still reverts, or that the "stop revising" signal interferes with learning to revise incorrect answers — would suggest the reversion problem is more fundamental, perhaps arising from the model's inability to recognize correctness rather than from training data bias.
Combining polynomial approximants with sparse autoencoders for overcomplete feature dictionaries. The paper notes in Section 5 (Future Work) that quadratic approximants provide an overcomplete basis of d² eigenvectors for a d-dimensional residual stream, and speculates that this overcompleteness could address polysemanticity similarly to SAEs. A concrete follow-up would train an SAE on transformer FFN activations, then compare the SAE's learned dictionary features to the eigenvectors of the quadratic approximant of the same FFN — measuring overlap, sparsity, and monosemanticity. If the quadratic eigenvectors show comparable interpretability to SAE features without requiring SAE training, this would be a significant practical advance: it provides a "free" overcomplete dictionary from a one-time analytic computation. If the eigenvectors are less interpretable than SAE features, the combination may still be useful — for instance, using the quadratic eigenvectors to initialize the SAE dictionary (warm start) or as a regularizer during SAE training. The experiment would use a small transformer (GPT-2 size) where the FFN hidden dimension d is in the hundreds, making quadratic approximants computationally feasible. The key measurement is neuron monosemy: for each eigenvector, compute its activation distribution on a corpus and measure whether it fires selectively for interpretable concepts, using the same metrics (e.g., top activating dataset examples) that the SAE literature uses.
Dynamic difficulty estimation that amortizes assessment into problem-solving. As noted in the Limitations (Section 6), the paper's compute-optimal policy requires expensive upfront difficulty estimation (2048 samples per question) that is not included in efficiency calculations. A concrete follow-up would implement an adaptive scheme: begin with a small number of parallel samples (e.g., 4–8), compute the PRM's average score on these initial attempts as a rough difficulty signal, and use this to select the strategy for the remaining budget. The experiment would compare this adaptive scheme to the static (pre-computed bins) scheme on a held-out test set, measuring total cost (initial assessment samples + strategy execution samples) vs. accuracy, and comparing to the static scheme's cost + accuracy. If adaptive difficulty estimation matches or exceeds the static scheme's accuracy at lower total cost, it addresses the primary practical limitation of the paper's framework. If it underperforms, the paper's difficulty estimation cost is not merely a measurement oversight but reflects a genuine information-theoretic requirement — you really do need 2048 samples to assess difficulty accurately enough for optimal allocation.
Testing the Master Theorem on GLU variants in pretrained transformers. The paper derives the polynomial approximation framework for GLUs (Section 3.5) but provides no empirical results for GLU architectures — all experiments use standard ReLU MLPs on MNIST. A strong follow-up would apply the framework to the FFN layers of a pretrained small transformer (GPT-2, Pythia-160M) that uses a GLU variant, compute linear and quadratic approximants under a N(0, I) input assumption for the residual stream, and measure: (a) the R² of linear and quadratic approximants for each FFN layer, (b) the interpretability of the top eigenvectors of the quadratic approximant (by projecting them into vocabulary space via the unembedding matrix, following Millidge & Black, 2022), and (c) whether SVD-based attacks on the linear approximant transfer to the full transformer (i.e., does projecting the residual stream onto the nullspace of the FFN's linear approximant's top singular vectors degrade next-token prediction?). This would test the paper's central claim — that polynomial approximants enable interpretability for off-the-shelf pretrained models — on an architecture that actually matters for current LLM research, not just on MNIST MLPs. A negative result (low R², non-transferring attacks, uninterpretable eigenvectors) would establish that the method's applicability is limited to shallow networks on simple data, while a positive result would immediately make the technique relevant to the large community of transformer interpretability researchers.
Practical Applications and Downstream Use Cases
Universal adversarial attack construction for model auditing. The SVD attack in Section 4.4 demonstrates that projecting inputs onto the nullspace of the linear approximant's top 10 singular vectors (only 1.3% of MNIST input dimensions) reduces a ReLU MLP to random chance accuracy while leaving images human-intelligible (Figure 4). This provides a recipe for constructing universal, low-rank adversarial perturbations for any pretrained model whose linear approximant can be computed: fit the linear approximant under a maximum-entropy Gaussian input assumption, compute SVD, and use the top singular vectors as an adversarial subspace. For model auditing and red-teaming, this is valuable because the attack is derived analytically from the model's weights, not from iterative gradient-based optimization against specific inputs — it reveals the model's intrinsic sensitivity directions rather than per-sample vulnerabilities. A security auditor with access to a pretrained image classifier's weights could compute this SVD attack in minutes (for models where input dimension is tractable for the linear approximant, i.e., up to thousands of dimensions), generate a universal projection matrix, and test whether the model relies on an unexpectedly low-dimensional subspace for its decisions — a finding that would indicate fragility to simple input perturbations. The key practical advantage over standard adversarial attack methods (FGSM, PGD) is that no forward passes through the model are needed to construct the attack — it is a pure weight-space computation.
Monitoring functional complexity during training for early stopping and hyperparameter tuning. The phase transition in linear FVU (Figure 2, sharp increase at 500–1000 steps) provides a training diagnostic that is orthogonal to standard metrics like validation loss or accuracy. A practitioner training an MLP (or potentially a small transformer, if the method scales) could periodically checkpoint the model, compute the linear and quadratic FVU under a Gaussian input assumption, and use the timing and magnitude of the linear FVU increase as a signal about what the network is learning. If linear FVU rises while validation accuracy plateaus, the network may be overfitting — learning higher-order structure that doesn't generalize. If linear FVU rises while validation accuracy continues improving, the network is genuinely transitioning to more complex, task-relevant computation. This could inform early stopping (stop when quadratic FVU begins rising, indicating potential overfitting to higher-than-quadratic structure) or hyperparameter choices (if linear FVU rises very early in training, the learning rate may be too high and the network is jumping to complex solutions before exploiting simple linear structure). For the specific setting of MNIST with a shallow MLP, the paper's Figure 1 suggests that features are most interpretable at intermediate training stages — if interpretability is a goal (e.g., for models deployed in regulated settings), the FVU trajectory could guide checkpoint selection to the point where the network has learned the task but hasn't yet overfit to noisy higher-order patterns.
Post-hoc eigendecomposition of pretrained classifiers for feature visualization without training auxiliaries. For any pretrained MLP classifier with tractable input dimension (e.g., image classifiers with inputs up to a few thousand dimensions, tabular data models, small-scale scientific models), the quadratic approximant's eigendecomposition provides interpretable feature visualizations for free — no sparse autoencoder training, no probe fitting, no additional optimization. A practitioner loads the pretrained weights, specifies the input distribution (Gaussian with training data mean and covariance, or a Gaussian mixture for class-conditional analysis), computes the quadratic approximant analytically (or via the hybrid SGD approach in Appendix B if exact computation is intractable), eigendecomposes the class-specific interaction matrices, and visualizes the top eigenvectors as input-space images. For domains where input dimensions have spatial or semantic meaning (images, spectrograms, feature vectors with named dimensions), this provides immediate insight into what input patterns each output class is sensitive to. The method is particularly attractive for small-scale scientific applications — e.g., a biologist who trained an MLP to classify cell morphology from a hundred-dimensional feature vector can apply this technique to visualize what feature combinations the model associates with each cell type, without needing to train an interpretability model or collect additional data. The >95% R² on MNIST (Figure 2) suggests that for simple classification tasks, the quadratic features will be both accurate and interpretable.
When to Prefer This Method
The paper articulates clear boundary conditions for its approach rather than positioning it against a specific named alternative at the method level. The decision framework is therefore conditional on when the Gaussian maximum-entropy assumption is appropriate and when the computational costs are tractable, rather than a "prefer polynomial approximants vs. prefer SAEs" tradeoff — the two answer different questions (model features vs. data features) and are potentially complementary.
Prefer polynomial approximants under maximum-entropy inputs when:
- The input distribution is approximately Gaussian or Gaussian mixture, making the maximum-entropy assumption a reasonable approximation of the true data (e.g., MNIST pixel distributions, tabular data with approximately normal features, or settings where you deliberately want to analyze the model independent of data distribution).
- Input dimension is small-to-moderate (
d ≲ 100for exact quadratic approximants,d ≲ 2000for exact linear approximants, or larger with the hybrid SGD approach in Appendix B if convex optimization is acceptable). - The goal is to understand features of the model itself — its inductive biases, its functional form, its sensitivity directions — rather than features it computes on a specific dataset.
- You need a causal validation of feature importance (via SVD attacks) and want to avoid training auxiliary interpretability models.
- The network's quadratic
R²is high empirically (e.g., >0.90 on your data, as the paper finds on MNIST), indicating the polynomial approximation is faithful.
The method is not suitable when:
- The input distribution is highly non-Gaussian (e.g., discrete tokens, multimodal distributions with complex structure, distributions where the Gaussian is a poor fit). The paper recommends initializing with the analytic solution and finetuning on empirical samples in this case.
- Input dimension is in the tens of thousands (transformer residual streams, high-resolution images), where even the linear approximant's
O(d³)matrix inversion becomes expensive and quadratic approximants are completely intractable without dimensionality reduction. - The network's function is genuinely high-degree — if quadratic
R²is low, "there is little that can be done to interpret the model using our approach in isolation" (Section 5, Limitations). - The goal is to understand what the model computes on a specific, non-Gaussian data distribution (for which SAEs or probes trained on that distribution are more appropriate).
