试验观察雅克比透镜 J-Lens

1# 官方仓库： https://github.com/anthropics/jacobian-lens
2# Fork 官方仓库后克隆个人仓库
3git clone https://github.com/你的用户名/jacobian-lens.git
4cd jacobian-lens

 1# 安装uv工具
 2pip install uv
 3
 4# 创建虚拟环境
 5uv venv
 6
 7# 激活环境
 8# Windows
 9.venv\Scripts\activate
10# Mac/Linux
11source .venv/bin/activate
12
13# 安装核心依赖
14uv pip install torch transformers accelerate datasets matplotlib
15
16# 本地可编辑安装项目依赖
17pip install -e .

 1from transformers import AutoModelForCausalLM, AutoTokenizer
 2
 3# 加载 Qwen3.5-4B 模型与分词器
 4model_name = "Qwen/Qwen3.5-4B"
 5model = AutoModelForCausalLM.from_pretrained(
 6    model_name,
 7    torch_dtype="auto",
 8    device_map="auto"
 9)
10tokenizer = AutoTokenizer.from_pretrained(model_name)

1lens_l(h) = unembed( J_l @ h )

1lens = jlens.fit(model, prompts=my_prompts, checkpoint_path="out/ckpt.pt")
2lens.save("out/jacobian_lens.pt")

 1# fitting.py
 2def jacobian_for_prompt(
 3    model: LensModel,
 4    prompt: str,
 5    source_layers: Sequence[int],
 6    *,
 7    target_layer: int | None = None,
 8    dim_batch: int = 8,
 9    max_seq_len: int = 128,
10    skip_first: int = SKIP_FIRST_N_POSITIONS,
11) -> tuple[dict[int, torch.Tensor], int, int]: 
12    """Compute the per-layer Jacobian estimator ``J_l`` for one prompt.
13    
14        Runs one forward pass on the prompt replicated ``dim_batch`` times along
15        the batch axis, retains the graph, then runs ``ceil(d_model / dim_batch)``
16        backward passes against it. Each backward computes ``dim_batch`` rows of
17        ``J_l`` at once: batch element ``b`` carries a one-hot cotangent at output
18        dimension ``dim_start + b``, set at every valid target position. See the
19        module docstring for the resulting estimator and how it relates to
20        a strict per-position Jacobian.
21    
22        Args:
23            model: The model to compute Jacobians for.
24            prompt: Input text.
25            source_layers: Layer indices ``l`` to compute ``J_l`` at.
26            target_layer: Layer to take gradients with respect to. Defaults to the
27                final layer; negative indices count from the end. In some cases,
28                targeting the penultimate layer can give a better-conditioned
29                ``J_l``.
30            dim_batch: Output dimensions computed per backward pass. Higher uses
31                more GPU memory (the prompt is replicated this many times); total
32                backward FLOPs are unchanged.
33            max_seq_len: Truncate the prompt to this many tokens.
34            skip_first: Leading positions to exclude; see :func:`valid_position_mask`.
35    
36        Returns:
37            ``(jacobians, seq_len, n_valid_positions)``. ``jacobians`` maps each
38            source layer to a ``[d_model, d_model]`` fp32 CPU tensor.
39        """
40        n_layers, d_model = model.n_layers, model.d_model
41        source_layers, target_layer = _check_layer_indices(
42            source_layers, target_layer, n_layers
43        )
44    
45        input_ids = model.encode(prompt, max_length=max_seq_len)
46        seq_len = input_ids.shape[1]
47        position_mask = valid_position_mask(seq_len, skip_first=skip_first)
48        n_valid_positions = int(position_mask.sum())
49        ...

1jacobians = {layer: torch.zeros(d_model, d_model, dtype=torch.float32) for layer in source_layers}
2n_passes = math.ceil(d_model / dim_batch)

 1# L154-L169
 2with (
 3    ActivationRecorder(
 4        model.layers,
 5        at=[*source_layers, target_layer],
 6        start_graph_at=min(source_layers),
 7    ) as recorder,
 8    torch.enable_grad(),
 9):
10    # One forward on the prompt replicated dim_batch times. The retained
11    # graph is reused for every backward pass below.
12    replicated_ids = input_ids.expand(dim_batch, -1)
13    model.forward(replicated_ids)
14    target_activation = recorder.activations[
15        target_layer
16    ]  # [dim_batch, seq_len, d_model]
17    source_activations = [recorder.activations[layer] for layer in source_layers]

1# L171-L175
2valid_positions = position_mask.nonzero(as_tuple=True)[0].to(
3    target_activation.device
4)
5batch_indices = torch.arange(dim_batch, device=target_activation.device)
6cotangent = torch.zeros_like(target_activation)

 1# L177-L221
 2for pass_idx, dim_start in enumerate(range(0, d_model, dim_batch)):
 3    n_dims_this_pass = min(dim_batch, d_model - dim_start)
 4    # One-hot cotangent at dim (dim_start + b) for batch element b,
 5    # at every valid target position. Yields rows dim_start..+n of J_l.
 6    cotangent.zero_()
 7    cotangent[
 8        batch_indices[:n_dims_this_pass, None],
 9        valid_positions[None, :],
10        dim_start + batch_indices[:n_dims_this_pass, None],
11    ] = 1.0
12    grads = torch.autograd.grad(
13        outputs=target_activation,
14        inputs=source_activations,
15        grad_outputs=cotangent,
16        retain_graph=(pass_idx < n_passes - 1),
17    )
18    for layer, grad in zip(source_layers, grads, strict=True):
19        # grad: [dim_batch, seq_len, d_model] on whatever device this
20        # layer lives on; mean over the valid positions -> dim_batch rows.
21        positions_on_device = valid_positions.to(grad.device, non_blocking=True)
22        rows = (
23            grad[:n_dims_this_pass, positions_on_device, :].float().mean(dim=1)
24        )
25        jacobians[layer][dim_start : dim_start + n_dims_this_pass, :] = (
26            rows.cpu()
27        )

 1# lens.py L52-L64
 2def save(self, path: str, *, dtype: torch.dtype = torch.float16) -> None:
 3    """Save to ``path``. Jacobians are stored as ``dtype`` (default fp16:
 4    halves file size; entries are O(1) so the range is not a constraint
 5    and fp16's extra mantissa bits beat bf16 here)."""
 6    torch.save(
 7        {
 8            "J": {layer: J.to(dtype) for layer, J in self.jacobians.items()},
 9            "n_prompts": self.n_prompts,
10            "source_layers": self.source_layers,
11            "d_model": self.d_model,
12        },
13        path,
14    )

 1import jlens
 2
 3# 拟合
 4lens = jlens.fit(model, prompts=my_prompts, checkpoint_path="out/ckpt.pt")
 5
 6# 保存最终 lens（默认 fp16，体积约为 fp32 一半）
 7lens.save("out/jacobian_lens.pt")
 8
 9# 加载
10lens = jlens.JacobianLens.load("out/jacobian_lens.pt")
11
12# 或从 HuggingFace
13lens = jlens.JacobianLens.from_pretrained("org/lens-repo", filename="model/lens.pt")

 1import transformers
 2import jlens
 3
 4# 1. 加载模型
 5hf = transformers.AutoModelForCausalLM.from_pretrained("org/model").cuda()
 6tok = transformers.AutoTokenizer.from_pretrained("org/model")
 7model = jlens.from_hf(hf, tok)
 8
 9# 2. 加载已拟合透镜
10lens = jlens.JacobianLens.load("out/jacobian_lens.pt")
11# 或
12lens = jlens.JacobianLens.from_pretrained("org/lens-repo", filename="model/lens.pt")
13
14# 3. 对新文本应用
15lens_logits, model_logits, input_ids = lens.apply(
16    model,
17    "Fact: The currency used in the country shaped like a boot is",
18    layers=[0, 10, 20],      # 可选：指定层；默认全部 source_layers
19    positions=[-2],          # 可选：指定位置；默认全部位置
20    max_seq_len=512,
21    use_jacobian=True,       # False 则退化为普通 logit lens（不做 J 运输）
22)
23
24# 4. 解读结果
25for layer, logits in sorted(lens_logits.items()):
26    top5 = logits[0].topk(5).indices  # positions=[-2] 时只有 1 行
27    print(f"Layer {layer}:", [tok.decode([t]) for t in top5])

 1# lens.py L66-L79
 2@classmethod
 3def load(cls, path: str) -> JacobianLens:
 4    """Load a lens previously written by :meth:`save`."""
 5    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
 6    if "J" not in checkpoint:
 7        raise ValueError(
 8            f"{path} is not a JacobianLens file "
 9            f"(found keys {sorted(checkpoint)!r}; a fit() checkpoint?)"
10        )
11    return cls(
12        jacobians=checkpoint["J"],
13        n_prompts=checkpoint["n_prompts"],
14        d_model=checkpoint["d_model"],
15    )

 1# lens.py L182-L194
 2if layers is None:
 3    layers = self.source_layers
 4out_of_range = sorted(l for l in set(layers) if not 0 <= l < model.n_layers)
 5if out_of_range:
 6    raise ValueError(...)
 7unknown = set(layers) - set(self.source_layers)
 8if use_jacobian and unknown:
 9    raise ValueError(
10        f"layers {sorted(unknown)} not in source_layers; "
11        f"fitted layers are {self.source_layers}"
12    )

1# lens.py L195-L201
2final_layer = model.n_layers - 1
3record_at = sorted(set(layers) | {final_layer})
4
5input_ids = model.encode(prompt, max_length=max_seq_len)
6with ActivationRecorder(model.layers, at=record_at) as recorder:
7    model.forward(input_ids)
8    activations = {i: recorder.activations[i].detach() for i in record_at}

1# lens.py L203-L206
2def select(layer: int) -> torch.Tensor:
3    """Residuals at the requested positions: ``[n_positions, d_model]``."""
4    full = activations[layer][0]  # [seq_len, d_model]
5    return (full if positions is None else full[list(positions)]).float()

 1# lens.py L208–L216
 2lens_logits: dict[int, torch.Tensor] = {}
 3for layer in layers:
 4    residual = select(layer)
 5    if use_jacobian:
 6        residual = self.transport(residual, layer)
 7    lens_logits[layer] = model.unembed(residual).float().cpu()
 8
 9model_logits = model.unembed(select(final_layer)).float().cpu()
10return lens_logits, model_logits, input_ids

1# lens.py L135-L143 
2def transport(self, residual: torch.Tensor, layer: int) -> torch.Tensor:
3    """Map a residual at ``layer`` into the final-layer basis: ``J_l @ h``."""
4    J_bar = self.jacobians[layer].to(residual.device)
5    return residual @ J_bar.T

 1import argparse
 2import gzip
 3import json
 4import os
 5import sys
 6from pathlib import Path
 7
 8import torch
 9import transformers
10
11import jlens
12from jlens.examples import EXAMPLES, load_wikitext_prompts, resolve_prompt
13from jlens.vis import build_page, compute_slice
14from jlens.data.zh_prompts import ZH_PROMPTS
15
16# ---------- 配置 ----------
17MODEL_NAME = "Qwen/Qwen3.5-4B"
18LENS_PATH = "qwen4b_jacobian_lens.pt"
19CHECKPOINT_PATH = "checkpoints/qwen4b_ckpt.pt"

 1def cmd_fit(args):
 2    """拟合透镜（支持断点恢复）。"""
 3    jlens.configure_logging()
 4
 5    print(f"[1/3] 加载模型 {MODEL_NAME} ...")
 6    hf_model = transformers.AutoModelForCausalLM.from_pretrained(
 7        MODEL_NAME,
 8        dtype=torch.bfloat16,
 9        #device_map="auto",
10        #device="mps",
11    ).to("mps")
12    tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)
13    model = jlens.from_hf(
14        hf_model, tokenizer, compile=args.compile, force_bos=True
15    )
16    print(f"  {model}")
17
18    print(f"[2/3] 加载 {args.prompts} 条 prompt ...")
19    prompts = load_wikitext_prompts(n_prompts=args.prompts, min_chars=600)
20    print(f"  已加载 {len(prompts)} 条")
21
22    print(f"[3/3] 拟合透镜 (dim_batch={args.dim_batch}, max_seq_len={args.max_seq_len}) ...")
23    print(f"  检查点: {args.checkpoint}")
24    print(f"  层选择: {args.layers or '所有层（除最后一层）'}")
25    print()
26
27    source_layers = None
28    if args.layers:
29        source_layers = [int(l) for l in args.layers.split(",")]
30
31    lens = jlens.fit(
32        model,
33        prompts, 
34        #ZH_PROMPTS,
35        dim_batch=args.dim_batch,
36        max_seq_len=args.max_seq_len,
37        source_layers=source_layers,
38        checkpoint_path=args.checkpoint,
39        checkpoint_every=1,
40        resume=True,
41    )
42
43    out = args.output or LENS_PATH
44    lens.save(out)
45    print(f"\n完成！透镜已保存到 {out}")
46    print(f"文件信息: {lens}")

 1def cmd_apply(args):
 2    """加载已拟合透镜并测试几条 prompt。"""
 3    print(f"加载模型 {MODEL_NAME} ...")
 4    hf_model = transformers.AutoModelForCausalLM.from_pretrained(
 5        MODEL_NAME, dtype=torch.bfloat16, device_map="auto"
 6    )
 7    tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)
 8    model = jlens.from_hf(hf_model, tokenizer)
 9
10    lens_path = args.lens or LENS_PATH
11    print(f"加载透镜 {lens_path} ...")
12    lens = jlens.JacobianLens.load(lens_path)
13    print(f"  {lens}")
14
15    # 选择要观察的层：1/4, 1/2, 3/4, 最后第二层
16    layers = [
17        model.n_layers // 4,
18        model.n_layers // 2,
19        model.n_layers // 4 * 3,
20        model.n_layers - 2,
21    ]
22
23    examples_to_run = EXAMPLES if args.all else [EXAMPLES[0]]
24    for example in examples_to_run:
25        print(f"\n{'='*60}")
26        print(f"示例: {example.section}")
27        prompt = resolve_prompt(example, tokenizer)
28        display = prompt[:80] + "..." if len(prompt) > 80 else prompt
29        print(f"Prompt: {display}")
30
31        jl_logits, model_logits, _ = lens.apply(
32            model, prompt, layers=layers, positions=[-2]
33        )
34        ll_logits, _, _ = lens.apply(
35            model, prompt, layers=layers, positions=[-2], use_jacobian=False
36        )
37
38        def top5(logits):
39            return [
40                tokenizer.decode([t], clean_up_tokenization_spaces=False)
41                for t in logits.topk(5).indices
42            ]
43
44        for layer in sorted(jl_logits):
45            print(f"  L{layer:>3}  J-lens:  {top5(jl_logits[layer][0])}")
46            print(f"         logit:   {top5(ll_logits[layer][0])}")
47        print(f"  model output:   {top5(model_logits[0])}")

 1def cmd_slice(args):
 2    """生成交互式切片 HTML 页面。"""
 3    print(f"加载模型 {MODEL_NAME} ...")
 4    hf_model = transformers.AutoModelForCausalLM.from_pretrained(
 5        MODEL_NAME, dtype=torch.bfloat16, device_map="auto"
 6    )
 7    tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)
 8    model = jlens.from_hf(hf_model, tokenizer)
 9
10    lens_path = args.lens or LENS_PATH
11    print(f"加载透镜 {lens_path} ...")
12    lens = jlens.JacobianLens.load(lens_path)
13
14    gloss = {}
15    gloss_path = Path("assets/qwen_gloss.json.gz")
16    if gloss_path.exists():
17        gloss = {
18            int(k): v
19            for k, v in json.load(gzip.open(gloss_path)).items()
20        }
21        print(f"  已加载 {len(gloss)} 条 token gloss")
22
23    example_map = {e.slug: e for e in EXAMPLES}
24    slug = args.example or "multihop"
25    example = example_map.get(slug, EXAMPLES[0])
26    prompt = resolve_prompt(example, tokenizer)
27    print(f"示例: {example.section}")
28    print(f"Prompt 长度: {len(prompt)} 字符")
29
30    slice_data = compute_slice(
31        model, lens, prompt,
32        layer_stride=2,
33        mask_display=True,
34    )
35
36    out_dir = Path(args.out_dir or f"slices/{slug}")
37    page, _, _ = build_page(
38        slice_data, prompt,
39        title=example.section,
40        description=example.description,
41        alt_token=gloss,
42        mode=args.mode,
43        out_dir=out_dir,
44    )
45
46    if args.mode == "embed":
47        out_file = out_dir / "index.html"
48        out_dir.mkdir(parents=True, exist_ok=True)
49        out_file.write_text(page)
50        print(f"页面已保存: {out_file}")
51        print("直接用浏览器打开或 notebook_iframe 加载")
52    else:
53        (out_dir / "index.html").write_text(page)
54        print(f"页面已保存: {out_dir / 'index.html'}")
55        print("启动方式: cd slices && python -m http.server 8000")

 1def main():
 2    parser = argparse.ArgumentParser(description="Qwen3.5-4B Jacobian Lens 拟合工具")
 3    sub = parser.add_subparsers(dest="command")
 4
 5    p_fit = sub.add_parser("fit", help="拟合透镜")
 6    p_fit.add_argument("--prompts", type=int, default=100)
 7    p_fit.add_argument("--dim-batch", type=int, default=8)
 8    p_fit.add_argument("--max-seq-len", type=int, default=128)
 9    p_fit.add_argument("--compile", action="store_true")
10    p_fit.add_argument("--layers", type=str, default=None)
11    p_fit.add_argument("--checkpoint", type=str, default=CHECKPOINT_PATH)
12    p_fit.add_argument("--output", "-o", type=str, default=None)
13
14    p_apply = sub.add_parser("apply", help="应用透镜测试 prompt")
15    p_apply.add_argument("--lens", type=str, default=None)
16    p_apply.add_argument("--all", action="store_true")
17
18    p_slice = sub.add_parser("slice", help="生成交互式切片页面")
19    p_slice.add_argument("--lens", type=str, default=None)
20    p_slice.add_argument("--example", type=str, default="multihop")
21    p_slice.add_argument("--mode", type=str, default="embed", choices=["embed", "fetch"])
22    p_slice.add_argument("--out-dir", type=str, default=None)
23
24    args = parser.parse_args()
25    if args.command == "fit":
26        cmd_fit(args)
27    elif args.command == "apply":
28        cmd_apply(args)
29    elif args.command == "slice":
30        cmd_slice(args)
31    else:
32        parser.print_help()
33
34
35if __name__ == "__main__":
36    main()

 1jacobian-lens ❯ uv run python fit_qwen4b.py fit --prompts 100
 2[1/3] 加载模型 Qwen/Qwen3.5-4B ...
 3Fetching 2 files: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 3442.19it/s]
 4Download complete: : 0.00B [00:00, ?B/s]              [transformers] The fast path is not available because one of the required library is not installed. Falling back to torch implementation. To install follow https://github.com/fla-org/flash-linear-attention#installation and https://github.com/Dao-AILab/causal-conv1d
 5Download complete: : 0.00B [00:00, ?B/s]
 6Loading weights: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 426/426 [00:00<00:00, 23670.58it/s]
 7  HFLensModel(Qwen3_5ForCausalLM, n_layers=32, d_model=2560)
 8[2/3] 加载 100 条 prompt ...
 9  已加载 100 条
10[3/3] 拟合透镜 (dim_batch=8, max_seq_len=128) ...
11  检查点: checkpoints/qwen4b_ckpt.pt
12  层选择: 所有层（除最后一层）
13
14[  2m46s +166.51s] fit: n_layers=32 d_model=2560, fitting 31 source layers (target=L31) on 100 prompts
15[ 12m00s +553.91s]   prompt 1/100  seq_len=128 n_valid=111  554s  max||J||/sqrt(d)=6.291  max_d_mean=nan
16[ 20m57s +536.89s]   prompt 2/100  seq_len=128 n_valid=111  536s  max||J||/sqrt(d)=4.775  max_d_mean=5.65e-01