🤖 machine learning

Size Doesn't Matter: Cosine-Scored Sparse Autoencoders

This paper argues that replacing the standard inner-product scoring in sparse autoencoders with a learned blend of cosine similarity and magnitude prevents high-norm tokens from dominating feature selection, thereby enabling the discovery of more human-interpretable concepts on normalized representations.

Original authors: Silen Naihin, Lev Stambler

Published 2026-06-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Silen Naihin, Lev Stambler

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Loud Voice" Mistake

Imagine you are trying to understand a conversation in a crowded room. You have a team of 65,000 tiny detectives (called features) listening to the speakers. Each detective is trained to spot a specific idea, like "politics," "coding," or "sadness."

In the standard way these detectives work (using Inner Product), they get excited based on two things:

What is being said: Does the topic match their specialty?
How loud the speaker is: Is the speaker shouting or whispering?

The problem is that in modern AI models, the "volume" of a word (its mathematical size, or norm) varies wildly. Some words are naturally huge, and some are tiny.

Because the standard detectives care about volume, a loud word (like a rare, complex token) makes every detective scream "I found it!" at the same time, even if the topic doesn't match. It's like a fire alarm going off because someone dropped a heavy box, not because there is a fire.

The AI model itself doesn't actually care about the volume; it only cares about the direction of the conversation (the meaning). But the standard detectives are wasting their time shouting about volume, leaving them no room to actually listen to the meaning.

The Solution: The "Cosine" Detective

The authors propose a new type of detective that ignores the volume and only listens to the direction (the angle) of the words. They call this Cosine Scoring.

Think of it like this:

Standard Detective: "I'm excited because the speaker is shouting AND they are talking about cats!"
Cosine Detective: "I'm excited only because they are talking about cats. I don't care if they are whispering or screaming."

To make this work, the authors created a flexible system where the detective can learn exactly how much to care about volume. They gave the detectives a "volume knob" (a parameter called $a$ ).

If $a=1$ , they care about volume (the old way).
If $a=0$ , they ignore volume completely (pure cosine).

What Happened When They Tried It?

The authors trained these new detectives on a massive language model (Qwen3-8B) and compared them to the old ones. Here is what they found:

1. The "Volume Knob" Turned Down
Even though the detectives were allowed to keep caring about volume, the training process automatically turned the volume knob down to almost zero. The system learned that direction is what matters, and volume is just noise. No single detective ever decided to care about volume more than 50%.

2. Better at Finding Real Ideas
When tested on whether the detectives could identify human concepts (like "programming" or "geography"), the new Cosine Detectives were 14.9% better than the old ones.

The Old Way: Many of the detectives were "volume detectors." They only fired when a word was loud, regardless of what the word meant. They wasted their slots.
The New Way: Because they stopped shouting about volume, they had room to find actual concepts. They filled the "dictionary" with useful ideas instead of just loud noises.

3. The "High-Norm" Trap
The paper found that the old system was terrible at handling "loud" words (high-norm tokens). When a loud word came in, the old system would reconstruct the sentence with a volume 9.5 times larger than reality, effectively breaking the model's understanding. The new system handled these loud words perfectly, keeping the volume realistic.

4. It's Not Just About "More Data"
The authors proved that this wasn't just because the new system was bigger or trained longer. Even when they gave the old system three times as many detectives, it still failed to find good concepts because its "scoring rule" was broken. The problem was the geometry of how they listened, not their capacity.

The Bottom Line

The paper argues that for AI models that use a specific type of normalization (called RMSNorm), the standard way of detecting features is flawed because it confuses loudness with meaning.

By switching to a Cosine Score (which measures alignment rather than magnitude), we get a dictionary of features that:

Matches the model's reconstruction quality (it doesn't break the model).
Finds human-recognizable concepts much more often.
Stops wasting space on "volume detectors."

The authors conclude that for any AI model using these normalized representations, Cosine Scoring should be the default setting, as it aligns the dictionary learning with what the model actually "reads."

Technical Summary: Size Doesn't Matter: Cosine-Scored Sparse Autoencoders

1. Problem Statement

Sparse Autoencoders (SAEs) are standard tools for mechanistic interpretability, designed to decompose model activations into a dictionary of human-interpretable features. Standard SAEs detect features via an inner product score, $\langle w_i, x \rangle$ , which scales with both the directional alignment of the input $x$ and the input's magnitude $\|x\|$ .

This scoring geometry is fundamentally mismatched with modern transformer architectures that employ pre-sublayer normalization (specifically RMSNorm). In these models, the sublayer reads the normalized activation $x/\|x\|$ , effectively discarding the magnitude information before computation. However, the SAE encoder, hooked on the residual stream, still observes the full magnitude. Consequently, the inner-product score conflates content alignment with token norm.

Under BatchTopK selection (the standard sparsity mechanism), this mismatch creates a pathological selection pressure: high-norm tokens inflate pre-activations for all features simultaneously, regardless of their directional alignment. These tokens disproportionately claim dictionary slots, causing the SAE to learn "norm detectors" (features that fire based on token magnitude) rather than "content detectors." This results in dictionaries dominated by features that encode a quantity the model itself ignores, undermining the reliability of sparse probing, feature steering, and circuit analysis.

2. Methodology

The authors propose replacing the standard inner-product encoder score with a learned-exponent cosine score. The new pre-activation $s_i(x)$ is defined as:

$s_i(x) = e^b \|x_c\|^a \cdot \cos(x_c, w_i) + b_{enc,i}$

Where:

$x_c = x - b_{dec}$ is the input centered on the decoder bias.
$w_i$ are unit-normalized encoder rows.
$a$ is a learnable exponent controlling the dependence on input magnitude.
$b$ is a learnable scale parameter.
The term $e^b \|x_c\|^a$ interpolates between pure cosine similarity ( $a=0$ ) and inner product ( $a=1$ ).

The paper investigates two parameter regimes:

Global $a$ : A single scalar exponent shared across all features.
Per-feature $a_i$ : An extension where each feature learns its own exponent ( $a_i = a_{base} + \delta_i$ ), allowing heterogeneity in norm dependence.

Key Design Choices:

Unit Normalization: Encoder rows are unit-normalized on every forward pass to ensure the score is a true cosine similarity when $a=0$ .
Initialization: $a$ is initialized to 0 (pure cosine) and $b$ to $\log \sqrt{d_{model}}$ .
Stability: To prevent "winner-take-all" cascades where features with slightly higher initial gradients grow their exponents to dominate high-norm tokens, the per-feature variant uses a base+delta parameterization ( $a_i = a_{base} + \delta_i$ ). This anchors the global scale while allowing individual feature offsets.

3. Key Contributions

Architecture

The paper introduces cosine-scored SAEs, a drop-in replacement for the standard encoder that removes the default assumption that magnitude is a useful signal for feature detection. By allowing the optimizer to learn the norm-dependence exponent $a$ , the architecture can recover the inner product if magnitude is genuinely useful, but defaults to direction-only scoring.

Empirical Results

Experiments on Qwen3-8B (Layer 18, 500M tokens) and Gemma-2-2B demonstrate that cosine-scored SAEs achieve:

Matched Reconstruction: Fraction of Variance Explained (FVE) and model-behavior preservation (KL/CE scores) are statistically indistinguishable from standard SAEs.
Significant Interpretability Gains: At matched reconstruction, the per-feature cosine variant improves sparse-probing top-1 accuracy by +14.9% across eight SAEBench datasets.
Feature Discovery: The improvement is not due to better encoding of shared features but rather the discovery of new, interpretable features. Standard SAEs fail to learn these features because their dictionary slots are occupied by norm detectors.
Learned Geometry: Training consistently drives the learned exponent $a$ toward zero. In the global regime, $a \approx 0.26$ ; in the per-feature regime, the mean $a_i \approx 0.08$ . No feature ever approaches the inner-product limit ( $a=1$ ), confirming that direction, not magnitude, is the useful signal.

Mechanism Identification

The paper isolates the cause of the standard SAE's failure to forward-pass TopK selection under norm inflation.

Gradient Reweighting Fails: Attempts to equalize gradient magnitudes across norm quartiles in the standard SAE (via loss reweighting) close only 1.9–6.8% of the performance gap. This confirms that the issue is the selection geometry (which features win the TopK competition), not the backward-pass gradient flow.
Norm-Conditioned Features: Standard SAEs allocate the majority of their capacity to features that fire overwhelmingly on high-norm tokens (86% of activations in the top quartile). These features encode magnitude, not content.
Reconstruction Pathology: On high-norm tokens, standard SAEs reconstruct with an output norm 9.5 $\times$ larger than the input, whereas cosine SAEs maintain a 1 $\times$ ratio, aligning with the model's actual readout.

4. Significance and Claims

The paper argues that cosine scoring should be the default for dictionary learning on normalized representations. The failure mode identified—where inner-product SAEs systematically undercount content features in favor of norm detectors—likely affects any SAE trained on post-RMSNorm sites.

Practical Impact: The proposed change is a minimal architectural modification (adding one or two scalar parameters) that yields more interpretable dictionaries at no reconstruction cost.
Scope and Limitations:
- The advantage is most pronounced under RMSNorm. On LayerNorm models (e.g., Pythia, Falcon), the benefit diminishes or disappears at deep layers, suggesting that LayerNorm's mean subtraction preserves magnitude-correlated structure that inner-product scoring can exploit.
- Sentiment is the only task where the standard SAE outperforms the cosine variant, consistent with the hypothesis that magnitude carries task-relevant signal in sentiment analysis.
- The results are robust across layers and models but are not universal; the magnitude of the gain scales with model dimension and depth.

The authors conclude that the "size doesn't matter" principle holds: for normalized representations, the magnitude of the activation is noise for feature detection, and the encoder score geometry must reflect this to recover meaningful, human-recognizable concepts.