Acceleration of Randomized Kaczmarz Method via the Johnson-Lindenstrauss Lemma

Introduction

The Kaczmarz method is a popular algorithm for solving overdetermined consistent systems of linear equations. Due to its simplicity and speed, it has been used in a variety of applications ranging from tomography to digital signal processing . The method uses a series of alternating projections to iteratively converge to the solution of $Ax=b$, and is therefore computationally feasible even for very large systems. Given an initial guess $x_{0}$ and denoting by $a_{1},\ldots,a_{m}$ the rows of the $m\times n$ matrix $A$, each iteration of the method orthogonally projects the current estimation onto the next hyperplane defined as the solutions to $\langle a_{i},x\rangle=b_{i}$, chosen in a cyclic fashion. The algorithm can be described by the iterations:

where $x_{k}$ is the $k^{th}$ iterate, $b[i]$ (here and throughout) denotes the $i$th coordinate of $b$, and $i=(k$ mod $m)+1$.

Although this technique has been used in practice for quite some time, theoretical guarantees on convergence were difficult to obtain . It is clear that by design of the algorithm, the convergence rate depends on the ordering of the rows in $A$. Therefore, poorly ordered rows can lead to slower convergence. To overcome this difficulty, the rows of $A$ can be selected in a random fashion. It has been observed that this randomized version of the algorithm improves the convergence rate , however only recently have theoretical results been obtained .

In , Strohmer and Vershynin propose at each iteration to randomly select a row of $A$ with probability proportional to the Euclidean norm of the row. The randomized Kaczmarz (RK) method can thus be described by

where $p(i)$ takes values in $\{1,\ldots,m\}$ with probabilities $\frac{\|a_{p(i)}\|_{2}^{2}}{\|A\|_{F}^{2}}$. Here and throughout, $\|A\|_{F}$ denotes the Frobenius norm of $A$ and $\|\cdot\|_{2}$ denotes the standard Euclidean norm or spectral norm for vectors or matrices, respectively. The selection rule used here is not optimal in general. The motivation for setting the rule according to the weight of the row norms is two-fold. First, it allows for a guarantee of expected exponential convergence for the Kaczmarz method . Second, it is a computationally efficient strategy since often these values will be known approximately or exactly, and will only need to be computed once. A selection rule of this type is of course also related to the idea of preconditioning the matrix $A$ by scaling its rows. Although other diagonal preconditioners may certainly perform better in general, finding such an optimal preconditioner is itself an optimization problem of high complexity. With the above selection strategy, the following exponential bound was shown in for the convergence in expectation of this randomized method:

The empirical and theoretical benefits of this approach lead one to ask whether it is also accurate in the more realistic case when noise is present. One may thus consider the (now possibly inconsistent) system $Ax\approx b+w$ where $w$ is an arbitrary error vector that has been added to the consistent system $Ax=b$. It is shown in that in this case we have exponential convergence to the solution within an error factor:

where $R$ is the same as above and $\gamma=\max_{i}\frac{|w[i]|}{\|a_{i}\|_{2}}$. It is also shown that this bound is sharp and is attained even for simple examples .

2 Modified approach

Implementation and Runtime

Since the projections in the algorithm are orthogonal, one easily sees that the optimal projection in the $k$th iteration would be the one that maximizes $\|x_{k+1}-x_{k}\|_{2}$, or equivalently, the term

for all $s_{i},s_{j}\in S$, where $C$ is an absolute constant.

Remark. The value of $C$ in which this lemma holds depends on the distribution from which $\Phi$ is created. When $\Phi$ is Gaussian, one has $C\leq 8$ .

In our setting, the Johnson-Lindenstrauss Lemma allows us to project the rows of $A$ as well as the estimations $x_{k}$ onto a space of substantially lower dimension. This will then let us approximately calculate the terms in (2.1) using far fewer operations, which we can use to decide on which hyperplane to project the current estimate. We note that the projection of the rows of $A$ will be performed offline, adding to the preprocessing time, whereas the projection of the estimation $x_{k}$ will be done at each iteration. We choose to use an analagous strategy as in the RK algorithm for our row selection; that is, using the weight of the row norms. This choice yields a worst case convergence rate which is the same as that guaranteed by the original RK method (see Remark 2 below). This leads to the following modified randomized Kaczmarz method, called Randomized Kaczmarz via Johnson-Lindenstrauss (RKJL) which can be summarized as follows.

Randomized Kaczmarz via Johnson-Lindenstrauss (RKJL)

2. In the Test stage of the algorithm, we see that in addition to approximating the $n$ inner products, we also exactly calculate the inner product of a randomly selected row and also the row that was chosen. This will guarantee that the convergence is not slowed by any drastic consequences of the error in the approximations, and does not affect the overall runtime.

3. We note that the initialization step may of course be computationally expensive, as is calculating the probabilities $p(i)$ in the standard version (1.1). However, this need only be done once, and thus this version of the algorithm will be beneficial for situations in which the same matrix $A$ is used in many problems. This is the case for many applications such as the wave-scattering problem and structural mechanics problems ; see for others.

We next turn to an analysis of this modified method. In Section 4 we provide numerical results demonstrating the improved convergence rate.

2 Runtime

Analytical Justification

We next analyze how the Johnson-Lindenstrauss Lemma is utilized by our method. We will assume here that the system is real-valued and homogeneous (ie. $Ax=0$), that the rows of $A$ all have unit norm, and that the initial guess $x_{0}$ also satisfies $\|x_{0}\|_{2}\leq 1$. These assumptions are of course not necessary, but will make the analysis simpler. We discuss the case where the row norms may be far from equal in Remark 3 below. We begin with an easy lemma which shows that the geometry of the vectors used in the RKJL method is approximately preserved.

Similarly we have that $\gamma_{i}\geq\langle a_{i},x_{k}\rangle-2\delta$, which completes the claim.

This shows that the terms $\gamma_{i}$ used for selection in the algorithm are approximately equal to the actual desired values $\langle a_{i},x_{k}\rangle$. Thus for $\delta$ small, the RKJL algorithm makes well educated decisions at each iteration which allows for quicker convergence (see Theorem 3.2 below). This also shows that when the estimation $x_{k}$ becomes very close to the true solution $x=0$, the error $\delta$ begins to dominate and improvements may no longer be expected. However, this does not pose a problem since it only occurs when the estimate is already approximately $x$.

It is clear from construction of the RKJL algorithm (and especially in light of Remark 2 above), that convergence using RKJL is at least as fast as the standard randomized version. Moreover, when the error produced by applying $\Phi$ is small, the RKJL method will project onto the “best” hyperplane out of those it selected in that iteration. Since the probability of choosing this “best” row when selecting only a single row is strictly less than the probability of choosing that row when a set of rows is selected, this implies that the only case in which RKJL would not provide a strictly faster convergence rate is when $\langle a_{i},x_{k}\rangle=\langle a_{j},x_{k}\rangle$ for all rows $a_{i}$, $a_{j}$ selected in the $k$th iteration.

Given a current estimate $x_{k}$, one can explicitly compare the expectation of the improvement the next estimation provides, for both the RKJL and standard randomized methods. First observe that if $P$ denotes the projection in the $k$th iteration, then $x_{k-1}-x_{k}$ resides in the kernel of $P$, and is thus orthogonal to the space onto which $P$ projects. This space contains $x_{k}-x$ since $x$ is the solution to all equations $Ax=b$ and $x_{k}=Px_{k-1}$. Therefore, $x_{k}-x$ and $x_{k-1}-x_{k}$ are orthogonal, implying that

The relation (3.1) shows that the larger $\|x_{k}-x_{k-1}\|_{2}$, the bigger the improvement made in that iteration. We thus fix an estimation $x_{k}$ and analyze the expectation of $\|x_{k}-x_{k+1}\|_{2}$ for the RKJL method versus the standard one. For convenience, we again consider the real and homegenous case (ie. when $b=0$), and assume the rows of $A$ have unit norms. We then have the following result.

Fix an estimation $x_{k}$ and denote by $x_{k+1}$ and $x_{k+1}^{*}$ the next estimations using the RKJL and the standard RK method, respectively. Set $\gamma_{j}^{*}=|\langle{a_{j}},{x_{k}}\rangle|^{2}$ and reorder these so that $\gamma_{1}^{*}\geq\gamma_{2}^{*}\geq\ldots\geq\gamma_{m}^{*}$. Then when $d=C\delta^{-2}\log n$,

are non-negative values satisfying $\sum_{j=1}^{m}p_{j}=1$ and $p_{1}\geq p_{2}\geq\ldots\geq p_{m}=0$.

Proof. Since we assume the rows of $A$ have unit norm and that $b=0$, we see that $\gamma_{J}$ is precisely the value of $\|x_{k+1}-x_{k}\|_{2}^{2}$ if the algorithm were to select row $J$. We begin by examining the $k$th RKJL iteration.

Let $\mathcal{L}$ denote the set of rows chosen in the selection step, so that $|\mathcal{L}|=n$. If exact geometry were preserved (i.e. $\delta=0$), then the method would simply select the index of the largest $\lambda_{i}^{*}$ contained in $\mathcal{L}$. However, due to the error induced by $\Phi$, even when the “best” row is selected to be in $\mathcal{L}$, the algorithm may not choose this row for the projection.

We thus define sets $T_{j}$ for $j=1,\ldots m$, which consist of rows which could be “confused” in this way,

From Lemma 3.1, if $j\in\mathcal{L}$ and $1,\ldots,j-1\notin\mathcal{L}$, then the worst row we could choose is one that would give $\|x_{k+1}-x_{k}\|_{2}^{2}=\mu_{j}$. Therefore,

Since each row of $A$ has equal norm, each $row$ of $A$ is equally likely to be selected in $\mathcal{L}$. Thus

Lemma 3.1 and the definition of $T_{j}$ guarantee that $\mu_{j}\geq\max(\gamma_{j}^{*}-2\delta,0)$. This along with the fact that $\sum p_{j}=1$ yields

Finally, since all the rows of $A$ have the same norm, the standard RK method selects each row uniformly at random, so that

Combining (3.3) with (3.2) and (3.1) we have

Remarks. 1. Theorem 3.2 gives a lower bound, which shows improvements in the “worst case”, when the error induced by the Johnson-Lindenstrauss projection causes the method to choose a row of $A$ in the worst way. Numerical experiments (as seen in the next section) demonstrate substantial improvements in the convergence rate.

2. Note that since the sequences $\{\gamma_{j}^{*}\}$ and $\{p_{j}\}$ are non-increasing and $\sum_{j=1}^{m}p_{j}=1=\sum_{j=1}^{m}\frac{1}{m}$, the sum $\beta=\sum_{j=1}^{m}\Big(p_{j}-\frac{1}{m}\Big)\gamma_{j}^{*}$ is non-negative. Furthermore, $\beta=0$ only when $\gamma_{1}^{*}=\gamma_{2}^{*}=\ldots=\gamma_{m}^{*}$. Indeed, if $\gamma_{i}^{*}>\gamma_{j}^{*}$ even for just one pair $i<j$, then for $\delta$ small enough, the RKJL method provides strict convergence improvement. Knowledge of $x_{k}$ and $A$ would allow one to precisely calculate the improvement.

3. The theorem is proven under the assumption that the rows of $A$ have the same norm. The same argument holds (with different values of $p_{j}$) still showing improvement when this assumption does not hold, and numerical experiments show similar results in either case. Although the analysis of exact improvement in these other cases may quickly become quite complicated, we recall that by construction the RKJL method offers overall improvement in any case.

It is also helpful to identify some particularly interesting scenarios, for example, when one or a few rows are substantially of larger norm. In this case the standard RK algorithm (noticing that each row is selected independently of the previous selections) will choose this row repeatedly with high probability. Clearly this may slow convergence (especially when these rows are highly correlated), and although there is still guaranteed exponential convergence, $R$ in this case is much larger so that the guaranteed rate is slower. In RKJL, although the guaranteed worst case rate is the same, these large rows are likely not to be selected when their contributions toward the solution is minimal. This will again speed up convergence, since in the language of Theorem 3.2, the $\lambda^{*}_{k}$ corresponding to the rows which are highly correlated with the previous projection will be such that $k\approx m$. This means that the same argument as in Theorem 3.2 will hold for all $\lambda^{*}_{j}$ with $j\leq k\approx m$. Although this means the improvement in RKJL may not be as substantial as in the case of equal normed rows, we will still see a large improvement.

These ideas highlight the fact that the selection strategy is not optimal in general. For example, if there are $k\ll m$ rows which are highly uncorrelated but have very small norm, and $m-k\approx m$ equal rows with substantially larger norm, neither RK nor RKJL will perform well. Of course if this were reversed, with the uncorrelated rows having large norm and the equal rows having small norm, the selection strategy will yield excellent convergence in both RK and especially RKJL. We emphasize again that the choice of the selection rule is certainly not optimal in general, but is computationally efficient and allows for provable expected exponential rate of convergence. Choosing an optimal selection strategy for a given system is itself a problem of high complexity. If one is using RKJL and investing preprocessing time to perform dimension reduction, one may also wish to simply normalize the rows to avoid such difficulties.

Theorem 3.2 implies the following corollary, showing the improved convergence using RKJL when exact geometry is preserved (i.e. when $\delta\rightarrow 0$).

Fix an estimation $x_{k}$ and denote by $x_{k+1}$ and $x_{k+1}^{*}$ the next estimations using the RKJL and the standard method, respectively. Set $\gamma_{j}^{*}=|\langle{a_{j}},{x_{k}}\rangle|^{2}$ and reorder these so that $\gamma_{1}^{*}\geq\gamma_{2}^{*}\geq\ldots\geq\gamma_{m}^{*}$. Then when exact geometry is preserved ($\delta\rightarrow 0$),

Numerical Results

We now demonstrate improved convergence using the RKJL method. The first experiment we run is in the computationally infeasible situation where we do not use the Johnson-Lindenstrauss projection, but simply choose the best row out of the randomly selected $n$ rows. This experiment will demonstrate the improved convergence using RKJL when the error $\delta$ induced by $\Phi$ goes to $0$. This is the best improvement one can hope for in RKJL. When the problem sizes grow very large, the effect of the $\delta^{2}$ term in (2.2) becomes minimal, so it may be realistic to take $\delta$ quite small. For these simulations we use a $60000\times 1000$ matrix with Bernoulli entries and use a homogeneous system with an initial estimate chosen uniformly at random on the sphere. We see in Figure 1 that the convergence in this scenario is significantly improved.

This work is partially supported by the NSF DMS EMSW21-VIGRE grant and the Israel Science Foundation under Grant no. 1081/07. We would also like to thank Emmanuel Candès and Thomas Strohmer for helpful suggestions.

References