probe/k_sparse_probing.py

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)