Statistical properties of BayesCG under the Krylov prior

Abstract

We analyse the calibration of BayesCG under the Krylov prior. BayesCG is a probabilistic numeric extension of the Conjugate Gradient (CG) method for solving systems of linear equations with real symmetric positive definite coefficient matrix. In addition to the CG solution, BayesCG also returns a posterior distribution over the solution. In this context, a posterior distribution is said to be ‘calibrated’ if the CG error is well-described, in a precise distributional sense, by the posterior spread. Since it is known that BayesCG is not calibrated, we introduce two related weaker notions of calibration, whose departures from exact calibration can be quantified. Numerical experiments confirm that, under low-rank approximate Krylov posteriors, BayesCG is only slightly optimistic and exhibits the characteristics of a calibrated solver, and is computationally competitive with CG.

Appendices

Appendix A: Auxiliary results

We present auxiliary results required for proofs in other sections.

The stability of Gaussian distributions implies that a linear transformation of a Gaussian random variable remains Gaussian.

Lemma A.1

(Stability of Gaussian Distributions [37, Section 1.2]) Let $X\sim {{\,\mathrm{{\mathcal {N}}}\,}}({\textbf{x}},\varvec{\Sigma })$ be a Gaussian random variable with mean ${\textbf{x}}\in {\mathbb {R}}^n$ and covariance $\varvec{\Sigma }\in {\mathbb {R}}^{n\times n}$. If ${\textbf{y}}\in {\mathbb {R}}^n$ and ${\textbf{F}}\in {\mathbb {R}}^{n\times n}$, then

$$\begin{aligned} Z= {\textbf{y}}+{\textbf{F}}X\sim {{\,\mathrm{{\mathcal {N}}}\,}}({\textbf{y}}+{\textbf{F}}{\textbf{x}}, {\textbf{F}}\varvec{\Sigma }{\textbf{F}}^T). \end{aligned}$$

The conjugacy of Gaussian distributions implies that the distribution of a Gaussian random variable conditioned on information that linearly depends on the random variable is a Gaussian distribution.

Lemma A.2

(Conjugacy of Gaussian Distributions [38, Section 6.1], [25, Corollary 6.21]) Let $X\sim {{\,\mathrm{{\mathcal {N}}}\,}}({\textbf{x}},\varvec{\Sigma }_x)$ and $Y\sim {{\,\mathrm{{\mathcal {N}}}\,}}({\textbf{y}},\varvec{\Sigma }_y)$. The jointly Gaussian random variable $\begin{bmatrix}X^T&Y^T\end{bmatrix}^T$ has the distribution

$$\begin{aligned} \begin{bmatrix} X\\ Y\end{bmatrix} \sim {{\,\mathrm{{\mathcal {N}}}\,}}\left( \begin{bmatrix} {\textbf{x}}\\ {\textbf{y}}\end{bmatrix},\begin{bmatrix}\varvec{\Sigma }_x &{} \varvec{\Sigma }_{xy} \\ \varvec{\Sigma }_{xy}^T &{} \varvec{\Sigma }_y\end{bmatrix} \right) , \end{aligned}$$

where $\varvec{\Sigma }_{xy} \equiv {{\,\textrm{Cov}\,}}(X,Y) = {{\,\mathrm{{\mathbb {E}}}\,}}[(X-{\textbf{x}})(Y-{\textbf{y}})^T]$ and the conditional distribution of $X$ given $Y$ is

$$\begin{aligned} (X\mid Y) \sim {{\,\mathrm{{\mathcal {N}}}\,}}(\overbrace{{\textbf{x}}+ \varvec{\Sigma }_{xy}\varvec{\Sigma }_y^\dagger (Y- {\textbf{y}})}^{\text {mean}}, \ \overbrace{\varvec{\Sigma }_x - \varvec{\Sigma }_{xy}\varvec{\Sigma }_y^\dagger \varvec{\Sigma }_{xy}^T}^{\text {covariance}}). \end{aligned}$$

We show how to transform a ${\textbf{B}}$-orthogonal matrix into an orthogonal matrix.

Lemma A.3

Let ${\textbf{B}}\in {\mathbb {R}}^{n\times n}$ be symmetric positive definite, and let ${\textbf{H}}\in {\mathbb {R}}^{n\times n}$ be a ${\textbf{B}}$-orthogonal matrix with ${\textbf{H}}^T{\textbf{B}}{\textbf{H}}= {\textbf{H}}{\textbf{B}}{\textbf{H}}^T = {\textbf{I}}$. Then

$$\begin{aligned} {\textbf{U}}\equiv {\textbf{B}}^{1/2}{\textbf{H}}\end{aligned}$$

is an orthogonal matrix with ${\textbf{U}}^T{\textbf{U}}= {\textbf{U}}{\textbf{U}}^T={\textbf{I}}$.

Proof

The symmetry of ${\textbf{B}}$ and the ${\textbf{B}}$-orthogonality of ${\textbf{H}}$ imply

$$\begin{aligned} {\textbf{U}}^T{\textbf{U}}= {\textbf{H}}^T{\textbf{B}}{\textbf{H}}= {\textbf{I}}. \end{aligned}$$

From the orthonormality of the columns of ${\textbf{U}}$, and the fact that ${\textbf{U}}$ is square follows that ${\textbf{U}}$ is an orthogonal matrix [39, Definition 2.1.3]. $\square$

Definition A.4

[39, Section 7.3] The thin singular value decomposition of the rank-p matrix ${\textbf{G}}\in {\mathbb {R}}^{m\times n}$ is

$$\begin{aligned} {\textbf{G}}= {\textbf{U}}{\textbf{D}}{\textbf{W}}^T, \end{aligned}$$

where ${\textbf{U}}\in {\mathbb {R}}^{m\times p}$ and ${\textbf{W}}\in {\mathbb {R}}^{n\times p}$ are matrices with orthonormal columns and ${\textbf{D}}\in {\mathbb {R}}^{p\times p}$ is a diagonal matrix with positive diagonal elements. The Moore–Penrose inverse of ${\textbf{G}}$ is

$$\begin{aligned} {\textbf{G}}^\dagger = {\textbf{W}}{\textbf{D}}^{-1}{\textbf{U}}^T. \end{aligned}$$

If a matrix has full column-rank or full row-rank, then its Moore–Penrose can be expressed in terms of the matrix itself. Furthermore, the Moore–Penrose inverse of a product is equal to the product of the Moore–Penrose inverses, provided the first matrix has full column-rank and the second matrix has full row-rank.

Lemma A.5

[40, Corollary 1.4.2] Let ${\textbf{G}}\in {\mathbb {R}}^{m\times n}$ and ${\textbf{J}}\in {\mathbb {R}}^{n\times p}$ have full column and row rank respectively, so ${{\,\textrm{rank}\,}}({\textbf{G}})={{\,\textrm{rank}\,}}({\textbf{J}})=n$. The Moore–Penrose inverses of ${\textbf{G}}$ and ${\textbf{J}}$ are

$$\begin{aligned} {\textbf{G}}^\dagger = ({\textbf{G}}^T{\textbf{G}})^{-1}{\textbf{G}}^T \quad \text {and} \quad {\textbf{J}}^\dagger = {\textbf{J}}^T({\textbf{J}}{\textbf{J}}^T)^{-1} \end{aligned}$$

respectively, and the Moore–Penrose inverse of the product equals

$$\begin{aligned} ({\textbf{G}}{\textbf{J}})^\dagger = {\textbf{J}}^\dagger {\textbf{G}}^\dagger . \end{aligned}$$

Below is an explicit expression for the mean of a quadratic form of Gaussians.

Lemma A.6

[41, Sections 3.2b.1–3.2b.3] Let $Z\sim {{\,\mathrm{{\mathcal {N}}}\,}}({\textbf{x}}_z,\varvec{\Sigma }_z)$ be a Gaussian random variable in ${\mathbb {R}}^n$, and ${\textbf{B}}\in {\mathbb {R}}^{n\times n}$ be symmetric positive definite. The mean of $Z^T{\textbf{B}}Z$ is

$$\begin{aligned} {{\,\mathrm{{\mathbb {E}}}\,}}[Z^T{\textbf{B}}Z]&= {{\,\textrm{trace}\,}}({\textbf{B}}\varvec{\Sigma }_z) + {\textbf{x}}_z^T{\textbf{B}}{\textbf{x}}_z. \end{aligned}$$

We show that the squared Euclidean norm of a Gaussian random variable with an orthogonal projector as its covariance matrix is distributed according to a chi-squared distribution.

Lemma A.7

Let $Z\sim {{\,\mathrm{{\mathcal {N}}}\,}}({\textbf{0}},{\textbf{P}})$ be a Gaussian random variable in ${\mathbb {R}}^n$. If the covariance matrix ${\textbf{P}}$ is an orthogonal projector, that is, if ${\textbf{P}}^2 = {\textbf{P}}$ and ${\textbf{P}}= {\textbf{P}}^T$, then

$$\begin{aligned} \Vert X\Vert _2^2= (X^TX) \sim \chi _p^2, \end{aligned}$$

where $p = {{\,\textrm{rank}\,}}({\textbf{P}})$.

Proof

We express the projector in terms of orthonormal matrices and then use the invariance of the 2-norm under orthogonal matrices and the stability of Gaussians.

Since ${\textbf{P}}$ is an orthogonal projector, there exists ${\textbf{U}}_1\in {\mathbb {R}}^{n\times p}$ such that ${\textbf{U}}_1{\textbf{U}}_1^T = {\textbf{P}}$ and ${\textbf{U}}_1^T{\textbf{U}}= {\textbf{I}}_p$. Choose ${\textbf{U}}_2\in {\mathbb {R}}^{n\times (n-p)}$ so that ${\textbf{U}}= \begin{bmatrix} {\textbf{U}}_1&{\textbf{U}}_2 \end{bmatrix}$ is an orthogonal matrix. Thus,

$$\begin{aligned} X^TX= X^T{\textbf{U}}{\textbf{U}}^TX= X^T{\textbf{U}}_1{\textbf{U}}_1^TX+ X^T{\textbf{U}}_2{\textbf{U}}_2^TX. \end{aligned}$$

(A1)

Lemma A.1 implies that $Y= {\textbf{U}}_1^TX$ is distributed according to a Gaussian distribution with mean ${\textbf{0}}$ and covariance ${\textbf{U}}_1^T{\textbf{U}}_1{\textbf{U}}_1^T{\textbf{U}}= {\textbf{I}}_p$. Similarly, $Z= {\textbf{U}}_2^TX$ is distributed according to a Gaussian distribution with mean ${\textbf{0}}$ and covariance ${\textbf{U}}_2^T{\textbf{U}}_1{\textbf{U}}_1^T{\textbf{U}}_2 = {\textbf{0}}$, thus $Z= {\textbf{0}}$.

Substituting $Y$ and $Z$ into (A1) gives $X^TX= Y^TY+ {\textbf{0}}^T{\textbf{0}}$. From $Y\sim {{\,\mathrm{{\mathcal {N}}}\,}}({\textbf{0}},{\textbf{I}}_p)$ follows $(X^TX)\sim \chi _p^2$. $\square$

Lemma A.8

If ${\textbf{A}}\in {\mathbb {R}}^{n\times n}$ is symmetric positive definite, and $M\sim {{\,\mathrm{{\mathcal {N}}}\,}}({\textbf{x}}_\mu \varvec{\Sigma }_\mu )$ and $N\sim {{\,\mathrm{{\mathcal {N}}}\,}}({\textbf{x}}_\nu ,\varvec{\Sigma }_\nu )$ are independent random variables in ${\mathbb {R}}^n$, then

$$\begin{aligned} {{\,\mathrm{{\mathbb {E}}}\,}}[\Vert M-N\Vert _{\textbf{A}}^2] = \Vert {\textbf{x}}_\mu -{\textbf{x}}_\nu \Vert _{\textbf{A}}^2 + {{\,\textrm{trace}\,}}({\textbf{A}}\varvec{\Sigma }_\mu ) + {{\,\textrm{trace}\,}}({\textbf{A}}\varvec{\Sigma }_\nu ). \end{aligned}$$

Proof

The random variable $M-N$ has mean ${{\,\mathrm{{\mathbb {E}}}\,}}[M-N] = {\textbf{x}}_\mu -{\textbf{x}}_\nu$, and covariance

$$\begin{aligned} \varvec{\Sigma }_{M-N}&\equiv {{\,\textrm{Cov}\,}}(M-N,M-N) \\&= {{\,\textrm{Cov}\,}}(M,M) + {{\,\textrm{Cov}\,}}(N,N) -{{\,\textrm{Cov}\,}}(M,N) - {{\,\textrm{Cov}\,}}(N,M) \\&= {{\,\textrm{Cov}\,}}(M,M) + {{\,\textrm{Cov}\,}}(N,N) = \varvec{\Sigma }_{\mu }+\varvec{\Sigma }_{\nu }, \end{aligned}$$

where the covariances ${{\,\textrm{Cov}\,}}(M,N) ={{\,\textrm{Cov}\,}}(N,M) = 0$ because $M$ and $N$ are independent. Now apply Lemma A.6 to $M-N$. $\square$

Appendix B: Algorithms

We present algorithms for the modified Lanczos method (Sect. B.1), BayesCG with random search directions (Sect. B.2), BayesCG with covariances in factored form (Sect. B.3), and BayesCG under the Krylov prior (Sect. B.4).

1.1 Modified Lanczos method

The Lanczos method [42, Algorithm 6.15] produces an orthonormal basis for the Krylov space ${\mathcal {K}}_g({\textbf{A}},{\textbf{v}}_1)$, while the modified version in Algorithm B.1 produces an ${\textbf{A}}$-orthonormal basis.

Algorithm B.1

Modified Lanczos Method

Algorithm B.1 reorthogonalizes the basis vectors ${\textbf{v}}_i$ with Classical Gram-Schmidt performed twice, see Lines 10 and 11. This reorthogonalization technique can be implemented efficiently and produces vectors that are orthogonal to machine precision [43, 44].

1.2 BayesCG with random search directions

The version of BayesCG in Algorithm B.2 is designed to be calibrated because the search directions do not depend on ${\textbf{x}}_*$, hence the posteriors do not depend on ${\textbf{x}}_*$ either [6, Section 1.1].

After sampling an initial random search direction ${\textbf{s}}_1\sim {{\,\mathrm{{\mathcal {N}}}\,}}({\textbf{0}},{\textbf{I}})$, Algorithm B.2 computes an ${\textbf{A}}\varvec{\Sigma }_0{\textbf{A}}$-orthonormal basis for the Krylov space ${\mathcal {K}}_m({\textbf{A}}\varvec{\Sigma }_0{\textbf{A}},{\textbf{s}}_1)$ with Algorithm B.1. Then Algorithm B.2 computes the BayesCG posteriors directly with (2) and (3) from Theorem 2.1. The numerical experiments in Sect. 5 run Algorithm B.2 with the inverse prior $\mu _0 = {{\,\mathrm{{\mathcal {N}}}\,}}({\textbf{0}},{\textbf{A}}^{-1})$.

Algorithm B.2

BayesCG with random search directions

1.3 BayesCG with covariances in factored form

Algorithm B.3 takes as input a general prior covariance $\varvec{\Sigma }_0$ in factored form, and subsequently maintains the posterior covariances $\varvec{\Sigma }_m$ in factored form as well. Theorem B.1 presents the correctness proof for Algorithm B.3.

Theorem B.1

Under the conditions of Theorem 2.1, if $\varvec{\Sigma }_0 = {\textbf{F}}_0{\textbf{F}}_0^T$ for ${\textbf{F}}_0\in {\mathbb {R}}^{n\times \ell }$ and some $m \le \ell \le n$, then $\varvec{\Sigma }_m = {\textbf{F}}_m{\textbf{F}}_m^T$ with

$$\begin{aligned} {\textbf{F}}_m = {\textbf{F}}_0\left( {\textbf{I}}- {\textbf{F}}_0^T{\textbf{A}}{\textbf{S}}_m({\textbf{S}}_m^T{\textbf{A}}{\textbf{F}}_0{\textbf{F}}_0^T{\textbf{A}}{\textbf{S}}_m)^{-1}{\textbf{S}}_m{\textbf{A}}{\textbf{F}}_0\right) \in {\mathbb {R}}^{n\times \ell },\qquad 1\le m\le n. \end{aligned}$$

Proof

Fix m. Substituting $\varvec{\Sigma }_0 = {\textbf{F}}_0{\textbf{F}}_0^T$ into (3) and factoring out ${\textbf{F}}_0$ on the left and ${\textbf{F}}_0^T$ on the right gives $\varvec{\Sigma }_m = {\textbf{F}}_0{\textbf{P}}{\textbf{F}}_0^T$ where

$$\begin{aligned} {\textbf{P}}&\equiv {\textbf{I}}- {\textbf{F}}_0^T{\textbf{A}}{\textbf{S}}_m({\textbf{S}}_m^T{\textbf{A}}{\textbf{F}}_0{\textbf{F}}_0^T{\textbf{A}}{\textbf{S}}_m)^{-1}{\textbf{S}}_m{\textbf{A}}{\textbf{F}}_0\\&=({\textbf{I}}- {\textbf{Q}}({\textbf{Q}}^T{\textbf{Q}})^{-1}{\textbf{Q}}^T) \qquad \text {where}\quad {\textbf{Q}}\equiv {\textbf{F}}_0^T{\textbf{A}}{\textbf{S}}_m. \end{aligned}$$

Show that ${\textbf{P}}$ is a projector,

$$\begin{aligned} {\textbf{P}}^2&= {\textbf{I}}- 2{\textbf{Q}}({\textbf{Q}}^T{\textbf{Q}})^{-1}{\textbf{Q}}^T + {\textbf{Q}}({\textbf{Q}}^T{\textbf{Q}})^{-1}{\textbf{Q}}^T{\textbf{Q}}({\textbf{Q}}^T{\textbf{Q}})^{-1}{\textbf{Q}}^T \\&= {\textbf{I}}- {\textbf{Q}}({\textbf{Q}}^T{\textbf{Q}})^{-1}{\textbf{Q}}^T = {\textbf{P}}. \end{aligned}$$

Hence $\varvec{\Sigma }_m = {\textbf{F}}_0{\textbf{P}}{\textbf{F}}_0^T = {\textbf{F}}_0{\textbf{P}}{\textbf{P}}{\textbf{F}}_0^T = {\textbf{F}}_m{\textbf{F}}_m^T$. $\square$

Algorithm B.3

BayesCG with covariances in factored form

1.4 BayesCG under the Krylov prior

We present algorithms for BayesCG under full Krylov posteriors (Sect. B.4.1) and under approximate Krylov posteriors (Sect. B.4.2).

1.4.1 Full Krylov posteriors

Algorithm B.4 computes the following: a matrix ${\textbf{V}}$ whose columns are an ${\textbf{A}}$-orthonormal basis for ${\mathcal {K}}_g({\textbf{A}},{\textbf{r}}_0)$; the diagonal matrix $\varvec{\Phi }$ in (5); and the posterior mean ${\textbf{x}}_m$ in (26). The output consists of the posterior mean ${\textbf{x}}_m$, and the factors ${\textbf{V}}_{m+1:g}$ and $\varvec{\Phi }_{m+1:g}$ for the posterior covariance.

Algorithm B.4

BayesCG under the Krylov prior with full posteriors

1.4.2 Approximate Krylov posteriors

Algorithm B.5 computes rank-d approximate Krylov posteriors in two main steps: (i) posterior mean and iterates ${\textbf{x}}_m$ in Lines 5–14; and (ii) factorization of the posterior covariance ${\widehat{\varvec{\Gamma }}}_m$ in Lines 16–26.

Algorithm B.5

BayesCG under the Krylov prior [2, Algorithm 3.1]