Source-linked AI summary
Modeling Point Clouds with Self-Attention and Gumbel Subset Sampling
Jiancheng Yang, Qiang Zhang, Bingbing Ni, Linguo Li, Jinxian Liu, Mengdie Zhou, Qi Tian
TL;DR
Point-cloud learning needs models that handle variable-sized, permutation-structured inputs while addressing limitations of heuristic hierarchical sampling. The paper introduces PATs with GSA and GSS, using differentiable training and discrete test-time sampling; experiments report effective, efficient performance across benchmarks.
Problem
Point-cloud networks must handle size-varying and permutation-invariant inputs, while FPS is task-dependent, low-dimensional, permutation-variant, and sensitive to outliers.
Method
PATs use parameter-efficient Group Shuffle Attention for point relations and Gumbel Subset Sampling for end-to-end, task-agnostic subset selection.
Results
Results on classification and segmentation benchmarks demonstrate the effectiveness and efficiency of the proposed methods.
Takeaways & Limitations
GSS produces soft continuous subsets during training and hard discrete selections during testing, supporting hierarchical point-cloud representations with lower computation cost.
Takeaways & Limitations
FPS-based sampling is constrained to low-dimensional Euclidean space and is sensitive to outliers, while the paper's position-embedding operation has O(N^2) complexity.
Abstract
from arXiv · showhide
Geometric deep learning is increasingly important thanks to the popularity of 3D sensors. Inspired by the recent advances in NLP domain, the self-attention transformer is introduced to consume the point clouds. We develop Point Attention Transformers (PATs), using a parameter-efficient Group Shuffle Attention (GSA) to replace the costly Multi-Head Attention. We demonstrate its ability to process size-varying inputs, and prove its permutation equivariance. Besides, prior work uses heuristics dependence on the input data (e.g., Furthest Point Sampling) to hierarchically select subsets of input points. Thereby, we for the first time propose an end-to-end learnable and task-agnostic sampling operation, named Gumbel Subset Sampling (GSS), to select a representative subset of input points. Equipped with Gumbel-Softmax, it produces a "soft" continuous subset in training phase, and a "hard" discrete subset in test phase. By selecting representative subsets in a hierarchical fashion, the networks learn a stronger representation of the input sets with lower computation cost. Experiments on classification and segmentation benchmarks show the effectiveness and efficiency of our methods. Furthermore, we propose a novel application, to process event camera stream as point clouds, and achieve a state-of-the-art performance on DVS128 Gesture Dataset.
1. Introduction
The paper introduces Point Attention Transformers for permutation-equivariant, size-varying point-cloud processing, combining self-attention with parameter-efficient Group Shuffle Attention. It also proposes Gumbel Subset Sampling as a differentiable, task-agnostic alternative to heuristic point sampling.
- Point clouds require networks that consume size-varying and permutation-invariant sets.
- PATs model relations between points using self-attention, replacing costly Multi-Head Attention with parameter-efficient Group Shuffle Attention.
- Hierarchical point-cloud reasoning benefits from sampling central subsets, but the sampling operation is identified as a bottleneck.
- FPS is task-dependent, operates in low-dimensional Euclidean space, and is permutation-variant and sensitive to outliers.
- GSS provides end-to-end trainable, task-agnostic, permutation-invariant subset sampling with soft training samples and hard test-time selections.
2. Preliminaries
The preliminaries review direct point-cloud processing, attention mechanisms, transformers, and discrete reparameterization. They motivate attention for modeling point relations and explain Gumbel-Softmax as a continuous relaxation for discrete sampling.
- Direct point-cloud networks should process size-varying and permutation-invariant sets, while early PointNet-style aggregation does not sufficiently capture point relations.
- PointNet++, Kd-Net, and DGCNN use hierarchical, tree-based, or dynamic-graph structures, but some methods do not satisfy both theoretical conditions simultaneously.
- Self-Attention: Attention uses input-dependent weights to linearly combine inputs, and self-attention lets inputs attend to every input element.
- Self-Attention: Multi-Head Attention aggregates inputs repeatedly under different linear transformations, increasing single-layer capacity over vanilla self-attention.
- Discrete Reparameterization: Gumbel-Softmax provides a continuous relaxation that enables smooth gradient estimation for categorical variables during training.
- Discrete Reparameterization: As τ approaches zero, the soft sample degenerates into Gumbel-Max; training uses differentiable samples, while testing uses discrete samples.
3. Point Attention Transformers
Point Attention Transformers combine ARPE embeddings with parameter-efficient Group Shuffle Attention to model point relations while preserving permutation equivariance. Classification adds hierarchical FPS or GSS down-sampling, whereas segmentation uses GSA layers without down-sampling.
- Overview: PATs embed each point with its absolute position and the relative positions of other points through the ARPE module.To reduce O(N^2) cost, ARPE uses sampled nearest neighbors with a dilated sampling rate for varying point counts.
- Classification: Classification alternates GSA blocks with FPS or GSS down-sampling, then applies shared MLPs and global average pooling to produce one class label.The element-wise loss trick computes cross-entropy over each MLP output before averaging.
- Segmentation: Segmentation omits down-sampling because it would lose information, using stacked GSA layers and a shared MLP to assign a label to every point.The output is pointwise, with one label per input point.
- Group Shuffle Attention: GSA models relations between points using non-linear self-attention, compact group transformations, channel shuffle, residual connections, and group normalization.Channel shuffle enables information flow between groups while retaining the multi-head design.
- Group Shuffle Attention: GSA is permutation-equivariant: applying a permutation to the input produces the same permutation of the output.The paper states this as GSA(P · X) = P · GSA(X).
- Gumbel Subset Sampling: GSS learns representative subsets end to end, using soft virtual points during training and hard discrete selections during testing.It is designed to be permutation-invariant, task-agnostic, and differentiable, unlike the stated limitations of FPS.
4. Applications
PATs achieve competitive classification with parameter-efficient models, strong segmentation trade-offs, and state-of-the-art event-camera gesture recognition while meeting real-time latency constraints.
- ModelNet40 Shape Classification: PAT models provide competitive accuracy with parameter efficiency and acceptable forward speed; down-sampling improves performance while reducing computation cost.GSS further improves FPS with negligible computational burden.
- S3DIS Indoor Scene Segmentation: On S3DIS, PAT achieves the best trade-off between segmentation performance and model size, outperforming all comparisons on Area 5 and all but PointCNN on AREAS mIoU.The method has significantly smaller model size on AREAS.
- S3DIS Indoor Scene Segmentation: Compared with PointCNN, PAT achieves higher mAcc on AREAS and higher mIoU and mAcc on Area 5, with better IoUs for most classes on Area 5.The comparison reports per-class IoU, mIoU, and mAcc across both evaluation settings.
- Event Camera Stream as Point Clouds: On DVS128 Gesture, PAT achieves state-of-the-art performance with strong parameter efficiency, and all point-cloud approaches remain within the 100ms maximum decision latency.A 256-point PAT performs similarly to its 1,024-point version; the dataset contains event-camera gesture streams represented as spatio-temporal point clouds.
5. Ablation Study
Ablations identify the components and design choices that improve PAT efficiency or accuracy. Grouping and channel shuffle are important for GSA, while element-wise loss mainly accelerates training.
- Element-wise Loss: Element-wise loss accelerates training but does not improve test-set performance.The comparison uses ModelNet40 training loss and evaluation results.
- Number of Groups: Eight groups provide the best trade-off between accuracy and model size, whereas one group significantly reduces accuracy and increases model size.Grouping is coupled to the multi-head attention design.
- Channel Shuffle vs. no Shuffle: Channel shuffle enables information flow across groups with negligible computational cost.The operation is parameter-free.
- Embedding Layer: ARPE consistently improves performance by approximately 0.8% on both 256- and 1,024-point inputs.The comparison is reported in Table 6 on ModelNet40.
- GSA vs. MHA: GSA outperforms comparable MHA variants in both parameter efficiency and accuracy.The comparison includes MHA with GSA's hidden size and MHA tuned to a comparable model size.
- Group Norm vs. Layer Norm: Group normalization with eight groups consistently outperforms layer normalization.The comparison is reported for the normalization variants in Table 6.
6. Conclusion
The paper presents PATs with GSA for point-cloud relations and GSS for learnable, task-agnostic subset selection. Results across benchmarks support the effectiveness and efficiency of these methods.
- 6. Conclusion: PATs use parameter-efficient GSA to learn relations between points.GSA is introduced as part of the point-cloud reasoning architecture.
- 6. Conclusion: GSS provides an end-to-end learnable and task-agnostic operation for sampling point subsets.The paper identifies GSS as a separate sampling contribution alongside PATs.
- 6. Conclusion: Results on several benchmarks demonstrate the effectiveness and efficiency of the proposed methods.The conclusion states this aggregate outcome without giving a single metric.
A. Proof of Permutation Equivariance of Group Shuffle Attention
The proof establishes that self-attention and GSA preserve point ordering transformations. Consequently, permuting the input points permutes the output in the same way.
- Self-Attention: Self-attention uses softmax over transformed attention scores, and softmax commutes with simultaneous row and column permutation.This yields P · softmax(A) · P^T = softmax(P · A · P^T).
- Non-linear Self-Attention: Non-linear self-attention is permutation-equivariant when the query is the input itself.For self-attention, Q = X.
- Non-linear Self-Attention: Element-wise nonlinearities preserve permutation equivariance because applying them after permutation equals permuting their outputs.The proof states σ(P · X) = P · σ(X).
- Group Shuffle Attention: GSA is permutation-equivariant: GSA(P · X) = P · GSA(X).The proposition states the result for any permutation matrix P.
- Group Shuffle Attention: GSA's element-wise operations do not alter the permutation-equivariance inherited from Attnσ.This connects the architectural construction to the preceding self-attention proof.
B. Proof of Permutation Invariance of Gumbel Subset Sampling
The proof shows that GSS is unaffected by permutations of the input points. Element-wise operations introduced by Gumbel-Softmax preserve this permutation property.
- Permutation Invariance: GSS is permutation-invariant: GSS(P · X) = GSS(X).The proposition applies to any permutation matrix P.
- Permutation Invariance: Gumbel-Softmax preserves permutation invariance because its added operations are element-wise.The proof applies this property to the sampling operation.