Voxel Set Transformer: A Set-to-Set Approach to 3D Object Detection from Point Clouds

In plain words

Most self-driving car detectors that read LiDAR either chop space into a grid and run convolutions on it, or keep the raw points and group them into small clusters. Grouping forces every cluster to the same size, so points get dropped or padded at random. Grid convolutions only see a small neighbourhood at a time.

VoxSeT splits space into voxels. Inside each voxel, the points send their information to a small, fixed set of learned summary slots, and then read back from those slots. This is a transformer-style attention step, but it runs on the few slots instead of on all pairs of points, so it stays cheap and no points are thrown away. The network then turns the result into a top-down map and predicts boxes. It is reported to be competitive on two driving benchmarks.

TL;DR

VoxSeT is a 3D LiDAR detector whose backbone is built from voxel-based set attention (VSA). VSA runs two cross-attentions per voxel through a small set of latent codes, so its cost is linear in the number of points. Points are never dropped or padded. The backbone feeds a bird's-eye-view (BEV) CNN and an anchor-based detection head. Results on KITTI and Waymo are author-reported; the paper calls them competitive, and some baselines beat it in individual columns tab. 1 tab. 3.

Problem

Point clouds are long, unevenly spread sequences, so full self-attention is too costly §1. Grouping-based transformers balance cluster sizes by dropping or padding points, which makes detections unstable and adds redundant work §1. Voxel-based convolutional attention is linear, but its attention field is small, which limits long-range context §1. Methods that compute attention on a reduced seed set are, according to the paper, limited to indoor scenes §1.

Method

The model has three parts §3.3:

1. A voxel feature extractor made of four VSA blocks. Each block uses voxels twice as large in x and y as the previous one, with z unchanged §3.2 voxset.py.
2. A BEV encoder: points are pooled onto the BEV grid and a shallow CNN with two strides, three convolutions per stride, concatenates its two scales §3.3 voxset.yaml.
3. A detection head. The paper uses an anchor-based head with a foreground-segmentation loss on the point features, following PointRCNN §3.3.

Inside a VSA block, the paper describes three steps §3.2:

• Encoder. Each point is projected to a key and a value. The latent codes act as queries, and a softmax is taken over the points of each voxel, so every voxel becomes k hidden vectors. Scatter sums give one hidden set per voxel eq. 6 eq. 7 eq. 8.
• Convolutional feed-forward network (ConvFFN). The voxel hidden features are placed on a sparse grid and passed through two depthwise convolutions with a non-linearity in between, so neighbouring voxels exchange information eq. 9.
• Decoder. Each point queries the hidden features of its own voxel, with its output computed from the enriched hidden set eq. 10 eq. 11 eq. 12.

Local coordinates of each point inside its voxel are encoded with a Fourier embedding and added before each VSA block §3.2. The paper's first stage on KITTI uses a 0.32 m by 0.32 m voxel and 4 m height, and the second uses 0.64 m by 0.64 m §4.1.

The paper's pseudocode (own wording, following the paper's equations; the code differs in places, see Paper vs code):

# Voxel-based set attention (VSA) block, one per scale  [§3.2] [eq. 6-12]
# Input:  feats (N, d) point features; voxel (N,) voxel index per point
# Output: (N, 2 * d) point features (the code concatenates, see Pitfalls)
def vsa_block(feats, voxel, voxel_coords, grid_shape, k):
    # Per-point MLP before the encoder  [code: voxset.py#L133-L142]
    x = pre_mlp(feats)                              # (N, d)
    # Encoder: a score per latent for every point, then a softmax over
    # the points that share a voxel  [§3.2] [eq. 6]
    # [code: voxset.py#L175]
    scores = score(x)                               # (N, k)
    weights = softmax_per_voxel(scores, voxel)      # (N, k)
    # Weighted sum of the point features inside each voxel gives the
    # hidden latents  [§3.2] [eq. 7] [eq. 8]  [code: voxset.py#L176-L181]
    latents = sum_per_voxel(weights[:, :, None] * x[:, None, :])
    latents = relu(latents)                         # (M, k * d)
    # ConvFFN: scatter to the voxel grid, convolve, gather back  [§3.2]
    # [eq. 9]  [code: voxset.py#L180-L184]
    grid = scatter_to_grid(latents, voxel_coords, grid_shape)
    grid = conv_ffn(grid)
    latents = gather_from_grid(grid, voxel_coords)  # (M, k * d)
    # Decoder: each point attends over the latents of its own voxel
    # [§3.2] [eq. 10] [eq. 11] [eq. 12]  [code: voxset.py#L187-L190]
    latents = broadcast_to_points(latents, voxel)   # (N, k, d)
    decoded = attention(query=x, key=latents, value=latents)
    # Output is the concatenation of input and decoded features, not a
    # sum  [code: voxset.py#L193]
    return concat(feats, decoded)

Key equations, as printed in the paper. Induced set attention, which VSA localises eq. 1–3:

H = \text{CrossAttention}(L, X) \in \mathbb{R}^{k \times d} eq. 1

\hat H = \text{FFN}(H) \in \mathbb{R}^{k \times d} eq. 2

O = \text{CrossAttention}(X, \hat H) \in \mathbb{R}^{n \times d} eq. 3

ConvFFN eq. 9:

\hat H_r = \text{DwConv}\big(\sigma(\text{DwConv}(\mathcal T(H_r, C_r); W_1)); W_2\big) eq. 9

Training loss as printed eq. 14. The printed formula has an unmatched closing parenthesis, and no weights are written:

\mathcal L = \mathcal L_{seg} + \frac{1}{N_p}(\mathcal L_{cls} + \mathcal L_{reg}) + \mathcal L_{dir} eq. 14

Pillar pooling for the BEV map eq. 13 uses a softmax over the points of a pillar as weights:

F^j = \sum_{m} X^j_m \, w^j_m, \qquad w^j_m = \frac{e^{X^j_m}}{\sum_{m} e^{X^j_m}} eq. 13

The paper describes its VSA module with a matrix-multiplication diagram and PyTorch-style steps fig. 2. The overall architecture is shown in fig. 3.

Hyperparameters

Results

Author-reported. Setup: KITTI and Waymo, 3D detection, car class on KITTI with IoU 0.7 §4.2. The paper trains on four GPUs §4.1. The Waymo table reports vehicle 3D mAP at IoU 0.7 §4.2.

Waymo validation, vehicle, 3D mAP tab. 1 §4.3:

KITTI test, car 3D AP (Easy / Moderate / Hard), two-stage row tab. 3 §4.4:

KITTI validation, 11 recall points, car 3D AP (Easy / Moderate / Hard), two-stage row tab. 4 §4.4:

KITTI validation, single-stage, 11 recall points, mAP (Easy / Moderate / Hard) tab. 2 §4.4:

Latency and memory, KITTI, 2080 Ti tab. 8: SECOND 48 ms and 6093 MB; PointPillars 22 ms and 1508 MB; VoxSeT single-stage 34 ms and 2381 MB.

Ablations, KITTI validation, RPN output, 11 recall points §4.5:

• ConvFFN replaced by a plain FFN: Easy / Moderate / Hard drop from 88.31 / 79.56 / 77.84 to 70.12 / 69.54 / 54.23 tab. 5.
• Latent codes per VSA block, 4 / 8 / 16: mAP 76.63 / 78.25 / 78.74 tab. 6.
• PointRCNN with PointNet++ backbone vs VoxSeT backbone: 88.52 / 78.95 / 77.81 vs 89.61 / 80.14 / 78.69 tab. 7.

Paper vs code

Code: the official repository at code_commit, built on OpenPCDet README.md. It has single-stage configurations only: a KITTI PointPillar-style detector with the VoxSeT backbone voxset.yaml, and a Waymo CenterPoint configuration centerpoint_voxset_1x.yaml.

Discrepancies found:

• Point cloud range. The paper gives KITTI x in [0, 70.4] m and y in [-40, 40] m §4.1. The code uses x in [0, 69.12] m and y in [-39.68, 39.68] m voxset.yaml. Waymo: the paper gives ±75.2 m §4.1, the code ±74.24 m centerpoint_voxset_1x.yaml.
• Learning rate and batch. The paper gives batch size 4 and learning rate 0.003 §4.1. The code trains with 4 samples per GPU on 4 GPUs (16 total) at learning rate 0.0015, and its own comment says the learning rate should change with batch size voxset.yaml.
• Latent codes are not learned parameters. The paper describes learned latent codes L and separate key and value projections §3.1 §3.2. The code has no latent parameter. A single linear layer produces one score per latent for each point, and the value is the per-point MLP output. The code comment says the latent codes are absorbed into this linear projection voxset.py voxset.py.
• Skip connection. The paper says each VSA module is wrapped in a residual block §3.3. The code concatenates the block input with the decoded features voxset.py. Block output width is therefore twice its input width. A post-MLP projects 256 channels to 128 voxset.py.
• ConvFFN. Eq. 9 has two depthwise convolutions eq. 9. The code uses two depthwise 3×3 convolutions with BatchNorm and ReLU, then a dense 1×1 convolution without a following activation voxset.py. A third depthwise convolution is commented out voxset.py. A ReLU is applied to the voxel hidden features before the ConvFFN voxset.py. The code uses a dense grid for the ConvFFN (.dense()), so memory grows with grid size voxset.py.
• Positional encoding. The paper describes a Fourier embedding with sin(f_k π x) and cos terms and bandwidth L = 64 §3.2. The code uses the DETR-style sinusoidal encoding with scale 2π and temperature 10000, 64 frequencies per axis, followed by a linear projection voxset.py. The offset is taken inside the finest voxel and reused by all four blocks voxset.py voxset.py.
• BEV pooling. The paper uses soft-pooling over a 0.36 m pillar §3.3 eq. 13. The code uses max-pooling over points in the first-level voxel (0.32 m) voxset.py, then the standard pillar scatter to BEV voxset.yaml.
• Loss. The paper states no normalisation for L_seg. The code uses a focal loss with α = 0.25 and γ = 2.0 point_head_template.py, normalised by the number of positive points point_head_template.py and summed point_head_template.py. Classification and regression losses are normalised by positive anchors per sample anchor_head_template.py. The direction loss is divided by batch size anchor_head_template.py. Weights are 1.0 (cls), 2.0 (loc), 0.2 (dir) voxset.yaml; the paper gives no weights.
• NMS and score thresholds. The paper states NMS at IoU 0.1 and keeps boxes with confidence above 0.3 §4.1. The code uses NMS 0.01 and score threshold 0.1 voxset.yaml voxset.yaml.
• Voxel count caps. The KITTI config sets MAX_NUMBER_OF_VOXELS to 16000 (train) and 40000 (test) voxset.yaml. These caps have no effect with MAX_POINTS_PER_VOXEL = -1: the processor returns early data_processor.py.
• Field of view. The paper keeps only points inside the camera field of view §4.1. The code does this with FOV_POINTS_ONLY kitti_dataset.yaml kitti_dataset.py.
• Two-stage model. The paper's two-stage results (Tables 1, 3, 4) use a second-stage RoI head. §3.3 names the LiDAR-RCNN RoI head §3.3, while Table 1 and §4.3 call the head "CT3D" tab. 1 §4.3. No configuration in the repository uses a second stage.
• Waymo head. The paper describes an anchor-based head §3.3 and does not say which head produced Table 1. The Waymo config uses CenterHead with 4 latent codes centerpoint_voxset_1x.yaml centerpoint_voxset_1x.yaml.
• KITTI validation numbers. The README reports single-stage car 3D AP at R11 of 88.87 / 78.77 / 77.58, pedestrian 60.25 / 55.55 / 50.19, cyclist 85.42 / 70.28 / 64.98 README.md. tab. 2 reports 88.45 / 78.48 / 77.07, 60.62 / 54.74 / 50.39, and 84.07 / 68.11 / 65.14. The README uses IoU 0.5 for pedestrian and cyclist README.md README.md. The paper does not state the pedestrian and cyclist IoU.
• Epochs for Waymo. The Waymo config uses 30 epochs centerpoint_voxset_1x.yaml. The paper gives no Waymo epoch count.

No discrepancy is claimed for the 100-epoch KITTI schedule or weight decay 0.01, which match §4.1 voxset.yaml voxset.yaml.

Pitfalls

• Eq. 14 is printed with an unmatched parenthesis and without loss weights. Use the code weights (1.0, 2.0, 0.2, 1.0) if you need to reproduce the results voxset.yaml.
• The paper's IoU bounds σ1 and σ2 for N_p are not stated. Use the per-class thresholds in the config voxset.yaml voxset.yaml voxset.yaml.
• Latent codes. The paper says they are learned and uses separate key and value projections; the code uses one linear score layer and no separate key projection. A faithful reimplementation of the paper's text would be a different network from the released code. Decide which one you are reproducing.
• Block width. The paper lists 16, 32, 64, 128 as feature dimensions §4.1. Because of concatenation, the code's input widths are 16, 32, 64, 128 and the outputs are 32, 64, 128, 256 voxset.py voxset.py.
• Input features. The first MLP receives absolute x, y, z and intensity voxset.py voxset.py. The paper does not say whether positions are absolute or relative.
• BatchNorm. The code uses eps 1e-3 and momentum 0.01 voxset.py. The paper gives no BatchNorm settings.
• Dense ConvFFN cost. The paper's ConvFFN cost counts grid cells (HWD/dxdydz) §3.2. The code densifies the voxel grid at every block voxset.py, so memory follows grid size, not only occupied voxels.
• Positional encoding. The paper uses sin(f_k π x) §3.2; the code uses the 2π-scaled encoding with temperature 10000 voxset.py. Follow the code if you want to match its behaviour; the paper does not state the temperature.
• Paper text vs its own tables. The text claims +5% over PV-RCNN on LEVEL_1 (table: 76.02 − 70.30 = 5.72), +2.4% over Voxel-RCNN on LEVEL_2 (table: 68.16 − 66.59 = 1.57), and +0.9% / +1.4% over VoTR-TSD (table: 1.07 / 2.25) §4.3 tab. 1. For KITTI test Moderate, the text says +0.4% over CT3D, the table gives 82.06 − 81.77 = 0.29 §4.4 tab. 3. The differences above are my arithmetic from the table values.
• Table references. Ablation text cites "Table \cite{pointnet++}" where it means the PointRCNN table §4.5 tab. 7. Table numbers here follow source order in the LaTeX; check against the compiled PDF if they differ.
• Latency. The timing is measured on a 2080 Ti tab. 8 but the training hardware is a different GPU §4.1. Reproduce timing on your own hardware.
• Two-stage gains need the RoI head. The gains in the two-stage rows come from a head not present in the repository §3.3. Do not attribute them to the backbone alone.
• Data augmentation. The paper cites other works for augmentation and gives no details §4.1. Use the config values voxset.yaml.

Lineage

• Builds on: induced set attention (the paper's [setTr]) for latent-code attention §3.1; voxel feature encoding of VoxelNet, SECOND and PointPillars for the voxel encoder §3.2; grid-based CNN detectors §1.
• Compared against: SECOND, PointPillars, PV-RCNN, Voxel-RCNN, CT3D and VoTR-TSD on KITTI and Waymo tab. 1 tab. 2 tab. 3 tab. 4.
• Superseded by: not stated in this wiki.