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

1. Il Problema

I filtri di Kalman a radice quadrata (Square-Root Kalman Filters) sono fondamentali per la stabilità numerica nei modelli a spazio di stato (SSM), poiché propagano le covarianze degli stati sotto forma di fattori di Cholesky triangolari inferiori ($S_{t|t}$) invece delle matrici di covarianza complete. Questo approccio evita che le matrici di covarianza perdano la proprietà di semi-definita positiva a causa di errori di arrotondamento.

Tuttavia, l'addestramento di questi modelli tramite metodi basati sul gradiente (come la massimizzazione della verosimiglianza marginale logaritmica) richiede l'uso della differenziazione automatica (AD). Il problema centrale identificato nell'articolo è che l'operazione fondamentale del filtro, ovvero la triangolarizzazione di una matrice $M$ (ottenuta tramite decomposizione QR), presenta due gravi ostacoli quando si tenta di calcolare i gradienti in modo "naive" (diretto):

1. Non unicità quando $m > n$: Se la matrice di input è "ampia" (più colonne che righe), il fattore ortogonale nella decomposizione QR non è unico. Esiste una famiglia continua di scelte valide, rendendo il gradiente (JVP - Jacobian-Vector Product) indefinito o non canonico.
2. Singolarità in caso di rango deficitario: Le formule standard per il gradiente della decomposizione QR coinvolgono l'inversa della matrice triangolare ($R^{-1}$). Se la matrice $M$ è a rango deficitario (comune quando le matrici di covarianza $U$ o $V$ sono singolari o degenerano), l'inversa diverge, rendendo il calcolo del gradiente impossibile.

Questi problemi si verificano frequentemente nei passi di previsione e aggiornamento del filtro di Kalman, specialmente in scenari con rumore di processo o osservazione singolare.

2. Metodologia

L'autore propone una soluzione basata su un'osservazione chiave: tutti gli output rilevanti per l'apprendimento del filtro (verosimiglianza marginale, media filtrata, covarianza filtrata) dipendono dal fattore triangolare $L$ solo attraverso la sua Gramiana $LL^\top = MM^\top$, e non direttamente dalla struttura triangolare di $L$.

Poiché la mappa $M \mapsto MM^\top$ è polinomiale e quindi liscia (smooth) ovunque, anche se la triangolarizzazione $T(M)$ non lo è, la funzione di perdita composta è liscia rispetto a $M$.

La metodologia si articola nei seguenti punti:

• Differenziale della Gramiana: Invece di differenziare attraverso la decomposizione QR, il metodo deriva un gradiente diretto partendo dall'identità differenziale della Gramiana:
  $$dL L^\top + L dL^\top = dM M^\top + M dM^\top$$
  L'obiettivo è trovare un "tensore surrogato" $dL$ che soddisfi questa equazione.
• Decomposizione in Sottospazi: Per risolvere l'equazione sopra anche quando $L$ non è invertibile (rango $r < n$), il metodo scompone il problema in due componenti:
  1. Componente dello Spazio delle Colonne: Gestisce le perturbazioni all'interno dello spazio generato da $L$. Utilizza la pseudoinversa di Moore-Penrose ($L^\dagger$) per costruire una soluzione unica per la parte triangolare inferiore.
  2. Componente dello Spazio Nullo: Gestisce le perturbazioni esterne allo spazio delle colonne di $L$. Viene aggiunta una correzione specifica ($dL_{null}$) che garantisce che l'identità della Gramiana sia soddisfatta anche per le direzioni fuori dallo spazio colonna.
• Linearità: Viene dimostrato che la mappa che trasforma il vettore tangente di input $dM$ nel vettore tangente surrogato $dL$ è lineare. Questo è cruciale per l'implementazione nella modalità inversa (Reverse Mode) dell'AD.

3. Contributi Chiave

I principali contributi del lavoro sono:

1. Formula Chiusa per JVP: Derivazione di una formula esplicita e a forma chiusa per il Jacobian-Vector Product (JVP) dell'operatore di triangolarizzazione. Questa formula è esatta per qualsiasi perdita dipendente dalla Gramiana, recupera le formule classiche basate su QR quando esistono, e si estende correttamente ai casi a rango deficitario.
2. Dimostrazione di Sufficienza della Gramiana: Una prova formale che tutti gli output critici del filtro di Kalman dipendono solo dalla Gramiana, giustificando matematicamente l'approccio di bypassare la triangolarizzazione diretta.
3. Derivazione Gratuita del VJP: Grazie alla linearità dimostrata della mappa JVP, il Vector-Jacobian Product (VJP) necessario per la retropropagazione (reverse mode) può essere ottenuto per semplice trasposizione automatica delle operazioni primitive, senza bisogno di derivare nuove formule complesse.
4. Implementazione Open Source: Fornitura di un'implementazione pratica compatibile con framework moderni come JAX e PyTorch.

4. Risultati Sperimentali

L'autore valida il metodo su un modello lineare Gaussiano a spazio di stato (LGSSM) con:

• Matrici di covarianza del rumore di osservazione $V(\alpha)$ che variano da rango pieno a rango uno (singolare) al variare di un parametro $\alpha$.
• Confronto tra l'approccio proposto e l'approccio "classico" basato su QR.

Risultati osservati:

• L'approccio classico basato su QR fallisce completamente quando il rango è deficitario, restituendo valori invalidi o divergenti per i gradienti.
• L'approccio proposto calcola gradienti corretti e coerenti con la verosimiglianza marginale logaritmica, anche in presenza di rango deficitario e dimensioni non corrispondenti.
• I vettori tangenti calcolati coincidono perfettamente con la pendenza della curva di verosimiglianza (come mostrato nella Figura 1 del paper).

5. Significato e Impatto

Questo lavoro risolve un collo di bottiglia fondamentale nell'apprendimento automatico basato su modelli a spazio di stato.

• Robustezza: Permette l'uso di filtri di Kalman a radice quadrata in scenari realistici dove le matrici di covarianza possono essere singolari o il rumore può essere a basso rango, senza interrompere l'addestramento.
• Stabilità Numerica: Mantiene i vantaggi numerici dei filtri a radice quadrata (stabilità a lungo termine) abilitando al contempo l'ottimizzazione tramite gradienti.
• Generalità: La tecnica delle "Gramian Differentials" è un approccio generale che può essere applicato ad altri operatori di fattorizzazione dove l'output è definito solo a meno di una trasformazione ortogonale, ma la perdita dipende solo dal prodotto simmetrico.
• Adozione Pratica: Essendo compatibile con JAX e PyTorch, questa metodologia rende immediatamente possibile l'addestramento end-to-end di modelli complessi a spazio di stato in applicazioni moderne di machine learning.