Fast and Stable Gradient Approximation for Bilinear Forms of Hermitian Matrix Functions

TL;DR

This paper introduces a fast, stable, forward-only gradient approximation method for bilinear forms of Hermitian matrix functions, reducing computational costs and avoiding reorthogonalization.

Contribution

The authors propose a novel forward-only gradient approximation that reuses Lanczos passes, offering improved stability and efficiency over existing methods.

Findings

The new method's error is proportional to the Lanczos residual norm.

It is unconditionally stable in tests without reorthogonalization.

The approach is faster than current state-of-the-art techniques.

Abstract

Objectives involving bilinear forms $u^\top f(A(\theta))v$ for Hermitian $A$ arise widely in scientific computing and probabilistic machine learning. For large matrices, Lanczos efficiently approximates these quantities, but differentiating them with respect to $\theta$ is challenging. Existing approaches either backpropagate through the Lanczos recurrence, requiring reorthogonalization for stability, or apply Arnoldi to an augmented block matrix of twice the original size. Both introduce extra computation and orthogonalization costs that can limit performance on modern hardware. We propose a forward-only gradient approximation that reuses the Lanczos pass and adds very minimal overhead in most cases. We prove that its error is proportional to the Lanczos residual norm, the same quantity controlling the forward approximation. Whereas a traditional adjoint-based calculation would be unstable without reorthogonalization, the new method appears unconditionally stable in our tests. It is also faster than existing state-of-the-art approaches.