recombination: reproduce the E4 "merge, don't average" finding in real weights

src/neural/recombine.py mirrors Layer-1 run_coverage but trains K_T specialist RNNs on
assignments from the exact shared-switch retention construction (K_T/rho/q clean; union
matches the closed form), then recombines the measured teacher distributions two ways:
mean (naive pooling) vs oracle-guided max-merge (per-mode strongest teacher, M2N2-style),
each followed by size-n resampling.

Result (8 reps): at rho=0, union rises 0.49->0.96 (supply matches closed form); analytic
surviving_max rises 0.043->0.087 while surviving_mean stays flat ~0.045 — the conservation
law (averaging cancels the union gain, max-merge realises it). At rho=1 (identical
teachers) union and max are flat. The lesson holds in the neural setting; trained-weight
columns show the same signs but noisier (smoothing inflates baseline; deep tail barely
clears n=200 resampling). torch-gated test added. 93 tests green.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Giorgio Gilestro 2026-07-04 21:49:44 +01:00
parent aca7b394a3
commit d22dd9d535
7 changed files with 263 additions and 7 deletions

110
src/neural/recombine.py Normal file
View file

@ -0,0 +1,110 @@
"""Multi-teacher recombination in real weights — the neural image of Layer-1 E4.
Layer-1 E4 (``knowledge.experiment.run_coverage``) showed the sharpest result of the study:
under **mean-mixture** distillation surviving tail coverage is *flat* in the teacher count
K_T (a conservation law averaging's 1/K_T dilution cancels the union gain), while a
**union-preserving max-merge** (à la M2N2) realises the benefit. This module tests whether
that survives when the teachers are *trained generative models* rather than analytic
distributions.
Faithful parallel to ``run_coverage``:
* teacher *assignments* come from the exact shared-switch retention construction
(``make_retention_matrix`` / ``make_correlated_teachers``), so K_T, rho and q are clean
knobs and the construction-level ``union_coverage`` matches the closed form exactly;
* each teacher is then **trained** on samples from its assigned distribution (real weights);
* the pupil's recombination is applied to the *measured* teacher distributions p_hat_k:
``mean`` (pool teacher outputs the naive distillation null) vs ``max`` (oracle-guided
union keep each mode's strongest teacher, the M2N2-style merge), each followed by the
pupil's size-n resampling. Surviving tail coverage under the two operators is the result.
"""
from __future__ import annotations
import itertools
import numpy as np
import pandas as pd
from knowledge.config import _sub
from knowledge.seeding import spawn_seeds
from knowledge.teachers import make_correlated_teachers, make_retention_matrix
from .config import ModelCfg, SyntheticCfg
from .models import make_model
from .oracle import ExactOracle
from .synthetic import make_mode_truth
def run_recombination(cfg: dict) -> pd.DataFrame:
"""Train K_T specialist models per grid point and compare mean vs max-merge coverage.
Args:
cfg (dict): Parsed experiment YAML with ``synthetic``, ``model``, a ``coverage``
block (``n`` resample size, ``q`` marginal retention, ``retain_thresh``,
optional ``region_specialisation``), a ``sweep`` (K_T x rho), ``seed`` and
``n_replicates``.
Returns:
pd.DataFrame: One row per (K_T, rho, replicate) with ``union_coverage``,
``surviving_mean``, ``surviving_max`` (trained teachers), plus the analytic
``surviving_mean_target`` / ``surviving_max_target`` from the untrained assignment
distributions as an E4 cross-check.
"""
syn = _sub(cfg["synthetic"], SyntheticCfg)
model_cfg = _sub(cfg["model"], ModelCfg)
cov = cfg["coverage"]
n, q = int(cov["n"]), float(cov["q"])
retain_thresh = float(cov.get("retain_thresh", 1e-3))
region_spec = bool(cov.get("region_specialisation", False))
td = make_mode_truth(syn)
tail_idx = np.flatnonzero(td.tail_mask)
T = tail_idx.size
oracle = ExactOracle(syn)
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"]))
def surviving(p_over_modes: np.ndarray, rng: np.random.Generator) -> float:
p = p_over_modes / p_over_modes.sum()
counts = rng.multinomial(n, p)
return float(np.mean(counts[tail_idx] > 0))
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"])
for rep, ss in enumerate(seeds):
child = int(ss.generate_state(1)[0])
rng = np.random.default_rng(child)
# (1) construction-level supply: exact union from the retention matrix.
R = make_retention_matrix(T, K_T, rho, q, np.random.default_rng(child + 1))
union = float(np.mean(R.any(axis=0)))
# (2) assigned teacher distributions (same construction), then TRAIN each.
targets = make_correlated_teachers(
td.p_star, td.tail_mask, K_T, rho, q,
region_assignment=td.regions, region_specialisation=region_spec, seed=child)
p_hats = np.empty((K_T, syn.K))
for k, target in enumerate(targets):
model = make_model(model_cfg, syn, oracle)
model.initialise(np.asarray(target), rng)
p_hats[k] = model.mode_distribution(rng)
targets = np.asarray(targets)
# (3) mean vs oracle-guided max-merge, on trained and on analytic teachers.
rows.append({
"experiment": cfg["experiment"], "K_T": K_T, "rho": rho, "q": q,
"replicate": rep, "union_coverage": union,
"surviving_mean": surviving(p_hats.mean(axis=0), rng),
"surviving_max": surviving(p_hats.max(axis=0), rng),
"surviving_mean_target": surviving(targets.mean(axis=0), rng),
"surviving_max_target": surviving(targets.max(axis=0), rng),
})
return pd.DataFrame(rows)