📊 statistics

Estimating the expected output of wide random MLPs more efficiently than sampling

This article proposes a sampling-free method that utilizes cumulants and Hermite expansions to efficiently estimate the expected outputs of broad random MLPs, achieving lower computational costs compared to traditional Monte Carlo sampling and superior accuracy for rare events.

Original authors: Wilson Wu, Victor Lecomte, Michael Winer, George Robinson, Jacob Hilton, Paul Christiano

Published 2026-05-07

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Wilson Wu, Victor Lecomte, Michael Winer, George Robinson, Jacob Hilton, Paul Christiano

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Estimating Averages

Imagine a huge, complex machine (a neural network) made up of thousands of gears and levers. You want to know: "If I feed this machine a random input, what is the average output it will produce?"

The standard way to answer this in the world of machine learning is Monte-Carlo simulation.

The old way: You feed the machine a random input and record the output. You do this 1,000 times. Then 10,000 times. Then 100,000 times. Finally, you take the average of all these results.
The problem: This is like trying to guess the average height of all people in a city by measuring one person after another. It works, but it is incredibly slow and computationally expensive. If you want a very precise answer, you have to run the machine millions of times.

The New Solution: The "Mechanical" Map

The authors of this paper propose a different approach. Instead of running the machine over and over, they want to calculate the answer directly by analyzing how the machine's gears are connected.

They call this Cumulant Propagation.

The Analogy: The Foggy Factory

Imagine the machine as a factory where raw materials (inputs) come in on one side and products (outputs) come out on the other.

The Input: The raw materials are somewhat "foggy" or uncertain (random).
The Process: As the materials travel through the factory, they are mixed, heated, and shaped by various machines (layers of the network).
The Goal: We want to know what the shape of the fog looks like at the very end of the factory.

The old way (Simulation): You send a single truck with raw materials through the factory and see what comes out. Then you send another truck. And another. You keep doing this until you have a good idea of the final shape.

The new way (Cumulant Propagation): Instead of sending trucks, you look at the factory's blueprint. You know exactly how the first machine mixes the fog. You know how the second machine stretches it.

The authors have developed a mathematical "lens" (using tools called cumulants and Hermite expansions) that allows them to track the shape of the fog as it travels through the factory without ever actually sending a truck through.
They track the "center" of the fog, how "spread out" it is, and how "lumpy" or "strange" it becomes. They pass these statistics from one machine to the next, updating the shape mathematically until they reach the end.

Why This Is a Big Deal

The paper shows that for wide networks (factories with very wide conveyor belts), this new method is much faster than the old simulation method.

Efficiency: To achieve the same level of accuracy, the new method uses significantly fewer "computational operations" (FLOPs). In some cases, it is 100 times faster.
Rare Events: The new method is particularly good at detecting rare events.
- Analogy: Imagine you want to know the probability that a specific, very rare defect occurs in the factory.
- Simulation: You could run the factory a million times and never see the defect. You would have to guess it is zero, or run it a billion times to see it once.
- New Method: Since it analyzes the mechanics of the factory, it can estimate the probability of this rare defect even if it has never actually occurred in a simulation. It is like looking at the blueprint and saying, "If the gears are aligned exactly this way, a defect could occur," without waiting for it to happen.

How It Works (The "Secret Recipe")

The paper relies on some clever mathematical tricks to make this possible:

Cumulants: Think of these as a way to describe the "shape" of the fog.
- The first cumulant is the average.
- The second is the spread (variance).
- The third and fourth describe how skewed or peaked the fog is.
- The authors track these shapes layer by layer.
Hermite Expansions: When the fog hits a nonlinear machine (like a ReLU activation that cuts off everything below zero), the shape gets distorted. The authors use a special mathematical series (similar to a Taylor series, but for shapes) to approximate how this distortion happens without doing the heavy lifting of a full simulation.
Factorization: To keep the math from becoming too heavy, they break down the complex shapes into smaller, manageable pieces (factors), similar to how one might break a huge puzzle into smaller sections to solve it faster.

What They Actually Claim

It works for random networks: The method is proven to work best on networks where the weights (the settings of the gears) are chosen randomly at the start.
It beats simulation: For wide networks, this method reaches a target level of accuracy with far fewer computer operations than running simulations.
It can train networks: Since the method provides a smooth, mathematical estimate (instead of a noisy average from simulations), it can be used to train a student network that mimics a teacher network. They call this "mechanistic distillation."
It helps with safety: Since it is better at estimating rare, unlikely events, this method could theoretically help train models that are less likely to make catastrophic errors (tail risks) that are too rare to be captured by standard simulations.

What It Is NOT

It is not a cure-all for every neural network. It works best on "wide" networks (many neurons) and is still being developed for very deep or narrow networks.
It does not yet replace simulation for all tasks; it is a specialized tool for estimating expected values in specific, well-behaved scenarios.

In short: The authors have found a way to calculate the answer to a complex probability question by analyzing the structure of the machine, rather than guessing the answer by running the machine millions of times.

Technical Summary: More Efficient Estimation of Expected Outputs of Wide Random MLPs than via Sampling

1. Problem Statement

The standard approach to estimating the expected loss of a neural network involves drawing samples from the input distribution, computing the loss for each, and forming the empirical average (Monte-Carlo sampling). Although this seems natural, this method can be suboptimal when the input distribution and the loss function are formally specified.

This contribution addresses the specific problem of estimating the expected output of a wide random Multilayer Perceptron (MLP) over Gaussian inputs. Specifically, the goal is, given a network with fixed weights $\theta$ drawn from a Gaussian initialization (a "quenched" average in the sense of statistical physics), to estimate:
$\mathbb{E}_{X \sim \mathcal{N}(0, I_n)} [M_\theta(X)]$
where $M_\theta$ is an MLP with width $n$ and $L$ hidden layers. The objective is to achieve a target mean squared error (MSE) of $O(\epsilon^2)$ with significantly fewer floating-point operations (FLOPs) than the $O(n^2/\epsilon^2)$ required for naive Monte-Carlo sampling.

2. Methodology

The authors propose sampling-free estimation procedures based on cumulant propagation. Instead of sampling inputs, the algorithms propagate the statistics of the activation distributions layer by layer analytically.

Core Mathematical Tools

Cumulants: The algorithms track the cumulants of the activation distributions. The $k$ -th order cumulants measure deviations from Gaussianity (where orders 1 and 2 correspond to the mean and covariance).
Hermite Expansions: To handle nonlinear activation functions $\phi$ , the authors expand $\phi$ in terms of Hermite polynomials with respect to the local Gaussian approximation of the pre-activation distribution.
Diagram Summation Formula: A key theoretical contribution is the Hermite-based diagram summation formula. This formula expresses the cumulants of the output of a nonlinear layer in terms of the cumulants of the input and the Hermite coefficients of the activation function. It generalizes covariance propagation (Wright et al., 2024) to arbitrary orders.

Algorithmic Variants

The contribution describes four versions of the estimator that trade off computational cost and accuracy:

Base Algorithm: Tracks cumulants up to a maximum order $K$ . It uses the diagram summation formula to propagate cumulants through linear and nonlinear layers.
Augmented Algorithm: Tracks additional harmonic components of the $(K+1)$ -th order cumulant tensor (specifically the full trace) to improve the constant factor in the MSE without increasing the asymptotic runtime order.
Factored Algorithm: Represents the $K$ -th order cumulant tensor in factored form (size $O(n^{K-1})$ ) rather than as a full tensor (size $O(n^K)$ ). This reduces the runtime of the linear step (matrix multiplication) by a factor of $n$ .
Factored Augmented: Combines both improvements.

Important Adjustments for Accuracy

To ensure that the error scales correctly with the network width $n$ , the authors introduce two critical adjustments:

Power Cumulants: Instead of expanding the activation function $\phi$ directly, they expand powers of $\phi$ (e.g., $\phi^2, \phi^3$ ). This enables the exact calculation of variances and higher moments of the activation output and handles cases where indices in the cumulant tensor are repeated.
Tracking the Full Trace: If the maximum cumulant order $K$ is odd, the algorithm tracks the full trace of the $(K+1)$ -th order cumulant tensor. This captures specific correlations that would otherwise be lost and ensures that the error scales as $O(n^{-K})$ .

3. Key Contributions and Results

Theoretical Results

The contribution provides rigorous theoretical bounds for networks with polynomial activation functions (and conjectures an extension to general functions):

Theorem 5.1 (Base Algorithm): For a fixed depth $L$ , the base algorithm achieves an MSE of $O(\epsilon^2)$ in time $O(n^2/\epsilon^2)$ . This matches the asymptotic complexity of Monte-Carlo sampling, albeit with a potentially better constant value.
Theorem 5.2 (Factored Algorithm): By using the factored representation, the algorithm achieves an MSE of $O(\epsilon^2)$ in time $O(n/\epsilon^2)$ , provided $\epsilon(n)^2 = O(1/n^2)$ . This represents a speedup by a factor of $n$ over Monte-Carlo sampling for sufficiently wide networks.
Error Scaling: The MSE scales as $O(1/n^K)$ , where $K$ is the maximum tracked cumulant order. The error grows with depth $L$ (conjectured as $O((L/n)^K)$ ), which constitutes a limitation compared to sampling, where the error is independent of depth.

Empirical Results

Experiments were conducted on random ReLU MLPs with widths up to 256 and depths up to 12.

Efficiency: The factored and factored augmented algorithms outperform Monte-Carlo sampling by factors between 10 and 1,000 across compute budgets spanning 7 orders of magnitude. For example, with 4 hidden layers and a width of 256, the methods outperform sampling by more than 100-fold.
Estimation of Low Probabilities: The methods excel at estimating the probabilities of rare events (tails of the distribution). While Monte-Carlo sampling fails when the true probability drops below $1/N$ (where $N$ is the number of samples), the cumulant propagation method maintains a low relative error for probabilities orders of magnitude smaller than $1/N$ .
Mechanistic Distillation: As a proof of concept, the authors used the differentiability of their estimations to train a "student" network to mimic a random "teacher" network (mechanistic distillation). Although this is somewhat less efficient in this specific setup than standard stochastic gradient descent (SGD), it demonstrates the feasibility of training via analytical loss estimation.

4. Significance and Claims

The authors position this work as a step towards mechanistic estimation, a paradigm where one reasons about the internal mechanisms of a network to derive properties of its output distribution, rather than relying on blind sampling.

Efficiency: The main claim is that for wide random networks, analytical methods can significantly outperform sampling in estimating expected outputs regarding computational efficiency.
Safety and Rare Events: The contribution highlights that the ability to accurately estimate low-probability events is crucial for safety. Standard SGD may never sample a "catastrophic" rare event during training, leading to models that ignore such risks. Mechanistic estimation could theoretically produce models robust to these rare but consequential events by analytically capturing tail probabilities.
Mechanistic Training: The authors propose a path towards "mechanistic training," where gradient descent is applied to analytical loss estimates. They acknowledge current limitations (guarantees apply primarily to random, untrained networks) but suggest that future work could adapt these methods to partially trained networks by treating weights as a mixture of "structure" and "randomness."

The contribution concludes that while the current results are specific to random MLPs at initialization, they offer a "first glimpse" into how deductive, mechanistic thinking could replace or complement sampling in machine learning, potentially leading to models with drastically reduced probabilities of catastrophic tail risks.