Inference from Real-World Sparse Measurements
Authors: Arnaud Pannatier
Kyle Matoba
François Fleuret
Venue: Accepted by TMLR ()

Abstract: Real-world problems often involve complex and unstructured sets of measurements, which occurs when sensors are sparsely placed in either space or time. Being able to model this irregular spatiotemporal data and extract meaningful forecasts is crucial. Deep learning architectures capable of processing sets of measurements with positions varying from set to set, and extracting readouts anywhere are methodologically difficult. Current state-of-the-art models are graph neural networks and require domain-specific knowledge for proper setup.

  We propose an attention-based model focused on robustness and practical applicability, with two key design contributions. First, we adopt a ViT-like transformer that takes both context points and read-out positions as inputs, eliminating the need for an encoder-decoder structure. Second, we use a unified method for encoding both context and read-out positions. This approach is intentionally straightforward and integrates well with other systems. Compared to existing approaches, our model is simpler, requires less specialized knowledge, and does not suffer from a problematic bottleneck effect, all of which contribute to superior performance.

  We conduct in-depth ablation studies that characterize this problematic bottleneck in the latent representations of alternative models that inhibit information utilization and impede training efficiency. We also perform experiments across various problem domains, including high-altitude wind nowcasting, two-day weather forecasting, fluid dynamics, and heat diffusion. Our attention-based model consistently outperforms state-of-the-art models in handling irregularly sampled data. Notably, our model reduces the root mean square error (RMSE) for wind nowcasting from 9.24 to 7.98 and for heat diffusion tasks from 0.126 to 0.084.

Published in Transactions on Machine Learning Research (04/2024)
Inference from Real-World Sparse Measurements
Arnaud Pannatier
arnaud. wmaaud..anmatiir@@ddoaoo
Idiap Research Institute, Martigny, Switzerland
École Polytechnique Fédérale de Lausonne, Lausonne, Swstzerland
Kyle Matoba
Idiap Research Institute, Martigny, Switzerland
Ecole Polytechnique Fédérale de Lausanne, Lausonne, Switzerland
François Fleuret
Université de Genève, Geneva, Switzerland
Reviewed on OpenReview: https: // opwnrevee.......... net/ por/formm? 1d=g910f0DRhr
Abstract
Real-world problems often involve complex and unstructured sets of measurements, which occur when sensors sparsely placed in either space time. Being able model this irregular spatiotemporal data and extract meaningful forecasts is crucial. Deep learning architectures capable of processing sets of mcasurements with positions varying from set to set, and ext racting readouts anywhere are methodologically difficult. Current state-of-the-art models are graph neural networks and require domain-spedte knowledge for proper setup,
We propose an attention-based model focused on robustness and practical applicability, with two key design contributions. First, we adopt a ViT-like transformer that takes both context points and read-out positions as inputs, eliminating the need for an encoder-decoder structure. Second, we use a unified method for encoding both context and read-out positions. This approach is intentionally straightforward and integrates well with other systems. Compared to existing approaches, our model is simpler, requires less specialized knowledge, and does not suffer from problematic bottleneck effect, all of which contribute to superior performance.
We conduct in-depth ablation studies that characterize this problematic bottlenock in the latent representations of alternative models that inhibit information utilization and impede training efficiency. We also perform experiments across various problem domains, including high-altitude wind nowcasting, two-day weather forecasting. fluid dynamics, and heat diffusion. Our attention-based model consistently outperforms statc-of-the-art models handling irregularly sampled data. Notably, our model reduces the rool mean square error (RMSE) for wind nowcasting from 9.24 7.98 and for heat diffusion tasks from 0.126 to 0.084.
1 Introduction
x
Deep learning (DL) has emerged as powerful tool for modeling dynamical systems by leveraging vast amounts of data available in ways that traditional solvers cannot. This has led to a growing reliance on DL models in weather forecasting, with state-of- the-art results precipitation nowcasting (Suman et al. 2002 Shi al. 2017) and performance on par with traditional partial differential equation (PDE) solvers in medium-term forecasting (Lam et al. 2022
food-forward
self-at eeff-attention
X.Ye X.Ye X.Ye X.Ye
X1 Xr
Figure Multi-layer Aelfnnittntoon

---

Published in Transactions on Machine Learning Research (04/2024)
Domain Knowledge Agnostic
Architecture
Performance
Simplicity
CNP
x
V
GEN
x
X
TFS
VV
MSA (Ours)
VV
VV
Table 1: Comparison between our approach and the different basclines. Multi-Layer Self- Self-Attention (MSA) achieves good performance while being simple to implement and does not require practitioner knowledge for proper setup.
However, these applications are currently limited to data represented as images or on regular grids, where models such convolutional nerwsrks graph neural nelworks used. In contrast, various real-world data often come from irregularly placed or moving sensors, which means custom architectures needed to handle it effectively.
An example that can benefit significantly from such an architecture is Air Traffic Control (ATC). ATC needs reliable weather forecasts to manage airspace efficiently. This is particularly true for wind conditions, as planes are highly sensitive to wind and deviations from the initial flight plan can be costly and pose safety hazards. DL models are a promising candidate for producing reliable wind forecasts as large amount of data is collected from airplanes that broadcast wind speed measurements with four-second frequency. A model that can effectively model wind speeds using data collected from airplanes, should be able to decode anywhere in space, as aim lo predict wind conditions at future locations of the airplane, conditioned on past measurements taken by that specific airplane or neighboring ones in a permutation invariant manner.
To meet these requirements, we introduce a Multi-layer Self-Attention model (MSA) and compare it to different baselines Table 1}: Conditional Neural Processes (CNP) (Garnelo et al. 2018). Graph Flement Networks (GEN) (Alet et al.) 2019) and transformer encoder-decoder baseline (TFS) that we developed. While all of these models possess the aforementioned characteristics, they cach adopt distinct strategies for representing measurements within the latent space. CNP models use single vector summary of encoded measures, while GEN models map the context to graph, based on their distance to its nodes. MSA keeps one latent vector per encoded measurement and can access them directly for forecasting, This latent representation is better, as it does not create a bottleneck. We show that due to that architectural choice, both baselines can fail, certain cases, to use information from the context they condition on.
Our approach offers better performance than its competitors and is conceptually simpler it does not require encoder-decoder structure. To evaluate the effectiveness of our approach, we conducted experiments high-altitude wind nowcasting, heat diffusion, fluid dynamics, and two-day weather forecasting. Several additional ablation studies show the impact of different architectural choices.
The main contributions of this work are summarized below:
We develop an attention-based model that can generate prediction anywhere in the space conditioned on set of measurements.
We propose a novel encoding scheme using a shared MLP encoder to map context and target positions, improving forecasting performance and enhancing the model's understanding of spatial patterns.
We evaluate our method on set of challenging tasks with data irregularly sampled in space: high-altitude wind nowcasting, two-day weather forecasting, heat diffusion, and fluid dynamics.
We examine the differences between models, and explain the impact of design choices such latent representation bottlenecks on the final performance of the trained models.
2

---

Published in Transactions on Machine Learning Research (04/2024)
2
Related Works
DL performance for weather forecasting has improved in recent years, with DL models increasingly matching surpassing the performance of traditional PDF-based systems. Initially applied to precipitation nowcasting based 2D radar images (Suman et al. 2021 Shi et al. traditional methods for longer forecast periods (Lam et al. 2022 These methods usually work with highly structured data. Radar precipitation data, for example, can be organized as images and analyzed using convolutional neural networks. For 3D regular spherical grid data, graph neural networks or spherical CNNs are employed (Lam et al.; 2022; Esteves ct al.) 2023 However, in our study, the data set is distributed sparsely in space, which hinders the use of these traditional architeclures. The use of DL for modeling dynamical systems, in general, has also seen recent advancements (Li et al.) 2021 Gupta & Brandstetter 2022] Pfaff et al. 2020 but most approaches in this field typically operate regnlarly-spaced data or irregular but lixed mesh.
2017). DL-based models have recently surpassed
Neural Processes GGarnelo et al. 2018 Kim et al. 2019), Graph Element Networks (Alet et al. 2019 and attention-bant models (Vaswani et al., 2017] are three DL-based approaches that are capable of modcling sets of data changing from set to set. In this study, we conduct a comparison of these models by selecting a representative architecture from each category. Additionally, attention-based models have been previously adapted for set classification tasks (Lee al. 2019), and here we adapt them to generate forecasts.
Pannatier et al. (022)) use bernel-based method for wind nowcasting based on flight data. This method incorporates distance metric with learned parameters to combine contexts for prediction at any spatial location. However, a notable limitation of this technique is that its forecasts are constrained to the convex hull of the imput measurements, preventing accurate extrapolation. We evaluate the efficacy of our method compared to this approach, along with the distinct outcomes obtained, in Appendix of the supplementary matcrial.
While previous studies have utilized transformers for modeling physical systems (Geneva & Zabaras 2022
time series (Li et al. 2019) or trajectory prediction (Girgis et al. 2022 Nayakanti et al. 2023 Yuan et al. 2021 these applications do not fully capture the specific structure of our particular domain, which involves relating two spatial processes al arbitrary points on a shared domain. Although we model temporal relationships, our approach lacks specialized treatment of time. Therefore, it does not support inherently time-based concepts like heteroskedasticity, time-series imputation, recurrence, or seasonality. Further details distinguishing our approach from other transformer- based applications are elaborated in Appendix C of the supplementary material.
3 Methodology
3.1 Context and Targets
The problem addressed in this paper is the prediction of target values given a context and a prediction target position. Data is in the form of pairs of vectors (x) and (x) where x, and x, are the position and Ye and yt are the measurements (or values), where we use e for context, & for target, x for spatial position and y for the corresponding vector value. The positions lie in the same underlying space X X₁ E X C X but the context, and (arget values do not necessarily. We define the corresponding spaces as y. I R and yr O C Rº, respectively, where X, O are integers that need not be equal. The dataset comprises multiple pairs of context and target sets, each possibly varying in length. We denote the length of the j-th context and target set as N and NÝ. All models take as input a set of context pairs (xx. 12 Nº as target positions, denoted {1x,) N²
l= as well
As an example, to transform data set of wind speed measurements into context and larget pair, we partitioned the data set into one-minute time segments and generated context and target sets with an intervening delay, as depicted in Figure 2, The underlying space, denoted by X, corresponds to 3D Euclidean space, with both I and O representing wind speed measurements in the x, y plane. The models are given set of context points at positions x, of value y. and query positions x, and are tasked to produce a corresponding
3

---

Published in Transactions on Machine Learning Research (04/2024)
value yı conditioned the context. Detailed descriptions of the data set, including illustrations of the different problems, and the respective spaces for other scenarios and the ablation study can be found in Table within the supplementary material.
3.2 Encoding Scheme
We propose in this section a novel encoding scheme for irregularly sampled data. Our approach leverages the fact that both the context measurements and target positions reside within a shared underlying space. To exploit this shared structure, we adopt a unified two-layer MLP encoder Φ for mapping both the context and larget position lo a latent. space representation. Then, we use a second MLP v to encode the context values and add them to the encoded positions when available. This differs from the approach proposed in Garnelo (2019) where both the context position and value concatenated and given to encoder, and the target position is encoded by another. The schemes contrasted as:
et al. (2018); Alet et al.
C. = Vc)
(1)
C. = x(xx) + myy)
(3)
C2 (xx,)
(2)
C, x(x))
(4)
Traditional methods
Proposed scheme
Where 6, 9, & are two hidden-layers MLPs, and e, et E E are context positions and values which are concatenated and encoded, and target position which is encoded respectively.
3.3 Multi-layer Self-Attention (MSA, Ours)
Our proposed model, Multi-layer Gelf-Atttntion (MSA) harnesses the advantages of attention- based models Figure 1. MSA maintains single latent reprosentation per input measurement and target position, which conveys the ability to propagate gradients easily and correct errors in training quickly. MSA can access and combine target position and context measurements at the same time, which forms a flexible latent representation, Our model is similar to transformer-encode as the backbone of ViT @Dsovitskiy et al. 2020 can be written as:
NY (L) (0))) (runsfrree--nooee(((e)]]] (())) ((8))) ((.))))
(5)
(6)
where y is a two hidden-layers ML.P that is used to generate readonts from the latent space. As a transformer- encoder outputs the same number of outputs as inputs, we define its output as (1)) ( cor- responding to the latent representations of the context and target positions respectively. However, only used generate the rcadouts.
(L)) NY
MSA does not. positional encoding for encoding the order of the inputs. This model is permutation equivariant due to the self-attention mechanism and it uses full attention, allowing each target feature to attend to all other targets and context measurements. MSA generates all the output in one pass in non-autoregressive way and the outputs of the model are only the units that correspond to the target positions, which are then used to compute the loss.
3.4 Baselines
Transformer(s) (TFS) We also modily an encoder-decoder transformer (TFS) model (Vaswani et al., 2017) to the task at hand. The motivation behind this was the intuitive appeal of the encoder-decoder stack for this specific problem. TFS in our approach deviates from the standard transformer in a few ways: Firstly, it does not employ causal masking in the decoder and secondly, the model forgoes the use of positional encoding for the sequence positions. It can be written as:
4

---

Published in Transactions on Machine Learning Research (04/2024)
(0,))] No Thanformo--hcodee((((//11 ((1))) No Transforme--Deccdd((((())) ((e)))))))))
(7)
(8)
(00)) N = (((())))
In comparison MSA, TFS uses an encoder-decoder architecture, which adds a layer of complexity. Moreover, it necessitates the propagation of error through two pathways, specifically through a cross-attention mechanism that lacks residual connection to the encoder inputs.
Graph Element Network(s) (GEN) Graph Element Networks (GEN) (Alet et al. 2019 represent an architecture that. leverages a graph 9 as a latent. representation. This graph comprises N nodes situated at positions Xn and connected by E edges. Selecting the nodes' positions and the graph's edges is critical, as these are additional parameters that require careful consideration. The node positions can be made learnable, allowing for their optimization during training through gradient descent. It's important note that the learning rate for these learnable parameters should be meticulously chosen ensure convergence, typically being smaller in magnitude compared to other parameters.
In the original work, edges are established using Delaunay triangulation and may be altered during training in response to shifts in node positions. The encoder function transforms measurements into a latent space. These are then aggregated to formulate the initial node values, influenced by their proximity to the nodes. Specifically, measurement at position x, impacts the initial value of node at position e, (with i 1,..., N) based on weighting function r(x;; e;), which assigns greater significance graph's nodes. In the original work, this weighting function is defined r(xe, e,) softmax (-ßl c,l)), where B is learnable parameter. The initial node values undergo processing through L, iterations of message passing. For model readouts, the same weighting function is employed, ensuring each node is weighted according to its distance from the query point. Overall, the model can be described as follows:
context points nearer to the
e, Σ Xn)ee (Xc:0c)EE
Vn E (I....,N}
(10)
Message-Passingg( Passing(((ee))))ss
(11)
(L.))
Σ r(xn, X,
Vi (1...,,
(12)
(x..).eC
(8 (
(13)
Where = ((x,, e Jii is the set of encoded context and their position and == {(x,: In) is the set of the graph nodes and their position.
GEN's inductive bias is that single latent vector summarizes small part of the space. As it includes a distance-basee based encoding and decoding scheme, the only way for the model to learn non-local patterns is through message passing. This model was originally designed with a simple message-passing scheme. But it can easily be extended to a broad family of graph networks by using different message-passing schemes, including ones with attention. We present some related experiments in Appendix K] the supplementary material.
Conditional Neural (rocess(cs) (CNP) CNP Garnelo et al. 2018 encodes the whole context as a single latent vector. They can be seen as a subset of GEN. Specifically, a CNP is GEN with graph with a single node and no message passing. While CNP possesses the desirable property of being able to model any permutation-invariant function (Zaheer et al. 2017), their expressive capability is constrained by the single node architecture (Kim al. 2019). Despite this, CNP serves as valuable baseline and is considerably less computationally intensive.
5

---

Published in Transactions on Machine Learning Research (04/2024)
Context
Dalay
Targets
13
MSA
12
Space
CNP
GKA
GEN
10
9
Time
10'
20'
1h 2h 4h
6h
Figure 2: Description of the context and target sets in the wind nowcasting case. The context set and the target set are time slices separated by a delay, which corresponds
Figure 3: Decay of precision in the wind nowcasting case. RMSE of the different models depending on the forecast duration (lower is better). We ran three experiments varying the pseudorandom number gen- erator seeds for each time window and each model to measure the standard deviation. The error does not increase drastically over the first two hours because of the persistence of the wind and the context values are good predictors of the targets in that regime.
the forecasting window. The underlying space is in that case X C R and the context values and target values both represent wind speed and belong to the same space I = 0 C R
1. (/ (0))) No {n (I. (e:)) 1 (6))) N: = =b((0))
(14)
(15)
(16)
Where 7j is simple linear projection.
4
Experiments
Our experiments aim benchmark the performance of our models on various data sets with irregularly sampled data. The first task focuses high-altitude wind nowcasting. The second task is on heat diffusion. Additionally, evaluate Darcy Flow equation and dynamic case modeling the Navier-Stokes equation in irregularly spaced setting. Finally, we compare the models on a weather forecasting task, utilizing irregularly sampled measurements from the ERA5 data set (eesbaach et al. 2023) to predict wind conditions two days ahead.
models fluid flows, considering both steady-state case governed by the
consists of wind speed measurements collected by airplanes with a sampling frequency of four seconds. We evaluate our models on this data set Table 2) and assess the models' performance in Figure 3 and against different metrics Table 3 to ensure the robustness of our approach. We select model configurations with approximately 100,000 parameters and run each model using three different random seeds in both cases. Our results indicate that atlention-based models consistently outperform other models for most forecast durations, except for in the 6-hour range. Notably, we found that the Gaussian Kernel Averaging (GKA) model used in previous work (Pannaticr al.] 2021) achieves satisfactory performance, despite its theoretical limitations, which we analyze Appendix of the supplementary material. Moreover, our findings suggest that attention-based models, particularly MSA and TFS, exhibit. superior performance
For the Wind Nowcasting Experiment, the data set,, described in Section 3.1
function of forecast duration, depicted

---

Published in Transactions on Machine Learning Research (04/2024)
Table 2: Validation RMSE of the High-Altitude Wind Nowcasting, Poisson, Navier-Stokes, and Darcy Flow equation and the weather forecasting task. Each model ran for 10, 2000, 1000, 100, and 100 epochs respectively on an NVIDIA GeForce GTX 1080 Ti. The low number of epochs for wind nowcasting is due to the amount of data which is considerably larger than in the other experiments. The standard deviation is computed over 3 runs. We present here the original implementation of CNP and GEN compared with TFS and MSA with sharing weights for the position. More details can be found in Table 2 of the supplementary material. We choose the configuration of the models so that every model has a comparable number of parameters. We underline the best models for each size and indicate in bold the best model overall.
# of Params.
Wind Nowcasting
Poisson Equation
Navier-Stokes Equation
Darcy Flow Equation
Architecture
ERA5
5k
11.91 0.78
0.33# 0.004
0.701 + 0.0023
0.0311+
0.0008
2. 2.129+ 0.0089
CNP
20k
10.19+ 1.83
0.32+ 0.003
0.672+ 0.0011
0.0295+ 0.0002
2.117+ 0.0018
100k
10.17# 1:24
0.33: 003
0.656 0.0007
0.0286
0 0.0011
2.110+ 0.0002
5k
11.02+ 5.19
0.12+ 0.008
0.604+ 0.0010
0,0304+
0.0003
2.132+ 0.0035
GEN
20k
9.98. 0.76
0.132 0 0.014
0.599 0.0008
0.0296: 0.0002
2.1244 0.0031
100k
9.561
0.162
0.5961 0.0005
0.0294
2.121 0.0005
5k
8.30 0.03
0.15+
0.604+ 0.0022
0.0275=
2.129 0.0082
TFS (Ours, baseline)
20k
8.20 0.04
0.09 o 0o0
0.5961 0.0008
0.0255&
0.0003
2.1094 0.0012
100k
8.38 0.13
0.18: 0.014
0.591 + 0.0012
0.0269: 0.0004
2.100 0.0011
5k
8.07 0.11
0.11 o 006
0.597 0.0011
0.0274
0.0011
2.125 0.0070
MSA (Ours)
20k
7.98 0.03
0.08 0.00a
0.689 A 0.0013
0.0259
0007 0.0264
2.707 0.0020
100k
8.18 0.34
0.10%
0.5891
2.098 0.0029
Table 3: Evaluation of the wind nowcasting task according to standard weather metrics, which are described in Appendix H of the supplementary material. The optimal value of the metric is indicated in the parenthesis. MSA is the best model overall, with the lowest absolute error, a near-zero systematical bias, and output values that have a similar dispersion to GEN. We added the results of the GKA model for comparison as it is the best-performing model in the original paper.
Relative BIAS, (0.0)
Relative BIAS, (0.0)
Model
RMSE (1)
0 MAE (1)
MAE (1)
rSTD (1.0)
NSE (t)
CNP
10.99: 0.75
25.55 1:22
9.224 0.33
0.004 0.08
-1.092 0.03
1.25: 0.07
-0.234 0.01
GEN
8.97
22.561
6.971
-0.021
-0.971
1.09
0.251
GKA
8.44# 0.01
21.894 0.02
6.65 0.02
-0.02+ 0.00
-1.78 0.02
1.13# 0.00
0.31 0.01
TFS (Ours)
7.99+ 0.15
22.17+ 1.20
6.48+ 0.50
-2.21+ 2.67
1.17+ 0.04
0.43+
MSA (Ours)
7.36 0.06
20.48 0.48
5.674 0.11
0.00 0.02
-0.04 0.64
1.09: 0.02
0.55: 0.0s
in this sctup. We also observe that the GKA model performs well for short time horizons when most of the information the context is still up-to-date. However, as the time horizon increases, the GKA model's lack of flexibility is more apparent, and GEN is more competitive.
For the Heat Diffusion Experiment, we utilize the data set introduced in (Alet et al. 2019), derived from Poisson Equation solver. The data set consists of context measurements in the unit square corresponding to sink or source points, as well as points on the boundaries. The targets correspond to irregularly sampled heat measurements in the unit cube. Our approach offers significant performance improvements, reducing the root mean square error (RMSE) from 0.12 to 0.08, (MSE reduction of 0.016 to 0.007, in terms of the original metric) as measured against the ground truth Table 2
For the Fluid Flow Experiment, both data sets are derived from Lii al.) 2021 subsampled irregularly in space. In both cases, MSA outperforms the alternative Table 2. In the Darey Flow equation, the TFS model with 20k parameters exhibits the best performance, but this task proved to be relatively casier, and we hypothesize that the MSA model could not fully exploit this specific setup. However, it is worth mentioning that the performance of the MSA model was within a standard deviation of the TFS model.
We conducted a Two-Day Weather Forecasting Experiment utilizing ERA5 data set measurements. The data set consists of irregularly sampled measurements of seven quantities, including wind speed at different altitudes, heat, and cloud cover. Our goal is to predict wind conditions at 100 meters two days ahead based
7

---

Published in Transactions on Machine Learning Research (04/2024)
these measurements. MSA demonstrates its effectiveness in capturing the temporal and spatial patterns of weather conditions, enabling accurate predictions Table
To summarize, our experiments encompass & range of tasks including high-altitude wind nowcasting, heat diffusion, fluid modeling, and two-day weather forecasting. Across these diverse tasks and data sets, our proposed model consistently outperforms baseline models, showcasing their efficacy in capturing complex temporal and spatial patterns.
5 Understanding Failure Modes
We examine the limitations of CNP and GEN latent representation for encoding a context. Specifically, we focus on the bottleneck effect that arises in CNP from using a single vector encode the entire context, resulting in an underfitting problem Garnelo al. 2018 Kim al. 2019), and that applies similarly to GEN. To highlight this issue, propose three simple experiments. (1) We show in which case baselines are not able to pass information in the context that they use for conditioning, and why MSA and TFS are not suffering from this problem. (2) We show that maintaining independent latent representation helped to the correct attribution of perturbations. (3) We show that this improved latent representation leads to better error correction.
Error
MSA
0
64
TFS
0
0
MSA Egt
0
64
TFS -
GNG
Bottleneck
GEN
64
GEN
PER
0.01
0.01
0.01
0.05
49
GEN 6x6
36
GEN
GEN 5x5
25
NP
0.18
0.22
0.22
0.80
16
GEN 3x3
0.1
1.000 0
10
100
fo 1000
Random
GEN 2x2 +
f
Figure 4: The results of the information retrieval experiment, evaluated using MSE, are considered satisfactory if the MSE is below 0.01. The first three depict models without bottlenecks. The x-axis represents datasets organized by increasing frequency, with 'Random's as the extreme case where the context value is independent of its position. Mod- els with bottlemecks are sufficient when the learned function varies minimally space. Models in italics denote hybrid architectures: GNG represents 'GEN No Graph', maintaining a latent. measure per context, and PER indicates transformer with perceiver layer (Jaegl el al.] 2021), introducing a bottleneck.
Figure 5: Gradients on the last layer the encoder corresponding to an artificial error of 10.0 added to the second output. MSA maintains independent latent representation and gradients exclusively non-zero for the latent associated with the error. We compare it to different GEN models each initialized with graph corresponding to a regular grid of size i x with E (1,...,8, 8}. Due to the bottleneck effect, the gradients corresponding to one error are propagated across different latent vectors for GEN. Evenn when there enough latents (GEN 8 x 8), GEN still disperse attribution because their distance- based conditioning that does not allow for one-to- one mapping between targets and latents.

---

Published in Transactions on Machine Learning Research (04/2024)
5.1 Capacity of Passing Information from Context to Targets
Every model considered in this work encodes context information differently. Ideally, each should be capable of using every measure in their context efficiently. We will see that excessive bottlenecking in the latent space can make this difficult or impossible.
To demonstrate this result, we design a simple experiment in which each model encodes a set of 64 measures (xe: Ye), and is then tasked with duplicating the corresponding y: = Y. given the x; = X,- The training and validation set have respectively 10000 and 1 1000 pairs of sets of 64 examples. It is worth noting that the models have access to all the information they need to solve the task with near-zero MSE. We conducted several experiments, starting by randomly sampling 2D context positions x (x, y) from a Gaussian distribution and computing the associated deterministic smooth function:
Ye sin(Tfx) co)( E R
(17)
where f is frequency parameter that governs problem difficulty. The higher f is, the more difficult the function becomes, as local information becomes less informative about the output. We also consider as a harder problem to sample Ye randomly and independently from the position.
The results of this experiment, as shown in Figure indicate that the CNP and GEN models are less effective in learning this task at higher frequencies. This inefficiency is primarily due to phenomenon we define as a 'bottleneck': situation where single latent variable is responsible for representing two distinct context measurements. This bottleneck impedes the models' ability to distinguish and retrieve the correct target value. In contrast, models with independent latent representations, like MSA, are not subject to this limitation and thus retonstrate superior performance in learning the task. To ensure that the effect arises from a bottleneck in the architectures, we created two hybrid models, trying to stay as close as possible to the original models but removing the bottleneck from GEN and adding one for transformer-basee GNG (for GEN No Graph), which adapts GEN and instead of relying on a common graph, creates one based on the measure position with one node per measure. Edges are artificially added between neighboring measures which serve as the base structure for L steps of message-pasiin.. This latent representation is computationally expensive as it requires the creation of graph per set of measurements, but it does not create a bottleneck in the latent. representation. We found that GNG is indeed able to learn the task at hand. We then followed the reverse approach and artificially added a bottleneck in the latent representation of allention-based models by using Perceiver Layer (Jaegle et al., 2021] with P learned latent vectors instead of the standard self- attentino in the transformer encoder (and call the resultant model PER). When is smaller than the number of context measurements, it creates bottleneck and PER does not succeed in learning the task. If the underlying space is smooth enough, GEN, CNP, and PER are capable of reaching perfect accuracy on this task as they can rely on neighboring values to retrieve the correct information. This experiment demonstrates that MSA and TFS can use their independent latent representations to efficiently retriove context information regardless of the level of discontinuity in the underlying space, while models with bottlenecks, such as CNP and GEN, are limited in this regard and perform better when the underlying space smooth.
models. We call the first one
5.2 Improving Error Correction
In the following analysis, we explore how independent. latent representations can enhance error correction during training.
To do so, we conduct an ablation study focusing on the hypothesis that maintaining independent representations is beneficial during training. We test this hypothesis by restricting the ablation to straightforward scenario, where the models have perfect accuracy on all their outputs but onc. To do so start by pretraining the models to zero error validation set in smooth case - 110). We then pass a fixed random validation batch to the different models and consider this output as the target for the experiment. We then deliberately add a perturbation c 10.0 to one of the target values and backpropagate the error through the model, observing how the models adapt to this perturbation and the consequent impact on their latent variables.
9

---

Published in Transactions on Machine Learning Research (04/2024)
In Figure 5) present the norm of gradients at the last encoder layer of the latent variables for differ- ent models, specifically MSA, TFS, and GEN, when correcting an error during training. It shows how MSA and TFS models experience a non-zero gradient only a single latent variable, whereas the GEN model exhibits changes across multiple latent vari- ables. This disparity suggests that models like MSA and TFS, which maintain one latent variable per output, are less constrained during training. They need only to adjust the specific latent corresponding to the error without affecting others. Conversely, models with entangled representations, as scen in GEN, must manage the dual task of correcting the error while ensuring other dependent outputs remain unchanged.
Increase in # of Gradient Updates
150%
100%
50%
0%
MSA
TFS
GEN 8x8
GEN 7x7
GEN 6x6
GEN 5x5
GEN 4x4
GEN 3x3
GEN 2x2
GEN 1x1
Figure 6: Comparison of the number of gradient up- dates required to correct an artificial error, with respect to MSA (lower is better). The y-axis represents the increase in percentage in the number of steps required to reach a perfect accuracy with respect to MSA. We compared MSA to different GEN each initialized with graph corresponding to a regular grid of size i x i with E {1. It can be observed that all GEN take more steps to correct the same mistake, and the more entangled the latent reprosentation is, the more time it requires correct the problem.
We think that the additional constraints might lead to slower training and to test. that hypothesis, we measure the number of backpropagation passes re- quired return zero error on the validation set Figure 6.. We see that the number of latent vari- ables reduces, it leads to increase in the number of steps required to correct the error. This result suggests that the entanglement of latent variables in GEN can lead to slower training and that maintaining independent latent representations can improve error correction. We make the hypothesis that this scenario extends to general training dynamics, proposing that models with independent latent representations can correct errors more efficiently due to their reduced constraints.
5.3 Encoding scheme
In this section, we evaluate the novel encoding scheme presented in /Section 3.2) and we present the results in Table 4. We found that it reduces the RMSE from 8.47 to 7.98 in the wind nowcasting task and enables the MSA model to achieve the best performance with an RMSE of 0.08 for the Poisson Equation. Sharing the same mapping for positions is the appropriate inductive bias for encoding positions, it eliminates the need to learn the same transformation twice. Since our data is irregularly sampled in space, the positioning of measurements and target positions significantly influence the prediction, as demonstrated in additional experiments, Appendices I and J in the supplementary material. We think that sharing the position mapping can link information from the context and target positions, which helps the model to understand better how the space is shaped.
6 Conclusion
In this work, we introduced an attention- based model to handle the challenges of wind nowcasting data. We demonstrated that the proposed attention-based model was able to reach the best performance for high-altitude wind prediction and other dynamical systems, such as weather forccasting. heat diffusion, and fluid dynamics when working with data irregularly sampled in space. We then explained why atlention-based models were capable of outperforming other models on that task and provided an in-depth examination of the differences between models, providing explanations for the impact of design choices such as latent representation bottlenecks on the final performance of the trained models.
Our work builds upon well-established attention models, which have demonstrated their versatility and efficacy in various domains. Although the core model esscntially a vanilla transformer, our architccture required careful adaptation to suit our specific requirements. We designed our model to be set-to-set rather than sequence-to-sequeece, handling data in non-causal and non-autoorerressvve manner, and generating continuous
10

---

Published in Transactions on Machine Learning Research (04/2024)
Table 4: Relative performance improvement with shared encoding for position across various models. This improvement is calculated as Errorse sbared models without shared encoding and Errorshared is the mean error across sizes for models utilizing shared encoding. Detailed values are available in Table 5 Scores exceeding 100% indicate superior performance without shared encoding. Generally, shared encoding enhances model efficiency, typically by small margin, but significantly in some cases like the Poisson Equation. The shared encoding approach is likely advantageous as an inductive bias for position encoding. By naturally treating the positions of context and targets similarly, it can effectively use this information to develop more robust representations.
Erroraharsd where Error shared represents the mean error across sizes for
Wind Nowcasting
Poisson Equation
Navier-Stokes Equation
Darcy Flow Equation
Architecture
ERA5
CNP GEN
100.7%
142.1%
154.1%
193.0%
100.6%
95.2%
63.5%
101.9%
77.9%
101.3%
TFS (Ours, baseline)
95.3%
79.5%
99.5%
99.3%
101.0%
MSA (Ours)
93.7%
24.8%
99.4%
97.2%
100.9%
values for regression. The success of influential models like BERT (Devlin et al. 2019). GPT (Radford let al., 2018), VIT (Dosovitskiy al. 2020)) and Whisper (Radford al. 20022), also closely resembles the original implementation by Vaswani et al. (2017). which further supports the effectiveness of the transformer framework across different tasks and domains.
Finally, our model's scalability is currently limited by its quadratic complexity in the context size. Although this limitation does not pose a problem in our particular use cases, it can impede the scaling of applications.
This is a significant challenge that affects all transformer-basee Recent devclopments tackle this challenge inchude flash-attention (Dao al. 2022), efficient transformers (Katharopoulos et al.) 2020), and quantization techniques (Deltmers al.j 2022), which can address this problem, enhancing the feasibility of our approach for large-scale applications.
models and has garnered considerable attention.
Acknowledgement Arnaud Pannatier supported by SkySoft ATM for the project "MALAT: Machine Learning for Air Traffic" and the Swiss Innovation Agency Innosuisse under grant number 32432.1 IP-ICT. Kyle Matoba was supported by the Swiss National Science Foundation under grant number FNS-188758 "CORTT".
References
Ferran Alet, Adarsh Keshav Jeewajee, Maria Bauza Villalonga, Alberto Rodriguez, Tomas Lozano-Perew and Leslic Kaelbling. Graph element networks: adaptive, structured computation and memory. In International Conference on Machine Learning, pp. 212 222. PMLR, 2019.
Tri Dao, Daniel Y Fu, Stefano Ermon, Atri Rudra, and Christopher Re Flashattention: Fast and memory- efficient. exact attention with TO-awareness. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https:// openreview openreview.net/forummn wwt.://eooormccccc/aff555555
Tim Dettmers, Mike Lewis, Youncs Belkada, and Luke Zettlemoyer. GPT3.int8() 8-bit matrix multiplication for transformers at scale. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https:///pennrviie..nn http:://ppparreeeeeeeeeeerrrrrr fittrrrunddt
aXiGWqBoxaD
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional
transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chopter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers). pp. 4171-4186. Association for Computational Linguistics, 2019. doi: 10. 1.118863/v 1111111///1///// 1423.. URL htttpp::/badi.rrg/101166666
11

---

Published in Transactions on Machine Learning Research (04/2024)
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.1192: 2020.
Carlos Esteves, Jean-Jacques E. Slotine, and Ameesh Makadia. Scaling spherical cnns. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), International Conference on Machine Learning, ICML 2023, 28-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. //proceedings.. iacinnssmml.. mlr. Srtess/v222/esseevee
PMLR, 2023. URL https:
html
Matthias Fey and Jan E. Lenssen Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Leurning on Graphs and Manifolds, 2019.
Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Rezende, and S. M. Ali Eslami. Conditional neural processes. In Jennifer Dy and Andreas Krause (cds.). Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1704 1713. PMLR, 10 15 Jul 2018. URL htttss:/proceeddings wwtiang..m.....................e ptwps//800gaareelllooo
Nicholas Geneva and Nicholas Zabaras. Transformers for modeling physical systems. Neurol Networks, 146: 272 289, 2022.
G. Ghiggi, V. Humphrey, S. T. Seneviratne, and L.. Gudmundsson. Grun: an observation-based global gridded runoff dataset from 1902 to 2014. Earth System Science Data, 11(4):1655 1674, 2019. doi: 10.5194/cssd-11 11-15---1-655-0019 URL httpss:/essd oopprn/css.org/aattte//////////2
Roger Girgis, Florian Golemo, Felipe Codevilla, Martin Weiss, Jim Aldon D'Souza, Samira Ebrahimi Kahou, Felix Heide, and Christopher Pal. Latent variable sequential set transformers for joint multi- agent motion prediction. In International Conference on Learning Representations, 2022. URL https :
//openreview.net// wwwp:eeefooooumdddbuuuee 2D0
Jayesh K. Gupta and Johannes Brandstetter. Towards milti-ppatiotemppoalssaaal generalized pde modeling, 2022. URL https://aaxii..rrrtt wwtp://wrwr.abb//.....e. 15616
Hans Hersbach, Barbara Bell, Paul Berrisford, Giovanni Biavati, András Horányi, Joaquín Muñoz Sabater, Julien Nicolas, Carole Peubey, Razvan Radu, Ivan Rozum, Dries Schepers, Adrian Simmons, Cornel Soci, Dick Dec, and Jcan-Noël Thépaut. ERA5 hourly data on single levels from 1940 to present. Copernicus Climate Change Service (C3S) Climate Data Store (CDS), 2023. Accessed on 17-05-2023.
Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Tonescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, et al. Perceiver io: A general architecture for structured inputs & outputs. arXiv preprint orXic:2107.14799 2021.
Steeven Janny, Aurélien Bénéteau, Madiha Nadri, Julie Digne, Nicolas THOME, and Christian Wolf. EAGLE: Large-scale learning of turbulent fluid dynamics with mesh transformers. In The Eleventh International Conference on Learning Representations, 2023. URL http:: httpps:/opeerevieet net tet/foruumi1dmmf114Ap..
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In Hal Daumé III and Aarti Singh (eds.), Proceed- ings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Reseorch, pp. 5156-5165. PMLR, 13-18 Jul 2020. URL URL http::://ppooeeddgngeeeee//e/eee mmr. ptt.s//ess//v1// katharopoulos20a..tnll
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processcs. In International Conference on Learning Representations, 2019. URI. https::/openreevie http:p///pppaeeeeeeeeeee//////ss wwt..e/ooruum//ieesssssppp.CCCC
12

---

Published in Transactions on Machine Learning Research (04/2024)
Remi Lam, Alvaro Sanchez-Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Alexander Pritzel, Suman Ravuri, Timo Ewalds, Ferran Alet, Zach Faton-Rosen, Weihua Hu, Alexander Merose, Stephan Hoyer, George Holland, Jacklynn Stott, Oriol Vinyals, Shakir Mohamed, and Peter Battaglia. Graphcast: Learning skillful medium-range global weather forocasting, 2022. URL httpss///axiin /trg//orv/voo//////////
abs/2212.12790
Juho Lee, Yoonho Lee, Junglaek Kim, Adam R Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer. In International Conference on Machine Learning, 2019.
Shiyang Li, Xiaoyong Jin, Yao Xuan, Xiyou Zhou, Wenhu Chen, Yu-Xiang Wang, and Xifeng Yan. Enhancing the locality and breaking the memory bottleneck of transformer on time series forocasting. Advances in neural information processing systems, 32, 2019.
Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Represenlations, 2021. URL.
net/fforumiid-<eP9NNVYY.
Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar. Fourier neural operator with learned deformations for pdes on general goometrics, 2022a.
Zongyi Li, Daniel Zhengyu Huang, Burigede Liu, and Anima Anandkumar. Fourier neural operator with learned deformations for pdes on general geometries, 2022b.
Nigamaa Nayakanti, Rami Al-Rfou, Aurick Zhou, Kratarth Goel, Khaled S. Refaat, and Benjamin Sapp. Wayformer: Motion forecasting via simple & elficient attention networks. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 1980-2987, 2023. doi: 10.1109 1LLL009/ICA8899...00 2.223
10160609.
Arnaud Pannatier, Ricardo Picatoste, and François Fleuret. Efficient wind speed nowcasting with gpu- accelerated nearest neighbors algorithm, 2021. URL http::://arrrii..orr/bb/b11111.1 wrtp./////a......//.bbpg 10408.
Tobias Pfaff, Meire Fortunato, Alvaro with graph networks. In International Conference on Learning Representations, 2020.
and Peter Battaglia. Learning mesh-based simulation
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving lan- guage understanding by generative pre-training. Technical report, OpenAI, 2018. Available: https://s3-us-west- amazonaws, cww/ com/openai coom/opena-assetaaa assets/ /ewp://rcrroooeeee//gggguguuuaee
munguage_unersstaddin wawpgngappppp..e pdf Accessed: 16-01-2024.
Alec Radford, Jong Wook Kim, Tao Xu, Greg Brockman, Christine McLeavey, and llya Sutskever. Robust speech recognition via large-scale weak supervision. arXiv preprint arXiv:221204483 04356, 2022.
Xingjian Shi, Zhihan Gao, Leonard Lausen, Hao Wang, Dit-Yan Yeung, Wai-kin Wong, and Wang-chun Woo. Deep Learning for Precipitation Nowcasting: A Benchmark and A New Model. Advances in Newral Information Processing Systems, 2017.
R. Suman, L. Karel, W. Matthew, K. Dmitry, L. Remi, M. Piotr, F. Megan, A. Maria, K Sheleem, M. Sam, P. Rachel, M. Amol, C. Aidan, B. Andrew, S. Karen, H. Raia, R. Niall, C. Ellen, A. Alberto, and M. Shakir. Skilful precipitation nowcasting using deep generative models of radar. Nature 597(7878): 672 677, 09 2021. ISSN 1476-4687. doi: (0.1038//411586-21--0333 341586-021-03854-
URL https:///oii org/10.1033
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Omry Yadan. Hydra- a framework for elegantly configuring complex applications. Github, 2019. URL hctp://cccbbbbreeerar///yyaaaaat
13

---

Published in Transactions on Machine Learning Research (04/2024)
Ye Yuan, Xinshuo Weng, Yanglan Ou, and Kris M. Kitani. Agentformer: Agent-aware transformers for socio-temporal multi-agent forecasting. In Proceedings of the IEEE/CVF International Conference Computer Vision (ICCV), pp. 9813 9823, October 2021.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabás Póezos, Ruslan Salakhutdinov, and Alexan- der J. Smola. Deep sets. CoRR, abs/1708.06114, 2017. URL
A Detailed results
Here, present breakdown of the detailed results for the different models with and without shared positional encoding, and presented in terms of the specific metrics reported in their original study. In (Tale 2 of the main paper, we provide the results for GEN and CNP based on their original implementations. which do not involve sharing the position encoding and we translated all metries to RMSE. For TFS and MSA, we include the results obtained with shared position encoding. To ensure comprehensiveness and to differentiate the enhancod performance attributed the latent reprosentation from that resulting from the improved data encoding approach, we also apply the shared encoding technique to GEN and CNP.
In most cases, observe that employing shared encoding for position tends to improve the performance of all models, except for CNP where it scems to decrease the overall performance in some cases. Additionally, it seems that MSA surpasses its competitors in the majority of instances, consistently exhibiting superior performance across various data sizes. Moreover, when the shared encoding for positions fails to achieve the optimal performance, the model's performance is generally comparable to that without shared encoding, often falling within standard deviation range. We also give in Table 6 detailed breakdown of the results for the training set.
B Dataset descriptions
B.1 Wind nowcasting
The wind nowcasting experiment uses the same dataset as Pannaaier et al. 2021 It is available at: https://zenodo. org/record/5074237 It is an important task for ATC as it involves the crucial prediction of high-allitude winds using real-time data transmitted by airplanes. This prediction is essential for ensuring efficient airspace management the airspace. It is important to note that in the vertical direction (2), the majority of wind measurements are typically taken at elevated altitudes, specifically between 4,000 and 12,000 meters. Regarding the horizontal dimensions, the airspace extends over 600 km from north to south and approximately 500 km from east to west. As ground-based measurements heights, our dependence lies the mcasurements gathered directly from airplanes. As planes do not record the wind in the direction, the inpul space corresponds to wind speed in the x, plane Il C R2, and the output space is the wind speed measured later, O C R²2 Here the underlying space is European airspace represented as X C R33 this particular sctup, the models need to extrapolate in time, based on a set of the last measurements.
not accessible at these
Airplanes measure wind speed with a sampling frequency of four seconds. We split the dataset into time slices, where each slice contains minute of data, such that each time slice contains between 50 and 1500 data points. We tried different time intervals but noticed that having longer time slice did not improve the quality of the forecasts. The objective for the different models is to output prediction of the wind at different query points 30 minutes later. We evaluate performance using the Root Mean Square Error (RMSE) metric, as in previous work.
B.2 Poisson Equation
This experiment uses the same datasel as Alet el al. 2019 It is available at The Pois-
https:://githin wwtp://oeee/eeee/ggeeeaaaa httpr/rgraapppeeeetttteeeeee element networks tete/teeeemmeeee wateee/mmatee/eeee data son equation models the heat diffusions over the unit square with sources represented by o(x,y) E RR and fixed boundary conditions w R. The equation given by:
14

---

Published in Transactions on Machine Learning Research (04/2024)
Table 5: Results of the High-Alti Wind Nowcasting, Poisson, Navier-Stokes and Darey Flow equation and the weather forecasting task. Each model ran for respectively 10, 2000, 1000, 100, and 100 epochs on an NVIDIA GeForce GTX 1080 Ti. The low number of epochs for wind nowcasting due to the amount of data which is considerably larger than in the other experiments. The standard dcviation is computed over 3 runs. We choose the configuration of the models so that every model has a comparable number of parameters. We underlying the best models for each size and we put in bold the best model overall. In this table, we kept the results in the same metric as originally reported. The first half the table corresponds to the case without the novel encoding scheme, as opposed to the second one. last part of the table reports the results of GeoFNO another interesling baseline, which we discussed detail in Appendix D In Table 2 of the main paper, we recomputed all metrics in terms of Root Mean Square Error (RMSE).
Validasion
Nowcasting
Poisson Equation
Navier-Stokes Equation
Darey Equation
(1)
(4)
(1) B-4
Default encoding for positiins
0.21
5k
8.84
2.92
.017 1. .0111
.355 1: .012
9:25
0.18
4.47
0.02
GEN
20k
8.24
0.35
020 I .0044
.3h 1 .007
6:74 1: 0.09
4.44
0.00
100k
8.23
0.44
.048 1 .0399
.353
.005
8.06 1 0. 0.00
4.46
0.00
5.70
0.14
0228
7:46
4.46 0.00
TFS (Ouns)
20k
5.70
0.06
.014
0019
6.09
0:14
1:38 0.01
100%
8.07
0.07
.018 + .0045
.300
.0111
7.47
0.24
4.42 0.01
8.40
0.10
.0054
T-74
(Ours)
9.47
0.12
6.70
100k
9.98
0.29
.000 1: .0081
350 1: 0013
7:33
0.22
4.41 0.01
5k
10.38
0.19
.186 + .0025
.706
0.0003
17.26
0.04
4.63 0.01
CNP
20k
10:12
0.21
.180 + .0011
.706
0.0001
17.06
0.06
4.47
0.00
100k
10.26
0.24
.158 + .0086
.706
0.0001
16.87 #: 0.02
4.43 0.01
9.79 0.28
+ .0023
0.0008
6.88
4.64
GEN
0.47
7.19
5k
6.30 0.03
.025 4 .0121
365 0-0025
7-58 1: D
4.53 0.01
TFS (Ours)
20k
5.20 0.04
.010 + .0013
.356
0.0009
0.04 + 0.14
4.45 0.01
100k
8.36 0.13
.055 t .0044
349
0.0015
7:25 0.19
4.41 0.00
8:07
0.11
t .0044
357
7:51
4.52 0.09
(Ours)
20k
7.98
+ 0.03
.007
.347 + 0.0016
6.74
0.98
4.44 + 0.01
100k
0.11 8.18
1 .0017
0.0008
6.87
0.24
1.10 0.01
Different encoding seleeme
GeoFNO
1.5M
0.40
.0036
0.01
f AΦ(Z, y) (0 if (x, v) E 0
if (x.z) E (0,
Breets
(18)
It should be noted that the boundary constant and sources function ((I u) conditions can change for each sample.
The dataset used in (Alet et al.
2019) uses three dimensions for the context values y. E I R cither sources values ((x, y) = μ inside the domain encoded Xx, = (x, v), (n, 0,0) or boundary conditions d(z,y) on the boundaries, encoded Xe. (x, (0, 1). The target space is one-dimensional
O C IR and corresponds to the solution of the Poisson equation at that point. A sample is depicted in Figure
Figure 7: A sample of the Poisson equation dataset (Alet et al. 2019 represents the targets, The con- text values are comprised of points the boundaries, and points on the sources and sink represented by which corresponds to their ther- mal coefficient.
We evaluate all models for three different sizes on this particular sctup and find that MSA with the novel encoding scheme outperforms other models by significant. margin as be seen in Table 2 We initialized
15

---

Published in Transactions on Machine Learning Research (04/2024)
Table 6: Corresponding training metrics
Tralning Model
Wind Nowcasting TMME (J)
Polsson Equation MSR (1)
Navier-Stokes Equation
Darey Flow Equation
ER A
Sixe
(L)
RMSE: (G) E-t
MSR (I)
Defonit cacodiag
11.94
0.78
.177 + .0179
.88
.0033
9:02
0.64
4.35
0.01
CNP
20k.
10:10
1.88
.084 + .0077
.48
.0015
7:88
0.16
4.26
0.01
10.17
1.23
.91 .0040
.0010
8.83
4.17
0.01
5k
11:02
5.19
011 + .0066
0012
8.49
0.17
4.32
0.02
GEN
20k
9.96
0.78
.000 + .0066
.544
0007
1.76
0.11
4.24 0.01
100k
8:50
0.21
.011 1 00122
.39
0005
T:01
0.36
0.01 4.10
(Ours)
9.56
0.01
.02
.0022
850
0057
6.01
4.26
0.02
MBA (Ours)
20k
7.94
0.08
.000
00022
.341 = 0042
0.72
0.11
4.13
0.01
100k
6.07
0.02
0000
0001
.310 1: 0015
4.01
0.17
4.19
0.01
Sharing encoding
5k
10.41
0.00
.154 + .0010
.703
.0001
16.04
D. 06
4.35
0.00
8:48
0.00
.130 0004
.001
08
1:30
0.01
100k
8.80
0.05
.107 1 .0008
.0001
16:36
1.23
0.01
5k
10.04
0.13
005 + .0044
.007
5.83
4.42
0.01
GEN
20k
9.76
0.08
.003 + .0007
.322
.002
6.65
0.06
4.36
0.01
100k
9.94
0.02
.001 1 .0000
.33
.005
4.07
0.22
4.34
0.02
(Ouns)
.0000
0.01
5k
9.54
0.08
.008 + .0024
.865
.0013
6.74
0.01
4.35
0.08
MSA (Ours)
20k
7.78
0.02
.002 .0083
.442 + .0044
6.70
4.26
0.01
8.58
0.05
0000
.0000
.0005
4:78
0.02
Ditferent encoding selemme
11.68
0.30
.008 4: 00000
325 :: 00070
7:50
4.18
0.02
20k
11.50
0.20
.003 1 00055
.47 1: .1000
7.02
4.07
100k
10:72
0.27
.00
0001
3: .0068
5.87
3:75
1.581
0.98
.0002
2.60
GEN with a 7 x 7 regularly initialized grid on the [0, 1]3 as the original work (Aet et al., 2019
B.3 Navier-Stoker Equation
Similarly to the darcy flow task, we adapted the Navicr-Stokes equation as described in Li et al 2021 | to an
irregular setup.
The Navier-Stokes equation describes a real fluid and is described with the following PDE
O,w(x, t) + u(x,t) Vw(x,t) vAw(a,t) f(z), (0, V vu(xz,)) = 0, (0, 1)2
if (0, if l E [0.
(19)
w)x = w
if E (0, 1)2
For more detailed information regarding the notation, please refer to the original work by Li al. (Li et al. 2021). specifically section 5.3.
In this study, our objective is to forecast the future state of the vorticity quantity, specifically 50 timesteps ahead, based on measurements of vorticity at different spatial locations. This approach differs from the original work, where the model was provided with initial vorticity grids and required predict the complete system evolution.
To create dataset, subsampled the complete dataset, which consisted of the evolution of the Navier- Stokes equation solved by a numerical solver. The full datasct had dimensions of 1000 x 1024 x 1024.
For our purposes, we selected pairs of slices that were separated by 50 timesteps and performed spatial subsampling. We took 64 context measurements and 256 targets as our subsampled data points. This experiment adapts the datasct available in (Li et al. 2021 It is available at /https://githul: com/ neural-operatorn neurzl-operator/fourieie neurallopereaaeee and can be processed with the code given to rearrange it the irregular setup. We present sample of the dataset in Figure 9
16

---

Published in Transactions on Machine Learning Research (04/2024)
50
100
100
200
300
200
50
100
150
200
100
200
300
400
520
100
100
200
300
400
:00
100
150
200
5:0
Figure 8: sample the Darcy Flow dataset Li
Figure 9: A sample of the Navier-Stokes dataset. (Li et al. 20021). Blue dots correspond to the context set and orange dots to the targets.
el al. 2021). Blue dots correspond to the context set and orange dots to the targets.
B.4 Darcy Flow
We evaluated the performance of various models in solving the Darcy Flow equation on the unit grid with mull boundary conditions, as described in (Li et al. 2021). Specifically, we aimed to predict the value the function u, given the diffusion function a, with the two related implicitly through the PDE given by:
- V (a(r)Vu(x)) =1 if x E (0,1)² u(x)=
(20)
if E 000,
The dataset used in our study is adapted from that used in Lii el al. 2021 and originally generated by a traditional high-resolution PDE solver. The dataset consists of a 1024 x 1024 grid, which was subsampled uniformly at random and arranged into context-target pairs. The context is comprised of the evaluations of the diffusion coefficient: (Xe.Ye) = ((I, y), a(x, y)), and the target is the solution of the Darcy Flow equation at position, yx) = ((x, v), ())
The results are presented in Table 2. they are coherent with the rest of the experiment and show that the MSA and TFS models outperform all competing models. We used the same initialization for GEN for the Poisson Equation.
This experiment adapts the dataset available in (Li et al. 2021 It is available at htttp::/ggtthiii hom. heural-operz couralloperatoo/fourii ott.:p/peerrrooooooeereaaaroooro and can be processed with the code provided to rearrange it in the irregular setup.
B.5 Two-Days Weather Forecasting
In this task, we want evaluate our models on the task of two-day weather forecasting based on irregularly sampled data in space. The requested data collection description focuses climate reanalysis from the
17

---

Published in Transactions on Machine Learning Research (04/2024)
Dataset
dim(X)
dim
dim (0)
# Train. points
# Val. points
URL
Pannatier
2021
Wind Nowcasting
3
26.956.857
Datasct page
2019
Poisson Equation
2
3
409.600
102.400
Gitiubb Repostory
Navicr-Stokes Equation
2
1
48, 28..883.0
11.673.600
Githish Repnstery
al.
et al. 2021
Darcy Flow Equation
2
1
409.600
102.400
Githubb Repostory
ERA5
2
7
2
8.6488701
2,107.392
Hersbach et Al 2023
Random
2
1
1
640.000
64.000
Handomly geuerated
Sine
64.000
Randomly generated
Table 7: Description of the different datasets used in this study.
ERA5 datasct, which is available through the Copernicus Climate Data Store (CDS). The dataset contains various parameters related lo wind, surface pressure, lemperature, and cloud cover.
To access the data, please visit the following URL: clamaate..cimmm..... http:///.....m....ooo..oeeeeee/a
We selected the following variables:
10m u-component wind: This refers to the eastward wind component at a height of 10 meters above the surface.
10m v-component of wind: This represents the northward wind component at a height of 10 meters above the surface.
100m u-component of wind: This denotes the eastward wind component at height of 100 meters above the surface.
100m v-component of wind: This signifies the northward wind component height of 100 meters above the surface.
Surface pressure: This represents the atmospheric pressure at the Earth's surface.
2m temperature: This indicates the air temperature at height of 2 meters above the surface.
Total cloud cover: This refers the fraction of the sky covered by clouds. The data collection includes measurements for all times of the day and all days available in the dataset. However, for training purposes, selected the data from the year 2000. For validation, we use the data from the year 2010, specifically the months of January and September.
Next, we proceeded extract the corresponding grib files. To create our dataset, we paired time slices that had time difference of two days, equivalent to 48 timesteps since ERA5 data has onc-hour resolution over the whole world.
For the context slice, we performed subsampling at 1024 different positions and retained all seven variables mentioned carlier. As for the target slico, we subsamplod it at 1024 different positions and kept only the two components of the wind at 100m.
C Specific Aspects of Our Approach Compared to Existing Works
In our research, we address a distinct challenge that sets our work apart from conventional mesh- based simulators (Pfaff et al. 2020 Janny et al. 2023 Li et al. (022a)). Our focus is on extrapolating sets of sparse and variably-positioned measurements, a context markedly different from the consistent mesh structures employed in traditional approaches. Unlike mesh-based models where values accessible at fixed points
18

---

Published in Transactions on Machine Learning Research (04/2024)
an irregular mesh, our method tackles scenarios where the quantity and locations of measurements change dynamically from one data set to another. This unique aspect of our work necessitates departure from the conventional mesh framework. Mesh-based simulators, while effective in their domains, are not equipped to handle the sporadic and non-uniform nature of our data. For instance, in applications such as airspace monitoring, data points represent measurements taken only where planes are present, leaving vast. areas without data. Our methodology, therefore, diverges fundamentally from mesh-based techniques, as it requires innovative handling of sparse and irregular data inputs, rather than relying on a fixed, uniform mesh structure.
Furthermore, our approach markedly differs from trajectory prediction models (Yuan et al. 2021 Nayakanti
et al. 2023 Gingis et al. 2022). While these models excel in forecasting future points along established trajec- tories, our work diverges in both intent and application. Our primary concern is not the sequential prediction of trajectories but rather the interpretation and forecasting of phenomena in a spatially and temporally sparse environment. In the context of wind noweasting, for example, our model excels at processing limited and scattered data points from multiple trajectories to generate comprehensive forecast. This capability to assimilate sparse measurements from varied locations and synthesize them into coherent prediction sets our work apart from traditional trajectory-focusee models. These models, such as Wayformer (Nayakanti et al. 2023) and Agentformer Yuan et al. 2021 are primarily designed for time series forecasting and trajectory completion.
Comparison with GeoFNO
The Fourier Neural Operator (FNO) family (Li et al. 2021 and its variant GeoFNO (Li et al. 2022bl, have emerged as significant tool in resolving Partial Differential Equations (PDEs) across various benchmarks. Even though the GeoFNO model was specifically designed for irregular meshes, setting it apart from other models evaluated, it is recognized for its adept handling of irregular geometries and flexibility in adapting to different types of Partial Differential Equations (PDEs), making it a strong baseline for our experiments.
We evaluated GeoFNO on our specific dataset and problem. We tried both the original model configuration, consisting of 1.5 million parameters, and its scaled-down versions with 5k, 20k, and 100k parameters, found the 1.5 million parameters to be over-parametrized. The outcomes of these tests be found in
Table 5
Despite GeoFNO's overall robustness, we found that it encounters challenges in this particular setup. We hypothesize that the model struggled with sparse irregular data, as needs to learn to map a set of data to a grid before applying a regular FNO layer. Sparsity, in this context, might lead to several parts of the grid being empty thus leading to information loss or dispersion. Secondly, the experiment's unique requirement of handling variable measurement positions, which is different from the usual consistent yet irregular meshing, introduced an additional layer of complexity in learning stable mapping. These factors potentially contributed to the performance inconsistencies observed in our experiments. For instance, in the Poisson sctup, where data points were more evenly distributed, GcoFNO's performance occasionally matched that of the MSA model in some randomized trial runs, as indicated in Figure 12e) However, in more demanding scenarios like wind nowcasting, GeoFNO consistently underperformed in comparison to both MSA and GEN models.
E Attention matrix for Wind Nowcasting
For all models, can plot the norm of the input gradients to see how changing given value would impact the prediction Figure 10). We see that Transformers seem to find tradc-off between taking into account the neighboring nodes and the global context.
F
Comparison with Graph Kernel Averaging
Smart averaging strategies such as GKA are limited by design to the range of the values in their context. However, these baselines are useful they are not prone to overfitting and in the case of time series forecasting
19

---

Published in Transactions on Machine Learning Research (04/2024)
Context, Targets
(b) MSA (Ours)
(c) GKA Pannatier
2021
(d) GEN Alet et al. 2019
(e) CNP Garnelo et al. 2018
(f) TES (Ours)
Figure 10: Displaying the importance given to points in the context to the prediction for the different models for a given query point (orange). We use the norm of the input gradients for that purpose and highlight the context values that have the largest gradient with respect the output. The opacity of the dots corresponds to the relative magnitude of the input gradients compared to other points in the context.
20

---

Published in Transactions on Machine Learning Research (04/2024)
Greel
100%
Centot
25- Turds
75%
50%
-25-
25%
10
0%
Targets
MSA
TFS
GKA
GEN
CNP
1540
-158-
(c) Percentage of the prediction values that are contained in the convex hull (a) Context, targets, and forecast in (b) Context and largets in the 3D of the measurement in the context in wind speed space.
space.
percent.
Context and targets represented in wind speed and 3D space. Both the context and targets correspond to 1-minute slices of data. The targets capture wind measurements recorded 30 minutes following the context data. Additionally, the forecasts generated by the MSA and GKA models depicted. Notably, the GKA model is confined to predicting within the convex hull of the context values, which is visually highlighted in grey. Conversely, the MSA model is unrestricted by this limitation and can make predictions beyond the convex hull of the context values.
Figure 11:
(if the underlying variable has some persistence), the last values would in general be the best predictor of the next value. We hypothesize that over sufficiently short horizons, these methods would achieve performance similar to that from more complicated models, but when the prediction horizon rises, this design limitation would become more pronounced.
Since, for extrinsic forecasting, target points need not lie in the convex hull context points, the greater flexibility of MSA and TFS can make correct forecasts that are impossible with GKA. Figure 11a demonstttss this fact, with plot of the target variable, the context points, and the context convex hull in wind speed space.
But greater capacity means also more failure modes. As we saw in the metrics of Table 3 of the main paper, MSA and TFS are the only models to outperform GKA whereas other models, even if they are more flexible, fail to beat this bascline. On average 27% of the predictions were outside the convex hull of the context.
Figure 11c we analyze the percentage of measurements outside the convex hull produced by the different. models. We can see that GKA is limited 100% of the measurements lie within the convex hull whereas the other models can compensate and predict values outside it.
G Decoder and conditioning function
In the main paper, we presented a comparison the key distinctions between our attention-based architoctures. However, could be that the performance differences observed could be attributed to minor variations in the model architecture design. For the sake of completeness, in this section, provide experimental results to address this aspect by comparing the remaining differences between GEN and TFS. Specifically, we focus on analyzing the impact of the number of decoder layers and the conditioning function. It is important to note that this analysis does not apply to MSA, as it does not. possess an encoder-decoder structure.
GEN and TFS differ not only their latent representation but also on two other points: (1) the way that they access the encoded information in the decoder and (2) how much computing power used in the decoder stack.
21

---

Published in Transactions on Machine Learning Research (04/2024)
Table 8: Results of the ablation of the conditioning function used to combine the latents with the target position. We adapted the architecture so that they used both a distance-basse conditioning function, which combines the query position with a weighted average of the nearest latents and standard cross-attention. We show in italics hybrid architectures where we had to switch the default conditioning method. For GEN and CNP replacing the conditioning method does not impact the performance, but for TFS switching to a distance based conditioning function impacts the performance drastically.
Distance- based
Cross-
Model
attention
CNP
.115 .015
.199 1 0014
GEN
.024 + .004
.022 + .006
TFS
042 1 .023 .020 1 .011
Table Results of the decoder-layer ablation. One difference between TFS and GEN is that by default TFS multiple layers in its decoder whereas GEN uses one, and delegates all processing to the encoder. We see that having multiple decoder layers helps the transformer whereas it impacts the performance of the GEN.
# Decoder layers
Model
1
4
GEN
.021 i .002
.040 1 .012
.106 1 .004
TFS
.020 + .011
.019 = .005
.013 + .001
In all three models, we tried using cither a distance-based function or cross-attention to combine the target position with the encoded latents. CNP concatenates the same context vector to all larget queries, which we model as equivalent to averaging the latents based on the distance in the case that there is only one latent.
Specifically, each latent associated with a position x the underlying space with corresponding value ly- Conditioning is made by averaging the latent features based on their distance with the query Xt. Its formula is given by:
2
r(dist(xi, xz))
(21)
IElatents
Where r ::R- - [0, 1] is a function that maps distances to weights, and x; represents the position associated with the latent nodes with value ly- This vector is then concatenated to the larget position x, and fed to the decoder which is MLP in this case.
GEN uses this distance-based aggregation function by default whereas TFS uses cross-attention. We tried using cross-allention for CNP and GEN and using the same distance-based aggregating function for TFS using the distance between context position and target position as a reference for this scheme. For GEN and CNP using cross-attention gives approximately the same results, but using distanco-based conditioning for TFS hurts the performance significantly.
We also ran ablations that created hybrid cases for both transformers and GEN and concluded that adding decoding layers helped TFS reach better performance. Adding layers to the GEN decoder drastically reduces model performance, which seems to be coherent with the fact that GEN was shown to be more prone to overfitting in the results section.
Finally, we want to highlight that MSA outperforms all the models presented in this section by a large margin Table?
H Wind nowcasting metrics
This section defines the metrics used for wind nowcasting.
22

---

Published in Transactions on Machine Learning Research (04/2024)
Root Mean Square Error (RMSE) It takes the square root of the square distance between the prediction and the output. It has the advantage of having the same units the forecasted values.
RMSE((, = ΣK (& tx)
(22)
N
Mean Absolute Error for angle (angle MAE) It is interesting and sometimes more insightful to decompose the errors made by the models in their angular and norm components. This is the role of this metric and the following.
angle MAF((, [I(a(t) (()))l
(23)
where a f (1))
arctan (y,))
360
2x
Mean Absolute Error for norm (norm MAE) Following the explanation of the last paragraph:
norm MAE(E,
Σ || & 1l2 t lla
(24)
Relative Bias (rel BIAS) x,y Additionally, used weather metrics in Ghiggi et al. 2019). The relative bias measures if the considered model under overestimates one quantity. It is defined as:
rel BIAS(E, mcan(f4 tx)
(25)
mean(fx)
Ratio of standard deviation (rSTD) The ratios standard deviation indicate whether the dispersion of the output of the model matches the target distribution. It has an optimal value of
std (f) = std(t)
(26)
NSE The last domain metric used is the Nash-Sutelifle efficiency (NSE), which compares the error of the model with the average the target data. A score of ideal and a negative score indicates that the model was worse than the average prediction on average.
MSE(f, t) NSE t) =1- MSE(t, mean(t))
(27)
Influence of the tightness of the measurements the performances
To evaluate the influence of the tightness of the measurements on the results we designed the following experiment: We start with pair of context (xx, ye) and target points y.). We pick one point from the context and restrict the context to small number of measurements close to this point. Then, calculate the error made by the model when it is given only this restricted context and rank them by their distance to the context's center. As the dimensions (longitude, latitude, altitude) differ, we first normalized the data along each axis. We repeat this process and average the results.
We see that the error increases with the distance to the context. Now in practice the context not restricted and contains points all over the space (so the mean minimum normalized distance to the target points is usually often below 1.0)
23

---

Published in Transactions on Machine Learning Research (04/2024)
Table 10: Influence of the tightness the measurements on the performance. We restrict the context to measurements that are close to the query point and then measure the error as a function of the distance from this query point averaged over the whole dataset. We estimate the standard deviation using three random soeds.
Normalized Dist
MSA
TFS
0.5-1.0
5.56 + 0.49
7.73 + 0.37
1.0-1.5
6.62 + 0.35
7.35 0.53
1.5-2.0
7.57 + 0.43
7.87
0.36
2.0 2.5
9.52 + 0.47
9.48 0.66
2.5-eni end
10.35 + 0.49
10.53 + 0.61
We assessed general uncertainty by training the model using various random seeds, resulting in an ensemble of models from which we can determine the standard deviation and add it to the results table of the previous experiment.
J Influence of the target positions
In this section, assess the impact of the target position the training of the MSA model.
The context data for the Poisson equation is heavily prescribed (where context points belonged to the source or sink or boundary conditions), while for the Darey Flow equation, we constructed a similarly-shaped data set by subsampling the data set from (Alet et al. 2019) to create the context and targets. Thus, the Darcy flow test problem gives an ideal testbed for examining the influence of target positions on the training. To examine the impact of the context position, we designed an additional experiment where we sampled the context from the bottom left quadrant of the unit cube and the largets from the upper right quadrants of the unit cube. We conduct our analysis for small (5k parameters) and large (100k parameters) models. We report the test loss as a function of the percentage of the context from the upper-right quadrants (where all targets are located, the remaining from the bottom-left quadrant).
Table 11: Influence of the target position on the training. We report the test error for the MSA model at two different sizes. We devised an experiment in which we extracted the context from the bottom left quadrant of the unit cube, while the targets were sampled from the upper right quadrants of the same unit cube.
0%
2%
25%
50%
100%
MSA 5k
0.14
0.15
0.14
0.09
0.03
MSA 100k
0.14
0.14
0.13
0.09
0.03
In the first column (0% of the data from the first quadrant), we see that the MSA model struggled to learn when the context and targets are disjoint. The error improves with greater overlaps between context and target, with p 100% corresponding to the case in the paper (albeit dynamics hold for both model sizes. We believe that this is due to the Darcy Flow simulation being highly location-dependentt, as can be seen in some examples from the dataset Figure 8
one-quarter of the data). Similar
Regarding the training performance, we noticed no substantial difference in the total time to a solution, they all took approximately 1000 epochs in all cases. Larger models were always able overfit the data.
K Different Message Passing Scheme
We try here different messago-passing schemes even one using attention, to help us demonstrate the fact that the whole family of models that encodes the space a single graph suffers from the same bottlenecking effect. Here the results of the wind nowcasting task:
24

---

Published in Transactions on Machine Learning Research (04/2024)
Table 12: Influence the target position on the training. We report the train error for the MSA model at two different sizes. Both models were able to overfit the data.
0%
2%
25%
50%
100%
MSA 5k
0.07
0.07
0.07
0.05
0.02
MSA 100k
0.00
0.00
0.00
0.00
0.00
Table 13: Performance of GEN with different message passing schemes for the wind nowcasting task. Various message-passee schemes, including those with al tention, were explored using PyTTorhh Geometric (Fey & Lenssen 2019). The best-perfrrming GEN model with the default message- peagge-spsingggge scheme indicated with a reference line. While certain schemes showed improvements over GEN's performance, the overall performance remained relatively consistent and was still outperformed by MSA. For additional information on the different message-passiing schemes and related references, please refer to the PyTorch Geometric documentation.
Model
Score
TransformerConv
9.28
GATv2Conv
9.34
GeneralConv
9.37
GATConv
9.77
SAGEConv
9.78
ARMAConV
9.84
TAGConv
9.87
SGConv
9.95
SuperGATCCov
10.10
GCNConv
10.13
LEConv
10.61
GENConv 23.98
Although cerlain message-passing schemes enhance the performance of the GEN model, it still falls short of the performance achieved by MSA. We attribute this difference to the bottlenecking effect caused by the latent graph. For additional information on the various messago-passinn schemes and relevant references, please consult the Py Torch Geometric documentation available at http:://pptorrdr-emmerriiiii. readthhhods io/ em/latest/noouues//m htw.:hhtmmllolooocccm Foo.
L Hyperparameter sensitivity analysis
In this section, analyze the robustness of MSA to hyperparameter changes. We focus on the number of layers Figure 12a] and the hidden dimension Figure 12c in the case of the Poisson equation. In the Table 5) we report three different sizes of the different models, and we tried to keep them as balanced as possible. The exact configurations of the models can be found the config folder of the code, however, for the sake of completeness, we grid-searched over the mentioned hyperparameters while randomly sampling the learning rate.
We also did it for the width hyperparameter of GcoFNO, in the case of the Poisson equation Figure 12c and the wind nowcasting experiment Figure 12g.. In the case of the Poisson experiment, the randomized trial helps in finding the best hyperparameter, Icading to a model that is on par with MSA, we suspect that in this setup, the context positions are more evenly distributed in space, which might lead to setup that does not suffer as much from the flaws listed in Apppendix DJ however in the case of the wind experiment, the randomized experiment did not help as much and the model is still lacking behind MSA and GEN. However, it might be possible to GeoFNO this special case of dataset irregularly sampled in space is an interesting avenue for future work.
25

---

Published in Transactions on Machine Learning Research (04/2024)
12
04
01:
42h
GCM
1640
0.000C
CA02
0.00
(a) Impact of number of layers on the performance of MSA in the Poisson Experiment.
(b) Validation loss function the learning rate when varying the number of layers for MSA in the Poisson Experiment.
285
2.0
115
Cebobotts
Learning naa
(c) Impact of the hidden dimension on the performance of MSA in the Poisson Experiment.
(d) Validation loss in function of the learning rate when varying the hidden dimension for MSA in the Poisson Experiment.
(e) Impact of the width hyperparameter the perfor- mance of GeoFNO for the Poisson equation.
(I) Validation loss in function of the learning rale when varying the width hyperparameler for the Poisson equa- tion.
226
(g) Impact of width hyperparameter on the perfor- mance of GeoFNO for the wind nowcasting task.
(h) Validation loss in function the learning rate when varying the width hyperparameter for the wind nowcast- ing task.
Figure 12: Different hyperparameter sensitivity analysis for the Poisson equation and the wind nowcasting task.
26

---

Published in Transactions on Machine Learning Research (04/2024)
M Broader Impacts
We do not anticipate any significant adverse effects on socicty due to this work. Generally, having an additional method for weather nowcasting may lead to an increase in computational resources required, but this is a common consequence of many deep learning systems.
We believe that enhancing the reliability of weather forecasts and dynamical systems has a more positive impact on society. By improving the accuracy and precision of these predictions, we can provide valuable information for various sectors, such agriculture, transportation, disaster management, and overall planning and decision-making processes.
N Limitations
Our method, like other iivenntionbbased models, suffers from quadratic scaling. In the case of MSA, it is slightly more computationally intensive compared to TFS due to the combination of target and context. in the same sequence. This results in an attention mechanism that scales with O((N, + N.))), which is asymptotically equivalent to the scaling of transformers, which is O(N + N + N,N,), but in practice introduces an additional N&N, term. However, in all our experiments, this computational overhead did not pose a problem. We acknowledge the issue in the scaling of the model and refer possible solutions outlined in the related work to address this challenge.
Another drawback of this study is the absence of a comparison between the models and other established models used for modeling dynamical systems on grid, as discussed in the work by Li et al. (2021) ). We anticipate that our approach may not perform as effectively as competing models on grids due to two reasons. Firstly, the previously mentioned scaling issue becomes significant. when dealing with larger grid sizes, such as 1024 x 1024 grids. Additionally, our approach lacks certain inductive biases that aid in grid-based modeling. Nonetheless, we believe that attention-basod models still hold promisc for grid-based systems, as demonstratod by the success of Vision Transformers (ViT) in other domains Dosovitskiy et al. 2020 However, tokenization may be necessary for their effective implementation.
O Experiment and training details
In this supplementary material, we provide the code required to process the dataset and reproduce the experiments. We use Hydra configuration manager (Yadan 2019) and include the precise configuration for each case. The experiments were executed multiple times a GPU cluster, but each experiment can be conducted on a single GPU in a relatively short timeframe, ranging from few hours to maximum of few days.
P Amount of Compute needed to replicate the experiments
Training smaller models (5k, 20k) usually lakes a few hours. Training larger models (except for CNP which is considerably faster) runs in at most two GPU days. We estimate the number of GPU days to replicate all experiments for all models be on the order of 100 GPU days.
Q Reproductibility
We provide the link to the dataset and the whole code base for processing it and running the experiments.
R Licenses
All concerned databases are openly accessible on the web and have permissive licenses, we give a link to each dataset
Table
27


Reviews:
[{"review_id": "EbmM2dqruy", "reviewer": "Reviewer_8Rzh", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n- The approach is well-motivated \n- The empirical analysis is done on a broad range of problem domains including publically available datasets\n- The analyzed baseline models sufficiently demonstrated the advantages of the proposed approach\n- Additional analysis is proposed for context retrieval and more saliently illustrates the limitations that arise from bottlenecks in the representation of models for sparse data. \n\n**Weaknesses:**\n- The presentation is somewhat messy. Specifically, the notation introduced in (1)-(4) operates per measurement, then in (5) $c_e$ is over the whole set. In (5) $c_l$ does not seem to have a role. \n- The explanation of GEN (10)-(13) seems superficial and insufficient to understand the method without referring to the original publication. Such a half-detailed presentation does not play a good role in the paper.  \n- For CNP equation (15) $t_l$ is undefined.\n- In the experiments section GKA method is mentioned. It is unclear why. It is not a baseline in the paper.\n- The figures are poorly distributed throughout the paper. The reader needs to search for the figures and find them far from where they are referenced. \n- Table 2. What does size refer to? The number of parameters of the model?\n- The captions of the figures and the axis sometimes miss relevant information. E.g. Figure 3 is about wind forecasting, this is not indicated, nor is the metric in the vertical axis. \n- Table 4 is referred to in the main body, but the table is actually in the appendix.\n- What is the metric used in Figure 4?\n- Why introduce the GNG and the PER model in the 5.1? What is the role of these models in this analysis?\n- The term disentangled representation is used in this case when referring to a representation that this distributed over space. This term is already used in the field of \"Learning disentangled representation\" which has a different meaning. I think this can be misleading and a more precise term should be used. \n- I fail to understand Figure 5 and Section 5.2. It is not clear to me how the perturbation is done and the error is backpropagated. I do not understand what the implication of this analysis is. \n- There are several grammatical errors throughout the paper.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The work proposes a new method for modeling an unstructured set of measurements sparsely placed in space and time. The proposed architecture is transformer-based and addresses the challenge of bottlenecks in the learned representation in existing methods for such data. \n\nThe performance of the proposed method is validated on four problem domains: high-altitude wind nowcasting, two-day weather forecasting, fluid dynamics, and heat diffusion. \n\nThe results show improved performance against two existing methods; Graph Element Networks and Conditional Neural Processes. Moreover, the authors included in the analysis a transformer-based method with a representation bottleneck as an additional baseline to demonstrate the advantage of avoiding the representational bottleneck.  \n\nIn addition to the analysis of real-world datasets and simulated dynamical systems, the authors present an additional experiment where the model is used to do context retrieval. The goal of this experiment is to show that sufficiently capable models can perfectly retrieve data from context while models that have a representational bottleneck struggle to achieve the same performance when the data has high spatial frequencies. \n\nFinally, the authors present an analysis of the capability of the different architectures to do error correction when exposed to localized perturbations.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Please address or respond to all the indicated weaknesses given above. Particularly the perturbation analysis and Figure 5. \n- Overall the current version of the paper does not meet the criteria for \"clearly written\" and as such I will answer the question on \"Claims And Evidence\" with a \"No\" at this time.", "broader_impact_concerns": "No concerns about ethical implications.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "mLaPDIOc5H", "reviewer": "Reviewer_dDd7", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The issue of learning from irregularly sampled data is important. In cases like robotic control, the data is multi-modal, asynchronous, and irregularly sampled. It is imperative to develop an efficient approach to manage such circumstances.\n- The proposed method is simple yet effective, enhancing performance across a variety of diverse tasks.\n- This study earnestly addresses the challenge of comprehending the mechanisms behind their model. Through detailed experimentation, the bottleneck effect has been identified as the pivotal element impacting the model's performance.\n\nWeaknesses:\n- (Main Concern) While I appreciate the authors' endeavor to explore the failure modes, there appears to be a lack of continuity in the logical sequence: The analysis in section 5 focusing on a separate information retrieval task rather than the initial ones. Detailed experiments demonstrate the bottleneck effect and the advantages of the disentangled latent representation. The information retrieval task necessitates that the models exhibit a precise recollection of the data. Nevertheless, it is questionable that whether the initial tasks demand such a strong requirement. To be concrete, the gap in the logic is *why the ability of accurately storing data enhances the performance in tasks such as the wind nowcasting.* \n  - Typically, DL models  strive to avoid memorizing all the data to prevent issues like overfitting. \n  - In addition, when the model mixes two context inputs (instead of remembering them independently), the resulting fused representation could be beneficial (e.g., see works in the multimodal learning literatures)\n  - Therefore, the aforementioned logical gap needs to be carefully addressed.\n- Regarding the discussion in subsection 5.2:\n  - the gradient-based analysis only concern the well-trained model (i.e., pre-trained to zero error), a model during training can have different behaviors. It may be too strong to use the word \"independently\"\n  - In Fig.6 the GEN5X5 case is slightly lower than the GEN6X6 case, indicting the underlying training dynamic could be too complex to be summarized solely by the degree of representation entanglement.\n- Regarding the presentation:\n  - In subsection 3.1 regarding symbol notation, it would be more intuitive to adopt $x_c, y_c$ and $x_t, y_t$ and not $c_x, c_y$ or $t_x, t_y$. It is much more intuitive to use $x$ for features and $y$ for targets/labels with subscripts such as $t$ and $c$ indicating their identities. \n  - From subsection 3.2 to 3.4, symbols have not been adequately defined. As an example, $\\gamma(\\cdot)$, which presumably denotes a FC readout module, are not explicated (Interestingly, I also note in section 5.2, $\\gamma$ is used to represent the output perturbation). \n  - Tables 2 and 3 lack citation in the main body of text, while Table 4, situated in the supplementary material, is referenced. This continually disrupts the reading rhythm.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This study introduces an attention-based model to address the issue of learning from irregularly sampled data. When compared with baseline models, the proposed model improves performance in a wide range of diverse tasks. Moreover, a detailed analysis is conducted to illuminate the reasons behind the superior performance of proposed models over others.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Please address the question: why does the ability of accurately storing data enhance the performance in tasks such as the wind nowcasting? \n- Improve the presentation.\n\n(for details pls refer to the weaknesses)", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "LBzjm6iW2b", "reviewer": "Reviewer_hrQX", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n-\tThis paper is well-written.\n\n-\tExtensive experiments and detailed analysis are provided to evaluate the model's effectiveness.\n\n-\tThe experiments to test the representation disentanglement are interesting.\n\n## Weaknesses\n-\tThe proposed MSA model is not significantly over the canonical Transformer (TFS). Also, with Section 5.1-5.3, I can understand why Transformer-based models outperform other models. But, what about the gap between TFS and MSA? More explanations are expected. For example, why removing Decoder and shared encoder can benefit the model performance?\n\n-\tThe hyperparameter Sensitivity test is missing, such as the number of layers and deep representation channels.\n\n-\tMore baselines are also expected. Geo-FNO [1] can also handle the irregular mesh, which should be compared.\n\n[1] Fourier Neural Operator with Learned Deformations for PDEs on General Geometries, arXiv 2022", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a simple model Multi-layer Self-Attention (MSA) to recover the complete field from sparse observations. MSA firstly employs a shared encoding layer to project the input observations and queries into deep representations and then infer the values on the query position with self-attention. Experimentally, MSA surpasses canonical Transformers, GNNs and CNPs in Wind Nowcasting and Poisson, Navier Stokes and Darcy Flow equation and the weather forecasting task.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "In general, I think this paper is well-qualified. More detailed explanations and comparisons to more baselines are expected.", "broader_impact_concerns": "This paper only focuses on technical designs. Thus, there are no ethical risks and extra concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Problem statement is interesting and relevant. Both physics based and machine learning baselines are evaluated. The model is simple but effective.
Decision: Accept (poster)