Auditing a published weight edit with TransformerBridge¶
This notebook examines the Qwen3.5-0.8B edit in Experiment 8 of Sparks of Weight-Space Programming by Rikiya Takehi, Jayden Teoh, Alex Zhang and Omar Khattab.
We reconstruct the released p50c checkpoint change, fit its dominant direction, translate the exact weight difference into a hook intervention, and run a coefficient sweep with controls. The default path uses CPU and fp32, with 32-token agreement checks and 64-token sweep windows. The larger bf16 transcript experiment uses the same sweep at a larger budget.
The experiment logs and edit artifacts are credited to their authors under CC BY 4.0. Our reconstruction below uses the published data, rather than executing code from that repository. See the source license.
1. Pinned inputs and runtime¶
The first run downloads the checkpoint (roughly 2 GB). The remaining runs reuse the Hugging Face cache. Allow about 6 GiB for the Python process. Time measurements below separate artifact preparation, loading, and forward validation; download time depends on the connection.
[1]:
# NBVAL_IGNORE_OUTPUT
import time
notebook_started = time.perf_counter()
import hashlib
import io
import json
import platform
import tempfile
import urllib.request
from contextlib import ExitStack, contextmanager
from pathlib import Path
import numpy as np
import pandas as pd
import psutil
import torch
import torch.nn.functional as F
import transformers
from huggingface_hub import snapshot_download
from safetensors import safe_open
from transformer_lens.model_bridge import TransformerBridge
from transformer_lens.utilities import temporarily_swap_parameter
MODEL_ID = "Qwen/Qwen3.5-0.8B"
MODEL_REVISION = "2fc06364715b967f1860aea9cf38778875588b17"
LOGS_REVISION = "4d09ba6d6a7e891dd8063974dfe81b3d78e3494f"
LAYERS = (15, 19)
MAX_TOKENS = 32
torch.set_num_threads(2)
started = time.perf_counter()
timings = {}
print(
{
"python": platform.python_version(),
"torch": torch.__version__,
"transformers": transformers.__version__,
"device": "cpu",
"dtype": "float32",
}
)
{'python': '3.12.13', 'torch': '2.11.0+cpu', 'transformers': '5.13.0', 'device': 'cpu', 'dtype': 'float32'}
2. Reconstruct the published tensors¶
The p50c artifact stores sparse replacement values, not a complete difference matrix. We read the original bf16 tensors, replace the specified entries, and verify both byte hashes against the authors’ manifest. Hashing happens before conversion to fp32.
The two small downloaded artifacts also have pinned SHA-256 hashes. Cached artifacts receive the same checks as fresh downloads.
[2]:
# NBVAL_IGNORE_OUTPUT
ARTIFACT_HASHES = {
"manifest.json": "83185c1db46e3dafb34af645f8c5ce033c09cc99d4512b9e1b8e3aba9d269225",
"edit.npz": "e92a92aafd1f1068b4fd94aa0fbe081ba3ddea2047ee7bcac7161d861be299e7",
}
artifact_cache = (
Path(tempfile.gettempdir()) / "transformerlens-weight-edit-audit" / LOGS_REVISION
)
artifact_cache.mkdir(parents=True, exist_ok=True)
def sha256(data):
return hashlib.sha256(data).hexdigest()
def read_artifact(filename):
destination = artifact_cache / filename
if destination.exists():
data = destination.read_bytes()
else:
url = (
f"https://raw.githubusercontent.com/RikiyaT/lm_understand-logs/"
f"{LOGS_REVISION}/exp8_math/runs/p50c/edit/{filename}"
)
with urllib.request.urlopen(url, timeout=60) as response:
data = response.read()
if sha256(data) != ARTIFACT_HASHES[filename]:
raise ValueError(f"Artifact checksum mismatch: {filename}")
if not destination.exists():
destination.write_bytes(data)
return data
manifest = json.loads(read_artifact("manifest.json"))
archive = np.load(io.BytesIO(read_artifact("edit.npz")), allow_pickle=False)
snapshot = Path(
snapshot_download(
MODEL_ID,
revision=MODEL_REVISION,
allow_patterns=["*.json", "*.safetensors", "*.txt", "*.jinja"],
max_workers=1,
)
)
index = json.loads(
(snapshot / "model.safetensors.index.json").read_text(encoding="utf-8")
)
expected_names = {
f"model.language_model.layers.{layer}.mlp.down_proj.weight" for layer in LAYERS
}
assert {item["name"] for item in manifest["tensors"]} == expected_names
edits = {}
for number, item in enumerate(manifest["tensors"]):
assert item["form"] == "sparse" and item["dtype"] == "bfloat16"
with safe_open(
str(snapshot / index["weight_map"][item["name"]]), framework="pt"
) as reader:
base = reader.get_tensor(item["name"])
assert base.dtype == torch.bfloat16 and list(base.shape) == item["shape"]
base_bits = base.view(torch.uint16).numpy().copy()
assert sha256(base_bits.tobytes()) == item["base_sha256"]
indices = archive[f"t{number}_idx"]
values = archive[f"t{number}_val"]
assert indices.ndim == values.ndim == 1 and indices.size == values.size
assert np.unique(indices).size == indices.size
assert (indices >= 0).all() and (indices < base.numel()).all()
edited_bits = base_bits.copy()
edited_bits.reshape(-1)[indices] = values
assert sha256(edited_bits.tobytes()) == item["edited_sha256"]
changed_entries = int(np.count_nonzero(base_bits != edited_bits))
assert changed_entries == item["changed_entries"]
edited = torch.from_numpy(edited_bits.view(np.int16)).view(torch.bfloat16)
layer = int(item["name"].split(".layers.")[1].split(".")[0])
columns = torch.from_numpy(np.unique(indices % base.shape[1])).long()
assert columns.numel() == item["changed_columns"] == 64
edits[layer] = {
"base": base.float(),
"edited": edited.float(),
"columns": columns,
"delta": (edited.float() - base.float())[:, columns],
"changed_entries": changed_entries,
}
assert (
sum(e["changed_entries"] for e in edits.values())
== manifest["changed_entries_total"]
)
archive.close()
timings["artifacts_and_checkpoint_seconds"] = time.perf_counter() - started
print(
{
layer: {
"columns": len(e["columns"]),
"changed_bf16_entries": e["changed_entries"],
}
for layer, e in edits.items()
}
)
{15: {'columns': 64, 'changed_bf16_entries': 65260}, 19: {'columns': 64, 'changed_bf16_entries': 64928}}
3. Recover the dominant directional update¶
For HF’s output projection, W has shape [d_model, d_mlp]. Fit the selected columns to
\[\Delta W_{:,S} \approx -\alpha u (u^\top W_{:,S}), \qquad \|u\|=1.\]
The fitted directional component is multiplied by 1 - alpha: alpha=1 deletes it; alpha near 6 reverses and amplifies it. The quantized checkpoint has a small residual, which we measure explicitly. The sign chosen for the singular vector does not change the edit.
[3]:
# NBVAL_IGNORE_OUTPUT
fit_started = time.perf_counter()
geometry = []
for layer, edit in edits.items():
delta = edit["delta"].double()
u, singular_values, _ = torch.linalg.svd(delta, full_matrices=False)
direction = u[:, 0]
projection = direction @ edit["base"][:, edit["columns"]].double()
alpha = -((direction @ delta) @ projection / projection.square().sum()).item()
fitted_delta = -alpha * torch.outer(direction, projection)
edit.update(
direction=direction.float(), alpha=alpha, fitted_delta=fitted_delta.float()
)
geometry.append(
{
"layer": layer,
"alpha": alpha,
"direction_multiplier": 1 - alpha,
"leading_energy_fraction": (
singular_values[0].square() / singular_values.square().sum()
).item(),
"relative_fit_residual": (delta - fitted_delta).norm().item()
/ delta.norm().item(),
}
)
timings["rank_one_fit_seconds"] = time.perf_counter() - fit_started
display(pd.DataFrame(geometry))
| layer | alpha | direction_multiplier | leading_energy_fraction | relative_fit_residual | |
|---|---|---|---|---|---|
| 0 | 15 | 6.000321 | -5.000321 | 0.999995 | 0.002347 |
| 1 | 19 | 6.000173 | -5.000173 | 0.999991 | 0.002944 |
4. Exact weight differences as hooks¶
If the down projection receives activation a, replacing its weights changes its output by
\[\Delta y = a_{S}\,\Delta W_{:,S}^{\top}.\]
We first check this identity on a small float64 linear example. Then we validate the exact difference on Qwen in fp32, separately from the fitted direction.
The hook pair reads mlp.hook_post (the down projection’s input) and changes mlp.hook_out. Each layer uses activations from the current forward, including earlier interventions.
[4]:
# NBVAL_IGNORE_OUTPUT
generator = torch.Generator().manual_seed(0)
a = torch.randn(2, 8, generator=generator, dtype=torch.float64)
w = torch.randn(4, 8, generator=generator, dtype=torch.float64)
delta = torch.randn(4, 8, generator=generator, dtype=torch.float64)
bias = torch.randn(4, generator=generator, dtype=torch.float64)
torch.testing.assert_close(
F.linear(a, w + delta, bias),
F.linear(a, w, bias) + F.linear(a, delta),
rtol=1e-12,
atol=1e-12,
)
print("Float64 linear identity passed")
Float64 linear identity passed
[5]:
# NBVAL_IGNORE_OUTPUT
load_started = time.perf_counter()
available_gib = psutil.virtual_memory().available / 2**30
print(f"Available RAM before loading: {available_gib:.1f} GiB")
bridge = TransformerBridge.boot_transformers(
str(snapshot),
device="cpu",
dtype=torch.float32,
).eval()
layers = bridge.original_model.model.language_model.layers
for layer, edit in edits.items():
parameter = layers[layer].mlp.down_proj.weight
assert parameter.dtype == torch.float32
assert torch.equal(parameter.detach(), edit["base"])
timings["model_load_seconds"] = time.perf_counter() - load_started
print("Pinned checkpoint loaded on CPU in fp32")
Available RAM before loading: 4.6 GiB
Pinned checkpoint loaded on CPU in fp32
[6]:
# NBVAL_IGNORE_OUTPUT
@contextmanager
def replaced_weights(replacements):
with ExitStack() as stack:
for layer, value in replacements.items():
stack.enter_context(
temporarily_swap_parameter(layers[layer].mlp.down_proj.weight, value)
)
yield
def difference_hooks(use_fit=False):
pending, hooks = {}, []
for layer, edit in edits.items():
columns = edit["columns"]
difference = edit["fitted_delta"] if use_fit else edit["delta"]
def keep(activation, hook, layer=layer, columns=columns):
pending[layer] = activation[..., columns]
def change(output, hook, layer=layer, difference=difference):
correction = F.linear(pending.pop(layer), difference)
return output + correction
hooks.extend(
[
(f"blocks.{layer}.mlp.hook_post", keep),
(f"blocks.{layer}.mlp.hook_out", change),
]
)
return hooks
@torch.no_grad()
def forward_capture(tokens, hooks=()):
captured = {}
def capture(activation, hook):
captured[hook.name] = activation.detach().clone()
names = [f"blocks.{layer}.mlp.hook_out" for layer in LAYERS]
with bridge.hooks(fwd_hooks=list(hooks) + [(name, capture) for name in names]):
logits = bridge(tokens)
return logits, captured
def error_summary(actual, expected):
difference = (actual - expected).abs()
return {"mean_abs": difference.mean().item(), "max_abs": difference.max().item()}
texts = [
"First compute the value, then check the result. " * 4,
"Let x be the unknown. Substitute x into the equation. " * 4,
]
inputs = [
bridge.tokenizer(text, add_special_tokens=False, return_tensors="pt").input_ids[
:, :MAX_TOKENS
]
for text in texts
]
validation_started = time.perf_counter()
comparisons = []
for input_number, tokens in enumerate(inputs):
baseline, _ = forward_capture(tokens)
with replaced_weights({layer: e["edited"] for layer, e in edits.items()}):
swapped, swapped_components = forward_capture(tokens)
assert not torch.allclose(swapped, baseline, rtol=1e-4, atol=1e-4)
exact, exact_components = forward_capture(tokens, difference_hooks())
fitted, _ = forward_capture(tokens, difference_hooks(use_fit=True))
assert torch.isfinite(fitted).all()
for name in swapped_components:
comparisons.append(
{
"input": input_number,
"tokens": tokens.shape[1],
"comparison": f"exact hook vs swapped: {name}",
**error_summary(exact_components[name], swapped_components[name]),
}
)
torch.testing.assert_close(
exact_components[name], swapped_components[name], rtol=1e-4, atol=1e-4
)
torch.testing.assert_close(exact, swapped, rtol=1e-4, atol=1e-4)
for label, actual, expected in [
("exact hook vs swapped", exact, swapped),
("fitted hook vs exact hook", fitted, exact),
("edited vs baseline", swapped, baseline),
]:
comparisons.append(
{
"input": input_number,
"tokens": tokens.shape[1],
"comparison": label,
**error_summary(actual, expected),
}
)
def token_nll(logits):
return F.cross_entropy(logits[0, :-1].float(), tokens[0, 1:], reduction="none")
swapped_nll, exact_nll = token_nll(swapped), token_nll(exact)
torch.testing.assert_close(exact_nll, swapped_nll, rtol=1e-4, atol=1e-4)
comparisons.append(
{
"input": input_number,
"tokens": tokens.shape[1],
"comparison": "exact hook vs swapped: token NLL",
**error_summary(exact_nll, swapped_nll),
}
)
restored, _ = forward_capture(tokens)
torch.testing.assert_close(restored, baseline, rtol=0, atol=0)
# Exercise restoration when the experiment body raises.
try:
with replaced_weights({layer: e["edited"] for layer, e in edits.items()}):
raise ValueError("restoration probe")
except ValueError as error:
assert str(error) == "restoration probe"
for layer, edit in edits.items():
assert torch.equal(layers[layer].mlp.down_proj.weight.detach(), edit["base"])
timings["forward_validation_seconds"] = time.perf_counter() - validation_started
display(pd.DataFrame(comparisons))
| input | tokens | comparison | mean_abs | max_abs | |
|---|---|---|---|---|---|
| 0 | 0 | 32 | exact hook vs swapped: blocks.15.mlp.hook_out | 8.780711e-09 | 3.054738e-07 |
| 1 | 0 | 32 | exact hook vs swapped: blocks.19.mlp.hook_out | 2.741603e-08 | 2.160668e-07 |
| 2 | 0 | 32 | exact hook vs swapped | 1.104146e-06 | 1.049042e-05 |
| 3 | 0 | 32 | fitted hook vs exact hook | 3.838236e-04 | 3.693104e-03 |
| 4 | 0 | 32 | edited vs baseline | 7.012547e-01 | 9.933302e+00 |
| 5 | 0 | 32 | exact hook vs swapped: token NLL | 9.273570e-07 | 7.629395e-06 |
| 6 | 1 | 32 | exact hook vs swapped: blocks.15.mlp.hook_out | 8.276834e-09 | 3.278255e-07 |
| 7 | 1 | 32 | exact hook vs swapped: blocks.19.mlp.hook_out | 2.772856e-08 | 4.470348e-07 |
| 8 | 1 | 32 | exact hook vs swapped | 1.139451e-06 | 9.059906e-06 |
| 9 | 1 | 32 | fitted hook vs exact hook | 3.841585e-04 | 4.105568e-03 |
| 10 | 1 | 32 | edited vs baseline | 4.792581e-01 | 7.829380e+00 |
| 11 | 1 | 32 | exact hook vs swapped: token NLL | 1.048819e-06 | 4.172325e-06 |
The assertion above applies to the exact difference, not the fitted rank-one edit. The error table reports both comparisons. The restored baseline also matches after normal completion and an exception.
5. Run the coefficient sweep¶
The following cells generate both result tables: selected-column interventions, whole-output removal, random-direction and random-column controls, followed by a coefficient sweep. The default is CPU/fp32 with two transcripts from one pinned shard and 64-token early/late windows. Change the settings at the top of the next cell to CUDA/bf16, ten shards, 40 samples and 1,024 tokens for the full-size run.
Each copy token ends an 8-gram already seen in its window. NLL is teacher-forced, and each table entry gives equal weight to each window containing that token category. An empty category is reported as missing, not zero; the counts below show which categories contribute. Random controls use the same coefficient and column count, with update norms determined by their chosen directions and columns.
The source files are pinned and SHA-256 checked. The returned sweep_record retains sample IDs, file hashes, per-window measurements, summaries and timings for reuse.
[7]:
# NBVAL_IGNORE_OUTPUT
import gzip
# Full-size run: cuda/bfloat16, 10 parts, 40 samples, 1,024-token windows.
SWEEP_DEVICE = "cpu"
SWEEP_DTYPE = torch.float32
TRANSCRIPT_PARTS = 1
SWEEP_SAMPLES = 2
WINDOW_TOKENS = 64
MIN_TRANSCRIPT_TOKENS = 32_000
COEFFICIENTS = (0.5, 1.0, 2.0, 3.0, 4.0, 6.0, 8.0, 10.0)
SEED = 0
SOURCE_HASHES = (
"7d153ba1036f3c4e5487effeedbee57bd595bfb5ed93019d843186a24f710d44",
"7452777e70dc791cc95247e8aa1bca409871c467bbb4f14b4402eae32e4fe751",
"5e73113409c5092a266390fbe6f4bf1ac60a136d064880928594d7767fb5b38f",
"1233548bdf0198ed2b04a4e6f129c5819413f8a52ba9272520a2f2bbe2a1b938",
"8fbcb6e9facbba93cc5af749f187f02ba9d9d1e89d0adba0078ac757fa5ed50e",
"cac46d5b0c1d7269e73a76b2c5ca668fdefde2a5defb623bb7a2df1b8afb20a4",
"c65c554b11aba8e4a19097d77a6b558556971f896a9c0897b497da45a60ee4fc",
"67f7bde178a9e274bab4389499c8c77c439c2ed832e1369a5bea8cd0f37346bb",
"7691ba5f0f07b2c9b83a68586a167b7621a680c7a565aa27098fd4a5fa00cd16",
"a9d7aa126016439f6e4e7629f2275b15d277e5dc1fb594d9462c170eb6de9ca1",
)
assert 1 <= TRANSCRIPT_PARTS <= len(SOURCE_HASHES)
assert 8 < WINDOW_TOKENS <= MIN_TRANSCRIPT_TOKENS
assert SWEEP_SAMPLES > 0
sweep_started = time.perf_counter()
if SWEEP_DEVICE == "cpu" and SWEEP_DTYPE == torch.float32:
sweep_bridge = bridge
else:
sweep_bridge = TransformerBridge.boot_transformers(
str(snapshot), device=SWEEP_DEVICE, dtype=SWEEP_DTYPE
).eval()
if SWEEP_DEVICE == "cuda":
torch.cuda.reset_peak_memory_stats()
# Draw controls before samples so the full-size run reproduces the recorded IDs.
rng = np.random.default_rng(SEED)
sweep_edits = {}
for layer, edit in edits.items():
direction = edit["direction"].to(SWEEP_DEVICE)
random = torch.randn(
direction.shape, generator=torch.Generator().manual_seed(layer)
).to(SWEEP_DEVICE)
random -= (random @ direction) * direction
pool = np.setdiff1d(np.arange(edit["base"].shape[1]), edit["columns"].numpy())
sweep_edits[layer] = {
"W": edit["base"].to(SWEEP_DEVICE),
"columns": edit["columns"].to(SWEEP_DEVICE),
"direction": direction,
"alpha": edit["alpha"],
"random_direction": random / random.norm(),
"random_columns": torch.from_numpy(
rng.choice(pool, len(edit["columns"]), replace=False)
).to(SWEEP_DEVICE),
}
def transcript_rows(parts):
rows, sources = [], []
for part in range(parts):
name = f"math500_s{part}of10"
relative = (
f"exp8_math/eval/unedited/parts/{name}/math500_first_samples.jsonl.gz"
)
destination = artifact_cache / f"{name}.jsonl.gz"
if destination.exists():
data = destination.read_bytes()
else:
url = f"https://raw.githubusercontent.com/RikiyaT/lm_understand-logs/{LOGS_REVISION}/{relative}"
with urllib.request.urlopen(url, timeout=60) as response:
data = response.read()
assert sha256(data) == SOURCE_HASHES[part], f"Transcript checksum: {name}"
if not destination.exists():
destination.write_bytes(data)
sources.append({"name": name, "sha256": sha256(data)})
rows.extend(
json.loads(line)
for line in gzip.decompress(data).decode("utf-8").splitlines()
)
return rows, sources
def copy_mask(sequence):
seen, mask = set(), np.zeros(len(sequence), dtype=bool)
for position in range(7, len(sequence)):
key = tuple(sequence[position - 7 : position + 1])
mask[position] = key in seen
seen.add(key)
return mask[1:]
assert copy_mask(list(range(8)) * 3).sum() == 9
assert not copy_mask(list(range(24))).any()
rows, source_files = transcript_rows(TRANSCRIPT_PARTS)
tokenized = [
sweep_bridge.tokenizer(row["text"], add_special_tokens=False).input_ids
for row in rows
]
eligible = np.flatnonzero(
[len(tokens) >= MIN_TRANSCRIPT_TOKENS for tokens in tokenized]
)
assert len(eligible) >= SWEEP_SAMPLES
picked = np.sort(rng.choice(eligible, SWEEP_SAMPLES, replace=False))
selected_ids = [rows[i]["id"] for i in picked]
windows = {
"late": [tokenized[i][-WINDOW_TOKENS:] for i in picked],
"early": [tokenized[i][:WINDOW_TOKENS] for i in picked],
}
masks = {
group: [copy_mask(tokens) for tokens in sequences]
for group, sequences in windows.items()
}
display(
pd.DataFrame(
[
{
"group": group,
"id": sample_id,
"copy_tokens": int(mask.sum()),
"other_tokens": int((~mask).sum()),
}
for group, group_masks in masks.items()
for sample_id, mask in zip(selected_ids, group_masks)
]
)
)
| group | id | copy_tokens | other_tokens | |
|---|---|---|---|---|
| 0 | late | test/prealgebra/930.json | 40 | 23 |
| 1 | late | test/counting_and_probability/51.json | 0 | 63 |
| 2 | early | test/prealgebra/930.json | 0 | 63 |
| 3 | early | test/counting_and_probability/51.json | 0 | 63 |
[8]:
# NBVAL_IGNORE_OUTPUT
def sweep_hooks(alpha, control=None, whole=False):
hooks, pending = [], {}
for layer, edit in sweep_edits.items():
direction = (
edit["random_direction"] if control == "direction" else edit["direction"]
)
columns = edit["random_columns"] if control == "columns" else edit["columns"]
strength = edit["alpha"] if alpha is None else alpha
projection = direction @ edit["W"][:, columns]
def keep(activation, hook, layer=layer, columns=columns):
pending[layer] = activation[..., columns].float()
def change(
output,
hook,
layer=layer,
direction=direction,
projection=projection,
strength=strength,
):
coefficient = (
output.float() @ direction if whole else pending.pop(layer) @ projection
)
return (output.float() - strength * coefficient[..., None] * direction).to(
output.dtype
)
if not whole:
hooks.append((f"blocks.{layer}.mlp.hook_post", keep))
hooks.append((f"blocks.{layer}.mlp.hook_out", change))
return hooks
conditions = {
"base": [],
"fitted_alpha": sweep_hooks(None),
"alpha_1": sweep_hooks(1.0),
"whole_alpha_1": sweep_hooks(1.0, whole=True),
"random_direction": sweep_hooks(None, control="direction"),
"random_columns": sweep_hooks(None, control="columns"),
}
for alpha in COEFFICIENTS:
conditions.setdefault(f"alpha_{alpha:g}", sweep_hooks(alpha))
def synchronize():
if SWEEP_DEVICE == "cuda":
torch.cuda.synchronize()
@torch.no_grad()
def measure_sweep():
results, condition_seconds = {}, {}
for name, hooks in conditions.items():
synchronize()
condition_started = time.perf_counter()
results[name] = {}
for group, sequences in windows.items():
measured = []
for sample_id, sequence, mask in zip(selected_ids, sequences, masks[group]):
tokens = torch.tensor([sequence], device=SWEEP_DEVICE)
with sweep_bridge.hooks(fwd_hooks=hooks):
logits = sweep_bridge(tokens)[0, :-1]
targets = tokens[0, 1:]
nll = (
torch.cat(
[
F.cross_entropy(
logits[start : start + 256].float(),
targets[start : start + 256],
reduction="none",
)
for start in range(0, len(targets), 256)
]
)
.cpu()
.numpy()
)
assert np.isfinite(nll).all()
measured.append(
{
"id": sample_id,
"all": float(nll.mean()),
"copy": float(nll[mask].mean()) if mask.any() else None,
"noncopy": float(nll[~mask].mean()) if (~mask).any() else None,
"copy_tokens": int(mask.sum()),
"tokens": len(targets),
}
)
results[name][group] = measured
synchronize()
condition_seconds[name] = time.perf_counter() - condition_started
return results, condition_seconds
synchronize()
forward_started = time.perf_counter()
results, condition_seconds = measure_sweep()
timings["sweep_forward_seconds"] = time.perf_counter() - forward_started
def summarize_results(results):
def mean_available(values, metric):
available = [row[metric] for row in values if row[metric] is not None]
return float(np.mean(available)) if available else None
return {
name: {
group: {
metric: mean_available(values, metric)
for metric in ("all", "copy", "noncopy")
}
for group, values in groups.items()
}
for name, groups in results.items()
}
summary = summarize_results(results)
assert (
summarize_results({"base": {"late": [{"all": 1.0, "copy": None, "noncopy": 1.0}]}})[
"base"
]["late"]["copy"]
is None
)
def delta_nll(name, group, metric):
value, baseline = summary[name][group][metric], summary["base"][group][metric]
return value - baseline if value is not None and baseline is not None else np.nan
interventions = {
"fitted_alpha": "Fitted directional hook",
"alpha_1": "α=1, selected columns",
"whole_alpha_1": "α=1, whole MLP output",
"random_direction": "Random direction",
"random_columns": "Random columns",
}
display(
pd.DataFrame(
[
{
"intervention": label,
"copy ΔNLL, late": delta_nll(name, "late", "copy"),
"other ΔNLL, late": delta_nll(name, "late", "noncopy"),
"all ΔNLL, early": delta_nll(name, "early", "all"),
}
for name, label in interventions.items()
]
)
)
display(
pd.DataFrame(
[
{
"alpha": alpha,
"copy ΔNLL, late": delta_nll(f"alpha_{alpha:g}", "late", "copy"),
}
for alpha in COEFFICIENTS
]
)
)
timings["sweep_total_seconds"] = time.perf_counter() - sweep_started
sweep_record = {
"model_revision": MODEL_REVISION,
"logs_revision": LOGS_REVISION,
"python": platform.python_version(),
"torch": torch.__version__,
"transformers": transformers.__version__,
"device": SWEEP_DEVICE,
"dtype": str(SWEEP_DTYPE),
"window_tokens": WINDOW_TOKENS,
"samples": SWEEP_SAMPLES,
"forward_passes": len(conditions) * len(windows) * SWEEP_SAMPLES,
"transcript_parts": TRANSCRIPT_PARTS,
"seed": SEED,
"eligible": len(eligible),
"selected_ids": selected_ids,
"source_files": source_files,
"condition_seconds": condition_seconds,
"summary": summary,
"per_window": results,
"sweep_seconds": timings["sweep_forward_seconds"],
"total_seconds": timings["sweep_total_seconds"],
}
if SWEEP_DEVICE == "cuda":
sweep_record["peak_cuda_GiB"] = torch.cuda.max_memory_allocated() / 2**30
print(
{
key: sweep_record[key]
for key in (
"device",
"dtype",
"samples",
"window_tokens",
"sweep_seconds",
"total_seconds",
)
}
)
| intervention | copy ΔNLL, late | other ΔNLL, late | all ΔNLL, early | |
|---|---|---|---|---|
| 0 | Fitted directional hook | 0.077371 | -0.004118 | -0.001147 |
| 1 | α=1, selected columns | 0.011263 | -0.001068 | -0.001247 |
| 2 | α=1, whole MLP output | 0.020979 | 0.007922 | -0.000449 |
| 3 | Random direction | 0.000151 | 0.001033 | 0.000654 |
| 4 | Random columns | -0.001511 | 0.003145 | 0.001026 |
| alpha | copy ΔNLL, late | |
|---|---|---|
| 0 | 0.5 | 0.005530 |
| 1 | 1.0 | 0.011263 |
| 2 | 2.0 | 0.023279 |
| 3 | 3.0 | 0.035934 |
| 4 | 4.0 | 0.049153 |
| 5 | 6.0 | 0.077366 |
| 6 | 8.0 | 0.108848 |
| 7 | 10.0 | 0.144491 |
{'device': 'cpu', 'dtype': 'torch.float32', 'samples': 2, 'window_tokens': 64, 'sweep_seconds': 44.44451929999923, 'total_seconds': 45.229604799998924}
Recorded larger-window experiment¶
A timed rerun completed on 2026-11-06, using the cells above with the full-size settings: an NVIDIA GeForce RTX 5060 Laptop GPU, bf16, ten transcript shards, 40 samples and 1,024-token early/late windows. PyTorch was 2.11.0+cu128, Transformers 5.13.0 and TransformerLens a70b1b65. Values are equal-weight means over windows, in nats/token.
The 13-condition run took 445.4 seconds for model loading, sampling and measurement, including 418.8 seconds for the sweep; peak allocated CUDA memory was 2.63 GiB. The six table conditions and seven additional coefficients used 1,040 forward passes. Selected problem IDs, source-file hashes, versions, condition timings, executed-cell source hash and summary values are included in this notebook’s weight_edit_audit_run metadata. The source hash covers the cells named in
source_hash_cells, before applying the larger-budget settings.
Intervention |
Copy-token ΔNLL, late windows |
Other-token ΔNLL, late windows |
All-token ΔNLL, early windows |
|---|---|---|---|
Fitted directional hook, α≈6, selected 64 columns |
+0.0600 |
+0.0067 |
+0.0091 |
α=1, same columns |
+0.0012 |
−0.0006 |
+0.0003 |
α=1, whole MLP output |
+0.0017 |
+0.0052 |
+0.0014 |
Random direction, same α and columns |
+0.0000 |
+0.0003 |
−0.0000 |
Random columns, same α and direction |
−0.0001 |
−0.0012 |
+0.0005 |
The same sample set gives this coefficient response for repeated tokens in late windows:
α |
0.5 |
1 |
2 |
3 |
4 |
6 |
8 |
10 |
|---|---|---|---|---|---|---|---|---|
ΔNLL |
+0.0005 |
+0.0012 |
+0.0030 |
+0.0063 |
+0.0126 |
+0.0599 |
+0.2480 |
+0.5900 |
These measurements use isolated teacher-forced windows, whereas the authors’ benchmark evaluates generated answers from full problem prompts.
6. Reusing the audit¶
For another edit, retain the original-value checks, compute the exact difference first, and compare it with a structured approximation. A coefficient sweep then answers a more precise question: how does that measured update affect the selected inputs? Keep sampling IDs, per-window results, device, versions and timings with the experiment.
[9]:
# NBVAL_IGNORE_OUTPUT
timings["total_seconds"] = time.perf_counter() - notebook_started
print(timings)
process = psutil.Process()
memory = process.memory_info()
peak_bytes = getattr(memory, "peak_wset", getattr(memory, "rss", 0))
print(
{
"process_memory_GiB": peak_bytes / 2**30,
"memory_measure": "Windows peak working set"
if hasattr(memory, "peak_wset")
else "current RSS",
}
)
{'artifacts_and_checkpoint_seconds': 0.10887820000061765, 'rank_one_fit_seconds': 0.027607399999396876, 'model_load_seconds': 2.3563274999978603, 'forward_validation_seconds': 6.945805199997267, 'sweep_forward_seconds': 44.44451929999923, 'sweep_total_seconds': 45.229604799998924, 'total_seconds': 60.873208799996064}
{'process_memory_GiB': 5.430072784423828, 'memory_measure': 'Windows peak working set'}