From e9ca86c21e3d0e3ece33253f1629fb5a9de02f60 Mon Sep 17 00:00:00 2001 From: wuyang <5700876+banisherwy@user.noreply.gitee.com> Date: Thu, 30 Jul 2026 11:58:34 +0800 Subject: [PATCH] research: lock AttnRes spike analyzer --- experiments/k3/attnres_spike/analyze.py | 751 ++++++++++++++++++++++++ 1 file changed, 751 insertions(+) create mode 100644 experiments/k3/attnres_spike/analyze.py diff --git a/experiments/k3/attnres_spike/analyze.py b/experiments/k3/attnres_spike/analyze.py new file mode 100644 index 0000000..4907eec --- /dev/null +++ b/experiments/k3/attnres_spike/analyze.py @@ -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()