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

Abstract: By far the most common way to estimate an expected loss in machine learning is to draw samples, compute the loss on each one, and take the empirical average. However, sampling is not necessarily optimal. Given an MLP at initialization, we show how to estimate its expected output over Gaussian inputs without running samples through the network at all. Instead, we produce approximate representations of the distributions of activations at each layer, leveraging tools such as cumulants and Hermite expansions. We show both theoretically and empirically that for sufficiently wide networks, our estimator achieves a target mean squared error using substantially fewer FLOPs than Monte Carlo sampling. We find moreover that our methods perform particularly well at estimating the probabilities of rare events, and additionally demonstrate how they can be used for model training. Together, these findings suggest a path to producing models with a greatly reduced probability of catastrophic tail risks.

• The paper introduces a deterministic cumulant propagation framework that analytically propagates low-order cumulants to estimate MLP outputs with error scaling as O(n^(-K)).
• It leverages Hermite expansions and diagrammatic summation to manage non-linear activations and reduce computational cost compared to traditional Monte Carlo methods.
• Empirical results show order-of-magnitude improvements in efficiency and accuracy, particularly for rare event estimation in wide random MLPs.

Efficient Deterministic Estimation of Expected Outputs in Wide Random MLPs

Motivation and Problem Setting

The standard approach to evaluating expected losses or output statistics in neural networks is empirical: draw samples from the input distribution, run them through the network, and average the results (Monte Carlo (MC) sampling). While simple, this method is generally sub-optimal in terms of computational efficiency, especially in high-dimensional networks and when rare-event estimation is necessary. This paper addresses the deterministic estimation of expected outputs of wide, randomly initialized multilayer perceptrons (MLPs) given a formal input distribution (specifically, Gaussian), with the goal of achieving lower mean-squared error (MSE) per computational cost than MC sampling.

The problem is formulated as follows: Given a specific instance of an MLP with input, output, and hidden dimension $n$, $L$ hidden layers, and specified activation functions $\phi_1,\dots,\phi_L$, with weights drawn i.i.d. from $\mathcal{N}(0, c/n)$, estimate $$\mathbb{E}_{X \sim \mathcal{N}(0, I_n)}[M_\theta(X)]$$ (the "quenched" mean), for given weights and specified computational budget.

Methodological Framework

The approach centers on deterministic, sample-free estimation using high-dimensional moment methods and asymptotic techniques rooted in statistical physics and probability theory. The key technical innovation is the propagation of low-order cumulants of pre- and post-activation distributions through the network, leveraging network width $n$ as a scaling parameter to control error. The approach is formalized in a cumulant propagation framework, based on:

• Cumulant Expansion: The activation statistics of each layer are propagated by estimating their cumulants up to a chosen order $K$. Given that for Gaussian random variables, higher cumulants vanish, deviations from Gaussianity (central in finite width networks) are captured by nonzero higher cumulants.
• Hermite Expansions: Nonlinear activation functions are handled via Hermite expansions, allowing analytic propagation of cumulant estimates through polynomial or analytic nonlinearities.
• Diagrammatic Summation: The computation of output cumulants is reduced to a combinatorial diagrammatic expansion, generalizing the standard Isserlis/formula and covariance propagation recurrences.

The algorithm tracks a minimal set of cumulant statistics such that the total MSE incurred by omitting higher-order terms scales as $O(n^{-K})$ for “well-behaved” activations and sufficiently large $n$.

Algorithmic Details and Variants

Several algorithmic variants are proposed, structured around choices of cumulant order, tensor representation, and computational/accuracy trade-offs:

• Basic Cumulant Propagation: Tracks cumulants up to order $K$, combining multilinear propagation rules (for linear layers) with Hermite-based recurrence relations (for nonlinearities). Achieves MSE $L$0 in $L$1 FLOPs.
• Augmented Versions: Includes additional cumulant projections (e.g., traces of higher cumulant tensors) for improved finite-$L$2 accuracy.
• Factorized Representation: Leverages the symmetric and partially factorized structure of high-order cumulant tensors to reduce memory and computation from $L$3 to $L$4, in exchange for quadratic dependence in depth $L$5.

These algorithms can be directly differentiated with respect to network parameters, enabling their use as estimators in mechanistic training or distillation.

Theoretical Results

A central theoretical contribution is the derivation and proof of asymptotic error bounds and computational complexity for the cumulant propagation algorithm family.

Main claims:

• For polynomial activations and fixed depth $L$6, the deterministic estimator achieves target MSE $L$7 in $L$8 time, versus $L$9 for MC sampling—an asymptotic improvement by a factor of $\phi_1,\dots,\phi_L$0.
• For general analytic activation functions, analogous behavior is conjectured and empirically confirmed, with additional technical caveats regarding Hermite expansion truncation.

The error analysis relies on bounding the contraction of cumulant tensors through each layer, formalized in a general diagrammatic expansion. The depth dependence is worse for deterministic estimation than MC for very deep networks, but in practically relevant regimes (depth moderate, width large), the deterministic approach dominates.

Empirical Results

Extensive experiments substantiate the theory across a range of $\phi_1,\dots,\phi_L$1, $\phi_1,\dots,\phi_L$2, activation functions (ReLU, GELU, tanh), and cumulant orders $\phi_1,\dots,\phi_L$3.

Key empirical findings:

• Efficiency: For ReLU MLPs of width up to 256, deterministic cumulant propagation (factorized version) achieves order-of-magnitude lower MSE per FLOP than MC sampling across a broad computational regime, with improvements up to $\phi_1,\dots,\phi_L$4–$\phi_1,\dots,\phi_L$5 for moderate depths.
• Rare Event Estimation: For estimation of events with probabilities $\phi_1,\dots,\phi_L$6 (with $\phi_1,\dots,\phi_L$7 the sample size of MC estimate), deterministic methods outperform all stochastic and Gaussian-extrapolation baselines, maintaining useful relative error where MC degenerates.
• Mechanistic Distillation: Deterministic estimators can be utilized to train student networks by minimizing analytically estimated expected loss. While not quite matching SGD’s efficiency on non-random networks, this demonstration opens a path for future research on analytic training.
• Scaling Law Confirmation: Empirically validated scaling of MSE as $\phi_1,\dots,\phi_L$8, with clear benefits from increasing cumulant order $\phi_1,\dots,\phi_L$9 at moderate and large widths.

Implications and Future Directions

The results have significant implications for both practical training and theoretical understanding of wide MLPs.

• Training for Extreme-Value Reliability: By efficiently estimating expected loss in rare-event regimes, these methods offer the potential to train networks that are robust to tail risks and failure modes unseen in finite samples—an important consideration for safety-critical systems.
• Mechanistic Training Paradigms: The ability to analytically propagate uncertainty and optimize “mechanistic” loss estimates could pave the way for training protocols wherein sample inefficiency is side-stepped, provided one can extend the theory to post-initialization or adapt to non-Gaussian input distributions.
• Extension to Learned or Structured Networks: While the present results rigorously apply to random (or “maximally unstructured”) weights at initialization, the methodology suggests avenues for approximate analytic propagation in partially trained or structured networks, via decomposition into structured/random components.
• Algorithmic Generalizations: Beyond MLPs, similar strategies may be applicable to other architectures wherever input-output relations admit sufficient analytic description (e.g., convolutional or residual networks at initialization).

Depth and activation dependence remain as limitations: For very deep networks, the error scaling for deterministic estimators is suboptimal, and algorithmic improvements are needed to match MC in these cases.

Conclusion

This work demonstrates that analytic, deterministic estimation using cumulant propagation and Hermite expansions can substantially outperform MC sampling for wide random MLPs, both in theory and practice, especially when rare event probabilities or low MSE are required. These results elucidate the connection between asymptotic network theory and practical algorithm design, and offer promising foundations for future research on sample-free training methods, rare event risk mitigation, and mechanistic interpretability in neural networks.

Reference:

"Estimating the expected output of wide random MLPs more efficiently than sampling" (2605.05179)