Dynestyx: A Probabilistic Programming Library for Dynamical Systems

Riepilogo Tecnico: Dynestyx – Una Libreria di Programmazione Probabilistica per Sistemi Dinamici

Problema
I modelli a spazio di stato (SSM) rappresentano la formalizzazione standard per il trattamento bayesiano dei sistemi dinamici, con applicazioni critiche nella neuroscienza, nell'elaborazione dei segnali e nel machine learning. Nonostante la loro importanza teorica e l'utilità pratica, i sistemi dinamici si sono rivelati difficili da integrare nei moderni linguaggi di programmazione probabilistica (PPL). Di conseguenza, i metodi di inferenza allo stato dell'arte rimangono spesso confinati in codebase bespoke, creando attrito per i professionisti che tentano di seguire un rigoroso "workflow bayesiano". Questa separazione ostacola il confronto sistematico degli algoritmi di inferenza e limita l'accessibilità di metodi avanzati per la stima sia degli stati latenti che dei parametri del modello.

Metodologia
Gli autori introducono dynestyx, una libreria di programmazione probabilistica costruita come estensione di NumPyro. Dynestyx è progettata per fornire un supporto di prima classe agli SSM pur aderendo al principio PPL della "separazione delle preoccupazioni" (separation of concerns).

• Interfaccia Unificata: La libreria consente agli utenti di specificare prior arbitrari sia per sistemi dinamici a tempo discreto che a tempo continuo attraverso un'unica interfaccia unificata.
  • Tempo discreto: Definito da $x_{t+1} \sim p(x_{t+1} | x_t, u_t, t)$.
  • Tempo continuo: Definito da equazioni differenziali stocastiche (SDE) $dx_t = f(x_t, u_t, t) dt + L(x_t, u_t, t) d\beta_t$.
  • Osservazioni: Accoppiate allo stato tramite $y_t \sim p(y_t | x_t, u_t, t)$.
• Ortogonalizzazione dei Componenti: Dynestyx disaccoppia tre componenti distinti:
  1. Specificazione del Modello: Una funzione Python che definisce il prior congiunto $p(\theta, \phi)$ e la struttura dinamica (usando DynamicalModel).
  2. Stima della Verosimiglianza Marginale (ML): L'utente seleziona un algoritmo di inferenza (ad es., Filtro di Kalman, Particle Filter, Ensemble Kalman Filter) per approssimare l'intrattabile verosimiglianza marginale $p(y_{1:T} | \theta, \phi)$.
  3. Inferenza dei Parametri: Il log-posterior risultante viene passato alle routine di inferenza standard di NumPyro (ad es., HMC, NUTS, Inferenza Variazionale, SG-MCMC).
• Implementazione tramite Effect Handlers: La libreria utilizza gli effect handlers (ispirati a Plotkin e Pretnar, 2009) per interpretare l'istruzione dsx.sample. Ciò consente alla definizione del modello di rimanere agnostica rispetto allo specifico discretizzatore numerico o all'algoritmo di filtraggio utilizzato, differendo tutte le approssimazioni alla fase di inferenza.
• Supporto per Effetti Misti: La libreria implementa sistemi dinamici a effetti misti tramite una primitiva plate, consentendo l'inferenza su strutture dati gerarchiche compatibili con qualsiasi combinazione di stimatori di ML e metodi di inferenza dei parametri.

Contributi Chiave

1. Supporto SSM di Prima Classe in PPL: Dynestyx colma il divario tra la complessa teoria dei sistemi dinamici e la moderna programmazione probabilistica, rendendo accessibili i metodi allo stato dell'arte all'interno di un workflow standard di NumPyro.
2. Flessibilità Combinatoria: Ortogonalizzando la specificazione del modello, la stima della ML e l'inferenza dei parametri, dynestyx abilita un vasto spazio di progettazione di combinazioni algoritmiche. Gli utenti possono scambiare algoritmi di filtraggio (ad es., EKF, UKF, Particle Filter) e metodi di inferenza dei parametri (ad es., HMC, Variational Bayes) senza riscrivere il codice del modello.
3. Nuove Combinazioni di Metodi: La libreria implementa combinazioni di stimatori e metodi di inferenza che, a conoscenza degli autori, non sono apparse in letteratura. Queste includono SSM continuo-discreti accoppiati con MCMC informato dal gradiente (ad es., CD EnKF con MALA) e varie altre configurazioni ibride, come dettagliato nella tabella di mappatura del paper.
4. Integrazione con Ecosistemi Esistenti: Dynestyx si integra con librerie software esistenti come dynamax, cd-dynamax e cuthbert per fornire una vasta gamma di algoritmi di stima dello stato, inclusi il filtro di Kalman, lo smoothing RTS e il particle filtering differenziabile.

Risultati e Capacità
Il paper dimostra che dynestyx riesce a unificare la specificazione di processi generatori di dati idealizzati con strategie di inferenza flessibili. Le capacità chiave includono:

• Stima di Stato e Parametri: La capacità di eseguire l'inferenza su dati a effetti misti e generare stime sia per gli stati latenti che per i parametri con una quantificazione del principio di incertezza.
• Confronto Algoritmico: Il framework facilita il confronto diretto di diverse composizioni di metodi di inferenza, permettendo ai professionisti di identificare "regole empiriche" (rules of thumb) e patologie in modelli complessi, flessibili o potenzialmente errati.
• Compatibilità: I modelli risultanti sono pienamente compatibili con le routine JAX generiche, inclusi l'Inferenza Variazionale, HMC/NUTS e l'MCMC a gradiente stocastico.

Significato
Gli autori posizionano dynestx come uno strumento che rimuove l'attrito attualmente associato all'applicazione di workflow bayesiani ai sistemi dinamici. Fornendo una rappresentazione unificata per gli SSM, la libreria non solo rende i metodi avanzati più accessibili ai professionisti, ma crea anche un naturale banco di prova per l'innovazione metodologica. La capacità di comporre e confrontare sistematicamente gli algoritmi di inferenza permette la scoperta di nuovi insight riguardo al comportamento di modelli complessi, guidando potenzialmente ulteriori progressi nel campo dei sistemi dinamici bayesiani.