Unveiling Decision-Making in LLMs for Text Classification : Extraction of influential and interpretable concepts with Sparse Autoencoders

TL;DR

This paper presents a novel SAE-based model ClassifSAE tailored for text classification, leveraging a specialized classifier head and incorporating an activation rate sparsity loss, and shows that ClassifSAE improves both the causality and interpretability of the extracted features.

Abstract

Sparse Autoencoders (SAEs) have been successfully used to probe Large Language Models (LLMs) and extract interpretable concepts from their internal representations. These concepts are linear combinations of neuron activations that correspond to human-interpretable features. In this paper, we investigate the effectiveness of SAE-based explainability approaches for sentence classification, a domain where such methods have not been extensively explored. We present a novel SAE-based model ClassifSAE tailored for text classification, leveraging a specialized classifier head and incorporating an activation rate sparsity loss. We benchmark this architecture against established methods such as ConceptShap, Independent Component Analysis, HI-Concept and a standard TopK-SAE baseline. Our evaluation covers several classification benchmarks and backbone LLMs. We further enrich our analysis with two novel metrics for measuring the precision of concept-based explanations, using an external sentence encoder. Our empirical results show that ClassifSAE improves both the causality and interpretability of the extracted features.

Figures (16)

• Figure 1: Examples of concepts discovered by from the internals of GPT-J fine-tuned on AG News.
• Figure 2: Architecture of our model. A classifier is trained jointly with the SAE to replicate the original LLM prediction. A low dimensionality of $\textbf{z}_{\text{class}}$ incentivizes the model to extract a small number of distinct task-relevant features.
• Figure 3: $(k)$ as a function of the number $k$ of shared top-activating concepts between sentence pairs. Concepts are learned from sentence-level hidden states in the penultimate transformer block of GPT-J fine-tuned on one of the four classification tasks. We consider $p=5$ concepts for each sentence.
• Figure 4: 2D Principal Component Analysis (PCA) fitted on sentence-level hidden-state activations extracted from the residual stream of the penultimate transformer block of LLaMA 3.1 Instruct tasked on two classification datasets: AG News (top) and IMDB (bottom). The colored circles stand for the concepts learned by . Their size is proportional to their mean activation over the dataset. The proportion of color is representative of the normalized class score for each concept. The triangle symbols depict the class prototypes activations.
• Figure 5: Averaged accuracy deterioration $\Delta \text{Acc}^{\text{global}}$ as a function of the ablation level $p$ of class-specific features segments $(\mathcal{F}_c)_{c \in C}$. Results are reported for concepts computed from hidden states extracted at the residual stream exiting the penultimate transformer block of Pythia-1B fine-tuned on AG News.