Elamur Research⌕
☰
home / papers / 2024 / Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks

Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks

Paper · 2024 · Jongho Park, Jaeseung Park, Zheyang Xiong et al.

Summary

In-context learning (ICL) has mostly been studied in Transformers. This paper asks whether attention-free state-space models (SSMs), chiefly Mamba and also S4, can acquire the same ability. Large-scale ICL is hard to test cheaply, so the authors follow the Garg et al. (2022) setup: small models are trained from scratch to solve synthetic function classes from in-prompt examples. The suite covers several regression variants, decision trees, sparse parity, Chain-of-Thought I/O and vector-valued MQAR (a retrieval task). It is complemented by formal-language benchmarks (GINC, RegBench).

Mamba roughly matches Transformers on standard regression and solves sparse parity, which Transformers fail. However, it lags on decision trees and breaks down on MQAR retrieval. To combine both strengths, the authors propose MambaFormer, a hybrid that substitutes Mamba blocks for MLP blocks, begins with a Mamba layer and needs no positional encoding. It performs competitively across every task. The work suggests SSMs face no fundamental barrier to ICL apart from retrieval, and it points to hybrids as a promising design direction.

Key points

  • Nine synthetic ICL tasks (e.g., linear, sparse linear and 2NN regression, decision tree, outlier regression, sparse parity, vector MQAR) were used. Models were mostly trained for 500,000 iterations.
  • On sparse parity, Transformers did no better than random guessing, while Mamba succeeded. On vector-valued MQAR, Mamba could not reliably retrieve values.
  • On RegBench after 120 epochs, accuracy was: Transformer 98.9, Mamba 87.8, MambaFormer 99.8, Standard Hybrid 99.9. The hybrids also converged faster.
  • On GINC, ICL accuracy was: LSTM 96.4, Mamba 87.1, Standard Hybrid 85.0, Transformer 84.2, MambaFormer 79.6.
  • Code is available at github.com/krafton-ai/mambaformer-icl. The authors note that the study is limited to small models and mostly non-language tasks.

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

 Institutions

🎓University of MichiganAnn Arbor · United States94 works · 18,548 citations🎓Seoul National UniversitySeoul · South Korea47 works · 46,072 citations🎓University of Wisconsin–MadisonMadison · United States60 works · 49,213 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksArts and Humanities · 17,959 worksLanguage and Linguistics · 17,789 works

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

Cite

Cite as Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee and Dimitris Papailiopoulos. 2024. Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks. arXiv:2402.04248.

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

@misc{park2024can,
  title = {{Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks}},
  author = {Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee and Dimitris Papailiopoulos},
  year = {2024},
  doi = {10.48550/arxiv.2402.04248},
  eprint = {2402.04248},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2402.04248}
}

References

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

  1. Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. Transformers learn to implement preconditioned gradient descent for in-context learning. arXiv preprint arXiv:2306.00297, 2023. 3
  2. Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2022. 1, 3
  3. Akyürek, E., Wang, B., Kim, Y., and Andreas, J. In-context language learning: Architectures and algorithms. arXiv preprint arXiv:2401.12973, 2024. URL https://arxiv.org/abs/2401. 12973. 3, 4, 6, 13, 15 ↗
  4. Arora, S., Eyuboglu, S., Timalsina, A., Johnson, I., Poli, M., Zou, J., Rudra, A., and Ré, C. Zoology: Measuring and improving recall in efficient language models. arXiv preprint arXiv:2312.04927, 2023. URL https://arxiv.org/abs/2312.04927. 3, 4, 6, 10 ↗
  5. Bai, Y., Chen, F., Wang, H., Xiong, C., and Mei, S. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306.04637, 2023. 1 — [arXiv:2306.04637](https://arxiv.org/abs/2306.04637) find
  6. Barak, B., Edelman, B., Goel, S., Kakade, S., Malach, E., and Zhang, C. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems, 35:21750–21764, 2022. 12 find
  7. Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020. 3
  8. Bhattamishra, S., Ahuja, K., and Goyal, N. On the ability and limitations of transformers to recognize formal languages. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 7096–7116. Association for Computational Linguistics, 2020. doi: 10.18653/v1…
  9. Bhattamishra, S., Patel, A., Blunsom, P., and Kanade, V. Understanding in-context learning in transformers and llms by learning to learn discrete functions. arXiv preprint arXiv:2310.03016, 2023. URL https://arxiv.org/abs/2310.03016. 3, 6, 10, 12, 24 — [arXiv:2310.03016](https://arxiv.org/abs/2310.03016) ↗
  10. Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020. 1
  11. Chan, S., Santoro, A., Lampinen, A., Wang, J., Singh, A., Richemond, P., McClelland, J., and Hill, F. Data distributional properties drive emergent in-context learning in transformers. Advances in Neural Information Processing Systems, 35:18878–18891, 2022. 1 find
  12. Dai, D., Sun, Y., Dong, L., Hao, Y., Ma, S., Sui, Z., and Wei, F. Why can gpt learn in-context? language models secretly perform gradient descent as meta-optimizers. In Findings of the Association for Computational Linguistics: ACL 2023, pp. 4005–4019, 2023. 3
  13. Dao, T., Fu, D. Y., Saab, K. K., Thomas, A. W., Rudra, A., and Ré, C. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212.14052, 2022. URL https://arxiv.org/abs/2212.14052. 2, 4 ↗
  14. Deletang, G., Ruoss, A., Grau-Moya, J., Genewein, T., Wenliang, L. K., Catt, E., Cundy, C., Hutter, M., Legg, S., Veness, J., et al. Neural networks and the chomsky hierarchy. In The Eleventh International Conference on Learning Representations, 2022. 14
  15. Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C. The pile: An 800gb dataset of diverse text for language modeling, 2020. 2 find
  16. Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35: 30583–30598, 2022. 1, 2, 3, 5, 20, 24
  17. Giannou, A., Rajput, S., Sohn, J.-y., Lee, K., Lee, J. D., and Papailiopoulos, D. Looped transformers as programmable computers. In International Conference on Machine Learning, pp. 11398–11442. PMLR, 2023. 1 find
  18. Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023. URL https://arxiv.org/abs/2312.00752. 1, 2, 4, 5, 9, 11, 12, 15, 20 ↗
  19. Gu, A., Goel, K., Gupta, A., and Ré, C. On the parameterization and initialization of diagonal state space models. In Advances in Neural Information Processing Systems, volume 35, pp. 35971–35983, 2022a. 3, 4 find
  20. Gu, A., Goel, K., and Ré, C. Efficiently modeling long sequences with structured state spaces. In The Tenth International Conference on Learning Representations, ICLR 2022, 2022b. URL https://openreview.net/forum?id=uYLFoz1vlAC. 2, 3 ↗
  21. Hahn, M. Theoretical limitations of self-attention in neural sequence models. Transactions of the Association for Computational Linguistics, 8:156–171, 2020. 10 find
  22. Hahn, M. and Goyal, N. A theory of emergent in-context learning as implicit structure induction. arXiv preprint arXiv:2303.07971, 2023. URL https://arxiv.org/abs/2303.07971. 13 — [arXiv:2303.07971](https://arxiv.org/abs/2303.07971) ↗
  23. Hochreiter, S. and Schmidhuber, J. Long short-term memory. Neural computation, 9(8):1735–1780, 1997. 4 find
  24. Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020. 1, 5, 21
  25. Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are rnns: Fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, volume 119 of Proceedings of Machine Learning Research, pp. 5156–5165, 2020. URL… find
  26. Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv:1412.6980, 2014. 5, 20 arXiv preprint Kol, G., Raz, R., and Tal, A. Time-space hardness of learning sparse parities. In Proceedings of the 49th Annual ACM SIGACT Symposium on Theory of Computing, pp. 1067–1080, 2017. 10
  27. Li, Y., Ildiz, M. E., Papailiopoulos, D., and Oymak, S. Transformers as algorithms: Generalization and stability in in-context learning, 2023a. 1 find
  28. Li, Y., Sreenivasan, K., Giannou, A., Papailiopoulos, D., and Oymak, S. Dissecting chain-of-thought: Compositionality through in-context filtering and learning. In Thirty-seventh Conference on Neural Information Processing Systems, 2023b. 6, 21 find
  29. Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., and Zhang, C. Transformers learn shortcuts to automata. In The Eleventh International Conference on Learning Representations, 2022. 1 find
  30. Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., and Zhang, C. Exposing attention glitches with flip-flop language modeling. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. 4 find
  31. Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. In International Conference on Learning Representations, 2018. 22 find
  32. Mahankali, A., Hashimoto, T. B., and Ma, T. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.03576, 2023. 3
  33. Mehta, H., Gupta, A., Cutkosky, A., and Neyshabur, B. Long range language modeling via gated state spaces. arXiv preprint arXiv:2206.13947, 2022. 12 — [arXiv:2206.13947](https://arxiv.org/abs/2206.13947) find
  34. Min, S., Lewis, M., Zettlemoyer, L., and Hajishirzi, H. Metaicl: Learning to learn in context. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 2791–2809, 2022a. 1, 3 find
  35. Min, S., Lyu, X., Holtzman, A., Artetxe, M., Lewis, M., Hajishirzi, H., and Zettlemoyer, L. Rethinking the role of demonstrations: What makes in-context learning work? arXiv preprint arXiv:2202.12837, 2022b. URL https://arxiv.org/abs/2202.12837. 1 ↗
  36. Muennighoff, N., Rush, A. M., Barak, B., Scao, T. L., Tazi, N., Piktus, A., Pyysalo, S., Wolf, T., and Raffel, C. Scaling data-constrained language models. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum?id= j5BuTrEj35. 1, 21 ↗
  37. Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., and et al., A. C. In-context learning and induction heads. arXiv preprint arXiv:2209.11895, 2022. URL https://arxiv.org/abs/2209.11895. 3, 4, 6 ↗
  38. Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Cao, H., Cheng, X., Chung, M., Grella, M., Kiran, K. G., et al. Rwkv: Reinventing rnns for the transformer era. arXiv preprint arXiv:2305.13048, 2023. URL https://arxiv.org/abs/2305.13048. 2, 4 ↗
  39. Pilault, J., Fathi, M., Firat, O., Pal, C., Bacon, P.-L., and Goroshin, R. Block-state transformers. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. 12
  40. Poli, M., Massaroli, S., Nguyen, E., Fu, D. Y., Dao, T., Baccus, S., Bengio, Y., Ermon, S., and Ré, C. Hyena hierarchy: Towards larger convolutional language models. arXiv preprint arXiv:2302.10866, 2023. URL https://arxiv.org/abs/2302.10866. 2, 4 ↗
  41. Press, O., Zhang, M., Min, S., Schmidt, L., Smith, N. A., and Lewis, M. Measuring and narrowing the compositionality gap in language models. arXiv preprint arXiv:2210.03350, 2022. 12 find
  42. Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019. 20 find
  43. Ravi, S. and Larochelle, H. Optimization as a model for few-shot learning. In International conference on learning representations, 2016. 1, 3 find
  44. Raz, R. Fast learning requires good memory: A time-space lower bound for parity learning. In 2016 IEEE 57th Annual Symposium on Foundations of Computer Science (FOCS), pp. 266–275. find
  45. Schaeffer, R., Miranda, B., and Koyejo, S. Are emergent abilities of large language models a mirage? In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum?id=ITw9edRDlD. 1 ↗
  46. Schmidhuber, J., Zhao, J., and Wiering, M. Shifting inductive bias with success-story algorithm, adaptive levin search, and incremental self-improvement. Machine Learning, 28:105–130, 1997. 1 find
  47. Sun, Y., Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, F. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621, 2023. URL https://arxiv.org/abs/2307.08621. 2, 4 ↗
  48. Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023. 22
  49. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, pp. 5998–6008, 2017. URL https://proceedin…
  50. PMLR, 2023a. 1 von Oswald, J., Niklasson, E., Schlegel, M., Kobayashi, S., Zucchet, N., Scherrer, N., Miller, N., Sandler, M., Vladymyrov, M., Pascanu, R., et al. Uncovering mesa-optimization algorithms in transformers. arXiv preprint arXiv:2309.05858, 2023b. URL https://arxiv.org/abs/2309. 05858. 1… — [arXiv:2309.05858](https://arxiv.org/abs/2309.05858) ↗
  51. Wang, S., Li, B. Z., Khabsa, M., Fang, H., and Ma, H. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020. 3 find
  52. Wei, J., Tay, Y., Bommasani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M., Zhou, D., Metzler, D., et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.07682, 2022. 1
  53. Xie, S. M., Raghunathan, A., Liang, P., and Ma, T. An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations, 2021. 3, 4, 6, 13, 14
  54. Yang, L., Lee, K., Nowak, R., and Papailiopoulos, D. Looped transformers are better at learning learning algorithms, 2023a. 1 find
  55. Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, Y. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06635, 2023b. 2, 4 find
  56. Yu, W., Luo, M., Zhou, P., Si, C., Zhou, Y., Wang, X., Feng, J., and Yan, S. Metaformer is actually what you need for vision. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10819–10829, 2022. 10 find
  57. Zhai, S., Talbott, W., Srivastava, N., Huang, C., Goh, H., Zhang, R., and Susskind, J. An attention free transformer. arXiv preprint arXiv:2105.14103, 2021. 2 — [arXiv:2105.14103](https://arxiv.org/abs/2105.14103) find
  58. Zhou, H., Bradley, A., Littwin, E., Razin, N., Saremi, O., Susskind, J., Bengio, S., and Nakkiran, P. What algorithms can transformers learn? a study in length generalization. arXiv preprint arXiv:2310.16028, 2023. 1
  59. Zuo, S., Liu, X., Jiao, J., Charles, D., Manavoglu, E., Zhao, T., and Gao, J. Efficient long sequence modeling via state space augmented transformer. arXiv preprint arXiv:2212.08136, 2022. 12 — [arXiv:2212.08136](https://arxiv.org/abs/2212.08136) find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/park-2024-mamba-learn-learn-comparative-study/