Variational Free Energy Pivot Selection for Pivoted Cholesky

Technical Summary: Variational Free Energy Pivot Selection for Pivoted Cholesky

Problem Statement
Low-rank approximations of symmetric positive definite (SPD) matrices are fundamental to numerical linear algebra, particularly in Gaussian Process (GP) regression where the kernel matrix $K$ is approximated by a rank-$r$ factorization $\tilde{K} = LL^\top$. Classical methods, such as greedy and randomly pivoted Cholesky (RPCholesky), select pivots sequentially to minimize the algebraic trace-norm error of the residual ($\text{tr}(K - \tilde{K})$). However, in many applications, the matrix $K$ is not evaluated in isolation but within a nonlinear matrix functional. Specifically, in GP regression, the quality of the approximation is determined by the Variational Free Energy (VFE), denoted $L(\tilde{K})$, which serves as a lower bound on the log marginal likelihood.

The VFE functional combines three distinct terms: a log-determinant term (complexity), a quadratic data-fit term, and a trace penalty term. Existing pivot selection rules, which rely solely on residual diagonal magnitudes, target only the trace-norm error. Consequently, they neglect the log-determinant and data-fit components, potentially leading to systematically suboptimal approximations for the specific inference objective. The paper addresses the question: How should pivot selection be performed when the goal is to optimize a matrix functional rather than purely algebraic approximation error?

Methodology
The authors propose a new pivot selection rule, termed $\Delta$-VFE pivoted Cholesky, which directly maximizes the exact one-step change in the VFE functional under Cholesky-consistent rank-1 updates.

1. Derivation of the Gain Formula:
   The authors derive a closed-form expression for the exact change $\Delta L$ in the VFE $L(\tilde{K})$ induced by a rank-1 update $\tilde{K}_{i+1} = \tilde{K}_i + uu^\top$. Using the Matrix Determinant Lemma and the Sherman-Morrison formula, they decompose the gain into three additive components:
   $$\Delta L = -\frac{1}{2} \left( \log(1 + a) - \frac{b^2}{1 + a} - \frac{1}{\sigma^2_\varepsilon} \|u\|^2_2 \right)$$
   where:

   • $a = u^\top \Sigma_i^{-1} u$ represents the complexity contribution (log-determinant term), measuring the interaction of the update with the regularized inverse metric.
   • $b = u^\top \Sigma_i^{-1} y$ represents the data-fit contribution, measuring the alignment between the update direction and the current prediction residual.
   • $\|u\|^2_2$ represents the trace contribution (residual variance reduction), which corresponds to the exact one-step trace decrease rather than the diagonal proxy used in classical methods.

2. Cholesky-Consistent Updates:
   To ensure the approximation remains a valid Cholesky factorization, the update vector $u$ is restricted to the form $u_p = t_p / \sqrt{(T_i)_{p,p}}$, where $t_p$ is the $p$-th column of the residual matrix $T_i$. This restriction ensures the residual remains positive semi-definite and the diagonal entry of the selected pivot becomes zero.

3. Algorithmic Realization:
   Evaluating $\Delta L$ for all $n$ candidates at every step would be computationally prohibitive ($O(n^2r^2)$). To maintain efficiency, the method employs a batched sampling strategy:

   • At each iteration, a batch of $s$ candidate pivots is sampled from the set of unused indices.
   • The sampling probability is proportional to the residual diagonal entries, identical to the distribution used in RPCholesky.
   • The regularized inverse $\Sigma_i^{-1} = (\tilde{K}_i + \sigma^2_\varepsilon I)^{-1}$ is applied to vectors efficiently using the Woodbury matrix identity and incremental updates to an auxiliary matrix $B = (I + L^\top L / \sigma^2_\varepsilon)^{-1}$.
   • This approach allows the evaluation of the gain for $s$ candidates with a total cost of $O(snr^2)$, matching the asymptotic complexity of RPCholesky up to the batch factor $s$.

Key Contributions

• Objective-Aware Pivot Rule: The paper derives the first pivot selection criterion that explicitly maximizes the per-step gain of the Variational Free Energy, rather than the trace-norm error.
• Theoretical Guarantees: The authors prove that any sequence of Cholesky-consistent rank-1 updates yields a monotonically non-decreasing sequence of VFE values. The proposed rule selects the pivot that maximizes this guaranteed gain at each step.
• Complexity Efficiency: The method achieves the same asymptotic complexity as randomized pivoted Cholesky ($O(snr^2)$) by leveraging incremental Woodbury updates and batched candidate evaluation.
• Hierarchical Interpretation: The work establishes a three-level hierarchy of pivot rules:
  1. Level 1: Classical greedy/RPCholesky (uses diagonal proxy for trace reduction).
  2. Level 2: Exact trace-reduction maximizer (uses full residual column, ignores data/complexity).
  3. Level 3: $\Delta$-VFE (uses full residual column, data vector, and regularized inverse to optimize the full VFE).

Results
Numerical experiments on the Abalone and QM7 datasets demonstrate the following:

• VFE Optimization: $\Delta$-VFE consistently achieves a tighter bound on the log marginal likelihood (smaller relative VFE gap) compared to both greedy and RPCholesky, particularly at low to moderate ranks ($r \le 1024$).
• Approximation Quality: The method maintains trace-norm approximation quality comparable to RPCholesky. It does not sacrifice standard kernel approximation accuracy to gain VFE improvements.
• Predictive Performance: In GP regression tasks, the improved VFE translates to reduced prediction errors (RMSE and MAE) at low to moderate ranks.
• Batch Size Sensitivity: Increasing the batch size $s$ yields diminishing returns in performance but consistently improves the VFE gap and predictive accuracy. Even small batch sizes (e.g., $s=8$ or $s=16$) provide substantial gains over RPCholesky.

Significance
The paper claims that the proposed method bridges the gap between algebraic matrix approximation and task-specific inference objectives. By deriving an exact, closed-form gain for the VFE, the authors provide a greedy pivot rule that is explicitly aligned with the probabilistic inference goal of Gaussian processes. The results suggest that for applications where the matrix functional includes data-fit or complexity terms, optimizing the trace-norm error alone is insufficient, and objective-aware selection yields superior practical performance without incurring significant computational overhead. The method is presented as a direct generalization of RPCholesky that incorporates additional information about the objective function.