Beyond Sampling: Analytical Estimation of Neural Network Outputs at Initialization

Paper: 2605.05179 Authors: Wilson Wu, Victor Lecomte, Michael Winer, George Robinson, Jacob Hilton, Paul Christiano Categories: cs.LG, cond-mat.dis-nn, stat.ML

The Gap

Machine learning relies on estimating expected losses by drawing samples and averaging. This works, but it’s computationally expensive and particularly bad at estimating rare events—exactly where catastrophic failures hide. Prior work on neural network theory at initialization (neural tangent kernels, mean field theory) gives us distributional insights but doesn’t directly solve the computational problem: we still run samples through networks to estimate expectations.

The gap: we understand the limiting distributions of wide networks theoretically, but we haven’t exploited this understanding to skip the sampling step entirely.

Problem: Need E[f(x)] for network f, input distribution p(x)
   |
   v
Current approach: Sample x1...xn ~ p(x), compute mean(f(xi))
   |                                    |
   +-----> Expensive for rare events <--+
   +-----> Scales poorly with network depth
   
This paper's path:
   |
   v
Assumption: Network is wide, at initialization
   |
   v
Method: Track activation distributions analytically layer-by-layer
   |
   v
Evidence: FLOP count < sampling for target MSE, better tail estimation
   |
   v
Conclusion: Can estimate expectations without running samples

The Increment

One sentence: Before this paper, estimating a network’s expected output meant running samples through it; after, you can compute approximate activation distributions analytically and skip sampling for wide networks at initialization.

Core Mechanism

The method treats each layer’s activations as random variables and propagates their distributions forward analytically. At layer 0, you have Gaussian inputs. At layer 1, you compute the distribution of activations after applying weights (Gaussian-distributed at init) and nonlinearities. The key: represent these distributions not as samples but as mathematical objects—specifically, using cumulants (moments of the log-characteristic function) and Hermite polynomial expansions.

For each layer, the algorithm: (1) takes the previous layer’s activation distribution, (2) applies the linear transformation (which is tractable because weights are Gaussian), (3) applies the nonlinearity by expanding it in Hermite polynomials and computing how cumulants transform, (4) outputs a new distribution representation. The final layer’s distribution gives you the expected output directly—no samples needed.

The Hermite expansion is crucial: any reasonable function can be written as a weighted sum of Hermite polynomials, and these polynomials have known behavior under Gaussian convolution. So instead of sampling to see what ReLU or tanh does to a Gaussian-ish distribution, you compute it algebraically.

Layer l-1 activations:          Layer l activations:
  Distribution D(l-1)              Distribution D(l)
  (cumulants k1,k2,...)            (cumulants k1',k2',...)
         |                                ^
         |                                |
         +---> Linear: W*x + b -----------+
         |     (Gaussian weights)         |
         |                                |
         +---> Nonlinearity: σ(·) --------+
               (Hermite expansion)
               
Final layer D(L) ---> E[output] directly computed

Think of it like weather forecasting versus actually waiting for the weather. Sampling is like running the actual weather day by day and averaging. This paper’s method is like having equations for atmospheric dynamics—you simulate the probability distributions forward without living through each possible weather timeline. Each layer is a time step, and instead of tracking individual air molecules (samples), you track probability clouds (distributions). The Hermite polynomials are like basis functions for describing cloud shapes, and cumulants are compact summaries of those shapes. You push these clouds through the network’s transformations algebraically, and at the end, you read off the expected outcome from the final cloud’s shape.

Key Concepts

• Cumulants: Imagine you want to describe a probability distribution. You could list all possible outcomes and their probabilities (exhaustive but unwieldy), or you could give summary statistics. Mean and variance are the first two moments—they tell you center and spread. Cumulants are an alternative set of summaries that have nicer mathematical properties: the cumulants of a sum of independent variables just add up. For a Gaussian, only the first two cumulants (mean and variance) are nonzero. For other distributions, higher cumulants capture skewness, heavy tails, etc. In this paper, cumulants compactly represent how activations are distributed at each layer, and they compose nicely through the network’s operations.

• Hermite Polynomials: These are a sequence of polynomials (H₀(x)=1, H₁(x)=x, H₂(x)=x²-1, …) that are orthogonal with respect to Gaussian weight. Any smooth function can be written as a weighted sum of Hermite polynomials—like a Fourier series but for Gaussians. Why useful here? When you apply a nonlinearity σ to a Gaussian random variable, you can expand σ in Hermite polynomials and compute the output distribution’s cumulants directly from the expansion coefficients. It’s like having a lookup table: “If input is Gaussian with variance v, and nonlinearity is ReLU, then output cumulants are…”. No sampling required.

• Width Regime: Neural network theory often studies the “infinite width limit” where each layer has infinitely many neurons. In this limit, the central limit theorem kicks in hard: activations become Gaussian, and the network’s behavior is governed by a kernel (the neural tangent kernel). This paper operates in the wide-but-finite regime: wide enough that distributional approximations are accurate, but not so wide that everything collapses to a Gaussian. The method exploits width to justify truncating Hermite expansions and cumulant series—higher-order terms vanish as width grows.

Framework Shift

Before (mainstream approach):        After (this paper):

Input x ~ p(x)                       Input distribution p(x)
    |                                    |
    v                                    v
Sample x1, x2, ..., xn            Represent p(x) as cumulants
    |                                    |
    v                                    v
Run each xi through network       Propagate distribution analytically:
f(x1), f(x2), ..., f(xn)            Layer 1: D1 = σ(W0*D0 + b0)
    |                                 Layer 2: D2 = σ(W1*D1 + b1)
    v                                    ...
Compute mean: (1/n)Σf(xi)             Layer L: DL
    |                                    |
    v                                    v
Estimate E[f(x)]                   Compute E[f(x)] from DL directly

Cost: O(n * network_flops)         Cost: O(layers * cumulant_ops)
Rare events: needs huge n          Rare events: read from tail of DL

From sample-then-average to distribution-then-query, the core shift is replacing stochastic simulation with deterministic propagation of probability measures.

Expert Assessment

Problem choice: Real gap. Sampling is genuinely expensive, especially for safety-critical applications where you care about 1-in-10,000 events. The problem sits at the intersection of neural network theory (which has exploded post-NTK) and practical initialization/training concerns. Not manufactured—people actually want cheaper ways to estimate expectations.

Method maturity: Clever insight, not brute force. The use of Hermite expansions is well-motivated by the Gaussian structure at initialization. However, the method is limited to initialization (weights haven’t moved yet) and wide networks (so approximations hold). The authors acknowledge this—it’s a first step, not a complete solution. A simpler approach might be importance sampling or variance reduction techniques, but those still require running samples. This is genuinely different.

Experimental integrity: Baselines are fair—they compare against standard Monte Carlo. The FLOP counts are carefully tracked, which is good. The experiments on rare event estimation are compelling. One concern: the method’s advantage grows with network width, but the experiments mostly show moderate widths. Would like to see more exploration of where the crossover happens (when does this beat sampling?). Also, the training experiments are preliminary—showing it’s possible, not that it’s practical yet.

Writing quality: The paper is dense. The main text rushes through the method to get to results. The appendix has the real meat, but it’s a slog. The section on Hermite expansions could be rewritten with more intuition upfront—right now it’s “here’s the math, trust us it works.” A clearer roadmap in the intro (why cumulants? why Hermite? why does width help?) would elevate the whole paper. The related work section undersells how different this is from prior neural network theory work.

Verdict: weak accept — Solid theoretical contribution with promising empirical results, but limited to a narrow regime (wide networks at init) and needs more work to be practically useful. The rare event estimation angle is the most compelling application.

Takeaways

• Hermite expansions for nonlinearities: If you’re ever propagating distributions through nonlinear functions and the input is approximately Gaussian, consider expanding the nonlinearity in Hermite polynomials. The coefficients tell you how moments/cumulants transform. This trick transfers beyond neural networks—anywhere you have Gaussian inputs and smooth nonlinearities.

• Cumulants over moments: When tracking how distributions evolve through compositions of operations, cumulants often compose more cleanly than raw moments. They’re underused in ML outside of specific theory subfields.

• Rare event estimation: The paper’s strongest practical contribution. If you care about tail probabilities (adversarial robustness, safety), analytical methods can give you better estimates than sampling with reasonable compute budgets. Even if you can’t use this exact method, the principle—exploit structure to avoid sampling—is worth stealing.

• Width as a computational resource: We usually think of width as a capacity/expressiveness knob. This paper shows it’s also a knob for making certain computations tractable. Wider networks have simpler distributional behavior, which can be exploited algorithmically.