diff --git a/Makefile b/Makefile index 6b01304..4b1d027 100644 --- a/Makefile +++ b/Makefile @@ -18,12 +18,14 @@ test: ## correctness tests + scientific-validation tests (the spine layer1: ## run experiments E1-E6 + the learning-kernel bridge (analytic) for e in E1 E2 E3 E4 E5 E6 E7 E8 E9 E10 E11 E12 E12_nk kernel_sharpen kernel_smooth; do uv run python -m knowledge.experiment configs/layer1/$$e.yaml; done -neural: ## run Layer 1.5 synthetic neural experiments (excludes the heavy MNIST tier) - for c in configs/neural/*.yaml; do case "$$c" in *mnist*) ;; \ +neural: ## run Layer 1.5 synthetic neural experiments (excludes the MNIST/torchvision tiers) + for c in configs/neural/*.yaml; do case "$$c" in *mnist*|*speciation_real*) ;; \ *) uv run python -m neural.experiment "$$c" ;; esac; done -mnist: ## run the real-MNIST confirmation tier (needs env-mnist; downloads MNIST) +mnist: ## run the torchvision tiers: MNIST collapse + E13 real-weight speciation (needs env-mnist) uv run python -m neural.experiment configs/neural/mnist_collapse.yaml + uv run python -m neural.experiment configs/neural/speciation_real.yaml + uv run python -m neural.experiment configs/neural/speciation_real_cliff.yaml env-llm: ## add the LLM stack for the Layer-2 prototype (GPU; transformers/peft) uv sync --extra dev --extra neural --extra llm diff --git a/configs/neural/speciation_real.yaml b/configs/neural/speciation_real.yaml new file mode 100644 index 0000000..d533ac4 --- /dev/null +++ b/configs/neural/speciation_real.yaml @@ -0,0 +1,28 @@ +experiment: speciation_real +kind: speciation_real +seed: 13 +n_replicates: 3 + +# E13 — REAL-WEIGHT model speciation (the real-weight image of E12). Two small no-BN MLPs are forked +# from a shared MNIST base and trained independently; we merge them (weight averaging) and measure the +# linear-mode-connectivity barrier BEFORE and AFTER Git Re-Basin permutation alignment. The barrier +# alignment removes is a coordinate artefact; the RESIDUAL it cannot remove is the reproductive-isolation +# signal. Three conditions decompose it: `shared` (same task, shared fork — no barrier, trivially +# mergeable), `independent` (same task, different init — large barrier that alignment REMOVES: coordinate +# artefact, residual ~0), `conflict` (conflicting label maps — large barrier alignment CANNOT remove: +# residual = real speciation). The residual after alignment is the discriminator. Falsifier: alignment +# fails to remove the independent-init barrier (then residual isn't meaningful), or conflict shows no +# residual. In-house deterministic weight-matching (scipy); statistically reproducible (seeded torch). + +speciation_real: + sizes: [784, 512, 512, 10] + conditions: [shared, independent, conflict] + t_div: [100, 200, 400, 800, 1600] + base_steps: 500 + lr: 0.05 + batch: 128 + n_eval: 2000 + data_root: data + +output: + dir: results/speciation_real diff --git a/configs/neural/speciation_real_cliff.yaml b/configs/neural/speciation_real_cliff.yaml new file mode 100644 index 0000000..14bdb64 --- /dev/null +++ b/configs/neural/speciation_real_cliff.yaml @@ -0,0 +1,24 @@ +experiment: speciation_real_cliff +kind: speciation_real +seed: 13 +n_replicates: 3 + +# E13 (cliff) — the real-weight reproductive-isolation curve, mirroring E12's isolation cliff. Two MLPs +# forked from a shared MNIST base; child B learns a CONFLICTING label map on a fraction `conflict_frac` +# of the classes (the rest agree). Sweeping the conflict fraction traces the residual (after-alignment) +# barrier from 0 (no conflict = same species) to its maximum (full conflict = fully isolated) — the +# real-weight image of "the isolation cliff rises with epistasis". Alignment is applied throughout, so +# the curve is the residual that permutation cannot explain away. + +speciation_real: + sizes: [784, 512, 512, 10] + conflict_fracs: [0.0, 0.2, 0.4, 0.6, 0.8, 1.0] + t_div_fixed: 800 + base_steps: 500 + lr: 0.05 + batch: 128 + n_eval: 2000 + data_root: data + +output: + dir: results/speciation_real_cliff diff --git a/figures/plot_speciation_real.py b/figures/plot_speciation_real.py new file mode 100644 index 0000000..82bd8b5 --- /dev/null +++ b/figures/plot_speciation_real.py @@ -0,0 +1,67 @@ +"""E13 figure — real-weight model speciation with Git Re-Basin. + +(A) The barrier decomposition per condition: the linear-mode-connectivity error barrier between two +merged MLPs, split into the part permutation alignment REMOVES (coordinate artefact) and the RESIDUAL it +cannot (reproductive isolation). `shared` ≈ 0; `independent` (same task, different init) is almost all +removable (residual ≈ 0 — same species, different basis); `conflict` (conflicting tasks) is almost all +residual (real isolation). (B) The isolation cliff: residual barrier vs the fraction of conflicting +classes — the real-weight image of E12's cliff, after alignment (so it is not a coordinate artefact). + +Usage: python figures/plot_speciation_real.py +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +sys.path.insert(0, str(Path(__file__).parent)) +from _figlib import load_bundle, savefig # noqa: E402 + + +def main() -> None: + dec, _ = load_bundle("results/speciation_real") + cliff, _ = load_bundle("results/speciation_real_cliff") + + fig, axes = plt.subplots(1, 2, figsize=(13, 5)) + + # Panel A: removable (coordinate artefact) vs residual (isolation), stacked, per condition. + ax = axes[0] + order = [c for c in ["shared", "independent", "conflict"] if c in set(dec["condition"])] + g = dec.groupby("condition").agg(removable=("removable", "mean"), + residual=("residual", "mean")).reindex(order) + x = np.arange(len(order)) + ax.bar(x, g["removable"], 0.6, label="removable by alignment\n(coordinate artefact)", color="#9ecae1") + ax.bar(x, g["residual"], 0.6, bottom=g["removable"], label="residual after alignment\n(reproductive isolation)", + color="#d62728") + ax.set_xticks(x); ax.set_xticklabels(order) + ax.set(ylabel="linear-mode-connectivity error barrier", + title="Merge barrier = coordinate artefact + residual isolation\n" + "(same task even across inits is coordinate; conflict is real)") + ax.legend(frameon=False, fontsize=8) + + # Panel B: the isolation cliff — residual barrier vs conflict fraction. + ax = axes[1] + cg = cliff.groupby("conflict_frac").agg(res_m=("residual", "mean"), res_s=("residual", "std"), + nai_m=("barrier_naive", "mean")).reset_index() + ax.plot(cg["conflict_frac"], cg["nai_m"], "--o", color="#999", lw=1.4, label="naive barrier") + ax.plot(cg["conflict_frac"], cg["res_m"], "-o", color="#d62728", lw=2, label="residual (after alignment)") + ax.fill_between(cg["conflict_frac"], cg["res_m"] - cg["res_s"], cg["res_m"] + cg["res_s"], + color="#d62728", alpha=0.15) + ax.set(xlabel="fraction of classes with conflicting labels", ylabel="error barrier", + ylim=(-0.02, None), + title="The reproductive-isolation cliff, in real weights\n" + "(residual rises with task conflict — not removable by alignment)") + ax.legend(frameon=False, fontsize=9) + + fig.suptitle("E13 — real-weight model speciation: what permutation alignment can and cannot merge", + y=1.02, fontsize=13) + fig.tight_layout() + savefig(fig, "results/speciation_real", "speciation_real") + + +if __name__ == "__main__": + main() diff --git a/results/speciation_real/README.md b/results/speciation_real/README.md new file mode 100644 index 0000000..74993d3 --- /dev/null +++ b/results/speciation_real/README.md @@ -0,0 +1,54 @@ +# E13 — Real-weight model speciation (Git Re-Basin residual) + +**Claim tested.** E12 predicts model *speciation* analytically: as two lineages diverge, recombination +(merging) fails, via Bateson–Dobzhansky–Muller incompatibilities. E13 confirms it in **real trained +weights**, and — decisively — separates the part of the incompatibility that is a mere **coordinate +artefact** (removable by permuting hidden units; Git Re-Basin, Ainsworth et al. 2022) from the +**residual** that permutation *cannot* remove, which is the true reproductive-isolation signal. This is +the experiment that answers the mode-connectivity reviewer: if alignment removes the barrier, it was a +coordinate artefact; the barrier that *survives* alignment is real speciation. + +**Setup.** Small no-BatchNorm MLPs (784–512–512–10) on MNIST — the clean Re-Basin regime. Two children +are forked from a shared base and trained independently; we weight-average them and measure the +**linear-mode-connectivity error barrier** before (`naive`) and after (`aligned`) in-house, deterministic +Git Re-Basin weight-matching (`neural/rebasin.py`, scipy `linear_sum_assignment`). Statistically +reproducible (seeded torch; NumPy/scipy alignment is deterministic). 3 replicates. + +### Results — the decomposition (mean over divergence, reps) +| condition | naive barrier | removable (coordinate) | **residual (isolation)** | +|---|---|---|---| +| `shared` (same task, shared fork) | 0.00 | 0.00 | **0.00** | +| `independent` (same task, different init) | 0.056 | 0.055 | **0.001** | +| `conflict` (conflicting label maps) | 0.496 | 0.000 | **0.496** | + +- **`independent`**: two nets trained *from different random inits* on the *same task* have a real naive + barrier — which alignment **removes ~98%** of (residual 0.001). Same species, different basis: the + incompatibility is a coordinate artefact. (This reproduces the canonical Git Re-Basin result and + proves our alignment works.) +- **`conflict`**: two nets that learned *conflicting* functions have a large barrier that alignment + **removes none** of (residual 0.496). Different species: genuine reproductive isolation. Because + alignment demonstrably works on `independent`, this residual cannot be dismissed as a failure to align. +- The **residual after alignment** is therefore the clean discriminator: ~0 for compatible models (even + independently trained), large only for functionally incompatible ones. + +### Results — the isolation cliff (`speciation_real_cliff/`) +Sweeping the fraction of classes on which child B learns a *conflicting* label map, the residual +(after-alignment) barrier rises monotonically — the real-weight image of E12's cliff: + +| conflict fraction | 0.0 | 0.2 | 0.4 | 0.6 | 0.8 | 1.0 | +|---|---|---|---|---|---|---| +| residual barrier | 0.00 | 0.13 | 0.19 | 0.28 | 0.40 | 0.49 | + +residual = naive at every point (alignment removes nothing in the conflict condition), so the cliff is +genuinely functional isolation, not a coordinate artefact. + +### Positioning +The incumbents each hold one piece: Git Re-Basin / Entezari (barriers are coordinate artefacts), +Frankle (the fork-instability protocol), Pari et al. 2024 (specialisation diverges representations, +route don't fuse), Zhou et al. 2026 (predict mergeability from divergence metrics). E13's contribution +is the synthesis they lack: a controlled decomposition where alignment cleanly partitions the merge +barrier into a **removable coordinate artefact** and a **residual reproductive-isolation** term that +rises with task conflict — the real-weight confirmation of E12's speciation prediction, and the direct +answer to "isn't this just a loss barrier / permutation artefact?" **Falsifier (not triggered):** +alignment failing to remove the independent-init barrier (then residual is meaningless), or conflict +showing no residual — instead alignment removed 98% of the former and 0% of the latter. diff --git a/results/speciation_real/manifest.json b/results/speciation_real/manifest.json new file mode 100644 index 0000000..8c6b8c6 --- /dev/null +++ b/results/speciation_real/manifest.json @@ -0,0 +1,18 @@ +{ + "experiment": "speciation_real", + "master_seed": 13, + "git_commit": "56f642e7f9ae01fe01d863bdffe98b226b9dbb7b", + "python": "3.14.5", + "libraries": { + "numpy": "2.5.0", + "scipy": "1.18.0", + "pandas": "3.0.3", + "pyarrow": "24.0.0", + "torch": "2.12.1", + "torchvision": "0.27.1" + }, + "rows": 45, + "results_sha256": "14b15c16ea8a43523fdc929641b8ad5741445435203d1511ecb698097da0a806", + "layer": "1.5", + "tier": "speciation_real" +} \ No newline at end of file diff --git a/results/speciation_real/resolved_config.yaml b/results/speciation_real/resolved_config.yaml new file mode 100644 index 0000000..a36b267 --- /dev/null +++ b/results/speciation_real/resolved_config.yaml @@ -0,0 +1,31 @@ +experiment: speciation_real +seed: 13 +n_replicates: 3 +source_config: + experiment: speciation_real + kind: speciation_real + seed: 13 + n_replicates: 3 + speciation_real: + sizes: + - 784 + - 512 + - 512 + - 10 + conditions: + - shared + - independent + - conflict + t_div: + - 100 + - 200 + - 400 + - 800 + - 1600 + base_steps: 500 + lr: 0.05 + batch: 128 + n_eval: 2000 + data_root: data + output: + dir: results/speciation_real diff --git a/results/speciation_real/speciation_real.pdf b/results/speciation_real/speciation_real.pdf new file mode 100644 index 0000000..80c8800 Binary files /dev/null and b/results/speciation_real/speciation_real.pdf differ diff --git a/results/speciation_real/speciation_real.png b/results/speciation_real/speciation_real.png new file mode 100644 index 0000000..68f5233 Binary files /dev/null and b/results/speciation_real/speciation_real.png differ diff --git a/results/speciation_real_cliff/manifest.json b/results/speciation_real_cliff/manifest.json new file mode 100644 index 0000000..d3fe0df --- /dev/null +++ b/results/speciation_real_cliff/manifest.json @@ -0,0 +1,18 @@ +{ + "experiment": "speciation_real_cliff", + "master_seed": 13, + "git_commit": "56f642e7f9ae01fe01d863bdffe98b226b9dbb7b", + "python": "3.14.5", + "libraries": { + "numpy": "2.5.0", + "scipy": "1.18.0", + "pandas": "3.0.3", + "pyarrow": "24.0.0", + "torch": "2.12.1", + "torchvision": "0.27.1" + }, + "rows": 18, + "results_sha256": "0252581d848376ad698f56d4e69edbcae40cf5a0f090203d6f723cfc5e25303c", + "layer": "1.5", + "tier": "speciation_real" +} \ No newline at end of file diff --git a/results/speciation_real_cliff/resolved_config.yaml b/results/speciation_real_cliff/resolved_config.yaml new file mode 100644 index 0000000..b8ec6f1 --- /dev/null +++ b/results/speciation_real_cliff/resolved_config.yaml @@ -0,0 +1,29 @@ +experiment: speciation_real_cliff +seed: 13 +n_replicates: 3 +source_config: + experiment: speciation_real_cliff + kind: speciation_real + seed: 13 + n_replicates: 3 + speciation_real: + sizes: + - 784 + - 512 + - 512 + - 10 + conflict_fracs: + - 0.0 + - 0.2 + - 0.4 + - 0.6 + - 0.8 + - 1.0 + t_div_fixed: 800 + base_steps: 500 + lr: 0.05 + batch: 128 + n_eval: 2000 + data_root: data + output: + dir: results/speciation_real_cliff diff --git a/src/neural/experiment.py b/src/neural/experiment.py index cc33547..542e2c9 100644 --- a/src/neural/experiment.py +++ b/src/neural/experiment.py @@ -188,6 +188,11 @@ def run_and_save(config_path: str | Path) -> Path: elif kind == "mnist_lineage": df, extra_manifest = run_mnist_experiment(cfg) # oracle provenance + confusion matrix grid = [{"label": label, "lineage_cfg": c} for label, c in _expand_mnist(cfg)] + elif kind == "speciation_real": + from .speciation_real import run_speciation_real # E13: real-weight speciation + Git Re-Basin + df = run_speciation_real(cfg, int(cfg["seed"])) + grid = None + extra_manifest = {"layer": "1.5", "tier": "speciation_real"} else: raise ValueError(f"unknown neural experiment kind {kind!r}") save_artifacts(cfg, df, out_dir, extra_libs=_EXTRA_LIBS, diff --git a/src/neural/rebasin.py b/src/neural/rebasin.py new file mode 100644 index 0000000..a6e4e67 --- /dev/null +++ b/src/neural/rebasin.py @@ -0,0 +1,110 @@ +"""Git Re-Basin weight-matching + linear-mode-connectivity barrier for small MLPs (in-house, deterministic). + +The real-weight speciation experiment (E13) needs to separate the part of two models' merge-incompatibility +that is a mere *coordinate artefact* (removable by permuting hidden units — Ainsworth et al., Git Re-Basin, +arXiv:2209.04836) from the *residual* incompatibility that permutation cannot fix (the true reproductive- +isolation / Dobzhansky–Muller signal). This module implements, for a plain MLP with no BatchNorm/residuals +(the clean Re-Basin regime): + +- ``weight_matching`` — align model B's hidden units to model A by per-layer coordinate descent, each layer + an exact linear-assignment problem (``scipy.optimize.linear_sum_assignment``). Deterministic given the two + weight sets and a seeded layer-visiting order; no data, no gradient — so the "coordinate artefact" claim is + not confounded with data. +- ``barrier`` — the linear-mode-connectivity loss/error barrier along the interpolation between two models. + +Weights are handled as a list of ``(W, b)`` numpy arrays, one per ``nn.Linear`` (``W`` is ``[out, in]``). +Permutations act on the *hidden* layers (outputs of all but the last linear), permuting rows of layer ``k`` +(and its bias) and columns of layer ``k+1``. +""" + +from __future__ import annotations + +import numpy as np +from scipy.optimize import linear_sum_assignment + + +def weight_matching(params_a: list, params_b: list, rng: np.random.Generator, + max_iter: int = 30) -> list: + """Permutations aligning B's hidden units to A (Git Re-Basin weight matching). + + Args: + params_a, params_b (list[tuple[np.ndarray, np.ndarray]]): ``(W, b)`` per linear layer; ``W`` is + ``[out, in]``. Both models must share architecture. + rng (np.random.Generator): seeds the (order-only) layer-visiting schedule — deterministic result. + max_iter (int): coordinate-descent sweeps; stops early at a fixed point. + + Returns: + list[np.ndarray]: permutation index arrays, one per hidden layer (length = #linear − 1). + """ + n_layers = len(params_a) + hidden = n_layers - 1 # permutable layers (not the output) + sizes = [params_a[k][0].shape[0] for k in range(hidden)] + perms = [np.arange(s) for s in sizes] # identity to start + + for _ in range(max_iter): + changed = False + for k in rng.permutation(hidden): + Wa_in, _ = params_a[k] + Wb_in, _ = params_b[k] + # incoming weights: B's columns already permuted by the previous hidden layer's perm + if k > 0: + Wb_in = Wb_in[:, perms[k - 1]] + cost = Wa_in @ Wb_in.T # (out_k, out_k): align A row i with B row j + # outgoing weights: layer k+1 columns correspond to these units; B's rows permuted by next perm + Wa_out = params_a[k + 1][0] + Wb_out = params_b[k + 1][0] + if k + 1 < hidden: + Wb_out = Wb_out[perms[k + 1], :] + cost = cost + Wa_out.T @ Wb_out # (in_{k+1}=out_k, out_k) + ri, ci = linear_sum_assignment(cost, maximize=True) + new_perm = ci[np.argsort(ri)] + if not np.array_equal(new_perm, perms[k]): + changed = True + perms[k] = new_perm + if not changed: + break + return perms + + +def apply_perms(params: list, perms: list) -> list: + """Return a copy of ``params`` with hidden-unit permutations applied (rows of k, columns of k+1).""" + out = [(W.copy(), b.copy()) for W, b in params] + hidden = len(params) - 1 + for k in range(hidden): + p = perms[k] + W, b = out[k] + out[k] = (W[p, :], b[p]) # permute this layer's output units + Wn, bn = out[k + 1] + out[k + 1] = (Wn[:, p], bn) # and the next layer's matching inputs + return out + + +def interpolate(params_a: list, params_b: list, alpha: float) -> list: + """Linear interpolation ``(1-alpha)*A + alpha*B`` of two parameter lists.""" + return [((1 - alpha) * Wa + alpha * Wb, (1 - alpha) * ba + alpha * bb) + for (Wa, ba), (Wb, bb) in zip(params_a, params_b)] + + +def barrier(params_a: list, params_b: list, loss_fn, alphas=None) -> dict: + """Linear-mode-connectivity barrier: worst-case excess loss on the interpolation path above the chord. + + Args: + params_a, params_b (list): endpoint parameter lists. + loss_fn (Callable): ``loss_fn(params) -> (loss, error)`` on a fixed eval set. + alphas (Sequence[float] | None): interpolation grid (default 21 points on [0, 1]). + + Returns: + dict: ``{loss_barrier, error_barrier, midpoint_loss_barrier}`` (excess over the endpoint chord). + """ + alphas = np.linspace(0, 1, 21) if alphas is None else np.asarray(alphas) + losses, errors = [], [] + for a in alphas: + L, E = loss_fn(interpolate(params_a, params_b, float(a))) + losses.append(L); errors.append(E) + losses, errors = np.array(losses), np.array(errors) + chord_L = (1 - alphas) * losses[0] + alphas * losses[-1] + chord_E = (1 - alphas) * errors[0] + alphas * errors[-1] + mid = len(alphas) // 2 + return {"loss_barrier": float(np.max(losses - chord_L)), + "error_barrier": float(np.max(errors - chord_E)), + "midpoint_loss_barrier": float(losses[mid] - chord_L[mid])} diff --git a/src/neural/speciation_real.py b/src/neural/speciation_real.py new file mode 100644 index 0000000..68bf7bf --- /dev/null +++ b/src/neural/speciation_real.py @@ -0,0 +1,165 @@ +"""E13 — real-weight model speciation: the merge-compatibility cliff in trained MLPs, with Git Re-Basin. + +The real-weight image of E12. Two small MLPs are forked from a shared base and trained independently; +we merge them (weight averaging) and measure the linear-mode-connectivity **barrier** — before and after +**permutation alignment** (:mod:`neural.rebasin`). The barrier alignment *removes* is a coordinate +artefact (Git Re-Basin); the barrier it *cannot* remove — the **residual** — is the true reproductive- +isolation / Dobzhansky–Muller signal. A ``condition`` knob sets how the children diverge, which is the +empirical question this experiment answers: + +- ``shared`` — both children keep training on the *same* task from the shared fork: they never leave the + basin, so there is ~no barrier at all (trivially mergeable — the low-divergence anchor). +- ``independent`` — same task, but each child trained from its *own random init* (the canonical Git + Re-Basin setting): a large naive barrier that alignment *removes* (residual ≈ 0) — the coordinate + artefact. "Same species, different basis." +- ``conflict`` — the children learn *conflicting* label maps (B's labels cyclically shifted): + genuinely incompatible functions on shared capacity. A large barrier that alignment *cannot* remove + (residual stays high) — true reproductive isolation. "Different species." + +The discriminating metric is the **residual** (barrier after alignment): ~0 for ``shared`` and +``independent`` (compatible — the incompatibility, if any, is coordinate), large for ``conflict``. + +Divergence is swept via post-fork training steps ``t_div``. Small no-BatchNorm MLPs on MNIST — the clean +Re-Basin regime. Statistically reproducible (seeded); NumPy/scipy alignment is deterministic. +""" + +from __future__ import annotations + +from typing import Any, Mapping + +import numpy as np +import pandas as pd + +from .rebasin import apply_perms, barrier, weight_matching +from .train import seed_everything + + +def _mlp(sizes, device): + import torch.nn as nn + layers = [] + for i in range(len(sizes) - 1): + layers.append(nn.Linear(sizes[i], sizes[i + 1])) + if i < len(sizes) - 2: + layers.append(nn.ReLU()) + return nn.Sequential(*layers).to(device) + + +def _get_params(model): + import torch.nn as nn + return [(m.weight.detach().cpu().numpy().copy(), m.bias.detach().cpu().numpy().copy()) + for m in model if isinstance(m, nn.Linear)] + + +def _set_params(model, params, device): + import torch + import torch.nn as nn + it = iter(params) + for m in model: + if isinstance(m, nn.Linear): + W, b = next(it) + m.weight.data = torch.tensor(W, dtype=torch.float32, device=device) + m.bias.data = torch.tensor(b, dtype=torch.float32, device=device) + + +def _train(model, X, y, steps, lr, batch, rng, device): + import torch + opt = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9) + lossf = torch.nn.CrossEntropyLoss() + model.train() + for _ in range(steps): + idx = rng.integers(0, len(X), size=batch) + xb = X[idx]; yb = y[idx] + opt.zero_grad(); loss = lossf(model(xb), yb); loss.backward(); opt.step() + + +def run_speciation_real(cfg: Mapping[str, Any], seed: int) -> pd.DataFrame: + """Fork-and-merge sweep over (condition x divergence x replicate); returns barrier decompositions.""" + import torch + import torchvision + + spec = dict(cfg["speciation_real"]) + sizes = list(spec.get("sizes", [784, 512, 512, 10])) + conditions = list(spec.get("conditions", ["shared", "independent", "conflict"])) + t_divs = list(spec.get("t_div", [50, 100, 200, 400, 800])) + base_steps = int(spec.get("base_steps", 300)) + lr, batch = float(spec.get("lr", 0.05)), int(spec.get("batch", 128)) + n_eval = int(spec.get("n_eval", 2000)) + reps = int(cfg.get("n_replicates", spec.get("reps", 3))) + device = "cuda" if __import__("torch").cuda.is_available() else "cpu" + + root = spec.get("data_root", "data") + tr = torchvision.datasets.MNIST(root, train=True, download=True) + te = torchvision.datasets.MNIST(root, train=False, download=True) + Xtr = (tr.data.float().reshape(-1, 784) / 255.0).to(device); ytr = tr.targets.to(device) + Xte = (te.data.float().reshape(-1, 784) / 255.0)[:n_eval].to(device); yte = te.targets[:n_eval].to(device) + lossf = torch.nn.CrossEntropyLoss() + + def make_loss_fn(model): + def loss_fn(params): + _set_params(model, params, device); model.eval() + with torch.no_grad(): + out = model(Xte) + L = float(lossf(out, yte)); E = float((out.argmax(1) != yte).float().mean()) + return L, E + return loss_fn + + def _conflict_labels(frac): + """B's labels: cyclically shift the first round(frac*10) classes (systematic conflict on them).""" + k = int(round(frac * 10)) + if k == 0: + return ytr + sel = np.arange(k); shifted = np.roll(sel, 1) + y2 = ytr.clone() + for c, c2 in zip(sel, shifted): + y2[ytr == int(c)] = int(c2) + return y2 + + def subset(cond, child, conflict_frac): + """(X, y) the child trains on for a given condition.""" + if cond == "conflict" and child == 1: + return Xtr, _conflict_labels(conflict_frac) # B learns a conflicting label map + return Xtr, ytr # shared / independent / conflict-A: normal task + + # Two modes: (1) conditions x t_div decomposition; (2) a conflict-fraction isolation cliff. + conflict_fracs = spec.get("conflict_fracs") + if conflict_fracs is not None: + sweep = [("conflict", int(spec.get("t_div_fixed", 800)), float(f)) for f in conflict_fracs] + else: + sweep = [(c, int(t), 1.0) for c in conditions for t in t_divs] + + rows: list[dict] = [] + for rep in range(reps): + for cond, t_div, conflict_frac in sweep: + ss = seed + 1000 * rep + hash((cond, t_div, conflict_frac)) % 997 + seed_everything(np.random.SeedSequence(ss)) + base = _mlp(sizes, device) + _train(base, Xtr, ytr, base_steps, lr, batch, np.random.default_rng(ss), device) + base_params = _get_params(base) + children = [] + for child in (0, 1): + Xc, yc = subset(cond, child, conflict_frac) + if cond == "independent": # each child from its OWN init (Git Re-Basin regime) + seed_everything(np.random.SeedSequence(ss + 100 * (child + 1))) + m = _mlp(sizes, device) + _train(m, Xc, yc, base_steps + t_div, lr, batch, + np.random.default_rng(ss + 17 * (child + 1)), device) + else: # shared fork, then independent divergence + m = _mlp(sizes, device); _set_params(m, base_params, device) + _train(m, Xc, yc, t_div, lr, batch, np.random.default_rng(ss + 17 * (child + 1)), device) + children.append(_get_params(m)) + pA, pB = children + probe = _mlp(sizes, device); loss_fn = make_loss_fn(probe) + b_naive = barrier(pA, pB, loss_fn) + perms = weight_matching(pA, pB, np.random.default_rng(ss + 5)) + b_aligned = barrier(pA, apply_perms(pB, perms), loss_fn) + accA, accB = 1 - loss_fn(pA)[1], 1 - loss_fn(pB)[1] + rows.append({ + "condition": cond, "t_div": t_div, "conflict_frac": conflict_frac, "replicate": rep, + "barrier_naive": b_naive["error_barrier"], + "barrier_aligned": b_aligned["error_barrier"], + "removable": b_naive["error_barrier"] - b_aligned["error_barrier"], + "residual": b_aligned["error_barrier"], + "loss_barrier_naive": b_naive["loss_barrier"], + "loss_barrier_aligned": b_aligned["loss_barrier"], + "parent_acc": (accA + accB) / 2.0}) + return pd.DataFrame(rows) diff --git a/tests/test_rebasin.py b/tests/test_rebasin.py new file mode 100644 index 0000000..2605492 --- /dev/null +++ b/tests/test_rebasin.py @@ -0,0 +1,66 @@ +"""Tests for E13's Git Re-Basin weight-matching + barrier (pure NumPy, always runnable).""" + +from __future__ import annotations + +import numpy as np +import pytest + +from neural.rebasin import apply_perms, barrier, interpolate, weight_matching + + +def _mlp(sizes, rng): + return [(rng.standard_normal((sizes[i + 1], sizes[i])), rng.standard_normal(sizes[i + 1])) + for i in range(len(sizes) - 1)] + + +def _forward(params, X): + h = X + for i, (W, b) in enumerate(params): + h = h @ W.T + b + if i < len(params) - 1: + h = np.maximum(h, 0) + return h + + +def test_weight_matching_recovers_a_known_permutation(): + # A random-but-functionally-identical permuted copy is in a different "basis"; weight matching must + # recover the permutation, so realigning makes the copy functionally identical to the original again. + rng = np.random.default_rng(0) + A = _mlp([8, 16, 16, 3], rng) + p1, p2 = rng.permutation(16), rng.permutation(16) + B = apply_perms(A, [p1, p2]) # functionally identical to A, permuted basis + X = rng.standard_normal((32, 8)) + assert np.allclose(_forward(A, X), _forward(B, X)) # permutation preserves the function + perms = weight_matching(A, B, np.random.default_rng(1)) + B_realigned = apply_perms(B, perms) + assert np.allclose(_forward(A, X), _forward(B_realigned, X), atol=1e-6) # recovered -> function matches + + +def test_weight_matching_is_deterministic(): + rng = np.random.default_rng(2) + A, B = _mlp([6, 10, 4], rng), _mlp([6, 10, 4], rng) + p1 = weight_matching(A, B, np.random.default_rng(3)) + p2 = weight_matching(A, B, np.random.default_rng(3)) + assert all(np.array_equal(a, b) for a, b in zip(p1, p2)) # deterministic given inputs + seed + + +def test_interpolate_endpoints_and_barrier_zero_for_identical(): + rng = np.random.default_rng(4) + A = _mlp([5, 8, 2], rng) + B = _mlp([5, 8, 2], rng) + assert np.allclose(interpolate(A, B, 0.0)[0][0], A[0][0]) + assert np.allclose(interpolate(A, B, 1.0)[0][0], B[0][0]) + X = rng.standard_normal((20, 5)); tgt = rng.standard_normal((20, 2)) + + def loss_fn(P): + L = float(((_forward(P, X) - tgt) ** 2).mean()); return L, L + b = barrier(A, A, loss_fn) # a model with itself: no barrier + assert b["loss_barrier"] == pytest.approx(0.0, abs=1e-9) + + +def test_apply_perms_preserves_function(): + rng = np.random.default_rng(5) + A = _mlp([4, 7, 7, 3], rng) + perms = [rng.permutation(7), rng.permutation(7)] + X = rng.standard_normal((16, 4)) + assert np.allclose(_forward(A, X), _forward(apply_perms(A, perms), X))