Elamur Research⌕
☰
home / papers / 2024 / B'MOJO: Hybrid State Space Realizations of Foundation Models with Eidetic and Fading Memory

B'MOJO: Hybrid State Space Realizations of Foundation Models with Eidetic and Fading Memory

Paper · 2024 · Luca Zancato, Arjun Seshadri, Yonatan Dukler et al.

Summary

Current sequence models store the past in one of two ways. Transformers keep an exact (eidetic) record of recent tokens, but only within a bounded context. State Space Models compress the whole history into a lossy (fading) state. Hybrids such as Jamba combine both, but they give little control over the balance and cannot widen the exact-memory span.

Building on Stochastic Realization Theory, the authors propose B'MOJO, a stackable module that unifies both memory types. It can also retrieve from an external, asynchronously updated store. Transformers, Mamba and Jamba turn out to be special cases of it. Instead of keeping tokens by recency, B'MOJO keeps the ones the model fails to predict well, a mechanism called Innovation Selection.

On synthetic associative recall, B'MOJO beats SSM and hybrid baselines. In language modeling up to 1.4B parameters, it matches the perplexity of similar-sized Transformers and SSMs while training faster. Models trained on sequences of at most 8K tokens keep improving on inputs up to 32K tokens. This matters for building models that use long, non-stationary inputs within fixed compute budgets.

Key points

  • The paper frames memory as four kinds: short-term eidetic "in-context", fading "in-state", structural "in-weights", and long-term eidetic "in-storage" (retrieval).
  • The speed gain is reported inconsistently: the abstract says up to 10% faster training, while the contributions section says 15%. The throughput figure covers 130M to 1.4B models compared with Mamba and Transformer implementations.
  • On long-context benchmarks (PG-19, SWDE, Scrolls-QAsper and Scrolls-NarraQA), the B'MOJO variants outperform Mamba. B'MOJO also beats the hybrid sliding-attention + SSM baseline on PG-19 and SWDE.
  • B'MOJO uses a 512-token sliding window versus Mistral's 1024. On short and medium-context tasks it still beats Mamba at 1.4B and roughly matches Mistral.
  • On length generalization, 790M models trained on 8K tokens improve up to 32K. Smaller models (up to 370M) trained on 16K tokens generalize to 64K. Transformer baselines fail this test.
  • Limitations: experiments stop at 1.4B parameters, and the Mamba-style fading component may generalize to longer lengths less well at larger scale.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksEngineering · 206 worksControl and Systems Engineering · 72 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 4,145 works

Cite

Cite as Luca Zancato, Arjun Seshadri, Yonatan Dukler, Aditya Golatkar, Yantao Shen, Benjamin Bowman, Matthew Trager, Alessandro Achille and Stefano Soatto. 2024. B'MOJO: Hybrid State Space Realizations of Foundation Models with Eidetic and Fading Memory. arXiv:2407.06324.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2407.06324 · arXiv 2407.06324

@misc{zancato2024b,
  title = {{B'MOJO: Hybrid State Space Realizations of Foundation Models with Eidetic and Fading Memory}},
  author = {Luca Zancato, Arjun Seshadri, Yonatan Dukler, Aditya Golatkar, Yantao Shen, Benjamin Bowman, Matthew Trager, Alessandro Achille and Stefano Soatto},
  year = {2024},
  doi = {10.48550/arxiv.2407.06324},
  eprint = {2407.06324},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2407.06324}
}

References

53 references, 12 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Alessandro Achille and Stefano Soatto. On the learnability of physical concepts: Can a neural network understand what’s real? arXiv preprint arXiv:2207.12186, 2022. — [arXiv:2207.12186](https://arxiv.org/abs/2207.12186) find
  2. Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall-throughput tradeoff. arXiv preprint arXiv:2402.18668, 2024.
  3. Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. arXiv preprint arXiv:2405.04517, 2024.
  4. Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432–7439, 2020.
  5. R Brockett. On the algebraic structure of bilinear systems. theory and applica-tions of variable structure systems, 1972. find
  6. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.
  7. Olivier Chapelle, Vladimir Vapnik, and Jason Weston. Transductive inference for estimating values of functions. Advances in Neural Information Processing Systems, 12, 1999. find
  8. Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457, 2018.
  9. Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344–16359, 2022. find
  10. Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402.19427, 2024. find
  11. Paolo D’Alessandro, Alberto Isidori, and Antonio Ruberti. Realization and structure theory of bilinear dynamical systems. SIAM Journal on Control, 12(3):517–535, 1974. find
  12. David LeRoy Elliott. Bilinear control systems: matrices in action, volume 169. Springer, 2009. find
  13. B. Bowman et al. Inductive learning of transductive inference:. ArXiv [pending], 2021. find
  14. Alessandro Favero, Luca Zancato, Matthew Trager, Siddharth Choudhary, Pramuditha Perera, Alessandro Achille, Ashwin Swaminathan, and Stefano Soatto. Multi-modal hallucination control by visual information grounding. arXiv preprint arXiv:2403.14003, 2024. — [arXiv:2403.14003](https://arxiv.org/abs/2403.14003) find
  15. Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212.14052, 2022. find
  16. Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Ani… find
  17. Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023. find
  18. Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022.
  19. Alberto Isidori. Nonlinear control systems: an introduction. Springer, 1985. find
  20. Andrew H Jazwinski. Stochastic processes and filtering theory. Courier Corporation, 2007. find
  21. Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023. find
  22. Thomas Kailath. Linear systems, volume 156. Prentice-Hall Englewood Cliffs, NJ, 1980. find
  23. Rudolph Emil Kalman. A new approach to linear filtering and prediction problems, 1960. find
  24. Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.
  25. Arthur J Krener. Bilinear and nonlinear realizations of input-output maps. SIAM Journal on Control, 13(4):827–834, 1975. find
  26. Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, et al. Jamba: A hybrid transformer-mamba language model. arXiv preprint arXiv:2403.19887, 2024.
  27. Anders Lindquist and Giorgio Picci. On the stochastic realization problem. SIAM Journal on Control and Optimization, 17(3):365–389, 1979. find
  28. Greta M Ljung and George EP Box. On a measure of lack of fit in time series models. Biometrika, 65(2):297–303, 1978. find
  29. Lennart Ljung. System identification toolbox: User’s guide. Citeseer, 1995. find
  30. Xuezhe Ma, Xiaomeng Yang, Wenhan Xiong, Beidi Chen, Lili Yu, Hao Zhang, Jonathan May, Luke Zettlemoyer, Omer Levy, and Chunting Zhou. Megalodon: Efficient llm pretraining and inference with unlimited context length. arXiv preprint arXiv:2404.08801, 2024. find
  31. Mireille Chaleyat Maurel and Dominique Michel. Results on the non existence of finite dimensional filters. Stochastics: An International Journal of Probability and Stochastic Processes, 13(1-2):83–102, 1984. find
  32. Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random-access infinite context length for transformers. arXiv preprint arXiv:2305.16300, 2023.
  33. Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amode… — In-context Learning and Induction Heads
  34. Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. arXiv preprint arXiv:1606.06031, 2016.
  35. Michael Poli, Armin W Thomas, Eric Nguyen, Pragaash Ponnusamy, Björn Deiseroth, Kristian Kersting, Taiji Suzuki, Brian Hie, Stefano Ermon, Christopher Ré, et al. Mechanistic design and scaling of hybrid architectures. arXiv preprint arXiv:2403.17844, 2024. find
  36. Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99–106, 2021. find
  37. Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. SCROLLS: Standardized CompaRison over long language sequences. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 12…
  38. RJ Solmonoff. A formal theory of inductive inference. i. II Information and Control, 7:224–254, 1964. find
  39. Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei Efros, and Moritz Hardt. Test-time training with self-supervision for generalization under distribution shifts. In International conference on machine learning, pages 9229–9248. PMLR, 2020. find
  40. Peter Van Overschee and Bart De Moor. Subspace identification for linear systems: Theory—Implementation—Applications. Springer Science & Business Media, 2012. find
  41. Vladimir Vapnik. Transductive inference and semi-supervised learning, 2006. find
  42. Terry A. Welch. A technique for high-performance data compression. Computer, 17(06):8–19, 1984. find
  43. Norbert Wiener. Extrapolation, interpolation, and smoothing of stationary time series: with engineering applications. The MIT press, 1949. find
  44. William M Wonham. On a matrix riccati equation of stochastic control. SIAM Journal on Control, 6(4):681–697, 1968. find
  45. Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, 2023. find
  46. Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06635, 2023. find
  47. Lili Yu, Dániel Simig, Colin Flaherty, Armen Aghajanyan, Luke Zettlemoyer, and Mike Lewis. Megabyte: Predicting million-byte sequences with multiscale transformers. Advances in Neural Information Processing Systems, 36, 2024. find
  48. L. Zancato, A. Achille, G. Paolini, A. Chiuso, and S. Soatto. Stacked residuals of dynamic layers for time series anomaly detection, 2022. find
  49. Luca Zancato, Alessandro Achille, Tian Yu Liu, Matthew Trager, Pramuditha Perera, and Stefano Soatto. Train/test-time adaptation with retrieval. In 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 15911–15921, 2023. find
  50. Luca Zancato and Alessandro Chiuso. A novel deep neural network architecture for non-linear system identification. IFAC-PapersOnLine, 54(7):186–191, 2021. 19th IFAC Symposium on System Identification SYSID 2021. find
  51. Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:1905.07830, 2019. find
  52. Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107– 115, 2021. find
  53. Jacob Ziv and Abraham Lempel. Compression of individual sequences via variable-rate coding. IEEE transactions on Information Theory, 24(5):530–536, 1978. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/zancato-2024-b-mojo-hybrid-state-space/