#!/usr/bin/env python3 """Validate, aggregate, and compact the reduced Attention Residuals study.""" from __future__ import annotations import argparse import hashlib import json import math import os import statistics from pathlib import Path from typing import Any, Iterable PROTOCOL_ID = "llm-atlas-k3-attnres-reduced-v1" ARCHITECTURES = ("baseline", "full", "block") SEEDS = (2026073001, 2026073002, 2026073003) REPLAY_FIELDS = ( "manifest", "model", "optimizer", "hashes", "evaluations", "training_history", "diagnostic", "environment", ) SMOKE_FIELDS = REPLAY_FIELDS def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--formal-dir", type=Path, required=True) parser.add_argument("--smoke-dir", type=Path, required=True) parser.add_argument("--replay", type=Path, required=True) parser.add_argument("--manifest", type=Path, required=True) parser.add_argument("--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 read_json(path: Path) -> dict[str, Any]: return json.loads(path.read_text()) 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: payload = json.dumps( value, ensure_ascii=False, sort_keys=True, separators=(",", ":") ).encode() return hashlib.sha256(payload).hexdigest() def write_json(path: Path, value: dict[str, Any]) -> None: 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" ) os.replace(temporary, path) def mean(values: Iterable[float]) -> float: return statistics.fmean(values) def elementwise_summary(rows: list[list[float]]) -> dict[str, list[float]]: length = len(rows[0]) if any(len(row) != length for row in rows): raise ValueError("array lengths do not match") return { "mean": [mean(row[index] for row in rows) for index in range(length)], "min": [min(row[index] for row in rows) for index in range(length)], "max": [max(row[index] for row in rows) for index in range(length)], } def coefficient_of_variation(values: list[float]) -> float: average = mean(values) variance = mean((value - average) ** 2 for value in values) return math.sqrt(variance) / average def verdict(deltas: list[float]) -> dict[str, Any]: average = mean(deltas) if all(delta < 0 for delta in deltas) and average <= -0.010: label = "directional support in this reduced protocol" elif all(delta > 0 for delta in deltas) and average >= 0.010: label = "directional concern in this reduced protocol" else: label = "inconclusive at this budget" return { "paired_deltas_bpc": deltas, "mean_delta_bpc": average, "min_delta_bpc": min(deltas), "max_delta_bpc": max(deltas), "same_direction": len({delta < 0 for delta in deltas}) == 1, "threshold_bpc": 0.010, "verdict": label, } def average_depth_weights(runs: list[dict[str, Any]]) -> dict[str, Any]: rows_by_run = [run["diagnostic"]["depth_weights"] for run in runs] layer_count = len(rows_by_run[0]) if any(len(rows) != layer_count for rows in rows_by_run): raise ValueError("depth-weight layer counts differ") rows = [] max_sources = 0 for layer in range(layer_count): source_count = rows_by_run[0][layer]["sources"] if any(rows[layer]["sources"] != source_count for rows in rows_by_run): raise ValueError("source count differs across seeds") weights = [ mean(rows_by_run[seed_index][layer]["mean_weights"][source] for seed_index in range(len(runs))) for source in range(source_count) ] entropies = [ rows_by_run[seed_index][layer]["entropy_mean"] for seed_index in range(len(runs)) ] rows.append( { "sublayer": layer + 1, "sources": source_count, "mean_weights": weights, "entropy_mean": mean(entropies), "entropy_min": min(entropies), "entropy_max": max(entropies), } ) max_sources = max(max_sources, source_count) output_rows = [run["diagnostic"]["output_weights"] for run in runs] output_source_count = output_rows[0]["sources"] output_weights = [ mean(row["mean_weights"][source] for row in output_rows) for source in range(output_source_count) ] return { "rows": rows, "max_sources": max_sources, "output": { "sources": output_source_count, "mean_weights": output_weights, "entropy_mean": mean(row["entropy_mean"] for row in output_rows), "entropy_min": min(row["entropy_mean"] for row in output_rows), "entropy_max": max(row["entropy_mean"] for row in output_rows), }, } def main() -> None: args = parse_args() manifest = read_json(args.manifest) if manifest["protocol_id"] != PROTOCOL_ID: raise ValueError("manifest protocol mismatch") runs: dict[tuple[str, int], dict[str, Any]] = {} formal_file_hashes: dict[str, str] = {} for seed in SEEDS: for architecture in ARCHITECTURES: path = args.formal_dir / f"{architecture}-{seed}.json" run = read_json(path) if run["protocol_id"] != PROTOCOL_ID: raise ValueError(f"protocol mismatch: {path}") if run["run_kind"] != "formal": raise ValueError(f"not a formal run: {path}") if run["architecture"] != architecture or run["seed"] != seed: raise ValueError(f"cell identity mismatch: {path}") if run["steps"] != 2000 or run["batch_size"] != 32: raise ValueError(f"formal budget mismatch: {path}") if run["target_bytes_seen"] != 16_384_000: raise ValueError(f"target byte count mismatch: {path}") if run["manifest"]["file_sha256"] != file_sha256(args.manifest): raise ValueError(f"manifest file hash mismatch: {path}") if run["manifest"]["formal_schedule_sha256"] != ( manifest["windows"]["formal_schedule_sha256"] ): raise ValueError(f"schedule mismatch: {path}") if run["evaluations"][-1]["step"] != 2000: raise ValueError(f"missing final evaluation: {path}") if any( not math.isfinite(value) for evaluation in run["evaluations"] for value in ( evaluation["cross_entropy_nats"], evaluation["bits_per_byte"], ) ): raise ValueError(f"non-finite evaluation: {path}") if len(run["diagnostic"]["layer_input_rms"]) != 32: raise ValueError(f"diagnostic depth mismatch: {path}") if len(run["diagnostic"]["core_parameter_grad_rms_by_block"]) != 16: raise ValueError(f"gradient depth mismatch: {path}") runs[(architecture, seed)] = run formal_file_hashes[path.name] = file_sha256(path) common_initial_exact = {} for seed in SEEDS: hashes = { architecture: runs[(architecture, seed)]["hashes"][ "initial_common_parameters" ] for architecture in ARCHITECTURES } common_initial_exact[str(seed)] = { "hashes": hashes, "exact": len(set(hashes.values())) == 1, } if not common_initial_exact[str(seed)]["exact"]: raise ValueError(f"common initialization mismatch for seed {seed}") by_seed = [] for seed in SEEDS: values = { architecture: runs[(architecture, seed)]["evaluations"][-1][ "bits_per_byte" ] for architecture in ARCHITECTURES } by_seed.append( { "seed": seed, "final_bpc": values, "full_minus_baseline": values["full"] - values["baseline"], "block_minus_baseline": values["block"] - values["baseline"], "block_minus_full": values["block"] - values["full"], } ) final = { "by_seed": by_seed, "means": { architecture: mean( runs[(architecture, seed)]["evaluations"][-1]["bits_per_byte"] for seed in SEEDS ) for architecture in ARCHITECTURES }, "full_contrast": verdict( [row["full_minus_baseline"] for row in by_seed] ), "block_contrast": verdict( [row["block_minus_baseline"] for row in by_seed] ), "block_minus_full": { "paired_deltas_bpc": [row["block_minus_full"] for row in by_seed], "mean_delta_bpc": mean(row["block_minus_full"] for row in by_seed), }, } evaluation_steps = [ evaluation["step"] for evaluation in runs[("baseline", SEEDS[0])]["evaluations"] ] curves = {} for architecture in ARCHITECTURES: curve = [] for index, step in enumerate(evaluation_steps): values = [ runs[(architecture, seed)]["evaluations"][index]["bits_per_byte"] for seed in SEEDS ] if any( runs[(architecture, seed)]["evaluations"][index]["step"] != step for seed in SEEDS ): raise ValueError("evaluation step mismatch") curve.append( { "step": step, "mean_bpc": mean(values), "min_bpc": min(values), "max_bpc": max(values), "by_seed": values, } ) curves[architecture] = curve timing = {} for architecture in ARCHITECTURES: cells = [runs[(architecture, seed)]["timing"] for seed in SEEDS] timing[architecture] = { "mean_step_ms": mean(cell["mean_ms"] for cell in cells), "median_step_ms": mean(cell["median_ms"] for cell in cells), "p95_step_ms": mean(cell["p95_ms"] for cell in cells), "mean_peak_allocated_bytes": mean( cell["peak_allocated_bytes"] for cell in cells ), "mean_peak_reserved_bytes": mean( cell["peak_reserved_bytes"] for cell in cells ), "by_seed": cells, } timing["relative_to_baseline"] = { architecture: { "step_time_ratio": timing[architecture]["mean_step_ms"] / timing["baseline"]["mean_step_ms"], "allocated_memory_ratio": timing[architecture][ "mean_peak_allocated_bytes" ] / timing["baseline"]["mean_peak_allocated_bytes"], } for architecture in ("full", "block") } parameters = { architecture: runs[(architecture, SEEDS[0])]["model"]["parameters"] for architecture in ARCHITECTURES } parameters["mixer_overhead_fraction_of_baseline"] = ( parameters["full"]["mixer"] / parameters["baseline"]["total"] ) traces = {} gradients = {} for architecture in ARCHITECTURES: architecture_runs = [runs[(architecture, seed)] for seed in SEEDS] traces[architecture] = { key: elementwise_summary( [run["diagnostic"][key] for run in architecture_runs] ) for key in ( "layer_input_rms", "branch_output_rms", "stream_state_rms", ) } gradient_rows = [ run["diagnostic"]["core_parameter_grad_rms_by_block"] for run in architecture_runs ] gradients[architecture] = { "by_block": elementwise_summary(gradient_rows), "cv_by_seed": [ coefficient_of_variation(row) for row in gradient_rows ], "mean_cv": mean(coefficient_of_variation(row) for row in gradient_rows), "first_last_ratio_by_seed": [ row[0] / row[-1] for row in gradient_rows ], "mean_first_last_ratio": mean(row[0] / row[-1] for row in gradient_rows), } mixers = { architecture: average_depth_weights( [runs[(architecture, seed)] for seed in SEEDS] ) for architecture in ("full", "block") } full_branch = traces["full"]["branch_output_rms"]["mean"] largest_index = max(range(len(full_branch)), key=full_branch.__getitem__) # Full output source 0 is the embedding; branch l is source l+1. largest_source_weight = mixers["full"]["output"]["mean_weights"][ largest_index + 1 ] uniform_output_weight = 1 / mixers["full"]["output"]["sources"] posthoc = { "label": "post-hoc descriptive callout; not a preregistered endpoint", "largest_full_branch_sublayer": largest_index + 1, "largest_full_branch_rms": full_branch[largest_index], "corresponding_final_output_weight": largest_source_weight, "uniform_final_output_weight": uniform_output_weight, "weight_over_uniform": largest_source_weight / uniform_output_weight, } replay = read_json(args.replay) formal_replay_source = runs[("block", 2026073001)] replay_exact = { field: formal_replay_source[field] == replay[field] for field in REPLAY_FIELDS } if not all(replay_exact.values()): raise ValueError(f"formal replay mismatch: {replay_exact}") smoke = {} for architecture in ARCHITECTURES: first_path = args.smoke_dir / f"{architecture}-2026073001-a.json" second_path = args.smoke_dir / f"{architecture}-2026073001-b.json" first = read_json(first_path) second = read_json(second_path) exact = {field: first[field] == second[field] for field in SMOKE_FIELDS} if not all(exact.values()): raise ValueError(f"smoke mismatch for {architecture}: {exact}") smoke[architecture] = { "fields": exact, "all_exact": True, "first_sha256": file_sha256(first_path), "second_sha256": file_sha256(second_path), } reproduction = { "schema_version": 1, "protocol_id": PROTOCOL_ID, "manifest_sha256": file_sha256(args.manifest), "formal_files": formal_file_hashes, "common_initial_parameters": common_initial_exact, "smoke": smoke, "formal_replay": { "architecture": "block", "seed": 2026073001, "fields": replay_exact, "all_numeric_and_hash_fields_exact": all(replay_exact.values()), "timing_exact_required": False, "timing_exact_observed": formal_replay_source["timing"] == replay["timing"], "formal_file_sha256": formal_file_hashes[ "block-2026073001.json" ], "replay_file_sha256": file_sha256(args.replay), }, } reproduction["canonical_sha256_without_self"] = canonical_sha256(reproduction) analysis = { "final_validation": final, "evaluation_curves": curves, "timing": timing, "parameters": parameters, "traces": traces, "gradients": gradients, "mixers": mixers, "posthoc": posthoc, "interpretation": { "primary": ( "Both Full and Block AttnRes satisfy the preregistered " "directional-support rule in this reduced protocol." ), "bounded_depth_pattern": ( "Block partial-state RMS resets every four residual sublayers; " "the complete 32-point vectors are reported." ), "gradient_boundary": ( "The preregistered core-parameter gradient RMS is not more " "uniform for AttnRes here; this metric and scale do not reproduce " "the paper's large-model gradient-magnitude result." ), "scope": ( "Reduced byte-level WikiText-2 mechanism probe; not a K3 " "checkpoint run, paper-scale reproduction, benchmark, or " "same-FLOP comparison." ), }, } raw = { "schema_version": 1, "protocol_id": PROTOCOL_ID, "manifest": manifest, "provenance": { "manifest_file_sha256": file_sha256(args.manifest), "formal_file_sha256": formal_file_hashes, "reproduction_sha256": reproduction[ "canonical_sha256_without_self" ], }, "formal_runs": [ runs[(architecture, seed)] for seed in SEEDS for architecture in ARCHITECTURES ], "analysis": analysis, "reproduction": reproduction, } raw["canonical_sha256_without_self"] = canonical_sha256(raw) compact = { "schema_version": 1, "protocol_id": PROTOCOL_ID, "dataset": { "repository": manifest["dataset"]["repository"], "revision": manifest["dataset"]["revision"], "train_bytes": manifest["dataset"]["splits"]["train"][ "concatenated_bytes" ], "schedule_sha256": manifest["windows"]["formal_schedule_sha256"], "validation_sha256": manifest["windows"][ "validation_tensor_sha256" ], }, "grid": { "architectures": list(ARCHITECTURES), "seeds": list(SEEDS), "runs": 9, "steps_per_run": 2000, "target_bytes_per_run": 16_384_000, "target_bytes_total": 9 * 16_384_000, }, "final_validation": final, "evaluation_curves": curves, "timing": timing, "parameters": parameters, "traces": traces, "gradients": gradients, "mixers": mixers, "posthoc": posthoc, "interpretation": analysis["interpretation"], "reproduction": reproduction, "source_sha256": raw["canonical_sha256_without_self"], } compact["canonical_sha256_without_self"] = canonical_sha256(compact) write_json(args.reproduction_output, reproduction) write_json(args.output, raw) write_json(args.compact_output, compact) print( json.dumps( { "output": str(args.output), "compact_output": str(args.compact_output), "reproduction_output": str(args.reproduction_output), "raw_sha256": file_sha256(args.output), "compact_sha256": file_sha256(args.compact_output), "reproduction_sha256": file_sha256(args.reproduction_output), "full": final["full_contrast"], "block": final["block_contrast"], "formal_replay": reproduction["formal_replay"], }, ensure_ascii=False, indent=2, ) ) if __name__ == "__main__": main()