Scalable Statistical Inference in Stochastic Gradient Descent

Abstract

Constructing confidence regions for stochastic gradient descent (SGD) ideally requires estimating the asymptotic covariance matrix, a severe computational bottleneck in high dimensions. Traditional cancellation-based batch means methods bypass this estimation but require inverting a sample batch covariance matrix. This introduces strict mathematical degeneracy when the parameter dimension exceeds the number of batches. To address this problem, we utilize equal batch size batch means method and propose a simultaneous, marginal-friendly framework. The proposed marginal statistics has a asymptotic Student's $t$-distribution, and eliminates the matrix inversion step, entirely circumventing high-dimensional degeneracy. To achieve valid simultaneous coverage, we present an algorithm utilizing wild bootstrap samples drawn from a statistic as a function of only the diagonals of the variance-covariance estimator, and to further incorporate the contribution of cross-dependencies, we introduce an efficient Quasi-Monte Carlo procedure utilizing a $t$-copula approximation. Additionally, we integrate a Lugsail variance estimator to aggressively correct finite-sample bias and under-coverage. The proposed methodology delivers interpretable, simultaneous hyper-rectangular confidence regions that are statistically robust, memory-efficient, and strictly scalable for high-dimensional inference. The theoretical results are supported by extensive numerical simulation analysis through various aspects of dimension, number of batches and error structure.