PDE-Driven Spatiotemporal Disentanglement

1 INTRODUCTION . The interest of the machine learning community in physical phenomena has substantially grown for the last few years ( Shi et al. , 2015 ; Long et al. , 2018 ; Greydanus et al. , 2019 ) . In particular , an increasing amount of works studies the challenging problem of modeling the evolution of dynamical systems , with applications in sensible domains like climate or health science , making the understanding of physical phenomena a key challenge in machine learning . To this end , the community has successfully leveraged the formalism of dynamical systems and their associated differential formulation as powerful tools to specifically design efficient prediction models . In this work , we aim at studying this prediction problem with a principled and general approach , through the prism of Partial Differential Equations ( PDEs ) , with a focus on learning spatiotemporal disentangled representations . Prediction via spatiotemporal disentanglement was first studied in video prediction works , in order to separate static and dynamic information ( Denton & Birodkar , 2017 ) for prediction and interpretability purposes . Existing models are particularly complex , involving either adversarial losses or variational inference . Furthermore , their reliance on Recurrent Neural Networks ( RNNs ) hinders their ability to model spatiotemporal phenomena ( Yıldız et al. , 2019 ; Ayed et al. , 2020 ; Franceschi et al. , 2020 ) . Our proposition addresses these shortcomings with a simplified and improved model by grounding spatiotemporal disentanglement in the PDE formalism . Spatiotemporal phenomena obey physical laws such as the conservation of energy , that lead to describe the evolution of the system through PDEs . Practical examples include the conservation of energy for physical systems ( Hamilton , 1835 ) , or the equation describing constant illumination in a scene ( Horn & Schunck , 1981 ) for videos that has had a longstanding impact in computer vision with optical flow methods ( Dosovitskiy et al. , 2015 ; Finn et al. , 2016 ) . We propose to model the evolution of partially observed spatiotemporal phenomena with unknown dynamics by leveraging a formal method for the analytical resolution of PDEs : the functional separation of variables ( Miller , 1988 ) . Our framework formulates spatiotemporal disentanglement for prediction as learning a separable solution , where spatial and dynamic information are represented in separate variables . Besides offering a novel interpretation of spatiotemporal disentanglement , it confers simplicity and performance compared to existing methods : disentanglement is achieved through the sole combination of a prediction objective and regularization penalties , and the temporal dynamics is defined by a learned Ordinary Differential Equation ( ODE ) . We experimentally demonstrate the applicability , disentanglement capacity and ∗Equal contribution . forecasting performance of the proposed model on various spatiotemporal phenomena involving standard physical processes and synthetic video datasets against prior state-of-the-art models . 2 RELATED WORK . Our contribution deals with two main directions of research : spatiotemporal disentanglement and the coupling of neural networks and PDEs . Spatiotemporal disentanglement . Disentangling factors of variations is an essential representation learning problem ( Bengio et al. , 2013 ) . Its cardinal formulation for static data has been extensively studied , with state-of-the-art solutions ( Locatello et al. , 2019 ) being essentially based on Variational Autoencoders ( VAEs ; Kingma & Welling , 2014 ; Rezende et al. , 2014 ) . As for sequential data , several disentanglement notions have been formulated , ranging from distinguishing objects in a video ( Hsieh et al. , 2018 ; van Steenkiste et al. , 2018 ) to separating and modeling multi-scale dynamics ( Hsu et al. , 2017 ; Yingzhen & Mandt , 2018 ) . We focus in this work on the dissociation of the dynamics and visual aspects for spatiotemporal data . Even in this case , dissociation can take multiple forms . Examples in the video generation community include decoupling the foreground and background ( Vondrick et al. , 2016 ) , constructing structured frame representations ( Villegas et al. , 2017b ; Minderer et al. , 2019 ; Liu et al. , 2019 ) , extracting physical dynamics ( Le Guen & Thome , 2020 ) , or latent modeling of dynamics in a state-space manner ( Fraccaro et al. , 2017 ; Franceschi et al. , 2020 ) . Closer to our work , Denton & Birodkar ( 2017 ) , Villegas et al . ( 2017a ) and Hsieh et al . ( 2018 ) introduced in their video prediction models explicit latent disentanglement of static and dynamic information obtained using adversarial losses ( Goodfellow et al. , 2014 ) or VAEs . Disentanglement has also been introduced in more restrictive models relying on data-specific assumptions ( Kosiorek et al. , 2018 ; Jaques et al. , 2020 ) , and in video generation ( Tulyakov et al. , 2018 ) . We aim in this work at grounding and improving spatiotemporal disentanglement with more adapted inductive biases by introducing a paradigm leveraging the functional separation of variables resolution method for PDEs . Spatiotemporal prediction and PDE-based neural network models . An increasing number of works combining neural networks and differential equations for spatiotemporal forecasting have been produced for the last few years . Some of them show substantial improvements for the prediction of dynamical systems or videos compared to standard RNNs by defining the dynamics using learned ODEs ( Rubanova et al. , 2019 ; Yıldız et al. , 2019 ; Ayed et al. , 2020 ; Le Guen & Thome , 2020 ) , following Chen et al . ( 2018 ) , or adapting them to stochastic data ( Ryder et al. , 2018 ; Li et al. , 2020 ; Franceschi et al. , 2020 ) . Most PDE-based spatiotemporal models exploit some prior physical knowledge . It can induce the structure of the prediction function ( Brunton et al. , 2016 ; de Avila Belbute-Peres et al. , 2018 ) or specific cost functions , thereby improving model performances . For instance , de Bézenac et al . ( 2018 ) shape their prediction function with an advection-diffusion mechanism , and Long et al . ( 2018 ; 2019 ) estimate PDEs and their solutions by learning convolutional filters proven to approximate differential operators . Greydanus et al . ( 2019 ) , Chen et al . ( 2020 ) and Toth et al . ( 2020 ) introduce non-regression losses by taking advantage of Hamiltonian mechanics ( Hamilton , 1835 ) , while Tompson et al . ( 2017 ) and Raissi et al . ( 2020 ) combine physically inspired constraints and structural priors for fluid dynamic prediction . Our work deepens this literature by establishing a novel link between a resolution method for PDEs and spatiotemporal disentanglement , thereby introducing a data-agnostic model leveraging any static information in observed phenomena . 3 BACKGROUND : SEPARATION OF VARIABLES . Solving high-dimensional PDEs is a difficult analytical and numerical problem ( Bungartz & Griebel , 2004 ) . Variable separation aims at simplifying it by decomposing the solution , e.g. , as a simple combination of lower-dimensional functions , thus reducing the PDE to simpler differential equations . 3.1 SIMPLE CASE STUDY . Let us introduce this technique through a standard application , with proofs in Appendix A.1 , on the one-dimensional heat diffusion problem ( Fourier , 1822 ) , consisting in a bar of length L , whose temperature at time t and position x is denoted by u ( x , t ) and satisfies : ∂u ∂t = c2 ∂2u ∂x2 , u ( 0 , t ) = u ( L , t ) = 0 , u ( x , 0 ) = f ( x ) . ( 1 ) Suppose that a solution u is product-separable , i.e. , it can be decomposed as : u ( x , t ) = u1 ( x ) · u2 ( t ) . Combined with Equation ( 1 ) , it leads to c2u′′1 ( x ) /u1 ( x ) = u ′ 2 ( t ) /u2 ( t ) . The left- and right-hand sides of this equation are respectively independent from t and x . Therefore , both sides are constant , and solving both resulting ODEs gives solutions of the form , with µ ∈ R and n ∈ N : u ( x , t ) = µ sin ( nπx/L ) × exp ( − ( cnπ/L ) 2 t ) . ( 2 ) The superposition principle and the uniqueness of solutions under smoothness constraints allow then to build the set of solutions of Equation ( 1 ) with linear combinations of separable solutions ( Le Dret & Lucquin , 2016 ) . Besides this simple example , separation of variables can be more elaborate . 3.2 FUNCTIONAL SEPARATION OF VARIABLES . The functional separation of variables ( Miller , 1988 ) generalizes this method . Let u be a function obeying a given arbitrary PDE . The functional variable separation method amounts to finding a parameterization z , a functional U , an entangling function ξ , and representations φ and ψ such that : z = ξ ( φ ( x ) , ψ ( t ) ) , u ( x , t ) = U ( z ) . ( 3 ) Trivial choices ξ = u and identity function as U , φ and ψ ensure the validity of this reformulation . Finding suitable φ , ψ , U , and ξ with regards to the initial PDE can facilitate its resolution by inducing separate simpler PDEs on φ , ψ , and U . For instance , product-separability is retrieved with U = exp . General results on the existence of separable solutions have been proven ( Miller , 1983 ) , though their uniqueness depends on the initial conditions and the choice of functional separation ( Polyanin , 2020 ) . Functional separation of variables finds broad applications . It helps to solve refinements of the heat equation , such as generalizations with an advection term ( see Appendix A.2 ) or with complex diffusion and source terms forming a general transport equation ( Jia et al. , 2008 ) . Besides the heat equation , functional separation of PDEs is also applicable in various physics fields like reactiondiffusion with non-linear sources or convection-diffusion phenomena ( Polyanin , 2019 ; Polyanin & Zhurov , 2020 ) , Hamiltonian physics ( Benenti , 1997 ) , or even general relativity ( Kalnins et al. , 1992 ) . Reparameterizations such as Equation ( 3 ) implement a separation of spatial and temporal factors of variations , i.e. , spatiotemporal disentanglement . We introduce in the following a learning framework based on this general method . 4 PROPOSED METHOD . We propose to model spatiotemporal phenomena using the functional variable separation formalism . We first describe our notations and then derive a principled model and constraints from this method . 4.1 PROBLEM FORMULATION THROUGH SEPARATION OF VARIABLES . We consider a distribution P of observed spatiotemporal trajectories and corresponding observation samples v = ( vt0 , vt0+∆t , . . . , vt1 ) , with vt ∈ V ⊆ Rm and t1 = t0 + ν∆t . Each sequence v ∼ P corresponds to an observation of a dynamical phenomenon , assumed to be described by a hidden functional uv ( also denoted by u for the sake of simplicity ) of space coordinates x ∈ X ⊆ Rs and time t ∈ R that characterizes the trajectories . More precisely , uv describes an unobserved continuous dynamics and v corresponds to instantaneous discrete spatial measurements associated to this dynamics . Therefore , we consider that vt results from a time-independent function ζ of the mapping uv ( · , t ) . For example , v might consist in temperatures measured at some points of the sea surface , while uv would be the complete ocean circulation model . In other words , v provides a partial information about uv and is a projection of the full dynamics . We seek to learn a model which , when conditioned on prior observations , can predict future observations . To this end , we posit that the state u of each observed trajectory v is driven by a hidden PDE , shared among all trajectories ; we discuss this assumption in details in Appendix C.1 . Learning such a PDE and its solutions would then allow us to model observed trajectories v. However , directly learning solutions to high-dimensional unknown PDEs is a complex task ( Bungartz & Griebel , 2004 ; Sirignano & Spiliopoulos , 2018 ) . We aim in this work at simplifying this resolution . We propose to do so by relying on the functional separation of variables of Equation ( 3 ) , in order to leverage a potential separability of the hidden PDE . Therefore , analogously to Equation ( 3 ) , we propose to formulate the problem as learning observation-constrained φ , ψ and U , as well as ξ and ζ , such that : z = ξ ( φ ( x ) , ψ ( t ) ) , u ( x , t ) = U ( z ) , vt = ζ ( u ( · , t ) ) , ( 4 ) with φ and ψ allowing to disentangle the prediction problem . In the formalism of the functional separation of variables , this amounts to decomposing the full solution u , thereby learning a spatial PDE on φ , a temporal ODE on ψ , and a PDE on U , as well as their respective solutions .

Summary: The authors present a generative model for videos where the latent trajectories have two components - a term without a slowness loss that represents "content" and a term with a slowness loss that represents "style". They present results on a dataset simulating the wave equation and on videos of moving MNIST digits and 3D chairs. The results are generally good, especially for long roll-outs, and they demonstrate something like disentangling by showing that the identities of the digits can be swapped in the moving MNIST data.