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

1. Problem Statement

Square-root Kalman filters are widely used for their numerical stability, propagating state covariances ($P_t$) in Cholesky-factor form ($S_t$) where $P_t = S_t S_t^\top$. This makes them ideal candidates for gradient-based parameter learning in State-Space Models (SSMs) using Automatic Differentiation (AD).

However, differentiating through the core operation of these filters—the triangularization operator $T(M)$, which computes a lower-triangular $L$ such that $LL^\top = MM^\top$ (typically via QR decomposition)—presents two critical failure modes in standard AD frameworks:

1. Non-uniqueness of the Orthogonal Factor: When the input matrix $M$ is "wide" ($m > n$), the orthogonal factor $Q$ in the QR decomposition ($M^\top = QR$) is not unique. It admits a continuous family of valid choices. Consequently, the map $M \mapsto L$ is not differentiable in the standard sense, leading to undefined or unstable gradients.
2. Singularity in Rank-Deficient Cases: Standard Jacobian formulas for QR decomposition involve the inverse of the triangular factor ($R^{-1}$). If the input matrix $M$ is rank-deficient (common when process noise $U$ or observation noise $V$ is singular, or when the state covariance degenerates), this inverse diverges, causing the gradient calculation to fail.

2. Methodology

The author proposes a novel approach that bypasses the direct differentiation of the QR decomposition entirely. Instead, the method leverages the observation that all relevant filter outputs (log-marginal likelihood, filtered means, and covariances) depend on the triangular factor $L$ only through the Gramian $\Sigma = LL^\top = MM^\top$.

Since the map $M \mapsto MM^\top$ is a smooth polynomial map, the composite loss function is smooth even where the triangularization $T(M)$ is not. The methodology proceeds as follows:

• Gramian Differential Identity: Instead of finding the unique differential of $L$, the method seeks a surrogate tangent $dL$ that satisfies the differential of the Gramian identity:
  $$dL L^\top + L dL^\top = dM M^\top + M dM^\top$$
  Any $dL$ satisfying this equation yields the exact Jacobian-Vector Product (JVP) for any loss function dependent on the Gramian.

• Two-Component Decomposition: To construct a well-defined $dL$ for any rank of $M$, the method decomposes the problem into two parts using the orthogonal projector $P = LL^\dagger$ (where $L^\dagger$ is the Moore-Penrose pseudoinverse):

  1. Column-Space Component ($dL_{col}$): Handles perturbations within the column space of $L$. It utilizes the QR structure to solve for a lower-triangular update involving $K = L^\dagger dM Q$.
  2. Null-Space Component ($dL_{null}$): Handles perturbations outside the column space of $L$. It is defined as $dL_{null} = (I - LL^\dagger) dM Q$.

• Closed-Form Solution: The total surrogate tangent is $dL = dL_{col} + dL_{null}$. This formula is valid for full-rank and rank-deficient matrices alike. When $L$ is full-rank, the null-space term vanishes, recovering classical results.

• Reverse-Mode Differentiation: The paper proves that the map $dM \mapsto dL$ is linear. This allows the Vector-Jacobian Product (VJP) required for reverse-mode AD (backpropagation) to be obtained via automatic transposition of the forward-mode operator, eliminating the need for manual derivation of complex backward formulas.

3. Key Contributions

1. Closed-Form Surrogate JVP: Derivation of a robust, closed-form formula for the Jacobian-Vector Product of the triangularization operator that works for both full-rank and rank-deficient inputs.
2. Gramian Sufficiency Proof: A formal proof demonstrating that all standard Kalman filter outputs (likelihood, mean, covariance) are functions of the Gramian $MM^\top$, ensuring that the surrogate tangent yields exact gradients for learning tasks.
3. Automatic VJP Generation: Establishment of the linearity of the derived operator, enabling seamless integration into modern AD frameworks (like JAX and PyTorch) via automatic transposition, avoiding manual implementation of backward passes.
4. Open-Source Implementation: Provision of a working implementation compatible with JAX.

4. Results

• Empirical Validation: The method was tested on a Linear Gaussian State-Space Model (LGSSM) with varying noise covariance rank (interpolating between full-rank and rank-one).
• Comparison:
  • Naive QR Approach: Failed completely, returning invalid values or diverging due to rank deficiency and non-uniqueness issues.
  • Proposed Method: Successfully computed consistent and correct gradients for the log-marginal likelihood across the entire range of noise ranks.
• Visual Confirmation: Figure 1 in the paper shows the log-marginal likelihood curve and the AD-computed tangents (arrows) aligning perfectly, confirming the mathematical correctness of the gradients even in degenerate cases.

5. Significance

This work resolves a fundamental bottleneck in training state-space models using square-root Kalman filters. By providing a mathematically rigorous and numerically stable way to compute gradients through rank-deficient and non-unique triangularization operations, it enables:

• Robust Learning: Training SSMs with singular noise covariances or degenerate states, which are common in real-world applications but previously required ad-hoc fixes or were impossible.
• End-to-End Differentiability: Facilitating the use of square-root filters in modern deep learning pipelines (e.g., JAX, PyTorch) without sacrificing numerical stability.
• General Applicability: The "Gramian Differential" concept offers a new paradigm for differentiating through non-smooth matrix factorizations where the loss depends only on the resulting Gramian.