MachineSex/figures/plot_collapse.py
Giorgio Gilestro ab3dc10587 Restructure: descriptive tier and experiment names, paper/manuscript
- paper/pnas -> paper/manuscript (venue-neutral)
- configs/layer1 -> configs/inheritance, src/knowledge -> src/inheritance
  (imported as `inheritance`), make layer1 -> make inheritance; layer2 alias dropped
- inheritance and trained-network bundles named after the manuscript figure
  they feed (fig2_grounding_sweep, figS3_rebaselining, ...), or descriptively
  where they feed none; configs keep their `experiment:` value so parquet
  hashes are unchanged, only output.dir moves
- figure scripts, SI figure sources, notebooks, REPRODUCING.md, README and the
  SI Methods/tables updated; make clean no longer deletes tracked manifests;
  reproduce.sh hashes the s{seed}/ layouts too

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Y64o8FKP7rCuXzC48pxpMm
2026-09-13 17:00:40 +01:00

84 lines
3.5 KiB
Python

"""`collapse` figure — model collapse in REAL RNN weights, arrested by grounding (↔ E1/C1).
The existence proof: a GRU trained each generation on the previous generation's own samples
loses the rare tail and drifts from truth (forward-KL climbs), and even a little grounding
arrests it. Forward-KL is the operative neural collapse metric (the RNN's smoothing keeps
spurious tail support alive, so H barely moves — see the `grounding` finding).
Four panels: (A) forward-KL trajectories (dry climbs, grounded suppressed); (B) H trajectories
(barely moves — smoothing resists H-collapse); (C) stationary forward-KL vs g; (D) tail-item
survival vs g. Reads only the committed bundle.
Usage: python figures/plot_collapse.py [results/collapse]
"""
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, mean_ci, savefig # noqa: E402
sys.path.insert(0, str(Path(__file__).parents[1] / "src"))
from inheritance.metrics import heterozygosity # noqa: E402
from neural.config import SyntheticCfg # noqa: E402
from neural.synthetic import make_mode_truth # noqa: E402
def main(results_dir: str = "results/collapse") -> None:
df, cfg = load_bundle(results_dir)
syn = SyntheticCfg(**cfg["synthetic"])
H_star = heterozygosity(make_mode_truth(syn).p_star)
g_values = sorted(df["g"].unique())
last = int(cfg["generations"] * 0.6)
stat = df[df["generation"] >= last]
colors = plt.cm.viridis(np.linspace(0, 0.85, len(g_values)))
fig, axes = plt.subplots(2, 2, figsize=(13, 9))
# Panel A: forward-KL trajectories (dry climbs, grounded suppressed).
ax = axes[0, 0]
for g, c in zip(g_values, colors):
s = df[df["g"] == g].groupby("generation")["forward_kl"].mean()
ax.plot(s.index, s.values, "-o", color=c, ms=3, label=f"g={g:g}")
ax.set(xlabel="generation", ylabel=r"forward-KL $D(p^*\Vert\hat p)$",
title="Collapse in weights: dry KL climbs, grounding holds it")
ax.legend(frameon=False, fontsize=9)
# Panel B: H trajectories (barely moves — smoothing resists H-collapse).
ax = axes[0, 1]
for g, c in zip(g_values, colors):
s = df[df["g"] == g].groupby("generation")["heterozygosity"].mean()
ax.plot(s.index, s.values, "-o", color=c, ms=3, label=f"g={g:g}")
ax.axhline(H_star, ls=":", color="gray", lw=1, label="$H^*$")
ax.set(xlabel="generation", ylabel="heterozygosity $H$",
title="H barely moves (RNN smoothing resists H-collapse)")
ax.legend(frameon=False, fontsize=9)
# Panel C: stationary forward-KL vs g.
ax = axes[1, 0]
kg, Km, Kci = mean_ci(stat, "g", "forward_kl")
ax.errorbar(kg, Km, yerr=Kci, fmt="o-", color="#1f77b4", capsize=3)
ax.set(xlabel="grounding fraction $g$", ylabel=r"stationary forward-KL",
title="Grounding lowers stationary divergence")
# Panel D: tail-item survival vs g.
ax = axes[1, 1]
tg, Tm, Tci = mean_ci(stat, "g", "tail_frac_alive")
ax.errorbar(tg, Tm, yerr=Tci, fmt="s-", color="#d62728", capsize=3)
ax.set(xlabel="grounding fraction $g$", ylabel="tail items alive",
title="Grounding lifts tail survival")
fig.suptitle("collapse — a trained GRU collapses under dry self-training; grounding arrests it "
f"($K$={syn.K}, $n$={cfg['dynamics']['n']})", y=1.0, fontsize=13)
fig.tight_layout()
savefig(fig, results_dir, "collapse")
if __name__ == "__main__":
main(*sys.argv[1:])