Info

Last Execution: 2026-09-03

Package	Version
nnsight	0.8.0
Python	3.12.13
torch	2.13.0+cu126
transformers	5.15.0.dev0

Concept Erasure¶

Introduction¶

A linear probe tells us that some concept — the gender of a name, the sentiment of a review, the language of a sentence — is linearly decodable from a model's residual stream. It does not tell us that the model uses it. Concept erasure is the intervention that turns the correlational claim into a causal one: we remove the concept's linear subspace from the residual stream, let the rest of the network run on the erased state, and watch what breaks.

Two methods are standard, and we run both:

• LEACE (Least-squares Concept Erasure, Belrose et al., 2023) computes the minimal affine edit under which no linear probe beats chance. For a binary concept it is rank 1, and it comes with a guarantee.
• INLP (Iterative Nullspace Projection, Ravfogel et al., 2020) repeatedly trains a probe and projects out its weight vector. It has no optimality guarantee, but it produces a nested family of projections — one fit gives you every rank — so it is the method that can answer "how many directions does this concept need?"

The part that makes this an nnsight tutorial rather than a scikit-learn one is where the projection is applied. Fitting an eraser on saved activations and multiplying the saved tensor by it proves nothing about the model: the forward pass already happened, the logits were already produced from the unerased state. To make a causal claim we have to apply the projection inside the forward pass, so that every downstream layer, and the unembedding, only ever see erased activations. In nnsight that is two lines inside a trace:

h = model.transformer.h[6].output
model.transformer.h[6].output = mu + (h - mu) @ M

Over the course of the notebook we hold ourselves to four standards, each of which is a step most write-ups skip:

1. The probe is real. A shuffled-label control sits at chance, and the concept is not leaking through sequence length or position.
2. The concept dies for a probe retrained on the erased activations, not merely for the probe we already had. Any projection defeats its own probe; that is not evidence of anything.
3. An unrelated capability survives, and the comparison is against a rank-matched random projection. Deleting directions always costs something; the question is whether a learned direction buys more erasure per unit of damage than an arbitrary one.
4. Erasing at one layer is not enough. Later layers rebuild the concept, and we show the refill.

We use gpt2 throughout, so the whole notebook runs in a few minutes on a single small GPU.

Setup¶

In [1]:

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

if is_colab:
    !pip install -U nnsight

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

In [2]:

from IPython.display import clear_output
import numpy as np
import pandas as pd
import torch
import nnsight
from nnsight import TransformersModel

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

from IPython.display import clear_output import numpy as np import pandas as pd import torch import nnsight from nnsight import TransformersModel import plotly.express as px import plotly.io as pio pio.renderers.default = "colab" if is_colab else "plotly_mimetype+notebook_connected+colab+notebook"

We load GPT-2 with an explicit task=. It is optional online, but TransformersModel cannot infer the task when the HuggingFace hub is unreachable (HF_HUB_OFFLINE=1), so passing it costs nothing and saves a confusing error later.

In [3]:

model = TransformersModel("openai-community/gpt2", task="text-generation",
                          device_map="auto", dispatch=True)
clear_output()
print(model)

model = TransformersModel("openai-community/gpt2", task="text-generation", device_map="auto", dispatch=True) clear_output() print(model)

GPT2LMHeadModel(
  (transformer): GPT2Model(
    (wte): Embedding(50257, 768)
    (wpe): Embedding(1024, 768)
    (drop): Dropout(p=0.1, inplace=False)
    (h): ModuleList(
      (0-11): 12 x GPT2Block(
        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
        (attn): GPT2Attention(
          (c_attn): Conv1D()
          (c_proj): Conv1D()
          (attn_dropout): Dropout(p=0.1, inplace=False)
          (resid_dropout): Dropout(p=0.1, inplace=False)
        )
        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
        (mlp): GPT2MLP(
          (c_fc): Conv1D()
          (c_proj): Conv1D()
          (act): NewGELUActivation()
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
    )
    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
  )
  (lm_head): Linear(in_features=768, out_features=50257, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
)

In [4]:

DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"

tokenizer = model.tokenizer
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

blocks = model.transformer.h
N_LAYERS = len(blocks)
D_MODEL = model.config.n_embd
print(f"{N_LAYERS} blocks, d_model = {D_MODEL}")

DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu" tokenizer = model.tokenizer tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "right" blocks = model.transformer.h N_LAYERS = len(blocks) D_MODEL = model.config.n_embd print(f"{N_LAYERS} blocks, d_model = {D_MODEL}")

12 blocks, d_model = 768

The concept: gender of a first name¶

We need a concept that GPT-2 genuinely represents and visibly uses, so that erasing it has something to break. Gender of an English first name is a good choice: a probe reads it easily from the residual stream, and the model's use of it is observable at the logits — after "Yesterday Sarah went to the store because" GPT-2 prefers " she" over " he", and the preference flips for a male name. That gives us a decoding measure and a behavioural measure of the same concept.

The dataset is 88 male and 88 female first names crossed with 12 sentence templates.

In [5]:

MALE = """James John Robert Michael William David Richard Joseph Thomas Charles Christopher Daniel
Matthew Anthony Donald Mark Paul Steven Andrew Kenneth George Joshua Kevin Brian Edward Ronald
Timothy Jason Jeffrey Ryan Jacob Gary Nicholas Eric Stephen Jonathan Larry Justin Scott Brandon
Frank Benjamin Gregory Samuel Raymond Patrick Alexander Jack Dennis Jerry Tyler Aaron Henry Douglas
Peter Adam Nathan Zachary Walter Kyle Harold Carl Jeremy Keith Roger Gerald Ethan Arthur Terry
Christian Sean Lawrence Austin Joe Noah Jesse Albert Bryan Billy Bruce Willie Jordan Dylan Alan
Ralph Gabriel Roy Juan Wayne Eugene Logan Randy Louis Russell Vincent Philip Bobby Johnny Bradley""".split()

FEMALE = """Mary Patricia Jennifer Linda Elizabeth Barbara Susan Jessica Sarah Karen Nancy Lisa Betty
Margaret Sandra Ashley Kimberly Emily Donna Michelle Carol Amanda Dorothy Melissa Deborah Stephanie
Rebecca Sharon Laura Cynthia Amy Kathleen Angela Shirley Anna Brenda Pamela Emma Nicole Helen Samantha
Katherine Christine Debra Rachel Carolyn Janet Catherine Maria Heather Diane Ruth Julie Olivia Joyce
Virginia Victoria Kelly Lauren Christina Joan Evelyn Judith Megan Andrea Cheryl Hannah Jacqueline
Martha Gloria Teresa Ann Sara Madison Frances Kathryn Janice Jean Abigail Alice Julia Judy Sophia
Grace Denise Amber Doris Marilyn Danielle Beverly Isabella Theresa Diana Natalie Brittany Charlotte
Marie Kayla""".split()

TEMPLATES = [
    "Yesterday {n} walked into the room and looked around before sitting down",
    "This morning {n} opened the letter carefully and read it twice before leaving",
    "Apparently {n} said that the meeting would start at nine, so everyone arrived",
    "After work {n} drove home slowly through the rain and thought about the whole day",
    "Somehow {n} had always wanted to learn the piano and finally started taking lessons",
    "Everyone agreed that {n} was the best person for the job at the company",
    "Then {n} picked up the phone, hesitated for a moment, and dialed the number",
    "When the storm hit {n} was already halfway across the field carrying a heavy box",
    "Later {n} spent the afternoon in the library reading about the history of Rome",
    "The teacher asked {n} to stay behind after class to discuss the written assignment",
    "Suddenly {n} laughed at the joke and went back to typing on the old keyboard",
    "Last summer {n} traveled to the coast and stayed in a small house by the water",
]

MALE = """James John Robert Michael William David Richard Joseph Thomas Charles Christopher Daniel Matthew Anthony Donald Mark Paul Steven Andrew Kenneth George Joshua Kevin Brian Edward Ronald Timothy Jason Jeffrey Ryan Jacob Gary Nicholas Eric Stephen Jonathan Larry Justin Scott Brandon Frank Benjamin Gregory Samuel Raymond Patrick Alexander Jack Dennis Jerry Tyler Aaron Henry Douglas Peter Adam Nathan Zachary Walter Kyle Harold Carl Jeremy Keith Roger Gerald Ethan Arthur Terry Christian Sean Lawrence Austin Joe Noah Jesse Albert Bryan Billy Bruce Willie Jordan Dylan Alan Ralph Gabriel Roy Juan Wayne Eugene Logan Randy Louis Russell Vincent Philip Bobby Johnny Bradley""".split() FEMALE = """Mary Patricia Jennifer Linda Elizabeth Barbara Susan Jessica Sarah Karen Nancy Lisa Betty Margaret Sandra Ashley Kimberly Emily Donna Michelle Carol Amanda Dorothy Melissa Deborah Stephanie Rebecca Sharon Laura Cynthia Amy Kathleen Angela Shirley Anna Brenda Pamela Emma Nicole Helen Samantha Katherine Christine Debra Rachel Carolyn Janet Catherine Maria Heather Diane Ruth Julie Olivia Joyce Virginia Victoria Kelly Lauren Christina Joan Evelyn Judith Megan Andrea Cheryl Hannah Jacqueline Martha Gloria Teresa Ann Sara Madison Frances Kathryn Janice Jean Abigail Alice Julia Judy Sophia Grace Denise Amber Doris Marilyn Danielle Beverly Isabella Theresa Diana Natalie Brittany Charlotte Marie Kayla""".split() TEMPLATES = [ "Yesterday {n} walked into the room and looked around before sitting down", "This morning {n} opened the letter carefully and read it twice before leaving", "Apparently {n} said that the meeting would start at nine, so everyone arrived", "After work {n} drove home slowly through the rain and thought about the whole day", "Somehow {n} had always wanted to learn the piano and finally started taking lessons", "Everyone agreed that {n} was the best person for the job at the company", "Then {n} picked up the phone, hesitated for a moment, and dialed the number", "When the storm hit {n} was already halfway across the field carrying a heavy box", "Later {n} spent the afternoon in the library reading about the history of Rome", "The teacher asked {n} to stay behind after class to discuss the written assignment", "Suddenly {n} laughed at the joke and went back to typing on the old keyboard", "Last summer {n} traveled to the coast and stayed in a small house by the water", ]

Two details in the templates are load-bearing, and both exist to stop the label leaking through something other than the model's representation of gender.

⚠️ Keep only names that are a single token with a leading space. GPT-2's tokenizer splits female first names into more pieces than male ones on average, so a mixed set makes sequences of different lengths for the two classes. Sequence length changes the positional embeddings and the index of the last real token, and a probe will happily read gender off that instead. It will not error; it will just give you a suspiciously high number at layer 0.

⚠️ The name is never sentence-initial. Every template starts with a filler word, so within one template every sequence has the same length and the same token at every position except the name.

We split by name, not by sentence: the 28 test names never appear in training. A split by sentence would let the probe memorise names rather than learn the concept.

In [6]:

single_token = lambda names: [n for n in names
                              if len(tokenizer.encode(" " + n, add_special_tokens=False)) == 1]

male, female = single_token(MALE), single_token(FEMALE)
n_per_class = min(len(male), len(female), 88)

rng = np.random.RandomState(0)
rng.shuffle(male); rng.shuffle(female)
male, female = male[:n_per_class], female[:n_per_class]

N_TEST = 28
test_m,  test_f  = male[:N_TEST], female[:N_TEST]
train_m, train_f = male[N_TEST:], female[N_TEST:]
print(f"{n_per_class} names per class -> {len(train_m)} train / {len(test_m)} test names per class")
print("train (male):  ", ", ".join(train_m[:8]))
print("test  (female):", ", ".join(test_f[:8]))

single_token = lambda names: [n for n in names if len(tokenizer.encode(" " + n, add_special_tokens=False)) == 1] male, female = single_token(MALE), single_token(FEMALE) n_per_class = min(len(male), len(female), 88) rng = np.random.RandomState(0) rng.shuffle(male); rng.shuffle(female) male, female = male[:n_per_class], female[:n_per_class] N_TEST = 28 test_m, test_f = male[:N_TEST], female[:N_TEST] train_m, train_f = male[N_TEST:], female[N_TEST:] print(f"{n_per_class} names per class -> {len(train_m)} train / {len(test_m)} test names per class") print("train (male): ", ", ".join(train_m[:8])) print("test (female):", ", ".join(test_f[:8]))

88 names per class -> 60 train / 28 test names per class
train (male):   Jordan, Albert, Carl, Willie, Juan, Vincent, Aaron, Jason
test  (female): Ashley, Janet, Cynthia, Cheryl, Betty, Beverly, Donna, Olivia

In [7]:

def build_set(names_m, names_f, templates=TEMPLATES):
    texts, labels = [], []
    for label, group in enumerate((names_f, names_m)):   # 0 = female, 1 = male
        for name in group:
            for t in templates:
                texts.append(t.format(n=name))
                labels.append(label)
    return texts, np.array(labels)

train_texts, train_y = build_set(train_m, train_f)
test_texts,  test_y  = build_set(test_m,  test_f)

print(len(train_texts), "train sequences /", len(test_texts), "test sequences")
print(repr(train_texts[0]))
print("token count for the 12 sequences of one name:",
      sorted({len(tokenizer.encode(t)) for t in train_texts[:12]}))

def build_set(names_m, names_f, templates=TEMPLATES): texts, labels = [], [] for label, group in enumerate((names_f, names_m)): # 0 = female, 1 = male for name in group: for t in templates: texts.append(t.format(n=name)) labels.append(label) return texts, np.array(labels) train_texts, train_y = build_set(train_m, train_f) test_texts, test_y = build_set(test_m, test_f) print(len(train_texts), "train sequences /", len(test_texts), "test sequences") print(repr(train_texts[0])) print("token count for the 12 sequences of one name:", sorted({len(tokenizer.encode(t)) for t in train_texts[:12]}))

1440 train sequences / 672 test sequences
'Yesterday Charlotte walked into the room and looked around before sitting down'
token count for the 12 sequences of one name: [12, 13, 14, 15, 16, 17]

Step 1: A probe for the concept¶

Erasure needs a direction, and a direction comes from a probe. We train an L2-regularised logistic regression on the residual stream — the output of a transformer block — and report it two ways:

• all positions: every non-pad token of every sequence is a training example. This is the distribution we will fit the erasers on.
• last token: only the final token of each sequence, which is the position the model actually predicts from.

Keep the distinction in mind; in the caveats we will find that it matters a great deal.

To collect activations we use tracer.cache. Three rules, all of which we rely on later:

• Always pass modules=. An unfiltered tracer.cache() records every module in the tree — about 150 of them for GPT-2 — and copies them all to the host synchronously. That is the difference between a three-second collection and a five-minute one.
• Declare the cache as the first statement in the trace body, before any .output read or assignment. This is the opposite of the forward-pass ordering rule that governs everything else, and skipping it raises ValueError: tracer.cache() must be declared before reading or modifying a model value. The cache still records post-intervention values — that is exactly what we need in Step 4.
• Wrap read-only collection in torch.no_grad(). A trace runs with autograd on by default, which costs a lot of memory for nothing when we are only reading.

The helper below already takes an erasers argument. We pass None for now; apply_erasure is defined in Step 3.

In [8]:

def tokenize(texts):
    enc = tokenizer(texts, return_tensors="pt", padding=True)
    return enc["input_ids"].to(DEVICE), enc["attention_mask"].to(DEVICE)


def collect(texts, erasers=None, layers=(), batch_size=128):
    """Run the model over `texts`, optionally erasing inside the forward pass.

    Returns (activations at each layer in `layers`, stacked over every non-pad
    token position; final-token logits for every sequence)."""
    acts = {l: [] for l in layers}
    last_logits = []
    modules = [blocks[l] for l in layers]

    for start in range(0, len(texts), batch_size):
        input_ids, attn_mask = tokenize(texts[start:start + batch_size])

        with torch.no_grad(), model.trace(input_ids, attention_mask=attn_mask) as tracer:
            # Declared first, but it observes post-intervention values.
            cache = tracer.cache(modules=modules, device=torch.device("cpu"),
                                 dtype=torch.float32) if modules else None
            if erasers:
                apply_erasure(erasers)              # defined in Step 3
            logits = model.lm_head.output.save()

        keep = attn_mask.cpu().bool()
        for l in layers:
            acts[l].append(cache[f"model.transformer.h.{l}"].output[keep].half())
        lengths = keep.sum(1)
        last_logits.append(logits[torch.arange(len(lengths)), lengths - 1].float().cpu())

    return {l: torch.cat(v) for l, v in acts.items()}, torch.cat(last_logits)

def tokenize(texts): enc = tokenizer(texts, return_tensors="pt", padding=True) return enc["input_ids"].to(DEVICE), enc["attention_mask"].to(DEVICE) def collect(texts, erasers=None, layers=(), batch_size=128): """Run the model over `texts`, optionally erasing inside the forward pass. Returns (activations at each layer in `layers`, stacked over every non-pad token position; final-token logits for every sequence).""" acts = {l: [] for l in layers} last_logits = [] modules = [blocks[l] for l in layers] for start in range(0, len(texts), batch_size): input_ids, attn_mask = tokenize(texts[start:start + batch_size]) with torch.no_grad(), model.trace(input_ids, attention_mask=attn_mask) as tracer: # Declared first, but it observes post-intervention values. cache = tracer.cache(modules=modules, device=torch.device("cpu"), dtype=torch.float32) if modules else None if erasers: apply_erasure(erasers) # defined in Step 3 logits = model.lm_head.output.save() keep = attn_mask.cpu().bool() for l in layers: acts[l].append(cache[f"model.transformer.h.{l}"].output[keep].half()) lengths = keep.sum(1) last_logits.append(logits[torch.arange(len(lengths)), lengths - 1].float().cpu()) return {l: torch.cat(v) for l, v in acts.items()}, torch.cat(last_logits)

Because we pool over token positions, we need to know which sequence (and therefore which label) each row of the stacked activation matrix came from, and which rows are final tokens.

In [9]:

def position_index(texts, batch_size=128):
    seq_of_row, is_last_row = [], []
    for start in range(0, len(texts), batch_size):
        _, attn_mask = tokenize(texts[start:start + batch_size])
        for b, length in enumerate(attn_mask.sum(1).cpu().numpy()):
            seq_of_row += [start + b] * length
            is_last_row += [p == length - 1 for p in range(length)]
    return np.array(seq_of_row), np.array(is_last_row)

TRAIN_SEQ, TRAIN_LAST = position_index(train_texts)
TEST_SEQ,  TEST_LAST  = position_index(test_texts)

y_train = train_y[TRAIN_SEQ]      # label of every token position
y_test  = test_y[TEST_SEQ]
print(f"{len(y_train)} train token positions, {len(y_test)} test token positions")

def position_index(texts, batch_size=128): seq_of_row, is_last_row = [], [] for start in range(0, len(texts), batch_size): _, attn_mask = tokenize(texts[start:start + batch_size]) for b, length in enumerate(attn_mask.sum(1).cpu().numpy()): seq_of_row += [start + b] * length is_last_row += [p == length - 1 for p in range(length)] return np.array(seq_of_row), np.array(is_last_row) TRAIN_SEQ, TRAIN_LAST = position_index(train_texts) TEST_SEQ, TEST_LAST = position_index(test_texts) y_train = train_y[TRAIN_SEQ] # label of every token position y_test = test_y[TEST_SEQ] print(f"{len(y_train)} train token positions, {len(y_test)} test token positions")

20880 train token positions, 9744 test token positions

The probe itself is a full-batch LBFGS logistic regression on the GPU. sklearn's LogisticRegression gives the same answer (we check that below), but we will fit this probe dozens of times during the rank sweep and the GPU version keeps the notebook to a few minutes.

In [10]:

def probe_accuracy(X_train, y_tr, X_test, y_te, l2=1e-3, steps=150):
    """L2 logistic regression, standardised with the train statistics. Returns test accuracy."""
    X_train = torch.as_tensor(np.asarray(X_train)).float().to(DEVICE)
    X_test  = torch.as_tensor(np.asarray(X_test)).float().to(DEVICE)
    y_tr_t  = torch.as_tensor(np.asarray(y_tr)).float().to(DEVICE)
    y_te_t  = torch.as_tensor(np.asarray(y_te)).float().to(DEVICE)

    mu, sd = X_train.mean(0, True), X_train.std(0, True).clamp(min=1e-6)
    X_train, X_test = (X_train - mu) / sd, (X_test - mu) / sd

    w = torch.zeros(X_train.shape[1], device=DEVICE, requires_grad=True)
    b = torch.zeros(1, device=DEVICE, requires_grad=True)
    opt = torch.optim.LBFGS([w, b], max_iter=steps, history_size=20, tolerance_grad=1e-9,
                            tolerance_change=1e-12, line_search_fn="strong_wolfe")
    loss_fn = torch.nn.BCEWithLogitsLoss()

    def closure():
        opt.zero_grad()
        loss = loss_fn(X_train @ w + b, y_tr_t) + l2 * (w * w).sum()
        loss.backward()
        return loss

    opt.step(closure)
    with torch.no_grad():
        return float((((X_test @ w + b) > 0).float() == y_te_t).float().mean())

def probe_accuracy(X_train, y_tr, X_test, y_te, l2=1e-3, steps=150): """L2 logistic regression, standardised with the train statistics. Returns test accuracy.""" X_train = torch.as_tensor(np.asarray(X_train)).float().to(DEVICE) X_test = torch.as_tensor(np.asarray(X_test)).float().to(DEVICE) y_tr_t = torch.as_tensor(np.asarray(y_tr)).float().to(DEVICE) y_te_t = torch.as_tensor(np.asarray(y_te)).float().to(DEVICE) mu, sd = X_train.mean(0, True), X_train.std(0, True).clamp(min=1e-6) X_train, X_test = (X_train - mu) / sd, (X_test - mu) / sd w = torch.zeros(X_train.shape[1], device=DEVICE, requires_grad=True) b = torch.zeros(1, device=DEVICE, requires_grad=True) opt = torch.optim.LBFGS([w, b], max_iter=steps, history_size=20, tolerance_grad=1e-9, tolerance_change=1e-12, line_search_fn="strong_wolfe") loss_fn = torch.nn.BCEWithLogitsLoss() def closure(): opt.zero_grad() loss = loss_fn(X_train @ w + b, y_tr_t) + l2 * (w * w).sum() loss.backward() return loss opt.step(closure) with torch.no_grad(): return float((((X_test @ w + b) > 0).float() == y_te_t).float().mean())

Now the baseline collection — every block, no intervention.

In [11]:

base_train, _ = collect(train_texts, layers=range(N_LAYERS))
base_test,  _ = collect(test_texts,  layers=range(N_LAYERS))

acc_all  = [probe_accuracy(base_train[l], y_train, base_test[l], y_test) for l in range(N_LAYERS)]
acc_last = [probe_accuracy(base_train[l][TRAIN_LAST], y_train[TRAIN_LAST],
                           base_test[l][TEST_LAST],   y_test[TEST_LAST]) for l in range(N_LAYERS)]

print("all positions:", " ".join(f"{a:.3f}" for a in acc_all))
print("last token:   ", " ".join(f"{a:.3f}" for a in acc_last))

base_train, _ = collect(train_texts, layers=range(N_LAYERS)) base_test, _ = collect(test_texts, layers=range(N_LAYERS)) acc_all = [probe_accuracy(base_train[l], y_train, base_test[l], y_test) for l in range(N_LAYERS)] acc_last = [probe_accuracy(base_train[l][TRAIN_LAST], y_train[TRAIN_LAST], base_test[l][TEST_LAST], y_test[TEST_LAST]) for l in range(N_LAYERS)] print("all positions:", " ".join(f"{a:.3f}" for a in acc_all)) print("last token: ", " ".join(f"{a:.3f}" for a in acc_last))

all positions: 0.911 0.894 0.899 0.907 0.917 0.909 0.918 0.921 0.923 0.926 0.930 0.930
last token:    0.979 0.957 0.952 0.973 0.966 0.960 0.976 0.978 0.984 0.987 0.994 0.996

In [12]:

fig = px.line(pd.DataFrame({"block": list(range(N_LAYERS)) * 2,
                            "accuracy": acc_all + acc_last,
                            "positions": ["all positions"] * N_LAYERS + ["last token"] * N_LAYERS}),
              x="block", y="accuracy", color="positions", markers=True,
              title="Gender probe accuracy on held-out names, by block")
fig.add_hline(y=0.5, line_dash="dot", annotation_text="chance")
fig.update_yaxes(range=[0.45, 1.02])
fig.show()

fig = px.line(pd.DataFrame({"block": list(range(N_LAYERS)) * 2, "accuracy": acc_all + acc_last, "positions": ["all positions"] * N_LAYERS + ["last token"] * N_LAYERS}), x="block", y="accuracy", color="positions", markers=True, title="Gender probe accuracy on held-out names, by block") fig.add_hline(y=0.5, line_dash="dot", annotation_text="chance") fig.update_yaxes(range=[0.45, 1.02]) fig.show()