Krylov subspace iteration

1 Power iteration

The simplest possible variant.

• Power iteration on Wikipedia

2 Lanczos decomposition

Lanczos decomposition is a handy approximation for matrices that are cheap to multiply because of some structure, but expensive to store. It can also be used to calculate an approximate inverse cheaply.

I learnt this trick from Gaussian process literature in the context of Lanczos Variance Estimates (LOVE) (Pleiss et al. 2018), although I believe it exists elsewhere.

Given some rank $k$ and an arbitrary starting vector $\mathit{b}$, the Lanczos algorithm iteratively approximates $\mathrm{K}\in {\mathbb{R}}^{n\times n}$ by a low-rank factorisation $\mathrm{K}\approx \mathrm{Q}\mathrm{T}{\mathrm{Q}}^{\mathrm{\top }}$, where $\mathrm{T}\in {\mathbb{R}}^{k\times k}$ is tridiagonal and $\mathrm{Q}\in {\mathbb{R}}^{n\times k}$ has orthogonal columns. Crucially, we do not need to form $\mathrm{K}$ to evaluate matrix-vector products $\mathrm{K}\mathit{b}$ for an arbitrary vector $\mathit{b}$. Moreover, with a given Lanczos approximand $\mathrm{Q},\mathrm{T}$ we may estimate $$\begin{array}{r}{\mathrm{K}}^{-1}\mathit{c}\approx \mathrm{Q}{\mathrm{T}}^{-1}{\mathrm{Q}}^{\mathrm{\top }}\mathit{c}.\end{array}$$ even for $\mathit{b}\ne \mathit{c}$. Say we wish to calculate ${(\mathrm{Z}{\mathrm{Z}}^{\mathrm{\top }}+{\sigma }^2\mathrm{I})}^{-1}\mathrm{B}$, with $\mathrm{Z}\in {\mathbb{R}}^{D\times N}$ and $N\ll D$.

We approximate the solution to this linear system using the partial Lanczos decomposition starting with probe vector $\mathit{b}=\bar{\mathrm{B}}$ and $\mathrm{K}=(\mathrm{Z}{\mathrm{Z}}^{\mathrm{\top }}+{\sigma }^2\mathrm{I})$.

This requires $k$ matrix-vector products of the form $$\begin{array}{r}\underset{\mathcal{O}(D^2)}{\underbrace{(\underset{\mathcal{O}(N{D}^2)}{\underbrace{\mathrm{Z}{\mathrm{Z}}^{\mathrm{\top }}}}+{\sigma }^2\mathrm{I})\mathit{b}}}=\underset{\mathcal{O}(ND)}{\underbrace{\mathrm{Z}\underset{\mathcal{O}(ND)}{\underbrace{({\mathrm{Z}}^{\mathrm{\top }}\mathit{b})}}}}+{\sigma }^2\mathit{b}.\end{array}$$ Using the latter representation, the required matrix-vector product may be found with a time complexity cost of $\mathcal{O}(ND)$. Space complexity is also $\mathcal{O}(ND)$. The output of the Lanczos decomposition is $\mathrm{Q},\mathrm{T}$ such that $(\mathrm{Z}{\mathrm{Z}}^{\mathrm{\top }}+{\sigma }^2\mathrm{I})\mathit{b}\approx \mathrm{Q}\mathrm{T}{\mathrm{Q}}^{\mathrm{\top }}\mathit{b}$. Then the solution to the inverse-matrix-vector product may be approximated by ${(\mathrm{Z}{\mathrm{Z}}^{\mathrm{\top }}+{\sigma }^2\mathrm{I})}^{-1}\mathrm{B}\approx \mathrm{Q}{\mathrm{T}}^{-1}{\mathrm{Q}}^{\mathrm{\top }}\mathrm{B}$. requiring the solution in $\mathrm{X}$ of the much smaller linear system $\mathrm{X}\mathrm{T}=\mathrm{Q}$. Exploiting the positive-definiteness of $\mathrm{T}$ we may use the Cholesky decomposition of $\mathrm{T}={\mathrm{L}}^{\mathrm{\top }}\mathrm{L}$ for a constant speedup over solving an arbitrary linear system. The time cost of the solution is $\mathcal{O}(D{k}^3)$, for an overall cost to the matrix inversions of $\mathcal{O}(NDk+D{k}^3)$.

Lifehack: Find derivatives of Lanczos iterations via pnkraemer/matfree: Matrix-free linear algebra in JAX. (Krämer et al. 2024).

3 Incoming

• Computing Stationary Distributions of a Discrete Markov Chain

4 References