SOTAVerified

Flash STU: Fast Spectral Transform Units

2024-09-16Code Available1· sign in to hype

Y. Isabel Liu, Windsor Nguyen, Yagiz Devre, Evan Dogariu, Anirudha Majumdar, Elad Hazan

arXivPDF

Code Available — Be the first to reproduce this paper.

Reproduce

Code

Abstract

This paper describes an efficient, open source PyTorch implementation of the Spectral Transform Unit. We investigate sequence prediction tasks over several modalities including language, robotics, and simulated dynamical systems. We find that for the same parameter count, the STU and its variants outperform the Transformer as well as other leading state space models across various modalities.

Tasks

State Space Models

Reproductions