Elamur Research⌕
☰
home / papers / 2026 / Understanding In-Context Learning Beyond Transformers: An Investigation of State Space and Hybrid Architectures

Understanding In-Context Learning Beyond Transformers: An Investigation of State Space and Hybrid Architectures

Paper · 2026 · Shenran Wang, Timothy Tin-Long Tse and Jian Zhu

Summary

Mechanistic work on in-context learning (ICL) has so far explained it mainly through particular attention heads in transformers, such as function vector (FV) heads and induction heads. It has been unclear whether state-space models (SSMs) and hybrid models rely on similar internals. The authors compare seven LLMs of roughly 1B parameters: transformers, the SSMs Mamba and Mamba2, and two hybrid designs, Hymba (parallel attention and Mamba) and Zamba2 (sequentially stacked layers).

Following Jin et al. (2025), they split knowledge-based ICL tasks into two groups:

  • Parametric knowledge retrieval, such as country–capital lookups.
  • Contextual knowledge understanding, such as hate-speech and stance classification.

They first run behavioral label-randomization probes, then use intervention methods: AIE-based FV head identification, steering and ablation.

All architectures behave similarly on the surface, but their mechanisms differ. FVs drive ICL in transformers, Mamba and the hybrids, yet seem less central in Mamba2. In hybrids, the important FVs sit mostly in self-attention layers. FVs matter mainly for retrieval-type tasks. The study shows that behavioral parity can hide divergent internals, so behavioral and mechanistic analyses should be combined.

Key points

  • Models: MAMBA-1.4B, MAMBA2-1.3B, GEMMA-3-1B-PT, LLAMA-3.2-1B, QWEN2.5-1.5B, HYMBA-1.5B-BASE and ZAMBA2-1.2B.
    • ZAMBA2-1.2B has 37 Mamba2 layers with 64 heads each and 6 attention layers with 32 heads each.
    • HYMBA-1.5B-BASE has 32 layers, each with 25 attention heads and 1 Mamba head.
  • Data:
    • 17 parametric-retrieval datasets adapted from Todd et al. (2024), with 30% held out for testing.
    • 16 contextual-understanding datasets from Min et al. (2022b), using their default Hugging Face splits.
  • Behavioral setup: k = 4, 8, 12, 16 or 32 demonstrations, averaged over 5 random seeds. All architectures were similarly robust to noisy labels.
  • Steering results:
    • Steering FV heads clearly helps on retrieval tasks but gives unstable gains on contextual tasks.
    • Steering the Mamba stream of either hybrid consistently hurts contextual-task performance.
    • Middle and later layers are the most effective steering targets.
  • Hymba specifics and code: In Hymba, the Mamba stream and non-FV heads still matter, because the two streams' outputs are averaged at each layer. Code is available at github.com/ShenranTomWang/ICL.

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 British ColumbiaVancouver · Canada98 works · 34,805 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 works

🧭 Topics 2Research subjects and languages the work is about.Language models & pretraining · 3,320 worksInterpretability · 130 works

Cite

Cite as Shenran Wang, Timothy Tin-Long Tse and Jian Zhu. 2026. Understanding In-Context Learning Beyond Transformers: An Investigation of State Space and Hybrid Architectures. Findings of the Association for Computational Linguistics: ACL 2026.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.findings-acl.39

@inproceedings{wang2026understanding,
  title = {{Understanding In-Context Learning Beyond Transformers: An Investigation of State Space and Hybrid Architectures}},
  author = {Shenran Wang, Timothy Tin-Long Tse and Jian Zhu},
  booktitle = {Findings of the Association for Computational Linguistics: ACL 2026},
  year = {2026},
  doi = {10.18653/v1/2026.findings-acl.39},
  url = {https://doi.org/10.18653/v1/2026.findings-acl.39}
}

References

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

  1. Ido Dagan, Oren Glickman, and Bernardo Magnini. 2006. The PASCAL recognising textual entailment challenge. In Machine learning challenges. evaluating predictive uncertainty, visual object classification, and recognising tectual entailment, pages 177–190. Springer. find
  2. Tri Dao and Albert Gu. 2024. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In International Conference on Machine Learning (ICML). find
  3. Ona de Gibert, Naiara Perez, Aitor García-Pablos, and Montse Cuadros. 2018. Hate Speech Dataset from a White Supremacy Forum. In Proceedings of the 2nd Workshop on Abusive Language Online (ALW2), pages 11–20, Brussels, Belgium. Association for Computational Linguistics.
  4. Roy Bar Haim, Ido Dagan, Bill Dolan, Lisa Ferro, Danilo Giampiccolo, Bernardo Magnini, and Idan Szpektor. 2006. The second PASCAL recognising textual entailment challenge. find
  5. Valerio Basile, Cristina Bosco, Elisabetta Fersini, Debora Nozza, Viviana Patti, Francisco Manuel Rangel Pardo, Paolo Rosso, and Manuela Sanguinetti. 2019. SemEval-2019 task 5: Multilingual detection of hate speech against immigrants and women in Twitter. In Proceedings of the 13th International Wor…
  6. Marco Bondaschi, Nived Rajaraman, Xiuying Wei, Kannan Ramchandran, Razvan Pascanu, Caglar Gulcehre, Michael Gastpar, and Ashok Vardhan Makkuva. 2025. From markov to laplace: How mamba in-context learns markov chains. arXiv preprint arXiv:2502.10178. — [arXiv:2502.10178](https://arxiv.org/abs/2502.10178) find
  7. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan… — Language Models are Few-Shot Learners
  8. Marie-Catherine De Marneffe, Mandy Simons, and Judith Tonhauser. 2019. The Commitment-Bank: Investigating projection in naturally occurring discourse. To appear in proceedings of Sinn und Bedeutung 23. Data can be found at https://github.com/mcdm/CommitmentBank/. ↗
  9. Thomas Diggelmann, Jordan Boyd-Graber, Jannis Bulian, Massimiliano Ciaramita, and Markus Leippold. 2020. Climate-fever: A dataset for verification of real-world climate claims. Preprint, arXiv:2012.00614.
  10. William B Dolan and Chris Brockett. 2005. Automatically constructing a corpus of sentential paraphrases. In Proceedings of the International Workshop on Paraphrasing.
  11. Xin Dong, Yonggan Fu, Shizhe Diao, Wonmin Byeon, Zijia Chen, Ameya Sunil Mahabaleshwarkar, Shih-Yang Liu, Matthijs Van Keirsbilck, Min-Hung Chen, Yoshi Suhara, Yingyan Lin, Jan Kautz, and Pavlo Molchanov. 2024. Hymba: A hybrid-head architecture for small language models.
  12. Ezra Edelman, Nikolaos Tsilivis, Benjamin Edelman, Eran Malach, and Surbhi Goel. 2024. The evolution of statistical induction heads: In-context learning markov chains. Advances in neural information processing systems, 37:64273–64311. find
  13. Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. 2022. What can transformers learn in-context? a case study of simple function classes. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS ’22, Red Hook, NY, USA. Curran Associates Inc. find
  14. Danilo Giampiccolo, Bernardo Magnini, Ido Dagan, and Bill Dolan. 2007. The third PASCAL recognizing textual entailment challenge. In Proceedings of the ACL-PASCAL workshop on textual entailment and paraphrasing, pages 1–9. Association for Computational Linguistics.
  15. Paolo Glorioso, Quentin Anthony, Yury Tokpanov, Anna Golubeva, Vasudev Shyam, James Whittington, Jonathan Pilault, and Beren Millidge. 2024a. The zamba2 suite: Technical report. Preprint, arXiv:2411.15242. — [arXiv:2411.15242](https://arxiv.org/abs/2411.15242) find
  16. Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, and Beren Millidge. 2024b. Zamba: A compact 7b ssm hybrid model. Preprint, arXiv:2405.16712. find
  17. Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, and et al. 2024. The llama 3 herd of models. Preprint, arXiv:2407.21783.
  18. Riccardo Grazzi, Julien Siems, Simon Schrodi, Thomas Brox, and Frank Hutter. 2024. Is mamba capable of in-context learning? arXiv preprint arXiv:2402.03170. — [arXiv:2402.03170](https://arxiv.org/abs/2402.03170) find
  19. Albert Gu and Tri Dao. 2023. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752. find
  20. Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. 2023. Pre-training to learn in context. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 4849–4870, Toronto, Canada. Association for Computational Linguistics. find
  21. Peter Hase, Harry Xie, and Mohit Bansal. 2021. The out-of-distribution problem in explainability and search methods for feature importance explanations. In Advances in Neural Information Processing Systems, volume 34, pages 3650–3666. Curran Associates, Inc. find
  22. Roee Hendel, Mor Geva, and Amir Globerson. 2023. In-context learning creates task vectors. In Findings of the Association for Computational Linguistics: EMNLP 2023, pages 9318–9333, Singapore. Association for Computational Linguistics.
  23. Evan Hernandez, Arnab Sen Sharma, Tal Haklay, Kevin Meng, Martin Wattenberg, Jacob Andreas, Yonatan Belinkov, and David Bau. 2023. Linearity of relation decoding in transformer language models. find
  24. Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. 2022. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations.
  25. Mingyu Jin, Kai Mei, Wujiang Xu, Mingjie Sun, Ruixiang Tang, Mengnan Du, Zirui Liu, and Yongfeng Zhang. 2025. Massive values in self-attention modules are the key to contextual knowledge understanding. In Forty-second International Conference on Machine Learning. find
  26. Guillaume Lample, Alexis Conneau, Marc’Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. 2018. Word translation without parallel data. In International Conference on Learning Representations. find
  27. Barak Lenz, Opher Lieber, Alan Arazi, Amir Bergman, Avshalom Manevich, Barak Peleg, Ben Aviram, Chen Almagor, Clara Fridman, Dan Padnos, Daniel Gissin, Daniel Jannai, Dor Muhlgay, Dor Zimberg, Edden M. Gerber, Elad Dolev, Eran Krakovsky, Erez Safahi, Erez Schwartz, Gal Cohen, Gal Shachaf, Haim Rozen… find
  28. Hector J Levesque, Ernest Davis, and Leora Morgenstern. 2011. The Winograd schema challenge. In AAAI Spring Symposium: Logical Formalizations of Commonsense Reasoning, volume 46, page 47. find
  29. Hongkang Li, Songtao Lu, Xiaodong Cui, Pin-Yu Chen, and Meng Wang. Understanding mamba in in-context learning with outliers: A theoretical generalization analysis. In High-dimensional Learning Dynamics 2025. find
  30. Yingcong Li, Xupeng Wei, Haonan Zhao, and Taigao Ma. 2024. Can mamba in-context learn task mixtures? In ICML 2024 Workshop on In-Context Learning. find
  31. P. Malo, A. Sinha, P. Korhonen, J. Wallenius, and P. Takala. 2014. Good debt or bad debt: Detecting semantic orientations in economic texts. Journal of the Association for Information Science and Technology, 65. find
  32. Marco Marelli, Stefano Menini, Marco Baroni, Luisa Bentivogli, Raffaella Bernardi, and Roberto Zamparelli. 2014. A SICK cure for the evaluation of compositional distributional semantic models. In Proceedings of the Ninth International Conference on Language Resources and Evaluation (LREC’14), pages…
  33. Clara H. McCreery, Namit Katariya, Anitha Kannan, Manish Chablani, and Xavier Amatriain. 2020. Effective transfer learning for identifying similar questions: Matching user questions to covid-19 faqs. Preprint, arXiv:2008.13546. — [arXiv:2008.13546](https://arxiv.org/abs/2008.13546) find
  34. Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. 2022a. 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, pages 2791–2809, Seattle, United States.… find
  35. Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. 2022b. Rethinking the role of demonstrations: What makes in-context learning work? In EMNLP.
  36. Saif Mohammad, Svetlana Kiritchenko, Parinaz Sobhani, Xiaodan Zhu, and Colin Cherry. 2016. Semeval-2016 task 6: Detecting stance in tweets. In Proceedings of the 10th International Workshop on Semantic Evaluation (SemEval-2016), pages 31–41.
  37. Ioannis Mollas, Zoe Chrysopoulou, Stamatis Karlos, and Grigorios Tsoumakas. 2020. Ethos: an online hate speech detection dataset. Preprint, arXiv:2006.08328. find
  38. Kim Anh Nguyen, Sabine Schulte im Walde, and Ngoc Thang Vu. 2017. Distinguishing antonyms and synonyms in a pattern-based neural network. In Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics: Volume 1, Long Papers, pages 76–85, Valencia, Spai… find
  39. Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. 2024a. Can mamba learn how to learn? a comparative study on in-context learning tasks. In Proceedings of the 41st International Conference on Machine Learning, pages 39793–39… find
  40. Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. 2024b. Can mamba learn how to learn? A comparative study on in-context learning tasks. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of… find
  41. Qwen, An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxi Yang, Jingren Zhou, Junyang Lin, Kai Dang, Keming Lu, Keqin Bao, Kexin Yang, Le Yu, Mei Li, Mingfeng Xu… find
  42. Liliang Ren, Yang Liu, Yadong Lu, yelong shen, Chen Liang, and Weizhu Chen. 2025. Samba: Simple hybrid state space models for efficient unlimited context language modeling. In The Thirteenth International Conference on Learning Representations.
  43. Laria Reynolds and Kyle McDonell. 2021. Prompt programming for large language models: Beyond the few-shot paradigm. In Extended Abstracts of the 2021 CHI Conference on Human Factors in Computing Systems, CHI EA ’21, New York, NY, USA. Association for Computing Machinery.
  44. Emily Sheng and David Uthus. 2020. Investigating societal biases in a poetry composition system. Preprint, arXiv:2011.02686. find
  45. Gemma Team, Aishwarya Kamath, Johan Ferret, Shreya Pathak, Nino Vieillard, Ramona Merhej, and et al. 2025. Gemma 3 technical report. Preprint, arXiv:2503.19786. find
  46. NVIDIA, :, Aaron Blakeman, Aarti Basant, Abhinav Khattar, Adithya Renduchintala, Akhiad Bercovich, Aleksander Ficek, Alexis Bjorlin, Ali Taghibakhshi, Amala Sanjay Deshmukh, and et al. 2025. Nemotron-h: A family of accurate and efficient hybrid mamba-transformer models. Preprint, arXiv:2504.03624.
  47. Eric Todd, Millicent L. Li, Arnab Sen Sharma, Aaron Mueller, Byron C. Wallace, and David Bau. 2024. Function vectors in large language models. In Proceedings of the 2024 International Conference on Learning Representations. ArXiv:2310.15213. find
  48. Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. 2022. In-context learning and induction heads. arXiv preprint arXiv:2209.11895.
  49. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances in neural information processing systems, 30.
  50. Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. 2023a. Interpretability in the wild: a circuit for indirect object identification in GPT-2 small. In The Eleventh International Conference on Learning Representations. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/wang-2026-understanding-context-learning-beyond-transformers/