wuyang/llm-atlas
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

research: audit reduced AttnRes study

Browse Source
This commit is contained in:
wuyang
2026-07-30 07:15:05 +08:00
parent 5f49906be6
commit d1d9d22bf3
6 changed files with 25838 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experiments/k3/attnres/README.md
+21
View File
@@ -41,3 +41,24 @@ python experiments/k3/attnres/train.py \
Raw parquet and checkpoints stay in the local cache. Frozen manifests, metric
JSON, analyses, code, checksums, and a compact website payload enter the public
repository.
## Validate and aggregate the complete study
After the nine formal cells, the preregistered replay, and the paired smoke runs
exist in the cache:
```bash
python experiments/k3/attnres/analyze.py \
--formal-dir /home/wuyang/.cache/llm-atlas/k3-attnres-reduced-v1/formal \
--smoke-dir /home/wuyang/.cache/llm-atlas/k3-attnres-reduced-v1/smoke \
--replay /home/wuyang/.cache/llm-atlas/k3-attnres-reduced-v1/replay/block-2026073001.json \
--manifest experiments/k3/attnres/manifest.json \
--output src/data/k3-attnres-reduced.json \
--compact-output src/data/k3-attnres-reduced-compact.json \
--reproduction-output experiments/k3/attnres/reproduction.json
```
The aggregator fails closed on protocol identity, grid completeness, byte
budget, schedule hashes, shared initialization, non-finite metrics, diagnostic
shape, smoke mismatch, or formal replay mismatch. Timing is recorded but is not
required to replay bit-for-bit.
experiments/k3/attnres/analyze.py
+536
View File
@@ -0,0 +1,536 @@
#!/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()
experiments/k3/attnres/reproduction.json
+109
View File
@@ -0,0 +1,109 @@
{
"canonical_sha256_without_self": "89a127e625ccefbd9b749e6ffca4ba72868e217e0c7a24773024bd2fafde427a",
"common_initial_parameters": {
"2026073001": {
"exact": true,
"hashes": {
"baseline": "af2724a1c34bcfd61d8a8bef402246430898c815e56b6e5c5257949a4eb0e7b1",
"block": "af2724a1c34bcfd61d8a8bef402246430898c815e56b6e5c5257949a4eb0e7b1",
"full": "af2724a1c34bcfd61d8a8bef402246430898c815e56b6e5c5257949a4eb0e7b1"
}
},
"2026073002": {
"exact": true,
"hashes": {
"baseline": "9fd4f04212ab5cea822f469902d8e80ecc368da329f3c20abacfe6b7a50ed523",
"block": "9fd4f04212ab5cea822f469902d8e80ecc368da329f3c20abacfe6b7a50ed523",
"full": "9fd4f04212ab5cea822f469902d8e80ecc368da329f3c20abacfe6b7a50ed523"
}
},
"2026073003": {
"exact": true,
"hashes": {
"baseline": "7a565cd353efdb3b95e9b8b1844581082c029991ea18d438b4b18566970c8c61",
"block": "7a565cd353efdb3b95e9b8b1844581082c029991ea18d438b4b18566970c8c61",
"full": "7a565cd353efdb3b95e9b8b1844581082c029991ea18d438b4b18566970c8c61"
}
}
},
"formal_files": {
"baseline-2026073001.json": "f5cddecace6a70ee3824f272811d2a6336ab6dc24ad1100b18292f4413692e9a",
"baseline-2026073002.json": "4c3b9e1ea2ed6fd3078f1212984327364832a1146bfcbdb3a6f2e61b589bfaa8",
"baseline-2026073003.json": "e5ce5a8944dad0c4811f8bf19b4cfec8ce4a9283858fd59cd739dbc37626707b",
"block-2026073001.json": "5df870369d9a86ccb4ba4191fbd1d6f3642893dd47a60f8f6d1143006bdfbdaf",
"block-2026073002.json": "be1f109a630e27c8469438e33b53806cdaecac2f712af4f885da83f00ef6843b",
"block-2026073003.json": "949ec51ca3a219a11f260171da01f737296241536f5e139cf07d444c27efba92",
"full-2026073001.json": "648cb25ea98da1868779733da155275e9a16aad5efbbeeced7812c19d75817d5",
"full-2026073002.json": "1fa715bbd81a3a04a5aa0ba0fc27feb8883c44a01e9f28e281a309a394447979",
"full-2026073003.json": "4e34fa35bbfec460cabc94bce7a08f44fc188d908705ae2e256500f00813b4e3"
},
"formal_replay": {
"all_numeric_and_hash_fields_exact": true,
"architecture": "block",
"fields": {
"diagnostic": true,
"environment": true,
"evaluations": true,
"hashes": true,
"manifest": true,
"model": true,
"optimizer": true,
"training_history": true
},
"formal_file_sha256": "5df870369d9a86ccb4ba4191fbd1d6f3642893dd47a60f8f6d1143006bdfbdaf",
"replay_file_sha256": "e74d3323e5fe31378bb8aad7a8efa2fb91995cd7224c158cf03006466cdea2a7",
"seed": 2026073001,
"timing_exact_observed": false,
"timing_exact_required": false
},
"manifest_sha256": "9778ade5b1c9dd7676d2cdc52b4e4e7ff5ae513cb56c667974e2422702f9dc2b",
"protocol_id": "llm-atlas-k3-attnres-reduced-v1",
"schema_version": 1,
"smoke": {
"baseline": {
"all_exact": true,
"fields": {
"diagnostic": true,
"environment": true,
"evaluations": true,
"hashes": true,
"manifest": true,
"model": true,
"optimizer": true,
"training_history": true
},
"first_sha256": "fd8b14f70a6a14978e65f9899b694164ba91753bf82eead436a40837c251e82c",
"second_sha256": "fd8b14f70a6a14978e65f9899b694164ba91753bf82eead436a40837c251e82c"
},
"block": {
"all_exact": true,
"fields": {
"diagnostic": true,
"environment": true,
"evaluations": true,
"hashes": true,
"manifest": true,
"model": true,
"optimizer": true,
"training_history": true
},
"first_sha256": "980a4a534e458199d95b5864b008a81f51b565e05a7b2f24a644b36d2134eccc",
"second_sha256": "980a4a534e458199d95b5864b008a81f51b565e05a7b2f24a644b36d2134eccc"
},
"full": {
"all_exact": true,
"fields": {
"diagnostic": true,
"environment": true,
"evaluations": true,
"hashes": true,
"manifest": true,
"model": true,
"optimizer": true,
"training_history": true
},
"first_sha256": "c6e32b737c7bf7fdc9b647650ebd7c9d2f9f8550f66a5b84f168f16fa7d2eacb",
"second_sha256": "c6e32b737c7bf7fdc9b647650ebd7c9d2f9f8550f66a5b84f168f16fa7d2eacb"
}
}
}