Uncertainty Calibration for Linear Solvers

Summary

CG gives an excellent posterior mean $A_M=\tilde{Y}{\tilde{Y}}^{\top }$; the remaining task is to attach a trustworthy covariance. The idea: choose the prior covariance parameter $W_0$ to equal $A_M$ on the explored subspace $\mathrm{span}(S)$ and a scaled identity $\omega I$ on its orthogonal complement, then estimate the single scale $\omega$ at runtime. Rayleigh regression predicts $\omega$ from CG’s own scalar by-products (the Rayleigh coefficients $a(m)=s_m^{\top }A{s}_m/s_m^{\top }{s}_m$), giving a genuinely probabilistic error estimate for arbitrary projections $Av$. A hard worst-case bound ($\omega \ge {\max }_{ij}[A{]}_{ij}$) gives a conservative but loose error estimate; a conjugate Gauss–Gamma prior (Ch. 22.5) gives a fully Bayesian scale.

Overview

From Conjugate Gradients as Probabilistic Inference and Computational Constraints on Probabilistic Solvers, the mean $A_M$ is fixed and good, but the naive empirical-Bayes covariance ($W_0=A_M$) collapses to zero on the null space of the observations. The final goal of Part III (Ch. 21) is a covariance that is (i) probabilistically consistent with $A_M$ (both arise from one generative model) and (ii) well-calibrated so it can serve as an uncertainty. Because the object is matrix-valued, different design criteria (diagonal vs off-diagonal error, projection error) pull towards different priors, so no single choice is perfect.

Main Content

The projection-complement covariance

Calibrated $W_0$ with a null-space scale

Set $W_0$ to act like $A$ on $\mathrm{span}(S)$ and estimate its effect on the complement using regularity assumptions:

$$W_0=\tilde{Y}{\tilde{Y}}^{\top }+(I-S(S^{\top }S{)}^{-1}{S}^{\top })\Omega (I-S(S^{\top }S{)}^{-1}{S}^{\top }),\text{\hspace{1em}(21.1)}$$

with a general spd $\Omega$. The projection matrices ensure $\Omega$ only acts on the space not covered by $A_M=\tilde{Y}{\tilde{Y}}^{\top }$, so any such $\Omega$ still reproduces the mean $A_M$ (20.3). For simplicity take a scalar $\Omega =\omega I$ (21.2): with no specific prior knowledge there is no way to prefer directions in the null-space of $S$. Then

$$W_M=W_0-W_{0}S(S^{\top }{W}_{0}S{)}^{-1}{S}^{\top }{W}_0=\omega (I-S(S^{\top }S{)}^{-1}{S}^{\top }).\text{\hspace{1em}(via 21.1)}$$

The scale $\omega$ scales the remaining uncertainty over the entire null-space of $S$ — the space CG has not yet explored.

The diagonal/off-diagonal trade-off of (20.4)–(20.6) persists: calibrated variance on the diagonal implies under-confidence off it, and vice versa, so $\omega$ strikes a balance.

Rayleigh regression: estimating $\omega$ from CG’s by-products

Rayleigh coefficients ^def-rayleigh

During a CG run on $A$, the projected values

$$a(m):=\frac{s_m^{\top }A{s}_m}{s_m^{\top }{s}_m}$$

(the $m$-th direction’s Rayleigh coefficient) are computed essentially for free (the term $s_m^{\top }A{s}_m$ is line 7 of Algorithm 17.2, up to rescaling). For spd $A$ with eigenvalues ${\lambda }_1\ge \cdots \ge {\lambda }_N$,

$${\lambda }_1\ge a(m)\ge {\lambda }_N\forall m,\text{and}{\sigma }^2\le \frac{v^{\top }Av}{v^{\top }v}\le \mathrm{tr}A\forall v.$$

So $\{a(m)\}$ carry (loose) spectral information about $A$ at no extra cost.

Empirically (SARCOS kernel-ridge example, Fig. 21.1): the collected projections decay rapidly in the first $\sim 50$ steps (an expanding subspace of relevant directions, related to but not equal to top eigenvectors), then a “kink”, then slow non-monotone decay. Since the $a(m)$ come for free and are scalar, cheap univariate regression on them predicts $\omega$ (and hence $v^{\top }Av$ for unseen $v$):

Rayleigh regression ^def-rayleigh-reg

Fit a parametric curve to the $M$ observed ${\log }_{10}a(m)$, e.g.

$$\hat{a}(m)={\sigma }^2+10^{{\xi }_1+{\xi }_{2}m}+10^{{\xi }_3+{\xi }_{4}m},\text{\hspace{1em}(21.4)}$$

with constants ${\xi }_1,\dots ,{\xi }_4$ found by least squares (the posterior mean of parametric Gaussian regression). Then estimate the average projection scale from stopping point $M$ to $N$:

$${\omega }_{\text{projections}}:=\frac{1}{N-M}{\int }_M^N\hat{a}(m)\mathrm{d}m.\text{\hspace{1em}(21.5)}$$

A more automatic GP-regression version on $\log a(m)$ is possible (Wenger & Hennig, 2020).

Predicting general matrix projections ^thm-projection

Under the posterior $p(A)=\mathcal{N}(A;A_M,W_M{\otimes }_{\ominus }{W}_M)$, the marginal over a projection $Av=(I\otimes v{)}^{\top }\vec{A}$ is Gaussian:

$$p(Av)=\mathcal{N}(Av;A_{M}v,\underset{=:{\Sigma }_v}{\underbrace{\frac{1}{2}(W_M{v}^{\top }{W}_{M}v+(W_{M}v)(v^{\top }{W}_M))}}).$$

With $\omega$ set by Rayleigh regression (SARCOS: $\omega \approx 0.02$ after $M=300$, vs radically larger $\omega =\mathrm{tr}A$), the standardised elements $z={\Sigma }_v^{-1/2}(Av-\mathbb{E}[Av])$ are close to standard-normal for random $v$ (Gaussian/uniform/binary), Fig. 21.2 — the posterior captures the two moments of $Av$ well without access to the distribution of $v$. This is a new, genuinely probabilistic error estimate emerging from the interpretation of linear solvers.

Why the good Gaussian fit is unsurprising: for $v_i\sim p_v$ i.i.d., $[Av{]}_i={\sum }_j[A{]}_{ij}{v}_j$ is approximately Gaussian by the CLT, with ${\mathbb{E}}_{p_v}([Av{]}_i)=\mathbb{E}[v]{\sum }_j[A{]}_{ij}$ and $\mathrm{var}=\mathrm{var}(v){\sum }_j[A{]}_{ij}^2$. Randomness “washes out structure”, leaving the maximum-entropy Gaussian the model already assumes — helpful here, but a sign the interesting structure has been removed.

The hard part: individual matrix elements

Predicting single elements $[A{]}_{ij}$ (deterministic, not smeared by randomness) is genuinely harder. The marginal is

$$p([A{]}_{ij}\mid Y,S)=\mathcal{N}([A{]}_{ij};[A_M{]}_{ij},\frac{1}{2}([W_M{]}_{ii}[W_M{]}_{jj}+[W_M{]}_{ij}^2)),$$

and (20.5)–(20.6) show no scalar $\omega$ (nor even a full spd $W_0$) makes this variance a tight prediction of the error on all elements simultaneously. Two options:

Two calibration regimes for elements ^def-two-regimes

• Hard upper bound: ${\mathrm{var}}_{p(A)}([A{]}_{ij})$ upper-bounds $[A-A_M{]}_{ij}^2$ if $\frac{1}{2}([W_0{]}_{ii}[W_0{]}_{jj}+[W_0{]}_{ij}^2)>|A-A_M{|}_{ij}$. For spd matrices ($|A{|}^2\le |A_{ii}||A_{jj}|$) this holds if $\omega \ge {\max }_{ij}[A{]}_{ij}$; an $\mathcal{O}(N)$ a-priori bound is $\omega ={\max }_i[A{]}_{ii}^2$ (SARCOS Gram matrix: known to be $1$). Guarantees a conservative bound but is loose for off-diagonal elements (the vast majority).
• Estimated average: $\omega ={\omega }_{\text{projections}}$ (21.5) gives no guarantee but captures the typical matrix scale — a more aggressive, often more useful error estimate for off-diagonal elements ($\sim 1$), though on the diagonal some outliers can have true/estimated error ratios beyond $10$ (Fig. 21.3).

Conjugate-prior scale inference (Ch. 22.5)

An alternative, fully Bayesian calibration uses the scalar-covariance prior $p_A(A\mid \alpha ,\beta )=\mathcal{N}(A;\alpha I,{\beta }^{2}I{\otimes }_{\ominus }I)$ (which by Theorem 19.16 is also CG-consistent) and places a Gauss–Gamma conjugate prior on the hyperparameters $(\alpha ,\beta )$.

Gauss–Gamma posterior on the scale (Eqs. 22.7–22.11) ^thm-gauss-gamma

With prior $\mathcal{N}(\alpha ;{\mu }_0,{\beta }^2/{\lambda }_0)\mathcal{G}({\beta }^{-2};a_0,b_0)$ and marginal likelihood $p(Y,S\mid \alpha ,\beta )=\mathcal{N}(\vec{Y};\alpha \vec{S},{\beta }^2(I\otimes S{)}^{\top }(I{\otimes }_{\ominus }I)(I\otimes S))$, the posterior sufficient statistics are

$${\mu }_N=\frac{{\lambda }_0{\mu }_0+M\bar{\alpha }}{{\lambda }_0+M},{\lambda }_N={\lambda }_0+M,a_N=a_0+\frac{1}{2}(NM-\frac{1}{2}(M^2-M)),$$ $$b_N=b_0+\frac{1}{2}(M{\hat{\beta }}^2+\frac{{\lambda }_{0}M}{{\lambda }_0+M}(\bar{\alpha }-{\mu }_0{)}^2),$$

where, writing the SVD $S=V\Sigma U^{\top }$,

$$\bar{\alpha }=\frac{1}{M}\sum _{i=1}^M(V^{\top }AV{)}_{ii},{\hat{\beta }}^2=\frac{1}{M}\sum _{ij=1}^M((V^{\top }AV{)}_{ij}^2-\frac{(V^{\top }AV{)}_{ii}(V^{\top }AV{)}_{jj}}{M}).$$

Interpretation: the sufficient statistics compute an empirical expectation over the elements of $A$ in the left-singular basis of $S$. $\bar{\alpha }$ is an empirical mean over the (rotated) diagonal; ${\hat{\beta }}^2$ is an empirical variance. When $S$ is orthonormal, $V=S$ and $V^{\top }AV=S^{\top }AS=Y^{\top }S$ directly. Both direct and iterative conjugate solvers thus project along a transformation of orthogonal directions — for CG, the sequence of residuals $r_i=A{x}_i-b$.

Examples

Why CG's step sizes calibrate the whole matrix

The Rayleigh coefficients $a(m)$ are literally line 7 of the solver — CG already computes them. A cheap hand-crafted regression (21.4) on these $M$ scalars finds the right scale for the $N-M$ directions of $A$ that the mean $A_M$ has not captured. The predicted variance ${\Sigma }_v$ for a projection $Av$ is then a fundamentally probabilistic error bar that “comes for free” — a concrete instance of the recurring thesis that a classical solver already carries the ingredients of its own uncertainty.

Conservative worst-case vs aggressive average (Fig. 21.3)

On the SARCOS Gram matrix, $\omega ={\max }_i[A{]}_{ii}^2=1$ makes the posterior std a valid but loose upper bound on element errors. Switching to ${\omega }_{\text{projections}}\approx 0.02$ gives well-scaled ($\sim 1$) error ratios for off-diagonal elements, at the cost of losing the guarantee and occasionally under-estimating diagonal errors by $>10\times$. Calibration is thus a deliberate choice between conservative bounds and realistic average error.

Connections

• Directly resolves the $W_M=0$ failure noted in Computational Constraints on Probabilistic Solvers.
• The Gauss–Gamma treatment is the linear-algebra instance of Hierarchical Inference in Gaussian Models (conjugate-prior hyperparameter inference, §6.2).
• The “under-confidence trade-off” mirrors the integration chapter’s choice of kernels with analytic means but imperfect calibration — see Convergence and Priors in Bayesian Quadrature and Lessons from Integration.
• Rayleigh regression’s spectral reading ties back to the Lanczos/Krylov structure of Classic Linear Solvers - A Review.

See Also

• Computational Constraints on Probabilistic Solvers — supplies the mean $A_M=\tilde{Y}{\tilde{Y}}^{\top }$ this note calibrates.
• Conjugate Gradients as Probabilistic Inference — the scalar prior underlying the conjugate-prior calibration.
• Hierarchical Inference in Gaussian Models — the general conjugate/empirical-Bayes machinery.
• Probabilistic Linear Solvers - Algorithmic Scaffold — element-variance formulas (20.4) and the symmetric-Kronecker covariance.