jason9464/activation-gradient-sae

Activation and gradient SAEs for Olmo-3-1025-7B

Sparse autoencoder (SAE) dictionaries trained on the SFT gradient and on the activation of the residual stream of Olmo-3-1025-7B, with analysis files, from a MATS 12.0 sprint project (2026-10-19 to 2026-10-30). The project asks how SFT teaches "diffuse" behaviours that no single training example explains.

The gradient SAEs decompose the gradient of the SFT loss with respect to the residual stream, dLoss/dh, of allenai/Olmo-3-1025-7B (the base checkpoint right before SFT) on allenai/Dolci-Think-SFT-7B. The loss is the summed cross-entropy over response tokens, and h is the output of decoder block L (model.model.layers[L], 0-indexed). Activation SAEs on h at the same places are the comparison. All dictionaries use a single seed (0), and the analysis files are exploratory.

The code is in the GitHub repo jason9464/MATS_sprint, which is private. The loading example at the end of this card needs only this Hub repo.

Layout

Path	Content
sae/<tag>.pt	92 SAE checkpoints
sae/<tag>.json	trainer output (config and evaluation)
sae/<tag>.metrics.json	metrics re-measured on the same evaluation rows, with the random-dictionary floor
sae/<tag>.random_lstsq.json	random dictionary with least-squares coefficients, same rows and support
baselines/pca_baseline.json	PCA-k reconstruction baseline (k = 64, 256, 1,024) per unit, layer and signal
baselines/pca_<unit>_l<L>_<arm>.pt	all eigenvalues, top 1,024 eigenvectors and the fit mean (30 files)
whitening.pt	per-layer whitening matrix W and means, shared by all three units
transforms/<unit>_transforms.pt	per-unit mu, cap and scale for each layer and signal (arms[arm][L])
jlens/<tag>_scan.pt	J-lens readouts of every feature's decoder direction (top-50 token ids for both signs, random-direction readouts, firing frequency)
jlens/*.json, jlens/*_textcross.pt	summary and comparison tables of the J-lens analysis, and the cross with top-firing texts
jlens/judge/, jlens/judge_gva/	items and raw outputs of the blind LLM judge
jlens/tok_position_check_rows.pt	rows of the firing-position check for token-level layer-23 gradient features

Dictionaries

Names are <unit>_l<layer>_<grad|act>_k<L0>.pt.

• Unit tok (every response position), chunk (mean over each full 32-token chunk of the response) or doc (mean over all response positions).
• Layer 7, 15, 19, 23 or 27. Target L0 64, 256 or 1024.
• Width 32,768 for tok and chunk, 4,096 for doc. Two width-32,768 doc references at layer 15 carry a suffix: doc_l15_grad_k64_w32768.pt, doc_l15_act_k64_w32768.pt.
• Matryoshka JumpReLU SAEs with nested prefixes at width/4, width/2 and width, trained for 60,000 steps at batch 4,096.

Unit	Documents	Rows
token	5,280	14,552,840
chunk	42,242	3,628,855
document	202,753	202,753

The holdout is 5% of documents. 2,242 "planted" documents (WildChat documents rewritten into a numbered-list template) are trained on and excluded from every metric.

Checkpoint keys

Key	Content
enc	fp16 [width, 4096]
dec	fp16 [4096, width], unit-norm columns
b_enc, theta	fp32 [width], theta is the per-feature JumpReLU threshold
b_dec	fp32 [4096]
groups	Matryoshka prefix ends, e.g. [8192, 16384, 32768]
transform	cap and scale of this unit, layer and signal
meta, cfg	tag, unit, layer, arm, width, target L0, and all trainer arguments
eval_train, eval_hold	FVU per prefix, L0, dead features, row count
eval_random	FVU of the untrained initial dictionary

whitening.pt has keys layers, W, grad_mu, act_mu, evals, n_tokens, damp_c. W[L] is fp32 [4096, 4096] and symmetric, and the means are [4096].

Input transform

SAEs do not see raw vectors.

• Gradient: x -> clip_norm(W (x - grad_mu), cap) / scale.
• Activation: x -> clip_norm(x - act_mu, cap) / scale.

W, grad_mu and act_mu come from whitening.pt at the checkpoint's layer, and cap and scale from the checkpoint's transform. W = U (Λ + 0.1 λ̄ I)^(-1/2) Uᵀ whitens the token-level gradient covariance, with λ̄ its mean eigenvalue. FVU is measured in this transformed space, so gradient and activation FVUs are errors on different targets. Inputs must be built like the training rows of the same unit.

Loading

import torch
from huggingface_hub import hf_hub_download

repo = "jason9464/mats-sprint-saes"
sae = torch.load(hf_hub_download(repo, "sae/doc_l23_grad_k64.pt"), map_location="cpu", weights_only=False)
wh = torch.load(hf_hub_download(repo, "whitening.pt"), map_location="cpu", weights_only=False)
L, arm = sae["meta"]["layer"], sae["meta"]["arm"]
cap, scale = sae["transform"]["cap"], sae["transform"]["scale"]
enc, dec = sae["enc"].float(), sae["dec"].float()  # [width, 4096], [4096, width]
b_enc, b_dec, theta = sae["b_enc"], sae["b_dec"], sae["theta"]


def to_sae_space(x):
    """Raw rows [n, 4096] (dLoss/dh for arm 'grad', h for arm 'act') -> SAE input."""
    if arm == "grad":
        x = (x.float() - wh["grad_mu"][L]) @ wh["W"][L]  # W is symmetric
    else:
        x = x.float() - wh["act_mu"][L]
    x = x * torch.clamp(cap / x.norm(dim=1, keepdim=True), max=1.0)
    return x / scale


def encode(x):
    z = (x - b_dec) @ enc.T + b_enc
    return z * (z > theta)  # JumpReLU with a per-feature threshold


def decode(f, k=None):
    """k = Matryoshka prefix size (one of sae["groups"]); None = full width."""
    k = k or f.shape[1]
    return f[:, :k] @ dec[:, :k].T + b_dec


x = to_sae_space(torch.randn(8, 4096))  # replace with real rows of the same unit
f = encode(x)
x_hat = decode(f)
print(tuple(f.shape), tuple(x_hat.shape), "mean L0", (f > 0).sum(1).float().mean().item())

Encoding is z = (x - b_dec) enc^T + b_enc followed by f = z * (z > theta). Decoding with prefix size k is f[:, :k] dec[:, :k]^T + b_dec. On the training server's holdout rows, this code reproduces the stored holdout FVU and L0 of doc_l23_grad_k64 (0.7115 and 74.8). Random input only checks shapes, and its L0 means nothing. To take a gradient decoder column d back to raw gradient space, use torch.linalg.solve(W, d). The J-lens analysis reads -W⁻¹d, the direction in which SFT pushes the residual stream.