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

TL;DR

This paper investigates in-context learning across different large language model architectures, revealing internal mechanisms and the importance of combining behavioral and mechanistic analyses for understanding LLM capabilities.

Contribution

It provides a detailed comparison of transformer, state-space, and hybrid models in ICL, uncovering different internal mechanisms and the role of function vectors.

Findings

FVs are mainly in self-attention and Mamba layers.

FVs are crucial for parametric knowledge retrieval.

Different architectures may perform similarly externally but differ internally.

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.