Linear Probing

This notebook demonstrates how to use Linear Probing to train classifiers on model representations.

Setup

import importlib.util

DEV = True

if importlib.util.find_spec("google.colab") is not None:
    MODE = "colab-dev" if DEV else "colab"
else:
    MODE = "local"

if MODE == "colab":
    %pip install -q tdhook
elif MODE == "colab-dev":
    !rm -rf tdhook
    !git clone https://github.com/Xmaster6y/tdhook -b main
    %pip install -q ./tdhook

Usage

Load model and prepare data

from transformers import AutoTokenizer, AutoModelForCausalLM
from tensordict import TensorDict
from datasets import load_dataset
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
import torch

from tdhook.latent.probing import Probing, ProbeManager

model = AutoModelForCausalLM.from_pretrained("gpt2")
model.eval()
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

dataset = load_dataset("scikit-learn/imdb", split="train")

num_train = 100
num_test = 10

shuffled_dataset = dataset.shuffle(seed=42)
train_data = shuffled_dataset.select(range(num_train))
test_data = shuffled_dataset.select(range(num_train, num_train + num_test))

train_texts = [item["review"] for item in train_data]
train_labels = [1 if item["sentiment"] == "positive" else 0 for item in train_data]

test_texts = [item["review"] for item in test_data]
test_labels = [1 if item["sentiment"] == "positive" else 0 for item in test_data]

train_encoded = [tokenizer.encode(text, max_length=512, truncation=True) for text in train_texts]
test_encoded = [tokenizer.encode(text, max_length=512, truncation=True) for text in test_texts]

max_len = max(len(seq) for seq in train_encoded + test_encoded)
pad_token_id = tokenizer.pad_token_id

train_input_ids = torch.tensor([seq + [pad_token_id] * (max_len - len(seq)) for seq in train_encoded])
test_input_ids = torch.tensor([seq + [pad_token_id] * (max_len - len(seq)) for seq in test_encoded])

Set up linear probing

def compute_metrics(predictions, labels):
    return {"accuracy": accuracy_score(labels, predictions)}


probe_manager = ProbeManager(LogisticRegression, {}, compute_metrics)

Train probe on training data and evaluate on test data

with Probing(
    "transformer.h.(0|5|10).mlp$", probe_manager.probe_factory, additional_keys=["labels", "step_type"]
).prepare(model, in_keys=["input_ids"], out_keys=["logits"]) as hooked_model:
    with torch.no_grad():
        train_td = TensorDict(
            {
                "input_ids": train_input_ids,
                "labels": torch.tensor(train_labels),
                "step_type": "fit",
            },
            batch_size=len(train_texts),
        )
        train_out = hooked_model(train_td)

        test_td = TensorDict(
            {
                "input_ids": test_input_ids,
                "labels": torch.tensor(test_labels),
                "step_type": "predict",
            },
            batch_size=len(test_texts),
        )
        test_out = hooked_model(test_td)

Display probe metrics

print("Training metrics:")
for key, value in probe_manager.fit_metrics.items():
    print(f"  {key}: {value}")

print("\nTest metrics:")
for key, value in probe_manager.predict_metrics.items():
    print(f"  {key}: {value}")

Training metrics:
  transformer.h.0.mlp_fwd: {'accuracy': 1.0}
  transformer.h.5.mlp_fwd: {'accuracy': 1.0}
  transformer.h.10.mlp_fwd: {'accuracy': 1.0}

Test metrics:
  transformer.h.0.mlp_fwd: {'accuracy': 0.5}
  transformer.h.5.mlp_fwd: {'accuracy': 0.4}
  transformer.h.10.mlp_fwd: {'accuracy': 0.6}