SAE features: reading a model's concepts

Note

Generated from notebooks/applications/sae_features.ipynb, which you can run yourself. The outputs below are the ones it produced.

A sparse autoencoder (SAE) re-expresses one layer’s residual stream as a long, mostly-zero vector of features. The hope is that each feature carries one concept, where a raw residual dimension carries a superposition of many.

Murano loads a pre-trained SAE from HuggingFace and reads the layer it was trained on straight off its config, so you never match hook sites by hand.

This notebook encodes five sentences through an SAE and asks what fired, and what the thing that fired means.

Key questions

• How sparse is a sparse autoencoder, in practice?
• Which feature best characterizes a whole sentence?
• A feature id is a number: how do I find out what concept it carries?

Structure

1. Set up
2. Encode prompts through the SAE
3. Confirm the code really is sparse
4. Find each sentence’s defining feature
5. Read what a feature means
6. Save

Model and data. Gemma 2 2B Instruct with the 16k gemma-scope SAE at layer 20. The model is gated on the Hub: accept its licence once, then hf auth login. Unlike the GPT-2 notebooks this one loads a 2B model, so expect a few GB of GPU memory. Five sentences, each naming a landmark in a different country.

Requirements. pip install murano-interp[sae]

1. Set up

SAEEncode reads its target layer from the SAE’s own config, so the layer is never passed twice.

from murano import MuranoModel, Pipeline, keys
from murano.steps import SAEEncode, SAEFeatureLabel, Save, top_sae_features_per_prompt
from murano.steps.prompts import LoadPrompts

OUTPUT_DIR = "murano_outputs/sae_features"

MODEL_ID = "google/gemma-2-2b-it"
SAE_RELEASE = "gemma-scope-2b-pt-res-canonical"
SAE_ID = "layer_20/width_16k/canonical"

model = MuranoModel(MODEL_ID)
print(f"{model.model_id}: {model.n_layers} layers")

google/gemma-2-2b-it: 26 layers

2. Encode prompts through the SAE

PROMPTS = [
    "The Eiffel Tower is located in the city of Paris",
    "The Colosseum is located in the city of Rome",
    "The Brandenburg Gate is located in the city of Berlin",
    "Senso-ji temple is located in the city of Tokyo",
    "Red Square is located in the city of Moscow",
]

results = Pipeline(
    [
        LoadPrompts(PROMPTS),
        SAEEncode(model, release=SAE_RELEASE, sae_id=SAE_ID),
    ]
).run()

record = results[keys.SAE_RECORD]
n_prompts, n_tokens, n_features = record.activations.shape
print(f"encoded {n_prompts} prompts at {record.hook}")
print(f"every token is now a {n_features}-feature vector")
print(f"activations {tuple(record.activations.shape)}  [prompts, tokens, features]")

encoded 5 prompts at L20.resid_post
every token is now a 16384-feature vector
activations (5, 13, 16384)  [prompts, tokens, features]

3. Confirm the code really is sparse

The claim in the name is checkable: count how many features fire on each real token.

active_per_token = (record.activations > 0).sum(dim=-1)
mean_active = active_per_token[record.attention_mask.bool()].float().mean().item()

print(f"{mean_active:.0f} of {n_features} features fire per token", end=" ")
print(f"({100 * mean_active / n_features:.1f}% active)")

751 of 16384 features fire per token (4.6% active)

4. Find each sentence’s defining feature

To read what a sentence is about, average each feature over the sentence’s tokens and take the largest. top_sae_features_per_prompt does that, and drops the few high-frequency features that fire on nearly every token and would otherwise win every average by default.

top_feats = top_sae_features_per_prompt(record, n=1, reduce="mean")
for prompt, feats in zip(record.texts, top_feats):
    print(f"{prompt}\n    -> top feature #{feats[0]}\n")

The Eiffel Tower is located in the city of Paris
    -> top feature #5516

The Colosseum is located in the city of Rome
    -> top feature #1780

The Brandenburg Gate is located in the city of Berlin
    -> top feature #16173

Senso-ji temple is located in the city of Tokyo
    -> top feature #1315

Red Square is located in the city of Moscow
    -> top feature #1116

5. Read what a feature means

A feature id is a number. To learn what it carries, SAEFeatureLabel takes the feature’s decoder direction and projects it through the unembedding: the tokens it pushes hardest are the tokens it promotes. This is the logit lens, applied to a feature rather than to a residual stream.

feat_ids = sorted({fid for feats in top_feats for fid in feats})
results = Pipeline([SAEFeatureLabel(model, feat_ids=feat_ids, k_tokens=1)]).run(results)
labels = results["feature_labels"]

for prompt, feats in zip(record.texts, top_feats):
    fid = feats[0]
    print(f"{prompt}\n    -> feature #{fid} means {labels.tokens[fid][0].strip()!r}\n")

The Eiffel Tower is located in the city of Paris
    -> feature #5516 means 'French'

The Colosseum is located in the city of Rome
    -> feature #1780 means 'Italian'

The Brandenburg Gate is located in the city of Berlin
    -> feature #16173 means 'German'

Senso-ji temple is located in the city of Tokyo
    -> feature #1315 means 'Japan'

Red Square is located in the city of Moscow
    -> feature #1116 means 'Russian'

Each sentence names a landmark in a different country, and each defining feature promotes that country’s name or nationality. The SAE has pulled one clean concept out of a residual stream that mixed dozens.

Note what this does not establish. A feature that promotes " France" at the unembedding is a feature the model uses to say “France”. Whether it is the feature the model uses to know France is a further, causal claim, and reading it off the decoder cannot settle it. sae_steering.ipynb makes that claim testable by intervening.

6. Save

run_dir = Save(
    output_dir="murano_outputs", run_name="sae_features", model_id=model.model_id
)(results)[keys.OUTPUT_DIR]

print("saved to:", run_dir)

saved to: murano_outputs/sae_features

What next

• sae_steering.ipynb turns a feature from something you read into something you write with.
• sae_enrichment.ipynb ranks thousands of features by how well they separate two classes.