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

1. Problem Statement

Large Language Models (LLMs) have achieved state-of-the-art performance in text classification, but their internal decision-making processes remain opaque ("black boxes"). While traditional interpretability methods (e.g., gradient-based attribution) exist, they often lack semantic coherence or stability.

Concept-based interpretability aims to extract high-level, human-understandable features (concepts) from the model's latent space. However, existing approaches face three main challenges in the context of text classification:

1. Completeness: Extracted concepts must be sufficient to reconstruct the model's original prediction.
2. Causality: Activating or ablating a concept must significantly alter the model's output probability.
3. Interpretability: Concepts should be "monosemantic" (representing a single clear idea) and align with human notions.
4. Efficiency: Current state-of-the-art methods (like HI-Concept) are computationally expensive and often require re-running the entire LLM during training.

Furthermore, most Sparse Autoencoder (SAE) research focuses on generative tasks or pre-trained models, leaving a gap in applying SAEs specifically to sentence-level classification where the goal is to isolate task-relevant features.

2. Methodology: ClassifSAE

The authors propose ClassifSAE, a supervised variant of the Sparse Autoencoder tailored for text classification. Unlike standard SAEs trained on massive, diverse corpora, ClassifSAE is trained directly on the specific classification dataset to uncover concepts relevant to that task.

Architecture

• Input: The hidden state ($h_\ell$) of a specific token in the LLM (the token preceding the label for decoder-only models, or the [CLS] token for encoder-only models). This token aggregates sentence-level information.
• Encoder-Decoder: The SAE learns a sparse representation $z$ of $h_\ell$.
  • $z = \sigma(W_{enc}h + b_{enc})$
  • $\hat{h} = W_{dec}z + b_{dec}$
• Joint Classifier: A key innovation is the addition of a lightweight classifier head ($g_\theta$) trained jointly with the SAE. This classifier takes a subset of the SAE features ($z_{class} \subset z$) and attempts to reproduce the LLM's predicted label.
  • This forces the SAE to concentrate task-relevant information into a specific subspace ($z_{class}$), while the remaining features handle reconstruction.

Training Loss

The total loss function combines three components:
$$L = \lambda_1 L_{TopK}^{SAE} + \lambda_2 L_{class} + \lambda_3 L_{sparse}^{feature}$$

1. Reconstruction Loss ($L_{TopK}^{SAE}$): Minimizes the error between the original hidden state and the reconstructed state, using a TopK activation function to enforce sparsity (fixing the number of active concepts per input).
2. Classification Loss ($L_{class}$): Cross-entropy loss between the joint classifier's prediction and the LLM's predicted label. This drives the selection of discriminative features.
3. Activation Rate Sparsity Loss ($L_{sparse}^{feature}$): A novel mechanism to prevent "feature collapse" (where a few features activate on almost all inputs). It penalizes features that activate on more than a threshold percentage ($\gamma$) of the batch, encouraging a balanced distribution of activation across features.

3. Key Contributions

A. Novel Metrics for Interpretability

The paper introduces two metrics to quantify interpretability without relying on human annotators or LLM-as-a-judge (which can be unstable):

1. ConceptSim: Measures the semantic coherence of a single concept. It clusters sentences activating a specific feature and calculates the average pairwise cosine similarity of their sentence embeddings (via Sentence-BERT). High scores indicate the concept is monosemantic.
2. SentenceSim: Measures the consistency of sentences sharing the same top-activated concepts. It evaluates if sentences with overlapping concept sets are semantically closer in the embedding space.

B. Evaluation Framework

The authors benchmark ClassifSAE against four baselines:

• TopK-SAE: Standard unsupervised SAE.
• ICA (Independent Component Analysis): A classic linear unmixing technique.
• ConceptSHAP: A method maximizing concept contribution to classifier completeness.
• HI-Concept: A state-of-the-art method explicitly optimizing for causal impact.

Evaluation covers 7 backbone LLMs (including BERT, DeBERTa, Pythia, GPT-J, Mistral, LLaMA) and 4 datasets (AG News, IMDB, TweetEval Offensive, TweetEval Sentiment).

4. Results

Performance Metrics

• Interpretability (ConceptSim): ClassifSAE consistently achieves the highest ConceptSim scores across all models and datasets, indicating superior monosemanticity and semantic coherence compared to all baselines.
• Causality ($\Delta f_{cond}$): ClassifSAE achieves the second-best causality scores, trailing only HI-Concept. However, HI-Concept achieves this at the cost of lower interpretability. ClassifSAE demonstrates a better trade-off, offering strong causal influence while maintaining high interpretability.
• Completeness (RAcc): ClassifSAE maintains high Recovery Accuracy (often >95%), proving that the extracted concepts retain sufficient information to reconstruct the model's decision.

Efficiency

• Training Time: ClassifSAE is significantly more efficient. On the largest model (LLaMA 3.1 8B), it requires up to 83% less training time than HI-Concept and 69% less than ConceptSHAP.
• Reason: Unlike baselines that require repeated forward passes through the entire downstream LLM layers to compute causal proxies, ClassifSAE only trains an autoencoder and a lightweight classifier on the hidden states, avoiding the need to load the full LLM during training.

Qualitative Analysis

Visualizations (PCA projections and word clouds) show that ClassifSAE discovers fine-grained concepts (e.g., distinguishing "Airlines" from "Currency" within the "Business" category) rather than collapsing all features into broad category labels.

5. Significance and Conclusion

This paper addresses a critical gap in mechanistic interpretability for classification tasks. By integrating a supervised classifier and a sparsity mechanism into the SAE framework, ClassifSAE successfully extracts features that are:

1. Task-Relevant: Specifically tuned to the classification objective.
2. Interpretable: Highly monosemantic and semantically coherent.
3. Causal: Capable of significantly influencing model predictions.
4. Efficient: Computationally feasible for large models.

The work suggests that while explicit causal optimization (as in HI-Concept) yields high causality scores, a joint training approach with sparsity constraints offers a more balanced and practical solution for understanding LLM decision-making in text classification. The proposed metrics (ConceptSim, SentenceSim) provide a reproducible, automated way to evaluate concept quality, moving the field away from reliance on qualitative human inspection.