Randomized Kaczmarz: How Should You Sample?

The randomized Kaczmarz method is a method for solving systems of linear equations:

(1) $\text{Find $x$ such that } Ax = b.$

Throughout this post, the matrix $A$ will have dimensions $n\times d$ . Beginning from an initial iterate $x_0 = 0$ , randomized Kaczmarz works as follows. For $t = 0,1,2,\ldots$ :

Sample a random row index $i_t$ with probability $\prob \{ i_t = j \} = p_j$ .
Update to enforce the equation $a_{i_t}^\top x = b_{i_t}$ holds exactly:
$x_{t+1} \coloneqq x_t + \frac{b_{i_t} - a_{i_t}^\top x_t}{\norm{a_{i_t}}^2} a_{i_t}.$
Throughout this post, $a_j^\top$ denotes the $j$ th row of $A$ .

What selection probabilities $p_j$ should we use? The answer to this question may depend on whether the system (1) is consistent, i.e., whether it possesses a solution $x$ . For this post, we assume (1) is consistent; see this previous post for a discussion of the inconsistent case.

The classical selection probabilities for randomized Kaczmarz were proposed by Strohmer and Vershynin in their seminal paper:

(1) $p_j = \frac{\norm{a_j}^2}{\norm{A}_{\rm F}^2} \quad \text{for } j = 1,2,\ldots,n.$

Computing these selection probabilities requires a full pass over the matrix, which can be expensive for the largest problems.¹ A computationally appealing alternative is to implement randomized Kaczmarz with uniform selection probabilities

(2) $p_j = \frac{1}{n} \quad \text{for } j = 1,2,\ldots,n.$

Ignoring computational cost, which sampling rule leads to faster convergence: (1) or (2)?

Surprisingly, to me at least, the simpler strategy (2) often works better than (1). Here is a simple example. Define a matrix $A \in \real^{20\times 20}$ with entries $A_{ij} = \min(i,j)^2$ , and choose the right-hand side $b\in\real^{20}$ with standard Gaussian random entries. The convergence of standard RK with sampling rule (1) and uniform RK with sampling rule (2) is shown in the plot below. After a million iterations, the difference in final accuracy is dramatic: the final relative error 0.00012 was uniform RK and 0.67 for standard RK!

Error for randomized Kaczmarz with both squared row norm sampling ("standard") and uniformly random rows on matrix with entries min(i,j)^2. Uniform randomized Kaczmarz achieves significantly smaller final error

In fairness, uniform RK does not always outperform standard RK. If we change the matrix entries to $A_{ij} = \min(i,j)$ , then the performance of both methods is similar, with both methods ending with a relative error of about 0.07.

Error for randomized Kaczmarz with both squared row norm sampling ("standard") and uniformly random rows on matrix with entries min(i,j). Uniform randomized Kaczmarz and standard randomized Kaczmarz achieve comparable final errors

Another experiment, presented in section 4.1 of Strohmer and Vershynin’s original paper, provides an example where standard RK converges a bit more than twice as fast as uniform RK (called “simple RK” in their paper). Still, taken all together, these experiments demonstrate that standard RK (sampling probabilities (1)) is often not dramatically better than uniform RK (sampling probabilities (2)), and uniform RK can be much better than standard RK.

Randomized Kaczmarz Error Bounds

Why does uniform RK often outperform standard RK? To answer these questions, let’s look at the error bounds for the RK method.

The classical analysis of standard RK shows the method is geometrically convergent

(3) $\expect\left[ \norm{x_t - x_\star}^2 \right] \le (1 - \kappa_{\rm dem}(A)^{-2})^t \norm{x_\star}^2.$

Here,

(4) $\kappa_{\rm dem}(A) = \frac{\norm{A}_{\rm F}}{\sigma_{\rm min}(A)} = \sqrt{\sum_i \left(\frac{\sigma_i(A)}{\sigma_{\rm min}(A)}\right)^2}$

is known as the Demmel condition number and $\sigma_i(A)$ are the singular values of $A$ . Recall that we have assumed the system $Ax = b$ is consistent, possessing a solution $x_\star$ . If there are multiple solutions, we let $x_\star$ denote the solution of minimum norm.

What about uniform RK? Let $D_A = \diag ( \norm{a_i}^{-1} : i=1,\ldots,n )$ denote a diagonal matrix containing the inverse row norms, and introduce the row-equilibrated matrix $D_A A$ . The row-equilibrated matrix $D_A A$ has been obtained from $A$ by rescaling each of its rows to have unit norm.

Uniform RK can then be related to standard RK run on the row-equilibrated matrix:

Fact (uniform sampling and row equilibration): Uniform RK on the system $Ax = b$ produces the same sequence of (random) iterates $\hat{x}_t$ as standard RK applied to the row-equilibrated system $(D_A A)x = D_A b$ .

Therefore, by (3), the iterates $\hat{x}_t$ of uniform RK satisfy the bound

(5) $\expect\left[ \norm{\hat{x}_t - x_\star}^2 \right] \le (1 - \kappa_{\rm dem}(D_A A)^{-2})^t \norm{x_\star}^2.$

Thus, at least using the error bounds (3) and (5), whether standard or uniform RK is better depends on which matrix has a smaller Demmel condition number: $A$ or $D_A A$ .

Row Equilibration and the Condition Number

Does row equilibration increase or decrease its condition number? What is the optimal way of scaling the rows of a matrix to minimize its condition number? These are classical questions in numerical linear algebra, and they were addressed in a classical 1969 paper of van der Sluis. These results were then summarized and generalized in Higham’s delightful monograph Accuracy and Stability of Numerical Algorithms. Here, we present answers to these questions using a variant of van der Sluis’ argument.

First, let’s introduce some more concepts and notation. Define the spectral norm condition number

$\kappa(A) \coloneqq \frac{\sigma_{\rm max}(A)}{\sigma_{\rm min}(A)}$

The spectral norm and Demmel condition numbers are always comparable $\kappa(A) \le \kappa_{\rm dem}(A)\le \sqrt{\min(n,d)}\cdot \kappa(A)$ . Also, let $\mathrm{Diag}$ denote the set of all (nonsingular) diagonal matrices.

Our first result shows us that row equilibration never hurts the Demmel condition number by much. In fact, the row-equilibrated matrix produces a nearly optimal Demmel condition number when compared to any row scaling:

Theorem 1 (Row equilibration is a nearly optimal row scaling). Let $A\in\real^{n\times d}$ be wide $n\le d$ and full-rank, and let $D_AA$ denote the row-scaling of $A$ to have unit row norms. Then
$\kappa_{\rm dem}(D_AA) \le \sqrt{n}\cdot \min_{D \in \mathrm{Diag}} \kappa (DA) \le \sqrt{n}\cdot \min_{D \in \mathrm{Diag}} \kappa_{\rm dem} (DA).$

By scaling the rows of a square or wide matrix to have unit norm, we bring the Demmel condition number to within a $\sqrt{n}$ factor of the optimal row scaling. In fact, we even bring the Demmel condition number to within a $\sqrt{n}$ factor of the optimal spectral norm condition number for any row scaling.

Since the convergence rate for randomized Kaczmarz is $\kappa_{\rm dem}(A)^{-2}$ , this result shows that implementing randomized Kaczmarz with uniform sampling yields to a convergence rate that is within a factor of $n$ of the optimal convergence rate using any possible sampling distribution.

This result shows us that row equilibration can’t hurt the Demmel condition number by much. But can it help? The following proposition shows that it can help a lot for some problems.

Proposition 2 (Row equilibration can help a lot). Let $A\in\real^{n\times d}$ be wide $n\le d$ and full-rank, and let $\gamma$ denote the maximum ratio between two row norms:
$\gamma \coloneqq \frac{ \max_i \norm{a_i}}{\min_i \norm{a_i}}.$
Then the Demmel condition number of the original matrix $A$ satisfies
$\kappa_{\rm dem}(A) \le \gamma \cdot \kappa_{\rm dem}(D_A A).$
Moreover, for every $\gamma\ge 1$ , there exists a matrix $A_\gamma$ where this bound is nearly attained:
$\kappa_{\rm dem}(A_\gamma) \ge \sqrt{1-\frac{1}{n}} \cdot \gamma \cdot \kappa_{\rm dem}(D_{A_\gamma}A_\gamma).$

Taken together, Theorem 1 and Proposition 2 show that row equilibration often improves the Demmel condition number, and never increases it by that much. Consequently, uniform RK often converges faster than standard RK for square (or short wide) linear systems, and it never converges much slower.

Proof of Theorem 1

We follow Higham’s approach. Each of the $n$ rows of $D_AA$ each have unit norm, so

(7) $\norm{D_AA}_{\rm F} = \sqrt{n}.$

The minimum singular value of $D_A A$ can be written in terms of the Moore–Penrose pseudoinverse $(D_A A)^\dagger = A^\dagger D_A^{-1}$ as follows

$\frac{1}{\sigma_{\rm min}(D_A A)} = \norm{A^\dagger D_A^{-1}}.$

Here, $\norm{\cdot}$ denotes the spectral norm. Then for any nonsingular diagonal matrix $D$ , we have

(8) $\frac{1}{\sigma_{\rm min}(D_A A)} = \norm{A^\dagger D^{-1} (DD_A^{-1})} \le \norm{A^\dagger D^{-1}} \norm{DD_A^{-1}} = \frac{\norm{DD_A^{-1}}}{\sigma_{\rm min}(DA)}.$

Since the matrix $DD_A^{-1}$ is diagonal its spectral norm is

$\norm{DD_A^{-1}} = \max \left\{ \frac{|D_{ii}|}{|(D_A)_{ii}|} : 1\le i \le n \right\}.$

The diagonal entries of $D_A$ are $\norm{a_i}^{-1}$ , so

$\norm{DD_A^{-1}} = \max \left\{ |D_{ii}|\norm{a_i} : 1\le i\le n \right\}$

is the maximum row norm of the scaled matrix $DA$ . The maximum row norm is always less than the largest singular value of $DA$ , so $\norm{DD_A^{-1}} \le \sigma_{\rm max}(DA)$ . Therefore, combining this result, (7), and (9), we obtain

$\kappa_{\rm dem}(D_AA) \le \sqrt{n} \cdot \frac{\sigma_{\rm max}(DA)}{\sigma_{\rm min}(DA)} = \sqrt{n}\cdot \kappa (DA).$

Since this bound holds for every $D \in \mathrm{Diag}$ , we are free to minimize over $D$ , leading to the first inequality in the theorem:

$\kappa_{\rm dem}(D_AA) \le \sqrt{n}\cdot \min_{D \in \mathrm{Diag}} \kappa (DA).$

Since the spectral norm condition number is smaller than the Demmel condition number, we obtain the second bound in the theorem.

Proof of Proposition 2

Write $A = D_A^{-1}(D_AA)$ . Using the Moore–Penrose pseudoinverse again, write

(10) $\kappa_{\rm dem}(A) = \norm{D_A^{-1}(D_AA)}_{\rm F} \norm{(D_A A)^\dagger D_A}.$

The Frobenius norm and spectral norm satisfy a (mixed) submultiplicative property

$\norm{BC}_{\rm F} \le \norm{B}\norm{C}_{\rm F}, \quad \norm{BC} \le\norm{B}\norm{C}.$

Applying this result to (1), we obtain

$\kappa_{\rm dem}(A) \le \norm{D_A^{-1}}\norm{D_AA}_{\rm F} \norm{(D_A A)^\dagger} \norm{D_A}.$

We recognize $\gamma = \norm{D_A^{-1}}\norm{D_A}$ and $\kappa_{\rm dem}(D_A A) = \norm{D_AA}_{\rm F} \norm{(D_A A)^\dagger}$ . We conclude

$\kappa_{\rm dem}(A) \le \gamma \cdot \kappa_{\rm dem}(D_A A).$

To show this bound is nearly obtained, introduce $A_\gamma = \diag(\gamma,\gamma,\ldots,\gamma,1)$ . Then $D_{A_\gamma} A_\gamma = I$ with $\kappa_{\rm dem}(D_{A_\gamma}A_{\gamma}) = \sqrt{n}$ and

$\kappa_{\rm dem}(A_\gamma) = \frac{\norm{A_{\gamma}}_{\rm F}}{\sigma_{\rm min}(A_\gamma)} = \frac{\sqrt{(n-1)\gamma^2+1}}{1} \ge \sqrt{n} \cdot \sqrt{1-\frac{1}{n}} \cdot \gamma.$

Therefore,

$\kappa_{\rm dem}(A_\gamma) \ge \sqrt{1-\frac{1}{n}} \cdot \gamma \cdot \kappa_{\rm dem}(D_{A_\gamma}A_\gamma).$

Practical Guidance

What does this theory mean for practice? Ultimately, single-row randomized Kaczmarz is often not the best algorithm for the job for ordinary square (or short–wide) linear systems, anyway—block Kaczmarz or (preconditioned) Krylov methods have been faster in my experience. But, supposing that we have locked in (single-row) randomized Kaczmarz as our algorithm, how should we implement it?

This question is hard to answer, because there are examples where standard RK and uniform RK both converge faster than the other. Theorem 1 suggests uniform RK can require as many as $n\times$ more iterations than standard RK on a worst-case example, which can be a big difference for large problems. But, particularly for badly row-scaled problems, Proposition 2 shows that uniform RK can dramatically outcompete standard RK. Ultimately, I would give two answers.

First, if the matrix $A$ has already been carefully designed to be well-conditioned and computing the row norms is not computationally burdensome, then standard RK may be worth the effort. Despite this theory suggesting it can do quite badly, it took a bit of effort to construct a simple example of a “bad” matrix where uniform RK significantly outcompeted standard RK. (On most examples I constructed, the rate of convergence of the two methods were similar.)

Second, particularly for the largest systems where you only want to make a small number of total passes over the matrix, expending a full pass over the matrix to compute the row norms is a significant expense. And, for poorly row-scaled matrices, sampling using the squared row norms can hurt the convergence rate. Based on these observations, given a matrix of unknown row scaling and conditioning or given a small budget of passes over the matrix, I would use the uniform RK method over the standard RK method.

Finally, let me again emphasize that the theoretical results Theorem 1 and Proposition 2 only apply to square or wide matrices $A$ . Uniform RK also appears to work for consistent systems with a tall matrix, but I am unaware of a theoretical result comparing the Demmel condition numbers of $D_AA$ and $A$ that applies to tall matrices. And for inconsistent systems of equations, it’s a whole different story.

Edit: After initial publication of this post, Mark Schmidt shared that the observation that uniform RK can outperform standard RK was made nearly ten years ago in section 4.2 of the following paper. They support this observation with a different mathematical justification

Randomized Kaczmarz Error Bounds

Row Equilibration and the Condition Number

Proof of Theorem 1

Proof of Proposition 2

Practical Guidance

Leave a Reply Cancel reply