M ETA -L EARNING D YNAMICS F ORECASTING U SING T ASK I NFERENCE
Anonymous authors
Paper under double-blind review
A BSTRACT
Current deep learning models for dynamics forecasting struggle with generalization.
They can only forecast in a speciﬁc domain and fail when applied to systems with different parameters, external forces, or boundary conditions.
We propose a model-based meta-learning method called DyAd which can generalize across heterogeneous domains by partitioning them into different tasks.
DyAd has two parts: an encoder which infers the time-invariant hidden features of the task with weak supervision, and a forecaster which learns the shared dynamics of the entire domain.
The encoder adapts and controls the forecaster during inference using adaptive instance normalization and adaptive padding.
Theoretically, we prove that the generalization error of such procedure is related to the task relatedness in the source domain, as well as the domain differences between source and target.
Experimentally, we demonstrate that our model outperforms state-of-the-art approaches on both turbulent ﬂow and real-world ocean data forecasting tasks.
I NTRODUCTION
Modeling dynamical systems with deep learning has shown great success in a wide range of systems from ﬂuid mechanics to neural dynamics (Tompson et al., 2017;
Chen et al., 2018; Kolter & Manek, 2019; Zoltowski et al., 2020; Li et al., 2021).
However, the main limitation of previous works is very limited generalizability.
Most approaches only focus on a speciﬁc system and train on past data in order to predict the future.
Thus a new model must be trained to predict a system with different dynamics.
Consider, for example, learning ﬂuid dynamics; shown in Fig.
1are two ﬂuid ﬂows with different degrees of turbulence.
Even though the ﬂows are governed by the same equations, the difference in buoyant forces would require two separate deep learning models to forecast.
Therefore, it is imperative to develop generalizable deep learning models for dynamical systems that can learn and predict well over a large heterogeneous domain.
time ?
forecasting on turbulent ﬂow.
The model needs to generalize to a ﬂow with a very different buoyant force.
Meta-learning (Thrun & Pratt, 1998; Baxter, 1998;
Finn et al., 2017), or learning to learn, improves generalization by learning multiple tasks from the environment.
Recent developments in meta-learning have been successfully applied to few-shot classiﬁcation (Munkhdalai & Yu, 2017), active learning (Yoon et al., 2018), and reinforcement learning (Gupta et al., 2018).
However, meta-learning in the context of forecasting high-dimensional physical dynamics has not been studied before.
The challenges with meta-learning dynamical systems are unique in that (1) we need to efﬁciently infer the latent representation of the dynamical system given observed time series data, (2) we need to account for changes in unknown initial and boundary conditions, and (3) we need to model the temporal dynamics across heterogeneous domains.
Our approach is inspired by the fact that similar dynamical systems may share time-invariant hidden features.
Even the slightest change in these features may lead to vastly different phenomena.
For example, in climate science, ﬂuids are governed by a set of differential equations called Navier-
Stokes equations.
Some features such as kinematic viscosity and external forces (e.g. gravity), are time-invariant and determine the ﬂow characteristics.
By inferring this latent representation, we can model diverse system behaviors from smoothly ﬂowing water to atmospheric turbulence.
Inspired by neural style transfer (Karras et al., 2019), we propose a model-based meta-learning method, called DyAd , which can rapidly adapt to systems with varying dynamics.
DyAd has two parts, an encoder g and a forecaster f .
The encoder maps different dynamical systems to timeinvariant hidden features representing constants of motion, boundary conditions, and external forces which characterize the system.
The forecaster f then takes the hidden representations and the past system states to forecast the future system state.
Controlled by the time-invariant hidden features, the forecaster has the ﬂexibility to adapt to a wide range of systems with heterogeneous dynamics.
Unlike gradient-based meta-learning techniques such as MAML (Finn et al., 2017), DyAd automatically adapts during inference using an encoder and does not require any retraining.
Similar to model-based meta-learning methods such as MetaNets (Munkhdalai & Yu, 2017), we employ a two-part design with an adaptable learner which receives task-speciﬁc weights.
However, for time series forecasting, since input and output come from the same domain, a support set of labeled data is unnecessary to deﬁne the task.
The encoder can infer the task directly from query input.
Our contributions include:
• A novel model-based meta-learning method ( DyAd ) for dynamic forecasting in large heterogeneous domains.
• An encoder capable of extracting the time-invariant hidden features of a dynamical system using time-shift invariant model structure and weak supervision.
• A new adaptive padding layer (AdaPad), designed for adapting to boundary conditions.
• Theoretical guarantees for DyAd on the generalization error of task inference in the source domain as well as domain adaptation to the target domain.
• Improved generalization performance on heterogeneous domains such as ﬂuid ﬂow and sea temperature forecasting
, even to new tasks outside the training distribution
. M ETHODS 2. M ETA - LEARNING IN DYNAMICS FORECASTING
Let x ∈ R
d be a d -dimensional state of a dynamical system governed by parameters ψ .
The problem of dynamics forecasting is that given a sequence of past states ( x t − l +1 , . . .
, x t ) , we want to learn a map f such that:
[Equation]
Here l is the length of the input series, and h is the forecasting horizon in the output.
Existing approaches for dynamics forecasting only predict future data for a speciﬁc system as a single task.
Here a task refers to forecasting for a speciﬁc system with a given set of parameters.
The resulting models often generalize poorly to different system dynamics.
Thus a new model must be trained to predict for each speciﬁc system.
To perform meta-learning, we identify each forecasting task by some parameters c ⊂ ψ , such as constants of motion, external forces, and boundary conditions.
We learn multiple tasks simultaneously and infer the task from data.
Here we use c for a subset of system parameters ψ , because we usually do not have the full knowledge of the system dynamics.
In the turbulent ﬂow example, the state x t is the velocity ﬁeld at time t .
Parameters c can represent Reynolds number, average vorticity, average magnitude, or a vector of all three.
Formally, let µ be the data distribution over X × Y representing the function f : X → Y where
X = R d × l and Y = R d × h .
Our main assumption is that the domain X can be partitioned into separate tasks X
= ∪ c ∈C X c , where X c is the domain for task c and C is the set of all tasks.
The data in the same task share the same set of parameters.
Let µ c be the conditional distribution over X c × Y for task c .
During training, the model is presented with data drawn from a subset of tasks { ( x, y ) : ( x, y ) ∼ µ c , c ∼ C } .
Our goal is to learn the function f : X → Y over the whole domain X which can thus generalize across all tasks c ∈ C .
To do so, we need to learn the map g : X → C taking x ∈ X c to c in order to infer the task with minimal supervision.
2. D Y A D : D YNAMIC A DAPTATION N ETWORK
We propose a model-based meta-learning approach for dynamics forecasting.
Given multiple forecasting tasks, we propose to learn the function f in two stages, that is, by ﬁrst inferring the task c from the input x , and then adapting to a specialized forecaster f c : X c → Y for each task.
An alternative is to use a single deep neural network to directly model f in one step over the whole domain.
But this requires the training set to have good and uniform coverage of the different tasks.
If the data distribution µ is highly heterogeneous or the training set is not sampled i.i.d.
from the whole domain, then a single model may struggle with generalization.
We hypothesize that by partitioning the domain into different tasks, the model would learn to pick up task-speciﬁc features without requiring uniform coverage of the training data.
Furthermore, by separating task inference and forecasting into two stages, we allow the forecaster to rapidly adapt to new tasks that never appeared in the training set.
As shown in Fig.
2, our model consists of two parts: an encoder g and a forecaster f .
We introduce z c as a time-invariant hidden feature for task c .
We assume that c depends linearly on the hidden feature for simplicity and easy interpretation.
We design the encoder to infer the hidden feature z c given the input x .
We then use z c to adapt the forecaster f to the speciﬁc task, i.e., model y =
f c ( x ) as y = f
( x, z c ) .
As the system dynamics are encoded in the input sequence x , we can feed the same input sequence x to a forecaster and generate predictions ˆ y =
f c ( x ) .
2. E NCODER N ETWORK
The encoder maps the input x to the hidden features z c that are time-invariant.
To enforce this inductive bias, we encode time-invariance both in the architecture and in the training objective.
Time-Invariant Encoder.
The encoder is implemented using 4 Conv 3D layers, each followed by
BatchNorm , LeakyReLU , and max-pooling.
Note that theoretically, max-pooling is not perfectly shift invariant since 2x2x2 max-pooling is equivariant to shifts of size 2 and only approximately invariant to shifts of size 1.
But standard convolutional architectures often include max-pooling layers to boost performance.
We convolve across spatial and temporal dimensions.
After that, we use a global mean-pooling layer and a fully connected layer to estimate the hidden feature ˆ z c .
The task parameter depends linearly on the hidden feature.
We use a fully connected layer to compute the parameter estimate ˆ c .
Since convolutions are equivariant to shift (up to boundary frames) and mean pooling is invariant to shift, the encoder is shift-invariant.
In practice, shifting the time sequence one frame forward will add one new frame at the beginning and drop one frame at the end.
This creates some change in output value of the encoder.
Thus, practically, the encoder is only approximately shift-invariant.
Encoder Training.
The encoder network g is trained ﬁrst.
To combat the loss of shift invariance from the change from the boundary frames, we train the encoder using a time-invariant loss.
Given two training samples ( x ( i ) , y ( i ) ) and ( x ( j ) , y ( j ) ) and their task parameters c , we have loss
L enc =
[Equation]
, where ˆ z ( i )
=
g ( x ( i ) ) and ˆ z ( j ) = g ( x ( j ) ) and ˆ c ( i ) =
W ˆ
z
( i ) c
+ b is an afﬁne transformation of z c .
The ﬁrst term (cid:107) ˆ c − c (cid:107) 2 uses weak supervision of the task parameters whenever they are available.
Such weak supervision helps guide the learning of hidden feature z c for each task.
While not all parameters of the dynamical system are known, we can compute approximate values in the datum c ( i ) based on our domain knowledge.
For example, instead of the Reynolds number of the ﬂuid ﬂow, we can use the average velocity as a surrogate for task parameters.
The second term (cid:107) ˆ z ( i ) c − ˆ z ( j ) c (cid:107) 2 is the time-shift invariance loss, which penalizes the changes in latent variables between samples from different time steps.
Since the time-shift invariance of convolution is only approximate, this loss term drives the time-shift error even lower.
The third term |(cid:107) ˆ z ( i ) c
(cid:107) − m |( m is a positive value) prevents the encoder from generating small ˆ z ( i )
c
due to time-shift invariance loss.
It also helps the encoder to learn more interesting z , even in the absence of weak supervision.
Hidden Features.
The encoder learns time-invariant hidden features.
These hidden features resemble the time-invariant dimensionless parameters (Kunes, 2012) in physical modeling, such as Reynolds number in ﬂuid mechanics.
The hidden features may also be viewed as partial disentanglement of the system state.
As suggested by Locatello et al.
; Nie et al.
(2020), our disentanglement method is guided by inductive bias and training objectives.
Unlike complete disentanglement, as in e.g. Massague et al.
(2020), in which the latent representation is factored into time-invariant and time-varying components, we focus only on time-shift-invariance.
Nonetheless, the hidden features can control the forecaster which is useful for generalization.
2. F ORECASTER N ETWORK
𝒙
FC z
AdaPad 𝒚
The forecaster incorporates the hidden feature z c from the encoder and adapts to the speciﬁc forecasting task f c = f ( · , z c ) .
In what follows, we use z for z c .
We use two specialized layers, adaptive instance normalization ( AdaIN ) and adaptive padding ( AdaPad ).
AdaIN has been used in neural style transfer (Karras et al., 2019; Huang & Belongie, 2017) to control generative networks.
Here, AdaIN may adapt for speciﬁc coefﬁcients and external forces.
We also introduce a new layer, AdaPad( x, z ) , which is designed for encoding the boundary conditions of dynamical systems.
In principle, the backbone of the forecaster network can be any sequence prediction model.
We use a design that is similar to ResNet for spatiotemporal sequences.
AdaIN.
We employ AdaIN to adapt the forecaster network.
Denote the channels of input x by x i and let µ ( x i ) and σ ( x i ) be the mean and standard deviation of channel i .
For each AdaIN layer, a particular style is computed s =
( µ i , σ i ) i =
Az + b , where the linear map A and bias b are learned weights.
Adaptive instance normalization is then deﬁned as y i =
σ i x i − µ ( x i ) σ ( x i ) + µ i .
In essence, the channels are renormalized to the style s .
For dynamics forecasting, the hidden feature z encodes data analogous to the various coefﬁcients of a differential equation and external forces on the system.
In numerical simulation of a differential equation these coefﬁcients enter as scalings of different terms in the equation and the external forces are added to the combined force equation (J.C.Butcher, 1996).
Thus in our context AdaIN , which scales channels and adds a global vector, is well-suited to injecting this information.
AdaPad.
To complement AdaIN , we introduce AdaPad , which encodes the boundary conditions of each speciﬁc dynamical system.
Generally when predicting dynamical systems, error is introduced along the boundaries since it is unknown how the dynamics interact with the boundary of the domain, and there may be unknown inﬂows or outﬂows.
In our method, the inferred hidden feature z may contain the boundary information.
AdaPad uses the hidden features to compute the boundary conditions via a linear layer.
Then it applies the boundary conditions as padding immediately outside the spatial domain in each layer, as shown in Fig.
Forecaster Training.
The forecaster network is trained separately after the encoder.
The kernels of the convolutions and the mappings of the AdaIN and AdaPad layers are all trained simultaneously as the forecaster network is trained.
Denote the true state as y and the predicted state as ˆ y , we compute the loss per time step (cid:107) ˆ y − y (cid:107) 2 for each example.
We accumulate the loss over different time steps and generate multi-step forecasts in an autoregressive fashion.
T HEORETICAL A NALYSIS
The high-level idea of our method is to learn a good representation of the dynamics that generalizes well across a heterogeneous domain, and then adapt this representation to make predictions on new tasks.
Our model achieves this by learning on multiple tasks simultaneously and then adapting to new tasks with domain transfer.
We prove that learning the tasks simultaneously as opposed to independently results in better generalization. We also ﬁnd theoretical support for training the encoder and forecaster separately.
See Appendix B for a longer treatment with proofs.
Suppose we have K tasks { c k } Kk =1 ∼ C , each of which is sampled from a continuous parameter space C .
Here c k are the parameters of the task, which can be inferred by the encoder.
For each task c k , we have a collection of data ˆ µ c k of size n , sampled from µ k , a shorthand for µ c k .
Multi-task Learning Error.
Our model performs multi-task representation learning (Maurer et al., 2016) with joint risk (cid:15) =
(1 /K )
(cid:80) k (cid:15) k , which is the average risk of each task (cid:15) k deﬁned separately.
Denote the corresponding empirical risks ˆ (cid:15) and ˆ (cid:15) k .
We can bound the true risk (cid:15) using the empirical risk ˆ (cid:15) and Rademacher complexity R ( F ) of the hypothesis class F .
The following theorem restates the main result in Ando et al.
(2005) with simpliﬁed notations.
Theorem 3.1.
(Ando et al.
Assume the loss is bounded l ≤ 1 / 2 .
Given n samples each from
K different forecasting tasks µ 1 , · · · , µ k , then with probability at least 1 − δ , the following inequality holds for each f ∈ F in the hypothesis class:
[Equation]
The following inequality compares the performance for multi-task learning to learning individual tasks.
Let R k ( F ) be the Rademacher complexity for F over µ k .
Lemma 3.2.
The Rademacher complexity for multi-task learning is bounded as
[Equation]
We can now compare Theorem 3.1 to the bound obtained by considering each task individually.
Proposition 3.3.
Assume the loss is bounded l ≤ 1 / 2 , then the generalization bound given by considering each task individually is
[Equation]
The upper bound in Theorem 3.1 is strictly tighter than that of Proposition 3.3 as the ﬁrst terms ˆ (cid:15) ( f ) are equal, the second term is smaller by Lemma 3.2 and the third is smaller since 1 / √ 2 nK ≤ 1 / √ 2 n .
This helps explain why our multitask learning forecaster has better generalization than learning each task independently.
The shared data tightens the generalization bound.
Ultimately, though a tighter upper bound suggests lower error, it does not strictly imply it.
We further verify this theory in our experiments by comparison to baselines which learn each task independently.
Encoder versus Forecaster Error. Error from DyAd may result from either the encoder g φ or the forecaster f θ .
Using (Redko et al., 2017) , we can bound the generalization error in terms of: (1) the empirical error of the encoder, (2) the empirical forecaster error on the source domains, and (3) the
Wasserstein distance W 1 between the source and target domains.
We use Wasserstein distance since different tasks may have disjoint support and thus inﬁnite KL divergence, however, samples from tasks c, c (cid:48) close in C may be close in W 1 ( µ c , µ c (cid:48) )
(see Appendix Fig. 11).
Our two-stage training scheme is motivated by this decomposition for achieving generalization.
Our hypothesis space has the form { x
(cid:55)→ f θ ( x, g φ ( x )) } where φ and θ are the weights of the encoder and forecaster respectively.
Let (cid:15)
X be the error over the entire domain X , that is, for all c .
Let (cid:15) enc ( g φ ) =
E x ∼X ( L 1 ( g ( x ) , g φ ( x )) be the encoder error where g : X → C is the ground truth. Let G = { g φ : X → C} be the task encoder hypothesis space. Denote the empirical risk of g φ by ˆ (cid:15) enc ( g φ ) .
Proposition 3.4.
Assume c
(cid:55)→ f θ ( · , c ) is Lipschitz continuous with Lipschitz constant γ uniformly in θ and l ≤ 1 / 2 .
Let λ c = min f ∈F ( (cid:15) c ( f ) + 1 /K (cid:80) Kk =1 (cid:15) c k ( f )) . For large n and probability ≥ 1 − δ , (cid:15) X ( f θ ( · , g φ ( · ))) ≤ γ ˆ (cid:15) enc ( g φ ) + 1 K
K (cid:88) k = ˆ (cid:15) c k ( f θ ( · , c k ))
+ E c ∼C (cid:34) W(cid:32) ˆ µ c , 1 K K (cid:88) k
= ˆ µ c k (cid:33) + λ c (cid:35) +2 γR ( G )+2 R ( F ) + ( γ + 1) (cid:112) log(1 /δ ) / (2 nK ) + (cid:112) 2 log(1 /δ ) (cid:16)(cid:112) 1 /n + (cid:112) 1 / ( nK ) (cid:17) . Although this result does not settle either the question of end-to-end versus pre-training or encoderforecaster versus monolithic model, it quantiﬁes the trade-offs these choices depend on. We empirically consider both questions in the experiments section.
R ELATED W ORK
Learning Dynamical Systems.
Deep learning models are gaining popularity for learning dynamical systems (Shi et al., 2017;
Chen et al., 2018; Kolter & Manek, 2019; Azencot et al., 2020a; Xie et al., 2018;
Tompson et al., 2017; Pfaff et al., 2021).
An emerging topic is physics-informed deep learning (Raissi et al., 2017;
Lutter et al., 2018; Azencot et al., 2020b; de Bezenac et al., 2018; Wang et al., 2020b; Ayed et al., 2019b;a; Li et al., 2021; Donà et al., 2021) which integrates inductive biases from physical systems to improve learning.
For example, Lutter et al.
(2018) encoded Euler-Lagrange equation into the deep neural nets but focus on learning low-dimensional trajectories.
Morton et al.
; Azencot et al.
(2020b) incorporated Koopman theory into the architecture.
Maziar Raissi (2019) used deep neural networks to solve PDEs with physical laws enforced in the loss functions.
Wang et al.
(2020a) proposed a hybrid approach by marrying two well-established turbulent ﬂow simulation techniques with deep learning to produce better prediction of turbulence.
However, these approaches deal with a speciﬁc system dynamics instead of the meta-learning problem in this work.
Multi-task learning and Meta-learning Multi-task learning (Vandenhende et al., 2021) focuses on learning shared representations from multiple related tasks. Architecture-based MTL methods can be categorized into encoder-focused (Liu et al., 2019) and decoder-focused (Xu et al., 2018). There are also optimization-based MTL methods, such as task balancing methods (Kendall et al., 2018). But
MTL assumes tasks are known a priori instead of inferring the task from data.
On the other hand, the aim of meta-learning (Thrun & Pratt, 1998) is to leverage the shared representation to fast adapt to unseen tasks.
Based on how the meta-level knowledge is extracted and used, meta-learning methods are classiﬁed into model-based (Munkhdalai & Yu, 2017; Alet et al., 2018; Oreshkin et al., 2019;
Seo et al., 2020; Zhou et al., 2021; Li et al., 2017), metric-based (Vinyals et al., 2016;
Snell et al., 2017) and gradient-based (Finn et al., 2017; Andrychowicz et al., 2016; Rusu et al., 2019; Grant et al., 2018;
Yao et al., 2019).
Most meta-learning approaches are not designed for forecasting with a few exceptions.
Oreshkin et al.
(2019) designed a residual architecture for time series forecasting with a meta-learning parallel.
Alet et al.
(2018) proposed a modular meta-learning approach for continuous control.
But forecasting physical dynamics poses unique challenges to meta-learning as we seek ways to encode physical knowledge into our model.
Style Transfer.
Our approach is inspired by neural style transfer techniques.
In style-transfer, a generative network is controlled by an external style vector though adaptive instance normalization between convolutional layers.
Our hidden representation bears afﬁnity with the “style” vector in style transfer techniques.
Rather than aesthetic style in images, our hidden representation encodes time-invariant features.
Style transfer initially appear in non-photorealistic rendering (Kyprianidis et al., 2012).
Recently, neural style transfer (Jing et al., 2019) has been applied to image synthesis (Gatys et al., 2016; Liu et al., 2021), videos generation (Ruder et al., 2016), and language translation (Prabhumoye et al., 2018).
For dynamical systems, Sato et al.
(2018) adapts texture synthesis to transfer the style of turbulence for animation.
Kim & Lee (2019) studies unsupervised generative modeling of turbulent ﬂows but for super-resolution reconstruction rather than forecasting.
Video Prediction.
Our work is also related to video prediction.
Conditioning on the historic observed frames, video prediction models are trained to predict future frames, e.g., (Mathieu et al., 2015;
Finn et al., 2016; Xue et al., 2016; Villegas et al., 2017; Oprea et al., 2020; Finn et al., 2016; Wang et al., 2017; 2021; Le Guen & Thome, 2020; Massague et al.).
There is also conditional video prediction (Oh et al., 2015) which achieves controlled synthesis.
Many of these models are trained on natural videos from unknown physical processes.
Our work is substantially different because we do not attempt to predict object or camera motions.
However, our method can be potentially combined with video prediction models to improve generalization.
E XPERIMENTS
We compare our model with a series of baselines on the multi-step forecasting with different dynamics.
We consider two testing scenarios: (1) dynamics with different initial conditions (test-future) and (2) dynamics with different parameters such as external force (test-domain).
The ﬁrst scenario evaluates the models’ ability to extrapolate into the future for the same task.
The second scenario estimates the capability of the models to generalize across different tasks.
We experiment on three datasets: synthetic turbulent ﬂows, real-world sea surface temperature and ocean currents data.
These are difﬁcult to forecast using numerical methods due to unknown external forces and complex dynamics not fully captured by simpliﬁed mathematical models.
For the weak supervision c , we use the mean vorticity for turbulence and ocean currents, and mean temperature for sea surface temperature.
We defer the details of the datasets and experiments to Appendix A.2.
[Table]
Baselines.
We include several SoTA baselines from meta-learning, as well as common methods for dynamics forecasting.
ResNet (He et al., 2016):
A widely adopted video prediction model (Wang et al., 2020b).
• U-net
(Ronneberger et al., 2015): Originally developed for biomedical image segmentation.
• ResNet/Unet-c
Above ResNet and Unet with an additional ﬁnal layer that generates task parameter c and trained with weak-supervision and forecasting loss altogether.
• PredRNN (Wang et al., 2017):
A state-of-the-art RNN-based spatiotemporal forecasting model.
• VarSepNet
(Donà et al., 2021):
A convolutional dynamics forecasting model based on spatiotemporal disentanglement.
• Mod-attn :
A modular meta-learning method which combines modules to generalize to new tasks (Alet et al., 2018) using attention.
• Mod-wt : A modular meta-learning variant which uses attention weights to combine the parameters of the convolutional kernels for new tasks.
• MetaNet (Munkhdalai & Yu, 2017):
A model-based meta-learning method which requires a few labels from test tasks as a support set to adapt.
• MAML (Finn et al., 2017):
A popular optimization-based meta-learning approach.
We replaced the classiﬁer in the original model with a ResNet for regression.
Note that both ResNet-c and Unet-c have access to task parameters c , but cannot adapt.
Mod-attn and Mod-wt also use a convolutional encoder to generate attention weights.
MetaNet requires samples from test tasks as a support set and MAML needs adaptation retraining on test tasks, while other models do not need any information from the test domains.
VarSepNet employs separation of variables through different loss terms but it is difﬁcult to ﬁnd the optimal hyperparameters for these terms.
To demonstrate the generality of DyAd , we experimented with both ResNet and
U-net as our forecaster.
See details about baselines in Appendix A.2.
Experiment Setup.
For all datasets, we use a sliding window approach to generate samples of sequences.
For test-future, we train and test on the same tasks but different time steps.
For testdomain, we train and test on different tasks with an 80-20 split.
All models are trained to make next step prediction given the previous steps as input.
We forecast in an autoregressive manner to generate multi-step ahead predictions.
All results are averaged over 3 runs with random initialization.
Apart from the root mean square error (RMSE), we also report the energy spectrum error (ESE) for ocean current prediction which quantiﬁes the physical consistency.
ESE indicates whether the predictions preserve the correct statistical distribution and obey the energy conservation law, which is a critical metric for physical consistency.
See details about energy spectrum in Appendix A.3.
5. E XPERIMENT R ESULTS
Prediction Performance.
Table 1 shows the RMSE of multi-step predictions on Turbulent Flows (20 steps), Sea Surface Temperature (10 steps), and Ocean Currents (10 step) in two testing scenarios.
We observe that DyAd makes the most accurate predictions in both scenarios across all datasets.
Comparing ResNet/Unet-c with DyAd , we observe the clear advantage of task inference with separate training.
VarSepNet achieves competitive performances on Ocean Currents (second best) through spatiotemporal disentanglement but cannot adapt to future tasks.
Table 1 also reports ESE on
Ocean Currents.
DyAd not only has small RMSE but also obtains the smallest ESE, suggesting it captures the statistical distribution of ocean currents well.
Figure 4 shows the target and the predicted velocity norm ( √ u 2 + v 2 ) by Unet-c , Modular-wt and DyAd at time step 1, 5, 10 for Turbulent Flows with buoyancy factors 9 and 21 respectively.
We can see that DyAd can generate realistic ﬂows with the corresponding characteristics while the baselines have trouble understanding and disentangling the buoyancy factor.
[Table] [Table]
Figure 5 shows DyAd , ResNet , U-net , PredRNN predictions on an ocean current sample in the future test set, and we see the shape of predictions by DyAd is closest to the target.
These results demonstrate that DyAd not only forecasts well but also accurately captures the physical characteristics of the system.
We also visualize the energy spectrum of target and predictions by ResNet , U-net and DyAd on two test sets of ocean currents in Figure 10, with DyAd being the closest to the target.
Ablation Study.
We performed an ablation study of
DyAd on the turbulence dataset to understand the contribution of each component, shown in Table 2.
We ﬁrst remove the encoder from DyAd while keeping the same forecaster network ( No_enc ).
The resulting model degrades but still outperforms ResNet .
This demonstrates the effectiveness of AdaIN and AdaPad for forecasting.
We also tested DyAd with AdaIN only ( No_AdaPad ), and the performance without AdaPad was slightly worse.
Another notable feature of our model is the ability to infer tasks with weakly supervised signals c .
It is important to have a c that is related to the task domain.
As an ablative study, we fed the encoder in DyAd with a random c , leading to Wrong_enc .
We can see that having the wrong supervision may hurt the forecasting performance.
We also trained the encoder and the forecaster in DyAd altogether ( End2End ) but observed worse performance.
This validates our hypothesis about the signiﬁcance of domain partitioning and separate training strategy.
We also tested 5 different alternatives to AdaIN for injecting the hidden feature z c into the forecaster, and reported the results in Table 5 of Appendix A.4.
Controllable Forecast.
DyAd infers the hidden features from data, which allows direct control of the latent space in the forecaster.
We tried varying the encoder input while keeping the forecaster input ﬁxed.
Figure 6 shows the forecasts from DyAd when the encoder is fed with ﬂows having different buoyancy factors c = 5 , 15 , 25 .
As expected, with higher buoyancy factors, the predictions from the forecaster become more turbulent.
This demonstrates that the encoder can successfully disentangle the latent representation of difference tasks, and control the predictions of the forecaster.
C ONCLUSION
We propose a model-based meta-learning method, DyAd to forecast physical dynamics.
DyAd uses an encoder to infer the parameters of the task and a prediction network to adapt and forecast giving the inferred task.
Our model can also leverage any weak supervision signals that can help distinguish different tasks, allowing the incorporation of additional domain knowledge.
On challenging turbulent ﬂow prediction and real-world ocean temperature and currents forecasting tasks, we observe superior performance of our model across heterogeneous dynamics.
Future work would consider non-grid data such as ﬂows on a graph or a sphere.
E THICS S TATEMENT
Our model allows for generalizable predictions of dynamical systems.
One example domain we apply it to is turbulent ﬂow predictions.
While there are many peaceful applications of turbulent ﬂow prediction, our method could also potentially be helpful for researching explosive weapons or for designing aircraft or missiles.
R EPRODUCIBILITY S TATEMENT
The implementation code of DyAd and baselines are included in the supplementary material.
The readme ﬁle includes the instructions of downloading PhiFlow and generating the turbulent ﬂow dataset.
The ocean currents and sea temperature can be downloaded manually from the link provided in the Appendix A.2.
The full proofs of the theorems in Section 3 can be found in Appendix B
A I MPLEMENTATION D ETAILS A. M ODEL D ESIGN
The prediction network ˆ y = f
( x, z ) is composed of 8 blocks.
Each block operates on a hidden state
h ( i ) of shape B × H × W × C in and yields a new hidden state h ( i +1) of the shape B × H × W × C out .
The ﬁrst input is h
= x and the ﬁnal output is computed from the ﬁnal hidden state as ˆ y =
Conv2D ( h (8) ) .
We deﬁne each block as a
[Equation]
as illustrated in Figure
[Figure]
[Figure]
A. E XPERIMENT D ETAILS A.2. D ATASETS
Turbulent Flow with Varying Buoyancy.
We generate a synthetic dataset of turbulent ﬂows with a numerical simulator, PhiFlow 1 .
It contains 64 × 64 velocity ﬁelds of turbulent ﬂows in which we vary the buoyant force acting on the ﬂuid from 1 to 25.
Each buoyant force corresponds to a forecasting task and there are 25 tasks in total.
We use the mean vorticity of each task as partial supervision c as we can directly calculate it from the data.
Vorticity can characterize formation and circular motion of turbulent ﬂows.
Sea Surface Temperature.
We evaluate on a real-world sea surface temperature data generated by the NEMO ocean engine (Madec et al., 2015) 2 .
We select an area from Paciﬁc ocean range from 01/01/2018 to 12/31/2020.
The corresponding latitude and longitude are (-150 ∼ -120, -20 ∼ -50).
This area is then divided into 25 64 × 64 subregions, each is a task since the mean temperature varies a lot along longitude and latitude.
For the encoder training, we use season as an additional supervision signal besides the mean temperature of each subregion.
In other words, the encoder should be able to infer the mean temperature of the subregion as well as to classify four seasons given the temperature series.
Ocean Currents.
We also experiment with the velocity ﬁelds of ocean currents from the same region and use the same task division as the sea surface temperature data set.
Similar to the turbulent ﬂow data set, we use the mean vorticity of each subregion as the weak-supervision signal.
A.2. B ASELINES .
For fair comparison, we set these models to have equal capacity as DyAd in terms of number of parameters.
Hyperparameters including learning rate, input length and the number of steps of accumulated loss for training are tuned on validation sets.
Modular-attn has a convolutional encoder f
that takes the same input x as each module M to generate attention weights, (cid:80)
[Equation]
mk =1 exp[ f ( x )( k )]
M l ( x ) .
Modular-wt also has the same encoder but to generate weights for combining the convolution parameters of all modules.
We use additional samples of up to 20% of the test set from test tasks.
MetaNet uses these as a support set.
MAML is retrained on these samples for 10 epoch for adaptation.
A.2. H YPERPARAMETER TUNING .
We tuned learning rate (1e-3 ∼ 1e-5), batch size (16 ∼ 64), the number of accumulated errors for backpropogation (2 ∼ 5), and hidden size (64 ∼ 512) of Modular Networks and Meta-Nets.
We ﬁxed the number of historic input frames as 20.
When we trained the encoder on turbulent ﬂows and sea surface temperature, we used α = 1 and β = 1 .
For ocean currents, we used α = 0 .
2 and β = 0 .
We performed all our experiments on 4 V100 GPUs.
A.2. M ODEL C APACITY .
Table 3 displays the number of parameters of each tuned model on the turbulent ﬂow dataset. [Table] [Table] [Table]
[Table]
The data are available at https://resources.marine.copernicus.eu/?option=com_ csw&view=
details&product_id
=GLOBAL_ANALYSIS_FORECAST_PHY_001_
A. T URBULENCE KINETIC ENERGY SPECTRUM
The turbulence kinetic energy spectrum E ( k ) is related to the mean turbulence kinetic energy as (cid:90) ∞ E ( k ) dk
= (( u (cid:48) ) + ( v (cid:48) ) 2 ) / 2 , ( u (cid:48) )
2 = 1 T T (cid:88) t = ( u ( t ) − ¯ u ) 2 , where the k is the wavenumber and t is the time step.
Figure 9 shows a theoretical turbulence kinetic energy spectrum plot.
The spectrum can describe the transfer of energy from large scales of motion to the small scales and provides a representation of the dependence of energy on frequency.
Thus, the Energy Spectrum Error can indicate whether the predictions preserve the correct statistical distribution and obey the energy conservation law.
A trivial example that can illustrate why we need
ESE is that if a model simply outputs moving averages of input frames, the accumulated RMSE of predictions might not be high but the ESE would be really big because all the small or even medium eddies are smoothed out.
A. A DDITIONAL R ESULTS
B T HEORETICAL A NALYSIS
The high-level idea of our method is to learn a good representation of the underlying dynamics from multiple tasks, and then transfer this representation to a target task (domain adaptation). [Table]
[Table]
Deﬁnition 1 (Forecasting task) .
Each forecasting task x t +1 = f
( x t , . . . ) is to learn a conditional distribution µ over the system states µ : p ( x t +1 | x t , . . . )
conditioned on the sequence of previous states where µ is a probability measure.
In our setting we have K tasks, each of which is sampled from a continuous, ﬁnite space { c k } ∼ C .
Let µ k be the corresponding conditional probability measure p ( x t , . . .
, x 1 | c k ) .
For each task c k , we have a collection of n series as realizations from the dynamical system
X k = { ( x t , . . .
, x 1 ; c k ) ( i ) } ni = sampled from µ k .
The semicolon here represents the system behavior in a speciﬁc domain c k .
Let
X = (cid:83)
k
X k be the union of samples over all tasks.
In practice, we often have some intuition of the variables that dictate the domain.
Therefore, we have two possible scenarios for the role of c in dynamical systems:
c fully distinguishes the task: the differences in X k can be completely explained by the differences in c k ; 2.
c partially distinguishes the task: a more realistic scenario where we only have partial knowledge of the domain.
There exist latent variables z (cid:48) that need to be inferred from raw data.
Together z
=
[ c, z (cid:48) ] can describe the behavior of the system in a domain.
We assume Scenario 1, which resembles the multi-task representation learning setting (Maurer et al., 2016) with joint true risk over all tasks (cid:15) and individual task true risk (cid:15) k deﬁned respectively
[Equation]
and corresponding empirical risks
[Equation]
where l is a loss function.
B. M ULTI -T ASK L EARNING E RROR
We want to bound the true loss (cid:15) using the empirical loss ˆ (cid:15) and Rademacher complexity of the hypothesis class F .
We can use the classic results from Ando et al.
Deﬁne empirical
Rademacher complexity for samples from all tasks as
[Equation]
where { σ ( i ) k } are independent binary variables σ ( i ) k ∈ {− 1 , 1 } .
The true Rademacher complexity is then deﬁned R ( F ) = E X ( ˆ R X ( F )) .
The following theorem restates the main result from Ando et al.
We simplify the statement by using Rademacher complexity rather than the set cover number argument used in the original proof.
Theorem B.1.
Ando et al.
Given data from K different forecasting tasks µ 1 , · · · , µ k and f in hypothesis class F , for some constant C with probability at least 1 − δ , the following inequality holds:
[Equation]
If we assume the loss is bounded l ≤ 1 / 2 , then we may take C = 1 / √ 2 .
Proof.
Consider { x ( i ) k } as independent random variables.
For a function φ that satisﬁes
[Equation]
by McDiarmid’s inequality, we have
[Equation]
Applying this inequality to the max difference Q ( X ) =
sup f ∈F [ (cid:15) ( f ) − ˆ (cid:15) ( f, X )] , then with probability at least 1 − δ , we have
[Equation]
where C is a constant depending on the bounds
c i .
If the loss l ≤ 1 / 2 , then | Q | ≤ 1 / 2 and so we can take c i = 1 leading to C = 1 / √ 2 .
A standard computation (see Mohri et al.
(2018), Theorem 3.3) using the law of total expectation shows E X [ Q ( X )] ≤ 2 R ( F ) , which ﬁnishes the proof.
We can use this to compare the generalization error of multi-task learning versus that of learning the individual tasks.
The following inequality compares the Rademacher complexity for multi-task learning to that of individual task learning.
Denote ˆ R X k
and
R k the empirical and true Rademacher complexity for F over µ k .
Lemma B.2.
The Rademacher complexity for multi-task learning is bounded R ( F ) ≤
[Equation]
Proof.
We compute the empirical Rademacher complexity, ˆ
[Equation]
(cid:17)(cid:17)(cid:33)(cid:35)
[Equation]
[Equation]
The ﬁrst inequality follows from the sub-additivity of the supremum function.
The next equality is due to the fact positive scalars commute with supremum, and by the linearity of expectation.
The expectation E σ reduces to the expectation E σ k over only those Rademacher variables appearing inside the expectation.
R k ( F ) is the Rademacher complexity of the function on the individual task k .
Taking expectation over all samples X gives the result.
It is instructive to compare the bound from Theorem B.1 with the generalization error bound obtained by considering each task individually.
Proposition B.3.
Assume n = n k for all tasks k
and the loss l is bounded l ≤ 1 / 2 , then the generalization bound given by considering each task individually is
[Equation]
which is strictly looser than the bound from Theorem B.1 under the same assumptions.
This result helps to explain why our multitask learning framework has better generalization than learning each task independently.
The shared data tightens the generalization bound.
Proof.
Applying the classical analog of Theorem B.1 to a single task, we ﬁnd with probability greater than 1 − δ ,
[Equation]
Averaging over all tasks yields
[Equation]
Since the loss l is bounded l ≤ 1 / 2 , we can take C =
C k = 1 / √ 2 , giving the generalization upper bound for the joint error of Equation 7.
By Lemma B.2 and the fact 1 / √ 2 nK ≤ 1 / √ 2 n , the bound in Theorem B.1 is strictly tighter.
Gap Between Bounds.
To quantify how much tighter the bound of Theorem B.1 is relative to
Proposition B.3, we compute the difference in upper bounds.
Assuming l ≤ 1 / 2 , the gap between the bounds is,
[Equation]
Simplifying gives [Equation] The ﬁrst term is positive by Lemma B.2 and second term is positive since √ K ≥ 1 . B. D OMAIN A DAPTATION E RROR
Since we test on c ∼ C outside the training set { c k } , we incur error due to domain adaptation from the source domains µ c 1 , . .
. , µ c K to target domain µ c with µ being the true distribution.
Denote the corresponding empirical distributions of n samples per task by ˆ µ c = 1 n c (cid:80) n c i =1 δ x ( i ) c .
For different c and c (cid:48) , the domains µ c and µ c (cid:48) may have largely disjoint support, leading to very high KL divergence.
However, if c and c (cid:48) are close, samples x c ∼ µ c and x c (cid:48) ∼ µ c (cid:48) may be close in the domain X with respect to the metric (cid:107) · (cid:107) X .
For example, if the external forces c and c (cid:48) are close, given x c ∼ µ c there is likely x c (cid:48) ∼ µ c (cid:48) such that the distance between the velocity ﬁelds (cid:107) x
c − x c (cid:48) (cid:107) is small.
This implies the distributions µ c and µ c (cid:48) may be be close in Wasserstein distance W 1 ( µ c , µ c (cid:48) ) .
The bound from Redko et al.
(2017) applies well to our setting as such:
Theorem B.4 (Redko et al.
(2017), Theorem 2) .
Let λ c = min f ∈F (cid:16) (cid:15) c
( f ) + 1 /K (cid:80)
Kk
[Equation]
There is N = N (dim( X )) such that for n > N , for any hypothesis f , with probability at least 1 − δ ,
[Equation]
Proof.
We apply Redko et al.
Theorem 2 to target domain µ T = µ c and joint source domain µ S = 1 /K
(cid:80)
Kk
=1 µ c k with empirical samples ˆ µ T =
ˆ µ c and ˆ µ S = 1 /K (cid:80)
Kk
=
ˆ µ c k .
B. E NCODER VERSUS P REDICTION N ETWORK E RROR
Our goal is to learn a joint hypothesis h over the entire domain X in two steps, ﬁrst inferring the task c and then inferring x t +1 conditioned on c .
Error from DyAd may result from either the encoder g φ or the prediction network f θ .
Our hypothesis space has the form { x
(cid:55)→ f θ ( x, g φ ( x )) } where φ and θ are the weights of the encoder and prediction network respectively.
Let (cid:15)
X be the error over the entire domain X , that is, for all c .
Let (cid:15) enc ( g φ ) =
E x ∼X ( L 1 ( g ( x ) , g φ ( x )) be the encoder error where g : X → C is the ground truth.
We state a result that decomposes the ﬁnal error into that attributable to the encoder and that to the prediction network.
Proposition B.5.
Assume c
(cid:55)→ f θ ( · , c ) is Lipschitz continuous with Lipschitz constant γ uniformly in θ .
Then we bound
[Equation]
where the ﬁrst term is the error due to the encoder incorrectly identifying the task and the second term is the error due the prediction network alone.
The hypothesis in the second term consists of the prediction network combined with the ground truth task label x
(cid:55)→ f θ ( x, g ( x )) .
Proof.
By the triangle inequality and linearity of expectation,
[Equation]
By Lipschitz continuity,
[Equation]
which, since g ( x ) = c and by linearity of expectation,
[Equation]
and by deﬁnition of (cid:15) enc and (cid:15) c ,
[Equation]
as desired.
By combining Theorem B.1, Proposition B.5, and Theorem B.4, we can bound the generalization error in terms of the empirical error of the prediction network on the source domains, the Wasserstein distance between the source and target domains, and the empirical error of the encoder.
Let G =
{ g φ : X → C} be the task encoder hypothesis space.
Denote the empirical risk of the encoder g φ with respect to X by ˆ (cid:15) enc ( g φ ) .
Proposition B.6.
Assuming the hypotheses of Theorem B.1, Proposition B.5, and Theorem B.4,
[Equation]
Proof.
We start with the bound of Proposition B.5,
[Equation]
By Theorem B.1 or Mohri et al.
(2018), Theorem 3.3, we can bound
[Equation]
In order to apply Theorem B.1 to the risk (cid:15) c and relate it to the empirical risk
[Equation]
we need to ﬁrst relate the error on the target domain back to the source domain of our empirical samples.
By Theorem B.4,
[Equation]
Applying Theorem B [Equation]
[Equation]
this is
[Equation]
Substituting equation 10 and equation 12 into equation 9 gives
[Equation]
Finally using linearity of the expectation over c ∼ C , removing it where there is no dependence on c , and rearranging terms gives the result.