Transformers for Modeling Physical Systems

1 INTRODUCTION . The transformer model ( Vaswani et al. , 2017 ) , built on self-attention , has largely become the stateof-the-art approach for a large set of neural language processing ( NLP ) tasks including language modeling , text classification , question answering , etc . Although more recent transformer work is focused on unsupervised pre-training of extremely large models ( Devlin et al. , 2018 ; Radford et al. , 2019 ; Dai et al. , 2019 ; Liu et al. , 2019 ) , the original transformer model garnered attention due to its ability to out-perform other state-of-the-art methods by learning longer-term dependencies without recurrent connections . Given that the transformer model was originally developed for NLP , nearly all related work has been rightfully confined within this field with only a few exceptions . Here , we focus on the development of transformers to model dynamical systems that can replace otherwise expensive numerical solvers . In other words , we are interested in using transformers to learn the language of physics . The surrogate modeling of physical systems is a research field that has existed for several decades and is a large ongoing effort in scientific machine learning . A surrogate model is defined as an approximate model of a physical phenomenon that is designed to replace an expensive computational solver that would otherwise be needed to resolve the system of interest . The key characteristic of surrogate models is their ability to model a distribution of initial or boundary conditions rather than learning just one solution . This is arguably essential for the justification of training a deep learning model versus using a standard numerical solver . The most tangible applications of surrogates are for optimization , design and inverse problems where many repeated simulations are typically needed . With the growing interest in deep learning , deep neural networks have been used for surrogate modeling a large range of physical systems in recent literature . Standard deep neural network architectures such as auto-regressive ( Mo et al. , 2019 ; Geneva & Zabaras , 2020a ) , residual/Euler ( González-Garcı́a et al. , 1998 ; Sanchez-Gonzalez et al. , 2020 ) , recurrent and LSTM based models ( Mo et al. , 2019 ; Tang et al. , 2020 ; Maulik et al. , 2020 ) have been largely demonstrated to be effective at modeling various physical dynamics . Such models generally rely exclusively on the past time-step to provide complete information on the current state of the system ’ s evolution . Particularly for dynamical systems , present machine learning models lack generalizable time cognisant capabilities to predict multi-time-scale phenomena present in systems including turbulent fluid flow , multi-scale materials modeling , molecular dynamics , chemical processes , etc . Thus currently adopted models struggle to maintain true physical accuracy for long-time 1Code available at : [ URL available after review ] . 2Supplementary videos available at : https : //sites.google.com/view/transformersphysx . predictions . Much work is needed to scale such deep learning models to larger physical systems that are of scientific and industrial interest . This work deviates from this pre-existing literature by investigating the use of transformers for the prediction of physical systems , relying entirely on selfattention to model dynamics . In the recent work of Shalova & Oseledets ( 2020 ) , such self-attention models were tested to learn single solutions of several low-dimensional ordinary differential equations . In this work , we propose a physics inspired embedding methodology to model a distribution of dynamical solutions . We will demonstrate our model on high-dimensional partial differential equation problems that far surpass the complexity of past works . To the authors best knowledge , this is the first work to explore transformer NLP architectures for the prediction of physical systems . 2 METHODS . When discussing dynamical systems , we are interested in systems that can be described through a dynamical ordinary or partial differential equation : φt = F ( x , φ ( t , x , η ) , ∇xφ , ∇2xφ , φ · ∇xφ , . . . ) , F : R× Rn → Rn , t ∈ T ⊂ R+ , x ∈ Ω ⊂ Rm , ( 1 ) in which φ ∈ Rn is the solution of this differential equation with parameters η , in the time interval T and spatial domain Ω with a boundary Γ ⊂ Ω . This general form can embody a vast spectrum of physical phenomena including fluid flow and transport processes , mechanics and materials physics , and molecular dynamics . In this work , we are interested in learning the set of solutions for a distribution of initial conditions φ0 ∼ p ( φ0 ) , boundary conditions B ( φ ) ∼ p ( B ) ∀x ∈ Γ or equation parameters η ∼ p ( η ) . This accounts for modeling initial value , boundary value and stochastic problems . We emphasize that this is fundamentally different , more difficult and of greater interest for most scientific applications compared to learning a single solution . To make this problem applicable to the use of transformer models , the continuous solution is discretized in both the spatial and temporal domains such that the solution of the differential equation is Φ = { φ0 , φ1 , . . .φT } ; φi ∈ Rn×d for which φi has been discretized by d points in Ω . We assume an initial state φ0 and that the time interval T is discretized by T time-steps with a time-step size ∆t . Hence , we pose the problem of modeling a dynamical system as a time-series problem . The machine learning methodology has two core components : the transformer for modeling dynamics and the embedding network for projecting physical states into a vector representation . Similar to NLP , the embedding model is trained prior to the transformer . This embedding model is then frozen and the entire data-set is converted to the embedded space in which the transformer is then trained as illustrated in Fig . 1 . During testing , the embedding decoder is used to reconstruct the physical states from the transformer ’ s predictions . 2.1 TRANSFORMER . The transformer model was originally designed with NLP as the sole application with word vector embeddings of a passage of text being the primary input ( Vaswani et al. , 2017 ) . However , recent works have explored using attention mechanisms for different machine learning tasks ( Veličković et al. , 2017 ; Zhang et al. , 2019 ; Fu et al. , 2019 ) and a few investigate the use of transformers for applications outside of the NLP field ( Chen et al. , 2020 ) . This suggests that self-attention and in particular transformer models may work well for any problem that can be posed as a sequence of vectors . In this work , the primary input to the transformer will be an embedded dynamical system , Ξ = { ξ0 , ξ1 , . . . .ξT } , where the embedded state at time-step i is denoted as ξi ∈ Re . Given that we are interested in the prediction of a physical time series , this motivates the usage of a language modeling architecture that is designed for the sequential prediction of words in a body of text . We select the transformer decoder architecture used in the GPT models ( Radford et al. , 2018 ; 2019 ) . Our model follows the GPT-2 architecture based on the implementation in the Hugging Face transformer repository ( Wolf et al. , 2019 ) , but is significantly smaller in size than these modern NLP transformers . This model consists of a stack of transformer decoder layers that use masked attention , as depicted in Fig . 2 . The inputs to the transformer are the embedded representation of the physical system with the sinusoidal positional encoding proposed in the original transformer ( Vaswani et al. , 2017 ) . To train the model , consider a data set of D embedded i.i.d . time-series D = { Ξi } D i=1 for which we can use the standard time-series Markov model ( language modeling ) log-likelihood : LD = D∑ i T∑ j − log p ( ξij |ξij−k , . . . , ξij−1 , θ ) . ( 2 ) k is the context window and θ are the model ’ s parameters . Contrary to the standard NLP approach which poses the likelihood as a softmax over a dictionary of tokens , the likelihood here is taken as a standard Gaussian between the transformer ’ s prediction and the target embedded value resulting in a L2 loss . This is due to the fact that the solution to most physical systems can not be condensed to a discrete finite set making tokenization into a finite dictionary not possible and thus a softmax approach not applicable . Training is the standard auto-regressive method used in GPT ( Radford et al. , 2018 ) , as opposed to the word masking ( Devlin et al. , 2018 ) , constrained to the embedded space . The physical states , φi , have the potential to be very high dimensional thus training the transformer in the lower-dimensional embedded space can significantly lower training costs . 2.2 EMBEDDING MODEL . The second major component of the machine learning methodology is the embedding model responsible for projecting the discretized physical state space into a 1D vector representation . In NLP , the standard approach is to tokenize then embed a finite vocabulary of words , syllables or characters using methods such as n-gram models , Byte Pair Encoding ( Gage , 1994 ) , Word2Vec ( Mikolov et al. , 2013a ; b ) , GloVe ( Pennington et al. , 2014 ) , etc . These methods allow language to be represented by a series of 1D vectors that serve as the input to the transformer . Clearly a finite tokenization and such NLP embeddings are not directly applicable to a physical system , thus we propose our own embedding method designed specifically for dynamical systems . Consider learning the generalized mapping between the system ’ s state space and embedded space : F : Rn×d → Re and G : Re → Rn×d . Naturally , multiple approaches can be used especially if the dimensionality of the embedded space is less than that of the state-space but this is not always the case . The primary approach that we will propose is a Koopman observable embedding which is a technique that can be applied universally to all dynamical systems . Considering the discrete time form of the dynamical system in Eq . 1 , the evolution of the state variables can be abstracted by φi+1 = F ( φi ) for which F is the dynamic map from one time-step to the next . The foundation of Koopman theory states that for any dynamical system of this form , there exists an infinite set of state observables , g ( φi ) , that evolve linearly in time such that : Kg ( φi ) , g ◦ F ( φi ) , ( 3 ) where K is known as the Koopman operator ( Koopman , 1931 ) which is time-invariant . Namely , the Koopman observables can evolve in time through continual matrix multiplication of with the Koopman operator : g ( φi+1 ) = Kg ( φi ) , g ( φi+2 ) = K2g ( φi ) , g ( φi+3 ) = K3g ( φi ) , . . . ( 4 ) in which Kn denotes n matrix products ( e.g . K3 = K · K · K ) . Modeling the dynamics of a system through the linear Koopman space can be attractive due to its simplification of the dynamics but also the potential physical insights it brings along with it . Spectral analysis of the Koopman operator can reveal fundamental dynamical modes that drive the system ’ s evolution in time . However , these observables , g ( φi ) , are typically unknown and are theoretically infinite dimensional . Thus Koopman theory can be viewed as a trade off between lifting the state space into observable space with more complex states but with simpler dynamics . For practical application , the Koopman operator and observables are finitely approximated . In recent years , various machine learning based approaches have been proposed to learn both the Koopman operator and observables approximated in a finite space for modeling , control and dynamical mode analysis ( Takeishi et al. , 2017 ; Li et al. , 2017 ; Lusch et al. , 2018 ; Korda & Mezić , 2018 ; Otto & Rowley , 2019 ; Korda et al. , 2020 ; Mezic , 2020 ) . While deep learning methods have enabled greater success with discovering Koopman observables and operators ( Lusch et al. , 2018 ; Otto & Rowley , 2019 ) , applications have still been limited to fairly simple systems and do not work for long time prediction . This is likely due to the approximation of the finite-dimensional Koopman observables , limited data and complete dependence on the discovered Koopman operator K to model the dynamics . Suggesting the prediction of a system through a single linear transform clearly has significant limitations and is fundamentally a naive approach from a machine learning perspective . In this work , we propose using approximate Koopman observations as a methodology to develop embeddings for the transformer model such that F ( φi ) , g ( φi ) . As seen in Fig . 3 , the embedding model follows a standard encoder-decoder model with the middle latent variables being the Koopman observables . Embedding model architectures for each numerical example are provided in Appendix A . We introduce a learnable Koopman operator which takes the form of a banded matrix that is symmetrical about the diagonal to encourage the discovery of dominate dynamical modes rather than high-frequency oscillations ( Lusch et al. , 2018 ; Otto & Rowley , 2019 ) . This learned Koopman operator is disposed of once training of the embedding model is complete . Given the data set of physical state time-series , DΦ = { Φi } D i , the Koopman embedding model is trained with the following loss : LDΦ = D∑ i=1 T∑ j=0 λ0MSE ( φij , G ◦ F ( φij ) ) ︸ ︷︷ ︸ Reconstruction +λ1MSE ( φij , G ◦ KjF ( φi0 ) ) ︸ ︷︷ ︸ Dynamics +λ2 ‖K‖22︸ ︷︷ ︸ Decay . ( 5 ) This loss function consists of three components : the first is a reconstruction loss which ensures a consistent mapping to and from the embedded representation . The second is the Koopman dynamics loss which pushes ξj to follow linear dynamics resulting in time-steps of similar dynamics to have similar embeddings . The last term decays the Koopman operator ’ s parameters to help force the model to discover lower-dimensional dynamical modes . In reference to traditional NLP embeddings , we believe our Koopman observable embedding has a motivation similar to Word2Vec ( Mikolov et al. , 2013a ) as well as more recent embedding methods such as context2vec ( Melamud et al. , 2016 ) , ELMo ( Peters et al. , 2018 ) , etc . Namely , these methods are based on the principle of context : words that are contextually related to each other have a similar embedding . The Koopman embedding has a similar objective encouraging physical realizations containing similar time-invariant dynamical modes to also have similar embeddings . Hence , our goal with the embedding model is to not find the true Koopman observables or operator , but rather leverage Koopman to enforce physical context .

Summary: The paper proposes to use transformers for modelling dynamical systems. The transformer is combined with a linear dynamical system to enforce Koopman features and is trained using the reconstruction and prediction loss. Finally, the proposed algorithm is applied to the different tasks with 1, 2 & 3 dimensions. On each simulated task the proposed algorithm marginally outperforms sufficient baselines.