You need to agree to share your contact information to access this model
This repository is publicly accessible, but you have to accept the conditions to access its files and content.
This lens is derived from Gemma 4 E4B's weights and is provided subject to the Gemma Terms of Use (https://ai.google.dev/gemma/terms). By requesting access you agree to those terms.
Log in or Sign Up to review the conditions and access this model content.
Jacobian lens for Gemma 4 E4B, fitted with the KV cache on
A Jacobian lens (Anthropic's jlens) for
google/gemma-4-E4B-it: 41 per-layer matrices J_l (layers 0-40,
d_model 2560, fp32) that carry a residual-stream vector at layer l into the final layer's basis, decoded with the
model's own unembedding.
Why this one: Gemma 4 E4B needs its cache to run correctly
jlens runs a model with use_cache=False. Gemma 4 E4B shares keys and values across its last 18 layers
(num_kv_shared_layers: 18) through the cache, so without one those layers compute something else. Measured on the
model's own chat prompt: with use_cache=False the model's next-token argmax agrees with the real model at 0.000
of positions ("X" where it says "Paris"); with use_cache=True, 1.000 (max logit difference 0.000).
This lens was fitted through the corrected forward. In practice the matrices come out close to a lens fitted the default way (on one real sequence, KL(model || lens) at layer 40 is 0.44 here against 0.43 for solarkyle/jspace-lenses's E4B lens). What matters more is reading: activations taken with the cache off are wrong from layer 24 on, whichever lens you use.
Recipe
- jlens's own estimator (
jlens.fit), unchanged: the average input-output Jacobian, the first 16 positions and the last excluded. - The first 100 WikiText-103 passages of at least 600 characters (
jlens.examples.load_wikitext_prompts), each with one<bos>written in (this tokenizer adds none to plain text), truncated to 128 tokens. - bf16 model,
dim_batch32, one A100 80 GB: 59 minutes, 42 GB peak. Saved in fp32 (jlens issue #6: the fp16 default can overflow large entries). - Model revision
ee0ef60; the weights file is byte-identical to revisionfee6332(sha256cfbd3d2f...).
Checks it passed
- The patched forward, hooked at the last layer and unembedded, reproduces the full model's logits exactly.
- The paper's two-hop example: at "boot" in "The currency used in the country shaped like a boot is", Italy is rank 0 at layers 14-18.
- Planted texts both ways: a text about a wolf and a dog raises a wolf/dog/collar word set over an island/volcano set at layer 36, and a text about a volcano does the reverse.
- KL(model output || lens) by layer on a real chat-format sequence: 13.3 (layer 10), 20.2 (20), 5.7 (30), 2.7 (36), 0.44 (40).
Load it, with the cache on
import types, torch, transformers, jlens
from huggingface_hub import hf_hub_download
tok = transformers.AutoTokenizer.from_pretrained("google/gemma-4-E4B-it")
hf = transformers.AutoModelForImageTextToText.from_pretrained("google/gemma-4-E4B-it", dtype=torch.bfloat16, device_map="cuda")
m = jlens.from_hf(hf, tok, force_bos=False)
# the E4B's shared-KV layers need a cache: without this line the read-out is of a different model
m.forward = types.MethodType(lambda self, input_ids: self._text_module(input_ids=input_ids, use_cache=True), m)
lens = jlens.JacobianLens.load(hf_hub_download("supremeDavid/jlens-gemma-4-e4b-cache-on", "lens.pt"))
logits, _, _ = lens.apply(m, "<bos>Fact: The currency used in the country shaped like a boot is", positions=[-2])
fit_lens.py in this repo is the script that fitted it.
Licence
The lens is derived from Gemma 4 E4B's weights and is provided subject to the Gemma Terms of Use. Using it requires the base model, whose licence you accept on its own page. The fitting code is jlens (Apache 2.0).