A Hierarchical Spatial Transformer for Massive Point Samples  in Continuous Space
Authors: Wenchong He
Zhe Jiang
Tingsong Xiao
Zelin Xu
Shigang Chen
Ronald Fick
MILES D MEDINA
Christine Angelini
Venue: NeurIPS 2023 poster (2023)
Keywords: Spatial representation learning
transformer
quadtree
efficiency

Abstract: Transformers are widely used deep learning architectures. Existing transformers are mostly designed for sequences (texts or time series), images or videos, and graphs. This paper proposes a novel transformer model for massive (up to a million) point samples in continuous space. Such data are ubiquitous in environment sciences (e.g., sensor observations), numerical simulations (e.g., particle-laden flow, astrophysics), and location-based services (e.g., POIs and trajectories). However, designing a transformer for massive spatial points is non-trivial due to several challenges, including implicit long-range and multi-scale dependency on irregular points in continuous space, a non-uniform point distribution, the potential high computational costs of calculating all-pair attention across massive points, and the risks of over-confident predictions due to varying point density. To address these challenges, we propose a new hierarchical spatial transformer model, which includes multi-resolution representation learning within a quad-tree hierarchy and efficient spatial attention via coarse approximation. We also design an uncertainty quantification branch to estimate prediction confidence related to input feature noise and point sparsity. We provide a theoretical analysis of computational time complexity and memory costs. Extensive experiments on both real-world and synthetic datasets show that our method outperforms multiple baselines in prediction accuracy and our model can scale up to one million points on one NVIDIA A100 GPU. The code is available at https://github.com/spatialdatasciencegroup/HST

A Hierarchical Spatial Transformer for Massive Point Samples in Continuous Space
Wenchong He
Zhe Jiang"
Tingsong Xiao Zelin Xu Shigang Chen
Department of Computer & Information Science & Engineering
University Florida {whe2, zhe. jiang, xiaotingsong, zelin.xu, egchen)@ufl.coo edu
Ronald Fick
Miles Medina Christine Angelini
Center for Coastal Solutions
University of Florida
(rfick, miles. medina)Qufl.edu christine. awbee..ngliiniin angelini@essia
Abstract
Transformers are widely used deep learning architectures. Existing transformers are mostly designed for sequences (texts time series), images videos, and graphs. This paper proposes a novel transformer model for massive (up to million) point samples in continuous space. Such data ubiquitous in environment sci- ences (e.g., sensor observations), numerical simulations (e.g.. particle-lade flow, astrophysics), and seaatio-bbseed services (e.g., POIs and trajectories). However, designing transformer for massive spatial points is non-trivial due to several challenges, including implicit long-range and multi-scale dependency irregular points continuous space, a non-uniform point distribution, the potential high com- putational costs of calculating all-pair attention across massive points, and the risks of over-confident predictions due to varying point density. To address these chal- lenges, we propose new hierarchical spatial transformer model, which includes multi-resolution representation learning within quad-tree hierarchy and efficient spatial attention via coarse approximation. We also design an uncertainty quantifi- cation branch to estimate prediction confidence related to input feature noise and point sparsity. We provide a theoretical analysis of computational time complexity and memory costs. Extensive experiments on both real-world and synthetic datasets show that our method outperforms multiple baselines prediction accuracy and our model can scale up to one million points on one NVIDIA A 1000 GPU. The code is available at ftttt::://// //tp:////tthbb com wttp://ppaiaalaaaacceeeegooopppp
Introduction
Transformers are widely used deep learning architectures. Existing transformers are largely designed for sequences (texts or time series), images or videos, and graphs [ proposes novel transformer model for massive (up to million) points in continuous space. Given set of point samples in continuous space with explanatory features and target response variables, the problem is to learn the spatial latent representation of point samples and to infer the target variable at any new point location.
o 7 This paper
Learning transformers for continuous-space points has broad applications. In environmental sciences, researchers interested in fusing remote sensing spectra with in-situ sensor observations at irregular
"Contact author
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

sample locations to monitor coastal water quality and air quality 8 9 10 In scientific computing,
researchers learn neural network surrogate speed up numerical simulations of particle-laden
flow [II or astrophysics [12]. For example, in cohesive sediment transport modeling, a transformer
surrogate can predict the force and torque of a large number of particles dispersed in the fluid and
simulate the transport of suspended sediment 13 [4] In location-based service, people are interested
in analyzing massive spatial point data (e.g, POIs, trajectories) to recommend new locations [65]6]]
However, the problem poses several technical challenges. First, implicit long-range and multi-scale dependency exists on irregular points in continuous space. For example, in coastal water quality monitoring, algae blooms in different areas are interrelated following ocean currents and sea surface wind (e.g. Lagrangian particle tracking). Second, point samples can be non- non-uniformn can-unfformly distributed with varying densities. Some areas can be covered with sufficient point samples while others may have only very sparse samples. Third, because of the varying sample density as well as feature noise and ambiguity, model inference at different locations may exhibit different degree of confidence. Ignoring this risks over-confident predictions. Finally, learning complex spatial dependency (e.g.. all-pair self-attention) across massive (millions) points has high computational cost.
To address these challenges, we propose new hierarchical spatial transformer model, which includes multi-resolution representation learning within a quad-tree hierarchy and efficient spatial attention via coarse approximation. We also design an uncertainty quantification branch to estimate prediction confidence related to input feature noise and point sparsity. We provide theoretical analysis of computational time complexity and memory costs. Extensive experiments on both real-world and synthetic datasets show that our method outperforms multiple baselines in prediction accuracy and our model can scale up to one million points on one NVIDIA A100 GPU.
2 Problem Statement
A spatial point sample is a data sample drawn from 2D continuous space, denoted as y(s)) si), where 1 < i S n, S; R2 is 2D spatial location coordinates (e.g., latitude and longitude), x(s:) E R x is a vector of m non- con-spatial explanatory features, and y(s;) is a target response variable (y(s;) E E for regression, yss))) E {0. 1} for binary classification). For example, in water quality monitoring, a spatial point sample consists of non-spatial explanatory features from spectral bands of an Earth imagery pixel, the spatial location of that pixel in longitude and latitude, and ground truth water quality level (e.g., algae count) at that location.
We aim to learn the target variable continuous-space function y : R R or {0, 1). Given set of point sample observations O = (s:), v(S.). {=)) where S; is irregularly sampled in 2D, our model learns the contiguous function y that can be evaluated at any new spatial points & {s: )=1- We define our model as learning the mapping from the observation samples to the continuous target variable function y(s) = OO, (s $). Thus, we formulate our problem as follows.
Input: Multiple training instances D = {0,} (sets of irregular points in continuous 2D space).
Output: A spatial transformer model f {v((), u(})} = f(x()), $, O2) for any j E [1,..., L], where s is any new sample location for inference, and x($) and ySS) are the explanatory features and output target variable for the new sample, respectively, and us$) ) is the uncertainty score corresponding to the prediction v()).
Objective: Minimize prediction errors and maximize the uncertainty quantification performance. Constraint: There exists implicit multi-scale and long-range spatial dependency structure between point samples in continuous space.
Note that supervise model training, we can construct the training instances by removing single point sample from each set O, as the new sample location and use its target variable as the ground truth.
Related Work
Attention based transformers are widely used deep learning architecture for sequential data (e.g., texts, time series) M1, images or videos [5 6 graphs 77. One main advantage
Transformer models:
of transformers the capability of capturing the long-range dependency between samples. One major computational bottleneck is the quadratic costs associated with the all-pair self-attention. Various
2

---

techniques have been developed to address this bottleneck. For sequential data, sparsity-base methods [7 T8 and low-rank-basee methods [9 20 21 have been developed. Sparsity-based
methods leverage various attention patterns, such as local attention, dilated window attention, or cross-partition attention to reduce computation costs. For low-rank-base
methods, they assume the attention matrix can be factorized into low-rank matrices. For vision transformers, patch-based [22] 23 or axial-based [24] 25 have been proposed to improve computational efficiency. Similarly, for P PS] and spectral-basee [29] 30 attention mechanisms were proposed to reduce computation complexity. However, these techniques require an explicit graph structure with fixed topology. To the best of our knowledge, existing transformer models cannot be applied to massive point samples in continuous space.
graph data, sampling-base -ased 126
Numerical Operator learning in continuous space: Neural operator learning aims to train a neural
network surrogate as the solver a family of partial differential equation (PDE) instances [3]. The surrogate takes initial or boundary conditions and predicts the solution function. Existing surrogate
models include deep convolutional neural networks (2]. graph neural operators
52 Fourier
neural operators B3 DeepONet [3I, NodeFormer [26 B5 and vision transformers 36 38
However, existing methods are mostly designed for regular grid or fixed graph node topology and thus cannot be applied to irregular spatial points. There are several methods for irregular spatial points in continuous space through implicit neural representation [39 40 41], or fixed-graph transformation [2)) but their neural networks are limited by only taking each sample's 's individual spatial coordinates without explicitly capturing spatial dependency.
Deep learning for spatial data: Extensive research exists on deep learning for spatial data. Deep convolutional neural networks (CNNs)
44 are often used for regular grids (e-g., satellite images, and graph neural networks (GNNs) +7 are used for irregular grids (e.g., meshes with irregular boundaries)
4
and global climate models)
45 46
4S 49 500 or spatial networks (e.g.. river or road However, CNNs and GNNs only capture local spatial dependency without long-range interactions. In recent years, the transformer architecture [| 54 has been widely used for spatial representation learning with long-range dependency, but existing transformer models are often and thus cannot be directly applied
networks) [5T 5 53
designed for regular grids (images, videos) [5] K. 23
25
22
to irregular point samples in continuous space.
4 Approach
This section introduces our proposed hierarchical spatial transformer (HST) model. Figure I shows the overall architecture with encoder and a decoder. The encoder learns multi-resolution representation of points via spatial pooling within a quad-tree hierarchy (quadtree pooling) and conducts efficient hierarchical spatial attention via point coarsening. The intuition is approximate the representation of faraway key points by the representation of coarse quadtree cells. The decoder makes inferences at new point location by traversing the quadtree and conducting cross-attention from this point to all other points. The decoder also contains an uncertainty quantification (UQ) branch to estimate the confidence of model prediction related feature noise and point sparsity. Note that our model differs from existing tree-based transformers [35] 56 for images or videos, as these methods require regular grid and cannot directly applied irregular points in continuous space.
4.1 Multi-resolution representation learning within a quadtree hierarchy
Learning multi-resolution latent representation of point samples is non-trivial due to the non-uniform (irregular) distribution of points in continuous space (instead of regular grid in images (22)). To address this challenge, we propose to use a quadtree to establish a multi-scale hierarchy, a continuous spatial positional encoding for point representation, and spatial pooling within quadtree hierarchy to learn multi-resolution representations.
4.1.1 Spatial representation of individual points (quadtree external nodes)
Continuous spatial positional encoding: A positional encoding is a continuous functional mapping
Φ: R + R from the 2D continuous space to 5 --dimensional encoded vector space. The encoding function needs to allow a potentially infinite number of possible locations in continuous space and the similarity between the positional encodings of two points should reflect their spatial proximity, i.c.,
3

---

Encoder
Grid Points
Quad-tree
Hierarchical
Quad-tree
Feed Forward
Spatial
Pooling
Attention
Fronsre Evnlasbbing
L layers
Decoder
Quad-treee
Portios Exeadinn
Hierarehical
Feed
Prediction
Spatial Attention
Forward
Uncertainty Quantification
Fastrre Jastedding
New Palnt
L2 layers
Figure 1: The overall architecture of our hierarchical spatial transformer model.
nearby samples tend to have a higher dot product similarity in their positional encoding. Common positional encodings based on discrete index numbers for sequence data insufficient. We propose to use multi-dimennional continuous space position encoding (57) as follows,
((s) (Sos([2, sin(S cos s), sin(Ay
where d is the encoding dimension, Q2 ~ (O, 2) is 1 by projection matrix following an i.idd
Gaussian distribution with a standard deviation The advantage of this encoding is that satisfies the following property, < ((11). ((82) >x k[]] $1 s2ll) = exp(-(s1 $2)) 1 82)). Here the hyperparameter E controls the spatial kernel bandwidth. Next, use (00;) to denote the positional encoding ((s;) for consistency.
Spatial representation:
We propose an initial spatial representation of individual points by concate- nating its continuous-space positional encoding and non-spatial feature embedding, i.e..
h(o,)) [+(o)); ((o.)
where Φ(o.) is a positional encoding, (b,) = w [x is the non-spatial embedding. W E R x (m++)) is the embedding parameter matrices, and is the dimension of concatenated representation. We denote the representation of all point samples (external quadtree nodes) matrix H,- =
Raxd, where each row is the representation of one point.
4.1.2 Spatial representation of coarse cells (quadtree internal nodes)
To learn multi-resolution representation of non-uniformly distributed points continuous space, we propose a spatial pooling operation within a quadtree hierarchy. A quadtree is a spatial index structure designed for large number of points in 2D continuous space. recursively partitions rectangular area into four equal cells until the number of points within cell falls below a maximum threshold. In a quadtree, an internal node a higher level represents an area at coarser spatial scale and nodes at different levels provide multi-resolution representation. Another advantage of using quadtree that can handle non-uniform point distribution. A subarea with denser points will have a deeper tree branch through continued recursive space partitioning.
Formally, given the set of point samples O {o in continuous space, we construct quadtree T. The quadtree has two kinds of node sets: an external node set E, which corresponds to the observed spatial point samples, and internal node set Z, which represents the spatial cells. quadtree has L levels, and all the nodes in level I form set R₁ = {ri. rk,, J. where r is the j-th node at level i, and ky is the total number of nodes in level 1. Given ((r)) (nodes on same hierarchical level under the same parent node), the ancestor node set as A(r)) (nodes on the path from the node r to the root). We denote the spatial representation of the node r/ as h(rr)) and L) and j E (1,...,k)) kr)-
node ry. we represent its sibling node set as
For example, in Figure a)) there are 11 input point samples with denser point distribution in the upper left corner and the lower right corner. Assuming maximum leaf node size of two samples, the

---

d
7
12
20
75
H
P
x
H
12
(r5)
(12)
(1
0g 09 10
Internal Node Feature Matrix
Sparse Pooling Matrix
<]
Oz 0g (5)6)))0))))
External Node Feature Matrix
(a) Space partitioning
(b) Quadtree hierarchy
(c) Quadtree pooling
Figure 2: An example of a quadtree and pooling operation.
corresponding quadtree is shown in Figure 2 b). which has 12 internal nodes (blue) and 11 external nodes (green). There are five different levels starting with level (the root node). Level has four internal nodes, i.e., R1 = (ri...,,'' ]. J. For instance, ri corresponds to the largest quad cell in the upper left corner. It has two non-empty children nodes level 2, one of which rī is an internal node and the other of which r3 is a leaf node linked to an external node r (also expressed as 0g). The sibling set (r))) is {r2), (the other two sibling nodes are empty cells and are thus ignored). The set of ancestors A(ri) is ".
Assume the total number of quadtree nodes is N, including n leaf nodes (point samples) and N internal nodes (coarse cells). We can compute the representation of each internal node by an average pooling of the representation of its children within the quadtree hierarchy (i.e., quadtree pooling).
Formally, the spatial representation H, for all quadtree internal nodes can be computed by sparse matrix multiplication, as shown Equation 3 where H, E IRaxd is the representation matrix of n point samples (external nodes). Hp E IR (g(ddd is the pooled representation matrix of N internal nodes, and P E R x is a sparse pooling matrix. Each row of is normalized and its non-zero values indicate all the corresponding external nodes (point samples) under an internal node. The computational structure shown by Figure C). We concatenate the internal node feature H, and external nodes feature H, form representation matrix E R-Nxd for all quadtree nodes.
Hp QuadtrePPooling(H. PH.,
4.2 Efficient Hierarchical Spatial Attention
The goal of the spatial attention layer is model the implicit long-range spatial dependency between all sample points. Computing all-pair self-attention across massive points (e.g.. million) is compu- tationally prohibitive due high time and memory costs. To reduce the computational bottleneck, we propose an efficient hierarchical spatial attention operation based on a coarse approximation of key points. Specifically, instead of computing the attention weight from a query point o, to all other points as keys, we only compute the weight from o, to selective subset of quadtree nodes. Our intuition is that for key points that are far away from the query point o;, we can use coarse cells (nodes) to approximate them, and the further away those points are, the coarser cells (upper-level nodes) we can use to approximate them.
Formally, for each query point (external node), we define its key node set as IC, which includes the point itself, its own siblings (points) well the siblings of its ancestors. That is, K(oo)) {o: S()))) U (S(r)]r E A(o;)). For example, in Figure 3 the key set of the external node or (also
noted rt) is {01, 02, r r2 r2.r r.).
In this way, we reduce the number of atten- tion weight calculations from all 11 points to 8 quadtree nodes. Particularly, we use one quadtree node r to approximate all the four points within it (Og to 011) since they are far away from O1. Based on the definition of key set, the spatial attention operator can be ex- pressed as Equation 4 below, where h; is the output representation for o,, K.; is the key node set o,, q, is the query vector of point o,, ky and vy are the key vector and value vector of
d.
P.
0
du
(rf)
(1)
0g 0g
010 Ou1
6
(2)))
dBo
0g
0g
(b)
Figure 3: An example of a quadtree (a) and the
selective key node of o, in red boxes (b).
5

---

the attended node ry. respectively, and d is the latent dimension of these vectors.
explq,k, /Vdjv, JEK. JJEK, exp(g,k) /Va)
We can express the spatial attention operator in matrix and tensor notations. Assume the spatial representation of all quadtree nodes from the prior layer is H E RNxx and the representation of external nodes (point samples) as H. The query matrix of all point samples can be computed by an embedding with learnable parameter matrix W, 4. i... Qo = W. For simplicity, we denote the query matrix Qo as Q. For each query point Og, its keys are subset of quadtree nodes. We denote their embedded key vectors in matrix K; E RaxiK,,, If we concatenate the key matrices for all queries, we get a 3D tensor K gdx|K,|xn Similarly, the corresponding value matrices can be concatenated into a 3D tensor V Rdx[K]/ The construction of 3D tensors can be implemented by the torch.gather( API in Pytorch. The corresponding cross-attentics weights can be calculated by matrix and tensor multiplications, as illustrated in Figure g We can see the difference between the
All Pair Afcaniin
proposed hierarchical spatial attention and the default all all-pair attention. In the all-pair self-attention (Figure R top), we would have to compute the dot product attention for all point entries, i.e. Q K.. whose size is X and can become too large (e.g.. n = 100, 000 for a million points). In our spatial attention layer, conduct a sparse self self-attention, i.e., only computing the attention weight of a sample point to its key set. other words, the corresponding key matrix for o, is a vertical slicing of K, which can be denoted as K f ;, RddxiKal, where [K]] is the maximum key node set size ([A) <<
N
d
K
N
V
Heerarchical Spatial Astention
[C]
IK
K
Figure 4: Sparse spatial attention with selective key set.
Time cost analysis: The proposed spatial attention computation cost depends on the uniformness of spatial point distribution, which determines how balanced the quadtree is. Assume that there are points and the maximum number of points in each quadtree node is M (quadtree threshold). We analyze two scenarios. the quadtree is completely balanced, the tree depth and the key set size for every external node is Oloo M )- Thus total computation cost is O(n (M log 37)) ))- In the worst case, the spatial points are highly nonuniform, then the quadtree structure will be highly imbalanced. The extreme tree depth is O " ). thus the attention computation is O(n ( [x + M)). However, such a worst-case scenario is unlikely in practice, as it would require samples to be concentrated within a single subgroup of the quadtree every level. For instance, for water quality monitoring, the sensors are sparsely distributed across a broad area to monitor multiple locations. In practical applications, M)) complexity. This is validated in our experiments.
takes (n (log
Memory cost analysis: We analyze the memory costs of the HST model theoretically. Assume the number of input point samples n. leaf node size threshold M, batch size B, and the number of head h, and hidden dimension The memory costs of HST are dominated by the hierarchical spatial attention layer. Its memory cost is O(B h d IKi)) per layer, where [G]] log + M for relatively balanced quadtree.
4.3 Decoder: inference on a new point with uncertainty quantification
Model inference (prediction): Given a test sample o, (xx, (x),8)). the decoder module predicts yt based on learned spatial representations of quadtree nodes H. Similar to the encoder, we use cross-attention between the test sample and its corresponding key node set in the quadtree in Equation 4 As shown in Figure I we first conduct a quadtree traversal of the test point until reaching the leaf node and then identify the key node set. Based on that, we can apply a hierarchical spatial cross-attention between the test location and its key node set, followed by a dense layer.
Uncertainty quantification (UQ): Due to the varying point density as well as feature noise and ambiguity, the model prediction at new location may come with different confidence levels. Intuitively, the prediction at a test location surrounded by many input point samples tends to be more confident. Many methods exist in UQ for deep learning (e.g., MC dropout and deep ensemble [58] 59 60 61). but few consider the uncertainty due to varying sample density. The most common

---

The UQ of GP is based on Equation 5 where Co E R Ixx= is the covariance vector between the test sample and all input point samples, C E R is the covariance matrix for input samples, and og is the self-variance. In GP, the covariance C is computed with a kernel function that reflects the location proximity (44] Although GP has good theoretical properties, it is inefficient for massive points due to expensive matrix inverse operation on a large covariance matrix, and it is unable to learn a non-linear representation for samples.
method for UQ in the continuous spatial domain is the Gaussian process (GP) 62 63..
o = C
Our proposed spatial transformer framework can be considered a generalization of a GP model. We use the dot-product cross- attentitn weights to approximate the covariance vector Co between the test location to all points, which reflects the dependency among point sample locations based on their non-linear embeddings, i.e., Co Hq: where qi is the embedded query vector for the test sample. However, this idea is insufficient for the approximation of the entire covariance matrix C across all points, since the inverse computation of full covariance matrix is very expensive. To overcome this challenge, we propose to directly approximate the precision matrix (C-1) based on an indicator matrix of selective key sets for all queries S E IR Xn, in which each column indicates the key node set of a query point (and thus specifies the dependency between a query point to all other quadtree nodes). Thus, we use S S/T, to approximate the precision matrix C-1. where T is hyper-paramene parameter to be calibrated by independent validation data with calibrated with the Expected Calibration Error (ECE) [65]. The intuition is that the precision matrix reflects the conditional independence structure among point samples (similar to the selective spatial attention in our model). Since S S Σ; s,s? (s; is column of S), can see that S S is summation of several sparse block-diagone -digonal matrices. This reflects our assumption in the quadtree that each external node (point sample) is conditionally independent of all other nodes given its key node set. Based on the above approximation, our uncertainty quantification method can be expressed by Equation 0 where K, SH is the key matrix corresponding the selective key set of the test sample.
U = oj (q.H') (sH")s""(Hqs))))
= 0 q (SSH)((S))
= =o
(K,g))/T
(6)
Note that approach shares similarities with the multipole graph neural operator model (MGNO) 32) MGNO employs a multi-scale low-rank matrix factorization technique to approximate the full kernel matrix across all samples. A key distinction lies in that MGNO uses neighborhood graph structure to approximate the dependency relationships. In contrast, our approach can capture the long-range interactions among samples in Euclidean space and we provide uncertainty quantification in prediction.
5
Experimental Evaluation
The goal is to compare our proposed spatial transformer with baseline methods in prediction accuracy and uncertainty quantification performance. All experiments were conducted on cluster installed with one NVIDIA A100 GPU (80GB GPU Memory). The candidate methods for comparison are listed below. The models' hyper-parameters are provided in the supplementary materials.
Gaussian Process (GP): We used the GP model based spatial location without using the explana-
tory feature 622 The prediction variance was used as the uncertainty measure. Deep Gaussian Process (Deep GP): We implememted hybrid Gaussian process neural network [66] with the
sample explanatory features and locations. The GP variance is used as the uncertainty. Spatial graph neural network (Spatial GNN): We first constructed a spatial graph based on each sample's kNN by spatial distance. Then we trained a GNN model 47 We used the MC-dropout [5X] method to quantify the prediction uncertainty. Multipole graph neural operator (MGNO): It belongs to the family of neural operator model for learning the implicit graph structure
32). NodeFormer: It is an efficient graph transformer model We used the code from its official website. Galerkin Transformer: It uses softmax-fres attention mechanism and acheieve linearlized transformer [35)] We quantify the prediction uncertainty based on MC-dropout. Hierarchical Spatial Transformer (HST): This is proposed method. We implemented it with Torch.
26
7

---

The prediction performance is evaluated with mean square enror (MSE) and mean absolute er- ror (MAE). The evaluation of UQ performance is challenging due to the lack of ground truth of uncertainty.
The quantitative evaluation metrics for UQ performance is Accuracy versus Uncertainty (AeU)) [67] We set accuracy uncertainty thresholds Te, Tow to group prediction accuracy and uncertainty into four categories Table T shows. RAC, TAU, Inc. nir represent the number of samples in the categories AC, AU, IC, IU, respectively. As Equation 7 shows, AvU measures the percentage of two categories AC and IU. A reliable model should provide a higher ArU measure (AvU E [0, 1)). The details on how we choose the thresholds are provided in the supplementary materials.
UQ evaluation metries:
PAC + "IU "AC RAU "IC "IU
AoU
However, AvU is usually biased by the accuracy of the model. Since models tend to have high confidence in accurate predictions. We propose an evaluation metric that evaluates uncertainty performance for accurate and inaccurate prediction separately. Specifically, computed AvU for accurate predictions and AvU, for inaccurate predictions with the following equations:
AvU'
AoU
(8)
MAC + YEAU
aic + AIU
In our evaluation, we computed the harmonic average of AvUA and AvU1: AvU to penalize the extreme cases instead of the arithmetic average.
A=AuU
AoU
Table I: Accuracy versus Uncertainty (AvU)
Unsoriaiitt
Certain
Uncettin
Accuracy
Accinii
Accunie (AI)))
[IC)
Anaunic Unsotton (AL)
Unceriain
Dataset description: We used three real- world datasets, including two water quality datasets collected from the Southwest Florida coastal area, and one sea-surface temperature and one PDE simulation dataset: Red tide dataset: The input data are satellite imagery obtained from the MODIS-Aqua sensor [68] and in-situ red tide data obtained from Florida Fish and Wildlife's (FWC) HAB Monitoring Database 69 We have 104, 100 sensor observations and we use sliding window with sequence length 400 to generate 103700 inputs. It is split into training, validation, and test sets with a ratio of 7 2. Turbidity dataset: We used the same satellite imagery features as in the red tide dataset. The ground truth samples measure the turbidity of the coastal water. It contains 13808 sensor observations. Darcy flow for PDE operator learning: The Darcy flow dataset [3 contain 100 simulated images with 241 x 241 resolution. For each image, we subsample 100 sets of point samples from the original image (each set has 400 nodes). Sea Surface Temperature: We used the sea surface temperature dataset of Atlantic ocean grid pixels.
7
We subsampled 400 point samples from the
Table 2: Comparison model performance on two real- world datasets and one simulation dataset
Model
Red ride
Tubicity
Darcy flow
MSL
MAE
MSL
MALI
MSL
MAL
GP
7,42 = 0.25
2.55 = 0.18
0.42 142:00.000
0.45 # 0.03
0.19 =: 0.03
0.33 2: 0.04
Deep
2.32
0.04
0.42 0.06
0.03
0.06
2.10
0.34
0.45 0.03
0.15 0.002
0.26 0.04
Al-uui mnaioonne
5,00 0.12
1.08 0.89
0.31 0.003
0.35 0.044
0.30
1 0.07
0.22 1: 0.005
MaNO
5.41
= 0.20
2.01 00110
0.322 2.222200 0:05
0.35 2 0.135
0.12
2: 0.103
0:21
2: 0.03
Nodelfomerr
5,34 - 0.10
2.06 + 0.06
0.52 0.04
0.38 + 0.04
0.16 0.08
0.29 0.04
0.34
0.05
0.08
method)
5.1 Comparison on prediction performance
We first compared the overall regression performance between the baseline models and our proposed HST model. The results on two datasets were summarized in Table 2 The first column summarizes the performance of the red tide dataset with MSE and MAE evaluation metrics. We can see that the GP model performed the worst due to the ignorance of the samples' features. The deep GP
8

---

model improved the GP performance from 7.42 to 6.23 by leveraging the neural network feature representation learning capability as well considering point sample correlation both spatial and feature proximity. The spatial GNN model performed well because the model considered local neighbor dependency structure and feature representation simultaneously. However, the spatial GNN model relies on proximity-basee graph structure so it ignores the spatial hierarchical structure and long-range dependency. MGNO baseline model performs slightly better than spatial GNN because the multi-level pressag-ppassing mechanism captures the long-range dependency structure. Our framework performs better because of the awareness of multi-scale spatial relationships, which is erucial for point samples' inference. Additionally, the performance of all-pair attention is better than the fix- -x--rrah GNN model because of the long-range modeling. However, the vanilla transformer models lack the hierarchical attention structure, which may cause imperfect attention weight score. In contrast, our model performed best with the lowest MAE of 5.25 because modeled the similar
interaction among point samples in the continuous multi-scale space efficiently. We results on the MAE metrics. For the second turbidity dataset, our model consistently performed the best in accuracy. For the PDE operator learning task, the results shown in Table 2. We proposed framework outperforms baselines except all-pair transformer but our model is more efficient. In summary, compared with recent SOTA graph transformer and neural operator learning methods, our framework performs best due to the capability of learning multi-scale spatial representations among massive point samples. Compared with the all all-pair transformer model, our HST reduces the time costs and makes a trade-off between efficiency and spatial granularity when calculating attention between points. More experiment results on the Sea Surface Temperature dataset are provided in the supplementary materials.
Sensitivity analysis: We also conducted sensitivity analysis of model to various hyper- parameters, including the quadtree leaf node size threshold M, spatial position encoding length scale a, the number of attention layers, and the embedding dimension the red-tide dataset. The results are summarized in Figure and the detailed analysis is in the Supplementary material. We can see that our model is generally stable with changes of the hyper-parameters.
6.0
6.0
5.3
5.5
5.5
5.0
6.0
5.0
5.0
2 4.3
5.3
4.5
4.5
4.0
5.0
4.0
6.0
33
33
Quedtree threshold
lengthscale
Attention
Embedoing dimension
(a) Quadtree leaf node size
(b) Length scale
(c) Attention layer
(d) Embedding dimension
Figure 5: Parameters sensitivity analysis.
5.2 Comparison on uncertainty quantification performance (UQ)
We use AvD in Equation to evaluate the performance of UQ for our proposed HST model versus baseline methods and the results were summarized in Table B The
bers in the table correspond to the number of samples in four categories: nAC, 7AL Pic, nuu We can see the base GP model had good certainty estimation for inaccurate predictions (65% are uncertain) but for the accurate predic- tions, the uncertainty estimation tended to be confident, only 14% were certain. This might be due to sparse area prediction being less confident but the long-range spatial correlation or feature similarity can improve prediction in the sparse sample area, but the GP model was unaware of such dependency. The deep GP model im- proved the accurate prediction confidence based on the feature representation learning but still performed worse than other models and gave 0.38 AvU score. The spatial GNN model was
Table 3: Comparison on uncertainty quantification
performance on Red tide dataset
Uncertiitty
Ceman
Uncernin
GP
Aocurate
2726
9144
1223
0.33
3332
4919
BAA
Deep
CPP
frsscumtt
Assualee
Centan
3316
X094
Uncertin
3509
5122
862
1.28
1138
Spatial
5152
Cenain
Transormer
Troccurate
Ausuratc
7165
3435
013
Conai
Linerttin
BST
Accurate
4342
$40
0.46
0.54
Tnoccurate
3374
6303

---

Contucaionn Cozte: Soown
1000
100
80
100 L25 159 178 200
100
200
200
600
500
(a) Computation cost per epoch for 64 training samples.
(h) Computation cost per epoch versus leaf node size threshold for 5K training samples.
(c) Computation cost per epoch versus hidden dimension for 5K training samples.
Figure Computation cost analysis
more confident in the accurate prediction, but model was over-confieent the inaccurate prediction, thus resulting in lower AvU' score (0.44). In contrast, uncertainty model improved both the accurate prediction confidence and inaccurate prediction uncertainty, and improved the overall AvU score to 0.49, because modeled the uncertainty coming from both feature space and sample density simultaneouslly The results validate our decoder attention capability in modeling the prediction uncertainty.
5.3 Analysis on computation and memory cost
We evaluate the computation costs on simulation dataset. The point samples' locations are uniformly distributed on the two-dimensional space, and we generate the input feature and target label based on the simulated Gaussian Process. The simulated number of point samples varies from 10 to 10 The training batch size is Other simulation parameters and training hyperparameters are provided in the Appendix. We compare our model with the vanilla all-pair attention transformer on both computation times. The computational time per epoch on the 64 training dataset shown in Figure ([ a). When the number of point samples increases 50K, the computation time and memory costs of the vanilla all-pair transformer increase dramatically and become out-of-memory (OOM) when further increasing the number 100K. However, model can be scaled to 1M point samples and trained with a reasonable time cost (1 hour). We also analyze the effect of the quadtree leaf node size threshold (the maximum number of point samples in the quadtree leaf) using 5K training samples. The computation time is shown in Figure 6 b). When the leaf node size threshold increase, the computation first decrease and then increase. Because when the quadtree depth decreases from 10 to 25, the quadtree depth decreases and can decrease the total number of query-key pairs. When the threshold increases from 50 to 200, the leaf node point sample will increase the computation. The computational costs with hidden feature dimensions are shown in Figure 6 c). It's observed that the computational costs scale linearly with respect the hidden feature dimension.
Conclusion and Future Work
This paper proposes novel hierarchical spatial transformer model that can model implicit spatial dependency on large number of point samples in continuous space. To reduce the computational bottleneck of all-pair attention computation, we propose a spatial attention mechanism based on hierarchical spatial representation in quadtree structure. confidence, we design an uncertainty quantification branch in the decoder. Evaluations of real- world remote sensing datasets for coastal water quality monitoring show that our method outperforms several baseline methods. In future work, we plan to extend our model to learn surrogate models for physics-simulation data on 3D mesh and explore improving the scalability of our model on multi-GPUs.
order reflect the varying degrees of
Acknowledgement
This material based upon work supported by the National Science Foundation (NSF) under Grant No. IIS-2147908, IIS IIS-2207072, CNS-1951974, OAC-2152085, the National Oceanic and
10

---

Atmospheric Administration grant NA19NES4320002 (CISESS) at the University of Maryland, and Florida Department of Envirommental Protection. Shigang Chen's work supported in part by the National Health Institute under grant R01 LM014027.
References
[1] 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.
[2] Limin Huang. Yu Jing, Hangyu Chen, Lu Zhang, and Yuliang Liu. A regional wind wave prediction surrogate model based on can deep learning network. Applied Ocean Research, 126: 103287, 2022.
[3] Anima Anandkumar, Kamyar Azizzadcneshcli, Kaushik Bhattacharya, Nikola Kovachki, Zongyi Li, Burigede Liu, and Andrew Stuart. Neural operator: Graph kernel network for partial differential equations. In ICLR 2020 Workshop on Integration of Deep Neural Modely and Disferemtial Equations, 2020.
[4] Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K Gupta, and Aditya Grover. Climax: A foundation model for weather and climate. arXiv preprint arXiv:2301 2023.
[5] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International conference on machine learning, pages 4055-4064. PMLR, 2018.
[6] Ze Liu, Jia Ning, Yue Cao, Yixuan Wei, Zheng Zhang, Stephen Lin, and Han Hu. Video swin transformer. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pages 3202-3211, 2022.
[7] Seongjun Yun, Minbyul Jeong. Rachyun Kim, Jaewoo Kang, and Hyunwoo J Kim. Graph transformer networks. Advances in neural information processing systems, 32, 2019.
[8] AG Dekker, ZZaurrvić- HJ Hoogenboom, and SWM Peters. Remote sensing, ecological water quality modelling and in situ measurements: a case study in shallow lakes. Hydrological Sciences Journal, 41(4):531 547, 1996.
[9] Jin Li and Andrew D Heap. A review of spatial interpolation methods for environmental scientists. Geoscience Australia, Record. Geoscience Australia, Canberra, 2008
[10] David W Wong, Lester Yuan, and Susan A Perlin. Comparison of spatial interpolation methods for the estimation of air quality data. Journal of Exposure Science & Environmental Epidemiology, 14(5):404-415, 2004.
[11] B Siddani and S Balachandar. Point-particle drag, lift, and torque closure models using machine learning: Hierarchical approach and interpretability. Physical Review Fluids, 8(1):01 ):014303, 2023.
[12] Volker Springel. High performance computing and numerical modelling, 2014.
[13] Eu Gene Chung, Fabián A Bombardelli, and S Geoffrey Schladow. Modeling linkages between sediment resuspension and water quality a shallow, eutrophic, wind-exposed lake. Ecological Modelling, 220(9- 1251 2009.
[14] Zhirui Deng, Qing He, Zeinab Safar, and Claire Chassagne. The role of algae in fine sediment flocculation: In-situ and laboratory measurements. Marine Geology, 1884 2019.
[15] Song Yang, Jiamou Liu, and Kaiqi Zhao. Getnext: trajectory flow map enhanced transformer for next poi recommendation. In Proceedings of the 45th International ACM SIGIR Conference on research and development in information retrieval, pages 1 1144-1153, 2022.
[16] Yanjun Qin, Yuchen Fang, Haiyong Luo, Fang Zhao, and Chenxing Wang. Next point-of-interest rec- ommendation with auto-correlation enhanced multi-modal transformer network. In Proceedings of the 45th International ACM SIGIR Conference on Research and Development Information Retrieval, pages 2612-2616, 2022.
[17] Iz Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint Xiv: 1550, 2020.
[18] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoe Le, and Ruslan Salakhutdi- nov. Transformer- xl: Attentive language models beyond a fixed-length context. arXiv preprint 1xiv:1901.02860 2019.
[19] Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee whye Teh. Set transformer:
A framework for attention based permutation-invann neural networks. In International conference on machine learning, pages 3744-3753, PMLR, 2019.

---

[20] Andrew Jaegle. Felix Gimeno, Andy Brock, Oriol Vinyals. Andrew Zisserman, and Joao Carreira. Perceiver: General perception with iterative attention. In International conference on machine learning, pages 4651- 4664. PMLR, 2021.
[21] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint DXXiv:200J.04455 2020.
[22] Shitao Tang, Jiahui Zhang, Siyu Zhu, and Ping Tan, Quadtree attention for vision transformers. arXiv preprint arXiv:2201.02760 2022.
[23] Zhi Cheng, Xiu Su, Xueyu Wang, Shan You, and Chang Xu. Sufficient vision transformer. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 190 200, 2022.
[24] Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidimensional transformers. arXiv preprint arXiv: 19/2. /2/80, 2019.
[25] Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsich. Dynamicvit: Efficient vision transformers with dynamic token sparsification. Advances in neural information processing systems, 14:13937-3394, 2021.
[26] Qitian Wu, Wentao Zhao, Zenan Li, David P Wipf, and Junchi Yan. Nodeformer: A scalable graph structure learning transformer for node classification. Advances in Neural Information Processing Systems, (127787--2740) 2022.
[27] Jinsong Chen, Kaiyuan Gao, Gaichao Li, and Kun He. Nagphormer: tokenized graph transformer for node classification in large graphs. In The Eleventh International Conference on Learning Representations, 2022.
[28] Zaixi Zhang. Qi Liu, Qingyong Hu, and Chee- Kong Lee. Hierarchical graph transformer with adaptive node sampling. Advances in Neural Information Processing Systems, 35:21171 171-21183, 2022.
[29] Devin Kreuzer, Dominique Beaini, Will Hamilton, Vincent Létourncau, and Prudencio Tossou. Rethinking graph transformers with spectral attention. Advances in Neural Information Processing Systems, 34:21618- 21629, 2021.
[30] Qitian Wu, Chenxiao Yang, Wentao Zhao, Yixuan He, David Wipf, and Junchi Yan. Difformer: Scalable (graph) transformers induced by energy constrained diffusion. arXiv preprint arXiv:2301.09472 2023.
[31] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew M Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces with applications to pdes. J. Mach. Learn. Res., 24(89): 2023.
[32] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Andrew Stuart, Kaushik Bhat- tacharya, and Anima Anandkumar. Multipole graph neural operator for parametric partial differential equations. Advances in Neural Information Processing Systems, 33:6755-6766, 2020.
[33] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli. Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprini arXiv:2010.088 088995 2020.
[34] John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catanzaro. Adaptive fourier neural operators: Efficient token mixers for transformers. arXiv preprint arX1::2111.13587. 2021.
[35] Shuhao Cao. Choose transformer: Fourier or galerkin. Advances in neural information processing xystems,
2021.
[36] Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjecv Raja, Ashesh Chattopadhyayy Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al. Fourcastnet: A global data-driven higo-resolution weather model using adaptive fourier neural operators. arXiv preprint arXiv:2202.1 1214, 2022.
[37] Kaifeng Bi, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu, and Qi Tian. Accurate medium-range global weather forecasting with 3d neural networks. Nature, 619(7970):533-522 2023.
[38] Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K Gupta, and Aditya Grover. Climax: foundation model for weather and climate, january 2023. arXiv preprint arXiv:2301. 10343, 2023.
[39] Shaowu Pan, Steven L Brunton, and Nathan Kutz. Neural implicit flow: mesh-agnostic dimensionality reduction paradigm of spatio-temporal data, 2022.
[40] Yuan Yin, Matthicu Kirchmeyer, Jean- lean-Yves Franceschi, Alain Rakotomamonjy, and Patrick Gallinari, Con- tinuous pde dynamics forecasting with implicit neural representations, arXiv preprint arXiv:2209.14855 2022.
[41] Peter Yichen Chen, Jinxu Xiang. Dong Heon Cho, Yue Chang, GA Pershing, Henrique Teles Maia, Maurizio Chiaramonte, Kevin Carlberg, and Eitan Grinspun. Crom: Continuous reduced-one order modeling of pdes using implicit neural representations. arXiv preprint arXiv:2206.088 02607, 2022.
12

---

[42] Zongyi Li, Nikola Borislavov Kovachki, Chris Choy, Boyi Li, Jean Kossaifi, Shourya Prakash Otta. Mohammad Amin Nabian, Maximilian Stadler, Christian Hundt, Kamyar Azizzadenesheli, et al. Geometry- informed neural operator for large-scale 3d pdes. arXiv preprint SiKi::230000088 2023.
[43] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25, 2012.
[44] Kaiming He, Xiangyu Zhang. Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[45] Krishna Karthik Gadiraju, Bharathkumar Ramachandra, 7exi Chen, and Ranga Raju Vatsavai. Multimodal deep learning based crop classification using multispectral and multitemporal satellite imagery. In Pro- ceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, 2020.
pages
[46] Yumin Liu, Auroop Ganguly, and Jennifer Dy. Climate downscaling using ynet: deep convolutional network with skip connections and fusion. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Duta Mining, pages 3145-3153, 2020.
[47] Thomas N Kipf and Max Welling. Semi cumi-supervieed classification with graph convolutional networks. arXiv preprint arXiv:100,002900 2016.
[48] Wenchong He, Zhe Jiang, Chengming Zhang, and Arpan Man Sainju. Curvanet: Geometric deep learning based on directional curvature for 3d shape analysis. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 2214-2224, 2020.
[49] Neng Shi, Jiayi Xu, Skylar W Wurster, Hanqi Guo, Jonathan Woodring, Luke P Van Rockel, and Han-Weii Shen. Gnn-surrogate: A hicrarchical and adaptive graph neural network for parameter space exploration of unstructured--mes -esh occan simulations. IEEE Transactions on Visualization and Computer Graphics, 28(6):230122313 13, 2022.
[50] Tingsong Xiao, Lu Zeng, Xiaoshuang Shi, Xiaofeng Zhu, and Guorong Wu. Dual-graph learning convolu- tional networks for interpretable alzheimer's disease diagnosis. In International Conference on Medical Image Computing and Compiter-Assiste Intervention, pages 406 466.445..66 Springer, 2022.
[51] Wenchong He, Arpan Man Sainju. Zhe Jiang, and Da Yan. Deep neural network for surface segmentation based on contour tree hierarchy. In Proceedings of the 2021 SIAM International Conference on Data Mining (SDM). pages 253-261 SIAM, 2021.
[52] Wenchong He, Arpan Man Sainju, Zhe Jiang, Da Yan, and Yang Zhou. Earth imagery segmentation on terrain surface with limited training labels: A semi-supervised approach based on ahysics-guides graph co-training. ACM Transactions on Intelligent Systems and Technology (TIST), 13(2):1-22, 2022.
[53] Xiaoyang Wang. Yao Ma, Yiqi Wang, Wei Jin, Xin Wang, Jiliang Tang, Caiyan Jia, and Jian Yu. Traffic flow prediction via spatial temporal graph neural network. In Proceedings of the weh conference 2020, pages 1082-1092, 2020.
[54] George Zerveas, Srideepika Jayaraman, Dhaval Patel, Anuradha Bhamidipaty, and Carsten Eickhoff. A transformer-base based framework for multivariate time series representation learning. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pages 2114-2124, 2021.
[55] Zhiruo Wang. Haoyu Dong. Ran Jia, Jia Li, Zhiyi Fu, Shi Han, and Dongmei Zhang. Tuta: Tree-based transformers for generally structured table pre-training. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pages 1780-1790, 2021.
[56] Shitao Tang, Jiahui Zhang, Siyu Zhu, and Ping Tan. Quadtree attention for vision transformers. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
[57] Yang Li, Si Si, Gang Li, Cho-Jui Hsich, and Samy Bengio. Learnable fourier features for multi-dimensional spatial positional encoding. Advances in Newral Information Processing Systems, $4:115816-15888 2021.
[58] Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pages 1050-1059. PMLR, 2016.
[59] Florian Wenzel, Jasper Snoek, Dustin Tran, and Rodolphe Jenatton. Hyperparameter ensembles for robustness and uncertainty quantification. Advances in Neural Information Processing Systems, 33:6514 14 6527, 2020.
[60] Dongxia Wu, Liyao Gao, Matteo Chinazzi, Xinyue Xiong, Alessandro Vespignani, Yi-An Ma, and Rose Yu. Quantifying uncertainty in deep spatiotemporal forecasting. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pages 1841-1851 2021.
[61] Wenchong He and Zhe Jiang. A survey on uncertainty quantification methods for deep neural networks: An uncertainty source perspective. arXiv preprint arXivv2302.13425, 2023.
13

---

[62] Sudipto Banerjee, Bradley P Carlin, and Alan E Gelfand. Hierarchical modeling and analysis for sputial data. Crc Press, 2014.
[63]
Wenchong He, Zhe Jiang, Marcus Kriby, Yiqun Xie, Xiaowei Jia, Da Yan, and Yang Zhou. Quantifying and reducing registration uncertainty of spatial vector labels on earth imagery. In Proceedings of the 28th ACM SIGKDD Conference Knowledge Discovery and Data Mining, pages 554 564, 2022.
[64]
Sudipto Banerjee, Bradley P Carlin, and Alan Gelfand. Hierarchical modeling and analysis for spatial data. Chapman and Hall/CRC, 2003.
[65] Volodymyr Kuleshov, Nathan Fenner, and Stefano Ermon. Accurate uncertainties deep learning using calibrated regression. In International conference on machine learning, pages 2796-288 2804, PMLR, 2018.
[66]
John Bradshaw, Alexander G de G Matthews, and Zoubin Ghahramani. Adversarial examples, uncertainty, and transfer testing robustness in gaussian process hybrid deep networks, arXiv preprint 177.00777 2017.
[67] Ranganath Krishnan and Omesh Tickoo. Improving model calibration with accuracy versus uncertainty optimization. Advances in Neural Information Processing Systems, 33:18237-18228, 2020.
[68]
[dataset] NASA Goddard Space Flight Center. Moderate-resolltion imaging spectroradiometer (modis)
aqua ocean level-2 data, 2021. Accessed Aug. 17, 2021.
[69]
[dataset] Florida Fish and Wildlife Conservation Commission. Hab monitoring database, 2021, Accessed Aug. 17, 2021.
[70] Emmanuel De Bézenac, Arthur Pajot, and Patrick Gallinari. Deep learning for physical processes: Incorporating prior scientific knowledge. Journal of Statistical Mechanics: Theary and Experiment. 2019(12):124009, 2019.
14


Reviews:
[{"review_id": "AI1NpzRj0z", "reviewer": "Reviewer_7rvT", "summary": "This paper proposes a hierarchical spatial transformer model for a large number of point samples in continuous space. The model is important for geoscience applications, such as water quality monitoring and air quality monitoring, and operator learning for numerical models. The novel idea includes continuous position-encoding and hierarchical attention layers, which make a trade-off between efficiency and spatial resolution and captures the interactions in multiple spatial scales. The proposed model also has uncertainty quantification that reflect the effect of sample spatial sparsity. The proposed method is compared against multiple baselines on different datasets with significant improvements. There are also computational experiments to show the efficiency of the method.", "questions": "1. Can you explain in more details how the proposed model can be used in numerical simulations (e.g., ocean current, multiphase flow) in the introduction?", "limitations": "No limitations discussed in the paper.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. The idea of hierarchical attention with a quad-tree structure for massive samples in the continuous space is quite novel. The idea makes a trade-off between computational efficiency and spatial approximation of latent representation. The spatial approximation in a quad-tree structure is well-motivated by the spatial autocorrelation effect. \n\n2. The proposed method has strong technical contributions. There is theoretical analysis of the time complexity of the proposed model and the gain in efficiency. \n\n3. There are extensive experimental evaluations against multiple state of the art methods on both real-world and synthetic datasets. The results show the better accuracy and uncertainty quantification of the proposed method.", "weaknesses": "1. There are some minor presentation issues that can be fixed. For example, Table 2 has a typo in the title. It should be “two real-world datasets and a synthetic dataset”.\n2. The application paragraph in the introduction can be more detailed to better highlight the impact of the model in operator learning for numerical simulation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer:\n\nThank you for your positive feedback and careful reading. We have done thorough proofreading and corrected the typos in the paper. For the second question about how the proposed model can be used in numerical simulations, we use multiphase flow as an example to explain. In this case, the input features consist of 3D coordinates of point samples in the continuous space as well as the initial pressure and velocity of those samples, the target output is the acceleration at any point location. We can add this into the paper."}, {"review_id": "NSZ78DpW4D", "reviewer": "Reviewer_Ma6V", "summary": "This paper proposes a quad-tree partition of irregularly distributed sample locations. This leads to an algorithm with O(NlogN) complexity.", "questions": "What is the vocab for the decoder? The solution space for regression is not discrete.", "limitations": "NA", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Quad-tree idea is nice, even though not original for irregular grids or pixels.", "weaknesses": "The targeted problem in the paper is regression. But the model is framed as an encoder-decoder. The paper does not explain why this seemingly unreasonable choice.\n\nThe O(NlogN) complexity has a large big-O constant because of sparse matrix operations which are known to be inefficient on GPUs.\n\nThere are not enough details on models and experiments. For example, the model parameters such as number of layers, hidden size and number of attention heads are not clearly specified. The training hyper parameters such as learning rate and its decay schedule are not explicitly stated. This makes it hard to determine the quality of the experiments, let alone replicating them.\n\nDatasets are all 2D with up to 1K sample points. These are well within the capacity of vanilla Transformer on A100, the GPU used in the paper. And the optimal leaf node size is 100. This results in a very shallow quad tree. So the evidence about the efficiency gain is not convincing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer:\n\n**W1 & Q1:the encoder-decoder module for regression and the vocab for the decoder.**\n\nResponse: The confusion may come from the naming. Although the encoder-decoder architecture was originally used in machine translation (discrete vocab), common encoder-decoder architectures, e.g., U-Net, transformer, have been widely used in regression problems such as temperature and precipitation prediction [1-4]. In a general sense, an encoder learns a latent feature representation that encodes the structural dependency within input data. A decoder predicts the target variable from encoded features. There is no restriction that the framework cannot be applied to regression. We hope this can address the concern and please comment if any questions.\n\n[1] Autoformer: Decomposition Transformers with Auto-Correlation for Long-Term Series Forecasting. NeurIPS 2021\n\n[2] Informer: Beyond Efficient Transformer for Long Sequence Time-Series Forecas. AAAI 2021\n\n[3] Spatiotemporal Swin-Transformer Network for Short Time Weather Forecasting. CIKM 2021\n\n[4] NumHTML: Numeric-Oriented Hierarchical Transformer Model for Multi-Task Financial Forecasting. AAAI 2022\n\n**W2: the O(NlogN) complexity has a large big-O constant**\n\nResponse: We agree that sparse matrix operations can be less efficient in GPU penalization when compared to normal dense matrix operations. However, there are already significant efforts in GPU optimization of sparse matrix multiplication (e.g., cuSPARSEt library). In our HST model, the large big-O constant is not a dominant factor, as shown in experiment results in Figure 4. For 400 points, our HST reduces the time costs from 25 mins to 5 mins. One note is that the time cost for the default all-pair transformer when the input size is 1000-point is not measurable because out-of-memory problem (please see our response below for this question). In contrast, our model can run 1000 points within twenty minutes for one epoch.  \n\n**W3: not enough details on models and experiments.**\n\nResponse: In fact, we had provided all the model hyper-parameters and training details in our supplementary material section 7.1. The description is quoted as following: \nOur HST model used six spatial attention layers in both the encoder and the decoder. We use 8 heads and the latent representation embedding dimension of each head was 64, and the quadtree leaf node size threshold was 20 by default. The batch size was $512$. For the training process, we used the MSE loss with a decaying learning rate that reduced the learning rate by half if the validation loss did not improve over five epochs (with an initial \nlearning rate of $10^{-4}$ and a minimum rate of $10^{-7}$). We also used early stopping with a patience of 10 epochs and a maximum of 50 epochs. The optimizer was Adam with $\\beta_1=0.9$ and $\\beta_2 = 0.98$. The $L_2$ regularization weight was $10^{-4}$. \n\n\n**W4: 1K samples within the capability of vanilla transformer**\n\nResponse: \nThis is a good question. First, we want to clarify that even 1K points will exceed the capacity of vanilla (all-pair) Transformer on A100 GPU. This is because the dominating GPU memory costs come from intermediate tensors in forward propagation. Please see an example below on the calculation of GPU memory cost with 1K point samples for both vanilla (all-pair) transformer and our HST transformer. \n\nMemory cost estimation of vanilla transformer: \n\nAssume $L = 1,000$ points, $H = 8$ attention heads, $N = 6$ block, the dimension of each head $D = 64$, and the batch size $B= 512$. The model stores results from forward attention activation (including self-attention, dense layer, and skip connection) ($M_{a}$) and backward gradient ($M_{g}$), and model parameters ($M_{m}$). The dominant factor is the attention layer. For each attention layer, given query and key tensors $Q,K \\in R^{L*D}$, the attention matrix from $QK$ multiplication will result in a tensor with the shape of $[L, L]$ for each attention head and each sample within the minibatch. Thus, the total memory cost for all heads in one minibatch in one attention layer is: $ M_{a} = BHL^{2} \\times (float32 \\ bytes)$. Multiplying this by $N$ attention layers, the cost is $4NBHL^{2} \\ bytes= 6\\times 512\\times 8\\times 4\\times 1000^{2} bytes \\approx  94 GB$. This exceeds A100 memory capacity (80GB), as shown in our results in Figure 4.\nIn contrast, our HST model significantly reduces the memory costs.  The reason is that we approximate point samples far away from the query points by coarse representation (in quadtree internal nodes) in the attention computation. The further away points are from the query points, the coarser cells we use. \n\nHST memory cost estimation: \n\nAssume same setup ($L = 1K$ points, $H = 8$ attention heads, $N = 6$ block, the dimension of each head $D = 64$, and the batch size $B= 512$) and a quadtree height $h=10$ and minimum leaf node size $T=100$. HST has the same number of query points. For each query point in a leaf node, the key set includes all the points in the same leaf node, as well as only a subset of internal quadtree nodes (the leaf’s siblings, its parent’s siblings, and its grandparent’s siblings, …, till the root node). This is illustrated in the red boxes in Fig.1 in global rebuttal box. In this example, the key set is reduced from 64 to 9 nodes. In general, the key set is reduced from $L$ to $T+3h$. Thus, the cost of $N$ attention layers is reduced from $NHL^{2}$ to $NHL(T+3h)$. The sparse indicator matrix for selective keyset attention calculation in HST has a shape of $\\mathbf{I} \\in R^{L \\times 2L}$, and its non-zero element count is $L(T+3h)$, thus the memory cost is $BNL(T+3h)\\times float32 \\ bytes $. Thus, the total memory cost is $2BNL(T+3h)\\times float32 \\ bytes\n= 2\\times 512\\times 6\\times 8\\times 1000\\times 130\\times 4 \\ bytes \\approx 24 GB$. Similarly, the time costs of HST are also lower than the vanilla transformer due to less attention pairs. See our responses above.\n\n"}, {"review_id": "Pk2hEiVUu6", "reviewer": "Reviewer_J8MW", "summary": "This paper proposes a hierarchical transformer to model a large number of irregular point samples in continuous space. This is achieved by a quad-tree hierarchy which could learn the multi-scale spatial representation. So the long-range interactions are recorded.\n\nThe experiment is performed in three real-world dataset (two water quality datasets and one sea-surface temperature prediction) and one simulation dataset.", "questions": "Could the author clarify how the adopted quadtree technique in the paper could solve uniformly distributed cases?\n\nLink to the weakness above, as Darcy flow and Sea Surface datasets are set, could the author claim why do we have to feed all point samples in a single feed? In other words, why don't we solve this task in all three datasets by following a classic 'classification setting'?\n\nCompared with the all-pair transformer, pros and cons?\n\nCould the reviewer know the number of parameters for all-pair transformer and other baselines in Table 2?\n\nAny guidance in choosing hyper-parameters for designing quadtree if needed?\n\nAs in Weakness, my main concern is first two questions.", "limitations": "NA", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This submission tries to solve the large number of issues by using quadtree. In my understanding, the quadtree is a classic technique for downsampling points so that the point number will be reduced to an affordable level.  And the quadtree could model the long-range interactions without the distribution assumption.\n\nThe proposed method reduced the computational cost to $O(N\\log N)$ or $O(NM)$.\n\nThe approach presentation is detailed and clear.", "weaknesses": "The adopted quadtree technique seems could only work well in the sparse case. As a result, the proposed method is not a general solution for all kinds of large numbers of point samples.\n\nThe experiment setting is not convincing in the reviewer's current understanding. In particular, Red Tide and Turbidity datasets have a large number of data samples, but it seems only to have one 'epoch'. In other words, The point number is large but the total data size seems very small. And for Darcy flow and Sea Surface, it seems each image only have 400 point samples, it is not a large number set in my understanding, the all-pair transformer could cover it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer:\n\n**W1 & Q1: how the quadtree technique could solve uniformly distributed cases.**\n\nThank you for the comment. In fact, our HST model does not require samples to be sparsely distributed for efficiency gain. For uniformly distributed point samples, the quadtree will be a balanced (complete) tree.  We can illustrate how our model works in this case with an example. Please see the example in Fig.1 of global rebuttal pdf (the topmost box). There are 8*8 point samples evenly distributed in 2D. Assume the minimum leaf node size is 4. The height of the quadtree is $h =log_4(64/4) = 2$, as Fig.1(b) shows. Consider the query sample of point 1, the key node set include 3 remaining points in the same leaf node (points 2-4), 3 internal nodes $(r_2^2 ... r_4^2)$, 3 internal nodes $(r_2^1 ... r_4^1)$. Thus, there are only 9 keys, much lower than the total number of points (64). In fact, our efficiency gain is better for uniformly distributed case, as the depth of the tree is lower. This significantly reduces the time costs.\n\n**W2 & Q2: the experiment with 400 point samples.**\n\nThis is a valid concern. First, we chose to randomly select 400 input points from each $64\\times 64$ image in the Darcy flow and Sea Surface dataset because drawing much denser input samples (e.g., over 1000) makes the problem trivial because of strong spatial autocorrelation. In this case, even a simple nearest-neighbor estimator can work well.  More importantly, even 400-point input samples is not as easy as it seems for the vanilla transformer model due to the high memory and time costs. The GPU memory costs for vanilla all-pair transformer with 1K point samples exceed the capability of A100 GPU due to the large size of intermediate tensors. This is confirmed in our experiment results (Figure 4 in the paper). \n\nFor detailed calculation, please see our calculation below.  Second, we can run all the point samples in a single feed, but the GPU memory cost is very high. We provide a memory consumption analysis below. Third, even with 400 point samples, our model can still improve the computational efficiency as Fig.4 in our paper shows. The computation time is reduced from 25 minutes to 5 minutes for one epoch. Our model can have slight improvement compared to all-pair transformer as Table 2 shows. \n\nMemory cost estimation of vanilla transformer: \n\nAssume $L = 1,000$ points, $H = 8$ attention heads, $N = 6$ block, the dimension of each head $D = 64$, and the batch size $B= 512$. The model stores results from forward attention activation (including self-attention, dense layer) ($M_{a}$) and backward gradient ($M_{g}$), and model parameters ($M_{m}$). The dominant factor is the attention layer. For each attention layer, given query and key tensors $Q,K \\in R^{L*D}$, the attention matrix multiplication will result in a tensor with the shape of $[L, L]$ for each attention head and each sample within the minibatch. Thus, the total memory cost for all heads in one minibatch in one attention layer is: $ M_{a} = BHL^{2} \\times (float32 \\ bytes)$. Multiplying this by $N$ attention layers, the cost is $4NBHL^{2} \\ bytes= 6\\times 512\\times 8\\times 4\\times 1000^{2} bytes \\approx  94 GB$. This exceeds A100 memory capacity (80GB), as shown in our experiment results in Figure 4.\nIn contrast, our HST model significantly reduces the memory costs to only 24 GB.  The reason is that we approximate point samples far away from the query points by coarse representation (in quadtree internal nodes) in the attention computation. The further away points are from the query points, the coarser cells we use in approximation.  Due to limited space, we do not provide  estimation here. The detailed derivation can be found in the rebuttal for reviewer 2.\n\n**Q3: pros and cons compared with all-pair transformer.**\n\nThis is a great question. We provide some analysis on the pros and cons of our model compared to the all-pair transformer. \nThe pros:\n1.\tAble to capture multi-scale effect with the hierarchical quad-tree structure: Our model learns a hierarchical spatial representation within the quadtree. The multi-scale effect is important for many scientific problems. This is reflected by the improved prediction accuracy in experiments.\n2.\tOur HST reduces the time and memory costs when compared with the all-pair transformer. Please see our analysis above.\n3.\tWe added uncertainty quantification branch to estimate the model’s confidence.  \n\nThe con:\n1.\tOur HST makes a trade-off between efficiency and spatial granularity when calculating attention between points. We approximate the set of key points by coarse cells in the quadtree based on distance to the current query point. \n\n**Q4: the number of parameters of baselines:** \n\nThanks for the great suggestion. We will add the number of parameters for the baselines and our model in Table 2. The numbers are as follows.  GP: 2 (RBF kernel parameters), Deep GP: 0.24 million, Spatial GCN: 0.79 million, All-pair transformer: 3.81 million, HST model: 3.81 million, NodeFormer: 2.2 million, Galerkin Transformer: 17.6 million\nNote that the large number of parameters in Galerkin Transformer is because we used the default code setup without tunning. We could set up the number of attention layers the same as other transformer models (six encoder-decoder layers). In this case, the number of Galerkin transformer parameters would be 4.4 million. \n\n**Q5. the guidance on choosing hyper-parameters for designing quadtree.**\n\nWe provide the quadtree hyperparameters in our supplementary material Section 7.1 and the sensitivity analysis. We chose the minimum leaf node size of 20 in experiments based on prior knowledge about the spatial dependency structure. The size of the leaf node can be determined based on the spatial homogeneity and efficiency needs. Roughly, the more spatially homogeneous the input point samples are and the more efficiency gain we need, the bigger leaf node cells we need to choose. "}, {"review_id": "CorZIi6AB8", "reviewer": "Reviewer_Y8FJ", "summary": "The authors proposed a hierarchical spatial transformer model for many irregular point samples in continuous spatial domain. Compared with existing methods, the proposed method can model implicit spatial dependency across irregular samples and in multiple scales in continuous space. The proposed model uses a quad-tree hierarchy to conduct efficient attention operations by approximating distant points in a coarse resolution. The model also includes an uncertainty quantification module to capture the varying prediction confidence in different input sample density. Evaluations on several datasets confirm that our method outperforms multiple baselines in accuracy and effectiveness of uncertainty quantification.", "questions": "• How are the thresholds of uncertain/certain predictions determined in UQ metric?", "limitations": "No limitations on societal impacts.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "Overall, this is a solid paper with novel technical contributions and extensive experimental evaluations. \n\n• The paper solves a significant problem of spatial representation learning for irregular samples in continuous space. The problem has many significant applications in environment sustainability. It is also important for learning surrogate models to speed up numerical simulations.\n\n• The technical novelty of the proposed hierarchical attention architecture is strong. The model uses a quad-tree to learn latent representation of point samples in different subareas in a multi-scale hierarchy. The attention layers use quad-tree to make spatial approximation to overcome the computational bottleneck. The model also has an uncertainty quantification module.\n\n• The proposed method has shown promising results on both real-world and synthetic datasets. The evaluation is solid with sensitivity analysis and computational experiments.", "weaknesses": "Although the experimental evaluation is extensive, it will be helpful to add some more on how the thresholds of uncertain/certain predictions are determined in UQ metric.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Question: how are the thresholds of uncertain/certain predictions determined in UQ metric?\n\nResponse: Thank you for your positive feedback. About how to choose thresholds of uncertainty in the UQ metric, we discussed them in the supplementary materials in section 7.1. Here we briefly describe how we choose the UQ threshold. We used a validation dataset for determining accurate/inaccurate and certain/uncertain thresholds. Specifically, the threshold for accuracy is determined by the average MAE loss on validated samples. Then the threshold for uncertainty is determined by the average uncertainty score for accurate and inaccurate prediction respectively. The detailed equation is provided in our supplementary materials section 7.1. "}]

Author rebuttal:
This PDF contains a figure illustration of evenly distributed input points. It illustrates how our model can reduce computational costs in this case by selecting a subset of keys in attention calculation.

Meta-review: This paper introduces a transformer architecture for spatial modeling problems with irregularly spaced sample points. The authors use a quad tree structure to reduce the computational requirements of the transformer, and introduce a decoder mechanism to quantify uncertainty. Overall, the proposed method, which contains many novel contributions, can be applied to many important application domains. A few concerns about the method include: (1) it will only be effective in certain settings (irregularly spaced points, not too many points), (2) it will have high computational costs in these other settings, and (3) the sparse operations cannot be significantly accelerated by GPUs. That being said, the efficacy of the proposed method, as well as the novel approach of combining quad trees with transformer modules, will be of great interest to the NeurIPS community and therefore I recommend acceptance.
Decision: Accept (poster)