An Isotropic Approach to Efficient Uncertainty Quantification with Gradient Norms

Abstract

Existing methods for quantifying predictive uncertainty in neural networks are either computationally intractable for large language models or require access to training data that is typically unavailable. We derive a lightweight alternative through two approximations: a first-order Taylor expansion that expresses uncertainty in terms of the gradient of the prediction and the parameter covariance, and an isotropy assumption on the parameter covariance. Together, these yield epistemic uncertainty as the squared gradient norm and aleatoric uncertainty as the Bernoulli variance of the point prediction, from a single forward-backward pass through an unmodified pretrained model. We justify the isotropy assumption by showing that covariance estimates built from non-training data introduce structured distortions that isotropic covariance avoids, and that theoretical results on the spectral properties of large networks support the approximation at scale. Validation against reference Markov Chain Monte Carlo estimates on synthetic problems shows strong correspondence that improves with model size. We then use the estimates to investigate when each uncertainty type carries useful signal for predicting answer correctness in question answering with large language models, revealing a benchmark-dependent divergence: the combined estimate achieves the highest mean AUROC on TruthfulQA, where questions involve genuine conflict between plausible answers, but falls to near chance on TriviaQA's factual recall, suggesting that parameter-level uncertainty captures a fundamentally different signal than self-assessment methods.

Figures (8)

• Figure 1: Normalized epistemic uncertainty on the XOR problem under three covariance assumptions, and the log ratio $\log(U_A / U_B)$. The identity produces spatially symmetric estimates that peak along the two decision boundaries. Each proxy inflates uncertainty in the half of input space absent from its data and suppresses it where its data is concentrated, despite the problem's underlying symmetry; the log ratio makes this asymmetry explicit. See \ref{['app:hessian-bias']} for experimental details and additional problems.
• Figure 2: Multiclass spirals uncertainty maps. Left two panels: epistemic uncertainty (MCMC vs. gradient norm $\|g\|^2$). Right two panels: aleatoric uncertainty (MCMC vs. point estimate). All maps are individually normalized to $[0, 1]$. Additional problems in \ref{['app:classification-additional']}.
• Figure 3: Correlation between our estimates and MCMC estimates as a function of model size (number of parameters) on a concentric rings problem. Both epistemic (blue) and aleatoric (orange) correlations follow a U-shaped trajectory, dipping at intermediate scales and recovering at larger model sizes. Full per-model results in \ref{['tab:scaling-full']}.
• Figure 4: Epistemic uncertainty distributions on the DistilBERT domain classification task under three covariance assumptions, and the log ratio between proxy estimates. Under the identity, both domains have similar uncertainty distributions. Each proxy shifts the relative uncertainty between domains: $H_A^{-1}$ (estimated from science data) suppresses uncertainty on science inputs relative to sports, while $H_B^{-1}$ does the reverse. The log ratio makes this asymmetry explicit, flipping sign between domains.
• Figure 5: Binary XOR uncertainty maps. Left two panels: epistemic uncertainty (MCMC vs. $\|g\|^2$). Right two panels: aleatoric uncertainty (MCMC vs. point estimate). All maps are individually normalized to $[0, 1]$.