Iterative solvers: conjugate gradient and Krylov methods

Problem

The direct methods factor A once and then solve by substitution. For an n × n dense matrix that costs about n³/3 operations and n² numbers of memory. At n = 10⁵ that is 3 × 10¹⁴ operations and 80 GB in float64, and machine learning meets far bigger systems: the Hessian of a model with 10⁸ parameters has 10¹⁶ entries, so it can never be stored, let alone factored. Yet second-order optimizers, Gaussian processes, ridge regression on sparse text features and spectral methods all need to solve Ax = b or find a few eigenvectors of A.

What these matrices do allow is cheap products. A sparse matrix with a few nonzeros per row multiplies a vector in O(nnz) time. XᵀX + αI multiplies a vector as Xᵀ(Xv) + αv without ever being formed. And automatic differentiation computes a Hessian-vector product Hv for about the cost of two gradients. An iterative solver is an algorithm that needs nothing else.

Intuition

For a symmetric positive-definite A, solving Ax = b is the same as minimizing the quadratic bowl f(x) = ½xᵀAx − bᵀx, because its gradient is Ax − b. The residual r = b − Ax is the negative gradient, so the obvious method is gradient descent with the best step length along r. That is steepest descent, and on a stretched bowl it zig-zags: each exact line search ends where the new gradient is perpendicular to the old one, so it keeps undoing part of its previous progress.

Conjugate gradient (CG) fixes this by choosing each new direction so that minimizing along it never spoils the minimization already done along the earlier ones. The directions are not perpendicular but A-conjugate, p_iᵀAp_j = 0: perpendicular after the bowl is stretched back into a round one. In n dimensions there are only n such directions, so CG reaches the exact answer in at most n steps, and usually is close long before.

Both runs computed by steepest_descent and conjugate_gradient in the page's helpers. The first step is identical, since CG's first direction is the residual. After that, steepest descent turns 90° every step, while CG's second direction p₂ satisfies p₁ᵀAp₂ = 9e-16, zero up to rounding.

Mechanics

Matrix-free: only A·v

Every method on this page touches A through one operation, v ↦ Av, plus a handful of vector additions and dot products. Each iteration costs one product and O(n) extra work, and the memory is a few vectors of length n. Libraries make this explicit: SciPy's LinearOperator, which scikit-learn wraps around XᵀX + αI, and a closure over autograd for a Hessian. How fast the method converges depends on A's eigenvalues, measured by the condition number κ = λ_max/λ_min.

Steepest descent

From x_k with residual r_k = b − Ax_k, step along r_k by the α that minimizes f on that line: α_k = r_kᵀr_k / r_kᵀAr_k. Then x_k+1 = x_k + α_kr_k and r_k+1 = r_k − α_kAr_k, reusing the one product. The exact line search makes consecutive residuals orthogonal, which is the zig-zag. The error, measured in the A-norm ‖e‖_A = √(eᵀAe), shrinks by at least a factor (κ − 1)/(κ + 1) per step. At κ = 100 that is 0.98: about 115 steps per factor of 10.

Conjugate gradient

CG (Hestenes and Stiefel, 1952) keeps the residual but builds each search direction from it and the previous direction:

1. r₀ = b − Ax₀, p₀ = r₀.
2. α_k = r_kᵀr_k / p_kᵀAp_k; x_k+1 = x_k + α_kp_k; r_k+1 = r_k − α_kAp_k.
3. β_k = r_k+1ᵀr_k+1 / r_kᵀr_k; p_k+1 = r_k+1 + β_kp_k.

That short recurrence is enough to make every new direction A-conjugate to all earlier ones, not just the last, because A is symmetric. Two consequences follow. After k steps, x_k is the point that minimizes ‖x − x*‖_A over everything reachable from x₀ in the Krylov subspace span{r₀, Ar₀, A²r₀, …, A^k−1r₀}, the best answer any method could build from k products. And in exact arithmetic CG terminates in at most n steps, or in as many steps as A has distinct eigenvalues. The general error bound is

‖x_k − x*‖_A ≤ 2 ((√κ − 1)/(√κ + 1))^k ‖x₀ − x*‖_A,

so the number of steps grows with √κ instead of κ: at κ = 100 the factor is 9/11 ≈ 0.82, about 12 steps per factor of 10. In floating point the conjugacy slowly decays and CG can need more than n steps, but it still converges, and in practice it is stopped long before n.

Preconditioning

If κ is the problem, change the system. With a matrix M ≈ A that is cheap to invert, CG on M⁻¹A converges at the rate set by κ(M⁻¹A). The algorithm changes in one line: each new residual is passed through z = M⁻¹r before it is used to build the direction. The simplest choice is Jacobi, M = diag(A), which fixes badly scaled variables for free. Stronger ones are incomplete Cholesky (Cholesky that keeps A's sparsity pattern), multigrid, and in optimization the diagonal or Kronecker-factored curvature estimates that Adam, Shampoo and K-FAC maintain. A good preconditioner is the main lever: it can turn thousands of iterations into tens.

The Krylov family

CG is one member of a family that all search the same Krylov subspace and differ in what matrix they accept and what they minimize:

• MINRES minimizes ‖r‖ for symmetric matrices that are indefinite, where CG's bowl has no bottom. GMRES does the same for any square matrix, at the price of storing every direction.
• CGLS and LSQR solve least squares min ‖Ax − b‖ as CG on AᵀA without forming it or squaring the condition number in the arithmetic; LSQR's damp argument adds a ridge penalty.
• Lanczos runs the same three-term recurrence to build an orthonormal basis in which A is tridiagonal; the eigenvalues of that small tridiagonal matrix converge to A's extreme eigenvalues within a few dozen products. It is CG's twin: CG's α and β are the entries of that tridiagonal matrix.
• LOBPCG (Knyazev, 2001) finds a block of extreme eigenpairs by minimizing the Rayleigh quotient xᵀAx/xᵀx over the current vectors, their residuals and the previous vectors, accepts a preconditioner, and is what torch.lobpcg implements.

Where they appear in ML

• Ridge regression. scikit-learn's Ridge(solver="sparse_cg") wraps the product v ↦ Xᵀ(Xv) + αv in a LinearOperator and hands it to SciPy's CG (_solve_sparse_cg); with more features than samples it switches to the n_samples × n_samples kernel form XXᵀ + αI. solver="lsqr" calls LSQR with damp = √α (_solve_lsqr). Both keep a sparse X sparse.
• Newton-CG. LogisticRegression(solver="newton-cg") computes each Newton step by solving H s = −g approximately with an inner CG over Hessian-vector products (_cg), stopping when the residual's 1-norm drops below min(0.5, √‖g‖₁)·‖g‖₁ and bailing out on zero or negative curvature, then line-searches along s (_newton_cg). This is truncated Newton: loose solves far from the optimum, tighter ones near it.
• Hessian-free optimization. The same idea for neural networks (Martens, 2010): Hessian- or Gauss-Newton-vector products from autograd, a damped CG inner loop, no Hessian ever stored. PyTorch exposes the product as autograd.functional.hvp and its transpose vhp. Influence functions and implicit differentiation solve Hx = v the same way.
• Gaussian processes. Exact GP inference needs (K + σ²I)⁻¹y and log det. GPyTorch replaces the Cholesky with batched, preconditioned CG plus Lanczos for the log-determinant, which turns O(n³) into products a GPU does well (Gardner et al., 2018).
• Spectral methods. Spectral clustering, Laplacian eigenmaps and the top Hessian eigenvalues used to study sharpness all need only a few extreme eigenpairs: torch.lobpcg computes the k largest or smallest of a dense or sparse symmetric matrix.

The dashed line is the 10⁻⁸ target. CG's residual is not monotone, since CG minimizes the error in the A-norm, not ‖r‖; it can rise for a few steps and still be making progress.

Worked example

Implementation

def conjugate_gradient(matvec, b, x, tol=1e-8, max_steps=1000, inv_diag=None):
    dot = lambda u, v: sum(p * q for p, q in zip(u, v))
    r = [bi - ai for bi, ai in zip(b, matvec(x))]
    z = r if inv_diag is None else [m * ri for m, ri in zip(inv_diag, r)]
    p, rz, r0 = list(z), dot(r, z), dot(r, r) ** 0.5
    for _ in range(max_steps):
        if dot(r, r) ** 0.5 <= tol * r0:
            break
        ap = matvec(p)                      # the only use of A
        alpha = rz / dot(p, ap)
        x = [xi + alpha * pi for xi, pi in zip(x, p)]
        r = [ri - alpha * api for ri, api in zip(r, ap)]
        z = r if inv_diag is None else [m * ri for m, ri in zip(inv_diag, r)]
        rz, rz_old = dot(r, z), rz
        p = [zi + (rz / rz_old) * pi for zi, pi in zip(z, p)]
    return x

A = [[3., 2.], [2., 6.]]
conjugate_gradient(lambda v: [sum(a * x for a, x in zip(row, v)) for row in A], [2., -8.], [-2., -2.])
# [2.0, -2.0] after two steps

Matrix-free with PyTorch: a Newton step on a loss, using Hessian-vector products from autograd.

import torch
from torch.autograd.functional import hvp

def loss(w):
    return torch.nn.functional.softplus(-y * (X @ w)).mean() + 1e-3 * w.square().sum()

w = torch.zeros(X.shape[1], dtype=torch.float64)
g = torch.func.grad(loss)(w)
step = conjugate_gradient(lambda v: hvp(loss, w, torch.tensor(v, dtype=w.dtype))[1].tolist(),
                          (-g).tolist(), [0.0] * len(w), tol=1e-4)

# extreme eigenpairs of a large symmetric matrix, dense or sparse
E, V = torch.lobpcg(S, k=4, largest=True)    # S: [m, m] with m >= 3 * 4

# the same solvers in scikit-learn
from sklearn.linear_model import LogisticRegression, Ridge
Ridge(alpha=1.0, solver="sparse_cg").fit(X_sparse, y)    # CG on X^T X + alpha I, never formed
LogisticRegression(solver="newton-cg").fit(X, y)         # inner CG on Hessian-vector products

Tradeoffs

• Direct or iterative. Below a few thousand unknowns with a dense matrix, a Cholesky or LU factorization is faster, exact and reusable for many right-hand sides. Iterative methods win when the matrix is sparse, structured or implicit, when an approximate answer is enough, or when memory rules out n².
• Convergence depends on the spectrum. CG is fast when κ is small or the eigenvalues are clustered, and slow otherwise; there is no fixed cost. Budget for a preconditioner before tuning anything else.
• Symmetric positive-definite only. CG can fail silently on an indefinite or non-symmetric A: pᵀAp ≤ 0 breaks its step. Use MINRES or GMRES there, and in optimization treat negative curvature as a signal, as scikit-learn's capped CG does.
• Floating point. Conjugacy decays with rounding, so CG may need more than n steps and Lanczos produces spurious copies of converged eigenvalues unless it re-orthogonalizes. Low precision makes both worse; these solvers usually run in float32 or float64.
• Stopping rules. The residual is not the error: a small ‖r‖ can hide a large error when κ is large. Inside an outer optimizer, solve loosely early and tightly late (truncated Newton), since accuracy beyond the outer step's needs is wasted.

Check yourself

Connections

References

• Hestenes and Stiefel. Methods of Conjugate Gradients for Solving Linear Systems. Journal of Research of the National Bureau of Standards 49(6), 409–436, 1952.
• Shewchuk. An Introduction to the Conjugate Gradient Method Without the Agonizing Pain. Carnegie Mellon University technical report CMU-CS-94-125, 1994. The 2 × 2 example above is its running example.
• Saad. Iterative Methods for Sparse Linear Systems, 2nd ed. SIAM, 2003: chapters 6 (Krylov subspaces, CG, GMRES) and 9–10 (preconditioning).
• Paige and Saunders. LSQR: An Algorithm for Sparse Linear Equations and Sparse Least Squares. ACM Transactions on Mathematical Software 8(1), 43–71, 1982.
• Knyazev. Toward the Optimal Preconditioned Eigensolver: Locally Optimal Block Preconditioned Conjugate Gradient Method. SIAM Journal on Scientific Computing 23(2), 517–541, 2001.
• Pearlmutter. Fast Exact Multiplication by the Hessian. Neural Computation 6(1), 147–160, 1994.
• Martens. Deep Learning via Hessian-Free Optimization. ICML, 2010.
• Gardner, Pleiss, Bindel, Weinberger and Wilson. GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration. NeurIPS, 2018.
• Royer, O'Neill and Wright. A Newton-CG Algorithm with Complexity Guarantees for Smooth Unconstrained Optimization. 2018 (arXiv 1803.02924), the capped CG cited in scikit-learn's _cg.
• scikit-learn 1.9.1: _cg, _newton_cg, _solve_sparse_cg, _solve_lsqr.
• PyTorch v2.15.0: torch.lobpcg, _lobpcg, autograd.functional.hvp, vhp.