Online Newton method speeds up and improves statistical inference

Inference for Newton Methods with Accelerated Sketch-and-Project via Random Scaling

Machine Learning

Summary

Solving complex math problems quickly is often slowed down by certain faster but costly steps. This paper studies a new way to speed up these steps using a smart sketching technique with acceleration, making computations faster. The authors also show how to use this to reliably understand the uncertainty in their solution as it progresses, without extra heavy calculations. Their method can give more accurate answers and helps make decisions confidently using these faster techniques.

What this means in practice

  • For machine learning engineers: Improve training speed and model uncertainty estimation in large-scale machine learning models using faster Newton-like methods with accelerated sketching.
  • For data science teams: Implement efficient online confidence assessments when fitting complex models without costly covariance calculations via a randomized rescaling test.

Authors

Xinchen Du, Elizaveta Rebrova, Michał Dereziński, Sen Na

Abstract

We study an online sketched Newton method that approximates the Newton direction at each step via a state-of-the-art sketching solver, called the generalized accelerated sketch-and-project solver (GAS), thereby mitigating the computational bottleneck of classical second-order methods. The GAS solver improves upon vanilla, unaccelerated sketch-and-project solvers by achieving accelerated convergence through Nesterov momentum updates, and accommodates a flexible projection metric whose proper choice further reduces computational cost. Building on this design, we establish asymptotic normality of the averaged sketched Newton iterates and characterize their limiting covariance matrix. The resulting covariance recovers that of the unaccelerated sketched Newton method under a specific choice of acceleration parameters, converges more rapidly (in the number of sketching steps) to the minimax-optimal covariance in general, and is smaller than that of the last iterate produced by the accelerated method. Finally, we strengthen these results by establishing a functional central limit theorem for the Newton iterates, which allows us to bypass explicit covariance estimation and develop an online inference procedure based on random scaling. Specifically, we construct a pivotal test statistic by appropriately rescaling the averaged iterates, so that its limiting distribution is free of any unknown parameters, enabling asymptotically valid online inference. Numerical experiments demonstrate superior performance of the proposed inference procedure.