On the Versatility of Sparse Autoencoders for In-Context Learning (2025)

Abstract: Sparse autoencoders (SAEs) are emerging as a key analytical tool in the field of mechanistic interpretability for large language models (LLMs). While SAEs have primarily been used for interpretability, we shift focus and explore an understudied question: “Can SAEs be applied to practical tasks beyond interpretability?” Given that SAEs are trained on billions of tokens for sparse reconstruction, we believe they can serve as effective extractors, offering a wide range of useful knowledge that can benefit practical applications. Building on this motivation, we demonstrate that SAEs can be effectively applied to in-context learning (ICL). In particular, we highlight the utility of the SAE-reconstruction loss by showing that it provides a valuable signal in ICL—exhibiting a strong correlation with LLM performance and offering a powerful unsupervised approach for prompt selection. These findings underscore the versatility of SAEs and reveal their potential for real-world applications beyond interpretability. Our code is available at https://github.com/ihcho2/SAE-GPS.

## introduction sparse autoencoders (saes), an entirely unsupervised approach, are rapidly gaining popularity as a key tool in mechanistic interpretability (mi), which aims to explain and understand how large language models (llms) process and generate responses @xcite . trained with a sparsity loss term to reconstruct the input (i.e., llm embedding) in a fully unsupervised manner, saes have been shown to be effective, to some extent, in decomposing embeddings into sparse features, where each feature corresponds to human-understandable, mono-semantic concept @xcite . although saes have been predominantly used for mi, there is no strong reason to assume they are limited to this domain foot_0 . thus, in this work, we explore the understudied question of whether saes can be practical for real-world applications beyond mi. promisingly, our findings reveal that saes are highly informative and can be effectively applied to in-context learning (icl), a widely used approach of llms. promisingly, our findings reveal that saes provide meaningful insights and can be effectively applied to in-context learning (icl), a widely used paradigm in large language models. icl-an emergent ability of llms to perform diverse tasks by learning on-the-fly from just a few exemplars within a single prompt-has become a dominant paradigm in natural language processing @xcite ). yet, understanding the underlying mechanisms of icl remains an open challenge @xcite . in this paper, we analyze icl behaviors of llms through the lens of saes, uncovering their interesting patterns and insights. our key motivation is that saes serve as effective extractors, trained on billions of tokens, that decompose llm embeddings into sparse, explicit features. we believe these well-trained extractors provide new sources of knowledge that can directly benefit multiple aspects of icl. in particular, we focus on the sae-reconstruction loss (sae-rl)-defined as the l2-norm between the input llm embedding and the sae's reconstructed embedding-and demonstrate that it offers valuable information to enhance the icl framework. we hypothesize that "the richer the knowledge contained in an embedding, the more challenging it becomes for saes to disentangle it into distinct features-resulting in a higher sae-rl". this, in turn, allows us to effectively measure the overall quality of the embeddings produced by a given prompt, making icl significantly more predictable. these findings highlight the versatility of saes, paving the way for exciting new avenues for exploration. ## related work in-context learning in llms in-context learning (icl) is a paradigm in which large language models (llms) are provided with a few exemplars within a single prompt and generate responses conditioned on these examples without any taskspecific fine-tuning @xcite . by leveraging pre-trained knowledge, icl enables efficient generalization across various tasks, making it highly adaptable in natural language processing. however, the precise mechanism behind icl remain debated and not fully understood @xcite @xcite . additionally, icl faces various practical challenges such as exemplar selection @xcite , exemplar order issues @xcite , and prompt sensitivity @xcite @xcite . in this paper, we examine icl through the lens of saes, aiming to uncover insights that not only deepen our understanding of icl but also enhance its performance by addressing the challenges mentioned above. sparse autoencoders in llm-interpretability saes have emerged as a key approach to sparse dictionary learning, which aims to decompose llm embeddings into sparse monosemantic features @xcite . they have proven surprisingly effective at uncovering ground truth features in controlled, toy experiment settings @xcite , although many challenges remain when applying them to real llms. nonetheless, due to their simplicity and effectiveness, saes have become one of the dominant baselines in mi. saes use a squared error reconstruction loss and a sparsity penalty: for the activation function σ, various techniques are used, such as jumprelu @xcite and topk-relu @xcite . ## sae-gps: sae-guided prompt search in this paper, we show that saes can add values to icl, making it more predictable in multiple ways. specifically, we verify that the sae-reconstruction loss (sae-rl) exhibits a strong correlation with llm performance across various tasks. based on this finding, we introduce sae-gps (sae-guided prompt search), which enables selecting effective prompts (e.g., a set of icl exemplars) on-the-fly in an entirely unsupervised manner. intuition the core intuition behind sae-gps is that the more knowledge an embedding contains, the more challenging it becomes for saes to decompose it effectively, resulting in a higher sae-rl. therefore, sae-rl allows comparing different prompts directly in an unsupervised manner. we verify this core intuition in section 5.1. methodology motivated by our intuition, we present sae-gps (sae-guided prompt search), an inference-time method that relies solely on sae-rl to accurately select effective prompts on-thefly. the overall architecture of sae-gps is illustrated in figure 1 . given a test instance x and k multiple candidate prompts p 1 , p 2 , . . . , p k , the goal is to select the prompt p * that yields the best performance. we utilize our key intuition that prompts with higher sae-rl (i.e., using more knowledge) tend to predict more accurately. accordingly, we run the llm m with each prompt p i up to layer l (i.e., g @xmath0:l [p i ; x]), then measure the sae-rl of the last tokens' embeddings (i.e., ||g i,last -sae(g i,last )|| 2 2 ). the prompt p that produces the highest sae-rl is selected, and further processed through the remaining layers for the final prediction (i.e., ŷ = m 1:l [ p ; x]). we show in section 5.2 that sae-gps consistently outperforms the baselines, and can also be effectively performed at layers below the final layer. this offers a unique computational benefit by requiring processing only up to layer l for the remaining (k -1) prompts, thereby reducing the computational burden. ## experimental settings datasets we select two widely used classification datasets in icl: aspect-based sentiment classification (absc) (semeval-14-laptops and restaurants @xcite ) and one widely used multiple-choice question-answering dataset: mmlu @xcite . we use the mul-tilingual version of mmlu provided by openai (openai, 2024) for multilingualism analysis. we create the validation set by selecting 300 instances per label from the training set, and the final evaluation is conducted on the test set. detailed explanations of each task can be found in appendix a. ## conclusion this paper investigates whether sparse autoencoders (saes) can benefit in-context learning (icl) in large language models (llms). we present several novel findings and approaches. first, we show that sae-reconstruction loss (sae-rl) is highly informative and can be leveraged to predict effective prompts during inference. we also demonstrate that sae-rl can distinguish between easy and challenging examples for llms. furthermore, sae-rl appears to be a valuable knowledge source across various domains-including multilingualism-thereby opening exciting avenues for future research. to the best of our knowledge, our work is among the first to explore the versatility of saes in icl, highlighting their potential and paving the way for further advancements in the field.