research: lock AttnRes spike analyzer
This commit is contained in:
wuyang
@@ -0,0 +1,751 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Aggregate and gate the preregistered AttnRes spike-path study."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import statistics
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
|
||||
PROTOCOL_ID = "llm-atlas-k3-attnres-spike-path-v1"
|
||||
SEEDS = (2026073001, 2026073002, 2026073003)
|
||||
POSITIONS = (
|
||||
"pre_attention_input",
|
||||
"attention_branch_output",
|
||||
"post_attention_state",
|
||||
"pre_mlp_input",
|
||||
"mlp_branch_output",
|
||||
"post_mlp_state",
|
||||
)
|
||||
MAIN_REDUCTIONS = (
|
||||
"element_rms",
|
||||
"token_rms_mean",
|
||||
"token_rms_median",
|
||||
"token_rms_p95",
|
||||
)
|
||||
ALL_REDUCTIONS = MAIN_REDUCTIONS + (
|
||||
"batch_mean_rms",
|
||||
"token_mean_rms",
|
||||
"global_l2",
|
||||
)
|
||||
MODES = (
|
||||
"learned",
|
||||
"detached_learned",
|
||||
"uniform_value_backward",
|
||||
)
|
||||
SPIKE_LAYERS = (21, 22, 23, 24, 25)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--raw-dir", type=Path, required=True)
|
||||
parser.add_argument("--manifest", type=Path, required=True)
|
||||
parser.add_argument("--aggregate-output", type=Path, required=True)
|
||||
parser.add_argument("--compact-output", type=Path, required=True)
|
||||
parser.add_argument("--reproduction-output", type=Path, required=True)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def file_sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for block in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def canonical_sha256(value: Any) -> str:
|
||||
return hashlib.sha256(
|
||||
json.dumps(
|
||||
value, ensure_ascii=False, sort_keys=True, separators=(",", ":")
|
||||
).encode()
|
||||
).hexdigest()
|
||||
|
||||
|
||||
def mean(values: Iterable[float]) -> float:
|
||||
return statistics.fmean(values)
|
||||
|
||||
|
||||
def average_ranks(values: list[float]) -> list[float]:
|
||||
ordered = sorted(range(len(values)), key=lambda index: (values[index], index))
|
||||
ranks = [0.0] * len(values)
|
||||
cursor = 0
|
||||
while cursor < len(ordered):
|
||||
end = cursor + 1
|
||||
while end < len(ordered) and values[ordered[end]] == values[ordered[cursor]]:
|
||||
end += 1
|
||||
average_rank = (cursor + 1 + end) / 2
|
||||
for offset in range(cursor, end):
|
||||
ranks[ordered[offset]] = average_rank
|
||||
cursor = end
|
||||
return ranks
|
||||
|
||||
|
||||
def pearson(left: list[float], right: list[float]) -> float:
|
||||
left_mean = mean(left)
|
||||
right_mean = mean(right)
|
||||
numerator = sum(
|
||||
(x - left_mean) * (y - right_mean)
|
||||
for x, y in zip(left, right)
|
||||
)
|
||||
left_square = sum((value - left_mean) ** 2 for value in left)
|
||||
right_square = sum((value - right_mean) ** 2 for value in right)
|
||||
if left_square == 0 or right_square == 0:
|
||||
raise RuntimeError("correlation is undefined for a constant vector")
|
||||
return numerator / math.sqrt(left_square * right_square)
|
||||
|
||||
|
||||
def spearman(left: list[float], right: list[float]) -> float:
|
||||
return pearson(average_ranks(left), average_ranks(right))
|
||||
|
||||
|
||||
def load_run(path: Path, *, expected_kind: str, expected_seed: int) -> dict[str, Any]:
|
||||
value = json.loads(path.read_text())
|
||||
if value["protocol_id"] != PROTOCOL_ID:
|
||||
raise RuntimeError(f"protocol mismatch: {path}")
|
||||
if value["run_kind"] != expected_kind or value["seed"] != expected_seed:
|
||||
raise RuntimeError(f"run identity mismatch: {path}")
|
||||
canonical = value.pop("canonical_sha256_without_self")
|
||||
if canonical_sha256(value) != canonical:
|
||||
raise RuntimeError(f"canonical hash mismatch: {path}")
|
||||
value["canonical_sha256_without_self"] = canonical
|
||||
if expected_kind != "smoke":
|
||||
if not value["round05_equivalence"]["passed"]:
|
||||
raise RuntimeError(f"Round 05 equivalence failed: {path}")
|
||||
if value["steps"] != 8000:
|
||||
raise RuntimeError(f"formal step count mismatch: {path}")
|
||||
return value
|
||||
|
||||
|
||||
def final_diagnostic(run: dict[str, Any]) -> dict[str, Any]:
|
||||
matches = [item for item in run["diagnostics"] if item["step"] == 8000]
|
||||
if len(matches) != 1:
|
||||
raise RuntimeError("final diagnostic missing or duplicated")
|
||||
return matches[0]
|
||||
|
||||
|
||||
def metric(
|
||||
run: dict[str, Any],
|
||||
*,
|
||||
position: str,
|
||||
reduction: str,
|
||||
mode: str = "learned",
|
||||
) -> dict[str, Any]:
|
||||
return final_diagnostic(run)["modes"][mode]["positions"][position][
|
||||
"reductions"
|
||||
][reduction]
|
||||
|
||||
|
||||
def compare_replay(formal: dict[str, Any], replay: dict[str, Any]) -> dict[str, Any]:
|
||||
fields = (
|
||||
"architecture",
|
||||
"depth",
|
||||
"seed",
|
||||
"steps",
|
||||
"batch_size",
|
||||
"target_bytes_seen",
|
||||
"manifest",
|
||||
"model",
|
||||
"optimizer",
|
||||
"hashes",
|
||||
"evaluations",
|
||||
"diagnostics",
|
||||
"training_history",
|
||||
"environment",
|
||||
"artifacts",
|
||||
"round05_equivalence",
|
||||
)
|
||||
checks = {field: formal[field] == replay[field] for field in fields}
|
||||
passed = all(checks.values())
|
||||
if not passed:
|
||||
raise RuntimeError(f"Round 06 replay mismatch: {checks}")
|
||||
compare_payload = {field: formal[field] for field in fields}
|
||||
return {
|
||||
"passed": True,
|
||||
"checks": checks,
|
||||
"frozen_compare_sha256": canonical_sha256(compare_payload),
|
||||
}
|
||||
|
||||
|
||||
def reduction_robustness(runs: list[dict[str, Any]], manifest: dict[str, Any]) -> dict[str, Any]:
|
||||
threshold = manifest["thresholds"]
|
||||
cells = []
|
||||
for run in runs:
|
||||
reference = metric(
|
||||
run,
|
||||
position="post_mlp_state",
|
||||
reduction="element_rms",
|
||||
)
|
||||
for reduction in MAIN_REDUCTIONS:
|
||||
candidate = metric(
|
||||
run,
|
||||
position="post_mlp_state",
|
||||
reduction=reduction,
|
||||
)
|
||||
stats = candidate["statistics"]
|
||||
rho = spearman(reference["values"], candidate["values"])
|
||||
checks = {
|
||||
"spike_contrast": (
|
||||
stats["spike_contrast"]
|
||||
>= threshold["spike_contrast"]
|
||||
),
|
||||
"top_five_overlap": (
|
||||
stats["top_five_spike_overlap"]
|
||||
>= threshold["top_five_min_overlap"]
|
||||
),
|
||||
"spearman": rho >= threshold["spearman_minimum"],
|
||||
}
|
||||
cells.append(
|
||||
{
|
||||
"seed": run["seed"],
|
||||
"reduction": reduction,
|
||||
"spike_contrast": stats["spike_contrast"],
|
||||
"top_five_layers": stats["top_five_layers"],
|
||||
"top_five_spike_overlap": stats[
|
||||
"top_five_spike_overlap"
|
||||
],
|
||||
"spearman_vs_element_rms": rho,
|
||||
"checks": checks,
|
||||
"passed": all(checks.values()),
|
||||
}
|
||||
)
|
||||
passed_count = sum(item["passed"] for item in cells)
|
||||
if passed_count == len(cells):
|
||||
verdict = "robust within the preregistered reduction family"
|
||||
elif passed_count == 0:
|
||||
verdict = "not robust at this threshold"
|
||||
else:
|
||||
verdict = "mixed"
|
||||
return {
|
||||
"verdict": verdict,
|
||||
"passed_cells": passed_count,
|
||||
"total_cells": len(cells),
|
||||
"cells": cells,
|
||||
}
|
||||
|
||||
|
||||
def visible_positions(runs: list[dict[str, Any]], manifest: dict[str, Any]) -> dict[str, Any]:
|
||||
threshold = manifest["thresholds"]["spike_contrast"]
|
||||
positions = []
|
||||
for position in POSITIONS:
|
||||
per_seed = []
|
||||
for run in runs:
|
||||
stats = metric(
|
||||
run, position=position, reduction="element_rms"
|
||||
)["statistics"]
|
||||
per_seed.append(
|
||||
{
|
||||
"seed": run["seed"],
|
||||
"spike_contrast": stats["spike_contrast"],
|
||||
"peak_layer": stats["peak_layer"],
|
||||
"peak_normalized": stats["peak_normalized"],
|
||||
"passed": stats["spike_contrast"] >= threshold,
|
||||
}
|
||||
)
|
||||
positions.append(
|
||||
{
|
||||
"position": position,
|
||||
"visible_3_of_3": all(item["passed"] for item in per_seed),
|
||||
"per_seed": per_seed,
|
||||
"mean_spike_contrast": mean(
|
||||
item["spike_contrast"] for item in per_seed
|
||||
),
|
||||
}
|
||||
)
|
||||
visible = [item["position"] for item in positions if item["visible_3_of_3"]]
|
||||
return {
|
||||
"criterion": f"step 8000 element-RMS spike contrast >= {threshold} in 3/3 seeds",
|
||||
"visible_positions": visible,
|
||||
"earliest_observed_tensor": visible[0] if visible else None,
|
||||
"verdict": "visible at one or more positions" if visible else "position-mixed",
|
||||
"positions": positions,
|
||||
}
|
||||
|
||||
|
||||
def intervention_effect(
|
||||
runs: list[dict[str, Any]],
|
||||
manifest: dict[str, Any],
|
||||
*,
|
||||
source_mode: str,
|
||||
target_mode: str,
|
||||
label: str,
|
||||
) -> dict[str, Any]:
|
||||
threshold = manifest["thresholds"]["material_relative_drop"]
|
||||
epsilon = manifest["thresholds"]["positive_denominator_epsilon"]
|
||||
per_seed = []
|
||||
for run in runs:
|
||||
source = metric(
|
||||
run,
|
||||
position="post_mlp_state",
|
||||
reduction="element_rms",
|
||||
mode=source_mode,
|
||||
)["statistics"]
|
||||
target = metric(
|
||||
run,
|
||||
position="post_mlp_state",
|
||||
reduction="element_rms",
|
||||
mode=target_mode,
|
||||
)["statistics"]
|
||||
if (
|
||||
not math.isfinite(source["spike_contrast"])
|
||||
or not math.isfinite(source["peak_normalized"])
|
||||
or source["spike_contrast"] <= epsilon
|
||||
or source["peak_normalized"] <= epsilon
|
||||
):
|
||||
raise RuntimeError("invalid intervention denominator")
|
||||
contrast_drop = (
|
||||
source["spike_contrast"] - target["spike_contrast"]
|
||||
) / source["spike_contrast"]
|
||||
peak_drop = (
|
||||
source["peak_normalized"] - target["peak_normalized"]
|
||||
) / source["peak_normalized"]
|
||||
per_seed.append(
|
||||
{
|
||||
"seed": run["seed"],
|
||||
"source_spike_contrast": source["spike_contrast"],
|
||||
"target_spike_contrast": target["spike_contrast"],
|
||||
"relative_drop_contrast": contrast_drop,
|
||||
"source_peak_normalized": source["peak_normalized"],
|
||||
"target_peak_normalized": target["peak_normalized"],
|
||||
"relative_drop_peak": peak_drop,
|
||||
"passed": (
|
||||
contrast_drop >= threshold and peak_drop >= threshold
|
||||
),
|
||||
}
|
||||
)
|
||||
passed = all(item["passed"] for item in per_seed)
|
||||
same_direction = all(
|
||||
item["relative_drop_contrast"] > 0
|
||||
and item["relative_drop_peak"] > 0
|
||||
for item in per_seed
|
||||
)
|
||||
opposite_direction = all(
|
||||
item["relative_drop_contrast"] < 0
|
||||
and item["relative_drop_peak"] < 0
|
||||
for item in per_seed
|
||||
)
|
||||
if passed:
|
||||
verdict = "material sensitivity at the preregistered threshold"
|
||||
elif same_direction:
|
||||
verdict = "same-direction but below the joint threshold"
|
||||
elif opposite_direction:
|
||||
verdict = "opposite direction in 3/3 seeds; no material reduction"
|
||||
else:
|
||||
verdict = "mixed"
|
||||
return {
|
||||
"label": label,
|
||||
"source_mode": source_mode,
|
||||
"target_mode": target_mode,
|
||||
"threshold": threshold,
|
||||
"passed_3_of_3": passed,
|
||||
"same_direction_3_of_3": same_direction,
|
||||
"opposite_direction_3_of_3": opposite_direction,
|
||||
"verdict": verdict,
|
||||
"per_seed": per_seed,
|
||||
"mean_relative_drop_contrast": mean(
|
||||
item["relative_drop_contrast"] for item in per_seed
|
||||
),
|
||||
"mean_relative_drop_peak": mean(
|
||||
item["relative_drop_peak"] for item in per_seed
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def intervention_summary(runs: list[dict[str, Any]], manifest: dict[str, Any]) -> dict[str, Any]:
|
||||
forward_gates = []
|
||||
for run in runs:
|
||||
diagnostic = final_diagnostic(run)
|
||||
forward_gates.append(
|
||||
{
|
||||
"seed": run["seed"],
|
||||
"passed": diagnostic["forward_identity_gate"]["passed"],
|
||||
}
|
||||
)
|
||||
if not all(item["passed"] for item in forward_gates):
|
||||
raise RuntimeError("a final forward identity gate failed")
|
||||
return {
|
||||
"scope": manifest["interventions"]["scope"],
|
||||
"forward_identity": forward_gates,
|
||||
"softmax_key_path": intervention_effect(
|
||||
runs,
|
||||
manifest,
|
||||
source_mode="learned",
|
||||
target_mode="detached_learned",
|
||||
label="global removal of all mixer softmax/query/key source-gradient paths",
|
||||
),
|
||||
"value_coefficients": intervention_effect(
|
||||
runs,
|
||||
manifest,
|
||||
source_mode="detached_learned",
|
||||
target_mode="uniform_value_backward",
|
||||
label="global replacement of learned value-backward coefficients with 1/N",
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def trajectory(runs: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
result = []
|
||||
for run in runs:
|
||||
evaluations = {
|
||||
item["step"]: item["bits_per_byte"]
|
||||
for item in run["evaluations"]
|
||||
}
|
||||
points = []
|
||||
for diagnostic in run["diagnostics"]:
|
||||
stats = diagnostic["modes"]["learned"]["positions"][
|
||||
"post_mlp_state"
|
||||
]["reductions"]["element_rms"]["statistics"]
|
||||
points.append(
|
||||
{
|
||||
"step": diagnostic["step"],
|
||||
"bits_per_byte": evaluations[diagnostic["step"]],
|
||||
"population_cv": stats["population_cv"],
|
||||
"spike_contrast": stats["spike_contrast"],
|
||||
"peak_layer": stats["peak_layer"],
|
||||
"peak_normalized": stats["peak_normalized"],
|
||||
}
|
||||
)
|
||||
result.append({"seed": run["seed"], "points": points})
|
||||
return result
|
||||
|
||||
|
||||
def mixer_associations(runs: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
rows = []
|
||||
for run in runs:
|
||||
diagnostic = final_diagnostic(run)
|
||||
gradients = metric(
|
||||
run,
|
||||
position="post_mlp_state",
|
||||
reduction="element_rms",
|
||||
)["statistics"]["normalized"]
|
||||
mixers = diagnostic["modes"]["learned"]["mixers"]
|
||||
for layer in range(1, 33):
|
||||
attention = mixers[2 * (layer - 1)]
|
||||
mlp = mixers[2 * (layer - 1) + 1]
|
||||
rows.append(
|
||||
{
|
||||
"seed": run["seed"],
|
||||
"layer": layer,
|
||||
"group": (layer - 1) // 4 + 1,
|
||||
"offset": (layer - 1) % 4 + 1,
|
||||
"normalized_gradient": gradients[layer - 1],
|
||||
"attention_latest": attention["latest_source_mass"],
|
||||
"attention_entropy": attention["normalized_entropy"],
|
||||
"mlp_latest": mlp["latest_source_mass"],
|
||||
"mlp_entropy": mlp["normalized_entropy"],
|
||||
"attention_max": attention["max_source_mass"],
|
||||
"mlp_max": mlp["max_source_mass"],
|
||||
}
|
||||
)
|
||||
|
||||
def correlations(selected: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
gradient = [row["normalized_gradient"] for row in selected]
|
||||
result = {"points": len(selected)}
|
||||
for key in (
|
||||
"attention_latest",
|
||||
"attention_entropy",
|
||||
"mlp_latest",
|
||||
"mlp_entropy",
|
||||
"attention_max",
|
||||
"mlp_max",
|
||||
):
|
||||
values = [row[key] for row in selected]
|
||||
result[key] = {
|
||||
"pearson": pearson(gradient, values),
|
||||
"spearman": spearman(gradient, values),
|
||||
}
|
||||
return result
|
||||
|
||||
target_layers = []
|
||||
for layer in range(19, 29):
|
||||
selected = [row for row in rows if row["layer"] == layer]
|
||||
target_layers.append(
|
||||
{
|
||||
"layer": layer,
|
||||
"group": selected[0]["group"],
|
||||
"offset": selected[0]["offset"],
|
||||
**{
|
||||
key: mean(row[key] for row in selected)
|
||||
for key in (
|
||||
"normalized_gradient",
|
||||
"attention_latest",
|
||||
"attention_entropy",
|
||||
"mlp_latest",
|
||||
"mlp_entropy",
|
||||
"attention_max",
|
||||
"mlp_max",
|
||||
)
|
||||
},
|
||||
"per_seed_normalized_gradient": [
|
||||
row["normalized_gradient"] for row in selected
|
||||
],
|
||||
}
|
||||
)
|
||||
return {
|
||||
"identity": "observational association; not causal attribution",
|
||||
"all_layers": correlations(rows),
|
||||
"layers_19_28": correlations(
|
||||
[row for row in rows if 19 <= row["layer"] <= 28]
|
||||
),
|
||||
"target_layers": target_layers,
|
||||
}
|
||||
|
||||
|
||||
def final_arrays(runs: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
values = []
|
||||
for run in runs:
|
||||
positions = {}
|
||||
for position in POSITIONS:
|
||||
positions[position] = {}
|
||||
for reduction in ALL_REDUCTIONS:
|
||||
item = metric(
|
||||
run, position=position, reduction=reduction
|
||||
)
|
||||
positions[position][reduction] = {
|
||||
"values": item["values"],
|
||||
"statistics": item["statistics"],
|
||||
}
|
||||
interventions = {}
|
||||
for mode in MODES:
|
||||
item = metric(
|
||||
run,
|
||||
position="post_mlp_state",
|
||||
reduction="element_rms",
|
||||
mode=mode,
|
||||
)
|
||||
interventions[mode] = {
|
||||
"values": item["values"],
|
||||
"statistics": item["statistics"],
|
||||
}
|
||||
values.append(
|
||||
{
|
||||
"seed": run["seed"],
|
||||
"positions": positions,
|
||||
"interventions": interventions,
|
||||
"final_mixers": final_diagnostic(run)["modes"]["learned"][
|
||||
"mixers"
|
||||
],
|
||||
}
|
||||
)
|
||||
return values
|
||||
|
||||
|
||||
def compact_final_arrays(runs: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
values = []
|
||||
for run in runs:
|
||||
position_element_rms = {
|
||||
position: metric(
|
||||
run,
|
||||
position=position,
|
||||
reduction="element_rms",
|
||||
)
|
||||
for position in POSITIONS
|
||||
}
|
||||
post_mlp_reductions = {
|
||||
reduction: metric(
|
||||
run,
|
||||
position="post_mlp_state",
|
||||
reduction=reduction,
|
||||
)
|
||||
for reduction in ALL_REDUCTIONS
|
||||
}
|
||||
interventions = {
|
||||
mode: metric(
|
||||
run,
|
||||
position="post_mlp_state",
|
||||
reduction="element_rms",
|
||||
mode=mode,
|
||||
)
|
||||
for mode in MODES
|
||||
}
|
||||
values.append(
|
||||
{
|
||||
"seed": run["seed"],
|
||||
"position_element_rms": position_element_rms,
|
||||
"post_mlp_reductions": post_mlp_reductions,
|
||||
"interventions": interventions,
|
||||
}
|
||||
)
|
||||
return values
|
||||
|
||||
|
||||
def write_json(path: Path, value: dict[str, Any]) -> None:
|
||||
value["canonical_sha256_without_self"] = canonical_sha256(value)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = path.with_suffix(path.suffix + ".tmp")
|
||||
temporary.write_text(
|
||||
json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
|
||||
)
|
||||
temporary.replace(path)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
manifest = json.loads(args.manifest.read_text())
|
||||
if manifest["protocol_id"] != PROTOCOL_ID:
|
||||
raise RuntimeError("manifest protocol mismatch")
|
||||
|
||||
formal_paths = {
|
||||
seed: args.raw_dir / f"formal-seed-{seed}.json"
|
||||
for seed in SEEDS
|
||||
}
|
||||
replay_path = args.raw_dir / f"replay-seed-{SEEDS[0]}.json"
|
||||
formal = [
|
||||
load_run(formal_paths[seed], expected_kind="formal", expected_seed=seed)
|
||||
for seed in SEEDS
|
||||
]
|
||||
replay = load_run(
|
||||
replay_path, expected_kind="replay", expected_seed=SEEDS[0]
|
||||
)
|
||||
replay_gate = compare_replay(formal[0], replay)
|
||||
|
||||
reduction = reduction_robustness(formal, manifest)
|
||||
positions = visible_positions(formal, manifest)
|
||||
interventions = intervention_summary(formal, manifest)
|
||||
trajectories = trajectory(formal)
|
||||
mixers = mixer_associations(formal)
|
||||
arrays = final_arrays(formal)
|
||||
|
||||
raw_files = {
|
||||
path.name: {
|
||||
"file_sha256": file_sha256(path),
|
||||
"canonical_sha256": run["canonical_sha256_without_self"],
|
||||
}
|
||||
for path, run in [
|
||||
*[(formal_paths[seed], run) for seed, run in zip(SEEDS, formal)],
|
||||
(replay_path, replay),
|
||||
]
|
||||
}
|
||||
reproduction = {
|
||||
"schema_version": 1,
|
||||
"protocol_id": PROTOCOL_ID,
|
||||
"raw_files": raw_files,
|
||||
"replay_gate": replay_gate,
|
||||
"artifacts": {
|
||||
"manifest": file_sha256(args.manifest),
|
||||
"runner": formal[0]["artifacts"]["runner_sha256"],
|
||||
"protocol": formal[0]["artifacts"]["protocol_sha256"],
|
||||
"scoping": formal[0]["artifacts"]["scoping_sha256"],
|
||||
"analyzer": file_sha256(Path(__file__)),
|
||||
},
|
||||
}
|
||||
write_json(args.reproduction_output, reproduction)
|
||||
|
||||
aggregate = {
|
||||
"schema_version": 1,
|
||||
"protocol_id": PROTOCOL_ID,
|
||||
"study": {
|
||||
"identity": manifest["study_identity"],
|
||||
"architecture": "block",
|
||||
"depth": 32,
|
||||
"seeds": list(SEEDS),
|
||||
"steps": 8000,
|
||||
"spike_layers": list(SPIKE_LAYERS),
|
||||
},
|
||||
"gates": {
|
||||
"round05_equivalence": [
|
||||
{
|
||||
"seed": run["seed"],
|
||||
**run["round05_equivalence"],
|
||||
}
|
||||
for run in formal
|
||||
],
|
||||
"replay": replay_gate,
|
||||
"forward_identity": interventions["forward_identity"],
|
||||
"initialization_negative_control": [
|
||||
{
|
||||
"seed": run["seed"],
|
||||
"passed": run["diagnostics"][0][
|
||||
"initialization_negative_control"
|
||||
]["passed"],
|
||||
}
|
||||
for run in formal
|
||||
],
|
||||
"loss_scale": [
|
||||
{
|
||||
"seed": run["seed"],
|
||||
"passed": run["diagnostics"][0][
|
||||
"loss_scale_gate"
|
||||
]["passed"],
|
||||
}
|
||||
for run in formal
|
||||
],
|
||||
},
|
||||
"verdicts": {
|
||||
"reduction_robustness": reduction,
|
||||
"visible_positions": positions,
|
||||
"interventions": interventions,
|
||||
},
|
||||
"trajectory": trajectories,
|
||||
"mixer_associations": mixers,
|
||||
"final_arrays": arrays,
|
||||
"runs": [
|
||||
{
|
||||
"seed": run["seed"],
|
||||
"canonical_sha256": run["canonical_sha256_without_self"],
|
||||
"final_bpc": run["evaluations"][-1]["bits_per_byte"],
|
||||
"final_model_state": run["hashes"]["final_model_state"],
|
||||
"final_optimizer_state": run["hashes"][
|
||||
"final_optimizer_state"
|
||||
],
|
||||
"timing": run["timing"],
|
||||
}
|
||||
for run in formal
|
||||
],
|
||||
"reproduction_canonical_sha256": reproduction[
|
||||
"canonical_sha256_without_self"
|
||||
],
|
||||
}
|
||||
write_json(args.aggregate_output, aggregate)
|
||||
|
||||
compact = {
|
||||
"schema_version": 1,
|
||||
"protocol_id": PROTOCOL_ID,
|
||||
"study": aggregate["study"],
|
||||
"gates": aggregate["gates"],
|
||||
"verdicts": aggregate["verdicts"],
|
||||
"trajectory": trajectories,
|
||||
"mixer_associations": mixers,
|
||||
"final_arrays": compact_final_arrays(formal),
|
||||
"runs": aggregate["runs"],
|
||||
"hashes": {
|
||||
"aggregate_canonical_sha256": aggregate[
|
||||
"canonical_sha256_without_self"
|
||||
],
|
||||
"reproduction_canonical_sha256": reproduction[
|
||||
"canonical_sha256_without_self"
|
||||
],
|
||||
},
|
||||
}
|
||||
write_json(args.compact_output, compact)
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"aggregate": str(args.aggregate_output),
|
||||
"compact": str(args.compact_output),
|
||||
"reproduction": str(args.reproduction_output),
|
||||
"reduction_verdict": reduction["verdict"],
|
||||
"earliest_observed_tensor": positions[
|
||||
"earliest_observed_tensor"
|
||||
],
|
||||
"softmax_key_path": interventions["softmax_key_path"][
|
||||
"verdict"
|
||||
],
|
||||
"value_coefficients": interventions["value_coefficients"][
|
||||
"verdict"
|
||||
],
|
||||
"replay_exact": replay_gate["passed"],
|
||||
},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user