SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhejun Zhang
zhejun.zhang@vision.ee.ethz.ch
Alexander Liniger
alex.liniger@vision.ee.ethz.ch
Christos Sakaridis
csakarid@vision.ee.ethz.ch
i@yf.io
The real-world deployment of an autonomous driving system requires its components to run on-board and in real-time, including the motion prediction module that predicts the future trajectories of surrounding traffic participants. Existing agent-centric methods have demonstrated outstanding performance on public benchmarks. However, they suffer from high computational overhead and poor scalability as the number of agents to be predicted increases. To address this problem, we introduce the K-nearest neighbor attention with relative pose encoding (KNARPE), a novel attention mechanism allowing the pairwise-relative representation to be used by Transformers. Then, based on KNARPE we present the Heterogeneous Polyline Transformer with Relative pose encoding (HPTR), a hierarchical framework enabling asynchronous token update during the online inference. By sharing contexts among agents and reusing the unchanged contexts, our approach is as efficient as scene-centric methods, while performing on par with state-of-the-art agent-centric methods. Experiments on Waymo and Argoverse-2 datasets show that HPTR achieves superior performance among end-to-end methods that do not apply expensive post-processing or model assembling. The code is available at https://github.com/zhejz/HPTR.
Motion prediction is an important component of modular autonomous driving stack [63]. As the downstream module of perception [22, 42, 48] and the upstream module of planning [3, 9, 39], the task of motion prediction [16, 60] is to predict the multi-modal future trajectories of other agents next to the self-driving vehicle (SDV) based on the heterogeneous observations, including for example high-definition (HD) maps, traffic lights and other vehicles, pedestrians and cyclists. This is an essential task because without accurate and real-time prediction results [30], the planning module cannot safely and comfortably navigate the SDV through highly interactive driving environments.
To achieve top performance on public motion prediction leaderboards [1, 58], state-of-the-art (SOTA) methods [36, 47, 52, 55] leverage agent-centric vectorized representations and Transformer-based network architectures. However, the good performance of these approaches comes at a cost of high computational overhead as illustrated in Figure 1. The most well-known problem of agent-centric
(a) Dense traffic scenario
(b) Agent-centric ROIs
(c) Online inference
Figure 1: To efficiently predict the multi-modal future of numerous agents in dense traffic (1a), HPTR minimizes the computational overhead by: (1b) Sharing contexts among target agents. (1c) Reusing static contexts during online inference. (1d) Avoiding expensive post-processing and ensembling.
(d) Trajectory aggregation
approaches is the poor scalability as the number of target agents grows in urban driving environments with dense traffic; for example the busy intersection in Figure 1a. Although the agent-centric regions of interest (ROI) in Figure 1b are largely overlapping, the same context is transformed to the coordinate system of each target agent and independently saved and processed. This causes a huge waste of computational resources, which is often neglected by prior works because their experiments focus on offline inference that queries the prediction only once given a scenario from the dataset. In contrast to offline inference, the motion prediction module of a real-world SDV is continuously queried with streaming inputs during online inference. For example, in Figure 1c, the prediction module is queried again shortly after Figure 1a. Although most contexts remain unchanged after this short period of time, existing methods would start inference from scratch without reusing the encoded features of static contexts. Moreover, prior works often predict a massive number of redundant trajectories and ensemble the outputs of numerous models, such that the final output can be adjusted in favor of prediction diversity during the post-processing, as illustrated in Figure 1d. Although these techniques significantly boost the performance on public benchmarks, they should be sparingly used on a real-world SDV where the computational resources are scarce and the raw predictions, rather than the heuristically aggregated ones, are preferred by the downstream planning module.
To address these problems, our first step is to represent everything as heterogeneous polylines [17]. Then we adopt the pairwise-relative representation [11, 28] and separate the local attribute from the global pose of each polyline. The local attribute specifies what the polyline actually is and it will be shared or reused if possible. The global pose specifies where the polyline is and it will be used to derive the pairwise-relative pose before being processed by the neural networks. Considering polylines as tokens, our second step is to introduce the K-nearest Neighbor Attention with Relative Pose Encoding (KNARPE) module, which allows Transformers [53] to aggregate the local contexts for each token via the local attributes and relative poses. While KNARPE enables context sharing among agents, we further propose the Heterogeneous Polyline Transformer with Relative pose encoding (HPTR), which emphasizes the heterogeneous nature of polylines in motion prediction tasks. Based on KNARPE, HPTR uses hierarchical Transformer encoders and decoders to separate the intra-class and inter-class attention, such that tokens can be updated asynchronously during online inference. More specifically, for static polylines, such as HD maps, their features will be reused, whereas for dynamic polylines, such as traffic lights and agents, their features will be updated on demand.
Our contributions are summarized as follows: (1) We introduce KNARPE, a novel attention mechanism that enables the pairwise-relative representation to be used by Transformer-based architectures. (2) Based on KNARPE we propose HPTR, a hierarchical framework that minimizes the computational overhead via context sharing among agents and asynchronous token update during online inference. (3) Compared to SOTA agent-centric methods, we achieve similar performance while reducing the memory consumption and the inference latency by  $80\%$ . By caching the static map features during online inference, HPTR can generate predictions for 64 agents in real time at 40 frames per second. Experiments on Waymo and Argoverse-2 datasets show that our approach compares favorably against other end-to-end methods which do not apply expensive post-processing and ensembling.
Motion prediction is a popular research topic because it is an essential component of modular autonomous driving stacks [63]. The task of motion prediction has been formulated in different ways. In this paper, we focus on the most popular one: the marginal motion prediction [6, 8, 27, 32, 33] where the multi-modal future trajectories are predicted individually for each target agent [1, 38, 58]. In contrast to the marginal formulation, joint motion prediction [19, 20, 35, 49] requires the multimodal futures to be simultaneously generated for all target agents from the same scenario [2, 57, 66]. Beyond the open-loop motion prediction tasks, behavior simulation [44, 50, 68] is formulated in a closed-loop fashion where the future agent trajectories are simulated by rolling out a learned policy [3, 59]. Numerous works have also investigated the possibility to combine motion prediction with other modules, such as perception and planning [7, 15, 23, 29], or to learn an end-to-end driving policy [24, 25, 67] mapping sensor measurements directly to the actions or motion plans of the SDV.
Vectorized representation proposed by VectorNet [17] is widely used in recent works because of its promising performance [47, 52, 55]. Depending on how the coordinate system is selected, the vectorized representation falls into three categories: agent-centric, scene-centric and pairwise-relative. The most popular one is the agent-centric representation [36, 47, 52, 55] that transforms all inputs to the local coordinates of each target agent. Despite its good performance, this approach suffers from poor scalability as the number of target agents grows. To reduce the computational overhead, scene-centric representation [37] shares the contexts among all target agents by transforming all inputs to a global coordinate system, which is not tied to any specific agent but to the whole scene. However, its performance is poor due to the lack of rotation and translation invariance. In this paper we use the pairwise-relative representation, which is less often discussed in prior works because it has not demonstrated any clear advantage. To the best of our knowledge, there are three works using the pairwise-relative representation: HDGT [28], GoRela [11] and HiVT [70]. HDGT and GoRela are based on Graph Neural Networks (GNNs), and their implementation requires message passing and GNN libraries. However, these libraries are often less efficiently implemented on Graphics Processing Units (GPUs) when compared to the basic matrix operations that Transformers rely on. HiVT augments the agent-centric encoders with a pairwise-relative interaction decoder in order to realize multi-agent prediction. In contrast to HiVT which uses agent-centric vectors and the standard Transformer, we formulate all inputs as pairwise-relative polylines and introduce the KNARPE attention mechanism. As a result, our HPTR demonstrates clear advantages in terms of accuracy and efficiency. Replacing the vanilla attention with our KNARPE, we can borrow ideas from other Transformer-based methods and adapt them to the pairwise-relative representation. For example, the hierarchical architecture of our HPTR is inspired by Wayformer [36]. More sophisticated architectures and techniques, such as goal-based decoding [21, 69], can also be incorporated into our framework to further boost the performance.
Rasterized representation was widely used in early works on motion prediction [34, 40, 45, 51]. These methods use convolutional neural networks (CNNs) to process the rasterized image of agent-centric ROI. To improve the inference efficiency, some works [5, 6, 10, 65] pre-compute the static map features and use rotated ROI alignment [26] to retrieve the local contexts around the target agent. In this paper, our KNARPE realizes this operation for pure Transformer-based architectures.
Transformers with attention mechanism [53] have achieved great success in natural language processing [13, 43] and computer vision tasks [4, 14, 64]. Inspired by the vision Transformers, we treat polylines as high-dimensional pixels; the local attribute corresponds to color channels, whereas the global pose corresponds to the pixel-wise location. We further extend the relative position encoding [12, 46] to higher dimensions and rename it to relative pose encoding (RPE) in this paper. Although the benefit of RPE is still controversial for other tasks [54, 61], its value for motion prediction is demonstrated in this paper. Instead of using RPE, previous motion prediction Transformers use everything as input, which is equivalent to concatenating the pixel-wise location to the RGB channels. This is a very uncommon practice from the perspective of vision Transformers.
In Sec. 3.1 we introduce the pairwise-relative polyline representation which enjoys the advantages of both the agent-centric and the scene-centric representation. Then in Sec. 3.2 we present the K-nearest Neighbor Attention with Relative Pose Encoding (KNARPE) which enables the pairwise-
(a) Polylines.
(b) Local attribute represented in the global pose's frame.
(c) Obtain relative poses.
Figure 2: Pairwise-relative polyline representation. MP: Map. TL: Traffic lights. AG: Agents.
relative polyline representation to be used by Transformers. Based on KNARPE, we propose the Heterogeneous Polyline Transformer with Relative pose encoding (HPTR) in Sec. 3.3. HPTR uses a hierarchical architecture to prevent redundant computations and to realize the asynchronous update of heterogeneous tokens. Finally in Sec. 3.4 we discuss our output representation and training strategies.
3.1 Pairwise-relative polyline representation
Based on prior works [11, 28], we formulate the pairwise-relative polyline representation as the third type of input representation for motion prediction. As shown in Figure 2a, all data relevant to motion prediction can be represented as an ordered list of consecutive vectors, i.e. polylines. By transforming the polyline to the coordinate frame of its global pose as done in Figure 2b, a polyline is fully described by a pair of global pose and local attribute. The global pose specifies the location and heading of the polyline in the global coordinate system, whereas the local attribute describes the polyline in its local coordinates, for example it is a yellow solid lane, 8 meters long, slightly curved to the right. In practice, it does not matter where the lane is on the earth; we only care about where the lane is relative to us. Hence, we obtain the relative poses from the global poses as shown in Figure 2c, before feeding them to the prediction network. Specifically, the global pose and local attribute of polyline  $i$  are denoted as  $(\mathbf{p}_i,\mathbf{u}_i)$ . The global pose  $\mathbf{p}_i$  has 3 degrees of freedom  $(x,y,\theta)$ , i.e. the 2D position and the heading yaw angle. The local attribute  $\mathbf{u}_i$  is derived from  $\mathbf{c}_i$ , the intrinsic characteristics of the polyline, and  $\mathbf{l}_i$ , the polyline vectors represented in the local coordinate.
The task of marginal motion prediction is to predict the future 2D positions individually for each target agent based on the static HD map (MP), the history trajectories of all agents (AG) and the traffic lights (TL). For each scenario to be predicted, we consider a maximum number of  $N_{\mathrm{MP}}$  map polylines,  $N_{\mathrm{TL}}$  traffic lights and  $N_{\mathrm{AG}}$  agents. We define  $t = 0$  to be the current step,  $\{T_{\mathrm{h}} - 1,\dots ,0\}$  to be the observed history steps and  $\{1,\ldots ,T_{\mathrm{f}}\}$  to be the predicted future steps. The static HD map are spatial polylines  $(\mathbf{p}_i^{\mathrm{MP}},\mathbf{l}_i^{\mathrm{MP}},\mathbf{c}_i^{\mathrm{MP}}), i\in \{1,\dots ,N_{\mathrm{MP}}\}$  where  $\mathbf{p}_i^{\mathrm{MP}}\in \mathbb{R}^3$  is its starting vector,  $\mathbf{l}_i^{\mathrm{MP}}\in \mathbb{R}^{N_{\mathrm{node}}\times 4}$  are the 2D positions and directions of the  $N_{\mathrm{node}}$  segments normalized against  $\mathbf{p}_i^{\mathrm{MP}}$ , and  $\mathbf{c}_i^{\mathrm{MP}}\in \mathbb{R}^{C_{\mathrm{MP}}}$  is the one-hot encoding of  $C_{\mathrm{MP}}$  different lane types. Polygonal elements, such as crosswalks, are converted to a group of parallel polylines across the polygon. Similar to the map, history trajectories of agents are represented as spatial-temporal polylines  $(\mathbf{p}_i^{\mathrm{AG}},\mathbf{l}_i^{\mathrm{AG}},\mathbf{c}_i^{\mathrm{AG}}), i\in \{1,\dots ,N_{\mathrm{AG}}\}$  where  $\mathbf{p}_i^{\mathrm{AG}}\in \mathbb{R}^3$  is the last observed agent pose,  $\mathbf{l}_i^{\mathrm{AG}}\in \mathbb{R}^{T_{\mathrm{h}}\times (6 + 3)}$  contains the history 2D positions, directions and velocities normalized against  $\mathbf{p}_i^{\mathrm{AG}}$  as well as the 1D speed, yaw rate and acceleration which do not need to be normalized, and  $\mathbf{c}_i^{\mathrm{AG}}\in \mathbb{R}^6$  is the 3D agent size and the one-hot encoding of 3 agent types (vehicle, pedestrian and cyclist). Following VectorNet [17], we use PointNet [41] with masked max-pooling to aggregate  $(\mathbf{l}_i^{\mathrm{MP}},\mathbf{c}_i^{\mathrm{MP}})$  into  $\mathbf{u}_i^{\mathrm{MP}}\in \mathbb{R}^D$  and  $(\mathbf{l}_i^{\mathrm{AG}},\mathbf{c}_i^{\mathrm{AG}})$  into  $\mathbf{u}_i^{\mathrm{AG}}\in \mathbb{R}^D$ , where  $D$  is the hidden dimension. Since current datasets contain only the detection results rather than the tracking results of traffic lights, we consider only the traffic lights observed at  $t = 0$  and represent them as singular polylines  $(\mathbf{p}_i^{\mathrm{TL}},\mathbf{c}_i^{\mathrm{TL}}), i\in \{1,\dots ,N_{\mathrm{TL}}\}$  where  $\mathbf{p}_i^{\mathrm{TL}}\in \mathbb{R}^3$  is the pose of the stop point and  $\mathbf{c}_i^{\mathrm{TL}}\in \mathbb{R}^{C_{\mathrm{TL}}}$  is the one-hot encoding of  $C_{\mathrm{TL}}$  different states of traffic lights. We use a multi-layer perceptron (MLP) to encode  $\mathbf{c}_i^{\mathrm{TL}}$  into  $\mathbf{u}_i^{\mathrm{TL}}\in \mathbb{R}^D$ .
Because the local attributes are normalized and the global poses are used to compute the relative poses before being consumed by the network, the pairwise-relative representation preserves the viewpoint invariance of the agent-centric representation. Additionally, since the local attributes have higher dimensions compared to the 3-dimensional global poses, sharing the local attributes enables the pairwise-relative representation to maintain the good scalability of the scene-centric representation.
However, so far this representation has only been exploited by GNNs, which are not comparable to Transformers in terms of accuracy and efficiency on the motion prediction benchmarks [1, 58].
3.2 KNARPE: K-nearest neighbors attention with relative pose encoding
After encoding the local attributes via the polyline-level encoders, the scenario is now described as  $(\mathbf{p}_i,\mathbf{u}_i), i\in \{1,\dots ,N\}$  where  $N = N_{\mathrm{MP}} + N_{\mathrm{AG}} + N_{\mathrm{TL}}$  is the total number of polylines. However, this representation cannot be handled by standard self-attention because on the one hand modeling the all-to-all attention is prohibitively expensive [36] due to the large  $N$ , and on the other hand the performance deteriorates when the input is scene-centric [37]. To address both problems, we introduce the K-nearest Neighbors Attention with Relative Pose Encoding (KNARPE). Similar to MTR [47], KNARPE limits the attention to the K-nearest neighbors of each token. Specifically,  $\kappa_i^K (d_{i1},\ldots ,d_{iN})\subseteq \{1,\ldots ,N\}$  is a set that contains the indices of  $K$  tokens closest to token  $i$ . For simplicity, we use the L2 distance to measure the distance  $d_{ij}$  between the global poses of token  $i$  and  $j$ . Instead of directly processing global poses, KNARPE uses the relative pose encoding (RPE), i.e. the positional encoding for pairwise-relative poses. Denoting  $\mathbf{r}_{ij} = (x_{ij},y_{ij},\theta_{ij})$  to be the global pose of token  $j$  represented in the coordinate of token  $i$ , i.e.  $\mathbf{p}_j$  transformed to the coordinate of  $\mathbf{p}_i$ , the RPE of  $\mathbf{r}_{ij}$  is computed using sinusoidal positional encoding (PE) and angular encoding (AE) [68],
$$
\mathrm {R P E} \left(\mathbf {r} _ {i j}\right) = \operatorname {c o n c a t} \left(\mathrm {P E} \left(x _ {i j}\right), \mathrm {P E} \left(y _ {i j}\right), \mathrm {A E} \left(\theta_ {i j}\right)\right), \tag {1}
$$
$$
\mathrm {P E} _ {2 i} (x) = \sin \left(x \cdot \omega^ {\frac {2 i}{D}}\right), \mathrm {P E} _ {2 i + 1} (x) = \cos \left(x \cdot \omega^ {\frac {2 i}{D}}\right), \tag {2}
$$
$$
\mathrm {A E} _ {2 i} (\theta) = \sin (\theta \cdot (i + 1)), \mathrm {A E} _ {2 i + 1} (\theta) = \cos (\theta \cdot (i + 1)), i \in \{0, \dots , D / 2 - 1 \}, \tag {3}
$$
where  $\omega$  is the base frequency. Following [46], the RPE is projected and added to the keys and values to obtain  $\mathbf{z}_i$ , the output of letting token  $i$  attend to its  $K$  neighbors  $\kappa_i^K$ ,
$$
\mathbf {z} _ {i} = \operatorname {K N A R P E} \left(\mathbf {u} _ {i}, \mathbf {u} _ {j}, \mathbf {r} _ {i j} \mid j \in \kappa_ {i} ^ {K}\right) = \sum_ {j \in \kappa_ {i} ^ {K}} \alpha_ {i j} \left(\mathbf {u} _ {j} \mathbf {W} ^ {v} + \mathbf {b} ^ {v} + \operatorname {R P E} \left(\mathbf {r} _ {i j}\right) \hat {\mathbf {W}} ^ {v} + \hat {\mathbf {b}} ^ {v}\right), \tag {4}
$$
$$
\alpha_ {i j} = \frac {\exp \left(e _ {i j}\right)}{\sum_ {k \in \kappa_ {i} ^ {K}} \exp \left(e _ {i k}\right)}, \quad e _ {i j} = \frac {\left(\mathbf {u} _ {i} \mathbf {W} ^ {q} + \mathbf {b} ^ {q}\right) \left(\mathbf {u} _ {j} \mathbf {W} ^ {k} + \mathbf {b} ^ {k} + \mathrm {R P E} \left(\mathbf {r} _ {i j}\right) \hat {\mathbf {W}} ^ {k} + \hat {\mathbf {b}} ^ {k}\right)}{\sqrt {D}}, \tag {5}
$$
where  $\alpha_{ij}$  are the attention weights,  $e_{ij}$  are the logits,  $W^{\{q,k,v\}}$ ,  $b^{\{q,k,v\}}$  are the learnable projection matrices and biases for query, key and value, and  $\hat{W}^{\{k,v\}}$ ,  $\hat{b}^{\{k,v\}}$  are the learnable projection matrices and biases for RPE. We do not apply RPE to query because doing this does not boost the performance in our experiments.
Efficient implementation of KNARPE can be achieved using basic matrix operations such as matrix indexing, summation and element-wise multiplication. See the appendix for more details. KNARPE allows the pairwise-relative representation to be used by pure Transformer-based architectures. Self-attention with KNARPE aggregates the local context for each token in the same way as CNN aggregates the context around each pixel via convolutional kernels. The pixel corresponds to the token, whereas the kernel size of a CNN corresponds to the number of neighbors of KNARPE. Cross-attention with KNARPE enables rotated ROI alignment of pre-computed features, e.g. the static map features. Previously, this was only possible for CNN-based [6, 10] and GNN-based [11, 28] methods.
3.3 HPTR: Heterogeneous polyline transformer with relative pose encoding
By replacing the standard multi-head attention [53] with our KNARPE, we construct Transformer encoders and decoders which can model the interactions between heterogeneous polylines via relative poses and local attributes. To address the marginal motion prediction task involving HD map, traffic lights and agents, we propose the Heterogeneous Polyline Transformer with Relative pose encoding (HPTR) as shown in Figure 3. Since the temporal dimension is eliminated by the polyline-level encoder [17], HPTR models only the spatial relationship between tokens from different classes.
Firstly in Figure 3a, the intra-class Transformer encoders build a block diagonal attention matrix that models the interactions within each class of tokens. Then in Figure 3b, the inter-class Transformer decoders enhance the traffic lights tokens by making them attend to the map tokens, whereas the agent tokens are enhanced by attending to both the traffic lights and map tokens. The intuition behind this is the hierarchical nature of traffic; first there is only the map, then the traffic lights are added
(a) Intra-class TF encoder
(b) Inter-class TF decoder
(c) All-to-all TF encoder
Figure 3: The hierarchical architecture of HPTR. Transformers are applied in sequential order from left to right, top to down. Some attentions are redundant and can be skipped for better efficiency. We propose the lower triangular attention matrix, which excludes the intra-TL, intra-AG and all-to-all self-attention (the transparent parts). By removing the redundant attentions, this specific lower triangular architecture enables asynchronous token update during online inference. TF: Transformer. Attn: Attention. MP: Map. TL: Traffic lights. AG: Agents.
(d) AC-to-all TF decoder