andyx10/jacobian-lens-mistral-small-24b-instruct-2501

Jacobian Lens for Mistral-Small-24B-Instruct-2501

Jacobian lens for mistralai/Mistral-Small-24B-Instruct-2501, from Anthropic's Verbalizable Representations Form a Global Workspace in Language Models, fitted with the reference implementation anthropics/jacobian-lens.

Contents

mistral_small24b_2501_lens.pt — a torch.save dict with keys J (dict layer -> [5120, 5120] fp16 tensor), n_prompts, source_layers, d_model. Load it with jlens.JacobianLens.load(...).

field	value
base model	mistralai/Mistral-Small-24B-Instruct-2501 (40 layers, d_model=5120)
source layers	0–38 (all 39 layers below the final)
target layer	39 (final)
Jacobian shape	[5120, 5120] per layer, fp16 on disk
file size	~1.90 GiB
fitting corpus	100 WikiText-103 prompts × 128 tokens (first 16 positions skipped as attention sinks)
estimator	dim_batch=16, mean over valid source positions
fit hardware	1× NVIDIA H100 80GB, 2.1 h (76 s/prompt)

Usage

import torch, transformers, jlens
from huggingface_hub import hf_hub_download

MODEL = "mistralai/Mistral-Small-24B-Instruct-2501"
hf = transformers.AutoModelForCausalLM.from_pretrained(
    MODEL, dtype=torch.bfloat16, device_map="cuda"
)
tok = transformers.AutoTokenizer.from_pretrained(MODEL)
model = jlens.from_hf(hf, tok)

path = hf_hub_download(
    "dormantx/jacobian-lens-mistral-small-24b-instruct-2501",
    "mistral_small24b_2501_lens.pt",
)
lens = jlens.JacobianLens.load(path)

# per-layer lens logits at the last position, J-lens vs logit-lens baseline
lens_logits, model_logits, _ = lens.apply(
    model, "Fact: The currency used in the country shaped like a boot is",
    positions=[-1], layers=[23, 27, 32], use_jacobian=True,
)

Note: the lens was fitted and applied with the checkpoint's default tokenizer. transformers ≥5.13 warns about a known Tekken regex quirk (fix_mistral_regex=True); readouts here are self-consistent with the default setting — only enable the flag if you also refit.

License

Apache-2.0, matching the reference implementation. The base model mistralai/Mistral-Small-24B-Instruct-2501 is subject to its own license (Apache-2.0).