ScatterFormer: Efficient Voxel Transformer with Scattered Linear Attention

TL;DR

ScatterFormer tackles efficient large-scale LiDAR perception by replacing group-based attention in windowed voxel transformers with a Scattered Linear Attention (SLA) that attends across windows as a single sequence ($O(N)$) and a Cross Window Interaction (CWI) module to enable cross-window connectivity without window shifting. SLA uses a chunk-wise GPU algorithm and SRAM-based data handling to process variable-length window sequences, while CWI blends features across windows through inception-like convolutions, yielding fast, memory-efficient attention. On Waymo and NuScenes, ScatterFormer achieves state-of-the-art accuracy with competitive latency (around 23 FPS), outperforming prior transformer- and CNN-based detectors. This work delivers a scalable, hardware-friendly attention mechanism for sparse 3D data and demonstrates substantial practical impact for real-time autonomous perception.

Abstract

Window-based transformers excel in large-scale point cloud understanding by capturing context-aware representations with affordable attention computation in a more localized manner. However, the sparse nature of point clouds leads to a significant variance in the number of voxels per window. Existing methods group the voxels in each window into fixed-length sequences through extensive sorting and padding operations, resulting in a non-negligible computational and memory overhead. In this paper, we introduce ScatterFormer, which to the best of our knowledge, is the first to directly apply attention to voxels across different windows as a single sequence. The key of ScatterFormer is a Scattered Linear Attention (SLA) module, which leverages the pre-computation of key-value pairs in linear attention to enable parallel computation on the variable-length voxel sequences divided by windows. Leveraging the hierarchical structure of GPUs and shared memory, we propose a chunk-wise algorithm that reduces the SLA module's latency to less than 1 millisecond on moderate GPUs. Furthermore, we develop a cross-window interaction module that improves the locality and connectivity of voxel features across different windows, eliminating the need for extensive window shifting. Our proposed ScatterFormer demonstrates 73.8 mAP (L2) on the Waymo Open Dataset and 72.4 NDS on the NuScenes dataset, running at an outstanding detection rate of 23 FPS.The code is available at \href{https://github.com/skyhehe123/ScatterFormer}{https://github.com/skyhehe123/ScatterFormer}.

Figures (6)

• Figure 1: Illustration of (a) group-based attention and (b) our scattered linear attention on variable-length sequences. The square in different colors represent voxels in different windows and the square with a cross represents the padding voxel.
• Figure 2: The macro design of ScatterFormer. The backbone comprises a Conditional Position Encoding (CPE) and six transformer blocks. Each block is composed by a Scattered Linear Attention (SLA) module, a Cross-Window Interaction (CWI) module, and a Feed-Forward Network (FFN).
• Figure 3: (a) The voxel features are sorted by windows and flattened to a single matrix, which is further partitioned into multiple chunks of uniform length and sent to SLA kernel for attention computation. (b) In SLA, we first allocate individual threads to each window, each of them iterates over all corresponding key and value chunks, calculates, and accumulates their products to obtain the hidden state matrix of each window. Then, multiple threads are allocated to these query chunks $q_i \in Q$, calculating the chunk-wise output by multiplying $q_i$ with the corresponding hidden state matrix.
• Figure 4: Comparison between between scatter-based implementation (yellow) and our chunk-wise implementation (blue).
• Figure 5: Runtime decomposition in DSVT dsvt and ScatterFormer backbones. The offsets generation means computing the starting address of each chunk in the flattened matrix.