[Paper Review] A Stochastic Newton Algorithm for Distributed Convex Optimization

Brian Bullins, Kumar Kshitij Patel+3 more|arXiv (Cornell University)|Oct 7, 2021

Stochastic Gradient Optimization Techniques44 references4 citations

TL;DR

This paper proposes FedSN, a stochastic Newton algorithm for distributed convex optimization that leverages stochastic Hessian-vector products and gradients to reduce communication rounds. It achieves faster convergence with fewer communications than first-order methods by solving Newton steps via one-shot averaging of quadratic approximations, with theoretical guarantees for quasi-self-concordant objectives like logistic regression.

ABSTRACT

We propose and analyze a stochastic Newton algorithm for homogeneous distributed stochastic convex optimization, where each machine can calculate stochastic gradients of the same population objective, as well as stochastic Hessian-vector products (products of an independent unbiased estimator of the Hessian of the population objective with arbitrary vectors), with many such stochastic computations performed between rounds of communication. We show that our method can reduce the number, and frequency, of required communication rounds compared to existing methods without hurting performance, by proving convergence guarantees for quasi-self-concordant objectives (e.g., logistic regression), alongside empirical evidence.

Motivation & Objective

• Address the challenge of high communication frequency in distributed optimization, especially for large-scale machine learning with millions of parameters.
• Reduce the number of communication rounds in distributed stochastic optimization without sacrificing convergence performance.
• Leverage second-order information (Hessian-vector products) to accelerate convergence beyond first-order methods like Local SGD.
• Design a practical algorithm that maintains computational efficiency by avoiding full Hessian computation while enabling scalable, distributed optimization.
• Provide theoretical convergence guarantees for quasi-self-concordant objectives, such as logistic regression, under the intermittent communication model.

Proposed method

• Propose Federated-Stochastic-Newton (FedSN), a distributed algorithm that approximates Newton steps using one-shot averaging of a quadratic subproblem.
• Reformulate the Newton step as minimizing a convex quadratic: $\min_{\Delta x} \frac{1}{2}\Delta x^{\top}\nabla^{2}F(x)\Delta x + \nabla F(x)^{\top}\Delta x$, which is solved via distributed averaging.
• Use stochastic gradient oracles and stochastic Hessian-vector product oracles as inputs, both computationally efficient for generalized linear models.
• Enable each machine to perform $K$ independent stochastic computations per round before communicating, reducing communication frequency.
• Apply one-shot averaging to solve the quadratic subproblem in a single round, avoiding iterative communication during subproblem solution.
• Ensure computational efficiency by exploiting the fact that Hessian-vector products for logistic regression are rank-one operations, computable in $\mathcal{O}(d)$ time per sample.

Experimental results

Research questions

• RQ1 Can second-order information be efficiently leveraged in a distributed, stochastic setting to reduce communication rounds?
• RQ2 Does using stochastic Hessian-vector products lead to faster convergence than first-order methods in the intermittent communication model?
• RQ3 Can the Newton step be approximated effectively via one-shot averaging without incurring high computational cost?
• RQ4 How does the convergence rate of FedSN compare to Local SGD and FedAc for quasi-self-concordant objectives?
• RQ5 What is the computational overhead of Hessian-vector product oracles relative to stochastic gradients in practice?

Key findings

• FedSN achieves exponential convergence decay in the optimization error, with a rate of $\frac{HB^2}{KR} + \frac{\sigma B}{\sqrt{MK}} + \frac{\rho B^2}{\sqrt{K}R}$, where $H$, $\sigma$, and $\rho$ are smoothness, gradient variance, and Hessian-vector product variance, respectively.
• Empirical results show that FedSN reduces the number of communication rounds required to reach a given accuracy compared to Local SGD and FedAc, especially in low-communication regimes.
• For logistic regression, the computational cost of Hessian-vector products is asymptotically similar to that of stochastic gradients, enabling efficient implementation.
• Runtime measurements show that FedSN-Lite is only slightly slower than Local SGD (6.67×10⁻⁵ sec/step vs. 6.39×10⁻⁵ sec/step), with minimal overhead.
• The algorithm maintains high accuracy even with a single round of communication, demonstrating the effectiveness of one-shot averaging for Newton steps.
• Theoretical guarantees are established for quasi-self-concordant objectives, including logistic regression, under the intermittent communication model with $M$ machines and $R$ communication rounds.