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

TL;DR

STREAM introduces a universal state-space model that explicitly encodes geometric structure, enabling efficient processing of sparse geometric data and achieving state-of-the-art results across various benchmarks.

Contribution

The paper proposes a novel geometric encoding in state-space models, improving performance on sparse geometric data tasks without complex preprocessing.

Findings

Achieves 100% accuracy on DVS128 Gestures dataset.

Improves baseline performance on ModelNet40 and ScanObjectNN datasets.

Efficiently computes interactions among points in O(N) steps.

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.