Dynestyx: A Probabilistic Programming Library for Dynamical Systems

Resumo Técnico: Dynestyx – Uma Biblioteca de Programação Probabilística para Sistemas Dinâmicos

Problema
Os modelos de espaço de estados (SSMs) servem como o formalismo padrão para o tratamento Bayesiano de sistemas dinâmicos, com aplicações críticas em neurociência, processamento de sinais e aprendizado de máquina. Apesar de sua importância teórica e utilidade prática, os sistemas dinâmicos têm se mostrado difíceis de integrar em linguagens de programação probabilística (PPLs) modernas. Consequentemente, os métodos de inferência de última geração permanecem frequentemente confinados a bases de código específicas, criando fricção para profissionais que tentam seguir um "fluxo de trabalho Bayesiano" rigoroso. Essa separação dificulta a comparação sistemática de algoritmos de inferência e limita a acessibilidade de métodos avançados para estimar tanto os estados latentes quanto os parâmetros do modelo.

Metodologia
Os autores introduzem o dynestyx, uma biblioteca de programação probabilística construída como uma extensão do NumPyro. O dynestyx é projetado para fornecer suporte de primeira classe para SSMs enquanto adere ao princípio de PPL de "separação de preocupações".

• Interface Unificada: A biblioteca permite que os usuários especifiquem priors arbitrários tanto para sistemas dinâmicos de tempo discreto quanto de tempo contínuo através de uma interface única e unificada.
  • Tempo discreto: Definido por $x_{t+1} \sim p(x_{t+1} | x_t, u_t, t)$.
  • Tempo contínuo: Definido por equações diferenciais estocásticas (SDEs) $dx_t = f(x_t, u_t, t) dt + L(x_t, u_t, t) d\beta_t$.
  • Observações: Acopladas ao estado via $y_t \sim p(y_t | x_t, u_t, t)$.
• Ortogonalização de Componentes: O dynestyx desacopla três componentes distintos:
  1. Especificação do Modelo: Uma função Python que define o prior conjunto $p(\theta, \phi)$ e a estrutura dinâmica (usando DynamicalModel).
  2. Estimativa da Verossimilhança Marginal (ML): O usuário seleciona um algoritmo de inferência (ex: Filtro de Kalman, Filtro de Partículas, Filtro de Kalman de Conjunto) para aproximar a verossimilhança marginal intratável $p(y_{1:T} | \theta, \phi)$.
  3. Inferência de Parâmetros: O log-posterior resultante é passado para rotinas de inferência padrão do NumPyro (ex: HMC, NUTS, Inferência Variacional, SG-MCMC).
• Implementação via Manipuladores de Efeito (Effect Handlers): A biblioteca utiliza manipuladores de efeito (inspirados em Plotkin e Pretnar, 2009) para interpretar a instrução dsx.sample. Isso permite que a definição do modelo permaneça agnóstica ao discretizador numérico ou algoritmo de filtragem específico utilizado, adiando todas as aproximações para a fase de inferência.
• Suporte a Efeitos Mistos: A biblioteca implementa sistemas dinâmicos de efeitos mistos via um primitivo de placa (plate primitive), permitindo a inferência sobre estruturas de dados hierárquicas compatíveis com qualquer combinação de estimadores de ML e métodos de inferência de parâmetros.

Principais Contribuições

1. Suporte de Primeira Classe para SSM em PPL: O dynestyx preenche a lacuna entre a teoria complexa de sistemas dinâmicos e a programação probabilística moderna, tornando métodos de última geração acessíveis dentro de um fluxo de trabalho padrão do NumPyro.
2. Flexibilidade Combinatória: Ao ortogonalizar a especificação do modelo, a estimativa de ML e a inferência de parâmetros, o dynestyx possibilita um vasto espaço de design de combinações algorítmicas. Os usuários podem trocar algoritmos de filtragem (ex: EKF, UKF, Filtros de Partículas) e métodos de inferência de parâmetros (ex: HMC, Inferência Variacional) sem reescrever o código do modelo.
3. Novas Combinações de Métodos: A biblioteca implementa combinações de estimadores e métodos de inferência que, até onde os autores sabem, não apareceram na literatura. Estas incluem SSMs contínuos-discretos pareados com MCMC informado por gradiente (ex: CD EnKF com MALA) e várias outras configurações híbridas, conforme detalhado na tabela de mapeamento do artigo.
4. Integração com Ecossistemas Existentes: O dynestyx integra-se com bibliotecas de software existentes, como dynamax, cd-dynamax e cuthbert, para fornecer uma ampla gama de algoritmos de estimativa de estado, incluindo filtragem de Kalman, suavização RTS e filtragem de partículas diferenciável.

Resultos e Capacidades
O artigo demonstra que o dynestyx unifica com sucesso a especificação de processos geradores de dados idealizados com estratégias de inferência flexíveis. As principais capacidades incluem:

• Estimativa de Estado e Parâmetro: A capacidade de realizar inferência sobre dados de efeitos mistos e gerar estimativas tanto para estados latentes quanto para parâmetros com quantificação de incerteza fundamentada.
• Comparação Algorítmica: O framework facilita a comparação direta de diferentes composições de métodos de inferência, permitindo que os profissionais identifiquem "regras de bolso" e patologias em modelos complexos, flexíveis ou potencialmente mal especificados.
• Compatibilidade: Os modelos resultantes são totalmente compatíveis com rotinas genéricas JAX, incluindo Inferência Variacional, HMC/NUTS e MCMC de gradiente estocástico.

Significância
Os autores posicionam o dynestyx como uma ferramenta que remove a fricção atualmente associada à aplicação de fluxos de trabalho Bayesianos a sistemas dinâmicos. Ao fornecer uma representação unificada para SSMs, a biblioteca não apenas torna os métodos avançados mais acessíveis aos profissionais, mas também cria um campo de teste natural para inovação metodológica. A capacidade de compor e comparar sistematicamente algoritmos de inferência permite a descoberta de novos insights sobre o comportamento de modelos complexos, potencialmente impulsionando novos avanços no campo dos sistemas dinâmicos Bayesianos.