Reward-Informed Sparse Autoencoders and the Solution-Completeness Confound

Tanvi Nagilla; Alexander Jameson; Daniel Manta; Shayaan Uddin

Source-linked AI summary

Reward-Informed Sparse Autoencoders and the Solution-Completeness Confound

Tanvi Nagilla, Alexander Jameson, Daniel Manta, Shayaan Uddin

arXiv:2608.26136v1Jun 25, 2026cs.CL

TL;DR

The paper asks whether reward-curated SAE features reveal reasoning quality or merely completion form. It builds an RI-SAE from high- and low-reward GRPO continuations and tests the separation with structural, text-classifier, and reward-blind SAE controls. The main result is that the separation is largely solution completeness, so reward filtering is useful only when interpreted against those controls.

Problem
Standard SAEs learn from undifferentiated text, leaving limited evidence about whether reward-guided features separate competent from incompetent reasoning.
Method
The paper filters GRPO continuations by reward and trains a standard JumpReLU SAE on the resulting good- and bad-reasoning activations.
Results
The control battery shows that good/bad SAE separation is largely solution completeness rather than reasoning quality, while a generic reward-blind SAE does not recover it.
Takeaways & Limitations
Reward filtering is a cheap, reusable way to direct SAEs toward reasoning-adjacent data, but candidate reasoning features require a completeness baseline.
Takeaways & Limitations
The evaluation uses small sets and single runs without seeds or error bars, and the generic-SAE control does not isolate reward filtering from in-domain training.

Abstract

from arXiv · show

Sparse autoencoders (SAEs) decompose language-model activations into sparse, interpretable features, and an appealing way to aim them at reasoning is to curate their data with a signal reinforcement learning already produces: the reward. We build such a reward-informed SAE (RI-SAE): we split GRPO trajectories into high-reward ("good") and low-reward ("bad") reasoning continuations, train a standard JumpReLU SAE on their activations, and then ask what the resulting good/bad separation actually measures. On Llama-3.1-8B a sparse subset of the 16,384 features does separate the classes (silhouette 0.79 on the selected features versus 0.005 for the full code), but a control battery shows the separation is largely solution completeness rather than reasoning quality: a TF-IDF text classifier already splits the classes (AUC 0.75--0.83), and three structural cues alone (length, a closed reasoning block, and a boxed answer) reach AUC 0.70 (99% of good versus 69% of bad completions are boxed). A generic SAE that never saw the reward does not separate the classes at all (silhouette 0.01, no discriminative features), so the 0.79 is in-sample fitting of this curated signal rather than structure that a reward-blind dictionary recovers. We therefore present the recipe and its control battery together: reward filtering is a cheap, label-free way to reuse RL signals for interpretability, but most of what it surfaces is completion form. Two discriminative features are still readable (symbolic mathematics; procedural and evaluative language), which we take as illustrative rather than as isolated reasoning.

1 Introduction

The paper proposes using reward-labeled GRPO continuations to focus SAEs on reasoning-related activations. Its central question is whether the resulting separation reflects reasoning quality or solution completeness.

Motivation: SAEs recover sparse, often monosemantic activation features, but standard training does not target competent versus incompetent reasoning.They reconstruct activations under a sparsity penalty while learning from undifferentiated text.
Motivation: RL rewards provide unused supervision for selecting activations associated with high- and low-quality reasoning continuations.GRPO rollouts carry scalar rewards that can guide SAE data curation without adding a new training objective.
Method: RI-SAE filters GRPO trajectories into good and bad continuations, then trains an otherwise standard SAE on their activations.The reward enters through data selection rather than the SAE objective or base-model retraining.
Main question: The control battery finds that good/bad separation largely measures solution completeness rather than reasoning quality.The paper reports TF-IDF AUC 0.75–0.83, structure-only AUC 0.70, and no separation from a generic reward-blind SAE.
Contributions: The paper presents reward filtering as a cheap, label-free interpretability recipe whose outputs must be evaluated against controls.Two interpretable features and a fine-tuning study are offered as illustrative directions rather than definitive reasoning diagnostics.

2 Method

The method curates a balanced corpus of high- and low-reward GRPO continuations, applies coherence and formatting filters, and trains a standard JumpReLU SAE on residual-stream activations.

Reward-based data curation: The corpus uses GRPO continuations with rewards in [−1, 3] and a <think>...</think><answer>...</answer> format.Non-English text is removed, and examples must contain at least 30 tokens.
Reward-based data curation: Good reasoning is labeled reward ≥2.0, while bad reasoning is labeled reward ≤0.5.Both classes must pass coherence checks excluding excessive n-gram repetition and canned refusals.
Sparse autoencoder: The SAE encodes residual activation x with JumpReLU thresholds and reconstructs it from a learned sparse code.The encoder computes z = ReLU(W_enc x + b_enc − θ), followed by reconstruction x̂ = W_dec z + b_dec.
Sparse autoencoder: Training uses the standard SAE objective with Adam at learning rate 3 × 10^−4 for 500 steps.The reward is absent from the objective; analyses use a 16,384-feature dictionary on layer-22 activations.

3 Results

The reward-informed SAE separates good- and bad-reward continuations only after selecting discriminative features, while controls show that separation largely reflects solution completeness and curated in-sample fitting. Two selected features have readable token associations, but a checkpoint experiment characterizes model progress rather than feature evolution.

Good/bad separation: 0.79 silhouette versus 0.005 for the full 16,384-feature code shows clean separation only in the selected SAE subspace.The selection criterion was per-feature silhouette > 0.1, so the result demonstrates a discriminative sparse subset rather than global SAE separation.
Solution completeness: AUC 0.83 cross-validated and 0.75 held out for TF-IDF, plus AUC 0.70 from three structural cues, reveal strong lexical and structural confounds.The structural cues are length, a closed reasoning block, and a boxed answer; stripping lexical content barely changes the text-classifier separation.
Solution completeness: 99% of good versus 69% of bad completions contain a boxed answer, while 100% versus 83% close the reasoning block.These formatting differences support interpreting much of the split as solution completeness rather than reasoning quality.
Generic SAE control: A generic reward-blind SAE reaches full-code silhouette 0.012, has no feature above the selection threshold, and yields selected-feature silhouette 0.02.This control indicates that the 0.79 separation requires an SAE trained on the curated set itself.
Interpretable features: Feature 15968 is associated with symbolic mathematics, while Feature 4205 is associated with procedural and evaluative language.Feature 4205 also activates on common function words, making its interpretation suggestive rather than definitive.
Toward monitoring fine-tuning: Exact match rises from 0.5% to ∼5% and MAE falls from ∼4950 to ∼100 across five Gemma-2-2B checkpoints, but no SAE is evaluated across checkpoints.The largest gains occur between steps 1200 and 1800, characterizing model training rather than feature sharpening.

4 Limitations

The paper treats reward labels as a proxy and limits interpretation of the selected-feature silhouette because the analysis is completeness-dominated and in-sample. Small evaluation sets, single runs, differing backbones, and unimplemented follow-up tests further bound the conclusions.

Scope and interpretation: Reward labels are a proxy rather than a distinction between reasoning and its absence, and the main separation is dominated by solution completeness.The selected-feature silhouette therefore cannot by itself establish reasoning structure.
Scope and interpretation: The 0.79 silhouette is selection-based rather than global, since the full code has silhouette 0.005.The analysis selects features for discrimination before measuring separation.
Evaluation limits: Evaluation sets are small, and results come from single runs without seeds or error bars, so they are presented as descriptive.
Unresolved tests: The generic-SAE control does not isolate reward filtering from in-domain training; an SAE trained on unfiltered in-domain data remains the cleanest test.Causal interventions and a concept-alignment metric are also left for future work.
Scope and interpretation: The feature analyses use Llama-3.1-8B, whereas the fine-tuning study uses Gemma-2-2B.

5 Conclusion

Reward filtering can direct SAEs toward reasoning-adjacent data, but the resulting separation must be evaluated against solution completeness because most of the signal reflects completion form rather than reasoning quality.

5 Conclusion: AUC 0.70 is achieved by length, a closed reasoning block, and a boxed answer, showing that completion structure explains much of the good/bad separation.The authors recommend comparing reward-informed features against a completeness baseline before calling them reasoning features.
5 Conclusion: Reward filtering is a cheap, reusable way to aim standard SAEs at reasoning-adjacent data, but it is not a free reasoning signal.The control battery is presented as part of the recipe rather than as an optional diagnostic.
5 Conclusion: A generic reward-blind SAE fails to separate the classes, motivating controls that distinguish reward-filtering effects from SAE-only effects.The conclusion proposes same-recipe unfiltered training, reward-weighted objectives, and checkpoint monitoring as next tests.

A Fine-tuning details and curves

Gemma-2-2B was fine-tuned on GSM8K with LoRA across five checkpoints and evaluated using exact match, numeric-answer MAE, and training/validation loss curves.

A Fine-tuning details and curves: Five Gemma-2-2B GSM8K checkpoints were evaluated on final-answer exact match and numeric-answer mean absolute error.The model was fine-tuned with LoRA and evaluated on the GSM8K test set.
A Fine-tuning details and curves: Exact match is measured by strict last-line string matching, while MAE captures convergence of numeric answers toward correct values.The two metrics therefore characterize different aspects of answer performance.
A Fine-tuning details and curves: Figure 3 reports MAE on numeric answers, exact match, and training/validation loss as model-performance curves.SAE-based feature tracking across checkpoints is left to future work.

B SAE configuration

The reported analyses use a JumpReLU SAE with a 16,384-feature dictionary trained on Llama-3.1-8B layer-22 activations under a sparse reconstruction objective.

B SAE configuration: The JumpReLU SAE uses learned per-feature thresholds initialized at 0.001 and unit-norm decoder rows.Its objective combines MSE with a 0.1∥z∥0 sparsity term.
B SAE configuration: 16,384 features are trained with Adam at learning rate 3 × 10^-4, batch size 16, for 500 steps.These settings define the dictionary width and optimization schedule used in the reported analyses.
B SAE configuration: The reported analyses use Llama-3.1-8B residual-stream activations from layer 22.Data curation uses reward thresholds 2.0 and 0.5, a 30-token minimum, and coherence filters.