Physics-aware Spatiotemporal Modules with Auxiliary Tasks for Meta-Learning

1 Introduction

Deep learning has recently shown promise to play a major role in devising new solutions to applications with natural phenomena, such as climate change [Drgona et al., 2019], ocean dynamics [Cosne et al., 2019], air quality [Soh et al., 2018], and so on. Deep learning techniques inherently require a large amount of data for effective representation learning, so their performance is significantly degraded when there are only a limited number of observations. However, in many tasks in the real-world physical systems we only have access to a limited amount of data. One example is air quality monitoring [Berman, 2017], in which the sensors are irregularly distributed over the space - many sensors are located in urban areas whereas there are much fewer sensors in vast rural areas. Another example is extreme weather modeling and forecasting, i.e., temporally short events (e.g., tropical cyclones [Racah et al., 2017]) without sufficient observations over time. Moreover, inevitable missing values from sensors [Tang et al., 2019] further reduce the number of operating sensors and shorten the length of fully-observed sequences. Thus, achieving robust performance from a few spatiotemporal observations in physical systems remains an essential but challenging problem.

Learning on a limited amount of data from physical systems can be considered as a few shot learning. While many meta-learning techniques [Ravi and Larochelle, 2017; Finn et al., 2017; Munkhdalai and Yu, 2017; Duan et al., 2017; Snell et al., 2017; Mishra et al., 2018] have been developed to handle this few shot learning setting, there are still some challenges for the existing meta-learning methods to be applied in modeling natural phenomena. First, it is not easy to find a set of similar meta-tasks, which provide shareable latent representations needed to understand targeted observations. For instance, while image-related tasks (object detection [He et al., 2017]) can take advantage of an image-feature extractor pre-trained by a large set of images [Deng et al., 2009] and well-designed architecture [He et al., 2016], there is no such large data corpus that is widely applicable for understanding natural phenomena. Second, unlike computer vision or natural language processing tasks where a common object (images or words) is clearly defined, it is not straightforward to find analogous objects in the spatiotemporal data. Finally, exact equations behind natural phenomena are usually unknown, causing the difficulty in reproducing the similar dataset via simulation. Although there have been some works [de Bezenac et al., 2018; Greydanus et al., 2019] improving data efficiency via explicitly incorporating PDEs as neural network layers when modeling spatiotemporal dynamics, it is hard to generalize for modeling different or unknown dynamics, which is ubiquitous in real-world scenario.

In this work, we propose physics-aware modules designed for meta-learning to tackle the few shot learning challenges in physical observations. One of fundamental equations in physics describing the transport of physical quantity over space and time is a continuity equation:

where ρ is the amount of the target quantity (u) per unit volume, J is the flux of the quantity, and σ is a source or sink, respectively. This fundamental equation can be used to derive more specific transport equations such as the convectiondiffusion equation, Navier-Stokes equations, and Boltzmann transport equation. Thus, the continuity equation is a starting point to model spatiotemporal (conservative) observations which are accessible from sensors. Based on the form of ρ and J in terms of u, Eq. 1 can be generalized as:

where the function F (·) describes how the target u is changed over time from its spatial derivatives. Inspired by the form of Eq. 2, we propose two modules: spatial derivative modules (SDM) and time derivative modules (TDM). Since the spatial derivatives such as ∇, ∇·, and ∇2 are commonly used across different PDEs, the spatial modules are PDE-independent and they can be meta-initialized from synthetic data. Then, the PDE-specific temporal module is trained to learn the unknown function F(·) from few observations in the real-world physical systems.

This approach can effectively leverage a large amount of simulated data to train the spatial modules as the modules are PDE-independent and thus mitigating the need for a large amount of real-world tasks to extract shareable features. In addition, since the spatial modules are universally used in physics equations, the representations from the modules can be conveniently integrated with data-driven models for modeling natural phenomena. Based on the modularized PDEs, we introduce a novel approach that marries physics knowledge in spatiotemporal prediction tasks with meta-learning by providing shareable modules across spatiotemporal observations in the real-world.

Our contributions are summarized below:

• Modularized PDEs and auxiliary tasks: Inspired by forms of PDEs in physics, we decompose PDEs into shareable (spatial) and adaptation (temporal) parts. The shareable one is PDE-independent and specified by auxiliary tasks: supervision of spatial derivatives.
• Physics-aware meta-learning: We provide a framework for physics-aware meta-learning, which consists of PDE-independent/-specific modules. The framework is flexible to be applied to the modeling of different and unknown dynamics.
• Utilizing synthetic data for shareable modules: We extract shareable parameters in the spatial modules from synthetic data, which can be easily generated.

2 Modularized PDEs andMeta-Learning

In this section, we describe how the physics equations for conserved quantities are decomposable into two parts and how the meta-learning approach tackles the task by utilizing synthetic data when the data are limited.

2.1 Decomposability of Variants of a Continuity Equation

In physics, a continuity equation (Eq. 1) describes how a locally conserved quantity such as temperature, fluid density, heat, and energy is transported across space and time. This equation underlies many specific equations such as the convection-diffusion equation and Navier-Stokes equations:

where the scalar u and vector field u are the variables of interest (e.g., temperature, flow velocity, etc.). A dot over a variable is time derivative. The common feature in these equations is that the forms of equations can be digested as [Bar-Sinai et al., 2019]:

where the right-hand side denotes a function of spatial derivatives. As the time derivative can be seen as a Euler discretization [Chen et al., 2018], it is notable that the next state is a function of the current state and spatial derivatives. Thus, knowing spatial derivatives at time t is a key step for spatiotemporal prediction at time t + 1 for locally conserved quantities. According to Eq. 3, the spatial derivatives are universally used in variants of Eq. 1 and only the updating function F(·) is specifically defined for a particular equation. This property implies that PDEs for physical quantities can be decomposable into two modules: spatial and temporal derivative modules.

2.2 Spatial Derivative Modules: PDE-independent Modules

Finite difference method (FDM) is widely used to discretize a d-order derivative as a linear combination of function values on a n-point stencil.

where n > d. According to FDM, it is independent for a form of PDE to compute spatial derivatives, which are input components of F(·) in Eq. 3. Thus, we can modularize spatial derivatives as PDE-independent modules. The modules that can be learnable as a data-driven manner to infer the coefficients (αi) have been proposed recently [Bar-Sinai et al., 2019]. The data-driven coefficients are particularly useful when the discretization in the n-point stencil is irregular and low-resolution where the fixed coefficients cause substantial numerical errors.

2.3 Time Derivative Module: PDE-specific Module

Once d-order derivatives are modularized by learnable parameters, the approximated spatial derivatives from the spatial modules are fed into an additional module to learn the function F(·) in Eq. 3. This module is PDE-specific as the function F describes how the spatiotemporal observations change. Since the exact form of a ground truth PDE is not given, the time derivative module is data-driven and will be adapted to observations instead.

2.4 Meta-Learning with PDE-independent/-specific Modules

Recently, [Raghu et al., 2020] investigate the effectiveness of model agnostic meta-learning (MAML, [Finn et al., 2017]) and it is found that the outer loop of MAML is more likely to learn parameters for reusable features rather than rapid adaptation. The finding that feature reuse is the predominant reason for efficient learning of MAML allows us to use additional information which is beneficial for learning better representations. Previously, the objective in meta-training has been considered to be matched with one in meta-test as the purpose of meta-learning is to learn good initial parameters applicable across similar tasks (e.g., image classification to image classification). We are now able to incorporate auxiliary tasks under a meta-learning setting to reinforce reusable features for a main task. As described in Sec. 2.1, the spatial modules are reusable across different observations, and thus, we can meta-initialize the spatial modules first with spatial derivatives provided by synthetic datasets. Then, we can integrate the spatial modules with the task-specific temporal module during meta-test to help adaptation of TDM on few observations. Since the spatial modules are trained by readily available synthetic datasets, a large number of real-world datasets for meta-training is not required.

3 Physics-aware Meta-Learning with Auxiliary Tasks

In this section, we develop a physics-aware meta-learning framework for the modularized PDEs. Fig. 1 describes the proposed framework and its computational process.

3.1 Spatial Derivative Module

As we focus on the modeling and prediction of sensor-based observations, where the available data points are inherently on a spatially sparse irregular grid, we use graph networks for each module φk to learn the finite difference coefficients. Given a graph G = (V, E) where V = {1, . . . ,N} and E = {(i, j) : i, j ∈ V}, a node i denotes a physical location xi = (xi, yi) where a function value ui = u(xi, yi) is observed. Then, the graph signals with positional relative displacement as edge features are fed into the spatial modules to approximate spatial derivatives by Alg. 1. The coefficients (ai, b(i,j)) on each node i and edge (i, j) are output of φ and Algorithm 1 Spatial derivative module (SDM)

Input: Graph signals ui and edge features eij = xj - xi on G where xi is a coordinate of node i.

Output: Spatial derivatives {ˆuk,i | i ∈ V and k ∈ K} where K = {∇x, ∇y , ∇2x , ∇2y }.

Require: Spatial derivative module {φk|k ∈ K}.

• 1: for k ∈ K do
• 2: {ak,i, bk,(i,j) | (i, j) ∈ E, k ∈ K} = φk({u} , {e} , G)
• 3: for i ∈ V do
• 4: ˆuk,i = ak,iui + ∑(j,i)∈E bk,(j,i)uj
• 5: end for
• 6: end for

Algorithm 2 Time derivative module (TDM)

Input: Graph signals u and approximated spatial derivatives ˆuk where k ∈ K on G. Time interval ∆t

Output: Prediction of signals at next time step ˆu(t)

Require: Time derivative module.

• 1: ˆut = TDM({ui, ˆuk,i | i ∈ V and k ∈ K})
• 2: ˆu(t) = u(t -1) + ˆut-1 · ∆t

they are linearly combined with the function values ui and uj. K denotes a set of finite difference operators. For example, if we set K = {∇x,∇y,∇2x ,∇2y }, we have 4 modules which approximate first/second order of spatial derivatives in 2-dimension, respectively.

3.2 Time Derivative Module

Once spatial derivatives are approximated, another learnable module is required to combine them for a target task. The form of line 2 in Alg. 2 comes from Eq. 3 and TDM is adapted to learn the unknown function F(·) in the equation. As our target task is the regression of graph signals, we use a recurrent graph network for TDM.

3.3 Meta-Learning with Auxiliary Objective

As discussed in Sec. 2.1, it is important to know spatial derivatives at time t to predict next signals at t+1 for locally conserved physical quantities. However, it is impractical to access the spatial derivatives in the sensor-based observations as they are highly discretized over space.

The meta-initialization with the auxiliary tasks from synthetic datasets is particularly important to address the challenge. First, the spatial modules can be universal feature extractors for modeling observations following unknown physics-based PDEs. Unlike other domains such as computer vision, it has been considered that there is no particular shareable architecture for learning spatiotemporal dynamics from physical systems. We propose that the PDE-independent spatial modules can be applicable as feature extractors across different dynamics as long as the dynamics follow a local form of conservation laws. Second, we can utilize synthetic data to meta-train the spatial modules as they are PDE-agnostic. This property allows us to utilize a large amount of synthetic datasets which are readily generated by numerical methods Algorithm 3 Meta-initialization with auxiliary tasks: Supervision of spatial derivatives Input: A set of meta-train task datasets D = {D1, . . . , DB} where Db = (Dtr b ,Dte b ).

where y(ak,·) i is an k-th auxiliary task label for the i-th node. Learning rate α and β.

Output: Meta-initialized spatial modules Φ.

• 1: Initialize auxiliary modules Φ = (φ1, . . . , φK)
• 2: while not converged do
• 3: for Db in D do
• 4: Φ′b = Φ-α∇Φ ∑Kk=1 Laux k (Dtr b ; φk)
• 5: end for
• 6: Φ←Φ-β∇Φ ∑Bb=1 ∑Kk=1 Laux k (Dte b ; φ′b,k )
• 7: end while

regardless of the exact form of PDE for target observations. Finally, we can provide a stronger inductive bias which is beneficial for modeling real-world observations but not available in the observations explicitly.

Alg. 3 describes how the spatial modules are metainitialized by MAML under the supervision of K different spatial derivatives. First, we generate values and spatial derivatives on a 2D regular grid from an analytical function. Then, we sample a finite number of points from the regular grid to represent discretized nodes and build a graph from the sampled nodes. Each graph signal and its discretization becomes input feature of a meta-train task and corresponding spatial derivatives are the auxiliary task labels. Fig. 2 visualizes graph signals and spatial derivatives for meta-initialization.

Once the spatial modules are initialized throughout metatraining, we reuse the modules for meta-test where the temporal module (the head of the network) are adapted on few observations from real-world sensors (Alg. 4). Although the standard MAML updates the body of the network (the spatial modules) as well, we only adapt the head layer (θ) as like almost-no-inner-loop method in [Raghu et al., 2020]. The task at test time is graph signal prediction and the temporal modules (θ) are adapted by a regression loss function L = ∑Tt=1 ||u(t) - ˆu(t)||2 on length T sequence (Dtrm) and evaluated on held-out (t > T) sequence (Dtem) with the adapted parameters.

Algorithm 4 Adaptation on meta-test tasks

Input: A set of meta-test task datasets D = {D1, . . . , DM} where Dm = (Dtrm, Dtem). Meta-initialized SDM (Φ). Learning rate α.

Output: Adapted TDM θ′m on the m-th task.

• 1: Initialize temporal modules (θ1, . . . , θM)
• 2: for Dm in D do
• 3: θ′m = θm -α∇θm L(Dtrm;Φ, θm)
• 4: end for

4 Spatial Derivative Modules: Reusable Modules

We have claimed that the spatial modules provide reusable features associated with spatial derivatives such as ∇xu,∇yu, and ∇2xu across different dynamics or PDEs. While it has been shown that the data-driven approximation of spatial derivatives is more precise than that of finite difference method [Seo et al., 2020; Bar-Sinai et al., 2019], it is not guaranteed that the modules effectively provide transferrable parameters for different spatial resolution, discretization, and fluctuation of function values. We explore whether the proposed spatial derivative modules based on graph networks can be used as a feature provider for different spatial functions and discretization.

We perform two sets of experiments: evaluate few-shot learning performance (1) when SDM is trained from scratch; (2) when SDM is meta-initialized. Fig. 2 shows how the graph signal and its discretization is changed over the different settings. If the number of nodes is large, it can provide spatially high-resolution and thus, the spatial derivatives can be more precisely approximated. Table 1 shows the parameters we used to generate synthetic datasets. Note that metatest data is designed to evaluate interpolation/extrapolation properties. Initial frequency decides the degree of fluctuation (In Fig. 2, the right one has higher F than that of the left one.). For each parameter combination, we generate 100 different snapshots based on a form in [Long et al., 2018]:

where λk,l, γk,l ∼ N (0, 0.02). The index i denotes the i-th node whose coordinate is (xi, yi) in the 2D space ([0, 2π] × [0, 2π]) and k, l are randomly sampled integers. From the synthetic data, the first- and second-order derivatives are analytically given and SDM is trained to approximate them.

The prediction results for spatial derivatives are shown in Table 2. The results show that the proposed module (SDM) is efficiently adaptable to different configuration on few samples from meta-initialized parameters compared to learning from scratch. The finding implies that the parameters for spatial derivatives can be generally applicable across different spatial resolution, discretization, and function fluctuation.

5 Experimental Evaluation

5.1 Preliminary: Which Synthetic Dynamics Need to Be Generated?

While Table 2 demonstrates that the PDE-independent representations are reusable across different configurations, it is still an open question: which topological configuration needs to be used to construct the synthetic dynamics? According to Table 2, the most important factor affecting error is an initial frequency (F), which determines min/max scales and fluctuation of function values, and it implies that the synthetic dynamics should be similarly scaled to a target dynamics. We use the same topological configuration in Table 1 to generate synthetic datasets for a task in Sec. 5.2 and Sec. 5.3.

5.2 Multi-step Graph Signal Generation

Task. We adopt a set of multi-step spatiotemporal sequence generation tasks to evaluate our proposed framework. In each task, the data is a sequence of L frames, where each frame is a set of observations on N nodes in space. Then, we train an auto-regressive model with the first T frames (T -shot) and generate the following L-T frames repeatedly from a given initial input (T-th frame) to evaluate its performance.

Datasets. For all experiments, we generate meta-train tasks with the parameters described in Table 1 and the target observations are 2 real-world datasets: (1) AQI-CO: national air quality index (AQI) observations [Berman, 2017]; (2) ExtremeWeather (EW): the extreme weather dataset [Racah et al., 2017]. For the AQI-CO dataset, we construct 12 metatest tasks with the carbon monoxide (CO) ppm records from the first week of each month in 2015 at land-based stations. For the extreme weather dataset, we select the top-10 extreme weather events with the longest lasting time from the year 1984 and construct a meta-test task from each event with the observed surface temperatures at randomly sampled locations. Since each event lasts fewer than 20 frames, each task has a very limited amount of available data. In both datasets, graph signals are univariate. Note that all quantities have fluidic properties such as diffusive and convection. More details are in the supplementary material.

Baselines. We evaluate the performance of a physics-aware architecture (PA-DGN) [Seo et al., 2020], which also consists of spatial derivative modules and recurrent graph networks (RGN), to see how the additional spatial information affects prediction performance for same architecture. Note that PA-DGN has same modules in PiMetaL and the difference is that PiMetaL utilizes meta-initialized spatial modules and PA-DGN is randomly initialized for learning from scratch on meta-test tasks. Additionally, the spatial modules in PADGN is replaced by finite difference method (FDM+RGN) to see if the numerical method provides better PDE-agnostic representations. The baselines and PiMetaL are trained on the meta-test support set only to demonstrate how the additional spatial information is beneficial for few-shot learning tasks.

Discussion. Table 3 shows the multi-step prediction performance of our proposed framework against the baselines on real-world datasets. Overall, PA-DGN and PiMetaL show similar trend such that the prediction error is decreased as longer series are available for few-shot adaptation. There are two important findings: first, with the similar expressive power in terms of the number of learnable parameters, the meta-initialized spatial modules provide high quality representations which are easily adaptable across different spatiotemporal dynamics in the real-world. This performance gap demonstrates that we can get a stronger inductive bias from synthetic datasets without knowing PDE-specific information. Second, the contribution of the meta-initialization is more significant when the length of available sequence is shorter (T = 5) and this demonstrates when the metainitialization is particularly effective. Finally, the finite difference method provides proxies of exact spatial derivatives and the representations are useful particularly when T = 5 but its performance is rapidly saturated and it comes from the gap between the learnable spatial modules and fixed numerical coefficients. The results provide a new point of view on how to utilize synthetic or simulated datasets to handle challenges caused by limited datasets.

5.3 Graph Signal Regression

Task, datasets, and baselines. [Defferrard et al., 2019] conducted a graph signal regression task: predict the temperature xt from the temperature on the previous 5 days (xt-5 : xt-1). We split the GHCN dataset spatially into two regions: (1) the USA (1,705 stations) and (2) Europe (EU) (703 stations) where there are many weather stations full functioning. In this task, the number of shots is defined as the number of input and output pairs to train a model.

As the input length is fixed, more variants of graph neural networks are considered as baselines. We concatenate the 5step signals and feed it into Graph convolutional networks (GCN) [Kipf and Welling, 2017], Graph attention networks (GAT) [Veliˇckovi´c et al., 2018], GraphSAGE [Hamilton et al., 2017], and Graph networks (GN) [Sanchez-Gonzalez et al., 2018] to predict next signals across all nodes.

Discussion. Table 4 shows the results of the graph signal regression task across different baselines and the proposed method. There are two patterns in the results. First, although in general we observe an improvement in performance for all methods when we move from the 5-shot setting to the 10shot setting, PiMetaL's performance yields the smallest error. Second, for the EU dataset, while 5-shot seems enough to achieve stable performance, it demonstrates that the PDEindependent representations make the regression error converge to a lower level. Overall, the experimental results prove that the learned spatial representations from simulated dynamics are beneficial for learning on limited data.

6 RelatedWork

Physics-informed learning. Since physics-informed neural networks are introduced in [Raissi et al., 2019], which find that a solution of a PDE can be discovered by neural networks, physical knowledge has been used as an inductive bias for deep neural networks. Advection-diffusion equation is incorporated with deep neural networks for sea-surface temperature dynamics [de Bezenac et al., 2018]. [Lutter et al., 2019; Greydanus et al., 2019] show that Lagrangian/Hamiltonian mechanics can be imposed to learn the equations of motion of a mechanical system and [Seo and Liu, 2019] regularize a graph neural network with a specific physics equation. Rather than using explicitly given equations, physics-inspired inductive bias is also used for reasoning dynamics of discrete objects [Battaglia et al., 2016] and continuous quantities [Seo et al., 2020]. [Long et al., 2018] propose a numeric-symbolic hybrid deep neural network designed to discover PDEs from observed dynamic data. To the best of our knowledge, we are the first to provide a framework to use the physics-inspired inductive bias under the meta-learning settings to tackle the limited data issue which is common for real-world data such as extreme weather events [Racah et al., 2017].

Meta-learning. The aim of meta-learning is to enable learning parameters which can be used for new tasks unknown at the time of learning, leading to agile models which adapt to a new task utilizing only a few samples. Based on how the knowledge from the related tasks is used, meta-learning methods have been classified as optimizationbased [Ravi and Larochelle, 2017; Finn et al., 2017], modelbased [Munkhdalai and Yu, 2017; Duan et al., 2017; Mishra et al., 2018], and metric-based [Snell et al., 2017]. Recently, another branch of meta-learning has been introduced to more focus on finding a set of reusable modules as components of a solution to a new task. [Alet et al., 2019] provide a framework, structured modular meta-learning, where a finite number of modules are introduced as task-independent modules and an optimal structure combining the modules is found from a limited number of data. To our knowledge, none of the above works provide a solution for modeling physics-related spatiotemporal dynamics where it is hard to generate enough tasks for meta-initialization.

7 Conclusion

In this paper, we propose a framework for physics-aware meta-learning with auxiliary tasks. By incorporating PDEindependent/-invariant knowledge from simulated data, the framework provide reusable features to meta-test tasks with a limited amount of data. Experiments show that auxiliary tasks and physics-aware meta-learning help construct reusable modules that improve the performance of spatiotemporal predictions when data is limited.

Acknowledgements

This work is supported in part by NSF Research Grant IIS1254206 and MINERVA grant N00014-17-1-2281, granted to co-author Yan Liu in her academic role at the University of Southern California. The views and conclusions are those of the authors and should not be interpreted as representing the official policies of the funding agency, or the U.S. Government.