STREAM: A Universal State-Space Model for Sparse Geometric Data

TL;DR

This work proposes to encode geometric structure explicitly into the parameterization of a state-space model, and achieves competitive results on a range of benchmarks from point-cloud classification to event-based vision and audio classification.

Abstract

Handling sparse and unstructured geometric data, such as point clouds or event-based vision, is a pressing challenge in the field of machine vision. Recently, sequence models such as Transformers and state-space models entered the domain of geometric data. These methods require specialized preprocessing to create a sequential view of a set of points. Furthermore, prior works involving sequence models iterate geometric data with either uniform or learned step sizes, implicitly relying on the model to infer the underlying geometric structure. In this work, we propose to encode geometric structure explicitly into the parameterization of a state-space model. State-space models are based on linear dynamics governed by a one-dimensional variable such as time or a spatial coordinate. We exploit this dynamic variable to inject relative differences of coordinates into the step size of the state-space model. The resulting geometric operation computes interactions between all pairs of N points in O(N) steps. Our model deploys the Mamba selective state-space model with a modified CUDA kernel to efficiently map sparse geometric data to modern hardware. The resulting sequence model, which we call STREAM, achieves competitive results on a range of benchmarks from point-cloud classification to event-based vision and audio classification. STREAM demonstrates a powerful inductive bias for sparse geometric data by improving the PointMamba baseline when trained from scratch on the ModelNet40 and ScanObjectNN point cloud analysis datasets. It further achieves, for the first time, 100% test accuracy on all 11 classes of the DVS128 Gestures dataset.

Figures (5)

• Figure 1: A unified view on point cloud and event stream modeling. A. Adjacency structure of point based methods B. The coordinates of sparse geometric data are irregularly spacedC. Point clouds and event streams are ordered by spatial and temporal axes, respectively, before the state-space model.
• Figure 2: STREAM encodes geometric information into the SSM parameters. A. Point cloud input. B. The STREAM module converts relative differences in coordinates into the $\Delta_k$ scale of the SSM. C. A STREAM module integrates pairwise spatial relationships based on an exponentially oscillating kernel.
• Figure 3: Point cloud processing pipeline inspired by PointMamba liang2024pointmambasimplestatespace. FPS: Farthest point sampling, kNN: k-Nearest Neighbors.
• Figure 4: Hierarchical subsampling architectures with two stages to handle long event streams.
• Figure 5: Example of a Scan operating on four inputs. The left most column represents the input pairs $(e^{\mathbf{A}\Delta_i}, \mathbf{B}_i\mathbf{u}_i)$, and the right most column represents the final results.

Theorems & Definitions (3)

• Proposition 1
• proof
• Remark 1