The Conjugate Gradient Method for Linear Equations

One of the most common tasks in numerical algorithms is to solve a linear equation—that is, find $x$ for which $$Ax=b$$ for a given matrix $A$ and vector $b$. This can be solved via Gaussian elimination, which generally has a high runtime ($O(n^3)$). We will show how to improve upon this using optimization ideas. This is one instance of a problem for which we can find an approximate solution much faster using calculus tools, compared with using a close-form exact algebraic solution.

We will make the following assumption: $A$ is symmetric and positive definite. This means $A$ has eigenvalues $0<\lambda_{1}\le \ldots \le \lambda_{n}$. In particular, $A$ is square, $A\in \mathbb{R}^{n\times n}$, and $x,b\in\mathbb{R}^{n}$.

As an Optimization Problem

Solving $Ax=b$ is equivalent to minimizing the “distance” $Ax-b$. However, $Ax-b$ is a vector, so it cannot be minimized directly. To transform $Ax-b$ into a meaningful scalar, several approaches are possible. We could minimize $\left\lVert Ax-b \right\rVert_{2}$ (the least-squares problem), or any other norm for that matter. Instead, we choose to take the inner product with $x$, leading to the quadratic form: $$f(x):=\frac{1}{2}x^{\top}Ax -x^{\top}b.$$ The gradient of $f$ is $\nabla f(x)=Ax-b$, while the Hessian matrix is $\nabla^2 f(x)=A$. It is a general fact that quadratic functions with a positive-definite Hessian matrix are strongly convex. We thus conclude:

Proof. As noted above, $\nabla f(x)=Ax-b$; hence $Ax=b$ iff $\nabla f(x)=0$. We know $\nabla f(x)=0$ is a necessary condition for $x$ to be a minimizer. However, it is also sufficient because $\nabla^{2}f=A$ is positive definite, implying $\nabla f$ can have at most a single root. That root must be a minimizer, as the curvature of $f$ implies there are no saddle points. $\blacksquare$

Since $f$ is convex, we can use a general solver, like Gradient Descent (GD), to find a minimizer. GD requires computing gradients, which in our case is just $\nabla f(x)=Ax-b$. Therefore, each GD iteration requires time $T_{A}$, which is the time needed to compute the matrix-vector product $Ax$ (this dominates the computation). Hopefully, if the number of iterations is small, we need a small number of such products. Put explicitly, each iteration in the GD method is: $$x_{t+1}=x_{t}- \eta_{t}\cdot \nabla f(x_{t})$$ where $\eta_{t}$ is called the step size at time $t$. The following theorem is a general result about the rate of convergence for gradient descent:

To understand the meaning of the bounds, note that $A^{-1}b=x^*$, and so $x-A^{-1}b=x-x^*$, i.e., the difference between the solution and the optimal solution. Moreover, note that $$ \begin{aligned} (x-A^{-1}b)^{\top}A(x-A^{-1}b)&= x^{\top} Ax - x^{\top}AA^{-1}b - b^{\top}(A^{-1})^{\top}Ax + b^{\top}(A^{-1})^{\top}AA^{-1}b \\ &= x^{\top}Ax- 2x^{\top} b+ b^{\top}x^* \\ &= 2f(x)-2f(x^*) \end{aligned} $$ where we used the fact $A^{-1}$ must also by symmetric and that $f(x^*)=-\frac{1}{2}b^{\top}x^{*}$. In other words, the error bound in the theorem is just (divided by $2$): $$(f(x)-f(x^*))\le \varepsilon^2\cdot f(x^*)$$ A proof can be found in [5].

Viewing GD with Krylov Subspaces

}_

Suppose we initiate $x_{0}=0$. Then $$\nabla f(x_{0})=Ax_{0}-b=-b.$$ Consequently, $$x_{1}=x_{0}-\eta_{0}\nabla f(x_{0})=\eta_{0}b,$$ and similarly, $$x_{2}=x_{1}-\eta_{1}(Ax_{1}-b)=\eta_{0}\cdot b-\eta_{1}\eta_{0}\cdot Ab+\eta_{1}\cdot b=(\eta_{0}+\eta_{1})b-\eta_{1}\eta_{0} Ab.$$ Therefore, a simple inductive argument shows that $$x_{t}\in \mathrm{Span}\set{b,Ab,\ldots ,A^{t-1}b}.$$ Denote the $t$-th order Krylov subspace for $A,b$ by $\mathcal{K}_{t}$. Then we see GD moves $x_{1},\ldots,x_{t}$ in $\mathcal{K}_{1},\ldots,\mathcal{K}_{t}$. In other words, $x_i\in \mathcal{K}_i$.

However, GD does not guarantee that $x_{t}$ is the minimizer of $f$ in $\mathcal{K}_{t}$. The idea of the Conjugate Gradient (CG) method is, at each step, to minimize over $\mathcal{K}_{t}$. In particular, $x^{*}$, the true minimizer, is the minimizer over $\mathcal{K}_{n}$.

Finding the Minimizer in a Krylov Subspace

Question. How can we quickly find the minimizer of $f$ in $\mathcal{K}_{t}$?

To answer this question, we start with the following observation:

Suppose we were given a basis $\set{p_{0},\ldots,p_{t-1}}$ for $\mathcal{K}_{t}$, with the following separability property: $$f\left(\sum_{i=0}^{t-1} \beta_{i}p_{i}\right)=\sum_{i=0}^{t-1}f(\beta_{i}p_{i}),$$ for every $\beta_{0},\ldots,\beta_{t-1}$. Then, to minimize $f$ over $\mathcal{K}_{t}$, it suffices to find $\alpha_{i}$ for all $i$, where $\alpha_i$ is the scalar minimizing $f(\alpha p_{i})$. In particular, when moving from $\mathcal{K}_{t-1}$ to $\mathcal{K}_{t}$, to minimize $f$ over $\mathcal{K}_{t}$ we only need to add $\alpha_{t-1}p_{t-1}$ to the previous minimizer.

This separability property clearly holds for linear functions, so for our objective function $f(x)=\frac{1}{2}x^{\top}Ax+b^{\top}x$ (ignoring the constant), only the quadratic part $x^{\top}Ax$ might be the problem. Note that the left hand side is $$\left(\sum_{i=0}^{t-1}\beta_{i}p_{i}\right)^{\top}A\left(\sum_{i=0}^{t-1}\beta_{i}p_{i}\right)=\sum_{0\le i,j\le t-1}\beta_{i}\beta_{j} p_{i}^{\top}Ap_{j}$$ while the right hand side is $\sum_{i=0}^{t-1}(\beta_{i}p_{i})^{\top}A(\beta_{i}p_{i})=\sum_{i} \beta_{i}^{2}p_{i}^{\top}Ap_{i}$. Hence, for equality to hold, we need the basis $\set{p_{0},\ldots,p_{t-1}}$ to be $A$-orthonormal:

Computing an $A$-Orthonormal Basis

Gram-Schmidt Iterations

We wish to iteratively compute an $A$-orthonormal basis $\set{p_{0},\ldots,p_{t-1}}$ for $\mathcal{K}_{t}$. One way to do this is via Gram-Schmidt iterations:

1. Given a new vector $v\notin \mathcal{K}_{t}$ (specifically $v=A^{t}b$),
2. Define $$p_{t}=v-\sum_{i\le t-1}\frac{v^{\top}Ap_{i}}{p_{i}^{\top}Ap_{i}}p_{i}.$$

Since $p_{i}^{\top}Ap_{j}=0$ when $i\not=j$, we obtain that $$p_{t}^{\top}Ap_{j}=v^{\top}A p_{j}-\sum_{i\le t-1}\frac{v^{\top}Ap_{i}}{p_{i}^{\top}Ap_{i}}p_{i}^{\top}Ap_{j}=v^{\top}Ap_{j}-\frac{v^{\top}Ap_{j}}{\cancel{p_{j}^{\top}Ap_{j}}}\cdot \cancel{ p_{j}^{\top}Ap_{j}}=0$$ for every $j\le t-1$. However, to compute $p_{t}$ in this manner we need $O(t)$ matrix-vector products. Moreover, we haven’t really used the fact that we are computing a basis for the Krylov subspace generated by $A$, which gives us extra structure to work with.

A Slight Adjustment

The ida: If we start with $p_0=b$ and choose $v$ to be $Ap_{t-1}$, the Gram-Schmidt iterations simplify and are faster to compute, involving only a constant number of matrix-vector products.

Formally, let $p_{0}=b$ be the first vector, and assume we have constructed an $A$-orthonormal basis for $\mathcal{K}_{t}$, denoted by $\set{p_{0},\ldots,p_{t-1}}$, where the induction hypothesis states $$\mathcal{K}_{i}=\mathrm{Span}\set{p_{0},\ldots,p_{i-1}}$$ for every $i=1,\ldots,t-1$. In particular, $Ap_{i-1}\in \mathcal{K}_{i+1}$. We want to construct a vector $p_{t}$ so that the statement above is true for $i=t+1$.

As discussed in the observation above, if $Ap_{t-1}\in \mathcal{K}_{t}$, then $\mathcal{K}_{t}=\mathcal{K}_{n}$ and so we have reached the end of the iteration. Therefore we may assume without loss of generality that $Ap_{t-1}\notin \mathcal{K}_{t-1}$. Define $$p_{t}=A p_{t-1}-\sum_{i\le t-1}\frac{(Ap_{t-1})^{\top} Ap_{i}}{p_{i}^{\top}Ap_{i}}p_{i}.$$

By definition, $Ap_{t-1}\in \mathcal{K}_{t+1}$, and also $\mathcal{K}_{t+1}=\mathrm{Span}\set{p_{0},\ldots,p_{t}}$ (using the induction hypothesis). By the induction hypothesis, for every $0\le j<t-1$ it holds that $Ap_{j}\in \mathcal{K}_{j+2}$. This implies $Ap_j$ is a linear combination of $p_{0},\ldots,p_{j+1}$, and in particular there are scalars $c_{0},\ldots,c_{j+1}$ such that $Ap_{j}=\sum_{i=0}^{j+1}c_{i}p_{i}$. Hence: $$(Ap_{t-1})^{\top}Ap_{j}=p_{t-1}A^{\top}Ap_{j}=p_{t-1}A^{\top}\left(\sum_{i=0}^{j+1}c_{i}p_{i}\right)=\sum_{i=0}^{j+1}c_{i}\cdot p_{t-1}^{\top}Ap_{i}.$$ Due to $A$-orthogonality, $p_{t-1}^{\top}Ap_{i} = 0$ unless $i = t-1$. Thus: $$(Ap_{t-1})^{\top}Ap_{j} = \begin{cases} 0 & j<t-2, \\ c_{t-1}p_{t-1}^{\top}Ap_{t-1} & j=t-2. \end{cases}$$ Therefore, the definition of $p_{t}$ is essentially a sum of three elements only. In other words, the definition simplifies to a three-term recurrence: $$p_{t}=Ap_{t-1}-\frac{p_{t-1}^{\top}A^{2}p_{t-1}}{p_{t-1}^{\top}Ap_{t-1}}p_{t-1}-\frac{p_{t-1}^{\top}A^{2}p_{t-2}}{p_{t-2}^{\top}Ap_{t-2}}p_{t-2},$$ which implies we need a constant number of matrix-vector products to compute the next element of the basis.

The Conjugate Gradient Algorithm

The choice of $\alpha_{t}$ is done to minimize $f(\alpha p_{t})=\frac{1}{2}\alpha^{2}p_{t}^{\top}Ap_{t}- \alpha b^{\top}p_{t}$, viewed as a function $\mathbb{R}\to \mathbb{R}$. The number of iterations needed for a certain prescribed precision is therefore dependent on the convergence of the conjugate gradient algorithm (which is always better than the gradient descent algorithm).

Polynomial Minimization

Let $x^{*}$ denote the minimizer, satisfying $Ax^{*}=b$. Then $$f(x^{*})=\frac{1}{2}(x^{*})^{\top}Ax^{*}-(x^{*})^{\top}Ax^{*}=-\frac{1}{2}(x^{*})^{\top}Ax^{*}.$$ Thus at iteration $t$ we have $$f(x_{t})-f(x^{*})=\frac{1}{2}(x_{t}^{\top}Ax_{t}+(x^{*})^{\top}Ax^{*}-2x_{t}^{\top}Ax^{*})=\frac{1}{2}(x_{t}-x^{*})^{\top}A(x_{t}-x^{*}),$$ using the symmetry of $A$. This is the derivation as in the error of the gradient descent theorem.

We note that evaluating a polynomial at a matrix gives a new matrix which is the weighted sum of matrix powers. Here $A^{0}=I$.

Let $q(X)=p(X)\cdot X-1$. Then $q(A)=(p(A)A-I)$ is a degree $t$ polynomial. Moreover, there is a bijection between $$\mathcal{P}_{t} =\set{P(X)\mid \deg P\le t-1}\leftrightarrow \set{Q(X)\mid \deg Q\le t, Q(0)=-1}=\mathcal{Q}_{t}.$$ This bijection is given by the transformation of $P\mapsto Q$, i.e., a polynomial $P$ is mapped to $Q(X)=X\cdot P(X)-1$.

Another way to view polynomials of matrices is through their diagonalization. Since $A$ is symmetric, it is orthogonally diagonalizable, meaning that $A=U \Lambda U^{\top}$ for a diagonal matrix $\Lambda$ (with the eigenvalues on the diagonal) and $U$ an orthogonal matrix ($UU^{\top}=U^{\top}U=I$). In this case, $$A^{2}=(U\Lambda U^{\top})^{2}=U \Lambda U^{\top}U\Lambda U^{\top}=U \Lambda^{2}U^{\top}.$$ Since $\Lambda$ is diagonal with $\set{\lambda_{1},\ldots,\lambda_{n}}$, $\Lambda^{2}$ is just the element-wise product of the diagonals. By induction, this is true for every $k$. Therefore, $$p(A)=\sum_{i}\gamma_{i}A^{i}=\sum_{i}\gamma_{i}U \mathrm{diag}(\lambda_{1}^{i},\ldots,\lambda_{n}^{i}) U^{\top}=U\left(\sum_{i}\gamma_{i} \mathrm{diag}(\lambda_{1}^{i},\ldots, \lambda_{n}^{i})\right)U^{\top}$$ which is equal to $U \mathrm{diag}(p(\lambda_{1}),\ldots,p(\lambda_{n}))U^\top$. In words, we apply the polynomial on the eigenvalues to obtain a new diagonal matrix.

Proof. Write $M=U\Lambda U^{\top}$ as above. Note that diagonal matrices commute, hence $$v^{\top}(p(M))^{\top}Mp(M)v=v^{\top}Up(\Lambda)U^{\top}U\Lambda U^{\top}Up(\Lambda)U^{\top}v=v^{\top}U(\Lambda\cdot p^{2}(\Lambda ))U^{\top}v.$$ Writing $v=\sum_{i}c_{i}u_{i}$, where $u_{i}$ are the columns of $U$ (the eigenvectors basis), we have $$v^{\top}p(M)^{\top}Mp(M)v=\sum_{i,j}c_{i}c_{j}u_{i}^{\top}U(\Lambda\cdot p^{2}(\Lambda))U^{\top}u_{j}=\sum_{i}c_{i}^{2} \lambda_{i}p(\lambda_{i})^{2},$$ where we used the orthogonality. On the other hand, $v^{\top}Mv=\sum_{i}c_{i}^{2}\lambda_{i}$. The claim follows because all the eigenvalues are non-negative. $\blacksquare$

Therefore: $$\left\lVert x_{t}-x^{*} \right\rVert_{A}^{2}=\min_{q(X)\in \mathcal{Q}_{t}}\left\lVert q(A)x^{*} \right\rVert_{A}^{2}\le \min_{q(X)\in \mathcal{Q}_{t}}\max_{i=1}^{n}\left|q(\lambda_{i})\right|^{2}\cdot \left\lVert x^{*} \right\rVert_{A}^{2}.$$ More specifically, we have shown that (by extending the set on which we maximize): $$\left\lVert x_{t}-x^{*} \right\rVert_{A}^{2}\le \left\lVert x^{*} \right\rVert_{A}^{2}\cdot \min_{q(X)\in \mathcal{Q}_{t}}\max_{\lambda\in [\lambda_{1},\lambda_{n}]}\left|q(\lambda)\right|^{2}.$$ In particular, every polynomial $q \in \mathcal{Q}_t$ gives an upper bound. For example, take $$q(X)=-\left(1-\frac{2X}{\lambda_{1}+\lambda_{n}}\right)^{t},$$ and note $q(0)=-1$ so $q(X)\in \mathcal{Q}_{t}$, while $|q(X)|^{2}$ is maximized over $[\lambda_{1},\lambda_{n}]$ when $\lambda=\lambda_{1}$. Thus $$q(\lambda_{1})^{2}=\left(1-\frac{2\lambda_{1}}{\lambda_{1}+\lambda_{n}}\right)^{2t}=\left(\frac{\lambda_{n}-\lambda_{1}}{\lambda_{1}+\lambda_{n}}\right)^{2t}=\left(\frac{\kappa-1}{\kappa+1}\right)^{2t},$$ which leads to a slightly better convergence rate compared to gradient descent.

Polynomial Minimax Evaluation

We have essentially arrived at the minimax problem of polynomial evaluation: $$\min_{q(X)\in \mathcal{Q}_{t}}\max_{\lambda\in [\lambda_{1},\lambda_{n}]}\left|q(\lambda)\right|^{2}.$$ The solution to a more general version of this problem is given by the Chebyshev polynomials, a central result in approximation theory.

The minimization property of the Chebyshev polynomials is:

For proof and background, see [6].

Thus, Chebyshev polynomials offer good candidates to obtain the upper bound for the convergence rate. In particular, note that minimizing the square magnitude is equivalent to minimizing the magnitude. All we have to deal with is mapping polynomials evaluated on $[\lambda_{1},\lambda_{n}]$ to $[-1,1]$. This is done easily by mapping an interval $$[a,b]\to [-1,1]\quad \text{via}\quad x\mapsto \frac{a+b-2x}{b-a}.$$ Note that $a\mapsto \frac{b-a}{b-a}=1$ while $b\mapsto \frac{a-b}{b-a}=-1$, and by linearity, every $a<x<b$ is mapped inside the interval. Define $$Q_{a,b,n}(x)=\frac{T_{n}\left(\frac{a+b-2x}{b-a}\right)}{T_{n}\left(\frac{a+b}{b-a}\right)}.$$ It is a degree $n$ polynomial which satisfies $$Q_{a,b,n}(0)=\frac{T_{n}(\frac{a+b}{b-a})}{T_{n}(\frac{a+b}{b-a})}=1.$$ Therefore $Q_{a,b,t} \in \mathcal{Q}_{t}$ (up to a sign flip, which does not affect magnitude). Using the fact the numerator has magnitude $\le 1$ on the interval gives us the estimate $$|Q_{\lambda_{1},\lambda_{n},t}(x)|\le \frac{1}{T_{t}(\frac{\lambda_{1}+\lambda_{n}}{\lambda_{n}-\lambda_{1}})}=T_{t}\left(\frac{\lambda_{n}/\lambda_{1}+1}{\lambda_{n}/\lambda_{1}-1}\right)^{-1}.$$ By a few more manipulations, we see that $$T_{t}\left(\frac{\kappa+1}{\kappa-1}\right)\ge \frac{1}{2}\left(\frac{\sqrt{\kappa}+1}{\sqrt{\kappa}-1}\right)^{t},$$ hence $$\left|Q_{\lambda_{1},\lambda_{n},t}(\lambda)\right|\le 2 \cdot \left(\frac{\sqrt{\kappa}-1}{\sqrt{\kappa}+1}\right)^{t}.$$ This proves the following theorem (to obtain the theorem, a few more algebraic steps are required, similar to the omitted proof of the GD theorem), which is essentially a quadratic improvement over gradient descent ($\sqrt{\kappa(A)}$ instead of $\kappa(A)$):

Clustered Eigenvalues

Consider the following scenario: $\lambda_{n}$ is very large, $\lambda_{1}$ is very small, and most of the eigenvalues lie concentrated in the interval $[a,b]$ where $a$ and $b$ are moderately sized. The condition number $\lambda_n/\lambda_1$ is enormous, so the convergence bound derived above may be poor. However, it turns out that by a nice algebraic trick, the number of iterations required turns out to be small.

Proof. Let $a=\lambda_1$. Using the polynomials $Q_{a,b,t}$ we can define a composite polynomial: $$q(x)=Q_{a,b,t}(x)\cdot \prod_{i=1}^{c} \left(1 - \frac{x}{\mu_i}\right)$$ where $\mu_{1},\ldots,\mu_{c}$ are the “outlier” eigenvalues outside of $[a,b]$. Note that $q(0)=1 \cdot \prod 1 = 1$, so $q \in \mathcal{Q}_{t+c}$. Furthermore:

1. For any outlier eigenvalue $\mu_i$, the term $(1 - \mu_i/\mu_i) = 0$, so $q(\mu_i)=0$.
2. For any $\lambda \in [a,b]$, the term $|1 - \lambda/\mu_i| \le 1$.
3. The term $|Q_{a,b,t}(\lambda)|$ is bounded by the Chebyshev decay rate determined by the ratio $b/a$ (the “effective” condition number).

Thus, the error of the algorithm is at most $\max_{j}\left|q(\lambda_{j})\right| \le \max_{\lambda\in [a,b]}\left|Q_{a,b,t}(\lambda)\right|$. Setting $t=O(\sqrt{b/a}\cdot \log(1/\varepsilon))$, we obtain $\varepsilon$-error. In particular, $q(x)$ has degree $t+c$, so we need $t$ iterations for the bulk spectrum plus $c$ iterations to “kill” the outliers. $\blacksquare$

This lemma can be extended to outliers outside of $[a,b]$ for general $a$ (i.e., $a>\lambda_1$). In this case one has to deal with the fact $(1-x/\mu)$ may be large in magnitude for small $\mu$. This might add a logarithmic (in the condition number) number of iterations.

References

1. Hestenes, M. R., & Stiefel, E. (1952), Methods of Conjugate Gradients for Solving Linear Systems .
2. Shewchuk, J. R. (1994), An Introduction to the Conjugate Gradient Method Without the Agonizing Pain .
3. Trefethen, L. N., & Bau III, D. (1997), Numerical Linear Algebra .
4. Nocedal, J., & Wright, S. (2006), Numerical Optimization .
5. Vishnoi, N.K. (2013), $Lx=b$ Laplacian Solvers and Their Algorithmic Applications .
6. Mudde, M.H. (2017), Chebyshev approximation .