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

TL;DR

The paper investigates how in-context learning (ICL) operates across transformers, state-space models, and hybrid LLMs, focusing on parametric knowledge retrieval and contextual knowledge understanding. It combines behavioral probing and mechanistic interpretability to show that function vectors (FVs) and FV heads predominantly reside in self-attention and Mamba components, with Mamba2 likely employing a different mechanism. FVs strongly support ICL in parametric retrieval but are less critical for contextual understanding, and top FV heads vary across task types and architectures. The study broadens ICL understanding beyond transformers and demonstrates the value of integrating behavioral and mechanistic analyses to reveal architecture-specific ICL mechanisms, with practical implications for designing and evaluating future LLMs.

Abstract

We perform in-depth evaluations of in-context learning (ICL) on state-of-the-art transformer, state-space, and hybrid large language models over two categories of knowledge-based ICL tasks. Using a combination of behavioral probing and intervention-based methods, we have discovered that, while LLMs of different architectures can behave similarly in task performance, their internals could remain different. We discover that function vectors (FVs) responsible for ICL are primarily located in the self-attention and Mamba layers, and speculate that Mamba2 uses a different mechanism from FVs to perform ICL. FVs are more important for ICL involving parametric knowledge retrieval, but not for contextual knowledge understanding. Our work contributes to a more nuanced understanding across architectures and task types. Methodologically, our approach also highlights the importance of combining both behavioural and mechanistic analyses to investigate LLM capabilities.

Figures (37)

• Figure 1: Hymba-1.5B-Base's AIE heatmap on a subset of parametric knowledge retrieval ICL. Top FV heads identified are much more concentrated in self-attention layers than in SSM layers. [X-axis: head number; Y-axis: layer number.]
• Figure 2: Hymba-1.5B-Base's AIE heatmap on a subset of contextual knowledge understanding ICL. The top FV heads are far less concentrated than in parametric knowledge retrieval ICL. [X-axis: head number; Y-axis: layer number.]
• Figure 3: Left: Performance of different models on parametric knowledge retrieval datasets, for $k$ = 16, on initial setting. Right: Performance of different models on contextual knowledge understanding datasets, for $k$ = 16, on initial setting.
• Figure 4: Performance gain of label flipping and gold conditions on contextual knowledge understanding datasets. Performance gain is the difference between the performance of gold or label flipping and the corresponding no demo conditions. Performance gain for other conditions are plotted in Figure \ref{['fig:performance_gain_all']} in Appendix \ref{['appendix:additional-behavioral-results']}.
• Figure 5: AIE heatmap for model heads on parametric knowledge retrieval datasets. [X-axis: head number; Y-axis: layer number.]