Saved in:
| Main Authors: | , , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2405.06147 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866913372886269952 |
|---|---|
| author | Parnichkun, Rom N. Massaroli, Stefano Moro, Alessandro Smith, Jimmy T. H. Hasani, Ramin Lechner, Mathias An, Qi Ré, Christopher Asama, Hajime Ermon, Stefano Suzuki, Taiji Yamashita, Atsushi Poli, Michael |
| author_facet | Parnichkun, Rom N. Massaroli, Stefano Moro, Alessandro Smith, Jimmy T. H. Hasani, Ramin Lechner, Mathias An, Qi Ré, Christopher Asama, Hajime Ermon, Stefano Suzuki, Taiji Yamashita, Atsushi Poli, Michael |
| contents | We approach designing a state-space model for deep learning applications through its dual representation, the transfer function, and uncover a highly efficient sequence parallel inference algorithm that is state-free: unlike other proposed algorithms, state-free inference does not incur any significant memory or computational cost with an increase in state size. We achieve this using properties of the proposed frequency domain transfer function parametrization, which enables direct computation of its corresponding convolutional kernel's spectrum via a single Fast Fourier Transform. Our experimental results across multiple sequence lengths and state sizes illustrates, on average, a 35% training speed improvement over S4 layers -- parametrized in time-domain -- on the Long Range Arena benchmark, while delivering state-of-the-art downstream performances over other attention-free approaches. Moreover, we report improved perplexity in language modeling over a long convolutional Hyena baseline, by simply introducing our transfer function parametrization. Our code is available at https://github.com/ruke1ire/RTF. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2405_06147 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | State-Free Inference of State-Space Models: The Transfer Function Approach Parnichkun, Rom N. Massaroli, Stefano Moro, Alessandro Smith, Jimmy T. H. Hasani, Ramin Lechner, Mathias An, Qi Ré, Christopher Asama, Hajime Ermon, Stefano Suzuki, Taiji Yamashita, Atsushi Poli, Michael Machine Learning Systems and Control We approach designing a state-space model for deep learning applications through its dual representation, the transfer function, and uncover a highly efficient sequence parallel inference algorithm that is state-free: unlike other proposed algorithms, state-free inference does not incur any significant memory or computational cost with an increase in state size. We achieve this using properties of the proposed frequency domain transfer function parametrization, which enables direct computation of its corresponding convolutional kernel's spectrum via a single Fast Fourier Transform. Our experimental results across multiple sequence lengths and state sizes illustrates, on average, a 35% training speed improvement over S4 layers -- parametrized in time-domain -- on the Long Range Arena benchmark, while delivering state-of-the-art downstream performances over other attention-free approaches. Moreover, we report improved perplexity in language modeling over a long convolutional Hyena baseline, by simply introducing our transfer function parametrization. Our code is available at https://github.com/ruke1ire/RTF. |
| title | State-Free Inference of State-Space Models: The Transfer Function Approach |
| topic | Machine Learning Systems and Control |
| url | https://arxiv.org/abs/2405.06147 |