Rajat's painless Conjugate Gradients

A geometric derivation of Conjugate Gradients

Author: Rajat Vadiraj Dwaraknath

Thanks to Anjan Dwaraknath for helpful discussions.

Problem setup §

We wish to solve the linear system $A{x}^{\ast }=b$ where $A\in {\mathbb{R}}^{n\times n}$ is a symmetric positive definite matrix. We use $x^{\ast }$ to denote the exact solution.

We are working in the computational model where we have access to the matrix $A$ only through matrix-vector products. That is, we have a method to compute $Av$ for any vector $v\in {\mathbb{R}}^n$. We also have access to the vector $b$.

Since $A$ is symmetric positive definite, it induces an inner product given by $⟨u,v{⟩}_A=u^{T}Av$. Although we don’t have access to the solution $x^{\ast }$, we do have access to $b=A{x}^{\ast }$. So, we can compute the expression $v^{T}b=v^{T}A{x}^{\ast }=⟨v,x^{\ast }{⟩}_A$ for any $v$. In other words, we can compute A-inner products $⟨x^{\ast },v{⟩}_A$ with the solution for any $v$.

Solution attempt §

Motivated by this, we can posit a method to solve $A{x}^{\ast }=b$ by working in the $A$-inner product as follows:

• Compute an $A$-orthogonal basis for ${\mathbb{R}}^n$ which we denote $\left\{p_1,p_2,\dots ,p_n\right\}$ via Gram-Schmidt in the $A$-inner product. Note that for this CG derivation, we don’t need these basis vectors to be normalized.
• $A$-project the solution $x^{\ast }$ using this basis:

$$x^{\ast }=\sum _{i=1}^n\frac{⟨x^{\ast },p_i{⟩}_A}{⟨p_i,p_i{⟩}_A}{p}_i$$

This also naturally leads to an approximation scheme by truncating the sum in the projection:

$$x_k:=\sum _{i=1}^k\frac{⟨x^{\ast },p_i{⟩}_A}{⟨p_i,p_i{⟩}_A}{p}_i$$

Therefore, the sequence of approximations $x_k$ can be interpreted as the $A$-projection of the solution $x^{\ast }$ onto to the sequence of increasing subspaces given by $\text{span}(p_1,\dots ,p_k)$. We can iteratively update the approximations by noticing that:

$$\begin{array}{rl}{x}_{k+1}& =\sum _{i=1}^{k+1}\frac{⟨x^{\ast },p_i{⟩}_A}{⟨p_i,p_i{⟩}_A}{p}_i\\ & =\sum _{i=1}^k\left(\frac{⟨x^{\ast },p_i{⟩}_A}{⟨p_i,p_i{⟩}_A}{p}_i\right)+\frac{⟨x^{\ast },p_{k+1}{⟩}_A}{⟨p_{k+1},p_{k+1}{⟩}_A}{p}_{k+1}\\ & =x_k+\frac{⟨x^{\ast },p_{k+1}{⟩}_A}{⟨p_{k+1},p_{k+1}{⟩}_A}{p}_{k+1}\end{array}$$

Since the approximations are projections, we can also use the variational characterization of projection as finding the vector in the subspace that is closest to $x^{\ast }$ in the $A$-norm:

$$x_k={\text{argmin}}_{x\in \text{span}(p_1,\dots ,p_k)}\Vert x-x^{\ast }{\Vert }_A$$

Notice that there is some freedom in the choice of $\{p_i\}$ in this method. We use this freedom in a specific way to arrive at the Conjugate Gradients method.

Connecting to Conjugate Gradients §

The conjugate gradients method does exactly the above procedure, but for a very specific choice of the orthogonal basis $\left\{p_1,p_2,\dots ,p_n\right\}$. Specifically, it requires that these vectors span the Krylov sequence of $A$ with starting vector $b$. More precisely, CG chooses $\{p_i\}$ such that

$$\text{span}(p_1,\dots ,p_k)=\mathcal{K}(A,b,k):={\mathcal{K}}_k\text{ for all }k$$

With this choice, the variational characterization of the successive approximations becomes:

$$x_k={\text{argmin}}_{x\in {\mathcal{K}}_k}\Vert x-x^{\ast }{\Vert }_A$$

which is exactly the starting definition of CG!

What remains now is to find an efficient way of computing the $A$-orthogonal basis $\{p_i\}$. It turns out that choosing the successive approximation subspaces to be the Krylov sequence allows us to compute $p_i$ using a short recurrence by connecting to the Lanczos iteration.

The three-term recurrence for $p_i$ §

To compute the $A$-orthogonal basis $\{p_i\}$, we can perform Gram-Schmidt in the $A$-inner product on some vectors $v_1,\dots ,v_n$ that span the Krylov sequence. However, Gram-Schmidt is pretty slow since to compute $p_k$ we need to $A$-project out the components of $v_k$ along $p_1,\dots ,p_{k-1}$ and each projection needs $O(\text{nnz}(A))$ time since we need to compute an $A$-inner product which requires multiplication by $A$. So the total time to compute the basis $\{p_1,\dots ,p_k\}$ is $O(\text{nnz}(A)k^2)$. It would be nice if we only needed to project out a few components instead of all $k$ at each step. We can achieve this with a smart choice of starting vectors $v_i$.

A good choice of starting vectors $v_i$ would be one where they are already close to being $A$-orthogonal. Putting the vectors into a matrix $V_k:=[v_1\dots v_k]$, we want the $A$-inner product of $V_k$ with its transpose to be close to a diagonal matrix. More precisely, we want $V_k^{T}A{V}_k$ to be close to a diagonal matrix.

We have seen such a $V_k$ in the context of the Lanczos iteration. Specifically, the vectors $q_1,\dots ,q_k$ generated by Lanczos span the Krylov sequence and also have the property that $Q_k^{T}A{Q}_k$ is a tridiagonal matrix. This means that

$$q_i^{T}A{q}_j=0\text{ for }|i-j|>1⟹q_i{\perp }_A{q}_j\text{ for }|i-j|>1$$

In other words, $q_k{\perp }_A{K}_{k-2}$ for all $i$.

Therefore, we choose $v_k=q_k$ for all $k$. When performing $A$-Gram Schmidt on $q_k$, we only need to project out the component of $q_k$ along $p_{k-1}$ since $q_k$ is already $A$-orthogonal to $p_{k-2},\dots ,p_1$. Therefore, we can write the Gram-Schmidt step as follows:

$$p_k=q_k-\frac{⟨q_k,p_{k-1}{⟩}_A}{⟨p_{k-1},p_{k-1}{⟩}_A}{p}_{k-1}$$

This step only requires a $O(\text{nnz}(A))$ compute since we are only doing one $A$-projection. Therefore, the total time to compute $p_1,\dots ,p_k$ is $O(\text{nnz}(A)k)$ which is much faster than the previous $O(\text{nnz}(A)k^2)$ time. Note that this includes the time to run the Lanczos iteration to generate $q_k$ since that also takes $O(\text{nnz}(A)k)$ time.

Therefore, the total time to compute the approximation $x_k$ is also $O(\text{nnz}(A)k)$ since we can iteratively update the approximations to obtain $x_k$ as described before, and each of these updates also only takes $O(\text{nnz}(A))$ time.

Bringing in the residuals §

This version of CG might seem a bit different than the usual implementation since there is no mention of the residual vectors $r_k:=b-A{x}_k$. We can easily bring these into the picture by noticing an important fact: the residuals $r_k$ are scaled versions of the vectors $q_{k+1}$ generated by Lanczos. More precisely, $r_k$ is parallel to $q_{k+1}$ for all $0\le k\le n-1$.

Note that there is an off-by-one on the indices between $r$ and $q$ simply because in Lanczos $q_1=b/\Vert b{\Vert }_2$ but the corresponding residual is $r_0=b-A\cdot 0=b$.

We prove this statement by showing that the $r_k$ form an orthogonal basis for the Krylov sequence and then use the fact that this orthogonal basis must be unique up to scaling to get the result.

First, observe that $r_k\in K_{k+1}$. This is because $r_k=b-A{x}_k$ and $b\in {\mathcal{K}}_1$ and

$$x_k\in {\mathcal{K}}_k⟹A{x}_k\in {\mathcal{K}}_{k+1}.$$

Now, we can rewrite the residual as

$$r_k=A(x^{\ast }-x_k)=A{e}_k$$

where $e_k=x^{\ast }-x_k$ is the error in the approximation $x_k$. Now, since $x_k$ is the $A$-projection of $x^{\ast }$ onto the subspace ${\mathcal{K}}_k$, we know by property of the projection that the error in this projection must be $A$-orthogonal to ${\mathcal{K}}_k$. That is,

$$e_k{\perp }_A{\mathcal{K}}_k⟹A{e}_k\perp {\mathcal{K}}_k⟹r_k\perp {\mathcal{K}}_k$$

Combining this with $r_k\in {\mathcal{K}}_{k+1}$ for all $k$ gives us:

$$r_i\perp r_j\text{ for }j<i$$

But the orthogonality relation is symmetric, so we can extend the result to:

$$r_i\perp r_j\text{ for }j\ne i$$

Therefore, $r_i$ is an orthogonal set of vectors and

$$\text{span}(r_0,\dots ,r_{k-1})={\mathcal{K}}_k.$$

This means that $r_k$ must be parallel to $q_{k+1}$ from Lanczos since we found an orthogonal basis for the Krylov sequence, and we can show by induction that this basis is unique up to scaling factors. We have therefore shown the required statement.

If we put the residuals into a matrix as follows: $R_k:=[r_0\dots r_{k-1}]$, the above result says that $R_k^{T}A{R}_k$ is a tridiagonal matrix.

Now, we can instead use $v_k=r_{k-1}$ for all $k$ as our starting vectors for Gram-Schmidt in the $A$-inner product when finding $p_k$. The short recurrence for $p_k$ can now be written as:

$$p_k=r_{k-1}-\frac{⟨r_{k-1},p_{k-1}{⟩}_A}{⟨p_{k-1},p_{k-1}{⟩}_A}{p}_{k-1}$$

Additionally, we can find $r_{k+1}$ either as $b-A{x}_{k+1}$, or by using the update for $x_{k+1}$ in terms of $x_k$ to get a recurrence as well:

$$\begin{array}{rl}{r}_{k+1}& =b-A{x}_{k+1}\\ & =b-A{x}_k-A\frac{⟨p_{k+1},x^{\ast }{⟩}_A}{⟨p_{k+1},p_{k+1}{⟩}_A}{p}_{k+1}\\ & =r_k-\frac{⟨p_{k+1},x^{\ast }{⟩}_A}{⟨p_{k+1},p_{k+1}{⟩}_A}A{p}_{k+1}\end{array}$$

Putting it all together §

We can now combine the short recurrence for $p_k$ with the iterative update for $x_k$ and $r_k$ to get a more familiar version of the Conjugate Gradients method (we replace $k$ with $k+1$ and also write the $A$-inner products explicitly).

$$\begin{array}{rl}{p}_{k+1}& =r_k-\frac{p_k^{T}A{r}_k}{p_k^{T}A{p}_k}{p}_k\\ x_{k+1}& =x_k+\frac{p_{k+1}^{T}b}{p_{k+1}^{T}A{p}_{k+1}}{p}_{k+1}\\ r_{k+1}& =r_k-\frac{p_{k+1}^{T}b}{p_{k+1}^{T}A{p}_{k+1}}A{p}_{k+1}\end{array}$$

with $x_0=0$, $r_0=b$, $p_1=b$.

Summary §

• We want to solve $A{x}^{\ast }=b$ by using only matrix-vector products. Since we can compute inner products with $b$, we can compute $A$-inner products with $x^{\ast }$.
• Suppose we have an A-orthogonal basis $p_1,\dots ,p_k$ for the Krylov sequence starting with $b$.
• We can obtain approximate solution $x_k$ by $A$-projecting $x^{\ast }$ onto $\text{span}(p_1,\dots ,p_k)={\mathcal{K}}_k$.
• To compute the basis $p_1,\dots ,p_k$, do Gram-Schmidt in the $A$-inner product with residuals $r_k$ as the starting vectors.
• This choice leads to a short recurrence for $p_k$ resulting in an efficient algorithm for computing $x_k$.