SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Wenchong He Zhe Jiang* Tingsong Xiao Zelin Xu Shigang Chen
Department of Computer & Information Science & Engineering University of Florida
{whe2, zhe.jiang, xiaotingsong, zelin.xu, sgchen}@ufl.edu
Ronald Fick Miles Medina Christine Angelini
Center for Coastal Solutions
{rfick, miles.medina}@ufl.edu, christine.angelini@essie.ufl.edu
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
Transformers are widely used deep learning architectures. Existing transformers are largely designed for sequences (texts or time series), images or videos, and graphs [1234567]. This paper proposes a novel transformer model for massive (up to a million) points in continuous space. Given a 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.
Learning transformers for continuous-space points has broad applications. In environmental sciences, researchers are interested in fusing remote sensing spectra with in-situ sensor observations at irregular
*Contact author
sample locations to monitor coastal water quality and air quality [8, 9, 10]. In scientific computing, researchers learn a neural network surrogate to speed up numerical simulations of particle-laden flow [11] 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, 14]. In location-based service, people are interested in analyzing massive spatial point data (e.g., POIs, trajectories) to recommend new locations [15, 16].
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-uniformly 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 a 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 a high computational cost.
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.
2 Problem Statement
A spatial point sample is a data sample drawn from 2D continuous space, denoted as  $\mathbf{o_i} = (\mathbf{x}(\mathbf{s}_i),y(\mathbf{s}_i),\mathbf{s}_i)$ , where  $1\leq i\leq n$ ,  $\mathbf{s}_i\in \mathbb{R}^2$  is 2D spatial location coordinates (e.g., latitude and longitude),  $\mathbf{x}(\mathbf{s}_i)\in \mathbb{R}^{m\times 1}$  is a vector of  $m$  non-spatial explanatory features, and  $y(\mathbf{s}_i)$  is a target response variable  $(y(\mathbf{s}_i)\in \mathbb{R}$  for regression,  $y(\mathbf{s}_i)\in \{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 as a continuous-space function  $y: \mathbb{R}^2 \to \mathbb{R}$  or  $\{0,1\}$ . Given a set of point sample observations  $\mathcal{O} = \{(\mathbf{x}(\mathbf{s}_i), y(\mathbf{s}_i), \mathbf{s}_i)\}_{i=1}^n$ , where  $\mathbf{s}_i$  is irregularly sampled in 2D, our model learns the contiguous function  $y$  that can be evaluated at any new spatial points  $\hat{\mathbf{s}} \notin \{\mathbf{s}_i\}_{i=1}^n$ . We define our model as learning the mapping from the observation samples to the continuous target variable function  $y(\hat{\mathbf{s}}) = f_{\theta}(\mathcal{O}, \mathbf{x}(\hat{\mathbf{s}}), \hat{\mathbf{s}})$ . Thus, we formulate our problem as follows.
Input: Multiple training instances  $\mathcal{D} = \{\mathcal{O}_j\}_{j=1}^L$  (sets of irregular points in continuous 2D space).  
Output: A spatial transformer model  $f: \{y(\hat{\mathbf{s}}), u(\hat{\mathbf{s}})\} = f(\mathbf{x}(\hat{\mathbf{s}}), \hat{\mathbf{s}}, \mathcal{O}_j)$  for any  $j \in [1, \dots, L]$ , where  $\hat{\mathbf{s}}$  is any new sample location for inference, and  $\mathbf{x}(\hat{\mathbf{s}})$  and  $y(\hat{\mathbf{s}})$  are the explanatory features and output target variable for the new sample, respectively, and  $u(\hat{\mathbf{s}})$  is the uncertainty score corresponding to the prediction  $y(\hat{\mathbf{s}})$ .
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 to supervise model training, we can construct the training instances by removing a single point sample from each set  $\mathcal{O}_j$  as the new sample location and use its target variable as the ground truth.
- Transformer models: Attention-based transformers are widely used deep learning architecture for sequential data (e.g., texts, time series) [1], images or videos [5, 6], graphs [7]. One main advantage of transformers is 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
techniques have been developed to address this bottleneck. For sequential data, sparsity-based methods [17, 18] and low-rank-based methods [19, 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-based 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 graph data, sampling-based [26, 27, 28] and spectral-based [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.
- Numerical Operator learning in continuous space: Neural operator learning aims to train a neural network surrogate as the solver of a family of partial differential equation (PDE) instances [31]. The surrogate takes initial or boundary conditions and predicts the solution function. Existing surrogate models include deep convolutional neural networks [2], graph neural operators [3, 32], Fourier neural operators [33, 34], DeepONet [31], NodeFormer [26, 35] and vision transformers [36, 37, 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 [42], but their neural networks are limited by only taking each sample'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) [43, 44] are often used for regular grids (e.g., satellite images, and global climate models) [45, 46], and graph neural networks (GNNs) [47] are used for irregular grids (e.g., meshes with irregular boundaries) [48, 49, 50] or spatial networks (e.g., river or road networks) [51, 52, 53]. However, CNNs and GNNs only capture local spatial dependency without long-range interactions. In recent years, the transformer architecture [1, 54] has been widely used for spatial representation learning with long-range dependency, but existing transformer models are often designed for regular grids (images, videos) [5, 6, 23, 24, 25, 22] and thus cannot be directly applied to irregular point samples in continuous space.
4 Approach
This section introduces our proposed hierarchical spatial transformer (HST) model. Figure [1] shows the overall architecture with an encoder and a decoder. The encoder learns a 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 to approximate the representation of faraway key points by the representation of coarse quadrtree cells. The decoder makes inferences at a new point location by traversing the quadrtree 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 to feature noise and point sparsity. Note that our model differs from existing tree-based transformers [55, 56] for images or videos, as these methods require a regular grid and cannot be directly applied to irregular points in continuous space.
4.1 Multi-resolution representation learning within a quadtree hierarchy
Learning a multi-resolution latent representation of point samples is non-trivial due to the non-uniform (irregular) distribution of points in continuous space (instead of a 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 a spatial pooling within a 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  $\phi : \mathbb{R}^2 \to \mathbb{R}^{\frac{d}{2}}$  from the 2D continuous space to a  $\frac{d}{2}$ -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.e.,
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 are insufficient. We propose to use a multi-dimensional continuous space position encoding [57] as follows,
$$
\phi (\boldsymbol {s}) \approx [ \cos (\Omega_ {1} \boldsymbol {s}), \sin (\Omega_ {1} \boldsymbol {s}), \dots , \cos (\Omega_ {\frac {d}{2}} \boldsymbol {s}), \sin (\Omega_ {\frac {d}{2}} \boldsymbol {s}) ], \tag {1}
$$
where  $d$  is the encoding dimension,  $\Omega_{i}\sim \mathcal{N}(\mathbf{0},\Sigma)$  is a 1 by 2 projection matrix following an i.i.d. Gaussian distribution with a standard deviation  $\sigma$ . The advantage of this encoding is that it satisfies the following property,  $< \phi (s_1),\phi (s_2) > \approx k(||s_1 - s_2||) = \exp \bigl \{-(s_1 - s_2)\bigr)^T\Sigma^{-1}(s_1 - s_2)$ . Here the hyperparameter  $\boldsymbol{\Sigma}$  controls the spatial kernel bandwidth. Next, we use  $\phi (\pmb {o}_i)$  to denote the positional encoding  $\phi (s_i)$  for consistency.
Spatial representation: We propose an initial spatial representation of individual points by concatenating its continuous-space positional encoding and a non-spatial feature embedding, i.e.,
$$
\mathbf {h} \left(\boldsymbol {o} _ {i}\right) = \left[ \psi \left(\boldsymbol {o} _ {i}\right); \phi \left(\boldsymbol {o} _ {i}\right) \right] \tag {2}
$$
where  $\phi (\pmb {o}_i)$  is a positional encoding,  $\psi (\pmb {o}_i) = \mathbf{W}\cdot [\mathbf{x}_i;y]$  is the non-spatial embedding,  $\mathbf{W}\in \mathbb{R}^{\frac{d}{2}\times (m + 1)}$  is the embedding parameter matrices, and  $d$  is the dimension of concatenated representation. We denote the representation of all point samples (external quadtree nodes) in a matrix  $\mathbf{H}_o = [h(\pmb {o}_1),\dots,h(\pmb {o}_n)]^T\in \mathbb{R}^{n\times d}$ , where each row is the representation of one point.
4.1.2 Spatial representation of coarse cells (quadtree internal nodes)
To learn a multi-resolution representation of non-uniformly distributed points in continuous space, we propose a spatial pooling operation within a quadtree hierarchy. A quadtree is a spatial index structure designed for a large number of points in 2D continuous space. It recursively partitions a rectangular area into four equal cells until the number of points within a cell falls below a maximum threshold. In a quadtree, an internal node at a higher level represents an area at a coarser spatial scale and nodes at different levels provide a multi-resolution representation. Another advantage of using a quadtree is that it 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  $\mathcal{O} = \{\mathbf{o}_i\}_{i=1}^N$  in continuous space, we construct a quadtree  $\mathcal{T}$ . The quadtree has two kinds of node sets: an external node set  $\mathcal{E}$ , which corresponds to the observed spatial point samples, and an internal node set  $\mathcal{I}$ , which represents the spatial cells. A quadtree has  $L$  levels, and all the nodes in level  $l$  form a set  $\mathcal{R}_l = \{r_1^l, \dots, r_{k_l}^l\}$ , where  $r_j^l$  is the  $j$ -th node at level  $l$ , and  $k_l$  is the total number of nodes in level  $l$ . Given one node  $r_j^l$ , we represent its sibling node set as  $\mathcal{S}(r_j^l)$  (nodes on the same hierarchical level under the same parent node), the ancestor node set as  $\mathcal{A}(r_j^l)$  (nodes on the path from the node  $r_j^l$  to the root). We denote the spatial representation of the node  $r_j^l$  as  $\mathbf{h}(r_j^l)$  and  $l \in \{1, \dots, L\}$  and  $j \in \{1, \dots, k_l\}$ .
For example, in Figure 2(a), there are 11 input point samples with denser point distribution in the upper left corner and the lower right corner. Assuming a maximum leaf node size of two samples, the
(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 0 (the root node). Level 1 has four internal nodes, i.e.,  $\mathcal{R}_1 = \{r_1^1,\dots,r_4^1\}$ . For instance,  $r_1^1$  corresponds to the largest quad cell in the upper left corner. It has two non-empty children nodes at level 2, one of which  $r_1^2$  is an internal node and the other of which  $r_2^2$  is a leaf node linked to an external node  $r_4^3$  (also expressed as  $o_5$ ). The sibling set  $\mathcal{S}(r_1^2)$  is  $\{r_2^2\}$  (the other two sibling nodes are empty cells and are thus ignored). The set of ancestors  $\mathcal{A}(r_1^2)$  is  $\{r_1^1,r^0\}$ .
Assume the total number of quadtree nodes is  $N$ , including  $n$  leaf nodes (point samples) and  $N - 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  $\mathbf{H}_p$  for all quadtree internal nodes can be computed by sparse matrix multiplication, as shown in Equation 3 where  $\mathbf{H}_o\in \mathbb{R}^{n\times d}$  is the representation matrix of  $n$  point samples (external nodes),  $\mathbf{H}_p\in \mathbb{R}^{(N - n)\times d}$  is the pooled representation matrix of  $N - n$  internal nodes, and  $\mathbf{P}\in \mathbb{R}^{(N - n)\times n}$  is a sparse pooling matrix. Each row of  $\mathbf{P}$  is normalized and its non-zero values indicate all the corresponding external nodes (point samples) under an internal node. The computational structure is shown by Figure 2(c). We concatenate the internal node feature  $\mathbf{H}_p$  and external nodes feature  $\mathbf{H}_o$  to form a representation matrix  $\mathbf{H}\in \mathbb{R}^{N\times d}$  for all quadtree nodes.
$$
\mathbf {H} _ {p} = \text {Q u a d t r e e P o o l i n g} \left(\mathbf {H} _ {o}\right) = \mathbf {P H} _ {o}, \tag {3}
$$
4.2 Efficient Hierarchical Spatial Attention
The goal of the spatial attention layer is to model the implicit long-range spatial dependency between all sample points. Computing all-pair self-attention across massive points (e.g., a million) is computationally prohibitive due to 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_i$  to all other points as keys, we only compute the weight from  $o_i$  to a selective subset of quadtree nodes. Our intuition is that for key points that are far away from the query point  $o_i$ , 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  $\mathcal{K}$ , which includes the point itself, its own siblings (points) as well as the siblings of its ancestors. That is,  $\mathcal{K}(\pmb{o}_i) = \{\pmb{o}_i\} \cup \mathcal{S}(\pmb{o}_i) \cup \{\mathcal{S}(r) | r \in \mathcal{A}(\pmb{o}_i)\}$ . For example, in Figure 3, the key set of the external node  $\pmb{o}_1$  (also
noted as  $r_1^4$ ) is  $\{\pmb{o}_1, \pmb{o}_2, r_2^3, r_3^3, r_2^2, r_2^1, r_3^1, r_4^1\}$ . In this way, we reduce the number of attention weight calculations from all 11 points to 8 quadtree nodes. Particularly, we use one quadtree node  $r_4^1$  to approximate all the four points within it ( $o_8$  to  $o_{11}$ ) since they are far away from  $o_1$ . Based on the definition of a key set, the spatial attention operator can be expressed as Equation below, where  $h_i$  is the output representation for  $o_i$ ,  $\mathcal{K}_i$  is the key node set of  $o_i$ ,  $q_i$  is the query vector of point  $o_i$ ,  $k_j$  and  $v_j$  are the key vector and value vector of
(a)
Figure 3: An example of a quadtree (a) and the selective key node set of  $\mathbf{o}_1$  in red boxes (b).
the attended node  $r_j$ , respectively, and  $d$  is the latent dimension of these vectors.
$$
\boldsymbol {h} _ {i} = \sum_ {j \in \mathcal {K} _ {i}} \frac {\exp \left(\boldsymbol {q} _ {i} \boldsymbol {k} _ {j} ^ {T} / \sqrt {d}\right) \boldsymbol {v} _ {j}}{\sum_ {j \in \mathcal {K} _ {i}} \exp \left(\boldsymbol {q} _ {i} \boldsymbol {k} _ {j} ^ {T} / \sqrt {d}\right)} \tag {4}
$$
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  $\mathbf{H} \in \mathbb{R}^{N \times d}$  and the representation of external nodes (point samples) as  $\mathbf{H}_o$ . The query matrix of all point samples can be computed by an embedding with learnable parameter matrix  $\mathbf{W}_q$ , i.e.,  $\mathbf{Q}_o = \mathbf{W}_q\mathbf{H}_o$ . For simplicity, we denote the query matrix  $\mathbf{Q}_o$  as  $\mathbf{Q}$ . For each query point  $\mathbf{o}_i$ , its keys are a subset of quadtree nodes. We denote their embedded key vectors in a matrix  $\mathbf{K}_i \in \mathbb{R}^{d \times |\mathcal{K}_i|}$ . If we concatenate the key matrices for all queries, we get a 3D tensor  $\hat{\mathbf{K}} \in \mathbb{R}^{d \times |\mathcal{K}_i| \times n}$ . Similarly, the corresponding value matrices can be concatenated into a 3D tensor  $\hat{\mathbf{V}} \in \mathbb{R}^{d \times |\mathcal{K}_i| \times n}$ . The construction of 3D tensors can be implemented by the torch.gather() API in Pytorch. The corresponding cross-attention
weights can be calculated by matrix and tensor multiplications, as illustrated in Figure 4. We can see the difference between the proposed hierarchical spatial attention and the default all-pair attention. In the all-pair self-attention (Figure 4 top), we would have to compute the dot product attention for all point entries, i.e.,  $\mathbf{Q}_o^T\cdot \mathbf{K}_o$ , whose size is  $n\times n$  and can become too large (e.g.,  $n = 1,000,000$  for a million points). In our spatial attention layer, we conduct a sparse self-attention, i.e., only computing the attention weight of a sample point to its key set. In other words, the corresponding key matrix for  $\mathbf{o}_i$  is a vertical slicing of  $\hat{\mathbf{K}}$ , which can be denoted as  $\hat{\mathbf{K}} [(:, :, i]\in \mathbb{R}^{d\times |\mathcal{K}_i|}$ , where  $|\mathcal{K}|$  is the maximum key node set size ( $|\mathcal{K}| < < n$ ).
Figure 4: Sparse spatial attention with selective key set.