Jun 1, 2025
Judge using SAE Features
Dhruv Yadav
🏆 3rd place by peer review
The key idea of this project was to explore model judgement using
Sparse Autoencoder (SAE) features for mathematical reasoning
tasks involving addition, multiplication, and subtraction operations.
We compared this performance against basic Chain-of-Thought
(CoT) based judgement using routing algorithms deployed by the
Martian SDK.
Our work addresses the Judge Model Development track by
developing a SAE feature-based routing system that achieves
comparable performance to traditional CoT approaches while
providing enhanced interpretability. We demonstrate that SAE
feature analysis can effectively guide model selection for
mathematical reasoning tasks, with our system showing 91%
routing accuracy compared to 98% for traditional CoT methods.
Critically, our SAE-based approach identified 7% of cases where
both models were inadequate based on feature analysis - insights
invisible to traditional routing methods.
This work advances mechanistic interpretability of routing systems
by providing transparent, feature-level explanations for routing
decisions, enabling users to understand why specific models are
chosen for a sample problem of mathematical operations.
Thank you for your submission. This is an interesting paper with some good results.
As you say SAEs give some insights, but have limitations and are costly. It’s hard to know how much to rely on SAE features being accurate representations of the model’s capability.
I liked this project a lot and think it has substantial potential for further development (I realize the hackathon was short). In particular, this expands on the autointerp methods that already exist to get good judge results.
I'd be interested to see further evolutions of this that use attribution graphs (once transcoders are trained for the relevant models). However, when implementing this one must be careful to ensure that the autointerp results (for the individual SAE features) are of sufficiently high quality. Some are vague and underspecified.
Constructive critique:
Strength:
Problem is well defined: can we route based on activating features in a prompt that is an interpretable approach to routing
The experiments are properly designed and executed on!
Shows that SAE feature matching allows a more diverse routing strategy and also that there are cases where neither models are aligned on the task (Cool insight and it relates to fallback mechanisms in routing that are important to report especially in a product setting where a customer would prefer a warning rather than incorrect answers)
Weakness:
Comparing the SAE feature detection approach to simpler baselines would make this investigation complete: e.g, Train linear probes by extracting hidden activations from a chosen model layer for each math prompt, labeling them with interpretable task types (e.g., addition, subtraction), and fitting a logistic regression classifier to predict these labels for use in routing decisions
Strong assumption but fair for the hackathon: For each input query, they assume that the features it would activate in the small (distilled) model for which Neuronpedia has SAEs available for are representative of the reasoning traits that would be activated in the full model.
SAEs are costly to obtain if they are not already available so including that cost in routing makes it a bad routing decision.
Cite this work
@misc {
title={
(HckPrj) Judge using SAE Features
},
author={
Dhruv Yadav
},
date={
6/1/25
},
organization={Apart Research},
note={Research submission to the research sprint hosted by Apart.},
howpublished={https://apartresearch.com}
}
