D ISENTANGLED GENERATIVE MODELS FOR ROBUST DYNAMICAL SYSTEM PREDICTION
Anonymous authors
Paper under double-blind review
A BSTRACT
Deep neural networks have become increasingly of interest in dynamical system prediction, but out-of-distribution generalization and long-term stability still remains challenging.
In this work, we treat the domain parameters of dynamical systems as factors of variation of the data generating process.
By leveraging ideas from supervised disentanglement and causal factorization, we aim to separate the domain parameters from the dynamics in the latent space of generative models.
In our experiments we model dynamics both in phase space and in video sequences and conduct rigorous OOD evaluations 1 .
Results indicate that disentangled VAEs adapt better to domain parameters spaces that were not present in the training data. At the same time, disentanglement can improve the long-term and out-of-distribution predictions of state-of-the-art models in video sequences.
I NTRODUCTION
The robust prediction of dynamical systems behaviour remains an open question in machine learning, and engineering in general.
The ability to make robust predictions is important not only for forecasting systems of interest like weather (Garg et al., 2021; Ravuri et al., 2021) but even more so because it enables innovations in ﬁelds like system control, autonomous planning (Hafner et al., 2018) and computer aided engineering (Brunton et al., 2020).
In this context, the use of deep generative models has recently gained signiﬁcant traction for sequence modelling (Girin et al., 2020).
Robustness of machine learning models can be considered along two axes: long-term prediction and out-of-distribution (OOD) performance.
Accurate long-term prediction can be notoriously difﬁcult in many dynamical systems, because error accumulation can diverge in ﬁnite time (Zhou et al., 2020; Raissi et al., 2019), a problem that even traditional solvers can suffer from.
More importantly, machine learning techniques are known to suffer from poor OOD performance (Goyal & Bengio, 2020), i.e. when they are employed in a setting they had not encountered during the training phase.
Before addressing the OOD problem, we must ﬁrst deﬁne what constitutes as OOD in dynamical systems.
We start by the observation that even simple dynamical systems, i.e the swinging pendulum or the 𝑛 -body system, can have multiple continuous parameters that affect their evolution.
These parameters can be manifested as differential equation coefﬁcients, boundary or initial conditions etc.
Our starting point is to consider distinct ranges of those parameters as separate domains.
Under this view, it becomes apparent why OOD prediction of dynamical systems can be hard: capturing the whole range of those parameters in a single training set is unrealistic (Fotiadis et al., 2020) and further inductive biases are required (Miladinovi´c et al., 2019; Bird & Williams, 2019; Barber et al., 2021).
From a dynamical system point of view, different parameters can produce widely different trajectories in phase space.
A motivating example can be bifurcations which occur when a small change in the parameters of a system causes a sudden qualitative change in its behaviour.
We focus on the inductive bias of disentangled representations for which the dynamics are separated from the domain parameters.
Many approaches based on the use of neural networks try to jointly learn the dynamics and the physical parameters, which results in convoluted representations and usually leads to overﬁtting (Bengio et al., 2012).
System identiﬁcation can be used to extract parameters, but reproducing our experiments at requires knowledge of the underlying system to be computationally effective (Ayyad et al., 2020).
We, instead, leverage advances in Variational Autoencoders (VAEs) (Kingma & Welling, 2014) that enable learning disentangled representations.
Disentanglement enables different latent variables to focus on different factors of variation of the data distribution, and has been applied in the context of image generation (Higgins et al., 2017; Kim & Mnih, 2018).
This can be extended to modelling dynamical systems by looking at disentanglement from a causal perspective: from all the generative models which can have the same marginal distribution, identify the one with the true causal factors.
To map this idea to sequence modelling we treat the domain parameters of a dynamical system as factors of variation.
Recent ﬁndings (Locatello et al., 2018) emphasize the vital role of inductive biases from models or data for useful disentanglement.
Unsupervised disentanglement, based on the assumption of domain stationarity, is a promising direction (Miladinovi´c et al., 2019; Li & Mandt, 2018).
Nevertheless, this leaves a wealth of ground truth domain parameters, which can be cheaply collected in simulated datasets.
This type of privileged information originating from simulations has been shown to be effective for domain adaptation in computer vision tasks (Saraﬁanos et al., 2017; Lee et al., 2018).
We thus use supervised disentanglement (Locatello et al., 2019) by treating the ground truth domain parameters from simulations as privileged information which, to the best of our knowledge, has not been applied to dynamical system prediction previously.
Contributions While others have treated domain parameters as factors of variation in the data distribution, our work is the ﬁrst, to the best of our knowledge, that explicitly uses privileged information from simulated data to disentangle those domain parameters from dynamics in a supervised way.
We furthermore conduct experiments both in the low-dimensional phase space of 3 dynamical systems and the high-dimensional video rendering of a swinging pendulum.
Disentanglement has, in the past, been mostly applied to VAEs because they are easily amenable to it.
The problem is that these models usually lack in competitive performance.
In the case of video sequences, we additionally apply disentanglement on a hybrid model with both stochastic and deterministic parts (Hafner et al., 2018).
In doing so, we not only assess disentanglement on a generative model outside boundaries of VAEs but furthermore we do it on a model which is considered state-of-the-art in long-term video prediction (Saxena et al., 2021).
In all cases, the prediction performance is assessed both in-distribution and also in OOD settings of increasing degrees of distribution shift.
To our understanding, this is the ﬁrst time such a rigorous OOD test is performed.
Our results in phase-space demonstrate that disentangled models can better capture the variability of dynamical systems compared to baseline models both in-distribution and OOD.
In modelling dynamics in video sequences, results indicate that disentanglement is beneﬁcial both for long-term prediction and OOD prediction.
Limitations This work focuses on dynamical system prediction.
While the results can potentially open up many applications in general time-series modelling, this is out of the scope of this work.
We prioritize to empirically study OOD downstream task performance and the inspection of the disentangled representations with appropriate metrics is left out of scope in this work.
R ELATED W ORK
VAEs and disentanglement Disentanglement aims to produce representations where separate factors of variation in the data are encoded into independent latent components.
This can be seen as ﬁnding the true causal model of the data.
While supervised disentanglement in generative models is a long-standing idea (Mathieu et al., 2016), information-theoretic properties can be leveraged to allow unsupervised disentanglement in VAEs (Higgins et al., 2017; Kim & Mnih, 2018).
The impossibility result from (Locatello et al., 2018) suggested that disentangled learning is only possible by inductive biases coming either from the model or the data.
Hence, the focus shifted back to semi- or weaklysupervised disentanglement approaches (Locatello et al., 2019; 2020).
While most of these methods focus on disentanglement metrics, we opt to directly assess using a downstream prediction task.
Disentanglement in sequence modelling While disentanglement techniques are mainly tested in a static setting, there is a growing interest in applying it to sequence dynamics.
Using a bottleneck based on physical knowledge, Iten et al.
(2018) learn an interpretable representation that requires conditioning the decoder on time, but it can return physically inconsistent predictions in OOD data (Barber et al., 2021).
Deep state-space models (SSMs) have also employed techniques for disentangling content from dynamics (Fraccaro et al., 2017; Li & Mandt, 2018), but, focus mostly on modelling variations in the content, failing to take dynamics into account.
In hierarchical approaches
Encoder x 1 x...x n
Decoder x n+1 x n+ ...x n+o z zk zd ...
Predictionloss KL-Divergence
Supervised disentanglement
Encoder
Decoder z zk zd ...
Predictionloss KL-Divergence
Supervised disentanglement
[Figure]
(Karl et al., 2017), different layers of latent variables correspond to different timescales: for example, in speech analysis for separating voice characteristics and phoneme-level attributes (Hsu et al., 2017).
In an approach similar to our work, Miladinovi ´ c et al.
(2019) separate the dynamics from sequencewide properties in dynamical systems like Lotka-Volterra, but do so in an unsupervised way which dismisses a wealth of cheap information and only assesses OOD generalization in a limited way.
Feed-forward models for sequence modelling Deep SSM models are difﬁcult to train as they require non-trivial inference schemes and a careful design of the dynamic model (Krishnan et al., 2015; Karl et al., 2017).
Feed-forward models, with necessary inductive biases, have been used for sequence modelling in dynamical systems (Greydanus et al., 2019; Fotiadis et al., 2020). In works like Hamiltonian Neural Networks Greydanus et al.
the domain is ﬁxed; together with Barber et al.
(2021), our work is an attempt in tackling domain variability.
Privileged information for domain adaptation.
Using privileged information during training has been shown to help with domain adaptation in computer vision tasks.
Using segmentation masks of simulated urban scenery can improve semantic segmentation on the target domain (Lee et al., 2018), while clip art data can help with domain transfer in an action recognition task (Saraﬁanos et al., 2017).
M ETHODS 3. V ARIATIONAL A UTOENCODERS
Variational autoencoders (VAEs) (Kingma & Welling, 2014) offer a principled approach to latent variable modeling by combining a variational inference model 𝑞 𝜙 ( z | x ) with a generative model 𝑝 𝜃 ( x | z ) .
As in other approximate inference methods, the goal is to maximize the evidence lower bound (ELBO) over the data:
[Equation]
The ﬁrst part of the ELBO is the reconstruction loss (in our case the prediction loss) and the second part is the Kullback-Leibler divergence that quantiﬁes how close the posterior is to the prior.
Design choices for the model We use an isotropic unit Gaussian prior 𝑝 ( z ) =
N ( z | 0 , I ) which helps to disentangle the learned representation (Higgins et al., 2017).
The approximate posterior (encoder) distribution is a Gaussian with diagonal covariance 𝑞 𝜙 ( z | x ) =
N ( z | µ 𝑧 , 𝚺 𝑧 ) , allowing a closed form KL-divergence, while the decoder has a Laplace distribution 𝑝 𝜃 ( x | z ) =
Laplace ( x | µ 𝑥 , 𝛾 I ) with constant diagonal covariance 𝛾 > 0 , which is tuned empirically.
This leads to an 𝐿 1 loss that provides improved results in some problems (Mathieu et al., 2018) and empirically works better in our case.
The parameters µ 𝑧 ≡ µ 𝑧 ( x ; 𝜙 ) , 𝚺 𝑧 ≡ diag [ σ 𝑧 ( x ; 𝜙 )] 2 , and µ 𝑥 ≡ µ 𝑥 ( z ; 𝜃 ) are computed via feed-forward neural networks. [Figure]
3. D ISENTANGLEMENT OF DOMAIN PARAMETERS IN LATENT SPACE
Apart from the disentanglement that stems from the choice of prior 𝑝 ( z ) , we explicitly disentangle part of latent space so that it corresponds to the domain parameters of each input sequence.
We achieve this by using a regression loss term L ξ ( z 1: 𝑘 , ξ ) between the ground truth factors of the domain parameters ξ ∈ R 𝑘
and the output of the corresponding latents, z 1: 𝑘 .
We, empirically, opted for an 𝐿 1 loss, corresponding to a Laplacian prior with mean ξ and unitary covariance.
Previous methods have reported that binary cross-entropy works better than 𝐿 2 (Locatello et al., 2019) but this does not ﬁt well in a setting like ours.
We hypothesize that BCE works better because of the implicit scaling.
To address this, we propose applying a function G( 𝜇 𝑧 𝑖 ) which linearly scales the 𝜇 𝑧 𝑖 between the min and max values of the corresponding factor of variation:
[Equation]
where 𝜉 𝑖 are the domain parameters and their corresponding minimum and maximum values of domain parameters from the training set.
In all cases, the regression term is weighted by a parameter 𝛿 which is empirically tuned.
Plugging these choices in results in the following loss function:
[Equation]
Using the reparameterization trick (Kingma & Welling, 2014)
[Equation]
the loss is amenable to optimzation by stochastic gradient descent, with batch size 𝑛 .
The model architecture can be seen in Figure 1(left)
[Equation]
3. D ISENTANGLEMENT FOR VIDEO DYNAMICS
We further investigate the effect of disentanglement in video sequence dynamics.
To this end, two generative models are used.
The ﬁrst is derived from the VAE formulation of the previous section and is called CNN-VAE and is similar to the VAE with the addition of a convolutional encoder and a decoder.
The encoder projects the input frames down to a low-dimensional space which can be thought as equivalent to the phase space of the system.
A VAE is applied in this projection to predict in the future coordinates in the ”phase space”.
The decoder then maps the predictions of the VAE back to pixel space.
The schematic of the model can be seen in Figure 1(right).
The second model we use is the Recurrent State Space Model (RSSM) which has been successfully used for planning (Hafner et al., 2018).
Since RSSM is a hybrid model combining deterministic and variational components, it allows us to assess disentanglement outside the limited scope of VAEs.
Furthermore, using a state-of-the-art model in long-term video prediction (Saxena et al., 2021), allows [Figure] [Figure] us to identify the limits of applying disentanglement in competitive models.
The loss function we use shares the same formulation as in the original work of Hafner et al. (2018) with the addition of the supervised disentanglement loss.
Since in the RSSM formulation there are latent variables for each time-step, we apply a disentanglement loss on all of them, which empirically is set to be 𝐿 2 :
[Equation]
(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32) (cid:123)(cid:122) (cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32) (cid:125)
[Table]
[Equation]
[Table] (cid:125)
[Table]
+ 𝛿 E 𝑞 ( s 𝒕 | o ≤ 𝒕 ) (cid:104)(cid:13)(cid:13)(cid:13)
ξ − s ( 1: 𝑘 ) 𝒕 (cid:13)(cid:13)(cid:13)(cid:105) (cid:124) (cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32) (cid:123)(cid:122) (cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32)(cid:32) (cid:125)
[Table]
(cid:170)(cid:174)(cid:174)(cid:174)(cid:174)(cid:174) (cid:172) (4)
Where o t is the observations, s t the stochastic latent variables at time 𝑡 , ξ are the 𝑘 dimensional domain parameters and 𝛿 tunes the supervised disentanglement strength.
E XPERIMENT - ODE P HASE S PACE D YNAMICS 4. D ATASETS
In the phase-space experiments we compare the models on three well studied dynamical systems, the swinging pendulum, the Lotka-Volterra equations used to model prey-predator populations, and the planar 3-body system:
Pendulum 1.
Pendulum  5 Prey 1.
P r e d a t
o r
Lotka-Volterra
Prey 1.
P r e d a t
o r
Lotka-Volterra  1 x 1.
y 3-body system   x 1.
[Figure]
[Figure]
[Equation]
The systems where chosen for varied complexity in terms of degrees of freedom, number of ODE equations and factors of variation.
For the pendulum we consider one factor of variation, its length 𝑙 ; Lotka-Volterra has 4 factors of variation 𝛼, 𝛽, 𝛾, 𝛿 and the 3-body system has also 4 factors of variation 𝐾 1 , 𝑚 1 , 𝑚 2 , 𝑚 3 .
Factors are drawn uniformly from a predetermined range which is the same between the training, validation and test sets.
To further assess the OOD prediction accuracy, we create two additional test sets with factor values outside of the original range.
We denote these datasets as OOD Test-set Easy and Hard, representing a smaller and bigger shift from the original range.
As a visual example, the distribution of the factors of variation for the Lotka-Volterra system is illustrated in Figure 9 of the Appendix.
The data were additionally corrupted with Gaussian noise.
Dataset details can be found on Table 1 of the Appendix.
4. M ODELS AND TRAINING
The main goal of these experiments is to assess whether OOD prediction can be improved by disentangling dynamical system parameters in the latent space of VAEs.
We opt to use simple models to allow more experiments and comparisons.
Our main baseline is the VAE upon which we propose two enhancements that leverage supervised disentanglement.
The ﬁrst VAE model, termed VAE-SD uses supervised disentanglement without a scaling function while the second model termed VAE-SSD uses an additional linear scaling function G( 𝜇 𝑧 𝑖 ) for the latent variable mean vector 𝜇 𝑧 𝑖 , as described in Section 3.2.
Another baseline is a multilayer perceptron (MLP) autoencoder which allows comparison with a deterministic counterpart of the VAE.
We also use supervised disentanglement on the latent neurons of the MLP, a model we refer to as MLP-SD.
This enables us to assess if the privileged information can improve deterministic models.
Lastly, we include an LSTM model, a popular choice for low dimensional sequence modelling (Yu et al., 2019), as a representative recurrent method.
Early experiments revealed a signiﬁcant variance on the performance of the models, depending on hyperparameters.
Under these conditions, we took various steps to make model comparisons as fair as possible.
Firstly, all models have similar capacity in terms of neuron count.
Secondly, we tune various hyperparameter dimensions, some of which are shared, while others are model-speciﬁc.
Third, we conduct a thorough grid search on the hyperparameters to avoid undermining a model (details can be found in Tables 3, 4 and 5 of the Appendix).
Lastly, we train the same number of experiments for all models which amounts to 1440 trained models in total, as summarized in Tableof the Appendix. [Figure]
4. R ESULTS
For each dynamical system we focus on the performance on the three test-sets, the in-distribution test set, which shares the same parameter distribution with the training set, and the two OOD test-sets (Easy and Hard), which represent an increasing parameter shift from the training data.
Models are compared on the cumulative Mean Absolute Error (MAE) between prediction and ground truth for the ﬁrst 200 time-steps.
We consider this to be sufﬁciently long-term as it is at least 20 times longer than the prediction horizon used during training.
Long predictions are obtained by re-feeding the model predictions back as input.
This approach has been shown to work well in systems where the dynamics are locally deterministic (Fotiadis et al., 2020).
A summary of the quantitative results can be found in Figures 3 & 4 and Table 8.
To account for the variability in the results, we present a summary of the best 5 runs of each model, selected by validation MAE.
We generally observe that model performance is correlated to the distribution shift of test-sets, and this is consistent for all systems and models.
The MAE is increasing as we move from the in-distribution test-set to the OOD
Easy and Hard test-sets.
Nevertheless, not all models suffer equally from the OOD performance drop.
Comparing the VAEs (Figure 3), we see that disentangled VAE models offer a substantial and consistent improvement over the VAE across all 3 dynamical systems.
The improvement is more pronounced for the OOD test-sets where the distribution shift is greater, a strong indication that disentanglement of domain parameters is an inductive bias that can lead to better generalization.
We also observe that VAE-SSD models the in-distribution data better that VAE-SD.
This seems to come at a slight overﬁtting cost, because the VAE-SD provides better OOD extrapolation in most cases.
This could be explained because the scaling function is dependent on min and max values of the factors of the training set.
The extra information allows the model to better capture the training data but sacriﬁces some generalization capacity.
On the other hand, disentanglement results for the MLP are mixed.
While in-distribution MLP-SD offers better results than the plain MLP, on the OOD test-sets, MLP-SD only performs favourably in the pendulum data.
Furthermore in Lotka-Volterra, MLP-SD models are very unstable, and this is a drawback that affects some VAE-SD model too (see Table 9 of the Appendix).
Probabilistic models seem better suited to capture the variation in the data.
The contrast between VAE-SD and MLP-SD illustrates that making use of privileged information and latent space disentanglement are not trivial and more work is needed to help us understand what works in practice and why.
Lastly, the LSTM (Figure 11 & Table 8 of the Appendix) is only comparable in the pendulum dataset and only for small
OOD shifts.
Qualitative predictions can be found in Figure 5.
E XPERIMENT - V IDEO S EQUENCE D YNAMICS
In the ﬁrst experiment we assessed supervised disentanglement for phase space prediction, where the states of the input trajectories are fully observable and only the domain parameters are unknown.
This experiment extends the idea of supervised disentanglement to pixel-space input and output, where the physical states have to be inferred by the model. [Figure] [Figure]
5. D ATASETS
The dynamical system we use in this experiment is the swinging pendulum, a common benchmark for modelling dynamics in video sequences (Brunton et al., 2020; Barber et al., 2021).
We consider 4 factors of variation, the length 𝑙 , gravity 𝑔 and initial angle 𝜃 and angular velocity 𝜔 .
Factors are drawn uniformly from a predetermined range.
As before, we create a test-set and two additional
OOD test-sets (Easy and Hard).
The OOD sets have length and gravity values outside of the original range, while the initial conditions 𝜃, 𝜔 are drawn from the same distribution.
The distribution of the factors of variation for the test-sets is illustrated in Figure 2.
The trajectories are ﬁrst computed in phase space using a numerical simulator and then rendered as video frames of 64 × 64 pixels.
More details about the dataset can be found in Section A.2 of the Appendix.
5. M ODELS AND T RAINING
In this experiment we use two different models CNN-VAE and RSSM.
CNN-VAE is described in
Section 3.3 and architectural details can be found in Section B.2.1.
During training the CNN-VAE the inner VAE is recursively used to predict, the number of recursions being a hyperparameter (Table 6 of the Appendix).
We found that this type of training leads to more stable long term predictions.
In total, 48 CNN-VAE models are trained half of which are with supervised disentanglement (CNN-VAE-SD).
The RSSM model is a generative model including both a stochastic and deterministic component.
We only use supervised disentanglement on the stochastic part, and term that model RSSM-SD.
Disentanglement is applied all four factors of variation of the domain, despite only length and gravity varying between datasets.
Detailed architectural and training details can be found in Section B.2.2 of the Appendix.
5. R ESULTS
Figure 7 shows the quality of predictions on video pendulum on two perceptual metrics: structural similarity (SSIM) and LPIPS (Zhang et al., 2018) as a function of predicted time distance.
We select the models which have the best cumulative metrics over the ﬁrst 800 timesteps on a validation set.
For the CNN-VAE, effects of disentanglement are more subtle.
We observe that, in-distribution, the disentangled CNN-VAE-SD has very similar quality when compared to the CNN-VAE.
For the
OOD datasets, though, disentanglement offers improved long-term predictions.
The improvement is more noticeable on LPIPS, and when the distribution shift is bigger (OOD Test-set Hard), indicating that disentanglement can help with OOD robustness.
For RSSM, we ﬁrst note that both models perform signiﬁcantly better than the CNN-VAE models, which is expected since they are considered competitive in long-term video prediction.
Disentanglement in RSSM seems to produce a tradeoff.
The plain RSSM model better in short-term prediction but its performance deteriorates with time, reaching VAE-CNN levels in all metrics.
On the other hand, the RSSM-SD model provides the best long-term scores in all metrics and all datasets.
Qualitative results in Figure 5 show that almost all models produce accurate short time predictions (approximately up to 200 time-steps).
This further stresses the importance of disentanglement for long-term performance.
In terms of
OOD robustness, disentanglement also appears to be helping.
While the RSSM-SD model lacks in short-term prediction quality on the in-distribution test-set, this performance gap closes as the
OOD test-sets get harder.
More speciﬁcally, on the in-distribution test-set the RSSM-SD overtakes
RSSM in SSIM after around 400 frames, while in the OOD Easy and Hard test sets, this happens around 350 and 250 time-steps respectively.
This narrowing gap indicates robustness improves with increasing distribution shifts.
The above ﬁndings are corroborated by peak signal-to-noise ratio (PSNR) comparisons (Figure 13 and Table 10 of the Appendix).
Furthermore, the qualitative results show that all models accurately capture the appearance of the pendulum even long-term.
Where they differ is on how well they capture the dynamics of the pendulum movement.
This could offer an explanation why disentangling the domain from the dynamics is important, and why in practice offers better long-term and out-of-distribution performance.
Overall, experiments suggest that supervised disentanglement can be used to model dynamical systems in video sequences, resulting in improved long-term and OOD performance.
C ONCLUSIONS
Using supervised disentanglement of domain parameters in generative models is a promising avenue for improving robustness.
Our experiments show that it can improve both OOD generalization and long-term prediction of dynamical systems.
This was demonstrated in phase-space with VAEs and also in video sequence modelling with state-of-the-art RSSMs.
By treating the domain parameters as factors of variation of the data and applying supervised disentanglement, several inductive biases are potentially enforced.
First, the model in addition to prediction also performs “soft” system identiﬁcation which acts as a regularizer.
Second, it creates an implicit hierarchy such that some latent variables correspond to sequence-wide domain parameters and the rest capture instant dynamics.
We speculate that this could additionally make the latent space more interpretable.
Third, if the model can correctly extract the parameters this mean that the prediction is based on both of them which is closer to how numerical integrators work, where the domain is known.
All of these could lead the model to learn the correct causal structure of the data.
Nevertheless, using privileged information for OOD robustness is not always straightforward and requires further exploration.
This is evident by the results of the MLP autoencoders which do not yield as consistent improvements.
A criticism of our method could be that cheap privileged information is not always available and/or depends on using simulated data.
Firstly, training on simulations is an increasingly appealing option because it is a cheap way to generate data to begin with.
This is, also, clearly demonstrated by the many advancements on techniques like sim2real (Peng et al., 2017) that try to bring models trained in simulated data to the real world.
So there seems to be no reason not to use the privileged information that comes with simulated data.
Under that light supervised disentanglement can provide a pathway for real world applications where robustness in dynamical system prediction is critical.
Applying the method to other datasets where there are more complex dynamic can increase its relevance.
Sequence-wide parameters could also be exploited through self-supervision.
R EPRODUCIBILITY S TATEMENT
We provide all the necessary code to reproduce our experiments at the anonymous repo https: //anonymous.4open.science/r/dis-dyn-systems (will be de-anonymized after the review process).
The repo contains code for generating all the datasets from scratch and also code for training all the models presented in this work.
The README also contains instructions on how to train the models.
The hyperparameters we have used are clearly and thoroughly presented in the Appendix.
These steps should signiﬁcantly help others reproduce our experiments.
For any further clariﬁcations, you are encouraged to contact the corresponding author(s)
A PPENDIX A D ATASETS A. P HASE SPACE
For simulations, we use an adaptive Runge-Kutta integrator with a timestep of 0 .
01 seconds.
Each simulated sequence has a different combination of factors of variation.
Simulation of the pendulum uses an initial angle 𝜃 which is randomly between 10 𝑜 − 170 𝑜 while the angular velocity 𝜔 is 0.
For the other two systems the initial conditions are always the same to avoid pathological conﬁgurations.
[Table]
[Figure]
[Figure]
A. V IDEO P ENDULUM
This data set contains image sequences of a moving pendulum under different conditions.
The positions of the pendulum are ﬁrst computed by a numerical simulator and then rendered in pixel space as frames of dimension 64 × 64 .
An example image sequence is shown in Figure 10.
For the simulations, we use an adaptive Runge-Kutta integrator with a timestep of 0 .
05 seconds.
The length of the pendulum, the strength of gravity and the initial conditions (position, momentum) are set to different values so that each trajectory slightly differs from the others.
The initial angle and initial velocity are drawn from the same uniform distribution for all data sets.
The initial angle ranges from 30 ◦ to 170 ◦ and the initial velocity ranges from − 2 𝑟𝑎𝑑
/ 𝑠 to 2 𝑟𝑎𝑑 / 𝑠 .
For training, validation and in-distribution testing set, the gravity ranges from 8 .
𝑚 2 / 𝑠 to 12 .
𝑚 2 / 𝑠 , and the pendulum length ranges from 1 . 20 𝑚 to 1 .
40 𝑚 .
In the easy out-of-distribution testing set, the gravity is between 12 . 0 − 12 .
5 𝑚 2 / 𝑠
and the pendulum length is between 1 . 40 − 1 . 45 𝑚 , while in the hard out-of-distribution testing set, the gravity is 12 . 5 − 13 . {+
+}
𝑚 2 / 𝑠
and the pendulum length is 1 . 45 − 1 .
50 𝑚 .
The distributions of these parameters are shown in Figure 2.
B T RAINING AND H YPERPARAMETERS B. P
[Figure]
During training the back-propagation is used after a single forward pass.
The input and output of the models are smaller than the sequence size, so to cover the whole sequence we use random starting points per batch, both during training and testing.
Both the VAE and MLP AE have an encoder with two hidden layers size 400,200 and a reverse decoder.
The LSTM model has two stack LSTM cells with hidden size of 100, which results on an equivalent number of neurons.
We used the Adam optimizer with 𝑏 1 = 0 .
9 and 𝑏 2 = 0 .
A scheduler for the learning rate was applied whose patience and scaling factor are hyperparameters.
Maximum number of epochs was set to 2000 but we employed also early stopping using a validation set which led to signiﬁcantly less epochs.
[Table]
[Table] [Table]
[Table]
V IDEO P ENDULUM B
[Table]
2. CNN-VAE M ODEL
Encoder has 4 layers convolutional layers with 32, 32, 64 and 64 maps respectively.
The ﬁlter size is 3, padding is 1 and stride is 2.
The last convolutional layer is ﬂattened as a 256-dimensional vector to become the inner VAE input.
The decoder 4 convolutional layers (64,64,32,32) with bi-linear upsampling.
Model input and out is 20 frames.
For the models without supervised disentanglement, a grid search is performed upon the 𝛽 value, the size of the latent space [Table] and the roll-out length during training.
For the models with supervised disentanglement, a grid search is performed upon the 𝛽 value, the size of the latent space, the time step of the data set, the roll-out length during training and the supervision multiplier.
The detailed search grid is summarised in Table 6.
Learning rate was 10 − 3 and an Adam optimizer ( 𝑏 1 = 0 .
9 and 𝑏 2 = 0 .
999 ) was used.
We also used early stopping upon the cumulative reconstruction loss for the ﬁrst 200 steps on a validation set with the max number of epochs being 1000.
[Table]
B.2. RSSM MODELS
For the RSSM model we follow the architecture parameters as described in Hafner et al.
Saxena et al.
For training we use sequences of 100 frames and batch size 100.
All models were trained for 300 epochs with a learning rate of 10 − 3 and an Adam optimizer ( 𝑏 1 = 0 .
9 and 𝑏 2 = 0 .
During testing the model uses 50 frames as context (input).
The parameters we tune appear in Table 7.
[Table]
C P HASE SPACE RESULTS
[Figure] [Figure]
OOD
[Figure]
OOD
[Figure]
[Table]
M A E
[Figure]
[Table]
body
[Figure]
[Table]
LSTM
[Table] [Table] [Table]
[Figure]
Pendulum
Pendulum
[Figure]
r
[Figure]
d a t o r Lotka
[Figure]
Volterra Prey 1. P r e d a t o r
Lotka-Volterra
[Figure]
x 1.
[Figure] [Figure]
3-body system
[Figure]
D V IDEO P ENDULUM R ESULTS
[Figure] [Figure]
OOD [Figure]
Test
[Figure]
set
OOD Easy
OOD
[Figure]
[Table]
[Figure]