Natively Trainable Sparse Attention for Hierarchical Point Cloud Datasets
This work addresses efficiency challenges in applying transformers to large-scale physical systems, though it appears incremental as it builds directly on the Erwin architecture.
The authors tackled the quadratic scaling problem of attention mechanisms in transformers for large physical systems by combining the Erwin architecture with Native Sparse Attention (NSA) adapted for non-sequential data, achieving performance that matches or exceeds the original Erwin model on cosmology, molecular dynamics, and air pressure datasets.
Unlocking the potential of transformers on datasets of large physical systems depends on overcoming the quadratic scaling of the attention mechanism. This work explores combining the Erwin architecture with the Native Sparse Attention (NSA) mechanism to improve the efficiency and receptive field of transformer models for large-scale physical systems, addressing the challenge of quadratic attention complexity. We adapt the NSA mechanism for non-sequential data, implement the Erwin NSA model, and evaluate it on three datasets from the physical sciences -- cosmology simulations, molecular dynamics, and air pressure modeling -- achieving performance that matches or exceeds that of the original Erwin model. Additionally, we reproduce the experimental results from the Erwin paper to validate their implementation.