Re-scopes Layer 2 into a cheaper, architecture-general neural collapse proof
before the LLM rung. Realises the same Wright–Fisher abstractions in real trained
generative models on a fully-synthetic sandbox with an exact oracle, reusing
knowledge.metrics/truth/seeding and the output contract so neural curves overlay
the Layer-1 analytic curves.
- src/neural/: synthetic token-grammar sandbox (lossless identity + stochastic
style), ExactOracle, HistogramModel bridge, generation loop, experiment runner
- HARD GATE passed: histogram lineage reproduces Layer 1 exactly (neutral decay,
exact H_eq, tracks run_lineage) — tests/test_neural_validation.py
- torch models: autoregressive RNN + MLP (VAE implemented, not yet fidelity-
passing); determinism seeding derived from the SeedSequence stream
- N0 bridge (neural g*=0.047 ≈ Layer-1 0.048), N1 collapse-in-weights, N2 phase
boundary, N5 architecture-generality (collapse + grounding-rescue in histogram
+ RNN + MLP). Manifests/configs committed; parquet gitignored, hashes tracked
- additive backward-compatible save_artifacts extension; Makefile neural targets
Finding: neural smoothing partially resists H-collapse, so forward-KL and tail
survival are the sharp neural collapse metrics (H is smooth, per Layer 1).
92 tests green. Remaining (tasks/todo.md): N4 merge, N2 refine, N3/N6, VAE
fidelity, MNIST tier, figures. LLM/LoRA rung and C3 deferred.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
295 lines
12 KiB
Python
295 lines
12 KiB
Python
"""Experiment runner: sweep a parameter grid x replicates, write reproducible artifacts.
|
|
|
|
One YAML per experiment (blueprint 2.7 / 4). ``run_experiment`` expands the declared
|
|
sweep grid, runs ``run_lineage`` for every (combo, replicate) at seeds derived from a
|
|
single master seed, and returns long-form results. ``save_artifacts`` writes the output
|
|
contract: ``results.parquet`` + ``resolved_config.yaml`` + ``manifest.json``. Figures are
|
|
regenerated separately from ``results.parquet`` alone.
|
|
|
|
CLI: python -m knowledge.experiment configs/layer1/E2.yaml
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import copy
|
|
import hashlib
|
|
import itertools
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
from importlib.metadata import version
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import yaml
|
|
|
|
from .lineage import run_lineage
|
|
from .seeding import spawn_seeds
|
|
from .teachers import make_correlated_teachers
|
|
from .truth import make_true_distribution
|
|
|
|
# Keys that make up a single-lineage configuration (everything else is experiment-level).
|
|
_LINEAGE_KEYS = ("truth", "dynamics", "generations", "metrics")
|
|
|
|
|
|
def _set_by_path(d: dict, path: str, value: Any) -> None:
|
|
"""Set a nested dict value by a dotted path, creating intermediate dicts."""
|
|
keys = path.split(".")
|
|
cur = d
|
|
for k in keys[:-1]:
|
|
cur = cur.setdefault(k, {})
|
|
cur[keys[-1]] = value
|
|
|
|
|
|
def _apply_param(lineage_cfg: dict, param: str, value: Any) -> dict:
|
|
"""Apply one swept parameter to a lineage config; return label columns to record.
|
|
|
|
Special-cases ``g`` (the grounding fraction): converts to an integer real-sample
|
|
budget ``m = round(n*g/(1-g))`` (so g=0 -> m=0), records both ``g`` and ``m``.
|
|
Any other param is a dotted path into the lineage config.
|
|
"""
|
|
if param == "g":
|
|
n = lineage_cfg["dynamics"]["n"]
|
|
g = float(value)
|
|
m = 0 if g <= 0.0 else int(round(n * g / (1.0 - g)))
|
|
lineage_cfg["dynamics"].setdefault("grounding", {})["m"] = m
|
|
return {"g": g, "m": m}
|
|
if param == "arm":
|
|
# A named arm bundling several overrides applied together (e.g. E6 varies
|
|
# grounding + remint settings jointly). value = {name, set: {dotted.path: v}}.
|
|
for path, v in value.get("set", {}).items():
|
|
_set_by_path(lineage_cfg, path, v)
|
|
return {"arm": value["name"]}
|
|
_set_by_path(lineage_cfg, param, value)
|
|
return {param.split(".")[-1]: value}
|
|
|
|
|
|
def expand_sweeps(cfg: dict) -> list[tuple[dict, dict]]:
|
|
"""Expand the sweep grid into (label, resolved_lineage_cfg) pairs.
|
|
|
|
``sweep`` is a list of ``{param, values}`` entries; the Cartesian product is taken
|
|
(so E4 can sweep K_T x rho). A missing/empty sweep yields a single run.
|
|
|
|
Returns:
|
|
list[tuple[dict, dict]]: One (label-columns, lineage-config) pair per grid point.
|
|
"""
|
|
base = {k: copy.deepcopy(cfg[k]) for k in _LINEAGE_KEYS if k in cfg}
|
|
sweeps = cfg.get("sweep", [])
|
|
if isinstance(sweeps, dict):
|
|
sweeps = [sweeps]
|
|
if not sweeps:
|
|
return [({}, base)]
|
|
params = [s["param"] for s in sweeps]
|
|
value_lists = [list(s["values"]) for s in sweeps]
|
|
combos: list[tuple[dict, dict]] = []
|
|
for values in itertools.product(*value_lists):
|
|
lin = copy.deepcopy(base)
|
|
label: dict = {}
|
|
for param, val in zip(params, values):
|
|
label.update(_apply_param(lin, param, val))
|
|
combos.append((label, lin))
|
|
return combos
|
|
|
|
|
|
def run_experiment(cfg: dict) -> pd.DataFrame:
|
|
"""Run the full experiment: every grid point x every replicate.
|
|
|
|
Replicate seeds are derived once from the master seed and *reused across grid points*,
|
|
so comparisons across sweep values are paired (shared drift noise) — variance
|
|
reduction, and a pure function of the master seed.
|
|
|
|
Args:
|
|
cfg (dict): Parsed experiment YAML (experiment, seed, n_replicates, truth,
|
|
dynamics, generations, metrics, sweep, output).
|
|
|
|
Returns:
|
|
pd.DataFrame: Long-form results; one row per (combo, replicate, generation).
|
|
"""
|
|
name = cfg["experiment"]
|
|
master = int(cfg["seed"])
|
|
n_rep = int(cfg["n_replicates"])
|
|
combos = expand_sweeps(cfg)
|
|
seeds = spawn_seeds(master, n_rep)
|
|
|
|
frames: list[pd.DataFrame] = []
|
|
for label, lineage_cfg in combos:
|
|
for rep, ss in enumerate(seeds):
|
|
df = run_lineage(lineage_cfg, ss)
|
|
for col, val in label.items():
|
|
df[col] = val
|
|
df["replicate"] = rep
|
|
frames.append(df)
|
|
out = pd.concat(frames, ignore_index=True)
|
|
out.insert(0, "experiment", name)
|
|
return out
|
|
|
|
|
|
def run_coverage(cfg: dict) -> pd.DataFrame:
|
|
"""E4 runner: multi-teacher recombination coverage (blueprint 2.5-E4 / 2.7.1).
|
|
|
|
A single distillation step, not a lineage. For each (K_T, rho) grid point and
|
|
replicate: build K_T correlated teachers (marginal retention q, pairwise
|
|
correlation rho), then measure two coverages of the tail:
|
|
|
|
* ``union_coverage`` — fraction of tail items retained by >=1 teacher (the
|
|
construction-level union U(K_T, rho, q); must match the closed form). This is the
|
|
recombination *supply*.
|
|
* ``surviving_mean`` / ``surviving_max`` — fraction of tail items that survive the
|
|
pupil's size-n resampling (+ optional grounding m) under two recombination
|
|
operators: ``mean`` (blueprint mean-mixture distillation) and ``max`` (union-
|
|
preserving model-merge, à la M2N2). Under ``mean`` the union gain is diluted by
|
|
1/K_T and (in the rare-tail linear regime) is exactly cancelled — expected pupil
|
|
tail mass is conserved at q·(tail mass) regardless of K_T, so surviving is flat.
|
|
Under ``max`` each item keeps its strongest teacher, so surviving rises with K_T
|
|
and with decorrelation. The gap ``union - surviving`` is the tail recombination
|
|
supplied but drift/dilution re-erased.
|
|
|
|
Matched budget: the pupil draws n samples total from the combined teachers
|
|
(equivalently n/K_T each), so more teachers != more data.
|
|
"""
|
|
truth, cov = cfg["truth"], cfg["coverage"]
|
|
td = make_true_distribution(
|
|
truth["K"], truth["R"], truth["tail"], truth["tail_frac"], truth["zipf_s"], 0,
|
|
tail_threshold=truth["tail_threshold"],
|
|
)
|
|
tail_idx = np.flatnonzero(td.tail_mask)
|
|
n, q = int(cov["n"]), float(cov["q"])
|
|
retain_thresh = 1e-8 # dropped tails sit at ~tail_floor (1e-9); retained at ~p*_j
|
|
|
|
sweeps = cfg["sweep"]
|
|
if isinstance(sweeps, dict):
|
|
sweeps = [sweeps]
|
|
params = [s["param"] for s in sweeps]
|
|
value_lists = [list(s["values"]) for s in sweeps]
|
|
seeds = spawn_seeds(int(cfg["seed"]), int(cfg["n_replicates"]))
|
|
|
|
rows: list[dict] = []
|
|
for combo in itertools.product(*value_lists):
|
|
d = dict(zip(params, combo))
|
|
K_T, rho = int(d["K_T"]), float(d["rho"])
|
|
g = float(d.get("g", cov.get("g", 0.0))) # g may be swept or fixed
|
|
m = 0 if g <= 0.0 else int(round(n * g / (1.0 - g)))
|
|
for rep, ss in enumerate(seeds):
|
|
child = int(ss.generate_state(1)[0])
|
|
teachers = np.asarray(make_correlated_teachers(
|
|
td.p_star, td.tail_mask, K_T, rho, q, seed=child))
|
|
retained = teachers[:, tail_idx] > retain_thresh # (K_T, T)
|
|
union = float(np.mean(retained.any(axis=0)))
|
|
surviving = {}
|
|
for offset, combine in ((1, teachers.mean), (2, teachers.max)):
|
|
p = combine(axis=0)
|
|
p = p / p.sum()
|
|
rng = np.random.default_rng(child + offset)
|
|
counts = rng.multinomial(n, p)
|
|
if m > 0:
|
|
counts = counts + rng.multinomial(m, td.p_star)
|
|
surviving[offset] = float(np.mean(counts[tail_idx] > 0))
|
|
rows.append({
|
|
"experiment": cfg["experiment"], "K_T": K_T, "rho": rho,
|
|
"replicate": rep, "union_coverage": union,
|
|
"surviving_mean": surviving[1], "surviving_max": surviving[2],
|
|
"g": g, "q": q,
|
|
})
|
|
return pd.DataFrame(rows)
|
|
|
|
|
|
def _git_commit() -> str | None:
|
|
try:
|
|
return subprocess.check_output(
|
|
["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL, text=True
|
|
).strip()
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _content_hash(path: Path) -> str:
|
|
h = hashlib.sha256()
|
|
h.update(path.read_bytes())
|
|
return h.hexdigest()
|
|
|
|
|
|
def save_artifacts(cfg: dict, df: pd.DataFrame, out_dir: Path,
|
|
extra_libs: tuple[str, ...] = (),
|
|
extra_manifest: dict | None = None,
|
|
grid: list | None = None) -> None:
|
|
"""Write the reproducibility output contract (blueprint 2.7 / 4).
|
|
|
|
Writes ``results.parquet``, ``resolved_config.yaml`` (the fully-expanded config), and
|
|
``manifest.json`` (library versions, master seed, git commit, content hash).
|
|
|
|
Args:
|
|
cfg (dict): The parsed experiment config.
|
|
df (pd.DataFrame): The long-form results.
|
|
out_dir (Path): Output directory.
|
|
extra_libs (tuple[str, ...]): Extra library names to record versions for (e.g.
|
|
``torch``, ``torchvision`` for Layer 1.5). Missing libraries are skipped, so a
|
|
caller can pass optional deps unconditionally.
|
|
extra_manifest (dict | None): Extra key/value pairs to merge into the manifest
|
|
(e.g. model architecture, oracle checkpoint hash, determinism flags).
|
|
grid (list | None): Pre-expanded ``[{label, lineage_cfg}, ...]`` to record in the
|
|
resolved config. If None, it is computed via ``expand_sweeps`` for the Layer-1
|
|
``lineage`` kind (a caller with a different schema, e.g. Layer 1.5, passes its
|
|
own expanded grid here).
|
|
"""
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
results_path = out_dir / "results.parquet"
|
|
df.to_parquet(results_path, index=False)
|
|
|
|
resolved = {
|
|
"experiment": cfg["experiment"],
|
|
"seed": cfg["seed"],
|
|
"n_replicates": cfg["n_replicates"],
|
|
"source_config": cfg,
|
|
}
|
|
if grid is not None:
|
|
resolved["grid"] = grid
|
|
elif cfg.get("kind", "lineage") == "lineage":
|
|
resolved["grid"] = [
|
|
{"label": label, "lineage_cfg": lineage_cfg}
|
|
for label, lineage_cfg in expand_sweeps(cfg)
|
|
]
|
|
(out_dir / "resolved_config.yaml").write_text(yaml.safe_dump(resolved, sort_keys=False))
|
|
|
|
libraries: dict[str, str] = {}
|
|
for lib in ("numpy", "scipy", "pandas", "pyarrow") + tuple(extra_libs):
|
|
try:
|
|
libraries[lib] = version(lib)
|
|
except Exception: # optional dep not installed -> omit rather than crash
|
|
pass
|
|
manifest = {
|
|
"experiment": cfg["experiment"],
|
|
"master_seed": cfg["seed"],
|
|
"git_commit": _git_commit(),
|
|
"python": sys.version.split()[0],
|
|
"libraries": libraries,
|
|
"rows": int(len(df)),
|
|
"results_sha256": _content_hash(results_path),
|
|
}
|
|
if extra_manifest:
|
|
manifest.update(extra_manifest)
|
|
(out_dir / "manifest.json").write_text(json.dumps(manifest, indent=2))
|
|
|
|
|
|
def run_and_save(config_path: str | Path) -> Path:
|
|
"""Load an experiment YAML, run it, and write artifacts. Returns the output dir."""
|
|
config_path = Path(config_path)
|
|
cfg = yaml.safe_load(config_path.read_text())
|
|
out_dir = Path(cfg.get("output", {}).get("dir", f"results/{cfg['experiment']}"))
|
|
df = run_coverage(cfg) if cfg.get("kind") == "coverage" else run_experiment(cfg)
|
|
save_artifacts(cfg, df, out_dir)
|
|
return out_dir
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> None:
|
|
parser = argparse.ArgumentParser(description="Run a Layer-1 experiment from a YAML config.")
|
|
parser.add_argument("config", help="Path to configs/layer1/EX.yaml")
|
|
args = parser.parse_args(argv)
|
|
out_dir = run_and_save(args.config)
|
|
print(f"wrote artifacts to {out_dir}/")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|