Exploring Task Performance with Interpretable Models via Sparse Auto-Encoders

TL;DR

This paper introduces a method using sparse autoencoders to interpret large language models, extracting meaningful features, identifying misunderstandings, and improving task performance through prompt reformulation.

Contribution

It presents a novel approach combining dictionary learning and autoencoders to interpret LLMs and enhance downstream task accuracy.

Findings

Extracts monosemantic features from polysemantic neurons

Identifies internal misunderstandings in models

Improves downstream task performance with prompt reformulation

Abstract

Large Language Models (LLMs) are traditionally viewed as black-box algorithms, therefore reducing trustworthiness and obscuring potential approaches to increasing performance on downstream tasks. In this work, we apply an effective LLM decomposition method using a dictionary-learning approach with sparse autoencoders. This helps extract monosemantic features from polysemantic LLM neurons. Remarkably, our work identifies model-internal misunderstanding, allowing the automatic reformulation of the prompts with additional annotations to improve the interpretation by LLMs. Moreover, this approach demonstrates a significant performance improvement in downstream tasks, such as mathematical reasoning and metaphor detection.