Robust Automatic Differentiation of Square-Root Kalman Filters via Gramian Differentials

TL;DR

This paper introduces a robust method for differentiating square-root Kalman filters by leveraging Gramian differentials, ensuring stable gradients even with rank-deficient matrices, which enhances parameter learning in state-space models.

Contribution

It derives a closed-form chain-rule for differentiating through the Gramian of the matrix, resolving issues with non-uniqueness and divergence in standard methods, and extends it to rank-deficient inputs.

Findings

Exact gradient computation for Kalman log-marginal likelihood

Stable differentiation through Gramian-based methods

Extension to rank-deficient matrices using pseudoinverse and null-space correction

Abstract

Square-root Kalman filters propagate state covariances in Cholesky-factor form for numerical stability, and are a natural target for gradient-based parameter learning in state-space models. Their core operation, triangularization of a matrix $M \in \mathbb{R}^{n \times m}$, is computed via a QR decomposition in practice, but naively differentiating through it causes two problems: the semi-orthogonal factor is non-unique when $m > n$, yielding undefined gradients; and the standard Jacobian formula involves inverses, which diverges when $M$ is rank-deficient. Both are resolved by the observation that all filter outputs relevant to learning depend on the input matrix only through the Gramian $MM^\top$, so the composite loss is smooth in $M$ even where the triangularization is not. We derive a closed-form chain-rule directly from the differential of this Gramian identity, prove it exact for the Kalman log-marginal likelihood and filtered moments, and extend it to rank-deficient inputs via a two-component decomposition: a column-space term based on the Moore--Penrose pseudoinverse, and a null-space correction for perturbations outside the column space of $M$.