giorgio/MachineSex
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

epistasis_predicts: functional conflict, not weight geometry, predicts merge failure pre-merge

Browse source 
The decisive experiment from the external review. 39 LoRA parent pairs
(0.5B, 3 seeds) on three axes decorrelated by construction: conflict
(contradictory conventions on shared prompts, private budgets fixed),
compat (same prompts, SAME convention — overlap without conflict), and
duration (weight divergence, zero conflict). Six pre-merge predictors;
primary outcome = merge penalty (parent potential − merged achieved).

League table (Spearman vs penalty, n=39): functional measures predict
(dis_raw +0.460, epi_conf +0.446, p<0.005); geometry collapses
(delta_cos +0.03, delta_l2 +0.17 n.s.); gradient alignment weak (−0.35);
performance ~0. The first grid's apparent geometry win (+0.60) was an
overlap/volume artifact — the compat control axis (added for exactly
this) exposed and killed it: same overlap and data volume, zero penalty.
Honest riders in the README: confidence weighting does not beat raw
disagreement as a rank predictor (pre-registered internal prediction not
confirmed; it does double the conflict/compat level contrast), and
|rho|~0.45 is bounded by 0.5B merge-outcome noise (7B is the firm-up).

Also: micro-batched gradient accumulation (OOM fix on the shared 16GB
GPU), exact r-space LoRA-delta geometry (brute-force-verified test,
151 green), systemd-run runbook lesson (tmux dies with the SSH session
scope on this box).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BkRLcc18rwT2Lysu6PbG7v
This commit is contained in:
Giorgio Gilestro 2026-09-06 17:19:46 +01:00
parent 5a23ddaf2a
commit 287d2326cc
16 changed files with 622 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

251
src/llm/epistasis.py Normal file
View file

@ -0,0 +1,251 @@
"""The decisive experiment: does an operational EPISTASIS measure predict merge success PRE-merge?
The external review's bar (2026-08-11): population-genetic quantities must *predict*, not re-describe —
forecast merge success **before merging**, and beat existing predictors. This module builds a grid of
parent pairs in which functional conflict and weight divergence are **decorrelated by construction**
(the conflict knob adds contradictory-convention data at fixed private budget; the duration knob grows
divergence with zero conflict), computes a battery of **pre-merge predictors**, then merges and
measures the outcome.
The theory-derived predictor — and its built-in ablation. A raw functional-disagreement rate between
two parents is, by our own theory, the WRONG measure: disjoint specialists disagree hugely (one knows,
one is ignorant) yet merge perfectly — harmless *complementation*. The Dobzhansky–Muller analogue is
**bilateral confident contradiction**: both parents confidently produce different answers to the same
input. So the operational epistasis measure is confidence-weighted disagreement,
epi_conf = E_probe[ conf_A · conf_B · 1(ans_A != ans_B) ], conf = exp(mean token logprob),
and the theory makes an internal, falsifiable prediction: ``epi_conf`` should predict merge failure
where raw disagreement (``dis_raw``) mispredicts. Baseline predictors from the ML literature:
**gradient alignment** at the shared base (cf. arXiv:2601.22285), **LoRA-delta geometry** (cosine and
L2 distance, computed exactly in r×r space), and a **performance-based** predictor (cross-family
accuracy). Primary outcome (pre-registered): ``merge_penalty`` = parent-potential − merged-achieved,
the hybrid-load analogue. Pre-registered falsifier: if gradient/geometry predictors match or beat
``epi_conf``, the paper's "epistasis, not divergence, sets the cliff" claim stays analytic-only.
"""
from __future__ import annotations
from pathlib import Path
import numpy as np
import pandas as pd
from .evaluate import generate, load_model
from .merge import load_specialists, make_merge
from .specialise import _encode, train_lora_on_tasks
from .speciation import make_convention_tasks
from .tasks import _normalise, make_tasks, verify
# ---------------------------------------------------------------- pre-merge predictors
def generate_with_confidence(model, tok, prompts, max_new_tokens=24, batch_size=32, device="cuda"):
"""Greedy batched generation returning (answers, confidences).
Confidence = exp(mean logprob of the generated tokens up to the first EOS) — a per-answer
self-certainty in [0, 1], measurable without any ground truth.
"""
import torch
outs, confs = [], []
for i in range(0, len(prompts), batch_size):
chunk = prompts[i:i + batch_size]
msgs = [[{"role": "user", "content": p}] for p in chunk]
enc = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt",
return_dict=True, padding=True).to(device)
with torch.no_grad():
gen = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False,
pad_token_id=tok.pad_token_id, output_scores=True,
return_dict_in_generate=True)
L = enc["input_ids"].shape[1]
seq = gen.sequences
for j in range(len(chunk)):
toks = seq[j, L:]
lps = []
for t, tid in enumerate(toks):
if int(tid) in (tok.eos_token_id, tok.pad_token_id):
break
logp = torch.log_softmax(gen.scores[t][j].float(), dim=-1)[int(tid)]
lps.append(float(logp))
outs.append(tok.decode(toks, skip_special_tokens=True).strip())
confs.append(float(np.exp(np.mean(lps))) if lps else 0.0)
return outs, np.array(confs)
def lora_delta_inner(A1, B1, A2, B2) -> float:
"""Frobenius inner product <B1@A1, B2@A2> computed in r×r space (never materialises the deltas)."""
import torch
return float(torch.trace((B1.T @ B2) @ (A2 @ A1.T)))
def delta_geometry(dir_a: str, dir_b: str) -> dict:
"""Cosine similarity and L2 distance between two LoRA adapters' weight deltas (CPU, exact)."""
import torch
from safetensors.torch import load_file
def pairs(d):
sd = load_file(str(Path(d) / "adapter_model.safetensors"))
mods = {}
for k, v in sd.items():
if "lora_A" in k:
mods.setdefault(k.replace("lora_A", "@"), {})["A"] = v.float()
elif "lora_B" in k:
mods.setdefault(k.replace("lora_B", "@"), {})["B"] = v.float()
return mods
ma, mb = pairs(dir_a), pairs(dir_b)
keys = sorted(set(ma) & set(mb))
i12 = sum(lora_delta_inner(ma[k]["A"], ma[k]["B"], mb[k]["A"], mb[k]["B"]) for k in keys)
i11 = sum(lora_delta_inner(ma[k]["A"], ma[k]["B"], ma[k]["A"], ma[k]["B"]) for k in keys)
i22 = sum(lora_delta_inner(mb[k]["A"], mb[k]["B"], mb[k]["A"], mb[k]["B"]) for k in keys)
n1, n2 = np.sqrt(max(i11, 1e-12)), np.sqrt(max(i22, 1e-12))
return {"delta_cos": i12 / (n1 * n2), "delta_l2": float(np.sqrt(max(i11 + i22 - 2 * i12, 0.0)))}
def gradient_alignment(base_name: str, tasks_a: list, tasks_b: list, k: int = 32,
device: str = "cuda", seed: int = 0) -> float:
"""Cosine between the two tasks' SFT gradients at the SHARED BASE (no adapters) — the
gradient-alignment predictor of the ML literature, computed on ``k`` examples per side."""
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
torch.cuda.empty_cache() # release the eval models' cache first
tok = AutoTokenizer.from_pretrained(base_name)
if tok.pad_token_id is None:
tok.pad_token = tok.eos_token
model = AutoModelForCausalLM.from_pretrained(base_name, dtype=torch.bfloat16).to(device)
targets = [p for n, p in model.named_parameters()
if any(t in n for t in ("q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj"))]
rng = np.random.default_rng(seed)
def grad_vec(tasks, micro=4):
# Micro-batched gradient accumulation: identical total gradient (losses weighted to a mean
# over the k examples), 1/8 the activation memory of one big backward — OOM-safe on a shared
# 16 GB GPU.
model.zero_grad(set_to_none=True)
sel = [tasks[i] for i in rng.choice(len(tasks), size=min(k, len(tasks)), replace=False)]
for s in range(0, len(sel), micro):
enc = [_encode(tok, t, device) for t in sel[s:s + micro]]
L = max(len(ids) for ids, _ in enc)
ii = torch.full((len(enc), L), tok.pad_token_id, dtype=torch.long)
ll = torch.full((len(enc), L), -100, dtype=torch.long)
am = torch.zeros((len(enc), L), dtype=torch.long)
for b, (ids, lab) in enumerate(enc):
ii[b, :len(ids)] = torch.tensor(ids); ll[b, :len(lab)] = torch.tensor(lab)
am[b, :len(ids)] = 1
out = model(input_ids=ii.to(device), attention_mask=am.to(device), labels=ll.to(device))
(out.loss * (len(enc) / len(sel))).backward()
return torch.cat([p.grad.detach().float().flatten().cpu() for p in targets])
ga, gb = grad_vec(tasks_a), grad_vec(tasks_b)
cos = float((ga @ gb) / (ga.norm() * gb.norm() + 1e-12))
del model
torch.cuda.empty_cache()
return cos
# ---------------------------------------------------------------- the experiment
def run_epistasis_experiment(cfg: dict) -> pd.DataFrame:
"""Grid of parent pairs (conflict × duration axes) → pre-merge predictors + merged outcomes."""
import torch
name = cfg["experiment"]
base = cfg["base_model"]
fam_a, fam_b = cfg.get("family_a", "strings"), cfg.get("family_b", "arith")
n_train, n_test = int(cfg.get("n_train", 400)), int(cfg.get("n_test", 80))
epochs = int(cfg.get("epochs", 3))
n_probe = int(cfg.get("n_probe_each", 30))
grad_k = int(cfg.get("grad_k", 32))
lora = cfg.get("lora", {})
r, alpha = int(lora.get("r", 16)), int(lora.get("alpha", 32))
seed = int(cfg["seed"])
root = Path(cfg.get("adapters_dir", "models/llm")) / "epistasis"
# Fixed evaluation + probe sets (identical across the grid and across seeds).
test_a = make_tasks(fam_a, n_test, seed=1000)
test_b = make_tasks(fam_b, n_test, seed=1001)
amb_asc = make_convention_tasks(n_test, seed=5000, convention="asc")
amb_desc = make_convention_tasks(n_test, seed=5000, convention="desc")
# Probes: a broad mix drawn WITHOUT knowledge of where the conflict lives (ambiguous + both
# private families), so the disagreement measures are not handed the answer.
probe_prompts = ([t.prompt for t in make_convention_tasks(n_probe, seed=7000, convention="asc")]
+ [t.prompt for t in make_tasks(fam_a, n_probe, seed=7001)]
+ [t.prompt for t in make_tasks(fam_b, n_probe, seed=7002)])
def acc(model, tok, tasks):
outs = generate(model, tok, [t.prompt for t in tasks])
return float(np.mean([verify(o, t) for o, t in zip(outs, tasks)]))
grid = ([("conflict", float(f)) for f in cfg.get("conflict_fracs", [])]
+ [("compat", float(f)) for f in cfg.get("compat_fracs", [])]
+ [("duration", int(d)) for d in cfg.get("durations", [])])
rows: list[dict] = []
for mode, x in grid:
if mode in ("conflict", "compat"): # add-design: private budget fixed
# "compat" is the control that separates functional from geometric predictors: the SAME
# shared ambiguous prompts, but BOTH children learn the SAME convention — task overlap
# (which inflates delta similarity) without any conflict. A geometry predictor that is
# really an overlap detector mispredicts here; a functional-conflict measure does not.
conv_b = "desc" if mode == "conflict" else "asc"
n_conv = int(round(x * n_train))
tasks_a = (make_tasks(fam_a, n_train, seed=seed)
+ make_convention_tasks(n_conv, seed=seed + 50, convention="asc"))
tasks_b = (make_tasks(fam_b, n_train, seed=seed + 1)
+ make_convention_tasks(n_conv, seed=seed + 50, convention=conv_b))
ep = epochs
else: # duration: pure private, epochs swept
tasks_a = make_tasks(fam_a, n_train, seed=seed)
tasks_b = make_tasks(fam_b, n_train, seed=seed + 1)
ep = int(x)
da = train_lora_on_tasks(base, tasks_a, str(root / "a"), epochs=ep, r=r, alpha=alpha,
seed=seed)
db = train_lora_on_tasks(base, tasks_b, str(root / "b"), epochs=ep, r=r, alpha=alpha,
seed=seed)
# --- pre-merge: per-parent probe answers/confidences + parent outcomes ------------------
par = {}
for d, label in ((da, "a"), (db, "b")):
m, tok = load_model(base, adapter_dir=d)
ans, conf = generate_with_confidence(m, tok, probe_prompts)
par[label] = {
"ans": [_normalise(o) for o in ans], "conf": conf,
"fam_a": acc(m, tok, test_a), "fam_b": acc(m, tok, test_b),
"coh": max(acc(m, tok, amb_asc), acc(m, tok, amb_desc)),
}
del m; torch.cuda.empty_cache()
neq = np.array([a != b for a, b in zip(par["a"]["ans"], par["b"]["ans"])], dtype=float)
dis_raw = float(neq.mean())
epi_conf = float(np.mean(par["a"]["conf"] * par["b"]["conf"] * neq))
geo = delta_geometry(da, db)
grad_cos = gradient_alignment(base, tasks_a, tasks_b, k=grad_k, seed=seed)
cross_perf = float(np.mean([par["a"]["fam_b"], par["b"]["fam_a"]]))
# --- merge and measure -------------------------------------------------------------------
model, tok = load_specialists(base, [da, db])
make_merge(model, 2, "soup", "soup")
mo = {"fam_a": acc(model, tok, test_a), "fam_b": acc(model, tok, test_b),
"coh": max(acc(model, tok, amb_asc), acc(model, tok, amb_desc))}
del model; torch.cuda.empty_cache()
potential = float(np.mean([par["a"]["fam_a"], par["b"]["fam_b"],
max(par["a"]["coh"], par["b"]["coh"])]))
achieved = float(np.mean([mo["fam_a"], mo["fam_b"], mo["coh"]]))
rows.append({
"experiment": name, "mode": mode, "x": float(x),
"epi_conf": epi_conf, "dis_raw": dis_raw, "grad_cos": grad_cos,
"delta_cos": geo["delta_cos"], "delta_l2": geo["delta_l2"], "cross_perf": cross_perf,
"merge_fam_a": mo["fam_a"], "merge_fam_b": mo["fam_b"], "merge_coh": mo["coh"],
"merged_private": float(np.mean([mo["fam_a"], mo["fam_b"]])),
"merged_overall": achieved,
"parent_potential": potential,
"merge_penalty": potential - achieved,
"route_private": float(np.mean([par["a"]["fam_a"], par["b"]["fam_b"]])),
"pa_fam_a": par["a"]["fam_a"], "pa_fam_b": par["a"]["fam_b"],
"pb_fam_a": par["b"]["fam_a"], "pb_fam_b": par["b"]["fam_b"],
"pa_coh": par["a"]["coh"], "pb_coh": par["b"]["coh"],
})
return pd.DataFrame(rows)

8
src/llm/experiment.py
View file

@ -263,8 +263,14 @@ def run_speciation_dispatch(cfg: dict) -> pd.DataFrame:
return run_speciation_experiment(cfg)
def run_epistasis_dispatch(cfg: dict) -> pd.DataFrame:
from .epistasis import run_epistasis_experiment # local import: torch-heavy
return run_epistasis_experiment(cfg)
_RUNNERS = {"llm_merge": run_merge_experiment, "llm_moe": run_moe_experiment,
"llm_directed": run_directed_experiment, "llm_speciation": run_speciation_dispatch}
"llm_directed": run_directed_experiment, "llm_speciation": run_speciation_dispatch,
"llm_epistasis": run_epistasis_dispatch}
def run_and_save(config_path: str | Path) -> Path: