SGD/Kaczmarz

Analysis¶

The key result of Strohmer & Vershynin, 2009 is that randomized Kaczmarz converges linearly in expectation, at a rate governed by the ratio $\smin(\vec{A})^2 / \|\vec{A}\|_\F^2$.

Several remarks are in order:

• The quantity $\|\vec{A}\|_\F^2 / \smin(\vec{A})^2$ is sometimes called the Demmel condition number (squared) of $\vec{A}$. Since $\|\vec{A}\|_\F^2 \leq d \, \smax(\vec{A})^2$, Theorem 7.3 implies that

iterations suffice to guarantee $\EE[\|\vec{x}_t - \vec{x}^*\|^2] \leq \varepsilon \|\vec{x}_0 - \vec{x}^*\|^2$. Each iteration costs $O(d)$, so the total cost is $O(d^2 \cond(\vec{A})^2 \log(1/\varepsilon))$ operations — potentially sublinear in the size $nd$ of $\vec{A}$ when $n \gg d$.

• The rate is independent of $n$; the number of rows enters only through how the sampling probabilities are defined.

• The proof reveals why row-norm sampling is the natural choice: the $\|\vec{a}_s\|^2$ appearing in the sampling probability exactly cancels the $\|\vec{a}_s\|^{-2}$ in the projection, so the expected per-step progress becomes $\|\vec{A}\vec{e}\|^2 / \|\vec{A}\|_\F^2$, which can be bounded using the singular values of $\vec{A}$. This is the same cancellation that made row-norm sampling effective for approximate matrix multiplication.

Connection to stochastic gradient descent¶

Randomized Kaczmarz is closely related to stochastic gradient descent (SGD), the workhorse algorithm of modern machine learning. Consider the least-squares objective, written as a sum over the rows of $\vec{A}$:

Computing the full gradient $\nabla f(\vec{x})$ requires reading all of $\vec{A}$. Instead, SGD samples a random index $s$ from some distribution $(p_1, \ldots, p_n)$ and takes a step in the direction of an unbiased estimate of the gradient:

where $\eta > 0$ is a step size. Indeed, $\EE[ p_s^{-1} \nabla f_s(\vec{x}) ] = \sum_j \nabla f_j(\vec{x}) = \nabla f(\vec{x})$, regardless of the choice of sampling distribution.

Comparing (7.7) with (7.6), we see that randomized Kaczmarz is exactly SGD with row-norm importance sampling $p_j = \|\vec{a}_j\|^2 / \|\vec{A}\|_\F^2$ and the fixed step size $\eta = 1 / \|\vec{A}\|_\F^2$. This perspective, developed in Needell et al., 2016, explains several features of the method. For consistent systems, each $f_j$ is minimized (with value zero) at the solution $\vec{x}^*$, so the stochastic gradients vanish at $\vec{x}^*$ and a constant step size yields linear convergence. For inconsistent systems (i.e. genuine least-squares problems), the stochastic gradients do not vanish at the minimizer, and randomized Kaczmarz with a fixed step size only converges to within a ball around the least-squares solution whose radius depends on the residual; decaying step sizes or tail averaging are needed to converge to the minimizer itself.