Scaling sparse feature circuit finding for in-context learning

TL;DR

This work develops a scalable mechanistic interpretability framework for in-context learning by combining Sparse Autoencoders (SAEs) with Task Vector Cleaning (TVC) and an adapted Sparse Feature Circuits (SFC) pipeline. It identifies two core feature families—task-detection features that identify the current task and task-execution features that implement the task—demonstrating their causal roles within the Gemma-1 2B model through steering experiments and circuit analysis. The authors modify SFC to handle large ICL prompts and introduce token-position categorization, a revised loss, and faithfulness metrics, uncovering a network of SAE latents that strongly influence ICL and connect via attention and transcoder components. These findings advance mechanistic interpretability by showing how abstract task representations are decomposed into executable features that interact across transformer sublayers, offering a path toward controllable and safer ICL behaviors. The work also contributes practical tools and datasets that enable broader analysis of larger models and more complex ICL tasks.

Abstract

Sparse autoencoders (SAEs) are a popular tool for interpreting large language model activations, but their utility in addressing open questions in interpretability remains unclear. In this work, we demonstrate their effectiveness by using SAEs to deepen our understanding of the mechanism behind in-context learning (ICL). We identify abstract SAE features that (i) encode the model's knowledge of which task to execute and (ii) whose latent vectors causally induce the task zero-shot. This aligns with prior work showing that ICL is mediated by task vectors. We further demonstrate that these task vectors are well approximated by a sparse sum of SAE latents, including these task-execution features. To explore the ICL mechanism, we adapt the sparse feature circuits methodology of Marks et al. (2024) to work for the much larger Gemma-1 2B model, with 30 times as many parameters, and to the more complex task of ICL. Through circuit finding, we discover task-detecting features with corresponding SAE latents that activate earlier in the prompt, that detect when tasks have been performed. They are causally linked with task-execution features through the attention and MLP sublayers.

Figures (32)

• Figure 1: A diagram of the in-context learning circuit, showing task detection features (yellow) causing task execution features (blue) which cause the model to output the antonym (left $\to$ right). A more concrete circuit, along with texts these features activate on, can be seen in \ref{['img:circuit_clean']}.
• Figure 1: Activation masses for executor features across different token types, averaged across all tasks. We can notice they activate largely on arrow tokens.
• Figure 2: The effect on the Gemma 1 2B's task losses by steering with different kinds of reconstructed task vectors, at each layer. We see that cleaning performs similarly to the original task vector until layer 14. Average relative loss change measured as a post-steering relative loss change compared to 0-shot, averaged across all tasks.
• Figure 2: Activation masses for task-detection features across different token types, averaged across all tasks. We can notice that they activate almost exclusively on output tokens.
• Figure 3: Evaluation of task vectors after applying our TVC algorithm across different $L_1$ coefficient $\lambda$ values. Top: relative decrease in loss after steering (higher $\to$ better); bottom: fraction of retained active features (lower $\to$ better). Transparent lines represent different model and SAE combinations; solid lines show means across all of them; the results are averaged across tasks; ($x$-axis: $L_1$ coefficient $\lambda$). Further details in \ref{['app:sweeps']}.