E13: real-weight model speciation — the Git Re-Basin residual confirms E12
The real-weight image of E12, and the answer to the mode-connectivity reviewer. Small no-BN MLPs on MNIST, forked from a shared base and trained independently, are weight-averaged; we measure the linear-mode-connectivity barrier before and after in-house deterministic Git Re-Basin permutation alignment (neural/rebasin.py, scipy linear_sum_assignment), decomposing it into removable (coordinate artefact) and residual (reproductive isolation). kind: speciation_real. Result (3 reps): - shared (same task, shared fork): no barrier — trivially mergeable. - independent (same task, different init): naive 0.056, alignment removes 98% (residual 0.001) — the incompatibility is a coordinate artefact. - conflict (conflicting label maps): naive 0.496, alignment removes 0% (residual 0.496) — genuine reproductive isolation. Because alignment demonstrably works on the independent case, the conflict residual is real, not a failure to align. - Isolation cliff (speciation_real_cliff): residual rises 0.00->0.13->0.19->0.28-> 0.40->0.49 with the fraction of conflicting classes — the real-weight mirror of E12's cliff; residual==naive throughout (functional, not coordinate). rebasin.py sanity-gated (recovers a known permutation exactly). plot_speciation_real.py (2-panel), +4 pure-NumPy tests (142 green), README with honest positioning vs Git Re-Basin / Entezari / Frankle / Pari 2024 / Zhou 2026. Wired into make mnist (needs torchvision). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
parent
56f642e7f9
commit
01d87e504f
15 changed files with 620 additions and 3 deletions
8
Makefile
8
Makefile
|
|
@ -18,12 +18,14 @@ test: ## correctness tests + scientific-validation tests (the spine
|
||||||
layer1: ## run experiments E1-E6 + the learning-kernel bridge (analytic)
|
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
|
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)
|
neural: ## run Layer 1.5 synthetic neural experiments (excludes the MNIST/torchvision tiers)
|
||||||
for c in configs/neural/*.yaml; do case "$$c" in *mnist*) ;; \
|
for c in configs/neural/*.yaml; do case "$$c" in *mnist*|*speciation_real*) ;; \
|
||||||
*) uv run python -m neural.experiment "$$c" ;; esac; done
|
*) 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/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)
|
env-llm: ## add the LLM stack for the Layer-2 prototype (GPU; transformers/peft)
|
||||||
uv sync --extra dev --extra neural --extra llm
|
uv sync --extra dev --extra neural --extra llm
|
||||||
|
|
|
||||||
28
configs/neural/speciation_real.yaml
Normal file
28
configs/neural/speciation_real.yaml
Normal file
|
|
@ -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
|
||||||
24
configs/neural/speciation_real_cliff.yaml
Normal file
24
configs/neural/speciation_real_cliff.yaml
Normal file
|
|
@ -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
|
||||||
67
figures/plot_speciation_real.py
Normal file
67
figures/plot_speciation_real.py
Normal file
|
|
@ -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()
|
||||||
54
results/speciation_real/README.md
Normal file
54
results/speciation_real/README.md
Normal file
|
|
@ -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.
|
||||||
18
results/speciation_real/manifest.json
Normal file
18
results/speciation_real/manifest.json
Normal file
|
|
@ -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"
|
||||||
|
}
|
||||||
31
results/speciation_real/resolved_config.yaml
Normal file
31
results/speciation_real/resolved_config.yaml
Normal file
|
|
@ -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
|
||||||
BIN
results/speciation_real/speciation_real.pdf
Normal file
BIN
results/speciation_real/speciation_real.pdf
Normal file
Binary file not shown.
BIN
results/speciation_real/speciation_real.png
Normal file
BIN
results/speciation_real/speciation_real.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 136 KiB |
18
results/speciation_real_cliff/manifest.json
Normal file
18
results/speciation_real_cliff/manifest.json
Normal file
|
|
@ -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"
|
||||||
|
}
|
||||||
29
results/speciation_real_cliff/resolved_config.yaml
Normal file
29
results/speciation_real_cliff/resolved_config.yaml
Normal file
|
|
@ -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
|
||||||
|
|
@ -188,6 +188,11 @@ def run_and_save(config_path: str | Path) -> Path:
|
||||||
elif kind == "mnist_lineage":
|
elif kind == "mnist_lineage":
|
||||||
df, extra_manifest = run_mnist_experiment(cfg) # oracle provenance + confusion matrix
|
df, extra_manifest = run_mnist_experiment(cfg) # oracle provenance + confusion matrix
|
||||||
grid = [{"label": label, "lineage_cfg": c} for label, c in _expand_mnist(cfg)]
|
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:
|
else:
|
||||||
raise ValueError(f"unknown neural experiment kind {kind!r}")
|
raise ValueError(f"unknown neural experiment kind {kind!r}")
|
||||||
save_artifacts(cfg, df, out_dir, extra_libs=_EXTRA_LIBS,
|
save_artifacts(cfg, df, out_dir, extra_libs=_EXTRA_LIBS,
|
||||||
|
|
|
||||||
110
src/neural/rebasin.py
Normal file
110
src/neural/rebasin.py
Normal file
|
|
@ -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])}
|
||||||
165
src/neural/speciation_real.py
Normal file
165
src/neural/speciation_real.py
Normal file
|
|
@ -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)
|
||||||
66
tests/test_rebasin.py
Normal file
66
tests/test_rebasin.py
Normal file
|
|
@ -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))
|
||||||
Loading…
Add table
Add a link
Reference in a new issue