#!/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()