from collections import defaultdict from collections.abc import Sequence from pathlib import Path from typing import Callable import numpy as np import pandas as pd import torch from sae_lens import SAE from sklearn import metrics from sklearn.linear_model import LogisticRegression from torch import nn from tqdm.autonotebook import tqdm from transformer_lens import HookedTransformer from transformer_lens.hook_points import HookedRootModule from functools import partial # from sae_bench.evals.absorption.common import ( # PROBES_DIR, # RESULTS_DIR, # get_or_make_dir, # load_df_or_run, # load_dfs_or_run, # load_or_train_probe, # load_probe_data_split_or_train, # ) from .probing import LinearProbe, train_multi_probe from .util import batchify, DEFAULT_DEVICE from .util import get_sae_acts EPS = 1e-6 SPARSE_PROBING_EXPERIMENT_NAME = "k_sparse_probing" class KSparseProbe(nn.Module): weight: torch.Tensor # shape (k) bias: torch.Tensor # scalar feature_ids: torch.Tensor # shape (k) def __init__( self, weight: torch.Tensor, bias: torch.Tensor, feature_ids: torch.Tensor ): super().__init__() self.weight = weight self.bias = bias self.feature_ids = feature_ids @property def k(self) -> int: return self.weight.shape[0] def forward(self, x: torch.Tensor) -> torch.Tensor: filtered_acts = x[..., self.feature_ids] return filtered_acts @ self.weight + self.bias def to(self, device: torch.device | str): self.weight = self.weight.to(device) self.bias = self.bias.to(device) self.feature_ids = self.feature_ids.to(device) def train_sparse_multi_probe( x_train: torch.Tensor, # tensor of shape (num_samples, input_dim) y_train: torch.Tensor, # tensor of shape (num_samples, num_probes), with values in [0, 1] device: torch.device, l1_decay: float = 0.01, # l1 regularization strength num_probes: int | None = None, # inferred from y_train if None batch_size: int = 4096, num_epochs: int = 50, lr: float = 0.01, end_lr: float = 1e-5, l2_decay: float = 1e-6, show_progress: bool = True, verbose: bool = False, map_acts: Callable[[torch.Tensor], torch.Tensor] | None = None, probe_dim: int | None = None, ) -> LinearProbe: """ Train a multi-probe with L1 regularization on the weights. """ return train_multi_probe( x_train, y_train, num_probes=num_probes, batch_size=batch_size, num_epochs=num_epochs, lr=lr, end_lr=end_lr, weight_decay=l2_decay, show_progress=show_progress, verbose=verbose, device=device, extra_loss_fn=lambda probe, _x, _y: l1_decay * probe.weights.abs().sum(dim=-1).mean(), map_acts=map_acts, probe_dim=probe_dim, ) def train_k_sparse_probes( sae: HookedRootModule, train_labels: list[tuple[str, int]], # list of (token, letter number) pairs train_activations: torch.Tensor, # n_vocab X d_model ks: Sequence[int], map_acts: Callable[[torch.Tensor], torch.Tensor] | None = None, l1_decay: float = 0.01, batch_size_sae: int = 128, batch_size_lr: int = 4096, num_epochs: int = 50, device: torch.device | str = "cpu", ) -> dict[int, dict[int, KSparseProbe]]: # dict[k, dict[letter_id, probe]] """ Train k-sparse probes for each k in ks. Returns a dict of dicts, where the outer dict is indexed by k and the inner dict is the label. """ results: dict[int, dict[int, KSparseProbe]] = defaultdict(dict) with torch.no_grad(): labels = {label for _, label in train_labels} sparse_train_y = torch.nn.functional.one_hot( torch.tensor([idx for _, idx in train_labels]) ) if map_acts is None: map_acts = partial( get_sae_acts, sae=sae, batch_size=batch_size_lr, device=device, verbose=False, convert_to_cpu=False ) l1_probe = ( train_sparse_multi_probe( train_activations, sparse_train_y, l1_decay=l1_decay, num_epochs=num_epochs, batch_size=batch_size_lr, device=device, map_acts=map_acts, probe_dim=sae.cfg.d_sae, ) .float() .cpu() ) feature_acts = get_sae_acts( train_activations, sae, batch_size_sae, device, verbose=True, convert_to_cpu=True ).numpy() with torch.no_grad(): train_k_y = np.array([idx for _, idx in train_labels]) with tqdm(total=len(ks) * len(labels), desc="training k-probes") as pbar: for k in ks: for label in labels: # using topk and not abs() because we only want features that directly predict the label sparse_feat_ids = l1_probe.weights[label].topk(k).indices train_k_x = feature_acts[..., sparse_feat_ids] if k==1: train_k_x = train_k_x.reshape(-1, 1) # Use SKLearn here because it's much faster than torch if the data is small sk_probe = LogisticRegression( max_iter=500, class_weight="balanced" ).fit(train_k_x, (train_k_y == label).astype(np.int64)) probe = KSparseProbe( weight=torch.tensor(sk_probe.coef_[0]).float(), bias=torch.tensor(sk_probe.intercept_[0]).float(), # type: ignore feature_ids=sparse_feat_ids, ) results[k][label] = probe pbar.update(1) return results def train_l1_probe( sae: HookedRootModule, y_train: torch.Tensor, x_train: torch.Tensor, map_acts: Callable[[torch.Tensor], torch.Tensor] | None = None, l1_decay: float = 0.01, batch_size_lr: int = 4096, num_epochs: int = 50, device: torch.device | str = DEFAULT_DEVICE, ) -> LinearProbe: if map_acts is None: map_acts = partial( get_sae_acts, sae=sae, batch_size=batch_size_lr, device=device, verbose=False, convert_to_cpu=False ) l1_probe = ( train_sparse_multi_probe( x_train, y_train, l1_decay=l1_decay, num_epochs=num_epochs, batch_size=batch_size_lr, device=device, map_acts=map_acts, probe_dim=sae.cfg.d_sae, ) .float() .cpu() ) return l1_probe def get_probe_predictions( probe: LinearProbe | KSparseProbe, x: torch.Tensor, batch_size_lr: int = 4096, device: torch.device | str = DEFAULT_DEVICE, ) -> torch.Tensor: probe.eval() probe.to(device) predictions = [] with torch.no_grad(): for batch_x in batchify(x, batch_size_lr, show_progress=False): batch_x = batch_x.to(device) preds: torch.Tensor = probe(batch_x) predictions.append((torch.sigmoid(preds) > 0.5).cpu()) return torch.cat(predictions, dim=0) # shape (num_samples)