Info

Last Execution: 2026-08-19

Package	Version
nnsight	0.8
torch	2.13
transformers	5.15
datasets	5.0
plotly	6.9

Evaluating a Sparse Autoencoder¶

📗 You can find an interactive Colab version of this tutorial here.

A sparse autoencoder (SAE) is a dictionary of directions, trained to reconstruct a model's activations as a sparse sum of them. Dictionary Learning shows how to read features out of one. This tutorial asks the question you should ask before you read anything out of one: is this dictionary good enough to interpret?

The reason to ask is that an SAE cannot tell you when it is wrong. It is a pair of matrices with the right shapes. Attach it at the wrong layer, feed it activations it was never trained on, splice its reconstruction at the wrong positions — every one of those runs without error and produces feature activations that look entirely plausible. The only defence is numbers, and the numbers only mean something if you check them against something published.

We take google/gemma-scope-2b-pt-res — one of the Gemma Scope residual-stream SAEs, 16,384 features on google/gemma-2-2b — and put it through:

1. Hook-point identification. Six candidate sites, scored against the average L0 the SAE's own filename claims. Only one matches.
2. Attachment, as a real submodule inside the forward pass, checked against post-hoc application on the saved tensor.
3. The evaluation suite: L0, fraction of variance unexplained, dead features, and downstream loss recovered.
4. A diagnosis. Our first loss number comes out about seven times too high. Finding out why is the most useful thing in this notebook, and it takes about ninety seconds with the right control.
5. Controls that change how you read the metrics: an identity splice, matched-norm random noise, and a second SAE from a different family.

Setup¶

In [1]:

try:
    import google.colab
    is_colab = True
except ImportError:
    is_colab = False

if is_colab:
    !pip install -U nnsight datasets

try: import google.colab is_colab = True except ImportError: is_colab = False if is_colab: !pip install -U nnsight datasets

In [2]:

from IPython.display import clear_output
import numpy as np
import torch
import torch.nn.functional as F
import nnsight
from nnsight import TransformersModel
from nnsight.modeling.transformers import WrapperModule

import plotly.express as px
import plotly.io as pio
pio.renderers.default = "colab" if is_colab else "plotly_mimetype+notebook_connected+colab+notebook"

torch.manual_seed(0);

from IPython.display import clear_output import numpy as np import torch import torch.nn.functional as F import nnsight from nnsight import TransformersModel from nnsight.modeling.transformers import WrapperModule import plotly.express as px import plotly.io as pio pio.renderers.default = "colab" if is_colab else "plotly_mimetype+notebook_connected+colab+notebook" torch.manual_seed(0);

/home/localjadenfk/miniconda3/envs/ndif2/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

Loading the dictionary¶

Pretrained SAEs are not nn.Module trees. Gemma Scope ships each dictionary as a .npz of five arrays — W_enc, b_enc, W_dec, b_dec, and a JumpReLU threshold — so the loader is ours to write. That is twenty lines and it removes a dependency, which matters here: we want the SAE to be a module we control, for a reason that shows up two cells below.

The forward pass of a JumpReLU SAE is pre = x @ W_enc + b_enc, then acts = relu(pre) * (pre > threshold), then recon = acts @ W_dec + b_dec. The threshold is what makes it sparse: a feature must clear its own learned bar or it contributes nothing.

The two lines to look at are self.encode and self.decode. They are WrapperModules — identity modules that return their input unchanged. Routing the feature activations and the reconstruction through them costs nothing numerically and makes both of them hook points, so that once the SAE is attached we can read sae.encode.output from any trace. Bare parameters and plain Python methods are not hook points and never will be; if you want to read a value, route it through a module. We show what happens without the wrappers in Step 2.

In [3]:

class JumpReLUSAE(torch.nn.Module):
    """Gemma Scope's JumpReLU SAE, loaded straight from its `.npz` of parameters."""

    def __init__(self, params):
        super().__init__()
        # Identity modules, present only so their outputs are hookable.
        self.encode = WrapperModule()
        self.decode = WrapperModule()
        for name in ["W_enc", "b_enc", "W_dec", "b_dec", "threshold"]:
            setattr(self, name, torch.nn.Parameter(torch.from_numpy(params[name])))

    def features(self, x):
        """Feature activations. A plain method, so it also works outside a trace."""
        pre = x.to(self.W_enc.dtype) @ self.W_enc + self.b_enc
        return self.encode(torch.relu(pre) * (pre > self.threshold))

    def forward(self, x):
        acts = self.features(x)
        return self.decode(acts @ self.W_dec + self.b_dec).to(x.dtype)

class JumpReLUSAE(torch.nn.Module): """Gemma Scope's JumpReLU SAE, loaded straight from its `.npz` of parameters.""" def __init__(self, params): super().__init__() # Identity modules, present only so their outputs are hookable. self.encode = WrapperModule() self.decode = WrapperModule() for name in ["W_enc", "b_enc", "W_dec", "b_dec", "threshold"]: setattr(self, name, torch.nn.Parameter(torch.from_numpy(params[name]))) def features(self, x): """Feature activations. A plain method, so it also works outside a trace.""" pre = x.to(self.W_enc.dtype) @ self.W_enc + self.b_enc return self.encode(torch.relu(pre) * (pre > self.threshold)) def forward(self, x): acts = self.features(x) return self.decode(acts @ self.W_dec + self.b_dec).to(x.dtype)

The path we download names the layer, the dictionary width, and — crucially — the average L0 the authors measured: layer_12/width_16k/average_l0_82. That 82 is our anchor. It is the one number about this SAE that was published by somebody else, and reproducing it is what tells us we have wired the dictionary to the activations it was actually trained on.

In [4]:

from huggingface_hub import hf_hub_download

path = hf_hub_download(
    "google/gemma-scope-2b-pt-res",
    "layer_12/width_16k/average_l0_82/params.npz",
)
sae = JumpReLUSAE(np.load(path)).cuda()

N_FEATURES, D_MODEL = sae.W_dec.shape
print(f"{N_FEATURES} features x {D_MODEL} dims  ({N_FEATURES / D_MODEL:.1f}x expansion)")
print(sae)

from huggingface_hub import hf_hub_download path = hf_hub_download( "google/gemma-scope-2b-pt-res", "layer_12/width_16k/average_l0_82/params.npz", ) sae = JumpReLUSAE(np.load(path)).cuda() N_FEATURES, D_MODEL = sae.W_dec.shape print(f"{N_FEATURES} features x {D_MODEL} dims ({N_FEATURES / D_MODEL:.1f}x expansion)") print(sae)

16384 features x 2304 dims  (7.1x expansion)
JumpReLUSAE(
  (encode): WrapperModule()
  (decode): WrapperModule()
)

Now the model. We load gemma-2-2b in bfloat16, which fits the whole experiment in about 8 GB of GPU memory. The SAE stays in float32 — it is 150 MB and its arithmetic is where our metrics come from, so there is no reason to degrade it. (We check at the end that this mixed-precision choice is not doing anything to the numbers.)

In [5]:

model = TransformersModel(
    "google/gemma-2-2b",
    task="text-generation",
    device_map="auto",
    dtype=torch.bfloat16,
    dispatch=True,
)
clear_output()
print(model.model.layers[12])

model = TransformersModel( "google/gemma-2-2b", task="text-generation", device_map="auto", dtype=torch.bfloat16, dispatch=True, ) clear_output() print(model.model.layers[12])

Gemma2DecoderLayer(
  (self_attn): Gemma2Attention(
    (q_proj): Linear(in_features=2304, out_features=2048, bias=False)
    (k_proj): Linear(in_features=2304, out_features=1024, bias=False)
    (v_proj): Linear(in_features=2304, out_features=1024, bias=False)
    (o_proj): Linear(in_features=2048, out_features=2304, bias=False)
  )
  (mlp): Gemma2MLP(
    (gate_proj): Linear(in_features=2304, out_features=9216, bias=False)
    (up_proj): Linear(in_features=2304, out_features=9216, bias=False)
    (down_proj): Linear(in_features=9216, out_features=2304, bias=False)
    (act_fn): GELUTanh()
  )
  (input_layernorm): Gemma2RMSNorm((2304,), eps=1e-06)
  (post_attention_layernorm): Gemma2RMSNorm((2304,), eps=1e-06)
  (pre_feedforward_layernorm): Gemma2RMSNorm((2304,), eps=1e-06)
  (post_feedforward_layernorm): Gemma2RMSNorm((2304,), eps=1e-06)
)

Two things about this printout matter later.

First, the module path to a decoder block is model.model.layers[i] — the outer model is the nnsight wrapper, the inner one is the Gemma2Model. Second, on transformers 5 a decoder block's .output is a plain (batch, seq, hidden) tensor, not a tuple. Older SAE code writes layer.output[0] to unpack a tuple; here that silently selects batch row 0 and every number you compute afterwards is wrong. attn.output still is a tuple. When in doubt, print the shape.

A corpus¶

SAE metrics are averages over text, so we need text. We take wikitext-103, concatenate it, and cut it into 4,096 blocks of 128 tokens, each starting with <bos>. Half a million tokens is enough for stable L0 and variance numbers and — as we will see in Step 3 — is about the smallest corpus that gives an honest dead-feature count.

In [6]:

import datasets

def build_corpus(tokenizer, n_seq, seq_len=128):
    """`n_seq` blocks of `seq_len` tokens from wikitext-103, each starting with BOS."""
    ds = datasets.load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split="train")
    seqs, buffer = [], []
    for row in ds:
        if len(row["text"]) < 200:          # skip headings and blank lines
            continue
        buffer.extend(tokenizer(row["text"], add_special_tokens=False)["input_ids"])
        while len(buffer) >= seq_len - 1:
            seqs.append([tokenizer.bos_token_id] + buffer[:seq_len - 1])
            buffer = buffer[seq_len - 1:]
            if len(seqs) >= n_seq:
                return torch.tensor(seqs)
    return torch.tensor(seqs)

tokens = build_corpus(model.tokenizer, 4096)
clear_output()
print(tokens.shape)
print(repr(model.tokenizer.decode(tokens[0, :24])))

import datasets def build_corpus(tokenizer, n_seq, seq_len=128): """`n_seq` blocks of `seq_len` tokens from wikitext-103, each starting with BOS.""" ds = datasets.load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split="train") seqs, buffer = [], [] for row in ds: if len(row["text"]) < 200: # skip headings and blank lines continue buffer.extend(tokenizer(row["text"], add_special_tokens=False)["input_ids"]) while len(buffer) >= seq_len - 1: seqs.append([tokenizer.bos_token_id] + buffer[:seq_len - 1]) buffer = buffer[seq_len - 1:] if len(seqs) >= n_seq: return torch.tensor(seqs) return torch.tensor(seqs) tokens = build_corpus(model.tokenizer, 4096) clear_output() print(tokens.shape) print(repr(model.tokenizer.decode(tokens[0, :24])))

torch.Size([4096, 128])
'<bos> Senjō no Valkyria 3 : Unrecorded Chronicles ( Japanese : 戦場のヴァルキュリア3 ,'

Step 1: Establish the hook point, numerically¶

The repository name says layer_12. It does not say where in layer 12 — the residual stream entering the block, the residual stream leaving it, the MLP output, or the output of one of Gemma-2's four layernorms. All of those are (batch, seq, 2304) tensors, so all of them will multiply cleanly against a [2304, 16384] encoder matrix, and all of them will give you sparse-looking feature activations. Five of the six are wrong.

So we test all six against the published anchor. Two metrics decide it:

• L0 — the mean number of features active per token. This is the number the SAE's filename claims: 82.
• FVU, the fraction of variance unexplained — sum((x - x̂)²) / sum((x - mean(x))²). Zero is a perfect reconstruction; 1.0 means the SAE does no better than predicting the mean.

⚠️ Mask the special tokens before you reduce. Every metric below is computed over real tokens only — we drop position 0, which is <bos>. Gemma Scope's training protocol excludes BOS, EOS and padding, so those positions are off-distribution for the dictionary, and including them does not error, it just quietly changes your answer. We print both versions in the table below so you can see the size of the effect; the reason behind it is Step 5.

In [7]:

@torch.no_grad()
def sae_metrics(sae, x, chunk=2048):
    """L0, FVU and mean cosine similarity for a (..., d_model) activation tensor."""
    x = x.reshape(-1, x.shape[-1])
    mu = x.float().mean(0)
    n, l0, sse, sst, cos = 0, 0.0, 0.0, 0.0, 0.0
    for i in range(0, x.shape[0], chunk):          # chunked: features are 16k-wide
        xb = x[i:i + chunk].float()
        acts, recon = sae.features(xb), sae(xb)
        l0 += (acts > 0).float().sum().item()
        sse += ((xb - recon) ** 2).sum().item()
        sst += ((xb - mu) ** 2).sum().item()
        cos += F.cosine_similarity(xb, recon, dim=-1).sum().item()
        n += xb.shape[0]
        del acts, recon, xb
    return dict(L0=l0 / n, FVU=sse / sst, cos=cos / n)

@torch.no_grad() def sae_metrics(sae, x, chunk=2048): """L0, FVU and mean cosine similarity for a (..., d_model) activation tensor.""" x = x.reshape(-1, x.shape[-1]) mu = x.float().mean(0) n, l0, sse, sst, cos = 0, 0.0, 0.0, 0.0, 0.0 for i in range(0, x.shape[0], chunk): # chunked: features are 16k-wide xb = x[i:i + chunk].float() acts, recon = sae.features(xb), sae(xb) l0 += (acts > 0).float().sum().item() sse += ((xb - recon) ** 2).sum().item() sst += ((xb - mu) ** 2).sum().item() cos += F.cosine_similarity(xb, recon, dim=-1).sum().item() n += xb.shape[0] del acts, recon, xb return dict(L0=l0 / n, FVU=sse / sst, cos=cos / n)

We collect all six candidates in one forward pass. Two nnsight details are doing work here.

tracer.stop() aborts the forward pass once we have what we want. Layer 13 is the last site we read, so the remaining 12 blocks and the 256k-wide unembedding never run — that alone is most of the time and most of the memory of this cell.

And the sites must be requested in execution order. nnsight's trace body runs alongside the model, so asking for layers[11].output after you have already read layers[13].output raises OutOfOrderError — the model has gone past it. The error names the module and says exactly that, which makes it a five-second fix, but it is worth knowing that the order of your source lines is a claim about the order of the forward pass.

In [8]:

sites = [
    ("layers[11].output   (resid post 11)", lambda m: m.model.layers[11].output),
    ("layers[12].input    (resid pre 12) ", lambda m: m.model.layers[12].input),
    ("layers[12].mlp.output              ", lambda m: m.model.layers[12].mlp.output),
    ("layers[12].post_feedforward_ln     ", lambda m: m.model.layers[12].post_feedforward_layernorm.output),
    ("layers[12].output   (resid post 12)", lambda m: m.model.layers[12].output),
    ("layers[13].output   (resid post 13)", lambda m: m.model.layers[13].output),
]

saved = {}                                  # defined outside: trace-body locals do not survive
with torch.no_grad(), model.trace(tokens[:32].cuda()) as tracer:
    for name, get in sites:                 # execution order, not convenience order
        saved[name] = get(model).save()
    tracer.stop()                           # skip layers 14-25 and the unembedding

print(f"{'site':38s}  {'L0 (all)':>9s}  {'L0 (no BOS)':>11s}  {'FVU (all)':>10s}  {'FVU (no BOS)':>12s}")
for name, _ in sites:
    x = saved[name]
    a, b = sae_metrics(sae, x), sae_metrics(sae, x[:, 1:])
    print(f"{name:38s}  {a['L0']:9.1f}  {b['L0']:11.1f}  {a['FVU']:10.2f}  {b['FVU']:12.4f}")
    saved[name] = None

sites = [ ("layers[11].output (resid post 11)", lambda m: m.model.layers[11].output), ("layers[12].input (resid pre 12) ", lambda m: m.model.layers[12].input), ("layers[12].mlp.output ", lambda m: m.model.layers[12].mlp.output), ("layers[12].post_feedforward_ln ", lambda m: m.model.layers[12].post_feedforward_layernorm.output), ("layers[12].output (resid post 12)", lambda m: m.model.layers[12].output), ("layers[13].output (resid post 13)", lambda m: m.model.layers[13].output), ] saved = {} # defined outside: trace-body locals do not survive with torch.no_grad(), model.trace(tokens[:32].cuda()) as tracer: for name, get in sites: # execution order, not convenience order saved[name] = get(model).save() tracer.stop() # skip layers 14-25 and the unembedding print(f"{'site':38s} {'L0 (all)':>9s} {'L0 (no BOS)':>11s} {'FVU (all)':>10s} {'FVU (no BOS)':>12s}") for name, _ in sites: x = saved[name] a, b = sae_metrics(sae, x), sae_metrics(sae, x[:, 1:]) print(f"{name:38s} {a['L0']:9.1f} {b['L0']:11.1f} {a['FVU']:10.2f} {b['FVU']:12.4f}") saved[name] = None

site                                     L0 (all)  L0 (no BOS)   FVU (all)  FVU (no BOS)
layers[11].output   (resid post 11)         112.8         72.7        6.24        0.3146

layers[12].input    (resid pre 12)          112.8         72.7        6.24        0.3146
layers[12].mlp.output                        23.1         23.1     1110.63     1119.1135
layers[12].post_feedforward_ln              111.7        112.2        2.01        2.2255

layers[12].output   (resid post 12)         120.5         82.0        5.58        0.2119
layers[13].output   (resid post 13)         193.6        154.8        5.33        0.3750

layers[12].output gives L0 = 82.0 against a published average_l0_82. Three significant figures, on a site we did not tune. And it is unique: the immediate neighbours, one block earlier and one block later, give 72.7 and 154.8 — 11% and 89% away. The MLP output and the post-feedforward layernorm are not even close, with FVU above 1, meaning the SAE reconstructs them worse than a constant would.

Two things fall out of this table for free.

layers[12].input and layers[11].output are numerically identical, which is what you would hope: nnsight's .input on a decoder block is the pre-block residual stream. That is a free check on the API.

And look at the two FVU columns. Including <bos>, the correct site scores 5.58 — worse than predicting the mean. Excluding it, 0.212. One token in 128 moves the number by a factor of 26. Any FVU quoted for a Gemma Scope SAE without saying whether BOS is in or out is meaningless. Hold that thought.

In [9]:

LAYER = 12

LAYER = 12

Step 2: Attach the SAE into the forward pass¶

Everything so far applied the SAE post hoc, to a tensor we had already saved. That is fine for reading, but it cannot do the thing we actually need in Step 4 — put the reconstruction back into the model and see what the model does with it. For that the SAE has to live inside the forward pass.

In nnsight you assign it into the envoy tree. model.model.layers[12].sae = sae registers the SAE on the underlying module and mirrors it as a child envoy, so it now has a permanent path and its submodules have hookable locations. Then an edit routes the layer's output through it, on every trace, until we clear it.

hook=True is what makes the SAE's internals observable: it runs the attachment's full __call__ so its submodules' hooks fire, rather than just its forward.

In [10]:

PROMPT = "The Eiffel Tower is in the city of"

# Post-hoc: save the residual, apply the SAE outside the trace.
with torch.no_grad(), model.trace(PROMPT) as tracer:
    resid = model.model.layers[LAYER].output.save()
    tracer.stop()

acts_posthoc, recon_posthoc = sae.features(resid.float()), sae(resid.float())
print(tuple(acts_posthoc.shape), tuple(recon_posthoc.shape))

PROMPT = "The Eiffel Tower is in the city of" # Post-hoc: save the residual, apply the SAE outside the trace. with torch.no_grad(), model.trace(PROMPT) as tracer: resid = model.model.layers[LAYER].output.save() tracer.stop() acts_posthoc, recon_posthoc = sae.features(resid.float()), sae(resid.float()) print(tuple(acts_posthoc.shape), tuple(recon_posthoc.shape))

(1, 9, 16384) (1, 9, 2304)

Now attach the SAE and route the layer through it.

In [11]:

model.model.layers[LAYER].sae = sae
print(model.model.layers[LAYER].sae)

with model.edit(inplace=True):
    acts = model.model.layers[LAYER].output
    model.model.layers[LAYER].output[:] = model.model.layers[LAYER].sae(acts, hook=True)

model.model.layers[LAYER].sae = sae print(model.model.layers[LAYER].sae) with model.edit(inplace=True): acts = model.model.layers[LAYER].output model.model.layers[LAYER].output[:] = model.model.layers[LAYER].sae(acts, hook=True)

JumpReLUSAE(
  (encode): WrapperModule()
  (decode): WrapperModule()
)

The SAE now runs as part of the model, and sae.encode.output and sae.decode.output are readable from any later trace. Before trusting them we check them against the post-hoc route on the same prompt: if the two disagree at all, something about the attachment is wrong.

⚠️ Note the read order in the trace body: encode, then decode, then the model's logits. The block's own .output is served after the edit has finished running the SAE inside it, so asking for the layer output first would move the worker past the SAE and raise OutOfOrderError.

In [12]:

with torch.no_grad(), model.trace(PROMPT):
    acts_forward = model.model.layers[LAYER].sae.encode.output.save()
    recon_forward = model.model.layers[LAYER].sae.decode.output.save()
    logits_sae = model.output.logits.save()

print("feature activations, max |in-forward - post-hoc| :", (acts_forward - acts_posthoc).abs().max().item())
print("reconstruction,      max |in-forward - post-hoc| :", (recon_forward - recon_posthoc).abs().max().item())

with torch.no_grad(), model.trace(PROMPT): acts_forward = model.model.layers[LAYER].sae.encode.output.save() recon_forward = model.model.layers[LAYER].sae.decode.output.save() logits_sae = model.output.logits.save() print("feature activations, max |in-forward - post-hoc| :", (acts_forward - acts_posthoc).abs().max().item()) print("reconstruction, max |in-forward - post-hoc| :", (recon_forward - recon_posthoc).abs().max().item())

feature activations, max |in-forward - post-hoc| : 0.0
reconstruction,      max |in-forward - post-hoc| : 0.0

Bit-identical. Which route you take is a question of convenience, not of numbers — so from here on we use whichever is cheaper for the job.

The second check is behavioural. Gemma-2 completes this prompt with " Paris". If the SAE's reconstruction carries the information the model needs, the answer should survive being routed through a 16,384-way sparse bottleneck. If we instead zero the residual at layer 12, it should not.

In [13]:

model.clear_edits()                       # back to the unmodified model

with torch.no_grad(), model.trace(PROMPT):
    logits_clean = model.output.logits.save()

with torch.no_grad(), model.trace(PROMPT):
    model.model.layers[LAYER].output[:] = 0
    logits_zero = model.output.logits.save()

for name, lg in [("clean", logits_clean), ("through the SAE", logits_sae), ("zero-ablated", logits_zero)]:
    print(f"{name:16s} -> {model.tokenizer.decode(lg[0, -1].argmax())!r}")

model.clear_edits() # back to the unmodified model with torch.no_grad(), model.trace(PROMPT): logits_clean = model.output.logits.save() with torch.no_grad(), model.trace(PROMPT): model.model.layers[LAYER].output[:] = 0 logits_zero = model.output.logits.save() for name, lg in [("clean", logits_clean), ("through the SAE", logits_sae), ("zero-ablated", logits_zero)]: print(f"{name:16s} -> {model.tokenizer.decode(lg[0, -1].argmax())!r}")

clean            -> ' Paris'
through the SAE  -> ' Paris'
zero-ablated     -> '<pad>'

The reconstruction preserves the answer; destroying the same site does not. That is the weakest useful form of the drop-in test — a single prompt — and Step 4 replaces it with a real measurement.

Why the SAE needed those two WrapperModules

A faithful loader for a pretrained dictionary holds bare nn.Parameters and does the arithmetic in plain Python methods. Nothing in it is a hook point, and hook=True cannot help: it runs the attachment's __call__ so that its submodules fire, and a raw loader has no submodules. The cell below writes the same SAE without the wrappers — sharing the very same parameter tensors — and asks for its encoder output.

In [14]:

class RawJumpReLUSAE(torch.nn.Module):
    """The same SAE written the obvious way: raw parameters, plain methods, no submodules."""

    def __init__(self, source):
        super().__init__()
        for name in ["W_enc", "b_enc", "W_dec", "b_dec", "threshold"]:
            setattr(self, name, torch.nn.Parameter(getattr(source, name).data))   # shares storage

    def encode(self, x):                                  # a method, not a module
        pre = x.to(self.W_enc.dtype) @ self.W_enc + self.b_enc
        return torch.relu(pre) * (pre > self.threshold)

    def forward(self, x):
        return (self.encode(x) @ self.W_dec + self.b_dec).to(x.dtype)

model.model.layers[LAYER].raw_sae = RawJumpReLUSAE(sae)

with model.edit(inplace=True):
    acts = model.model.layers[LAYER].output
    model.model.layers[LAYER].output[:] = model.model.layers[LAYER].raw_sae(acts, hook=True)

try:
    with torch.no_grad(), model.trace(PROMPT):
        model.model.layers[LAYER].raw_sae.encode.output.save()
except AttributeError as e:
    print(f"AttributeError: {e}")

model.clear_edits()

class RawJumpReLUSAE(torch.nn.Module): """The same SAE written the obvious way: raw parameters, plain methods, no submodules.""" def __init__(self, source): super().__init__() for name in ["W_enc", "b_enc", "W_dec", "b_dec", "threshold"]: setattr(self, name, torch.nn.Parameter(getattr(source, name).data)) # shares storage def encode(self, x): # a method, not a module pre = x.to(self.W_enc.dtype) @ self.W_enc + self.b_enc return torch.relu(pre) * (pre > self.threshold) def forward(self, x): return (self.encode(x) @ self.W_dec + self.b_dec).to(x.dtype) model.model.layers[LAYER].raw_sae = RawJumpReLUSAE(sae) with model.edit(inplace=True): acts = model.model.layers[LAYER].output model.model.layers[LAYER].output[:] = model.model.layers[LAYER].raw_sae(acts, hook=True) try: with torch.no_grad(), model.trace(PROMPT): model.model.layers[LAYER].raw_sae.encode.output.save() except AttributeError as e: print(f"AttributeError: {e}") model.clear_edits()

AttributeError: 'function' object has no attribute 'output'

'function' object has no attribute 'output' — the value is computed, there is simply nowhere to hang a hook. This is the single reason our loader has those two identity modules in it. The same trick works for anything you want to observe that is not already a module: an intermediate in a custom loss, a probe's pre-activation, a step inside an adapter.

Step 3: The evaluation suite¶

Four numbers, and each one is there because it catches something the others miss.

• L0 — mean active features per token. Sparsity is the whole premise: a dictionary that fires 3,000 features per token explains the activation but has not decomposed it.
• FVU — fraction of variance unexplained. How much of the activation the SAE throws away.
• Dead-feature fraction — features that never fire on any token of the corpus. Dead features cost you dictionary capacity without showing up in L0 or FVU at all. An SAE can post excellent numbers on both while using a quarter of its dictionary.
• Loss recovered — what the model loses when it has to run on the reconstruction. This is the only metric that asks about the thing we actually care about, and it is the subject of Step 4.

We run the scan with the SAE attached but not splicing. The edit below calls the SAE on the layer's output and discards the result: the model is unchanged, the SAE runs alongside it, and sae.encode.output is readable. That separation matters — we want to measure the SAE's behaviour on the model's real activations, not on activations that have already been through the SAE at layer 12 and are therefore no longer the model's.

In [15]:

model.clear_edits()

with model.edit(inplace=True):
    # Call the SAE, keep the result. The model's own output is untouched.
    model.model.layers[LAYER].sae(model.model.layers[LAYER].output, hook=True)

model.clear_edits() with model.edit(inplace=True): # Call the SAE, keep the result. The model's own output is untouched. model.model.layers[LAYER].sae(model.model.layers[LAYER].output, hook=True)

The loop below is 4,096 sequences in batches of 16, reducing as it goes so that no 16,384-wide feature tensor is ever kept longer than one iteration. tracer.stop() again pays for itself: layer 12 is halfway through the model, so stopping there roughly halves both the time and the peak memory. On one A6000 this is about 30 seconds for 520,192 tokens.

In [16]:

import time

BATCH = 16
fired = torch.zeros(N_FEATURES, device="cuda")        # firings per feature
fired_at_pos1 = torch.zeros(N_FEATURES, device="cuda")
max_act = torch.zeros(N_FEATURES, device="cuda")      # per-feature max, positions >= 2
max_act_all = torch.zeros(N_FEATURES, device="cuda")  # per-feature max, positions >= 1
l0_by_position = torch.zeros(128, device="cuda")
n_tokens, sse, sum_x, sum_xx, cos_total = 0, 0.0, torch.zeros(D_MODEL, device="cuda"), 0.0, 0.0

torch.cuda.reset_peak_memory_stats()
start = time.time()
for i in range(0, tokens.shape[0], BATCH):
    with torch.no_grad(), model.trace(tokens[i:i + BATCH].cuda()) as tracer:
        acts = model.model.layers[LAYER].sae.encode.output.save()
        recon = model.model.layers[LAYER].sae.decode.output.save()
        resid = model.model.layers[LAYER].output.save()
        tracer.stop()

    on = acts > 0
    l0_by_position += on.sum(-1).float().sum(0)
    fired += on[:, 1:].sum((0, 1)).float()            # BOS excluded from here on
    fired_at_pos1 += on[:, 1].sum(0).float()
    max_act = torch.maximum(max_act, acts[:, 2:].amax((0, 1)))
    max_act_all = torch.maximum(max_act_all, acts[:, 1:].amax((0, 1)))

    x, xhat = resid[:, 1:].float(), recon[:, 1:]
    sse += ((x - xhat) ** 2).sum().item()
    sum_x += x.sum((0, 1)); sum_xx += (x ** 2).sum().item()
    cos_total += F.cosine_similarity(x, xhat, dim=-1).sum().item()
    n_tokens += x.shape[0] * x.shape[1]
    del acts, recon, resid, on, x, xhat

mean_resid = sum_x / n_tokens                          # kept for the mean-ablation baseline
sst = sum_xx - n_tokens * (mean_resid ** 2).sum().item()
l0_by_position /= tokens.shape[0]

print(f"{n_tokens:,} non-BOS tokens in {time.time() - start:.0f}s, "
      f"peak {torch.cuda.max_memory_allocated() / 2**30:.2f} GiB")

import time BATCH = 16 fired = torch.zeros(N_FEATURES, device="cuda") # firings per feature fired_at_pos1 = torch.zeros(N_FEATURES, device="cuda") max_act = torch.zeros(N_FEATURES, device="cuda") # per-feature max, positions >= 2 max_act_all = torch.zeros(N_FEATURES, device="cuda") # per-feature max, positions >= 1 l0_by_position = torch.zeros(128, device="cuda") n_tokens, sse, sum_x, sum_xx, cos_total = 0, 0.0, torch.zeros(D_MODEL, device="cuda"), 0.0, 0.0 torch.cuda.reset_peak_memory_stats() start = time.time() for i in range(0, tokens.shape[0], BATCH): with torch.no_grad(), model.trace(tokens[i:i + BATCH].cuda()) as tracer: acts = model.model.layers[LAYER].sae.encode.output.save() recon = model.model.layers[LAYER].sae.decode.output.save() resid = model.model.layers[LAYER].output.save() tracer.stop() on = acts > 0 l0_by_position += on.sum(-1).float().sum(0) fired += on[:, 1:].sum((0, 1)).float() # BOS excluded from here on fired_at_pos1 += on[:, 1].sum(0).float() max_act = torch.maximum(max_act, acts[:, 2:].amax((0, 1))) max_act_all = torch.maximum(max_act_all, acts[:, 1:].amax((0, 1))) x, xhat = resid[:, 1:].float(), recon[:, 1:] sse += ((x - xhat) ** 2).sum().item() sum_x += x.sum((0, 1)); sum_xx += (x ** 2).sum().item() cos_total += F.cosine_similarity(x, xhat, dim=-1).sum().item() n_tokens += x.shape[0] * x.shape[1] del acts, recon, resid, on, x, xhat mean_resid = sum_x / n_tokens # kept for the mean-ablation baseline sst = sum_xx - n_tokens * (mean_resid ** 2).sum().item() l0_by_position /= tokens.shape[0] print(f"{n_tokens:,} non-BOS tokens in {time.time() - start:.0f}s, " f"peak {torch.cuda.max_memory_allocated() / 2**30:.2f} GiB")

520,192 non-BOS tokens in 20s, peak 5.75 GiB

In [17]:

L0 = fired.sum().item() / n_tokens
FVU = sse / sst
dead = (fired == 0).sum().item()
rare = (fired / n_tokens < 1e-5).sum().item()

print(f"L0                      {L0:8.2f}      (published: 82)")
print(f"FVU                     {FVU:8.4f}")
print(f"variance explained      {1 - FVU:8.4f}")
print(f"mean cos(x, x_hat)      {cos_total / n_tokens:8.4f}")
print(f"dead features           {dead:8d} / {N_FEATURES}  = {100 * dead / N_FEATURES:.2f}%")
print(f"firing on < 1e-5 tokens {rare:8d} / {N_FEATURES}  = {100 * rare / N_FEATURES:.2f}%")

L0 = fired.sum().item() / n_tokens FVU = sse / sst dead = (fired == 0).sum().item() rare = (fired / n_tokens < 1e-5).sum().item() print(f"L0 {L0:8.2f} (published: 82)") print(f"FVU {FVU:8.4f}") print(f"variance explained {1 - FVU:8.4f}") print(f"mean cos(x, x_hat) {cos_total / n_tokens:8.4f}") print(f"dead features {dead:8d} / {N_FEATURES} = {100 * dead / N_FEATURES:.2f}%") print(f"firing on < 1e-5 tokens {rare:8d} / {N_FEATURES} = {100 * rare / N_FEATURES:.2f}%")

L0                         81.10      (published: 82)
FVU                       0.2049
variance explained        0.7951
mean cos(x, x_hat)        0.9211
dead features                400 / 16384  = 2.44%
firing on < 1e-5 tokens     1328 / 16384  = 8.11%

L0 = 81.10 against a published 82 — a gap of 1.1%. The residual gap is explained by context length: Gemma Scope measured over 1,024-token sequences and we used 128, and L0 rises with context (the next plot shows it). The dictionary explains 79.5% of the variance and 2.4% of its features are dead, which is healthy — hold on to that figure, because Step 9 shows an SAE where it is 73%.

In [18]:

fig = px.line(
    x=list(range(128)), y=l0_by_position.cpu(), log_y=True, markers=False,
    labels={"x": "token position in the sequence", "y": "mean L0 (log scale)"},
    title="Active features per token, by position",
)
fig.add_hline(y=82, line_dash="dash", annotation_text="published average L0 = 82")
fig.show()

fig = px.line( x=list(range(128)), y=l0_by_position.cpu(), log_y=True, markers=False, labels={"x": "token position in the sequence", "y": "mean L0 (log scale)"}, title="Active features per token, by position", ) fig.add_hline(y=82, line_dash="dash", annotation_text="published average L0 = 82") fig.show()