Refining Covariance Matrix Estimation in Stochastic Gradient Descent Through Bias Reduction

1. Problem Statement

The paper addresses the challenge of online covariance matrix estimation for the Stochastic Gradient Descent (SGD) algorithm.

• Context: In large-scale machine learning, SGD is used to minimize a convex objective function $F(x) = \mathbb{E}[f(x, \xi)]$. Under specific conditions (polynomially decaying step sizes $\eta_i = \eta i^{-\alpha}$ with $\alpha \in (0.5, 1)$), the averaged SGD iterates $\bar{x}_n$ converge asymptotically to a normal distribution: $\sqrt{n}(\bar{x}_n - x^*) \Rightarrow N(0, \Sigma)$.
• The Limiting Covariance ($\Sigma$): The matrix $\Sigma$ takes a "sandwich" form: $\Sigma = A^{-1} S A^{-1}$, where $A = \nabla^2 F(x^*)$ (Hessian) and $S = \mathbb{E}[\nabla f \nabla f^T]$ (gradient noise covariance).
• The Challenge: To perform statistical inference (e.g., constructing confidence intervals), one must estimate $\Sigma$.
  • Plug-in estimators require computing the Hessian $A$, which is computationally expensive ($O(d^3)$) and often unavailable in high-dimensional or non-differentiable settings.
  • Bootstrap methods require multiple resampled runs, violating the "single-pass" efficiency of SGD.
  • Existing Hessian-free methods (Batch-Means): While they avoid second-order derivatives, they suffer from slow convergence rates (specifically $O(n^{(\alpha-1)/4})$) due to inherent bias in their formulation.

The authors ask: Can we achieve a faster convergence rate for covariance estimation using only a single SGD trajectory, without Hessian information?

2. Methodology: The De-biased Estimator

The authors propose a novel fully online, de-biased covariance estimator that eliminates the need for second-order derivatives while significantly improving accuracy.

A. Core Intuition: Bias Reduction

Standard batch-means estimators estimate variance by analyzing the variability of local batches. However, due to the high correlation in SGD iterates, these estimators suffer from significant bias.
The authors derive an alternative identity for the variance of the sample mean $\bar{X}_n$:
$$\sigma_n = \frac{1}{n} \sum_{i=1}^n \mathbb{E} \left[ 2(X_i - \mathbb{E}X_i) \sum_{k=1}^i (X_k - \mathbb{E}X_k) - (X_i - \mathbb{E}X_i)^2 \right]$$
Instead of adjusting weights on a biased sample variance, they construct an estimator that directly approximates this theoretical expansion.

B. The Estimator Formulation

For the multivariate SGD iterates $\{x_i\}$, the proposed de-biased estimator $\hat{\Sigma}_n$ is defined as:
$$\hat{\Sigma}_n = \frac{1}{n} \sum_{i=1}^n \left[ (x_i - \bar{x}_n) \left( \sum_{k=i-\ell_i+1}^i x_k - \ell_i \bar{x}_n \right)^T + \left( \sum_{k=i-\ell_i+1}^i x_k - \ell_i \bar{x}_n \right) (x_i - \bar{x}_n)^T - (x_i - \bar{x}_n)(x_i - \bar{x}_n)^T \right]$$
where $\ell_i$ is the batch size at step $i$.

C. Block-Based Batching Scheme

To ensure the estimator is fully online (recursive) and satisfies theoretical requirements, the authors introduce a Block-Based Batching Scheme:

• Batch Size Growth: The batch size $\ell_i$ must grow at a rate of $O(i^\alpha \log i)$.
• Implementation: They define a sequence of indices $\{a_m\}$ where a new block starts when the current block length reaches $\lfloor i^\alpha \log i \rfloor$.
• Recursive Update: The algorithm maintains the sum of the current block and the previous block, allowing $\hat{\Sigma}_n$ to be updated at each iteration with $O(d^2)$ computational cost (minimal for a $d \times d$ matrix).

3. Key Contributions

1. Novel De-biased Estimator: Introduced a fully online estimator that removes the bias inherent in traditional batch-means methods without requiring Hessian computations.
2. Improved Convergence Rate: Proved that the Mean Squared Error (MSE) of the estimator converges at a rate of $O(n^{-(1-\alpha)/2} \sqrt{\log n})$.
   • This is a significant improvement over existing Hessian-free methods (Batch-Means), which converge at $O(n^{-(1-\alpha)/4})$.
   • For $\alpha \approx 0.5$, the rate approaches $O(n^{-1/2} \sqrt{\log n})$, which is near-optimal.
3. Computational Efficiency: The method operates on a single SGD trajectory with an update cost of $O(d^2)$, making it scalable for high-dimensional problems where $O(d^3)$ Hessian-based methods are infeasible.
4. Theoretical Guarantees: Established consistency and tight upper bounds on the estimation error under general assumptions (strong convexity, Lipschitz gradients, bounded noise moments).

4. Results

Theoretical Results

• Theorem 7: Under standard assumptions (strong convexity, smoothness), the proposed estimator $\hat{\Sigma}_n$ converges to the true limiting covariance $\Sigma$ with an MSE bound of $O(n^{-1+\alpha} \log n)$.
• Tightness: The authors demonstrate via a mean estimation model (Theorem 3) that this rate is tight and cannot be generally improved for Hessian-free estimators.

Empirical Results

Numerical experiments were conducted on Linear, Logistic, and Expectile regression models with varying dimensions ($d=5, 20, 50$).

• Accuracy: The de-biased estimator consistently achieved lower Frobenius norm errors compared to the state-of-the-art Online Batch-Means (BM) estimator.
• Convergence: Log-log plots showed a steeper slope for the de-biased method, confirming the faster convergence rate.
• Confidence Intervals: When used to construct 95% confidence intervals for model parameters, the de-biased estimator yielded empirical coverage rates closer to the nominal 95% (e.g., ~93% vs. ~88% for Online BM in some settings), demonstrating superior reliability for statistical inference.

5. Significance

This work bridges a critical gap between computational efficiency and statistical rigor in online learning:

• Practical Impact: It enables practitioners to perform valid statistical inference (uncertainty quantification, confidence intervals) on models trained with SGD without the prohibitive cost of computing Hessians or running bootstraps.
• Theoretical Advancement: It resolves the long-standing issue of slow convergence in Hessian-free covariance estimation by introducing a principled bias-reduction technique tailored for the non-stationary, time-inhomogeneous nature of SGD iterates.
• Scalability: By maintaining $O(d^2)$ complexity, the method is applicable to modern high-dimensional machine learning tasks where traditional covariance estimation fails.

In summary, the paper provides a robust, efficient, and theoretically sound solution for estimating the asymptotic covariance of SGD, significantly enhancing the reliability of online statistical inference.