Randomly Pivoted Cholesky

Suppose $\vec{A}$ is positive definite and recall the Nyström approximation

When the columns of $\vec{\Omega}$ a subset of the columns of the identity, then $\vec{A}\vec{\Omega}$ corresponds to subsampling the columns of $\vec{A}$. In particular, let $S \subseteq \{1, \ldots, n\}$ be a tuple containing the $k$ pivots (columns of $\vec{A}$) we will observe; i.e. so that $\vec{\Omega} = \vec{I}[:,S]$. Then

In particular, since $\vec{A}$ is symmetric, we can compute $\vec{A} \langle \vec{\Omega}\rangle$ having only observed $\vec{A}(:,S)$, which contains just $kn$ entries of $\vec{A}$.

The key question is how best to choose the columns of $\vec{A}$. Ideally, we would like to choose $k$ columns so that the Nyström approximation is competitive with the best rank-$k$ approximation to $\vec{A}$.

Partial Cholesky factorization with pivoting¶

It turns out that, to compute (7.5) we can use a Cholesky factorization algorithm with pivoting, and stop after $k$ steps.

Note that the “textbook” Cholesky factorization algorithms maintain $\vec{A} - \vec{F}_i \vec{F}_i^\T$ directly. On the other hand, anticipating that we will terminate for some $k<n$, Algorithm 7.2 only computes the necessary parts of $\vec{A} - \vec{F}_i\vec{F}_i^\T$ as they are needed.

Adaptive pivoting¶

Nothing in Algorithm 7.2 requires that the pivot $s_i$ be chosen prior to step $i$! In particular, we can choose the $i$-th pivot adaptively, based on the approximation $\vec{A}\langle S_{i-1}\rangle = \vec{F}_{i-1}\vec{F}_{i-1}^\T$, where $S_i := (s_1, \ldots, s_{i-1})$. While computing the error $\vec{A} - \vec{A}\langle S_{i-1}\rangle$ would let us try to find the column that would reduce the error the most, we want to avoid looking at all of the entires of $\vec{A}$.

Amazingly, we can find good pivots without observing all of $\vec{A}$. Towards this end, note that:

1. The error $\vec{A} - \vec{A}\langle \vec{\Omega} \rangle$ of a Nyström approximation is positive definite.

2. For any positive definite $\vec{E}$, $\|\vec{E}\|\leq \|\vec{E}\|_\F \leq \tr(\vec{E})$.

By computing the $n$ diagonal entries of $\vec{A}$, we can keep track of $\operatorname{diag}(\vec{A} - \vec{A}\langle S_{i-1}\rangle )$ and use this to choose the pivot. One approach is to greedily choose the pivot as the largest entries of $\operatorname{diag}(\vec{A} - \vec{A}\langle S_{i-1}\rangle )$. However, this approach is has the tendency to focus on outlier entires. Instead, we can sample proprtional to the values $\operatorname{diag}(\vec{A} - \vec{A}\langle S_{i-1}\rangle )$. This results in the Randomly Pivoted Cholesky algorithm introduced in Chen et al., 2024