Optimal Non-Asymptotic Edgeworth Expansions for Multivariate Neural Network Outputs

Abstract: Finite-width fully connected neural networks with Gaussian-initialized weights deviate from their infinite-width Gaussian limit, exhibiting non-vanishing higher-order cumulants. We approximate these deviations, for a neural network evaluated in a finite number of inputs, using multidimensional Edgeworth expansions of arbitrary order $4m-1$, with $m\in\mathbb{N}$. Assuming that the corresponding Gaussian limit has an invertible covariance matrix and that the activation function is polynomially bounded, we establish a bound of order $n^{-m}$ on the total variation distance between the law of the true network output and its Edgeworth approximation, with matching lower bounds. As an application, we quantify the error in Bayesian posterior distributions when the prior is replaced by its Edgeworth expansion. Our results are more general and also apply to sequences of conditionally Gaussian vectors converging to a Gaussian vector with invertible covariance.

• The paper introduces optimal non-asymptotic Edgeworth expansions that characterize finite-width corrections in multivariate neural network outputs.
• It establishes matching upper and lower TV distance bounds, proving the optimality of the expansion rate order n^{-m}.
• The method enables refined Bayesian inference by accurately approximating output distributions beyond the infinite-width Gaussian limit.

Optimal Non-Asymptotic Edgeworth Expansions for Multivariate Neural Network Outputs

Introduction and Motivation

Finite-width effects in neural networks have a significant and quantifiable impact on their output distributions, particularly in regimes where asymptotic (infinite-width) limits fail to accurately describe practical models. Standard results show that as the width $n$ of a fully connected neural network increases, its output converges in distribution to a Gaussian process. However, at finite width, higher-order cumulants remain non-zero, producing systematic non-Gaussian corrections. The paper "Optimal Non-Asymptotic Edgeworth Expansions for Multivariate Neural Network Outputs" (2605.24072) rigorously characterizes these corrections for multidimensional network outputs through multivariate Edgeworth expansions of arbitrary order. The main focus is to derive optimal non-asymptotic bounds in total variation (TV) distance between the law of the finite-width output and its Edgeworth expansion, along with matching lower bounds, and to demonstrate applications for Bayesian posterior inference with neural networks.

Theoretical Framework and Edgeworth Expansion

The analysis begins by extending classical Central Limit Theorem (CLT) techniques to the high-dimensional outputs of neural networks. While the infinite-width limit yields a Gaussian law for the outputs under Gaussian-initialized weights, the finite-$n$ corrections are systematically described by multidimensional Edgeworth expansions. The Edgeworth expansion leverages the sequence of cumulants of the neural network output and expresses the probability density as a sum of the reference Gaussian density and polynomial perturbations (involving Hermite polynomials), organized as a power series in $n^{-1/2}$.

For a conditionally Gaussian random vector $Z$ (representing the network output at finite width) with covariance matrix $A^{\oplus n}$ and reference Gaussian $G$ with covariance $K^{\oplus n}$, the Edgeworth expansion to order $4m-1$ constructs a signed measure $\gamma_{Z,G,m}$ whose density approximates the true law to order $n^{-m}$ in TV distance. The explicit structure of the expansion accommodates arbitrary dimension and arbitrary order $n$0, utilizing combinatorial constructs over Hermite polynomials and expectations of functions of the covariance difference $n$1.

Non-Asymptotic and Optimal Total Variation Bounds

A primary theoretical contribution is the derivation of matching upper and lower non-asymptotic bounds for the TV distance between the true multivariate network output law and its Edgeworth expansion:

• For general conditionally Gaussian vectors (and specifically neural network outputs at initialization), for any order $n$2, the TV distance satisfies

$n$3

for explicit constant $n$4 independent of width $n$5, under conditions on invertibility and boundedness of the covariance and activation function moments.

• Matching Lower Bound: When the Edgeworth expansion uses the expected conditional covariance as the reference, the TV distance is also bounded below by $n$6 for some $n$7. This demonstrates that the obtained rate is tight and optimal.

Numerical experiments demonstrate that even for moderate widths (e.g., $n$8), finite-order Edgeworth expansions provide drastic improvements over the infinite-width Gaussian approximation, and only for very small widths do oscillatory high-order terms limit the approximation quality.

Figure 1: Comparison between the estimated density of a neural network output (NN) and its Gaussian and Edgeworth (para-Gaussian) approximations for various widths and expansion orders.

Figure 2: Signed pointwise error between the Monte Carlo KDE of the NN output and its Gaussian/Edgeworth approximations at increasing widths, highlighting improved accuracy and the behavior of higher-order corrections.

Application to Bayesian Inference with Neural Networks

The work applies these non-asymptotic expansions to quantify errors incurred in Bayesian supervised learning if one substitutes the true prior distribution induced by the network output with its finite-order Edgeworth expansion. Specifically:

• For networks with Gaussian-initialized weights and biases, substituting the true network prior with the Edgeworth expansion yields a posterior distribution approximating the true posterior up to $n$9 in TV distance.
• In scenarios with Gaussian likelihoods (e.g., regression with $n^{-1/2}$0 loss), the resulting posterior admits explicit Hermite-polynomial expansions and normalization formulas, enabling principled uncertainty quantification and predictive inference that are finer than those available from infinite-width Gaussian process limits.

These results indicate that Edgeworth-corrected priors are theoretically justified and, for wide but finite networks, outperform both the standard Gaussian and existing heavy-tailed prior proposals (e.g., Student-$n^{-1/2}$1 priors [SWG14, PAPGGR23]) in terms of captured non-Gaussianity and approximation error.

Practical and Theoretical Implications

The Edgeworth approach elaborated in this work provides several key advances:

• Rigorous finite-width quantification: The results give explicit, dimension- and order-dependent error bounds, facilitating rigorous uncertainty estimates in both forward generative modeling and Bayesian inference with realistic (finite) networks.
• High order and multivariate expansions: The generalized structure allows arbitrary expansion order and output dimension, which is essential for multivariate prediction and uncertainty modeling in real-world tasks.
• Optimal bounds: The matching lower and upper bounds in TV distance establish that the Edgeworth expansion provides the best possible rate for this type of finite-size approximation, eliminating gaps between approximate and true output laws at a given order.

Theoretical implications extend to generalized weights beyond fully Gaussian settings. The same rates are shown to hold (under mild assumptions) for conditionally Gaussian neural networks with non-Gaussian weights [C26].

Potential Future Directions

Possible extensions include allowing simultaneous width and depth scaling, explicit dependence on network hyperparameters, studying other architectures (CNNs, ResNets) where conditional independence and permutation structures may differ, and integrating expansion-based corrections into scalable Bayesian deep learning pipelines and uncertainty quantification toolkits.

Conclusion

This work achieves a rigorous, order-optimal non-asymptotic characterization of finite-width corrections for multivariate neural network outputs via multidimensional Edgeworth expansions, with explicit, optimal control of the TV distance to the true output law. The analysis enables refined Bayesian inference and principled uncertainty quantification for wide, finite networks—bridging the gap between infinite-width Gaussian process limits and practical, finite models. The combination of explicit theory, sharp error control, and practical applicability provides a new benchmark for the mathematical analysis of neural network distributions and Bayesian approaches in high-dimensional machine learning (2605.24072).