Interpretable Reward Model via Sparse Autoencoder

TL;DR

SARM tackles the opacity of scalar reward models in RLHF by embedding a pretrained sparse autoencoder (SAE) into the reward model, projecting LLM hidden activations into a sparse, monosemantic feature space. A two-stage training pipeline—sequence-level SAE pretraining followed by standard reward-model training—enables direct feature-level attribution and enables precise, weight-based steering of reward preferences without requiring multidimensional labels. Empirical results show that SARM achieves superior alignment benchmarks compared to baselines while preserving interpretability, including successful targeted manipulation of safety-related features and robust performance at smaller scales. This approach offers a practical path to interpretable and controllable reward modeling with meaningful downstream impact in RLHF deployments.

Abstract

Large language models (LLMs) have been widely deployed across numerous fields. Reinforcement Learning from Human Feedback (RLHF) leverages reward models (RMs) as proxies for human preferences to align LLM behaviors with human values, making the accuracy, reliability, and interpretability of RMs critical for effective alignment. However, traditional RMs lack interpretability, offer limited insight into the reasoning behind reward assignments, and are inflexible toward user preference shifts. While recent multidimensional RMs aim for improved interpretability, they often fail to provide feature-level attribution and require costly annotations. To overcome these limitations, we introduce the Sparse Autoencoder-enhanced Reward Model (SARM), a novel architecture that integrates a pretrained Sparse Autoencoder (SAE) into a reward model. SARM maps the hidden activations of LLM-based RM into an interpretable, sparse, and monosemantic feature space, from which a scalar head aggregates feature activations to produce transparent and conceptually meaningful reward scores. Empirical evaluations demonstrate that SARM facilitates direct feature-level attribution of reward assignments, allows dynamic adjustment to preference shifts, and achieves superior alignment performance compared to conventional reward models. Our code is available at https://github.com/schrieffer-z/sarm.

Figures (3)

• Figure 1: Multidimensional reward model. It projects the final token through a regression layer, generating multidimensional reward scores. Despite their interpretability improvements, existing multidimensional methods exhibit two primary limitations: (1) limited feature-level interpretability, and (2) significantly increased annotation costs.
• Figure 2: Overview of the SARM framework. In Stage 1, we pretrain a sparse autoencoder (SAE) on sequence-level hidden states from Layer $l$ of a pretrained LLM using a general corpus. This step extracts a set of abstract, monosemantic, and interpretable features from the LLM's latent space. In Stage 2, we attach the pretrained SAE encoder back to Layer $l$, freeze its parameters, and train a learnable linear head for reward modeling on preference data. At inference time, SARM produces reward scores that are explicitly attributed to interpretable SAE features.
• Figure 3: Effect of manipulating a safety-related feature on reward distribution. (a) Manipulating a safety-related feature leads to a clear rightward shift in the reward distribution, indicating successful preference manipulation. (b) The distribution remains largely unchanged on the complement of target datasets, suggesting that the manipulating effect does not influence irrelevant attributes.