Statistical Properties of BayesCG under the Krylov prior

Matrices are represented in bold uppercase, such as $\mathbf{A}$; vectors in bold lowercase, such as $\mathbf{b}$; and scalars in lowercase, such as $m$.

The $m\times m$ identity matrix is $\mathbf{I}_{m}$, or just $\mathbf{I}$ if the dimension is clear. The Moore-Penrose inverse of a matrix $\mathbf{A}$ is $\mathbf{A}^{\dagger}$, and the matrix square root is $\mathbf{A}^{1/2}$ (Higham08, Chapter 6).

Probability distributions are represented in lowercase Greek, such as $\mu_{m}$; and random variables in uppercase, such as $X$. The random variable $X$ having distribution $\mu$ is represented by $X\sim\mu$, and its expectation by $\operatorname{\mathbb{E}}[X]$.

The Gaussian distribution with mean $\mathbf{x}\in\mathbb{R}^{n}$ and covariance $\mathbf{\Sigma}\in\mathbb{R}^{n\times n}$ is denoted by $\operatorname{\mathcal{N}}(\mathbf{x},\mathbf{\Sigma})$, and the chi-squared distribution with $f$ degrees of freedom by $\chi_{f}^{2}$.

We review the computation of posterior distributions for BayesCG under general priors (Theorem 2.1), and present a pseudo code for BayesCG (Algorithm 2.1).

Given an initial guess $\mathbf{x}_{0}$, BayesCG Cockayne:BCG solves symmetric positive definite linear systems (1) by computing iterates $\mathbf{x}_{m}$ that converge to the solution $\mathbf{x}_{*}$. In addition, BayesCG computes probability distributions that quantify the uncertainty about the solution at each iteration $m$. Specifically, for a user-specified Gaussian prior $\mu_{0}\equiv\operatorname{\mathcal{N}}(\mathbf{x}_{0},\mathbf{\Sigma}_{0})$, BayesCG computes posterior distributions $\mu_{m}\equiv\operatorname{\mathcal{N}}(\mathbf{x}_{m},\mathbf{\Sigma}_{m})$, by conditioning a random variable $X\sim\mu_{0}$ on information from $m$ search directions $\mathbf{S}_{m}$.

Algorithm 2.1 represents the iterative computation of the posteriors from (Cockayne:BCG, Propositions 6 and 7), (RICO21, Theorem 2.7). To illustrate the resemblance of BayesCG and the Conjugate Gradient method, we present the most common implementation of CG in Algorithm 2.2.

BayesCG (Algorithm 2.1) computes specific search directions $\mathbf{S}_{m}$ with two additional properties:

1. 1.

   They are $\mathbf{A}\mathbf{\Sigma}_{0}\mathbf{A}$-orthogonal, which means that $\mathbf{\Lambda}_{m}=\mathbf{S}_{m}^{T}\mathbf{A}\mathbf{\Sigma}_{0}\mathbf{A}\mathbf{S}_{m}$ is diagonal (Cockayne:BCG, Section 2.3), thus easy to invert.

2. 2.

   They form a basis for the Krylov space (BCG:Supp, Proposition S4)

   $$\operatorname{range}(\mathbf{S}_{m})=\mathcal{K}_{m}(\mathbf{A}\mathbf{\Sigma}_{0}\mathbf{A},\mathbf{r}_{0})\equiv\operatorname{span}\{\mathbf{r}_{0},\mathbf{A}\mathbf{\Sigma}_{0}\mathbf{A}\mathbf{r}_{0},\ldots,(\mathbf{A}\mathbf{\Sigma}_{0}\mathbf{A})^{m-1}\mathbf{r}_{0}\}.$$

After defining the Krylov space of maximal dimension (Definition 2.3), we review the ideal but impractical Krylov prior (Definition 2.4), and discuss its construction (Lemma 2.6) and properties (Theorem 2.7).

The Krylov prior is a Gaussian distribution whose covariance is constructed from a basis for the maximal dimensional CG Krylov space.

The posterior distributions from BayesCG under the Krylov prior depend on submatrices of $\mathbf{V}$ and $\mathbf{\Phi}$,

$$\begin{split}\mathbf{V}_{i:j}&\equiv\begin{bmatrix}\mathbf{v}_{i}&\cdots&\mathbf{v}_{j}\end{bmatrix}\\ \mathbf{\Phi}_{i:j}&\equiv\operatorname{diag}\begin{pmatrix}\phi_{i}&\cdots&\phi_{j}\end{pmatrix},\qquad 1\leq i\leq j\leq g,\end{split}$$

(6)

where $\mathbf{V}_{1:g}=\mathbf{V}$, $\mathbf{\Phi}_{1:g}=\mathbf{\Phi}$, and $\mathbf{V}_{j+1:j}=\mathbf{\Phi}_{j+1:j}=\mathbf{0}$, $1\leq j\leq n$.

Under suitable assumptions, BayesCG (Algorithm 2.1) produces the same iterates as CG (Algorithm 2.2).

Explicit construction of the ideal Krylov prior, followed by explicit computation of the Krylov posteriors in Algorithm 2.1 is impractical, because it is more expensive than solving the linear system (1) in the first place. That is the reason for introducing practical, approximate Krylov posteriors.

We dispense with the explicit computation of the Krylov prior, and instead compute a low-rank approximation of the final posterior (Definition 2.8) by running $d$ additional iterations. The corresponding CG-based implementation of BayesCG under approximate Krylov posteriors is relegated to Algorithm B.5 in Appendix B.

In contrast to the full Krylov posteriors, which reproduce the error as in (8), approximate Krylov posteriors underestimate the error (RICO21, Section 3.4),

$$\operatorname{trace}(\mathbf{A}\widehat{\mathbf{\Gamma}}_{m})=\operatorname{trace}(\mathbf{\Phi}_{m+1:m+d})=\|\mathbf{x}_{*}-\mathbf{x}_{m}\|_{\mathbf{A}}^{2}-\|\mathbf{x}_{*}-\mathbf{x}_{m+d}\|_{\mathbf{A}}^{2},$$

(9)

where $\|\mathbf{x}_{*}-\mathbf{x}_{m+d}\|_{\mathbf{A}}^{2}$ is the error after $m+d$ iterations of CG. The error underestimate $\operatorname{trace}(\mathbf{A}\widehat{\mathbf{\Gamma}}_{m})$ is equal to (StrakosTichy, Equation(4.9)), and it is more accurate when convergence is fast. Fast convergence makes $\operatorname{trace}(\mathbf{A}\widehat{\mathbf{\Gamma}}_{m})$ a more accurate estimate because fast convergence implies that $\|\mathbf{x}_{*}-\mathbf{x}_{m+d}\|_{\mathbf{A}}^{2}\ll\|\mathbf{x}_{*}-\mathbf{x}_{m}\|_{\mathbf{A}}^{2}$, and this, along with (9), implies that $\operatorname{trace}(\mathbf{A}\widehat{\mathbf{\Gamma}}_{m})\approx\|\mathbf{x}_{*}-\mathbf{x}_{m}\|_{\mathbf{A}}^{2}$ (StrakosTichy, Section 4).

We review the $p$-Wasserstein distance (Definition 3.1), extend the 2-Wasserstein distance to the $\mathbf{A}$-Wasserstein distance weighted by a symmetric positive definite matrix $\mathbf{A}$ (Theorem 3.4), and derive the $\mathbf{A}$-Wasserstein distance between approximate and full Krylov posteriors (Theorem 3.5).

The $p$-Wasserstein distance is a metric on the set of probability distributions.

In the special case $p=2$, the $2$-Wasserstein or Fréchet distance between two Gaussian distributions admits an explicit expression.

We generalize the 2-Wasserstein distance to the $\mathbf{A}$-Wasserstein distance weighted by a symmetric positive definite matrix $\mathbf{A}$.

We derive an explicit expression for the $\mathbf{A}$-Wasserstein distance analogous to the one for the 2-Wasserstein distance in Lemma 3.2.

We are ready to derive the $\mathbf{A}$-Wasserstein distance between approximate and full Krylov posteriors.

Expand the column dimension of $\mathbf{V}_{m+1:g}$ from $g-m$ to $n$ by adding an $\mathbf{A}$-orthogonal complement $\mathbf{V}_{m}^{\perp}\in\mathbb{R}^{n\times(n-g+m)}$ to create an $\mathbf{A}$-orthogonal matrix

$\displaystyle\widetilde{\mathbf{V}}$

$\displaystyle\equiv\begin{bmatrix}\mathbf{V}_{m+1:g}&\mathbf{V}_{m}^{\perp}\end{bmatrix}\in\mathbb{R}^{n\times n}$

with $\widetilde{\mathbf{V}}^{T}\mathbf{A}\widetilde{\mathbf{V}}=\mathbf{I}_{n}$. Analogously expand the dimension of the diagonal matrices by padding with trailing zeros,

	$\displaystyle\widetilde{\mathbf{\Phi}}_{m+1:g}$	$\displaystyle\equiv\operatorname{diag}\begin{pmatrix}\phi_{m+1}&\cdots&\phi_{g}&\mathbf{0}_{1\times(n-g+m)}\end{pmatrix}\in\mathbb{R}^{n\times n},$
	$\displaystyle\widetilde{\mathbf{\Phi}}_{m+1:m+d}$	$\displaystyle\equiv\operatorname{diag}\begin{pmatrix}\phi_{m+1}&\cdots&\phi_{m+d}&\mathbf{0}_{1\times(n-d)}\end{pmatrix}\in\mathbb{R}^{n\times n}.$

Factor the covariances in terms of the above square matrices,

$$\mathbf{\Gamma}_{m}=\widetilde{\mathbf{V}}\widetilde{\mathbf{\Phi}}_{m+1:g}\widetilde{\mathbf{V}}^{T}\quad\text{and}\quad\widehat{\mathbf{\Gamma}}_{m}=\widetilde{\mathbf{V}}\widetilde{\mathbf{\Phi}}_{m+1:m+d}\widetilde{\mathbf{V}}^{T}.$$

Substitute the factorizations into (3.4), and compute the $\mathbf{A}$-Wasserstein distance between $\eta_{m}$ and $\widehat{\eta}_{m}$ as

$$\left(W_{\mathbf{A}}(\eta_{m},\widehat{\eta}_{m})\right)^{2}=\operatorname{trace}(\mathbf{G})+\operatorname{trace}(\mathbf{J})-2\operatorname{trace}\left((\mathbf{G}^{1/2}\,\mathbf{J}\,\mathbf{G}^{1/2})^{1/2}\right),$$

(17)

where the congruence transformations of $\mathbf{\Gamma}_{m}$ and $\widehat{\mathbf{\Gamma}}_{m}$ are again Hermitian,

	$\displaystyle\mathbf{G}$	$\displaystyle\equiv\mathbf{A}^{1/2}\underbrace{\widetilde{\mathbf{V}}\widetilde{\mathbf{\Phi}}_{m+1:g}\widetilde{\mathbf{V}}^{T}}_{\mathbf{\Gamma}_{m}}\mathbf{A}^{1/2}=\mathbf{U}\widetilde{\mathbf{\Phi}}_{m+1:g}\mathbf{U}^{T},\qquad\mathbf{U}\equiv\mathbf{A}^{1/2}\widetilde{\mathbf{V}}$
	$\displaystyle\mathbf{J}$	$\displaystyle\equiv\mathbf{A}^{1/2}\underbrace{\widetilde{\mathbf{V}}\widetilde{\mathbf{\Phi}}_{m+1:m+d}\widetilde{\mathbf{V}}^{T}}_{\widehat{\mathbf{\Gamma}}_{m}}\mathbf{A}^{1/2}=\mathbf{U}\widetilde{\mathbf{\Phi}}_{m+1:d}\mathbf{U}^{T}.$

Lemma A.3 implies that $\mathbf{U}$ is an orthogonal matrix, so that the second factorizations of $\mathbf{G}$ and $\mathbf{J}$ represent eigenvalue decompositions. Commutativity of the trace implies

	$\displaystyle\operatorname{trace}(\mathbf{G})$	$\displaystyle=\operatorname{trace}(\widetilde{\mathbf{\Phi}}_{m+1:g})=\sum_{i=m+1}^{g}{\phi_{i}}$
	$\displaystyle\operatorname{trace}(\mathbf{J})$	$\displaystyle=\operatorname{trace}(\widetilde{\mathbf{\Phi}}_{m+1:m+d})=\sum_{i=m+1}^{m+d}{\phi_{i}}.$

Since $\mathbf{G}$ and $\mathbf{J}$ have the same eigenvector matrix, they commute, and so do diagonal matrices,

	$\displaystyle\mathbf{G}^{1/2}\mathbf{J}\mathbf{G}^{1/2}$	$\displaystyle=\mathbf{U}\widetilde{\mathbf{\Phi}}_{m+1:g}\widetilde{\mathbf{\Phi}}_{m+1:m+d}\mathbf{U}^{T}$
		$\displaystyle=\mathbf{U}\operatorname{diag}\begin{pmatrix}\phi_{m+1}^{2}&\cdots\phi_{m+d}^{2}&\mathbf{0}_{1\times(n-d)}\end{pmatrix}\mathbf{U}^{T}$

where the last equality follows from the fact that $\widetilde{\mathbf{\Phi}}_{m+1:g}$ and $\widetilde{\mathbf{\Phi}}_{m+1:m+d}$ share the leading $d$ diagonal elements. Thus

$\displaystyle\operatorname{trace}\left((\mathbf{G}^{1/2}\,\mathbf{J}\,\mathbf{G}^{1/2})^{1/2}\right)=\sum_{i=m+1}^{m+d}\phi_{i}.$

Substituting the above expressions into (17) gives

$\displaystyle\left(W_{\mathbf{A}}(\eta_{m},\widehat{\eta}_{m})\right)^{2}=\sum_{i=m+1}^{g}\phi_{i}+\sum_{i=m+1}^{m+d}\phi_{i}-2\sum_{i=m+1}^{m+d}\phi_{i}=\sum_{i=m+d+1}^{g}{\phi_{i}}.$

$\blacksquare$

Theorem 3.5 implies that the $\mathbf{A}$-Wasserstein distance between approximate and full Krylov posteriors is the sum of the CG steps sizes skipped by the approximate posterior, and this, as seen in (9) and (StrakosTichy, Equation (4.4)), is equal to the distance between the error estimate $\operatorname{trace}(\mathbf{A}\widehat{\mathbf{\Gamma}}_{m})$ and the true error $\|\mathbf{x}_{*}-\mathbf{x}_{m}\|_{\mathbf{A}}^{2}$. As a consequence, the approximation error decreases as the convergence of the posterior mean accelerates, or the rank $d$ of the approximation increases.

We interpret the Krylov prior as an ‘empirical Bayesian procedure’ (Theorem 3.7), and elucidate the connection between the random variables and the deterministic solution (Remark 3.8).

An empirical Bayesian procedure estimates the prior from data (Berger1985, Section 4.5). Our ‘data’ are the pairs of normalized search directions $\mathbf{v}_{i}$ and step sizes $\phi_{i}$, $1\leq i\leq m+d$, from $m+d$ iterations of CG. In contrast, the usual data for BayesCG are the inner products $\mathbf{v}_{i}^{T}\mathbf{b}$, $1\leq i\leq m$. However, if we augment the usual data with the search directions, which is natural due to their dependence on $\mathbf{x}_{*}$, then $\phi_{i}$ is just a function of the data.

From these data we construct a prior in an empirical Bayesian fashion, starting with a random variable

$$X=\mathbf{x}_{0}+\sum_{i=1}^{m+d}\sqrt{\phi_{i}}\mathbf{v}_{i}Q_{i}\in\mathbb{R}^{n},$$

where $Q_{i}\sim\operatorname{\mathcal{N}}(0,1)$ are independent and identically distributed scalar Gaussian random variables, $1\leq i\leq m+d$. Due to the independence of the $Q_{i}$, the above sum is the matrix vector product

$$X=\mathbf{x}_{0}+\mathbf{V}_{1:m+d}\,\mathbf{\Phi}_{1:m+d}^{1/2}\,Q$$

(18)

where $Q\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{m+d})$ is a vector-valued Gaussian random variable.

The distribution of $X$ is the empirical prior, while the distribution of $X$ conditioned on the random variable $Y\equiv\mathbf{V}_{1:m}^{T}\mathbf{A}X$ taking the value $\mathbf{V}_{1:m}^{T}\mathbf{b}$ is the empirical posterior. We relate these distributions to the Krylov prior.