File size: 6,242 Bytes
4ca4e4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | """Execute the authors' construction notebooks without modifying their bytes.
The arXiv source links to https://github.com/SprocketLab/hybrid-expressivity.
``source/official-code`` is a detached checkout of the last repository commit
before arXiv:2603.08859v1. This runner reads the code cells directly from the
two authored notebooks, executes their setup/weight-assignment cells unchanged
on CPU, and evaluates the constructed models on deterministic batches generated
by the authors' own Dataset implementation. Notebook display/plot cells are
intentionally not executed because they are not part of the forward pass.
This is construction validation only. It is not an attempted rerun of the
GPU-only learned-model training entry point, which hard-codes ``cuda``.
"""
from __future__ import annotations
import copy
import hashlib
import json
import os
import sys
from pathlib import Path
from typing import Any
import numpy as np
import torch
ROOT = Path(__file__).resolve().parent
CODE = ROOT / "source" / "official-code"
NOTEBOOKS = {
"claim3_native": (CODE / "constructions" / "construction_var_copy.ipynb", 18),
"claim4_native": (CODE / "constructions" / "construction_decode_recall.ipynb", 27),
}
OUT = ROOT / "outputs"
def sha256(path: Path) -> str:
return hashlib.sha256(path.read_bytes()).hexdigest()
def execute_notebook_prefix(path: Path, last_cell: int) -> dict[str, Any]:
"""Execute exact code-cell text through the construction's model cell."""
notebook = json.loads(path.read_text())
cells = notebook["cells"]
assert cells[0]["cell_type"] == "code"
# The two notebook magics only enable autoreload; skipping that display
# cell preserves every model/data/weight-assignment cell verbatim.
selected = [i for i in range(1, last_cell + 1) if cells[i]["cell_type"] == "code"]
namespace: dict[str, Any] = {"__name__": "__main__", "__file__": str(path)}
old_cwd = Path.cwd()
old_path = list(sys.path)
try:
os.chdir(path.parent)
sys.path.insert(0, str(path.parent))
np.random.seed(1729)
torch.manual_seed(1729)
for i in selected:
source = "".join(cells[i]["source"])
exec(compile(source, f"{path}::cell-{i}", "exec"), namespace)
finally:
os.chdir(old_cwd)
sys.path[:] = old_path
namespace["_executed_cells"] = selected
return namespace
def evaluate_last_position(ns: dict[str, Any], n_batches: int = 256) -> dict[str, Any]:
"""Use the already-authored Dataset/model objects and score final tokens."""
model = ns["model"]
dataset = ns["train_dataset"]
model.eval()
eligible = correct = 0
with torch.no_grad():
for i in range(n_batches):
batch = dataset[i]
logits = model(batch["input_ids"])
pred = torch.argmax(logits[:, -1], dim=-1)
valid = batch["mask"][:, -1].bool()
eligible += int(valid.sum())
correct += int((pred[valid] == batch["output_ids"][:, -1][valid]).sum())
return {
"n_batches": n_batches,
"examples_per_batch": int(ns["args"].train_batch_size),
"eligible_last_position_examples": eligible,
"correct_last_position_examples": correct,
"last_position_accuracy": correct / eligible if eligible else None,
}
def destructive_no_first_ssm(ns: dict[str, Any], n_batches: int = 256) -> dict[str, Any]:
"""Ablate the first authored SSM gate in-memory; source bytes stay intact."""
model = copy.deepcopy(ns["model"])
model.layers[0].layer.Delta.data.zero_()
dataset = ns["train_dataset"]
model.eval()
eligible = correct = 0
# Re-seed data generation so the control sees exactly the baseline inputs.
np.random.seed(1729 + 1)
with torch.no_grad():
for i in range(n_batches):
batch = dataset[i]
logits = model(batch["input_ids"])
pred = torch.argmax(logits[:, -1], dim=-1)
valid = batch["mask"][:, -1].bool()
eligible += int(valid.sum())
correct += int((pred[valid] == batch["output_ids"][:, -1][valid]).sum())
return {
"ablation": "first SimpleSSM Delta vector set to zero in memory",
"eligible_last_position_examples": eligible,
"correct_last_position_examples": correct,
"last_position_accuracy": correct / eligible if eligible else None,
}
def main() -> None:
OUT.mkdir(exist_ok=True)
git_head = (CODE / ".git").exists()
for output_name, (notebook, last_cell) in NOTEBOOKS.items():
# Baseline and destructive control receive identical generated batches.
ns = execute_notebook_prefix(notebook, last_cell)
np.random.seed(1729 + 1)
baseline = evaluate_last_position(ns)
control = destructive_no_first_ssm(ns)
payload = {
"kind": "direct_native_notebook_execution",
"official_repository": "https://github.com/SprocketLab/hybrid-expressivity",
"repository_checkout": "6be8f8fbc2169290af6f4ba5e4bd53a5c6485f7b",
"repository_is_git_checkout": git_head,
"notebook": str(notebook.relative_to(ROOT)),
"notebook_sha256": sha256(notebook),
"executed_code_cells": ns["_executed_cells"],
"device": "cpu",
"torch_version": torch.__version__,
"baseline": baseline,
"destructive_control": control,
"control_is_lower": (
control["last_position_accuracy"] < baseline["last_position_accuracy"]
if baseline["last_position_accuracy"] is not None
else None
),
"scope": (
"Construction cells plus deterministic last-position batches; "
"not a learned-model training rerun."
),
}
(OUT / f"{output_name}.json").write_text(json.dumps(payload, indent=2) + "\n")
print(
f"{output_name}: baseline={baseline['last_position_accuracy']:.6f} "
f"control={control['last_position_accuracy']:.6f} n={baseline['eligible_last_position_examples']}",
flush=True,
)
if __name__ == "__main__":
main()
|