Giorgio Gilestro
01d87e504f
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>
66 lines
2.6 KiB
Python
66 lines
2.6 KiB
Python
"""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))
|