Disentangled generative models for robust dynamical system prediction

1 INTRODUCTION . 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 fields 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 significant 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 difficult in many dynamical systems , because error accumulation can diverge in finite 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 first define 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 coefficients , 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 ( Miladinović 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 overfitting ( Bengio et al. , 2012 ) . System identification can be used to extract parameters , but 1Code for reproducing our experiments at : https : //anonymous.4open.science/r/ dis-dyn-systems/ 2Animated phase-space and video predictions are available at : https : //bit.ly/dis-dyn-systems 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 findings ( 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 ( Miladinović 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 ( Sarafianos et al. , 2017 ; Lee et al. , 2018 ) . We thus use supervised disentanglement ( Locatello et al. , 2019 ) by leveraging the ground truth domain parameters . To the best of our knowledge , using domain parameters information this way , has not been previously explored in the dynamical system prediction setting . Contributions While others have treated domain parameters as factors of variation in the data distribution , our work is the first , 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 . We additionally apply disentanglement on a more powerful , 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 first 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 beneficial 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 . 2 RELATED WORK . 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 finding 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 ( 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 , Miladinović 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 difficult 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 . ( 2019 ) the domain is fixed ; 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 ( Sarafianos et al. , 2017 ) . 3 METHODS . 3.1 VARIATIONAL AUTOENCODERS . Variational autoencoders ( VAEs ) ( Kingma & Welling , 2014 ) offer a principled approach to latent variable modeling by combining a variational inference model @ q ( 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 : Lq , \ ( x ) = E @ q ( z |x ) [ log ? \ ( x | z ) ] − ! ( @ q ( z | x ) | | ? ( z ) ) ( 1 ) The first part of the ELBO is the reconstruction loss ( in our case the prediction loss ) and the second part is the Kullback-Leibler divergence that quantifies 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 @ q ( z | x ) = N ( z | µI , I ) , allowing a closed form KL-divergence , while the decoder has a Laplace distribution ? \ ( x | z ) = Laplace ( x | µG , WI ) with constant diagonal covariance W > 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 µI ≡ µI ( x ; q ) , I ≡ diag [ σI ( x ; q ) ] 2 , and µG ≡ µG ( z ; \ ) are computed via feed-forward neural networks . 3.2 DISENTANGLEMENT 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ξ ( z1 : : , ξ ) between the ground truth factors of the domain parameters ξ ∈ R : and the output of the corresponding latents , z1 : : . 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 fit 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 ( ` I8 ) which linearly scales the ` I8 between the min and max values of the corresponding factor of variation : G ( ` I8 ) = ` I8 · ( max ( b8 ) −min ( b8 ) ) +min ( b8 ) ( 2 ) where b8 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 X which is empirically tuned . Plugging these choices in results in the following loss function : Lq , \ ( x ) =E @ q ( z |x1 : = ) [ 1 W ‖x=+1 : =+ > − µG ( z ; \ ) ‖1 ] + 3 log W ( Prediction loss ) + ‖σI ( x1 : = ; q ) ‖22 − log diag [ σI ( x1 : = ; q ) ] 2 + ‖µI ( x1 : = ; q ) ‖22 ( KL-Divergence ) ( 3 ) + X ξG − G ( µI1 : : ( x1 : = ; q ) ) 1 } ( Sup . disentangl . loss ) Using the reparameterization trick ( Kingma & Welling , 2014 ) , the loss is amenable to optimzation by stochastic gradient descent , with batch size = . The model architecture can be seen in Figure 1 ( left ) .

Summary: This paper introduces a supervised loss to encourage dynamical systems predictors to retain systems' parameters (e.g. appearing in the underlying ODE) in their latent space. It presents multiple experiments to evaluate the advantages of this approach, including better long-term forecasting ability as well as improved prediction performance for out-of-distribution parameters, that had not been seen during training.