SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $n$  points and the maximum number of points in each quadtree node is  $M$  (quadtree threshold). We analyze two scenarios. If the quadtree is completely balanced, the tree depth and the key set size for every external node is  $O(\log \frac{n}{M})$ . Thus total computation cost is  $O(n \cdot (M + \log \frac{n}{M}))$ . In the worst case, the spatial points are highly nonuniform, then the quadtree structure will be highly imbalanced. The extreme tree depth is  $O\left(\frac{n}{M}\right)$ , thus the attention computation is  $O(n \cdot \left(\frac{n}{M} + M\right))$ . 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 at every level. For instance, for water quality monitoring, the sensors are sparsely distributed across a broad area to monitor multiple locations. In practical applications, it takes  $O(n \cdot (\log \frac{n}{M} + M))$  complexity. This is validated in our experiments.
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  $d$ . The memory costs of HST are dominated by the hierarchical spatial attention layer. Its memory cost is  $O(B \cdot h \cdot d \cdot n \cdot |\mathcal{K}_i|)$  per layer, where  $|\mathcal{K}_i| = \log \frac{n}{M} + M$  for relatively balanced quadtree.
4.3 Decoder: inference on a new point with uncertainty quantification
Model inference (prediction): Given a test sample  $\mathbf{o}_t = (\mathbf{x}_t, \mathbf{s}_t)$ , the decoder module predicts  $y_t$  based on learned spatial representations of quadtree nodes  $\mathbf{H}$ . Similar to the encoder, we use cross-attention between the test sample and its corresponding key node set in the quadtree as in Equation 4. As shown in Figure 1, 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 a 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
method for UQ in the continuous spatial domain is the Gaussian process (GP) [62, 63]. The UQ of GP is based on Equation 5, where  $\mathbf{c}_0 \in \mathbb{R}^{1 \times n}$  is the covariance vector between the test sample and all input point samples,  $\mathbf{C} \in \mathbb{R}^{n \times n}$  is the covariance matrix for input samples, and  $\sigma_0^2$  is the self-variance. In GP, the covariance  $\mathbf{C}$  is computed with a kernel function that reflects the location proximity [64]. Although a 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.
$$
\sigma_ {t} ^ {2} = \sigma_ {0} ^ {2} - \mathbf {c} _ {0} ^ {T} \mathbf {C} ^ {- 1} \mathbf {c} _ {0} \tag {5}
$$
Our proposed spatial transformer framework can be considered a generalization of a GP model. We use the dot-product cross-attention weights to approximate the covariance vector  $\mathbf{c}_0$  between the test location to all points, which reflects the dependency among point sample locations based on their non-linear embeddings, i.e.,  $\mathbf{c}_0 = \mathbf{H}\mathbf{q}_t$ , where  $\mathbf{q}_t$  is the embedded query vector for the test sample. However, this idea is insufficient for the approximation of the entire covariance matrix  $\mathbf{C}$  across all points, since the inverse computation of a full covariance matrix is very expensive. To overcome this challenge, we propose to directly approximate the precision matrix  $(\mathbf{C}^{-1})$  based on an indicator matrix of selective key sets for all queries  $\mathbf{S} \in \mathbb{R}^{N \times n}$ , 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  $\mathbf{S}^T \cdot \mathbf{S} / T_u^2$  to approximate the precision matrix  $\mathbf{C}^{-1}$ , where  $T_u^2$  is a hyper-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  $\mathbf{S}^T \cdot \mathbf{S} = \sum_i \mathbf{s}_i \mathbf{s}_i^T$  ( $\mathbf{s}_i$  is a column of  $\mathbf{S}$ ), we can see that  $\mathbf{S}^T \cdot \mathbf{S}$  is a summation of several sparse block-diagonal 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 6, where  $\mathbf{K}_t = \mathbf{SH}$  is the key matrix corresponding to the selective key set of the test sample.
$$
\begin{array}{l} u _ {t} = \sigma_ {0} ^ {2} - \left(\mathbf {q} _ {t} ^ {T} \mathbf {H} ^ {T}\right) \mathbf {S} ^ {T} \mathbf {S} (\mathbf {H} \mathbf {q} _ {t}) / T _ {u} ^ {2} \\ = \sigma_ {0} ^ {2} - \mathbf {q} _ {t} ^ {T} (\mathbf {S H}) ^ {T} (\mathbf {S H}) \mathbf {q} _ {t} / T _ {u} ^ {2} \\ = \sigma_ {0} ^ {2} - \left(\mathbf {K} _ {t} \mathbf {q} _ {t}\right) ^ {T} \left(\mathbf {K} _ {t} \mathbf {q} _ {t}\right) / T _ {u} ^ {2} \tag {6} \\ \end{array}
$$
Note that our 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 a 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 a 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 on spatial location without using the explanatory feature [62]. The prediction variance was used as the uncertainty measure. Deep Gaussian Process (Deep GP): We implemented a 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 [58] method to quantify the prediction uncertainty. Multipole graph neural operator (MGNO): It belongs to the family of neural operator model [32]. NodeFormer: It is an efficient graph transformer model for learning the implicit graph structure [26]. We used the code from its official website. Galerkin Transformer: It uses softmax-free attention mechanism and achieve linearized transformer [35]. We quantify the prediction uncertainty based on MC-dropout. Hierarchical Spatial Transformer (HST): This is our proposed method. We implemented it with PyTorch.
The prediction performance is evaluated with mean square error (MSE) and mean absolute error (MAE). The evaluation of UQ performance is challenging due to the lack of ground truth of uncertainty.
UQ evaluation metrics: The quantitative evaluation metrics for UQ performance is Accuracy versus Uncertainty  $(AvU)[67]$ . We set accuracy uncertainty thresholds  $T_{ac}, T_{au}$  to group prediction accuracy and uncertainty into four categories as Table  $\square$  shows.  $n_{\mathrm{AC}}, n_{\mathrm{AU}}, n_{\mathrm{IC}}, n_{\mathrm{IU}}$  represent the number of samples in the categories AC, AU, IC, IU, respectively. As Equation  $\square$  shows,  $AvU$  measures the percentage of two categories AC and IU. A reliable model should provide a higher  $AvU$  measure  $(AvU \in [0,1])$ . The details on how we choose the thresholds are provided in the supplementary materials.
$$
A v U = \frac {n _ {\mathrm {A C}} + n _ {\mathrm {I U}}}{n _ {\mathrm {A C}} + n _ {\mathrm {A U}} + n _ {\mathrm {I C}} + n _ {\mathrm {I U}}} \tag {7}
$$
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, we computed  $AvU_{A}$  for accurate predictions and  $AvU_{I}$  for inaccurate predictions with the following equations:
$$
A v U _ {A} = \frac {n _ {\mathrm {A C}}}{n _ {\mathrm {A C}} + n _ {\mathrm {A U}}}, A v U _ {I} = \frac {n _ {\mathrm {I U}}}{n _ {\mathrm {I C}} + n _ {\mathrm {I U}}} \tag {8}
$$
In our evaluation, we computed the harmonic average of  $AvU_{A}$  and  $AvU_{I}$ :  $AvU = \frac{2*AvU_{A}*AvU_{I}}{AvU_{A} + AvU_{I}}$  to penalize the extreme cases instead of the arithmetic average.
Table 1: Accuracy versus Uncertainty (AvU)
Uncertainty
Certain	Uncertain
Accuracy	Accurate	Accurate Certain (AC)	Accurate Uncertain (AU)
Inaccurate	Inaccurate Certain (IC)	Inaccurate Uncertain (IU)
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 a sliding window with a sequence length 400 to generate 103700 inputs. It is split into training, validation, and test sets with a ratio of  $7:1: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 [33] contain 100 simulated images with  $241 \times 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 [70]. We subsampled 400 point samples from the grid pixels.
Table 2: Comparison of model performance on two real-world datasets and one simulation dataset
Model	Red tide	Turbidity	Darcy flow
MSE	MAE	MSE	MAE	MSE	MAE
GP	7.42 ± 0.25	2.55 ± 0.18	0.42 ± 0.02	0.46 ± 0.03	0.19 ± 0.03	0.33 ± 0.04
Deep GP	6.23 ± 0.42	2.32 ± 0.24	0.35 ± 0.04	0.42 ± 0.06	0.18 ± 0.03	0.31 ± 0.05
Spatial GNN	5.68 ± 0.07	2.19 ± 0.04	0.34 ± 0.02	0.46 ± 0.03	0.15 ± 0.02	0.26 ± 0.04
All-pair transformer	5.30 ± 0.12	1.98 ± 0.07	0.31 ± 0.03	0.35 ± 0.04	0.10 ± 0.02	0.22 ± 0.03
MGNO	5.41 ± 0.26	2.04 ± 0.10	0.32 ± 0.03	0.38 ± 0.05	0.12 ± 0.03	0.27 ± 0.03
NodeFormer	5.34 ± 0.10	2.05 ± 0.06	0.32 ± 0.04	0.38 ± 0.04	0.16 ± 0.03	0.29 ± 0.04
GalerkinTransformer	5.44 ± 0.17	2.11 ± 0.08	0.34 ± 0.03	0.42 ± 0.05	0.11 ± 0.02	0.25 ± 0.03
HST (Our method)	5.25 ± 0.11	1.97 ± 0.07	0.30 ± 0.03	0.35 ± 0.04	0.11 ± 0.02	0.23 ± 0.03
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
model improved the GP performance from 7.42 to 6.23 by leveraging the neural network feature representation learning capability as well as considering point sample correlation in 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-based 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 message-passing mechanism captures the long-range dependency structure. Our framework performs better because of the awareness of multi-scale spatial relationships, which is crucial for point samples' inference. Additionally, the performance of all-pair attention is better than the fix-graph GNN model because of the long-range modeling. However, the vanilla transformer models lack the hierarchical attention structure, which may cause an imperfect attention weight score. In contrast, our model performed best with the lowest MAE of 5.25 because it modeled the interaction among point samples in the continuous multi-scale space efficiently. We can see similar 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 are shown in Table 2. We can see our 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-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 a sensitivity analysis of our model to various hyperparameters, including the quadtree leaf node size threshold  $M$ , spatial position encoding length scale  $\sigma$ , the number of attention layers, and the embedding dimension on the red-tide dataset. The results are summarized in Figure 5 and the detailed analysis is in the Supplementary material. We can see that our model is generally stable with changes of the hyper-parameters.
(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  $AvU$  in Equation 7 to evaluate the performance of UQ for our proposed HST model versus baseline methods and the results were summarized in Table 3. The numbers in the table correspond to the number of samples in four categories:  $n_{\mathrm{AC}}, n_{\mathrm{AU}}, n_{\mathrm{IC}}, n_{\mathrm{IU}}$ .
We can see the base GP model had good uncertainty estimation for inaccurate predictions (65% are uncertain) but for the accurate predictions, 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 improved the accurate prediction confidence based on the feature representation learning but still performed worse than other models and gave a 0.38 AvU score. The spatial GNN model was
Table 3: Comparison on uncertainty quantification performance on Red tide dataset
Model	Accuracy	Uncertainty	AvUa/ AvUc	AvU
GP	Certain	Uncertain	0.33
Accurate	2726	9144	0.23
Inaccurate	3252	4919	0.60
Deep GP	Certain	Uncertain	0.38
Accurate	3316	8509	0.28
Inaccurate	3094	5122	0.62
Spatial GNN	Certain	Uncertain	0.40
Accurate	3152	3949	0.44
Inaccurate	8355	4585	0.35
Galerkin Transformer	Certain	Uncertain	0.38
Accurate	4515	4925	0.47
Inaccurate	7166	3435	0.32
HST	Certain	Uncertain	0.54
Accurate	4342	5302	0.46
Inaccurate	3874	6503	0.65
(a) Computation cost per epoch for 64 training samples.
(b) 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 6: Computation cost analysis
more confident in the accurate prediction, but
the model was over-confident in the inaccurate prediction, thus resulting in a lower  $AvU$  score (0.44). In contrast, our uncertainty model improved both the accurate prediction confidence and inaccurate prediction uncertainty, and improved the overall  $AvU$  score to 0.49, because it modeled the uncertainty coming from both feature space and sample density simultaneously. 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 a 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^{2}$  to  $10^{6}$ . The training batch size is 1. 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 is shown in Figure 6(a). When the number of point samples increases to  $50\mathrm{K}$ , 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 to  $100\mathrm{K}$ . However, our 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 decreases and then increases. 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 to the hidden feature dimension.
6 Conclusion and Future Work
This paper proposes a novel hierarchical spatial transformer model that can model implicit spatial dependency on a 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 a quadtree structure. In order to reflect the varying degrees of 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.
This material is based upon work supported by the National Science Foundation (NSF) under Grant No. IIS-2147908, IIS-2207072, CNS-1951974, OAC-2152085, the National Oceanic and
Atmospheric Administration grant NA19NES4320002 (CISESS) at the University of Maryland, and Florida Department of Environmental Protection. Shigang Chen's work is supported in part by the National Health Institute under grant R01 LM014027.
[62] Sudipto Banerjee, Bradley P Carlin, and Alan E Gelfand. Hierarchical modeling and analysis for spatial 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 on Knowledge Discovery and Data Mining, pages 554-564, 2022.
[64] Sudipto Banerjee, Bradley P Carlin, and Alan E Gelfand. Hierarchical modeling and analysis for spatial data. Chapman and Hall/CRC, 2003.
[65] Volodymyr Kuleshov, Nathan Fenner, and Stefano Ermon. Accurate uncertainties for deep learning using calibrated regression. In International conference on machine learning, pages 2796-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 arXiv:1707.02476, 2017.
[67] Ranganath Krishnan and Omesh Tickoo. Improving model calibration with accuracy versus uncertainty optimization. Advances in Neural Information Processing Systems, 33:18237-18248, 2020.
[68] [dataset] NASA Goddard Space Flight Center. Moderate-resolution 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 Bezenac, Arthur Pajot, and Patrick Gallinari. Deep learning for physical processes: Incorporating prior scientific knowledge. Journal of Statistical Mechanics: Theory and Experiment, 2019(12):124009, 2019.