Perceive, Attend, and Drive:
Learning Spatial Attention for Safe Self-Driving

Our proposed model extends upon the neural motion planner (NMP), which jointly solves the perception, prediction and planning problems for self-driving. In this section we briefly review NMP depicted in Figure 2, and refer the reader to  [51] for more details.

Input and backbone: NMP voxelizes LiDAR point clouds to a birds-eye-view (BEV) feature map and fuses them with $M$ channels of rasterized HDMap features to produce an input representation of size $(Z{T}^{\prime }+M)\times H\times W$, where $Z,H,W$ are the height and spatial dimensions and $T^{\prime }=10$ is the number of input LiDAR sweeps. The backbone consists of 5 blocks, with the first 4 producing multi-scale features that are concatenated and fed to the final block. Overall the backbone downsamples the spatial dimension by 4.

Multi-task headers: Given the features computed by the backbone, $X\in R^{128\times \frac{H}{4}\times \frac{W}{4}}$, NMP uses two separate headers for perception & prediction, and motion planning. The perception & prediction header consists of separate branches for classification and regression. The classification branch outputs a score for each anchor box at each spatial location over the feature map $X$, while the regression branch outputs regression targets for each anchor box, including targets for localization offset, size, and heading angle. The planning header consists of convolution and deconvolution layers to produce a cost volume $C\in R^{T\times \frac{H}{4}\times \frac{W}{4}}$ representing the cost for the self-driving vehicle to be at each location and time, with $T$ the fixed future planning horizon.

Planning inference: At inference time, NMP samples $N$ trajectories that are physically realizable, and chooses the lowest cost trajectory for the ego-car:

where the cost of a trajectory $\tau =(x_t,y_t{)}_{t=1}^T$ is the sum of all its waypoints in the cost volume:

We sample trajectories using a mixture of Clothoid, circle, and straight curves [43]. We refer the readers to [51] for more details on the sampling procedure.

In this section, we propose our sparse spatial attention module for self-driving, shown in Figure 3, which learns to save computation while performing well on the end task of driving safely to the goal.

Input and backbone: We exploit the same input representation as NMP and use the same perception, prediction, and planning headers. The NMP backbone is replaced with the state-of-the-art backbone network of PnPNet [24], which uses cross-scale blocks throughout to fuse BEV sensory input. Each cross-scale block consists of 3 parallel branches at different resolutions that downsample the feature map, perform bulk computations, and then upsample back to the backbone resolution, before finally fusing cross-scale features across all branches. There is an additional residual connection across each cross-scale block. The final output feature from the backbone consists of 128 channels at 4x downsampled resolution, which is forwarded to the planning and detection headers. In addition to the improved performance, PnPNet can be easily scaled for different computational budgets by varying the depth and width of the cross-scale blocks.

Existing attention-driven approaches [35] tackle only the perception task and use either a road mask obtained from map information or a vehicle mask produced by a different perception module. As a consequence, they waste computation on areas that will not affect the self-driving car. We instead propose a novel approach that is end-to-end trainable and performs computation selectively for planning a safe maneuver. As shown in Fig. 3, the learned attention mask then gates the backbone network, limiting computation to areas where attention is active. By using binary attention, we can leverage sparse convolution to improve the computational efficiency.

Generating binary attention: Computational efficiency has been shown to be one of the most prominent advantages of using the attention mechanism. For soft attention masks the computation is still dense across the entire activation map, and therefore no computation savings can be achieved. SBNet [35] showed that a sparse convolution operator can achieve significant speed-ups with a given discrete binary attention mask. Here, we would also like to exploit the computational benefit of sparse convolution by using discrete attention outputs.

We utilize a network to predict a scalar score for each spatial location, and binarize the score to represent our sparse attention map. For efficiency and simplicity, we use a small U-Net [38] with skip connections and two downsample/upsample stages. We would like to apply the generated attention back to the BEV features in the model backbone so as to sparsify the spatial information, allowing computation to be focused on the important regions only. We choose to do so in a residual manner [47, 35] to avoid deteriorating the features throughout the backbone. Let $x+F\left(x\right)$ denote the normal residual block. Our attention mask is multiplied with the input to the residual block as follows:

where $\odot$ denotes elementwise multiplication. See Fig. 3 for an illustration of our architecture.

Learning binary attention with Gumbel softmax: In order to learn the attention generator and backpropagate through the binary attention map, we make use of the Gumbel softmax technique [15, 30] since the step function is not differentiable, and using the standard sigmoid function suffers from a more severe bias-variance trade-off [15]. Let $i,j$ denote spatial coordinates, and $z_{i,j}$ the scalar output from the attention U-Net. We first add the Gumbel noise on the logits as follows:

where $g_{i,j}=-\log (-\log u)$, and $u$ is sampled from $\text{Uniform}[0,1]$. At inference time, hard attention $A_{i,j}$ can be obtained by comparing the logits,

During training, however, we would like to approximate the gradient by using the straight-through estimator [15, 3]. Hence, in the backward pass, the step function is replaced with a softmax function with a temperature constant $K$ (where $\tilde{A}$ is the underlying soft attention):

We train our sparse neural motion planner (including the attention) end-to-end with a multi-task learning objective that combines planning ($L_{\text{plan}}$) with perception & motion forecasting ($L_{\text{cls}}$, $L_{\text{reg}}$):

where $L_A$ is an ${\ell }_1$ loss, defined in Equation 17, that controls the sparsity of the attention mask, and $\parallel w{\parallel }_2^2$ is the standard weight decay term. Following [51, 24], we fix ${\lambda }_{\text{plan}}=0.001,{\lambda }_{\text{cls}}=1.0,{\lambda }_{\text{reg}}=0.5$.

Motion planning loss: The motion planning loss utilitizes the max-margin objective, where the ground-truth driving trajectory (performed by a human) should be of lower cost than other trajectories sampled by the model. Let $(x_t,y_t)$ be the groundtruth trajectory and let $c_t$ be the cost volume output by the model. We sample $N$ lowest cost trajectories based on the cost volume: $\{x_t^{\left(i\right)},y_t^{\left(i\right)}{\}}_{i=1}^N$, and penalize the maximum margin violation between the groundtruth and model samples:

where ${\Delta }_t^{\left(i\right)}$ is the task loss capturing spatial differences, and traffic violations $v_t$ encoded in binary:

Perception & prediction (PnP) loss: This loss follows the standard classification and regression objectives for object detection. The classification part uses binary cross-entropy:

where $\hat{y}$ is the predicted classification score between 0 and 1, and $y$ is the binary groundtruth. For each detected instance, the model outputs a bounding box, and a pair of coordinates and angles for each future step. We reparameterize the shift of a bounding box $(x,y,w,h,\theta )$ from an anchor bounding box $(x_a,y_a,w_a,h_a,{\theta }_a)$ in a 6-dimensional vector $\delta$:

A regression loss is then applied for the trajectory of the instance up to time $T$. For each spatial coordinate $(i,j)$, we sum up the losses of all bounding boxes $b$ where the ground truth belongs to this location:

with $\hat{\delta }$ the model predicted shifts and $\delta$ the ground truth shifts.

Auxiliary loss masking: Our overall objective is to achieve good performance in motion planning, so PnP is an auxiliary task. Since the majority of our computation happens within the attended area, intuitively the model should not be penalized as severely for mis-detecting objects not in the attended area. We therefore propose to use our spatial attention mask $A$ to re-weight the PnP losses as follows:

where ${\gamma }_1$ weights attended instances, and ${\gamma }_0$ weights all instances. We fix ${\gamma }_0=0.1$ and ${\gamma }_1=0.9$.

Attention sparsity loss: To encourage focused attention and high sparsity, we use an ${\ell }_1$ regularizer on the attention mask as follows. We control sparsity with ${\lambda }_A$ in Equation 9.

Training: To jointly train SA-NMP with attention, we use pretrained weights for the backbone and headers from training a SA-NMP without attention (dense) for two epochs. We train all our models with batch size 5 across 16 GPUs in parallel using Adam [19] optimizer. We use an initial learning rate of 0.0001 and a decay of 0.1 at 1.0 and 1.6 epoch(s), for a total of another 2.0 epochs.

Evaluation: To evaluate driving and safety performance, we focus on the following planning metrics which are accumulated over all 6 future timesteps (3s): Planning L2 is the L2 distance between waypoints of the predicted future ego trajectory and those of the ground-truth trajectory (characterized by human driving). Collision rate is the frequency of collisions between the planned ego trajectory and the ground truth trajectories of other actors in the scene. Lane violation rate measures the number of lane boundary violations by the planned ego trajectory. we do not evaluate this on nuScenes due to the inaccurate localization they provide,

Baselines: We compare our learned attention to baselines that are end-to-end trained using static attention masks obtained from priors. Road Mask covers the entire road as provided from the map data. Vehicle Mask strictly covers all detections in the input space, obtained from a PSPNet [53] trained for segmentation. Proximity Mask is a circular radius around the ego vehicle. Dense is not using sparse attention.

Quantitative results: With a sparse attention mask learned towards motion planning, we can leverage the sparsity in the network backbone to greatly reduce computational costs, while not only maintaining but improving model performance. In our experimental results, we use theoretical FLOPs to show the efficiency of our network, but this also translates to real time gains as SBNet [35] has been shown to leverage sparsity to achieve real speed-ups. The increase in efficiency from leveraging sparsity is shown in Table 1, where our learned attention model uses $\sim 80\%$ fewer FLOPs than Dense SA-NMP thanks to its 95% sparse attention mask. Also, even with an identical SA-NMP backbone as the baselines (except NMP), our model with learned attention performs better in all motion planning metrics, which indicates that focused backbone computation is greatly advantageous to the overall goal of safe planning. NMP+Road performs slightly better in Lane Violation due to the road mask attention focusing on all road/lane markings. However, this baseline uses more than double the FLOPs since its attention mask looks at the entire road surface at only 68.9% sparsity. From Fig. 4, our learned attention model clearly outperforms other baselines in collision rate and planning L2, across all computational budgets.

Qualitative results: In Figure 5, we show examples of our learned attention compared to baselines. As expected, our model focuses on the road and vehicles directly ahead; however, it also diverts some amount of attention to distant vehicles and road markings. This ability to dynamically distribute attention is likely why our model outperforms the baselines, which attend either indiscriminately (Dense and Road) or too selectively (Vehicle and Proximity). From the visualizations, we can better understand our model’s improved collision avoidance. Since the attention is dynamic, our model is more effective at anticipating other vehicles resulting in more cautious planning. This is illustrated by Columns A, B in Fig. 5, where our planned trajectory avoids future collisions with others. The failure cases mostly arise from rear-end collisions, one of which is Column D where all models are hit by the trailing vehicle. Note that this arises as we evaluate in open-loop. Our model focuses on surrounding vehicles and not enough on the open road to its right, which would give the option of making a right turn.

Sparsity of learned attention: Table 2 shows the result of varying ${\lambda }_A$ from Eq. 9, which weights the ${\ell }_1$ regularization term from Eq. 17, with other settings held constant. We found that overall motion planning performance improves with increased sparsity, or essentially more focused computation, and peaks at 95% sparsity.

Perception and prediction (PnP) loss reweighting: Table 3 shows results with varying ${\gamma }_1$ and ${\gamma }_0=1-{\gamma }_1$ from Eq. 16 which control the weighting of the PnP loss computed on actors inside vs. outside the attention mask. All other variables are fixed, including sparsity at 95%. As ${\gamma }_1$ increases, the learned attention is less restricted by detection performance on all actors, and is able to focus on only the most important actors and parts of the road, distributing attention towards improving motion planning performance. Note that ${\gamma }_1=1.0$ is an extreme case where PnP loss is computed only on actors within attention mask: the model learns to cheat by generating attention that avoids all actors resulting in no PnP learning signal, hence the poor performance. For our main experiments, we fix ${\gamma }_1=0.9$.

Detection performance: Since the overall goal is improved motion-planning with lighter computation, focusing on accurately detecting all actors indiscriminately would contradict the purpose of our learned sparse attention. We should not care as much about far away or irrelevant actors that have no effect on safe planning, and should instead focus our computation on important input regions. Table 4 compares detection performance between our learned attention and the baseline dense model evaluated on different subsets of actors in the scene. Full includes all actors in the input, while Attended Region is the subset of actors that lie within the attention mask. For evaluating the dense model, we use the attention mask generated by our learned model to get the Attended Region, ensuring that both models are evaluated on the same actor subsets in both settings. The results shows that our 95% sparse, learned attention model is better than the dense model at detecting actors within the attention mask, meaning that its performance is better focused on actors that it believes are important. This may explain the overall improved planning performance of our attention-driven models as demonstrated in the main quantitative and qualitative results.