experiment: implement AttnRes local path runner
This commit is contained in:
wuyang
@@ -0,0 +1,498 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Aggregate and gate preregistered Round 07 local-path results."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import statistics
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
PROTOCOL_ID = "llm-atlas-k3-attnres-local-path-v1"
|
||||
METRICS = ("spike_contrast", "peak_normalized")
|
||||
SUFFICIENCY_MODES = (
|
||||
"uniform_group_6_only",
|
||||
"uniform_group_7_only",
|
||||
"uniform_groups_6_7_only",
|
||||
"uniform_group_6_attention_only",
|
||||
"uniform_group_6_mlp_only",
|
||||
"uniform_group_7_attention_only",
|
||||
"uniform_group_7_mlp_only",
|
||||
"uniform_output_only",
|
||||
"uniform_depth_all",
|
||||
"uniform_all",
|
||||
)
|
||||
RESTORATION_MODES = (
|
||||
"uniform_except_group_6",
|
||||
"uniform_except_group_7",
|
||||
"uniform_except_groups_6_7",
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--manifest", type=Path, required=True)
|
||||
parser.add_argument(
|
||||
"--formal", type=Path, action="append", required=True
|
||||
)
|
||||
parser.add_argument("--replay", type=Path, required=True)
|
||||
parser.add_argument("--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 chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def canonical_sha256(value: Any) -> str:
|
||||
payload = json.dumps(
|
||||
value, ensure_ascii=False, sort_keys=True, separators=(",", ":")
|
||||
).encode()
|
||||
return hashlib.sha256(payload).hexdigest()
|
||||
|
||||
|
||||
def read_result(path: Path) -> dict[str, Any]:
|
||||
value = json.loads(path.read_text())
|
||||
if value["protocol_id"] != PROTOCOL_ID:
|
||||
raise RuntimeError(f"protocol mismatch: {path}")
|
||||
expected = value["canonical_sha256_without_self"]
|
||||
without_self = {
|
||||
key: item
|
||||
for key, item in value.items()
|
||||
if key != "canonical_sha256_without_self"
|
||||
}
|
||||
if canonical_sha256(without_self) != expected:
|
||||
raise RuntimeError(f"canonical self-hash mismatch: {path}")
|
||||
return value
|
||||
|
||||
|
||||
def replay_payload(value: dict[str, Any]) -> dict[str, Any]:
|
||||
cleaned = copy.deepcopy(value)
|
||||
for key in ("run_kind", "timing", "canonical_sha256_without_self"):
|
||||
cleaned.pop(key)
|
||||
return cleaned
|
||||
|
||||
|
||||
def primary_metrics(mode: dict[str, Any]) -> dict[str, float]:
|
||||
stats = mode["positions"]["post_mlp_state"]["reductions"][
|
||||
"element_rms"
|
||||
]["statistics"]
|
||||
return {name: float(stats[name]) for name in METRICS}
|
||||
|
||||
|
||||
def all_true(values: list[bool]) -> bool:
|
||||
return len(values) > 0 and all(values)
|
||||
|
||||
|
||||
def mixed_status(
|
||||
by_seed_metric: dict[str, dict[str, dict[str, Any]]],
|
||||
threshold: float,
|
||||
) -> dict[str, Any]:
|
||||
passed = []
|
||||
signs = []
|
||||
metric_passes = {metric: [] for metric in METRICS}
|
||||
for seed_values in by_seed_metric.values():
|
||||
for metric in METRICS:
|
||||
score = seed_values[metric]["score"]
|
||||
cell_passed = score is not None and score >= threshold
|
||||
passed.append(cell_passed)
|
||||
metric_passes[metric].append(cell_passed)
|
||||
if score is not None:
|
||||
signs.append(1 if score >= 0 else -1)
|
||||
reasons = []
|
||||
if any(passed) and not all(passed):
|
||||
reasons.append("seed_or_metric_pass_split")
|
||||
if metric_passes[METRICS[0]] != metric_passes[METRICS[1]]:
|
||||
reasons.append("metric_direction_split")
|
||||
if len(set(signs)) > 1:
|
||||
reasons.append("score_sign_split")
|
||||
return {
|
||||
"passed": all_true(passed),
|
||||
"threshold": threshold,
|
||||
"required_cells": len(passed),
|
||||
"passed_cells": sum(passed),
|
||||
"mixed": bool(reasons),
|
||||
"mixed_reasons": reasons,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
manifest = json.loads(args.manifest.read_text())
|
||||
if manifest["protocol_id"] != PROTOCOL_ID:
|
||||
raise RuntimeError("manifest protocol mismatch")
|
||||
if len(args.formal) != 3:
|
||||
raise RuntimeError("exactly three formal result paths are required")
|
||||
|
||||
formal_pairs = [(path, read_result(path)) for path in args.formal]
|
||||
formal_pairs.sort(key=lambda item: item[1]["seed"])
|
||||
expected_seeds = manifest["formal_seeds"]
|
||||
if [value["seed"] for _, value in formal_pairs] != expected_seeds:
|
||||
raise RuntimeError("formal seeds do not match manifest")
|
||||
for path, value in formal_pairs:
|
||||
if value["run_kind"] != "formal" or value["steps"] != 8000:
|
||||
raise RuntimeError(f"invalid formal cell: {path}")
|
||||
if not value["round06_equivalence"]["passed"]:
|
||||
raise RuntimeError(f"parent equivalence failed: {path}")
|
||||
for diagnostic in value["diagnostics"]:
|
||||
if not diagnostic["parent_learned_round06_exact"]:
|
||||
raise RuntimeError(f"parent diagnostic mismatch: {path}")
|
||||
if diagnostic["local_matrix"] is not None:
|
||||
for gate in ("forward_identity_gate", "endpoint_exactness"):
|
||||
if not diagnostic[gate]["passed"]:
|
||||
raise RuntimeError(f"{gate} failed: {path}")
|
||||
for mode in manifest["matrix_modes"]:
|
||||
if not diagnostic["local_matrix"][mode]["selector"][
|
||||
"passed"
|
||||
]:
|
||||
raise RuntimeError(f"selector failed: {path}:{mode}")
|
||||
step0 = value["diagnostics"][0]
|
||||
if (
|
||||
not step0["initialization_negative_control"]["passed"]
|
||||
or not step0["loss_scale_gate"]["passed"]
|
||||
):
|
||||
raise RuntimeError(f"step-0 control failed: {path}")
|
||||
|
||||
replay = read_result(args.replay)
|
||||
if (
|
||||
replay["run_kind"] != "replay"
|
||||
or replay["seed"] != expected_seeds[0]
|
||||
or replay["steps"] != 8000
|
||||
):
|
||||
raise RuntimeError("invalid replay cell")
|
||||
replay_exact = (
|
||||
replay_payload(formal_pairs[0][1]) == replay_payload(replay)
|
||||
)
|
||||
if not replay_exact:
|
||||
raise RuntimeError("formal seed1 and replay are not canonical exact")
|
||||
|
||||
thresholds = manifest["thresholds"]
|
||||
cells = []
|
||||
sufficiency_by_mode: dict[str, dict[str, dict[str, Any]]] = {
|
||||
mode: {} for mode in SUFFICIENCY_MODES
|
||||
}
|
||||
restoration_by_mode: dict[str, dict[str, dict[str, Any]]] = {
|
||||
mode: {} for mode in RESTORATION_MODES
|
||||
}
|
||||
|
||||
for path, value in formal_pairs:
|
||||
seed_key = str(value["seed"])
|
||||
final = value["diagnostics"][-1]["local_matrix"]
|
||||
mode_metrics = {
|
||||
mode: primary_metrics(final[mode])
|
||||
for mode in manifest["matrix_modes"]
|
||||
}
|
||||
global_metrics = {}
|
||||
for metric in METRICS:
|
||||
reference = mode_metrics["detached_learned"][metric]
|
||||
uniform_all = mode_metrics["uniform_all"][metric]
|
||||
log_gap = math.log(reference / uniform_all)
|
||||
relative_drop = (reference - uniform_all) / reference
|
||||
established = (
|
||||
reference
|
||||
> thresholds["positive_denominator_epsilon"]
|
||||
and uniform_all
|
||||
> thresholds["positive_denominator_epsilon"]
|
||||
and log_gap > 0
|
||||
and relative_drop
|
||||
>= thresholds["global_relative_drop_minimum"]
|
||||
)
|
||||
global_metrics[metric] = {
|
||||
"reference": reference,
|
||||
"uniform_all": uniform_all,
|
||||
"log_gap": log_gap,
|
||||
"relative_drop": relative_drop,
|
||||
"established": established,
|
||||
}
|
||||
for mode in SUFFICIENCY_MODES:
|
||||
sufficiency_by_mode[mode][seed_key] = {}
|
||||
for metric in METRICS:
|
||||
established = global_metrics[metric]["established"]
|
||||
score = (
|
||||
math.log(
|
||||
mode_metrics["detached_learned"][metric]
|
||||
/ mode_metrics[mode][metric]
|
||||
)
|
||||
/ global_metrics[metric]["log_gap"]
|
||||
if established
|
||||
else None
|
||||
)
|
||||
sufficiency_by_mode[mode][seed_key][metric] = {
|
||||
"score": score,
|
||||
"metric_value": mode_metrics[mode][metric],
|
||||
"global_gap_established": established,
|
||||
}
|
||||
for mode in RESTORATION_MODES:
|
||||
restoration_by_mode[mode][seed_key] = {}
|
||||
for metric in METRICS:
|
||||
established = global_metrics[metric]["established"]
|
||||
score = (
|
||||
math.log(
|
||||
mode_metrics[mode][metric]
|
||||
/ mode_metrics["uniform_all"][metric]
|
||||
)
|
||||
/ global_metrics[metric]["log_gap"]
|
||||
if established
|
||||
else None
|
||||
)
|
||||
restoration_by_mode[mode][seed_key][metric] = {
|
||||
"score": score,
|
||||
"metric_value": mode_metrics[mode][metric],
|
||||
"global_gap_established": established,
|
||||
}
|
||||
cells.append(
|
||||
{
|
||||
"seed": value["seed"],
|
||||
"path": str(path),
|
||||
"file_sha256": file_sha256(path),
|
||||
"canonical_sha256": value[
|
||||
"canonical_sha256_without_self"
|
||||
],
|
||||
"mode_metrics": mode_metrics,
|
||||
"global": global_metrics,
|
||||
"interaction_residual": {
|
||||
metric: (
|
||||
1.0
|
||||
- sufficiency_by_mode["uniform_output_only"][
|
||||
seed_key
|
||||
][metric]["score"]
|
||||
- sufficiency_by_mode["uniform_depth_all"][
|
||||
seed_key
|
||||
][metric]["score"]
|
||||
)
|
||||
for metric in METRICS
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
global_gap_passed = all_true(
|
||||
[
|
||||
cell["global"][metric]["established"]
|
||||
for cell in cells
|
||||
for metric in METRICS
|
||||
]
|
||||
)
|
||||
sufficiency_gates = {
|
||||
"groups_6_7": mixed_status(
|
||||
sufficiency_by_mode["uniform_groups_6_7_only"],
|
||||
thresholds["groups_6_7_sufficiency_minimum"],
|
||||
),
|
||||
"group_6": mixed_status(
|
||||
sufficiency_by_mode["uniform_group_6_only"],
|
||||
thresholds["single_group_material_minimum"],
|
||||
),
|
||||
"group_7": mixed_status(
|
||||
sufficiency_by_mode["uniform_group_7_only"],
|
||||
thresholds["single_group_material_minimum"],
|
||||
),
|
||||
"output_half_gap": mixed_status(
|
||||
sufficiency_by_mode["uniform_output_only"],
|
||||
thresholds["output_half_gap_minimum"],
|
||||
),
|
||||
}
|
||||
restoration_gates = {
|
||||
"groups_6_7": mixed_status(
|
||||
restoration_by_mode["uniform_except_groups_6_7"],
|
||||
thresholds["groups_6_7_restoration_minimum"],
|
||||
),
|
||||
"group_6": mixed_status(
|
||||
restoration_by_mode["uniform_except_group_6"],
|
||||
thresholds["single_group_material_minimum"],
|
||||
),
|
||||
"group_7": mixed_status(
|
||||
restoration_by_mode["uniform_except_group_7"],
|
||||
thresholds["single_group_material_minimum"],
|
||||
),
|
||||
}
|
||||
|
||||
branch_gates = {}
|
||||
for group in (6, 7):
|
||||
group_passed = sufficiency_gates[f"group_{group}"]["passed"]
|
||||
candidates = {}
|
||||
for branch, sibling in (("attention", "mlp"), ("mlp", "attention")):
|
||||
branch_mode = f"uniform_group_{group}_{branch}_only"
|
||||
sibling_mode = f"uniform_group_{group}_{sibling}_only"
|
||||
checks = []
|
||||
margins = []
|
||||
for seed in expected_seeds:
|
||||
seed_key = str(seed)
|
||||
for metric in METRICS:
|
||||
left = sufficiency_by_mode[branch_mode][seed_key][
|
||||
metric
|
||||
]["score"]
|
||||
right = sufficiency_by_mode[sibling_mode][seed_key][
|
||||
metric
|
||||
]["score"]
|
||||
margin = (
|
||||
left - right
|
||||
if left is not None and right is not None
|
||||
else None
|
||||
)
|
||||
margins.append(margin)
|
||||
checks.append(
|
||||
left is not None
|
||||
and left >= thresholds["branch_material_minimum"]
|
||||
and margin is not None
|
||||
and margin
|
||||
>= thresholds["branch_dominance_margin"]
|
||||
)
|
||||
candidates[branch] = {
|
||||
"passed": group_passed and all_true(checks),
|
||||
"group_gate_passed": group_passed,
|
||||
"passed_cells": sum(checks),
|
||||
"required_cells": len(checks),
|
||||
"margins": margins,
|
||||
}
|
||||
dominant = [
|
||||
branch
|
||||
for branch, gate in candidates.items()
|
||||
if gate["passed"]
|
||||
]
|
||||
branch_gates[f"group_{group}"] = {
|
||||
"passed": len(dominant) == 1,
|
||||
"dominant_branch": dominant[0] if len(dominant) == 1 else None,
|
||||
"exploratory_sufficiency_only": True,
|
||||
"candidates": candidates,
|
||||
}
|
||||
|
||||
localization_passed = (
|
||||
global_gap_passed
|
||||
and sufficiency_gates["groups_6_7"]["passed"]
|
||||
and restoration_gates["groups_6_7"]["passed"]
|
||||
)
|
||||
localization_status = (
|
||||
"established_at_preregistered_bidirectional_50pct_threshold"
|
||||
if localization_passed
|
||||
else (
|
||||
"one_sided_evidence_localization_not_established"
|
||||
if (
|
||||
sufficiency_gates["groups_6_7"]["passed"]
|
||||
!= restoration_gates["groups_6_7"]["passed"]
|
||||
)
|
||||
else "not_established_at_preregistered_threshold"
|
||||
)
|
||||
)
|
||||
|
||||
mode_means = {}
|
||||
for mode in manifest["matrix_modes"]:
|
||||
mode_means[mode] = {
|
||||
metric: statistics.fmean(
|
||||
cell["mode_metrics"][mode][metric] for cell in cells
|
||||
)
|
||||
for metric in METRICS
|
||||
}
|
||||
sufficiency_means = {
|
||||
mode: {
|
||||
metric: statistics.fmean(
|
||||
sufficiency_by_mode[mode][str(seed)][metric]["score"]
|
||||
for seed in expected_seeds
|
||||
)
|
||||
for metric in METRICS
|
||||
}
|
||||
for mode in SUFFICIENCY_MODES
|
||||
}
|
||||
restoration_means = {
|
||||
mode: {
|
||||
metric: statistics.fmean(
|
||||
restoration_by_mode[mode][str(seed)][metric]["score"]
|
||||
for seed in expected_seeds
|
||||
)
|
||||
for metric in METRICS
|
||||
}
|
||||
for mode in RESTORATION_MODES
|
||||
}
|
||||
|
||||
result = {
|
||||
"schema_version": 1,
|
||||
"protocol_id": PROTOCOL_ID,
|
||||
"study_identity": manifest["study_identity"],
|
||||
"manifest": {
|
||||
"path": str(args.manifest),
|
||||
"file_sha256": file_sha256(args.manifest),
|
||||
},
|
||||
"formal_cells": cells,
|
||||
"replay": {
|
||||
"path": str(args.replay),
|
||||
"file_sha256": file_sha256(args.replay),
|
||||
"canonical_sha256": replay[
|
||||
"canonical_sha256_without_self"
|
||||
],
|
||||
"formal_seed1_exact_excluding_run_kind_and_timing": replay_exact,
|
||||
},
|
||||
"scores": {
|
||||
"sufficiency": sufficiency_by_mode,
|
||||
"restoration": restoration_by_mode,
|
||||
},
|
||||
"means": {
|
||||
"mode_metrics": mode_means,
|
||||
"sufficiency": sufficiency_means,
|
||||
"restoration": restoration_means,
|
||||
},
|
||||
"gates": {
|
||||
"all_input_and_parent_gates_passed": True,
|
||||
"global_gap": {
|
||||
"passed": global_gap_passed,
|
||||
"required_cells": 6,
|
||||
"passed_cells": sum(
|
||||
cell["global"][metric]["established"]
|
||||
for cell in cells
|
||||
for metric in METRICS
|
||||
),
|
||||
},
|
||||
"sufficiency": sufficiency_gates,
|
||||
"restoration": restoration_gates,
|
||||
"localization": {
|
||||
"passed": localization_passed,
|
||||
"status": localization_status,
|
||||
"requires": (
|
||||
"groups 6+7 sufficiency and restoration both >=0.50 "
|
||||
"for 3/3 seeds and both metrics"
|
||||
),
|
||||
},
|
||||
"branch_dominance": branch_gates,
|
||||
},
|
||||
"limitations": [
|
||||
"same-forward diagnostic backward-rule sensitivity only",
|
||||
"reduced byte-level language model, not the Kimi K3 checkpoint",
|
||||
"effects are non-additive and are not contribution percentages",
|
||||
"group 7 includes layers 26-28 outside fixed spike set 21-25",
|
||||
"three-seed threshold gates are not population inference",
|
||||
],
|
||||
}
|
||||
result["canonical_sha256_without_self"] = canonical_sha256(result)
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = args.output.with_suffix(args.output.suffix + ".tmp")
|
||||
temporary.write_text(
|
||||
json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True)
|
||||
+ "\n"
|
||||
)
|
||||
temporary.replace(args.output)
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"output": str(args.output),
|
||||
"formal_cells": len(cells),
|
||||
"replay_exact": replay_exact,
|
||||
"global_gap": global_gap_passed,
|
||||
"localization": localization_status,
|
||||
"canonical_sha256": result[
|
||||
"canonical_sha256_without_self"
|
||||
],
|
||||
},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user