2024

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

Park, Jongho, Park, Jaeseung, Xiong, Zheyang et al.

Understand

State-space models (SSMs), such as Mamba (Gu & Dao, 2023), have been proposed as alternatives to Transformer networks in language modeling, by incorporating gating, convolutions, and input-dependent token selection to mitigate the quadratic cost of multi-head attention.

Although SSMs exhibit competitive performance, their in-context learning (ICL) capabilities, a remarkable emergent property of modern language models that enables task execution without parameter optimization, remain underexplored compared to Transformers.
In this study, we evaluate the ICL performance of SSMs, focusing on Mamba, against Transformer models across various tasks.
Our results show that SSMs perform comparably to Transformers in standard regression ICL tasks, while outperforming them in tasks like sparse parity learning.

Built on

Long short-term memory
Hochreiter, S. and Schmidhuber, J · 1997
Earlier work this paper cites.
Shifting inductive bias with success-story algorithm, adaptive levin search, and incremental self-improvement
Schmidhuber, J., Zhao, J., and Wiering, M · 1997
Earlier work this paper cites.
Adam: A method for stochastic optimization
Original
Kingma, D. P. and Ba, J · 2014
Earlier work this paper cites.
Optimization as a model for few-shot learning
Ravi, S. and Larochelle, H · 2016
Earlier work this paper cites.
Fast learning requires good memory: A time-space lower bound for parity learning
Raz, R · 2016
Earlier work this paper cites.
Time-space hardness of learning sparse parities
Kol, G., Raz, R., and Tal, A · 2017
Earlier work this paper cites.
Attention is all you need
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I · 2017
Earlier work this paper cites.
Decoupled weight decay regularization
Loshchilov, I. and Hutter, F · 2018
Earlier work this paper cites.
Language models are unsupervised multitask learners
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al · 2019
Earlier work this paper cites.
Longformer: The long-document transformer
Original
Beltagy, I., Peters, M. E., and Cohan, A · 2020
Earlier work this paper cites.
On the ability and limitations of transformers to recognize formal languages
Bhattamishra, S., Ahuja, K., and Goyal, N · 2020
Earlier work this paper cites.
Language models are few-shot learners
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al · 2020
Earlier work this paper cites.
The pile: An 800gb dataset of diverse text for language modeling, 2020
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 · 2020
Earlier work this paper cites.
Theoretical limitations of self-attention in neural sequence models
Hahn, M · 2020
Earlier work this paper cites.
Scaling laws for neural language models
Original
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D · 2020
Earlier work this paper cites.
Transformers are rnns: Fast autoregressive transformers with linear attention
Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F · 2020
Earlier work this paper cites.
Linformer: Self-attention with linear complexity
Original
Wang, S., Li, B. Z., Khabsa, M., Fang, H., and Ma, H · 2020
Earlier work this paper cites.
An explanation of in-context learning as implicit bayesian inference
Xie, S. M., Raghunathan, A., Liang, P., and Ma, T · 2021
Earlier work this paper cites.

Similar

An attention free transformer
Original
Zhai, S., Talbott, W., Srivastava, N., Huang, C., Goh, H., Zhang, R., and Susskind, J · 2021
Cited alongside, same era.
What learning algorithm is in-context learning? investigations with linear models
Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D · 2022
Cited alongside, same era.
Hidden progress in deep learning: Sgd learns parities near the computational limit
Barak, B., Edelman, B., Goel, S., Kakade, S., Malach, E., and Zhang, C · 2022
Cited alongside, same era.
Data distributional properties drive emergent in-context learning in transformers
Chan, S., Santoro, A., Lampinen, A., Wang, J., Singh, A., Richemond, P., McClelland, J., and Hill, F · 2022
Cited alongside, same era.
Hungry hungry hippos: Towards language modeling with state space models
Original
Dao, T., Fu, D. Y., Saab, K. K., Thomas, A. W., Rudra, A., and Ré, C · 2022
Cited alongside, same era.
Neural networks and the chomsky hierarchy
Deletang, G., Ruoss, A., Grau-Moya, J., Genewein, T., Wenliang, L. K., Catt, E., Cundy, C., Hutter, M., Legg, S., Veness, J., et al · 2022
Cited alongside, same era.
What can transformers learn in-context? a case study of simple function classes
Garg, S., Tsipras, D., Liang, P. S., and Valiant, G · 2022
Cited alongside, same era.
Efficiently modeling long sequences with structured state spaces
Gu, A., Goel, K., and Ré, C · 2022
Cited alongside, same era.
Transformers learn shortcuts to automata
Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., and Zhang, C · 2022
Cited alongside, same era.
Long range language modeling via gated state spaces
Original
Mehta, H., Gupta, A., Cutkosky, A., and Neyshabur, B · 2022
Cited alongside, same era.
Metaicl: Learning to learn in context
Min, S., Lewis, M., Zettlemoyer, L., and Hajishirzi, H · 2022
Cited alongside, same era.
In-context learning and induction heads
Original
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., and et al., A. C · 2022
Cited alongside, same era.

Then

Transformers as statisticians: Provable in-context learning with in-context algorithm selection
Original
Bai, Y., Chen, F., Wang, H., Xiong, C., and Mei, S · 2023
Later among the works it cites.
Understanding in-context learning in transformers and llms by learning to learn discrete functions
Original
Bhattamishra, S., Patel, A., Blunsom, P., and Kanade, V · 2023
Later among the works it cites.
Why can gpt learn in-context? language models secretly perform gradient descent as meta-optimizers
Dai, D., Sun, Y., Dong, L., Hao, Y., Ma, S., Sui, Z., and Wei, F · 2023
Later among the works it cites.
Looped transformers as programmable computers
Giannou, A., Rajput, S., Sohn, J.-y., Lee, K., Lee, J. D., and Papailiopoulos, D · 2023
Later among the works it cites.
Mamba: Linear-time sequence modeling with selective state spaces
Original
Gu, A. and Dao, T · 2023
Later among the works it cites.
A theory of emergent in-context learning as implicit structure induction
Original
Hahn, M. and Goyal, N · 2023
Later among the works it cites.
Exposing attention glitches with flip-flop language modeling
Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., and Zhang, C · 2023
Later among the works it cites.
One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention
Original
Mahankali, A., Hashimoto, T. B., and Ma, T · 2023
Later among the works it cites.
Scaling data-constrained language models
Muennighoff, N., Rush, A. M., Barak, B., Scao, T. L., Tazi, N., Piktus, A., Pyysalo, S., Wolf, T., and Raffel, C · 2023
Later among the works it cites.
Rwkv: Reinventing rnns for the transformer era
Original
Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Cao, H., Cheng, X., Chung, M., Grella, M., Kiran, K. G., et al · 2023
Later among the works it cites.
Block-state transformers
Pilault, J., Fathi, M., Firat, O., Pal, C., Bacon, P.-L., and Goroshin, R · 2023
Later among the works it cites.
Hyena hierarchy: Towards larger convolutional language models
Original
Poli, M., Massaroli, S., Nguyen, E., Fu, D. Y., Dao, T., Baccus, S., Bengio, Y., Ermon, S., and Ré, C · 2023
Later among the works it cites.
Are emergent abilities of large language models a mirage?
Schaeffer, R., Miranda, B., and Koyejo, S · 2023
Later among the works it cites.
Retentive network: A successor to transformer for large language models
Original
Sun, Y., Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, F · 2023
Later among the works it cites.
Llama: Open and efficient foundation language models
Original
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al · 2023
Later among the works it cites.
What algorithms can transformers learn? a study in length generalization
Original
Zhou, H., Bradley, A., Littwin, E., Razin, N., Saremi, O., Susskind, J., Bengio, S., and Nakkiran, P · 2023
Later among the works it cites.
In-context language learning: Architectures and algorithms
Original
Akyürek, E., Wang, B., Kim, Y., and Andreas, J · 2024
Closest in time.

Beyond the bibliography

alphaXiv searches the wider corpus for related work and actual follow-ups.

Open on alphaXiv

alphaXiv is searching for related work…

Reading the bibliography…