Natively Trainable Sparse Attention for Hierarchical Point Cloud Datasets

TL;DR

This paper introduces a novel approach combining the Erwin architecture with Native Sparse Attention to enable efficient, scalable transformer models for large physical datasets, achieving comparable or better performance.

Contribution

It adapts the NSA mechanism for non-sequential data and demonstrates its effectiveness on diverse large-scale physical science datasets.

Findings

Achieves performance comparable or superior to original Erwin model

Successfully adapts NSA for non-sequential data

Validates results by reproducing original experiments

Abstract

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.