Likelihood-Free Inference in State-Space Models with Unknown Dynamics
Abstract
We introduce a method for inferring and predicting latent states in the important and difﬁcult case of state-space models (SSM) where observations can only be simulated, and transition dynamics are unknown.
In this setting, the likelihood of observations is not available and only synthetic observations can be generated from a black-box simulator.
We propose a way of doing likelihood-free inference (LFI) of states and state prediction with a limited number of simulations.
Our approach uses a multi-output Gaussian process for state inference, and a Bayesian Neural
Network as a model of the transition dynamics for state prediction.
We improve upon existing LFI methods for the inference task, while also accurately learning transition dynamics.
The proposed method is necessary for modelling inverse problems in dynamical systems with computationally expensive simulations, as demonstrated in experiments with non-stationary user models.
Introduction
Conventional likelihood-free inference (LFI) methods [1, 2] estimate parameters θ of a static statistical model, given the dataset x ∗ and a black-box simulator g θ .
When the likelihood is intractable, these methods use synthetic datasets x ( i ) ∼ g θ ( x | θ ) produced by the simulator to assist the inference.
LFI has been successfully applied to identify parameters of complex real-world systems, such as ﬁnancial markets, [3, 4, 5], species populations [6, 7, 8] and cosmology models [9, 10, 11].
Many real-world systems, however, are time-dependent, and can be described as a state-space model (SSM) [12, 13], where observed measurements x t ∈ R n are emitted given a series of latent variables, the states, θ t ∈ R m .
When working with dynamical systems, one should account not only for an observation model, a simulator g θ , but for Markovian state transition dynamics h θ , as well, deﬁned as θ t +1 ∼ h
θ ( θ t +1 | θ t ) .
Therefore, traditional LFI methods, designed for static models, are ill-suited for dynamical systems, such as system simulation experiments in meteorology [14, 15], simulation-based forecasting in cosmology [16, 17] or adaptive modelling of human behaviour [18, 19].
Current LFI methods for dynamical systems [20, 21] proceed by assuming some simpliﬁed form of transition dynamics in an SSM.
Unless the SSM is both linear and Gaussian, state values are difﬁcult to estimate.
This issue has been addressed in a large and growing set of methods, including extended
Kalman ﬁlters (EKF) [22, 23, 24], GP-SSMs [25, 26], Bayes ﬁltering (BF) [27, 28] and sequential
Monte Carlo (SMC) methods [29, 30, 31].
Unfortunately, when the latent states have the additional requirement of being valid simulator parameters, these methods need a tractable likelihood that is not available when the observation model is black-box.
At the same time, LFI has been applied to SSMs, but instead of focusing on the problem we address here, the work has focused on more efﬁcient sampling-based methods [32], generation of a better matching statistics [33], on establishing theoretical convergence guarantees [34, 21, 35].
Despite these advances, inference with current LFI methods in SSMs is limited in cases when h t is non-linear and cannot be sampled or simpliﬁed, and when the simulators are computationally expensive.
In this paper, we introduce a method capable of obtaining likelihood-free state approximations and state predictions in discrete-time SSMs.
We assume that the state transition model is unknown, can be non-linear and the number of samples from the observation model is limited, when simulators are computationally expensive.
To solve the LFI problem in SSMs, we ﬁrst need to estimate simulator parameters for observed measurements with as few samples as possible, and then use these estimates to learn the transition dynamics.
We propose a method, which uses non-parametric multi-objective models for LFI of states, and non-parametric model of transition dynamics for learning state transition dynamics.
The contributions of the paper are:
• Focusing on problems where LFI has to be sample-efﬁcient, we improve upon the current LFI methods for the state inference task, while using samples from the state posteriors to accurately learn the state transition dynamics, as shown by comparisons with state-of-the-art SSMs inference techniques.
We demonstrate that the proposed method is especially suitable for applications with non-stationary user models.
• We propose sampling from transition dynamics model as a sample-efﬁcient strategy to determine where to run simulations next.
We show empirical comparisons with traditional techniques in settings where simulations are computationally expensive.
Likelihood-Free State Inference
For LFI, we follow a Bayesian optimization for LFI (BOLFI) approach [36], in which a Gaussian process (GP) surrogate is used for a discrepancy measure δ (e.g. Euclidean distance) between observed x ∗ and synthetic datasets x ( i ) .
This approach assumes that synthetic observations with the small discrepancy (smaller than a user-deﬁned threshold (cid:15) ) were produced by the simulator parameters that could plausibly replicate the observed datasets.
In this way, the likelihood can be approximated as p ( x
[Equation]
, where F ( · ) is a Gaussian CDF with the mean 0 and the variance 1, µ ( θ ) and ν ( θ ) are the mean and the standard deviation of a GP, and σ is a likelihood noise of the GP.
Because BOLFI actively chooses where to run simulations, its posterior approximation requires much fewer synthetic observations than other LFI methods.
However, BOLFI was not designed for SSMs, and hence, does not make use of any temporal information that is typical for
SSMs to improve quality of inference.
As an extension to BOLFI, we employ a multi-objective surrogate model for discrepancies δ t , thus considering multiple discrepancy objectives at the same time and leveraging information between consecutive states.
This approach allows using a discrepancy model of the previous state to infer the current state value, instead of simply discarding it.
There is an additional advantage of having a multi-objective surrogate in an SSM setting, namely we can have a much more ﬂexible surrogate for LFI than a traditional GP.
This does not need any additional data to ﬁt the surrogates, because all synthetic observations for discrepancy objectives can be shared across all state values.
When we have a new observation, we simply need to recalculate the discrepancy values for all synthetic datasets.
Once we have a trained surrogate for discrepancy objectives, we infer state posteriors p ( θ t | x t ) , similarly as in BOLFI.
There is an additional challenge for adapting multi-objective surrogates in SSMs, namely a high computational cost associated with considering too many objectives.
SSMs can potentially have hundreds of observations, and expanding the number of considered objectives may be detrimental for the performance of the surrogate.
We avoid this problem by limiting the number of objectives the surrogate can have.
Instead of considering all available time-steps as objectives, we propose to consider only L recent objectives by gradually including new objectives and discarding old ones that have little impact on current states.
The size of this moving window depends on how rapidly the transition dynamics change.
As the size L of the window grows, the model becomes less sensitive to the noise from the dynamics, at the cost of increased computations and decreased adaptability to the most recent state transitions.
The moving window reduces the number of objectives L considered at a time, making multi-objective modelling for SSMs feasible.
In the experiments, we used a linear model of coregionalization (LMC)
[37] as a surrogate, with a moving window size of 2.
Learning State Transition Dynamics
Once we have approximated the posteriors p ( θ t | x ∗ t ) with LFI, we can use their empirical samples to learn a state transition model.
This model should be able to train from noisy samples of LFI posterior approximations, and be ﬂexible enough to ﬁt arbitrary function families the dynamics may follow.
Thus, the transition model should be Bayesian and non-parametric.
Such a model would take the uncertainty associated with posterior approximations into account and be ﬂexible enough to follow possibly non-linear transition dynamics.
We propose to train this model in an autoregressive fashion by forming a training set from paired posterior samples.
For each SSM time interval, we group consecutive state posterior samples in a training set, and expand it when new state posteriors become available.
This way, the transition model does not need to be retrained when new observations present themselves and can be actively used throughout state inference for determining where to run simulations next.
This can be done, by sampling the predictive posterior p ( θ T +1 | x ∗ T ) from the trained model
ˆ h θ through p ( θ T +1 | x ∗ T ) = (cid:82) ˆ h
θ ( θ T +1 | θ T ) p ( θ T | x ∗ T ) d θ T .
Thus, the state transition model ˆ h θ does not inform state posteriors directly, but provides simulation candidates for the LFI surrogate.
In the experiments, we used a
Bayesian Neural Network (BNN) [38, 39] as a transition model for the state dynamics.
Experiments
We assess the quality of our method, LMC-BNN, for the inference and prediction tasks in a series of SSM experiments, where a simulator serves as the observation model g θ .
For the inference task, we compare the quality of our estimates against other LFI methods: BOLFI [36], SNPE [40], SNLE [41], SNRE [42].
We use a ﬁxed simulation budget for all these methods, with 20 simulations to initialize the models, and then with 2 additional simulations for each new time-step.
However, for the neural density estimation (NDE) approaches (SNPE , SNLE, SNRE), we provided all simulations at once, since they do not need additional simulations when observations change.
For the prediction task, we sampled trajectories from the transition model and evaluated them against trajectories from the true dynamics.
We performed these experiments in SSMs with simulators that have tractable likelihoods, providing them to the state-of-the-art SSM inference methods, GP-SSM [26, 43] and PR-SSM [44], while our method was still doing LFI.
For all methods, we provided 50 observations, and then sampled trajectories that had the same length of 50 time-steps.
We also supplied two variations of our method that differ only in the way the next simulations are sampled: LMC-BLR, where samples were taken from a Bayesian Linear Regression (BLR) model that was used as a local linearized model of the transition dynamics; and LMC-qEHVI, where a popular acquisition function for multitask BO, q-Expected Hypervolume Improvement (qEHVI) [45], was used to provide samples.
4.1 The State-Space Models
In this section, we present two case studies with non-stationary user models and three SSMs with tractable likelihoods.
In user modelling experiments, we simulated behavioural data from humans that completed two different tasks: uniform manifold approximation and projection (UMAP) [46]
[Table] parametrization and eye movement control for gaze-based selection.
Our goal in these experiments was to infer parameters of user models that generated the behavioural data, with three state dimensions in each user model.
In the experimental results, these case studies are referred to as UMAP and Gaze, respectively.
In addition to non-stationary user models [Table] we also experimented with three models with tractable likelihoods common in SSM literature: linear Gaussian (LG), non-linear non-Gaussian (NN) and stochastic volatility (SV) models.
In the LG model, the state transition dynamics and observation model are both linear, with high observational noise.
The NN model is a modiﬁed version (leaving only one unique solution for each observation) of a popular non-linear SSM [47].
Lastly, SV models are widely used for predicting volatility of asset prices [48, 49].
We use the speciﬁc model by [50].
The dimensionality of states in these models ranges from one to three, starting from LG and moving to SV.
More detailed descriptions of all SSMs can be found in the Appendix.
4.2 Results and Discussion
The results for the inference and prediction tasks are presented in Tables 1 and 2, respectively.
These tables report Root Mean Squared Error (RMSE), as a performance measure.
In the inference task, LMC-based methods clearly outperformed BOLFI and NDE approaches.
This indicates that considering multiple objectives at the same time was beneﬁcial for the state inference, and that the model actually leverages information from consecutive states without hindering the performance.
Additionally, it can be seen that LMC-based methods performed differently, which can be only attributed to how the next simulations were chosen, as the surrogate was exactly the same in all three methods.
As results show, having BNN as a model for state transition was beneﬁcial for experiments with non-stationary user models, while having BLR was more preferable for the simpler models.
This suggests that BLR is expressive enough to replicate simple transitions, but struggles with more complex ones, for which BNN was more suitable.
The comparisons with GP-SSMs and PR-SSMs for learning transition dynamics show that our method learns the dynamics more accurately.
The SSMs methods that were supposed to provide ground truth performance showed worse results than BLR and BNN approaches.
This can be only explained by the lack of observations for learning state transitions by the SSM methods, which also explains the high variance in the sampled trajectories from these methods.
As for comparisons between BLR and
BNN, BLR performs better only in LG and SV models, while BNN performs better in more complex case studies.
Moreover, it should be noted that trajectory sampling from BLR is possible only by retaining all local linearization, which is a far more limiting approach than having BNN.
Therefore,
BNN is more preferable as a model for transition dynamics.
In all experiments, we attribute the success of the proposed LMC-BNN method to a more ﬂexible multi-output surrogate, and a more efﬁcient way of choosing simulation candidates.
The LMC allows multi-ﬁdelity modelling and leveraging information from multiple consecutive time-steps, unlike standard GPs.
At the same time, samples from the transition model provide better candidates for simulations than the alternatives.
The ﬂexible surrogate along with adaptive acquisition make our method particularly suitable for online settings, where only a handful of samples are allowed per time-step. [Table] Acknowledgments and Disclosure of Funding
This work was supported by the Academy of Finland (Flagship programme: Finnish Center for
Artiﬁcial Intelligence FCAI; grants 328400, 319264, 292334) and UKRI Turing AI World-Leading
Researcher Fellowship EP/W002973/1. HP was also supported by European Research Council grant 742158 (SCARABEE, Scalable inference algorithms for Bayesian evolutionary epidemiology). Computational resources were provided by the Aalto Science-IT Project. A The State-Space Models with Tractable Likelihoods
In this section, we present details on the three SSMs with tractable likelihoods that were used to assess the quality of state transition models in the experiments. For all three SSMs, we deﬁne transition dynamics and an observation model along with priors for LFI. As observations, we used datasets of 10 points, and their mean and standard deviation as summary statistics. The objective in Bayesian
Optimization (BO) was the logarithm of Euclidean distance between the observed and simulated summary statistics. Linear Gaussian (LG). In the LG model, state transition dynamics (Figure 1a) and an observation model are both linear
[Equation]
with added Gaussian noise v t ∼ N (0 , 1 2 ) and w t ∼ N (0 , 10 2 ) .
The initial state value for the transition dynamics is θ 0 = 100 .
The prior for the state values is θ ∼ Unif (0 , 15) ∈ R .
Non-linear non-Gaussian (NN). The NN model is a popular non-linear SSM [47], where the transition dynamics (Figure 1b) and observation model are
[Equation]
with added Gaussian noise v t ∼ N (0 , 10) and w t ∼ N (0 , 10) . The initial state value for the noise standard deviation θ h, 0 is 0 with the prior θ h ∼ Unif ( − 30 , 30) ∈ R .
Stochastic volatility (SV) models are widely used for predicting volatility of asset prices [48, 49].
Here, we use the model by [50] that speciﬁes transition dynamics (Figure 1c) of volatility θ v as θ v,t +1 = ( z t − z t +1 + k (cid:88)
j = e j ) , z t +
[Equation]
with c 1: k ∼ Unif ( t, t + 1) , e 1: k ∼ Expon (0 . 5 / 0 .
25 2 ) and λ = 0 .
The random increases of volatility are regulated by the Poisson distributed variable k ∼ Poisson (0 .
5 λ 2 / 0 .
The observation model for the log-return of an asset x t ∈ R follows
[Equation]
where ε t ∼ N (0 , 1) ; θ µ = 0 and θ β = 0 remain the same, while the volatility θ v follows the transition dynamics, starting with the initial value for the volatility θ v, 0 of 1 .
We set the following priors for LFI of state values: θ µ ∼ Unif ( − 2 , 2) ∈ R
[Equation]
B Technical Details B.1 Eye Movement Control for Gaze-Based Selection
An observation model in the gaze selection experiment is a simulated environment, where a humansubject is modelled through a reinforcement learning agent.
In this environment, the agent searches for a target on a 2D display, where the target location, actions, observations, and beliefs of the agent are represented by two coordinates { c 1 , c 2 } , c 1 , c 2 ∈ [ − 1 , 1] on the display. At each episode of the environment e , the agent receives noisy observations of the target o e = N ( q , θ s × E ( e )) and moves
[Figure]
the gaze to a new location a e = N ( PPO ( o e , θ om × A ( e )) .
The beliefs b e about the target location q are updated according to
[Equation]
where σ ( o,e ) and σ ( b,e ) are observation and belief uncertainties respectively, A ( e ) is the amplitude, and E ( e ) is the eccentricity of the saccade at the episode e .
The user model was trained on000 episodes using the Proximal Policy Optimization (PPO) algorithm [51], provided by the Stable Baselines3 library [52]. We used default parameters, a multilayer perceptron policy and a clipping parameter set to 0.15.
The environment was implemented by Chen et al. [53] in Python with the Open AI gym framework [54].
In the UMAP parameterization model, the ground truth for state values is not available, because the human-subject cannot specify the ideal embeddings. Therefore, we approximate the ground truth by using ABC with rejection sampling.
Usually, this requires running millions of simulations for each time-step, but we make use of the weighted form of the preference score that allows us to use the same simulations across all time-steps.
We simulate 1 , 500 , 000 embeddings with state values sampled from the prior, and then calculate their corresponding U ( i ) and P ( i ) . For each time-step t , we calculate the preference scores δ and retain only 0 .
06% of those states that showed the lowest δ values. Finally, we use a Gaussian kernel density estimator (KDE) on the retained state values, and maximize the corresponding PDFs to ﬁnd the estimations of the ground truth. The bandwidth b of the kernel was calculated according to a Scott’s rule [55] b = n − 1 d +4 , where n is the number of data points and d is the number of dimensions.
The preference score was computed as a weighted sum between the relative validity U ( i ) and the classiﬁcation accuracy P ( i ) on the validation set. The relative validity U ( i ) is an approximation of the Density Based Cluster Validity (DBCV) score [56], which is often used as a quality measure of clustering. Intuitively, it shows how separable all the clusters are.
We use the HDBSCAN* package [57] and the HDBSCAN Boruvka KDTree [58] algorithm to cluster the points of the embeddings.
We set the smallest size grouping to 60, the density parameter of clusters to 10 and leave the rest parameters to their default values.
The classiﬁcation accuracy P ( i ) was calculated for the C-Support Vector Classiﬁer (SVC) [59, 60] with the Scikit-learn package [61] and default parameters.
The embeddings for the UMAP parameterization model were computed for the handwritten digit dataset [62]. The dataset was randomly split on the training and validation sets, resulting in 1198 and 599 8x8 pixel images of digits in 10 digit classes.
The UMAP algorithm was implemented with the UMAP-learn package [63]. B.3 Implementation Details of Methods LFI methods from Sections B.3.1-B.3.3 were integrated in the Engine for Likelihood-Free Inference (ELFI) [64] for application and further development.
GP surrogate . The implementation for the GP surrogate was provided by the GPy package [65]. It was initiated with zero mean function, and with the following priors for the RBF kernel lengthscale l , its variance σ 2 k , and the variance of the bias term σ 2 b : l ∼ Gamma (cid:18) θ max − θ min ,(cid:19) , σ 2 k ∼ Gamma (cid:18) max ( δ ( i ) ) ,(cid:19) , σ 2 b ∼ Gamma
[Equation] [Equation]
[Equation] [Equation] [Equation]
where θ min , θ max are the lower and upper bounds for θ ( i ) , ( θ ( i ) , δ ( i ) ) are initial training points, and 1 is an all-ones vector. The GP was minimized by using the SCG optimizer [66] on the GP negative log-likelihood with a maximum number of iterations of 50.
All inputs θ ( i ) for the GP were centred and normalized.
LCBSC acquisition .
The BOLFI implementation uses LCBSC [67, 68] as an acquisition function, which chooses points that minimize the lower conﬁdence bound (LCB)
[Equation]
where β t is the conﬁdence parameter, δ = 0 . 1 is the inverse exploration rate and d is an input dimension. Posterior sampling . The BOLFI posterior was obtained by weighting the prior samples and using corresponding unnormalized likelihoods as weights
[Equation]
with F ( · ) being a Gaussian CDF with the mean 0 and the variance 1, and where (cid:15) is the minimum of the GP surrogate mean function µ ( · ) obtained by using the L-BFGS-B optimizer [69] with a maximum number of iterations of 1000.
B.3.2 Muti-Objective LFI
LMC surrogate .
The LMC model was implemented in BoTorch [70].
Its latent GPs were initialized with linear mean functions µ ( x ) = kx
+ b , and RBF kernels.
The lengthscales of the kernels were parameterized in log scale, and initialized with 0 along with the constant b of the mean function.
For the RBF kernel, ARD was also enabled to assign separate lengthscales for each input dimension.
GP latents were also initialized with 50 inducing points uniformly sampled inside the input bounds for each latent GP.
The LMC training used Adam optimizer from the tensor computation package
PyTorch [71] with a learning rate of 0.1 to minimize the variational evidence lower bound (ELBO). The optimized parameters included LMC coefﬁcients, inducing points locations, and hyperparameters for the mean functions and kernels. We used the default training step size and 1000 epochs for training. All inputs for the LMC were centred and normalized.
qEHVI acquisition .
The qEHVI acquisition function used a Quasi-MC sampler [72] with scrambled
Sobol sequences [73] and a sample size of 128.
The reference point that was used for calculating the hypervolume for each objective was set to -5.
The acquisition points were acquired sequentially using successive conditioning [74] with a maximum number of restarts of 20, a batch size limit of 10, and a maximum number of iterations for the optimizer of 200.
Bayesian neural network (BNN)
[75] was built from stacked Bayesian layers implemented in torch zoo (BLiTZ) [39].
We used an architecture with 2 hidden layers, where each layer had 256 nodes.
During the training process, stochastic gradient descent [76] was used with a learning rate of 0.001 for minimizing the ELBO loss with squared L2 norm.
The loss was calculated based on 10 samples from the model, for each of 100 batches of training data in a single epoch.
The training data was randomly selected from previously stored approximated posterior samples from all states (1000 samples per each state) with replacement, resulting in a total of 1 000 000 points, where one point was a pair of consecutive state values.
The training data was updated each time the moving window moved.
Bayesian linear regression (BLR) is deﬁned as {+θ t +
 = θ Tt β + (cid:15) , where (cid:15) ∼ N (0 , σ 2 I ) . The hyperparameters σ ∈ R+} and β ∈ R m × m were inferred with maximum likelihood estimation (MLE) [77] of
[Equation]
The BLR model was implemented with the probabilistic programming package, PyMC3 [78] that used 100 random samples from three latest inference steps.
The pairs ( θ t − 1 , θ t ) and ( θ t − 2 , θ t − 1 ) pairs were used as inputs for training, and the normal distribution was chosen as a prior family of the BLR parameters.
The model was ﬁtted by using the No-U-Turn Sampler (NUTS) [79] with two chains of 2000 samples, 2000 tuning iterations, and a target acceptance rate of 0.85.
Posterior sampling . Similar to BOLFI, the posterior was sampled by using an importance-weighted resampling. Because the main model was implemented in PyTorch, we used the Adam optimizer to learn the threshold (cid:15) . It used a learning rate of 10 − 4 and 100 optimizing iterations.
The optimization started at the parameter point that produced a synthetic observation with the smallest discrepancy.
B.3.3 Sequential Neural Estimators
All three SNEs (SNPE, SNLE and SNRE) and their corresponding surrogate models were implemented in the simulation-based inference [80] and PyTorch [71] packages with default parameters.
In all three methods, Adam optimizer with the learning rate of 5 × 10 − 4 and the training batch size of 50 in 20 epochs were used for training the surrogate.
The total gradient norm was clipped in order to prevent exploding gradients at a value of 5.0, and z-score passing was used for surrogate model inputs and outputs.
For SNPE [81] and SNLE [41], the masked autoregressive ﬂow (MAF) surrogate was used for approximating the posterior p ( θ | x ) and likelihood, p ( x | θ ) respectively.
The MAF consisted oftransformations with 50 hidden features in each of 2 blocks.
Each autoregressive transformation had tanh activation along with batch normalization.
In contrast, SNRE [42] approximates the ratio p ( θ , x ) p ( θ ) p ( x ) , where a residual network [82] is used as a classiﬁer trained to approximate likelihood ratios.
The goal of the classiﬁer is to predict which of the ( θ , x ) pairs was sampled from p ( θ , x ) and which from p ( θ ) p ( x ) .
The residual network had 50 hidden features in 2 residual blocks with 10 ( θ , x ) pairs to use for classiﬁcation.
B.3.4 Recurrent State-Space Models with GP Transitions
In the experiments, we used two variants of recurrent state-space models with GP transitions: GPSSM with a variationally coupled dynamics and trajectories, in which the variational inference scheme for GP transition dynamics is used [43], and probabilistic recurrent state-space model PRSSM [44], which uses doubly stochastic variational inference for efﬁcient incorporation of latent state temporal correlations.
Both methods were implemented in the GPt package [43].
The GPs were using Matern32 kernels with linear mean functions, along with 50 randomly sampled inducing points.
The number of latent dimensions was set equal to the number of simulator parameters of the observation model.
The ELBO loss was calculated from 10 posterior samples and optimized with
Adam using a learning rate of 0 .
001 in 3000 iterations.