train.py

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import trackio
from model import TopKSparseAutoencoder, parameter_count
from safetensors.torch import save_file
from sklearn.decomposition import PCA
from torch.nn import functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

PROJECT_DIR = Path(__file__).resolve().parent
FOUNDRY_DIR = PROJECT_DIR.parents[1]
SNIP_DIR = FOUNDRY_DIR / "projects" / "snip-0.4m"
BASE_MODEL = SNIP_DIR / "artifacts" / "snip-0.4m-base"
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "snip-scope"
DATA_DIR = PROJECT_DIR / "data"
CONTEXT = 128


def read_texts(path: Path) -> list[str]:
    with path.open("r", encoding="utf-8") as handle:
        return [json.loads(line)["text"] for line in handle if line.strip()]


def token_blocks(texts: list[str], tokenizer) -> np.ndarray:
    tokens = []
    for text in texts:
        tokens.extend(tokenizer.encode(text, add_special_tokens=False))
        tokens.append(tokenizer.eos_token_id)
    usable = len(tokens) // CONTEXT * CONTEXT
    return np.asarray(tokens[:usable], dtype=np.int64).reshape(-1, CONTEXT)


@torch.inference_mode()
def extract_activations(
    language_model,
    blocks: np.ndarray,
    max_tokens: int,
) -> tuple[np.ndarray, np.ndarray]:
    activations = []
    token_ids = []
    language_model.eval()
    for start in range(0, len(blocks), 16):
        batch = torch.from_numpy(blocks[start : start + 16])
        outputs = language_model(
            input_ids=batch,
            output_hidden_states=True,
            use_cache=False,
        )
        activations.append(outputs.hidden_states[-1].reshape(-1, 96).numpy())
        token_ids.append(batch.reshape(-1).numpy())
        if sum(len(item) for item in activations) >= max_tokens:
            break
    return (
        np.concatenate(activations)[:max_tokens].astype(np.float32),
        np.concatenate(token_ids)[:max_tokens],
    )


def evaluate_sae(
    model: TopKSparseAutoencoder, activations: torch.Tensor
) -> tuple[dict, np.ndarray]:
    model.eval()
    with torch.inference_mode():
        reconstruction, features = model(activations)
    mse = float(F.mse_loss(reconstruction, activations))
    variance = float(torch.var(activations, unbiased=False))
    active = (features > 1e-7).sum(1).float()
    firing = (features > 1e-7).float().mean(0).numpy()
    return (
        {
            "reconstruction_mse": mse,
            "explained_variance": 1 - mse / variance,
            "mean_active_features": float(active.mean()),
            "median_active_features": float(active.median()),
            "dead_feature_fraction": float(np.mean(firing == 0)),
        },
        features.numpy(),
    )


def main() -> None:
    torch.manual_seed(2043)
    torch.set_num_threads(1)
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    language_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL)
    train_blocks = token_blocks(
        read_texts(SNIP_DIR / "data" / "train.jsonl"), tokenizer
    )
    eval_blocks = token_blocks(
        read_texts(SNIP_DIR / "data" / "eval.jsonl"), tokenizer
    )
    train_activations, _ = extract_activations(
        language_model, train_blocks, max_tokens=180_000
    )
    eval_activations, eval_tokens = extract_activations(
        language_model, eval_blocks, max_tokens=40_000
    )
    mean = train_activations.mean(0)
    std = train_activations.std(0).clip(1e-4)
    train_normalized = (train_activations - mean) / std
    eval_normalized = (eval_activations - mean) / std
    model = TopKSparseAutoencoder()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-6)
    rng = np.random.default_rng(2043)
    tensor = torch.from_numpy(train_normalized)
    trackio.init(
        project="snip-scope",
        name="topk-sae-v1",
        config={
            "source_model": "SNIP-0.4M base",
            "training_tokens": len(train_normalized),
            "heldout_tokens": len(eval_normalized),
            "dictionary_features": model.features,
            "top_k": model.top_k,
        },
    )
    history = []
    model.train()
    for step in range(1, 3_001):
        indices = rng.choice(len(tensor), 1_024, replace=False)
        batch = tensor[indices]
        reconstruction, features = model(batch)
        loss = F.mse_loss(reconstruction, batch) + 1e-5 * features.mean()
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        model.normalize_dictionary()
        if step % 100 == 0:
            record = {
                "training_step": step,
                "training_loss": float(loss.detach()),
                "batch_active_features": float(
                    (features > 1e-7).sum(1).float().mean()
                ),
            }
            history.append(record)
            trackio.log(record)
    metrics, eval_features = evaluate_sae(
        model, torch.from_numpy(eval_normalized)
    )
    pca = PCA(n_components=16, random_state=2043).fit(train_normalized)
    pca_reconstruction = pca.inverse_transform(pca.transform(eval_normalized))
    pca_mse = float(np.mean((pca_reconstruction - eval_normalized) ** 2))
    firing = (eval_features > 1e-7).mean(0)
    exemplar_rows = []
    for feature in range(model.features):
        best = np.argsort(eval_features[:, feature])[-5:][::-1]
        for rank, index in enumerate(best, start=1):
            exemplar_rows.append(
                {
                    "feature": feature,
                    "rank": rank,
                    "token_id": int(eval_tokens[index]),
                    "token": tokenizer.decode([int(eval_tokens[index])]),
                    "activation": float(eval_features[index, feature]),
                    "firing_rate": float(firing[feature]),
                }
            )
    report = {
        "model": "SNIP Scope top-k sparse autoencoder",
        "source_model": "SNIP-0.4M base final hidden layer",
        "parameters": parameter_count(model),
        "training_tokens": len(train_normalized),
        "heldout_tokens": len(eval_normalized),
        "input_dimension": 96,
        "dictionary_features": model.features,
        "top_k": model.top_k,
        "heldout": metrics,
        "pca_16_control": {
            "reconstruction_mse": pca_mse,
            "explained_variance": 1
            - pca_mse / float(np.var(eval_normalized)),
        },
        "training_history": history,
    }
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    save_file(model.state_dict(), ARTIFACT_DIR / "sae.safetensors")
    np.savez(
        ARTIFACT_DIR / "normalization.npz",
        mean=mean.astype(np.float32),
        std=std.astype(np.float32),
    )
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(report, indent=2), encoding="utf-8"
    )
    pd.DataFrame(exemplar_rows).to_parquet(
        DATA_DIR / "feature_exemplars.parquet", index=False
    )
    trackio.log(
        {
            "heldout_explained_variance": metrics["explained_variance"],
            "heldout_dead_feature_fraction": metrics["dead_feature_fraction"],
            "heldout_mean_active_features": metrics["mean_active_features"],
            "pca_16_explained_variance": report["pca_16_control"][
                "explained_variance"
            ],
        }
    )
    trackio.finish()
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()