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: In this paper, the authors propose a supervised approach to disentangle domain parameters from the dynamics in a deep latent variable model like VAE. Extending VAEs to dynamical systems is a relevant problem and has been a focus of interest in many recent works [1,2,3,4,5,6,7]. This paper identifies two issues for developing dynamical VAEs,   i) out of distribution generalisation  ii) long term trajectory prediction  The main contribution is to address the aforementioned issues using a supervised loss defined between latent variables and domain parameters. The authors present empirical experiments to support the idea.