State-Free Inference of State-Space Models: The *Transfer Function* Approach

Rom Parnichkun, Stefano Massaroli, Alessandro Moro, Jimmy T.H. Smith, Ramin Hasani, Mathias Lechner, Qi An, Christopher Re, Hajime Asama, Stefano Ermon, Taiji Suzuki, Michael Poli, Atsushi Yamashita
Proceedings of the 41st International Conference on Machine Learning, PMLR 235:39834-39860, 2024.

Abstract

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v235-parnichkun24a, title = {State-Free Inference of State-Space Models: The *{T}ransfer Function* Approach}, author = {Parnichkun, Rom and Massaroli, Stefano and Moro, Alessandro and Smith, Jimmy T.H. and Hasani, Ramin and Lechner, Mathias and An, Qi and Re, Christopher and Asama, Hajime and Ermon, Stefano and Suzuki, Taiji and Poli, Michael and Yamashita, Atsushi}, booktitle = {Proceedings of the 41st International Conference on Machine Learning}, pages = {39834--39860}, year = {2024}, editor = {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix}, volume = {235}, series = {Proceedings of Machine Learning Research}, month = {21--27 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v235/main/assets/parnichkun24a/parnichkun24a.pdf}, url = {https://proceedings.mlr.press/v235/parnichkun24a.html}, abstract = {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.} }
Endnote
%0 Conference Paper %T State-Free Inference of State-Space Models: The *Transfer Function* Approach %A Rom Parnichkun %A Stefano Massaroli %A Alessandro Moro %A Jimmy T.H. Smith %A Ramin Hasani %A Mathias Lechner %A Qi An %A Christopher Re %A Hajime Asama %A Stefano Ermon %A Taiji Suzuki %A Michael Poli %A Atsushi Yamashita %B Proceedings of the 41st International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2024 %E Ruslan Salakhutdinov %E Zico Kolter %E Katherine Heller %E Adrian Weller %E Nuria Oliver %E Jonathan Scarlett %E Felix Berkenkamp %F pmlr-v235-parnichkun24a %I PMLR %P 39834--39860 %U https://proceedings.mlr.press/v235/parnichkun24a.html %V 235 %X 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.
APA
Parnichkun, R., Massaroli, S., Moro, A., Smith, J.T., Hasani, R., Lechner, M., An, Q., Re, C., Asama, H., Ermon, S., Suzuki, T., Poli, M. & Yamashita, A.. (2024). State-Free Inference of State-Space Models: The *Transfer Function* Approach. Proceedings of the 41st International Conference on Machine Learning, in Proceedings of Machine Learning Research 235:39834-39860 Available from https://proceedings.mlr.press/v235/parnichkun24a.html.

Related Material