neural: real-MNIST external-validity tier (collapse + grounding)

Confirms model collapse and its arrest by grounding on REAL images, not
just the synthetic sandbox. A conv VAE (the canonical generative-collapse
model) is retrained each generation on its own generated digits, with a
fraction g of fresh real MNIST mixed in. Modes = digit class x stroke-
thickness bin (K=30, Zipf, ~18 tail modes); the oracle is a frozen CNN +
deterministic thickness at 98.5% mode accuracy (30x30 confusion matrix
recorded in the manifest as the measurement-noise floor).

Result (4 reps): dry (g=0) collapses to a single mode -- forward-KL
0.5->18, support 30->1, tail 1.0->0.06, H->0 -- while 10% grounding holds
all 30 modes (KL~0.6, full tail, H~0.9). Signs, not magnitudes (blueprint
3.5); the exact synthetic oracle stays the quantitative anchor. The VAE
needs ~10% grounding vs the synthetic histogram's ~5%, consistent with the
grounding finding that trained nets need more than the exact operator.

Plugs into the existing data-agnostic contract (metrics/grounding/output
reused verbatim): mnist_data (thickness bins, class x thickness bijection,
MnistSampler), mnist_oracle (ClassifierOracle + confusion matrix),
mnist_vae (ConvVAEGenerator), mnist_loop (run_mnist_lineage), kind=
mnist_lineage dispatch, MnistCfg/OracleCfg. Figures: plot_mnist (parquet-
only) + mnist_montage (eyeball diagnostic showing digits degenerate to one
blurry mode). make mnist / make env-mnist, kept out of the make neural
loop. 99 tests green (+5 torchvision-gated).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Giorgio Gilestro 2026-07-05 09:19:36 +01:00
parent 3b9f4f7893
commit 79bbc45f41
21 changed files with 2200 additions and 10 deletions

1
.gitignore vendored
View file

@ -6,3 +6,4 @@ __pycache__/
.pytest_cache/ .pytest_cache/
results/**/results.parquet results/**/results.parquet
models/ models/
data/

View file

@ -16,9 +16,11 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
**Done:** scaffold, the histogram bridge gate (reproduces Layer 1 exactly), `bridge` (neural **Done:** scaffold, the histogram bridge gate (reproduces Layer 1 exactly), `bridge` (neural
g*=0.047 ≈ Layer 1), `collapse` (in RNN weights), `grounding` (refined; sign confirmed, threshold g*=0.047 ≈ Layer 1), `collapse` (in RNN weights), `grounding` (refined; sign confirmed, threshold
softened by neural smoothing — see finding below), `architectures` (architecture-generality), softened by neural smoothing — see finding below), `architectures` (architecture-generality),
`recombination` (the E4 "merge, don't average" finding reproduced in real weights), and all five `recombination` (the E4 "merge, don't average" finding reproduced in real weights), all six
neural figures (`figures/plot_{bridge,collapse,grounding,architectures,recombination}.py`, wired neural figures, and the **real-MNIST external-validity tier** (`mnist_collapse`: a conv-VAE
into `make figures`). **Remaining:** `region_matched`, `remint`, the MNIST tier, VAE fidelity. The LLM/LoRA rung and the C3 vertical claim are deferred. Experiments are collapses to a single mode under dry self-training, ~10% grounding holds all 30 modes; frozen-CNN
oracle, confusion matrix recorded). **Remaining:** `region_matched`, `remint`, the synthetic-VAE
fidelity fix — all optional. The LLM/LoRA rung and C3 vertical claim stay deferred. The LLM/LoRA rung and the C3 vertical claim are deferred. Experiments are
named descriptively (`configs/neural/<name>.yaml`), not by code. named descriptively (`configs/neural/<name>.yaml`), not by code.
The two design documents are the source of truth for intent: The two design documents are the source of truth for intent:
@ -72,6 +74,8 @@ E4's whole purpose is to isolate the effect of teacher **decorrelation ρ**, so
**Finding (2026-07-05, neural `grounding`) — grounding's SIGN transfers to trained RNN weights, but the sharp `g*` does not; and tail-survival is the *wrong* neural collapse metric.** Re-ran the phase-boundary sweep at 18 replicates. Two results: (1) **forward-KL is the operative neural collapse metric, not H or tail-survival.** The RNN's smoothing inductive bias keeps *spurious* tail modes alive (it generalises to unseen codewords), so `tail_truth_mass_alive` is flat/**non-monotone** in g (dry 0.54 > most grounded points) and H stays ~0.770.85 of H\* throughout — neither shows a threshold. Stationary **forward-KL** falls monotonically (dry 2.08 → g=0.2: 0.75), significant at g≥0.05 (paired t up to 3.3; 89% of lineages improve at g=0.2). This *refines* the earlier "forward-KL AND tail survival" note: for a smoothing model, support-counting decouples from closeness-to-truth. (2) **The sharp `g*≪1` is an exact-operator feature, softened by neural inductive bias.** Half the achievable KL reduction closes by a *median-recovery* grounding g≈0.04 (bootstrap CI [0.004, 0.116]) — a striking echo of Layer-1's 0.048 — but full (95%) recovery needs g≈0.19, far more than the histogram bridge, because smoothing both caps dry collapse (KL~2, not ∞) and slows full recovery. So the quantitative `g*≪1` claim rests on the **histogram bridge** (g\*=0.047, exact reduction to Layer 1), which the trained RNN confirms in *sign* and softens in sharpness. Honest note: the pre-registered 95%-of-H\*/tail-survival falsifier is not met, but that is because those are the wrong metrics for a smoothing model, not because grounding fails — the blueprint §3.5 directional claim (grounding arrests collapse) holds robustly. Robustness fix landed alongside: a fully-degenerate RNN can emit only invalid codewords, so `measure_distribution` returns a terminal-collapse sentinel (fixation on the dominant mode) instead of crashing a long sweep. **Finding (2026-07-05, neural `grounding`) — grounding's SIGN transfers to trained RNN weights, but the sharp `g*` does not; and tail-survival is the *wrong* neural collapse metric.** Re-ran the phase-boundary sweep at 18 replicates. Two results: (1) **forward-KL is the operative neural collapse metric, not H or tail-survival.** The RNN's smoothing inductive bias keeps *spurious* tail modes alive (it generalises to unseen codewords), so `tail_truth_mass_alive` is flat/**non-monotone** in g (dry 0.54 > most grounded points) and H stays ~0.770.85 of H\* throughout — neither shows a threshold. Stationary **forward-KL** falls monotonically (dry 2.08 → g=0.2: 0.75), significant at g≥0.05 (paired t up to 3.3; 89% of lineages improve at g=0.2). This *refines* the earlier "forward-KL AND tail survival" note: for a smoothing model, support-counting decouples from closeness-to-truth. (2) **The sharp `g*≪1` is an exact-operator feature, softened by neural inductive bias.** Half the achievable KL reduction closes by a *median-recovery* grounding g≈0.04 (bootstrap CI [0.004, 0.116]) — a striking echo of Layer-1's 0.048 — but full (95%) recovery needs g≈0.19, far more than the histogram bridge, because smoothing both caps dry collapse (KL~2, not ∞) and slows full recovery. So the quantitative `g*≪1` claim rests on the **histogram bridge** (g\*=0.047, exact reduction to Layer 1), which the trained RNN confirms in *sign* and softens in sharpness. Honest note: the pre-registered 95%-of-H\*/tail-survival falsifier is not met, but that is because those are the wrong metrics for a smoothing model, not because grounding fails — the blueprint §3.5 directional claim (grounding arrests collapse) holds robustly. Robustness fix landed alongside: a fully-degenerate RNN can emit only invalid codewords, so `measure_distribution` returns a terminal-collapse sentinel (fixation on the dominant mode) instead of crashing a long sweep.
**Finding (2026-07-05, real-MNIST `mnist_collapse`) — collapse and grounding-rescue reproduce on real images.** External-validity tier: a small **convolutional VAE** (the canonical generative-collapse model) is retrained each generation on its own generated digits. Modes = digit class × stroke-thickness bin (K=30, Zipf, ~18 tail modes); the oracle is a **frozen CNN + deterministic thickness** at **98.5% mode accuracy** (its 30×30 confusion matrix is recorded in the manifest as the measurement floor). Result (4 reps): the **dry (g=0) lineage collapses to a single mode** — forward-KL 0.5→18, support 30→1, tail truth-mass 1.0→0.06, H→0 — while **10% grounding holds all 30 modes** (KL≈0.6, full tail, H≈0.9). The VAE needs ~10% grounding here vs the synthetic histogram's ~5%, consistent with the `grounding` finding that trained neural models need more grounding than the exact operator. **Confirmation-only (signs, not magnitudes; blueprint §3.5)** — the exact synthetic oracle stays the quantitative anchor. `figures/mnist_montage.py` is an eyeball diagnostic (re-runs a short dry lineage; NOT a parquet figure). Build gates passed: CNN mode accuracy 98.5%; VAE gen-0 recovers full 30/30 support (over-smooths frequencies, KL≈0.5, no prior hole — unlike the *synthetic*-codeword VAE, which is why the MNIST VAE works where that one didn't). The MNIST tier is heavy (torchvision `--extra mnist`, downloads MNIST, ~5 min): `make mnist`, kept out of the `make neural` loop.
## Build order (blueprint §7) — respect the gate ## Build order (blueprint §7) — respect the gate
1. Scaffold: repo layout (§5), container, pytest skeleton, config system, seeding utils. `make test` green. 1. Scaffold: repo layout (§5), container, pytest skeleton, config system, seeding utils. `make test` green.

View file

@ -1,7 +1,7 @@
# Layer 1 + Layer 1.5 automation. The uv venv (built from the committed uv.lock) is the # Layer 1 + Layer 1.5 automation. The uv venv (built from the committed uv.lock) is the
# reproducibility source of truth; every target runs inside it via `uv run`. # reproducibility source of truth; every target runs inside it via `uv run`.
.PHONY: env env-neural test layer1 layer2 neural figures clean .PHONY: env env-neural env-mnist test layer1 layer2 neural mnist figures clean
env: ## build .venv from the committed lockfile env: ## build .venv from the committed lockfile
uv sync --extra dev uv sync --extra dev
@ -9,14 +9,21 @@ env: ## build .venv from the committed lockfile
env-neural: ## add the Layer 1.5 torch stack (GPU; Stage C onward) env-neural: ## add the Layer 1.5 torch stack (GPU; Stage C onward)
uv sync --extra dev --extra neural uv sync --extra dev --extra neural
env-mnist: ## add torchvision for the real-MNIST confirmation tier
uv sync --extra dev --extra neural --extra mnist
test: ## correctness tests + scientific-validation tests (the spine of trust) test: ## correctness tests + scientific-validation tests (the spine of trust)
uv run pytest uv run pytest
layer1: ## run experiments E1-E6 layer1: ## run experiments E1-E6
for e in E1 E2 E3 E4 E5 E6; do uv run python -m knowledge.experiment configs/layer1/$$e.yaml; done for e in E1 E2 E3 E4 E5 E6; do uv run python -m knowledge.experiment configs/layer1/$$e.yaml; done
neural: ## run Layer 1.5 neural experiments (N-series) neural: ## run Layer 1.5 synthetic neural experiments (excludes the heavy MNIST tier)
for c in configs/neural/*.yaml; do uv run python -m neural.experiment $$c; done for c in configs/neural/*.yaml; do case "$$c" in *mnist*) ;; \
*) uv run python -m neural.experiment "$$c" ;; esac; done
mnist: ## run the real-MNIST confirmation tier (needs env-mnist; downloads MNIST)
uv run python -m neural.experiment configs/neural/mnist_collapse.yaml
layer2: neural ## alias: Layer 1.5 is the current Layer-2 deliverable (LLM rung deferred) layer2: neural ## alias: Layer 1.5 is the current Layer-2 deliverable (LLM rung deferred)

View file

@ -0,0 +1,59 @@
experiment: mnist_collapse
kind: mnist_lineage
seed: 20260705
n_replicates: 4
# (Layer 1.5 external-validity tier; maps to Layer-1 E1/E2 and neural collapse/grounding): does
# model collapse — and its rescue by grounding — appear on REAL MNIST images, not just the
# synthetic sandbox? A convolutional VAE (the canonical model in which generative collapse was
# first observed) is retrained each generation on the previous VAE's own generated images, plus a
# fraction g of fresh REAL MNIST images (grounding). Modes = digit class x stroke-thickness bin
# (K=30) resampled to a Zipf p*; a frozen CNN + deterministic thickness is the oracle (its
# confusion matrix, recorded in the manifest, is the measurement-noise floor). Expect (per E1/E2):
# the dry arm (g=0) collapses — rare modes die, forward-KL climbs, support shrinks — while a
# grounded arm holds the tail. This is confirmation-only: SIGNS, not magnitudes (blueprint 3.5);
# the exact synthetic oracle remains the anchor for every quantitative comparison. Falsifier: the
# dry VAE shows no diversity loss, or grounding fails to arrest it.
mnist:
K: 30
n_classes: 10
style_bins: 3 # K = 10 classes x 3 stroke-thickness bins
R: 1
tail: zipf
zipf_s: 1.5 # the rarest ~18/30 modes form a real tail (~9% of the mass)
tail_threshold: 1.0e-2
init: truth
data_root: data
model:
kind: convvae
latent: 32
epochs: 30
lr: 1.0e-3
batch_size: 256
beta: 1.0
n_eval: 10000 # generate-and-classify samples for the mode-distribution readout
oracle:
epochs: 5 # frozen digit CNN (~98.5% mode accuracy = the noise floor)
lr: 1.0e-3
batch_size: 256
cache: models/mnist_cnn.pt
dynamics:
n: 6000 # images the pupil VAE sees per generation (drift strength)
grounding: {m: 0, policy: proportional} # m overwritten per g by the sweep
generations: 15
metrics:
kl_floor: 1.0e-9
support_eps: 1.0e-9
sweep:
- param: g
values: [0.0, 0.1] # dry vs grounded (VAE collapse is strong; needs ~10% real, cf. grounding)
output:
dir: results/mnist_collapse

75
figures/mnist_montage.py Normal file
View file

@ -0,0 +1,75 @@
"""Eyeball diagnostic: watch a dry MNIST lineage collapse, generation by generation.
Unlike the `plot_mnist` figure (a pure function of committed parquet), this **re-runs** a short
dry VAE lineage and saves a grid of freshly-generated digits at a few generations, so the
collapse is visible directly early generations show varied digits, late generations degenerate
toward a single blurry mode. Diagnostic only; not part of the reproducible figure set.
Usage: python figures/mnist_montage.py [results/mnist_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__).parents[1] / "src"))
from neural.config import MnistCfg, ModelCfg, OracleCfg # noqa: E402
from neural.mnist_data import ( # noqa: E402
MnistSampler, assign_modes, load_mnist, make_mnist_truth,
)
from neural.mnist_oracle import build_oracle # noqa: E402
from neural.mnist_vae import ConvVAEGenerator # noqa: E402
_SNAP_GENS = [0, 4, 8, 12, 15] # generations to snapshot
_COLS = 12 # sample digits per row
def main(out_dir: str = "results/mnist_collapse") -> None:
cfg = MnistCfg()
data = load_mnist(cfg.data_root)
td = make_mnist_truth(cfg)
oracle, cuts, _ = build_oracle(cfg, OracleCfg(), data, seed=20260705)
modes = assign_modes(data.train_x, data.train_y, cuts, cfg)
sampler = MnistSampler(data.train_x, modes, cfg.K)
rng = np.random.default_rng(0)
mcfg = ModelCfg(kind="convvae", latent=32, epochs=30, lr=1e-3, batch_size=256,
beta=1.0, n_eval=10000)
n = 6000
model = ConvVAEGenerator(cfg, mcfg, oracle)
model.fit(sampler.draw(rng.multinomial(n, td.p_star), rng), rng) # gen 0 (dry)
snaps: dict[int, np.ndarray] = {}
g_max = max(_SNAP_GENS)
for t in range(0, g_max + 1):
if t in _SNAP_GENS:
snaps[t] = model.sample(_COLS, rng)[:, 0]
if t < g_max: # advance one dry generation
pupil = ConvVAEGenerator(cfg, mcfg, oracle)
pupil.fit(model.sample(n, rng), rng)
model = pupil
fig, axes = plt.subplots(len(_SNAP_GENS), _COLS,
figsize=(_COLS * 0.7, len(_SNAP_GENS) * 0.8))
for row, t in enumerate(_SNAP_GENS):
for col in range(_COLS):
ax = axes[row, col]
ax.imshow(snaps[t][col], cmap="gray_r", vmin=0, vmax=1)
ax.set_xticks([]); ax.set_yticks([])
if col == 0:
ax.set_ylabel(f"gen {t}", fontsize=9, rotation=0, ha="right", va="center")
fig.suptitle("Dry MNIST lineage collapses: varied digits (top) -> a single mode (bottom)",
fontsize=11)
fig.tight_layout()
out = Path(out_dir) / "mnist_montage.png"
fig.savefig(out, dpi=130, bbox_inches="tight")
fig.savefig(out.with_suffix(".pdf"), bbox_inches="tight")
print(f"wrote {out} and .pdf")
if __name__ == "__main__":
main(*sys.argv[1:])

80
figures/plot_mnist.py Normal file
View file

@ -0,0 +1,80 @@
"""`mnist_collapse` figure — collapse and grounding-rescue on REAL MNIST images.
External validity for the neural tier: a convolutional VAE retrained each generation on its own
generated digits collapses rare (class, thickness) modes die, forward-KL to the Zipf truth
climbs, support shrinks while a grounded arm (a fraction of fresh real MNIST images each
generation) holds the tail. Modes are read by a frozen CNN oracle whose mode accuracy (the
measurement-noise floor) is annotated from the run manifest. Signs, not magnitudes.
Four panels, dry (g=0) vs grounded, mean ± 95% CI across replicates: (A) forward-KL trajectories;
(B) support size (distinct modes alive); (C) tail truth-mass alive; (D) heterozygosity. Reads the
committed bundle (parquet) + manifest.json only.
Usage: python figures/plot_mnist.py [results/mnist_collapse]
"""
from __future__ import annotations
import json
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
sys.path.insert(0, str(Path(__file__).parents[1] / "src"))
from knowledge.metrics import heterozygosity # noqa: E402
from neural.config import MnistCfg # noqa: E402
from neural.mnist_data import make_mnist_truth # noqa: E402
def _traj(df, g, col):
"""Return (generations, mean, 95% half-width) of ``col`` for arm ``g`` across replicates."""
sub = df[df["g"] == g]
grp = sub.groupby("generation")[col]
gens = np.array(sorted(sub["generation"].unique()))
return gens, grp.mean().to_numpy(), 1.96 * grp.sem().to_numpy()
def main(results_dir: str = "results/mnist_collapse") -> None:
df, cfg = load_bundle(results_dir)
syn = MnistCfg(**cfg["mnist"])
H_star = heterozygosity(make_mnist_truth(syn).p_star)
manifest = json.loads((Path(results_dir) / "manifest.json").read_text())
oracle_acc = manifest.get("oracle_mode_accuracy", float("nan"))
g_dry, g_wet = min(df["g"].unique()), max(df["g"].unique())
arms = [(g_dry, "#d62728", f"dry (g={g_dry:g})"), (g_wet, "#2ca02c", f"grounded (g={g_wet:g})")]
fig, axes = plt.subplots(2, 2, figsize=(13, 9))
def panel(ax, col, title, ylabel, hline=None):
for g, c, lab in arms:
gens, m, ci = _traj(df, g, col)
ax.plot(gens, m, "-o", color=c, ms=3, label=lab)
ax.fill_between(gens, m - ci, m + ci, color=c, alpha=0.2)
if hline is not None:
ax.axhline(hline[0], ls=":", color="gray", lw=1, label=hline[1])
ax.set(xlabel="generation", ylabel=ylabel, title=title)
ax.legend(frameon=False, fontsize=9)
panel(axes[0, 0], "forward_kl", "Collapse: dry forward-KL climbs, grounding holds it",
r"forward-KL $D(p^*\Vert\hat p)$")
panel(axes[0, 1], "support_size", f"Support collapses (of K={syn.K} modes)",
"distinct modes alive", hline=(syn.K, f"$K$={syn.K}"))
panel(axes[1, 0], "tail_truth_mass_alive", "Rare tail dies dry, held by grounding",
"tail truth-mass alive")
panel(axes[1, 1], "heterozygosity", "Diversity collapses dry, held by grounding",
"heterozygosity $H$", hline=(H_star, "$H^*$"))
fig.suptitle("mnist_collapse — model collapse and grounding-rescue on REAL MNIST images "
f"(VAE; oracle mode acc {oracle_acc:.1%} = noise floor)", y=1.0, fontsize=13)
fig.tight_layout()
savefig(fig, results_dir, "mnist_collapse")
if __name__ == "__main__":
main(*sys.argv[1:])

View file

@ -0,0 +1,39 @@
# mnist_collapse — collapse and grounding-rescue on REAL MNIST images (external validity)
**Claim tested:** everything so far used a *synthetic* sandbox with a zero-error decoder oracle. Do
model collapse and its rescue by grounding also appear on **real images** with a **classifier**
oracle — i.e. is the effect real, not a synthetic artefact?
**Setup (Layer 1.5, real-data tier).** The generative model is a **convolutional VAE** (the model
in which generative collapse was first observed). Each generation a **fresh** VAE is trained from
scratch on the previous VAE's own generated digits, plus a fraction `g` of fresh **real** MNIST
images (grounding). `K = 30` **modes** = digit class × stroke-thickness bin (S=3), Zipf-resampled so
the rarest ~18 modes form a real tail. The **oracle** is a frozen CNN (digit class) + deterministic
thickness bin; its **mode accuracy ≈ 98.5%** (recorded in `manifest.json` with the full 30×30
confusion matrix) is the measurement-noise floor. Two arms — dry (`g = 0`) vs grounded (`g = 0.1`) —
`n = 6000` images/generation, 15 generations, 4 replicates.
### Symbols
- **mode** = (digit class, stroke-thickness bin); **`p*`** = Zipf truth over the 30 modes; **`p̂`** = the VAE's oracle-measured mode distribution.
- **`g`** = grounding fraction (share of real MNIST images each generation). **forward-KL** = distance from truth; **support** = distinct modes alive; **`H`** = diversity; **tail truth-mass alive** = fraction of the rare tail retained.
### The four panels (dry = red, grounded = green; band = 95% CI over 4 reps)
1. **Forward-KL.** Dry climbs from ~0.5 to **~18** (the VAE drifts far from truth); grounded stays
near the floor. Collapse is real on images.
2. **Support.** Dry collapses from all 30 modes to **~1** (the VAE ends up emitting a single blurry
mode); grounded holds all 30.
3. **Tail truth-mass alive.** Dry's rare tail is wiped out (→ 0.06); grounded keeps the whole tail.
4. **Heterozygosity.** Dry diversity → 0; grounded holds `H ≈ 0.9`.
See **`mnist_montage.png`** for the eyeball version: gen-0 digits are varied and recognisable; by
gen 1215 the dry lineage has degenerated into one blurry blob.
### Takeaway
Model collapse and its arrest by a small dose of real data **reproduce on real MNIST images with a
learned classifier oracle** — external validity for the whole Layer-1.5 story. Note the VAE needs
~10% grounding here (vs ~5% for the synthetic histogram), consistent with the `grounding` finding
that trained neural models need somewhat more grounding than the exact operator. This is
**confirmation-only** (signs, not magnitudes; blueprint §3.5) — the exact synthetic oracle remains
the anchor for every quantitative claim, and the oracle confusion matrix is the recorded noise floor.
**Falsifier (not triggered):** if the dry VAE had shown no diversity loss, or grounding had failed to
arrest it, the external-validity claim would fail.

View file

@ -0,0 +1,984 @@
{
"experiment": "mnist_collapse",
"master_seed": 20260705,
"git_commit": "3b9f4f789358a442d510bf2d609aab0986834ac5",
"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": 128,
"results_sha256": "d808003e0e01247923996c94eff18c430308448e23bb5d102aa860701145131e",
"layer": "1.5",
"tier": "mnist",
"model_kind": "convvae",
"oracle_ckpt_sha256": "7ff2cc619861daacf8168277c83cde29483554c92d732a24db9fdaa446fa9ff8",
"oracle_mode_accuracy": 0.9847,
"oracle_class_accuracy": 0.9847,
"confusion_matrix": [
[
365,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
1,
0,
0,
0,
1,
0,
0,
1,
0,
1,
0,
0,
0,
0,
0
],
[
0,
302,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
2,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0
],
[
0,
0,
307,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0
],
[
0,
0,
0,
366,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
1,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0
],
[
0,
0,
0,
0,
373,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0
],
[
0,
0,
0,
0,
0,
390,
0,
2,
1,
0,
0,
0,
0,
0,
0,
0,
0,
0,
1,
1,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0
],
[
1,
0,
0,
0,
0,
1,
335,
0,
0,
1,
0,
0,
0,
1,
0,
0,
0,
0,
0,
0,
0,
1,
0,
1,
1,
0,
0,
0,
0,
0
],
[
1,
0,
0,
0,
0,
0,
0,
308,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
1,
0,
0,
0,
0,
0,
0
],
[
0,
0,
1,
0,
0,
1,
0,
0,
376,
0,
0,
2,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0
],
[
0,
0,
0,
0,
0,
0,
0,
0,
0,
327,
0,
0,
0,
0,
0,
5,
0,
0,
0,
0,
0,
1,
0,
0,
0,
0,
0,
0,
0,
0
],
[
1,
0,
0,
0,
0,
0,
1,
0,
0,
0,
303,
0,
0,
0,
0,
0,
2,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0
],
[
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
368,
0,
0,
0,
0,
0,
1,
0,
0,
0,
0,
0,
0,
0,
0,
1,
0,
0,
0
],
[
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
326,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
4,
0,
0
],
[
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
324,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
1,
0
],
[
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
324,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
3
],
[
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
269,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0
],
[
0,
0,
0,
0,
0,
0,
0,
0,
0,
1,
0,
0,
0,
0,
0,
0,
299,
0,
0,
1,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0
],
[
0,
2,
0,
0,
0,
0,
0,
0,
0,
0,
1,
0,
0,
0,
0,
0,
0,
318,
0,
1,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0
],
[
3,
0,
0,
1,
0,
1,
0,
0,
0,
0,
0,
0,
2,
1,
0,
3,
0,
0,
247,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0
],
[
3,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
1,
0,
0,
2,
0,
0,
315,
0,
0,
0,
0,
1,
0,
0,
0,
0,
0
],
[
0,
0,
1,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
1,
0,
0,
2,
0,
0,
374,
0,
0,
0,
0,
0,
0,
0,
0,
0
],
[
0,
0,
0,
1,
0,
0,
4,
0,
0,
2,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
327,
0,
0,
1,
0,
0,
0,
0,
0
],
[
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
335,
0,
1,
0,
0,
0,
2,
0
],
[
0,
1,
0,
0,
0,
0,
1,
2,
1,
0,
1,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
349,
0,
0,
0,
0,
0,
0
],
[
1,
0,
0,
0,
0,
0,
4,
0,
0,
1,
1,
0,
1,
0,
0,
4,
1,
0,
1,
2,
0,
0,
0,
0,
280,
0,
0,
0,
0,
0
],
[
2,
0,
0,
0,
0,
0,
0,
0,
0,
0,
2,
1,
0,
0,
1,
0,
2,
2,
0,
0,
0,
0,
0,
0,
0,
347,
0,
0,
0,
5
],
[
0,
0,
1,
0,
0,
1,
0,
0,
1,
0,
0,
5,
0,
0,
0,
0,
0,
1,
0,
0,
2,
0,
0,
2,
0,
0,
297,
0,
0,
6
],
[
0,
0,
0,
0,
2,
0,
0,
0,
0,
0,
0,
0,
1,
1,
0,
1,
0,
0,
0,
0,
0,
1,
0,
0,
1,
0,
0,
268,
0,
0
],
[
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
2,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
365,
0
],
[
0,
1,
1,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
2,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
363
]
]
}

Binary file not shown.

Binary file not shown.

After

Width:  |  Height:  |  Size: 278 KiB

Binary file not shown.

Binary file not shown.

After

Width:  |  Height:  |  Size: 182 KiB

View file

@ -0,0 +1,110 @@
experiment: mnist_collapse
seed: 20260705
n_replicates: 4
source_config:
experiment: mnist_collapse
kind: mnist_lineage
seed: 20260705
n_replicates: 4
mnist:
K: 30
n_classes: 10
style_bins: 3
R: 1
tail: zipf
zipf_s: 1.5
tail_threshold: 0.01
init: truth
data_root: data
model:
kind: convvae
latent: 32
epochs: 30
lr: 0.001
batch_size: 256
beta: 1.0
n_eval: 10000
oracle:
epochs: 5
lr: 0.001
batch_size: 256
cache: models/mnist_cnn.pt
dynamics:
n: 6000
grounding:
m: 0
policy: proportional
generations: 15
metrics:
kl_floor: 1.0e-09
support_eps: 1.0e-09
sweep:
- param: g
values:
- 0.0
- 0.1
output:
dir: results/mnist_collapse
grid:
- label:
g: 0.0
m: 0
lineage_cfg:
mnist:
K: 30
n_classes: 10
style_bins: 3
R: 1
tail: zipf
zipf_s: 1.5
tail_threshold: 0.01
init: truth
data_root: data
model:
kind: convvae
latent: 32
epochs: 30
lr: 0.001
batch_size: 256
beta: 1.0
n_eval: 10000
dynamics:
n: 6000
grounding:
m: 0
policy: proportional
generations: 15
metrics:
kl_floor: 1.0e-09
support_eps: 1.0e-09
- label:
g: 0.1
m: 667
lineage_cfg:
mnist:
K: 30
n_classes: 10
style_bins: 3
R: 1
tail: zipf
zipf_s: 1.5
tail_threshold: 0.01
init: truth
data_root: data
model:
kind: convvae
latent: 32
epochs: 30
lr: 0.001
batch_size: 256
beta: 1.0
n_eval: 10000
dynamics:
n: 6000
grounding:
m: 667
policy: proportional
generations: 15
metrics:
kl_floor: 1.0e-09
support_eps: 1.0e-09

View file

@ -58,6 +58,45 @@ class SyntheticCfg:
return self.id_len + self.style_len return self.id_len + self.style_len
@dataclass(frozen=True)
class MnistCfg:
"""Real-MNIST mode-truth: a Zipf ``p*`` over ``K = n_classes * style_bins`` modes.
A mode is ``(digit class, stroke-thickness bin)`` under the fixed bijection
``mode = class * style_bins + bin``. The first seven fields are the Layer-1 ``TruthCfg``
knobs (they build the Zipf ``p*`` over the ``K`` modes via ``make_true_distribution``);
the rest govern the image tier. Unlike ``SyntheticCfg`` there is no rendering grammar
observations are real images and the oracle is a frozen classifier.
"""
K: int = 30
R: int = 1
tail: str = "zipf"
zipf_s: float = 1.5 # steep enough that the rarest ~18/30 modes form a real tail
tail_frac: float = 0.5
tail_threshold: float = 1e-2 # modes with p* < 0.01 are "tail" (~9% of the mass)
init: str = "truth" # initial p_0 over modes: {uniform, truth}
n_classes: int = 10 # MNIST digit classes
style_bins: int = 3 # S: per-class stroke-thickness quantile bins (K = n_classes*S)
data_root: str = "data" # gitignored MNIST download dir
def __post_init__(self) -> None:
if self.K != self.n_classes * self.style_bins:
raise ValueError(
f"K ({self.K}) must equal n_classes*style_bins "
f"({self.n_classes}*{self.style_bins}={self.n_classes * self.style_bins})")
@dataclass(frozen=True)
class OracleCfg:
"""Frozen-classifier oracle training/caching (MNIST tier)."""
epochs: int = 5
lr: float = 1.0e-3
batch_size: int = 256
cache: str = "models/mnist_cnn.pt" # gitignored checkpoint; its hash goes in the manifest
@dataclass(frozen=True) @dataclass(frozen=True)
class ModelCfg: class ModelCfg:
"""The generative learner. ``kind`` selects the architecture behind a thin adapter. """The generative learner. ``kind`` selects the architecture behind a thin adapter.
@ -65,7 +104,7 @@ class ModelCfg:
Neural hyperparameters are ignored by the ``histogram`` bridge model. Neural hyperparameters are ignored by the ``histogram`` bridge model.
""" """
kind: str = "histogram" # {histogram, rnn, vae, mlp} kind: str = "histogram" # {histogram, rnn, vae, mlp, convvae}
hidden: int = 64 hidden: int = 64
embed: int = 16 embed: int = 16
epochs: int = 30 epochs: int = 30

View file

@ -28,6 +28,9 @@ from .generation_loop import run_generative_lineage
# Config groups that make up a single neural lineage (everything else is experiment-level). # Config groups that make up a single neural lineage (everything else is experiment-level).
_NEURAL_KEYS = ("synthetic", "model", "dynamics", "generations", "metrics", "n_eval") _NEURAL_KEYS = ("synthetic", "model", "dynamics", "generations", "metrics", "n_eval")
# Config groups for one MNIST lineage (the real-image tier; `mnist` replaces `synthetic`).
_MNIST_KEYS = ("mnist", "model", "dynamics", "generations", "metrics")
# Libraries recorded in the manifest on top of the Layer-1 core set (skipped if absent). # Libraries recorded in the manifest on top of the Layer-1 core set (skipped if absent).
_EXTRA_LIBS = ("torch", "torchvision") _EXTRA_LIBS = ("torch", "torchvision")
@ -91,12 +94,90 @@ def run_experiment(cfg: dict) -> pd.DataFrame:
return out return out
def _expand_mnist(cfg: dict) -> list[tuple[dict, dict]]:
"""Expand the MNIST g-sweep into (label, resolved-lineage) pairs (reuses ``_apply_param``)."""
base = {k: copy.deepcopy(cfg[k]) for k in _MNIST_KEYS if k in cfg}
sweeps = cfg.get("sweep", [])
if isinstance(sweeps, dict):
sweeps = [sweeps]
if not sweeps:
return [({}, base)]
params = [s["param"] for s in sweeps]
value_lists = [list(s["values"]) for s in sweeps]
combos: list[tuple[dict, dict]] = []
for values in itertools.product(*value_lists):
lin = copy.deepcopy(base)
label: dict = {}
for param, val in zip(params, values):
label.update(_apply_param(lin, param, val))
combos.append((label, lin))
return combos
def run_mnist_experiment(cfg: dict) -> tuple[pd.DataFrame, dict]:
"""Run every MNIST grid point x replicate; return (results, oracle-provenance manifest).
The frozen classifier oracle, per-mode real-image pools, and ``p*`` are built **once** and
shared across all arms/replicates (training the CNN and indexing the pools is expensive).
The oracle's confusion matrix over the real test set is returned for the manifest as the
measurement-noise floor.
Args:
cfg (dict): Parsed MNIST experiment YAML (``mnist``/``model``/``dynamics``/``oracle``
blocks, a ``sweep`` over ``g``, ``seed``, ``n_replicates``).
Returns:
tuple[pd.DataFrame, dict]: Long-form results and the oracle-provenance manifest dict.
"""
from knowledge.config import _sub
from .config import MnistCfg, OracleCfg
from .mnist_data import assign_modes, load_mnist, make_mnist_truth, MnistSampler
from .mnist_loop import run_mnist_lineage
from .mnist_oracle import build_oracle, confusion_summary
mnist_cfg = _sub(cfg["mnist"], MnistCfg)
oracle_cfg = _sub(cfg.get("oracle", {}), OracleCfg)
master = int(cfg["seed"])
n_rep = int(cfg["n_replicates"])
data = load_mnist(mnist_cfg.data_root)
td = make_mnist_truth(mnist_cfg)
oracle, cuts, ckpt_hash = build_oracle(mnist_cfg, oracle_cfg, data, seed=master)
modes = assign_modes(data.train_x, data.train_y, cuts, mnist_cfg)
sampler = MnistSampler(data.train_x, modes, mnist_cfg.K)
conf = confusion_summary(oracle, data, cuts, mnist_cfg)
combos = _expand_mnist(cfg)
seeds = spawn_seeds(master, n_rep)
frames: list[pd.DataFrame] = []
for label, lineage_cfg in combos:
for rep, ss in enumerate(seeds):
df = run_mnist_lineage(lineage_cfg, int(ss.generate_state(1)[0]), oracle, sampler, td)
for col, val in label.items():
df[col] = val
df["replicate"] = rep
frames.append(df)
out = pd.concat(frames, ignore_index=True)
out.insert(0, "experiment", cfg["experiment"])
manifest = {
"layer": "1.5", "tier": "mnist", "model_kind": cfg.get("model", {}).get("kind"),
"oracle_ckpt_sha256": ckpt_hash,
"oracle_mode_accuracy": conf["mode_accuracy"],
"oracle_class_accuracy": conf["class_accuracy"],
"confusion_matrix": conf["confusion"],
}
return out, manifest
def run_and_save(config_path: str | Path) -> Path: def run_and_save(config_path: str | Path) -> Path:
"""Load a neural experiment YAML, run it, and write artifacts. Returns the output dir.""" """Load a neural experiment YAML, run it, and write artifacts. Returns the output dir."""
config_path = Path(config_path) config_path = Path(config_path)
cfg = yaml.safe_load(config_path.read_text()) cfg = yaml.safe_load(config_path.read_text())
out_dir = Path(cfg.get("output", {}).get("dir", f"results/{cfg['experiment']}")) out_dir = Path(cfg.get("output", {}).get("dir", f"results/{cfg['experiment']}"))
kind = cfg.get("kind", "gen_lineage") kind = cfg.get("kind", "gen_lineage")
extra_manifest = {"layer": "1.5", "model_kind": cfg.get("model", {}).get("kind", "histogram")}
if kind == "recombination": if kind == "recombination":
from .recombine import run_recombination # the `recombination` experiment; lazy import from .recombine import run_recombination # the `recombination` experiment; lazy import
df = run_recombination(cfg) df = run_recombination(cfg)
@ -104,11 +185,13 @@ def run_and_save(config_path: str | Path) -> Path:
elif kind == "gen_lineage": elif kind == "gen_lineage":
df = run_experiment(cfg) df = run_experiment(cfg)
grid = [{"label": label, "neural_cfg": c} for label, c in expand_sweeps(cfg)] grid = [{"label": label, "neural_cfg": c} for label, c in expand_sweeps(cfg)]
elif kind == "mnist_lineage":
df, extra_manifest = run_mnist_experiment(cfg) # oracle provenance + confusion matrix
grid = [{"label": label, "lineage_cfg": c} for label, c in _expand_mnist(cfg)]
else: else:
raise ValueError(f"unknown neural experiment kind {kind!r}") raise ValueError(f"unknown neural experiment kind {kind!r}")
model_kind = cfg.get("model", {}).get("kind", "histogram")
save_artifacts(cfg, df, out_dir, extra_libs=_EXTRA_LIBS, save_artifacts(cfg, df, out_dir, extra_libs=_EXTRA_LIBS,
extra_manifest={"layer": "1.5", "model_kind": model_kind}, grid=grid) extra_manifest=extra_manifest, grid=grid)
return out_dir return out_dir

198
src/neural/mnist_data.py Normal file
View file

@ -0,0 +1,198 @@
"""Real-MNIST data for the secondary-confirmation tier.
Turns MNIST into a `K`-mode world with a Zipf tail so the *same* Wright-Fisher machinery
applies. A **mode** is ``(digit class, stroke-thickness bin)`` under the fixed bijection
``mode = class * style_bins + bin``; the Zipf ``p*`` over the ranked modes comes from Layer 1's
``make_true_distribution``. **Style = stroke thickness** (mean ink per image, an always-defined
pixel statistic), binned into per-class quantiles fit on the real training set so within each
class the thin/medium/thick variants split into equal-frequency style bins that the generator
must keep alive.
No physical resampling of MNIST: the Zipf enters through *sampling* the gen-0 training set and
each generation's grounding both draw real images whose modes follow ``p*`` from per-mode pools
(``MnistSampler``, the image analogue of ``synthetic.render_modes``).
"""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import numpy as np
from knowledge.truth import TrueDist, make_true_distribution
from .config import MnistCfg
def make_mnist_truth(cfg: MnistCfg) -> TrueDist:
"""Build the Zipf ``p*`` (and tail mask / regions) over the ``K`` MNIST modes.
Thin wrapper over Layer 1's ``make_true_distribution`` — so "mode", "tail", and "region"
are the *same objects* as everywhere else. Mode ``k`` ranks by descending ``p*`` and maps
to ``(class=k//style_bins, bin=k%style_bins)``.
Args:
cfg (MnistCfg): The MNIST mode-truth configuration.
Returns:
TrueDist: ``p_star`` (length ``K``), ``regions``, and ``tail_mask``.
"""
return make_true_distribution(cfg.K, cfg.R, cfg.tail, cfg.tail_frac, cfg.zipf_s, 0,
tail_threshold=cfg.tail_threshold)
@dataclass(frozen=True)
class MnistData:
"""Loaded MNIST split as float images in [0,1] plus integer labels.
Attributes:
train_x (np.ndarray): Train images, shape ``(N, 1, 28, 28)``, float32 in [0,1].
train_y (np.ndarray): Train digit labels, shape ``(N,)``.
test_x (np.ndarray): Test images, shape ``(M, 1, 28, 28)``.
test_y (np.ndarray): Test digit labels, shape ``(M,)``.
"""
train_x: np.ndarray
train_y: np.ndarray
test_x: np.ndarray
test_y: np.ndarray
def load_mnist(root: str | Path) -> MnistData:
"""Load MNIST (downloading on first use) as float images in [0,1].
Args:
root (str | Path): Download/cache directory (gitignored).
Returns:
MnistData: Train/test images (``(N,1,28,28)`` float32) and labels.
"""
from torchvision import datasets # lazy: only the MNIST tier needs torchvision
tr = datasets.MNIST(str(root), train=True, download=True)
te = datasets.MNIST(str(root), train=False, download=True)
to_x = lambda d: (d.data.numpy().astype("float32") / 255.0)[:, None, :, :]
return MnistData(to_x(tr), tr.targets.numpy().astype("int64"),
to_x(te), te.targets.numpy().astype("int64"))
def thickness(images: np.ndarray) -> np.ndarray:
"""Per-image stroke thickness = mean pixel intensity (ink), always defined.
Args:
images (np.ndarray): Images of shape ``(N, 1, 28, 28)`` (or ``(N, 28, 28)``).
Returns:
np.ndarray: Length-``N`` thickness values.
"""
x = np.asarray(images, dtype="float32")
return x.reshape(x.shape[0], -1).mean(axis=1)
def fit_thickness_thresholds(images: np.ndarray, labels: np.ndarray, cfg: MnistCfg) -> np.ndarray:
"""Fit per-class thickness quantile cut-points on real training images.
Within each class, the ``style_bins`` bins are equal-frequency (quantile) splits of
thickness so every (class, bin) mode is populated on real data.
Args:
images (np.ndarray): Train images ``(N,1,28,28)``.
labels (np.ndarray): Train digit labels ``(N,)``.
cfg (MnistCfg): Supplies ``n_classes`` and ``style_bins``.
Returns:
np.ndarray: Cut-points of shape ``(n_classes, style_bins - 1)`` (empty middle dim
collapses to shape ``(n_classes, 0)`` when ``style_bins == 1``).
"""
S = cfg.style_bins
th = thickness(images)
qs = np.linspace(0.0, 1.0, S + 1)[1:-1] # interior quantiles
cuts = np.zeros((cfg.n_classes, max(S - 1, 0)), dtype="float32")
for c in range(cfg.n_classes):
vals = th[labels == c]
if S > 1 and vals.size:
cuts[c] = np.quantile(vals, qs)
return cuts
def thickness_bin(images: np.ndarray, class_ids: np.ndarray, cuts: np.ndarray) -> np.ndarray:
"""Assign each image a per-class thickness bin in ``[0, style_bins)``.
Args:
images (np.ndarray): Images ``(N,1,28,28)``.
class_ids (np.ndarray): Length-``N`` class index used to pick each image's cut-points.
cuts (np.ndarray): Per-class cut-points ``(n_classes, S-1)`` from
:func:`fit_thickness_thresholds`.
Returns:
np.ndarray: Length-``N`` thickness-bin indices.
"""
th = thickness(images)
class_ids = np.asarray(class_ids, dtype="int64")
out = np.empty(th.shape[0], dtype="int64")
for i in range(th.shape[0]):
out[i] = int(np.digitize(th[i], cuts[class_ids[i]]))
return out
def assign_modes(images: np.ndarray, class_ids: np.ndarray, cuts: np.ndarray,
cfg: MnistCfg) -> np.ndarray:
"""Map (image, class) to a mode index ``class * style_bins + thickness_bin``.
Args:
images (np.ndarray): Images ``(N,1,28,28)``.
class_ids (np.ndarray): Length-``N`` class index (true label, or oracle prediction).
cuts (np.ndarray): Per-class thickness cut-points.
cfg (MnistCfg): Supplies ``style_bins``.
Returns:
np.ndarray: Length-``N`` mode indices in ``[0, K)``.
"""
bins = thickness_bin(images, class_ids, cuts)
return np.asarray(class_ids, dtype="int64") * cfg.style_bins + bins
class MnistSampler:
"""Per-mode pools of real images; draws grounding/gen-0 observations by mode counts.
The image analogue of ``synthetic.render_modes``: given a length-``K`` count vector it
returns that many real images with the requested modes (drawn with replacement so rare
modes never run dry).
Args:
images (np.ndarray): The real image bank ``(N,1,28,28)`` to draw from.
modes (np.ndarray): Length-``N`` true mode of each image.
K (int): Number of modes.
"""
def __init__(self, images: np.ndarray, modes: np.ndarray, K: int) -> None:
self.images = images
self.K = K
self._pools = [np.flatnonzero(modes == k) for k in range(K)]
self._empty = [k for k, p in enumerate(self._pools) if p.size == 0]
def draw(self, counts: np.ndarray, rng: np.random.Generator) -> np.ndarray:
"""Draw images for a per-mode count vector.
Args:
counts (np.ndarray): Length-``K`` non-negative counts.
rng (np.random.Generator): Random source (draws with replacement).
Returns:
np.ndarray: Images ``(sum(counts), 1, 28, 28)`` in draw order by mode.
"""
counts = np.asarray(counts, dtype="int64")
idx_parts = []
for k in range(self.K):
c = int(counts[k])
if c <= 0:
continue
pool = self._pools[k]
if pool.size == 0:
continue # unpopulated mode: skip (rare)
idx_parts.append(rng.choice(pool, size=c, replace=True))
if not idx_parts:
return self.images[:0]
idx = np.concatenate(idx_parts)
return self.images[idx]

106
src/neural/mnist_loop.py Normal file
View file

@ -0,0 +1,106 @@
"""The MNIST analogue of ``generation_loop.run_generative_lineage``.
Identical Wright-Fisher generational step drift (``n`` samples from the parent model) +
grounding (``m`` fresh real samples, ``g = m/(n+m)``) + refit but observations are **images**:
the parent VAE *generates* the drift images, and grounding *draws real MNIST images* from the
per-mode pools (``MnistSampler``) instead of rendering token sequences. The frozen classifier
oracle reads each generation's mode distribution, logged with the *same* metric schema as every
other tier (``neural.evaluate.measure_metrics``), so the MNIST curves overlay the synthetic ones.
The oracle, sampler, and ``p*`` are built once by the experiment runner and passed in (training
the classifier and indexing the image pools is expensive and shared across all arms/replicates).
"""
from __future__ import annotations
from typing import Any, Mapping
import numpy as np
import pandas as pd
from knowledge.config import MetricsCfg, _sub
from knowledge.step import allocate_m, structured_multinomial
from knowledge.truth import TrueDist, uniform_init
from .config import MnistCfg, ModelCfg, NeuralDynamicsCfg
from .evaluate import measure_metrics
from .mnist_data import MnistSampler
from .mnist_vae import ConvVAEGenerator
from .oracle import Oracle
def _make_mnist_model(model_cfg: ModelCfg, cfg: MnistCfg, oracle: Oracle):
"""Construct the image generative model for the MNIST tier."""
if model_cfg.kind == "convvae":
return ConvVAEGenerator(cfg, model_cfg, oracle)
raise ValueError(f"unknown MNIST model kind {model_cfg.kind!r} (expected convvae)")
def run_mnist_lineage(cfg: Mapping[str, Any], seed: int, oracle: Oracle,
sampler: MnistSampler, td: TrueDist) -> pd.DataFrame:
"""Run one MNIST lineage and return per-generation metrics (same schema as Layer 1).
Args:
cfg (Mapping): Resolved config with ``mnist``, ``model``, ``dynamics``, ``generations``,
and optional ``metrics`` blocks.
seed (int): Replicate seed (statistically reproducible for the VAE).
oracle (Oracle): Prebuilt frozen classifier oracle.
sampler (MnistSampler): Prebuilt per-mode real-image pools.
td (TrueDist): ``p*``, tail mask, regions over the ``K`` modes.
Returns:
pd.DataFrame: One row per generation 0..T with the standard metric columns.
"""
mcfg = _sub(cfg["mnist"], MnistCfg)
model_cfg = _sub(cfg["model"], ModelCfg)
dyn_raw = dict(cfg.get("dynamics", {}))
from knowledge.config import GroundingCfg, RemintCfg
dynamics = NeuralDynamicsCfg(
n=dyn_raw.get("n", NeuralDynamicsCfg.n),
grounding=_sub(dyn_raw.get("grounding", {}), GroundingCfg),
remint=_sub(dyn_raw.get("remint", {}), RemintCfg),
)
metrics = _sub(cfg.get("metrics", {}), MetricsCfg)
generations = int(cfg.get("generations", 20))
p_star = td.p_star
tail_mask, regions, R = td.tail_mask, td.regions, mcfg.R
n = dynamics.n
grounding = dynamics.grounding
exercised = np.asarray(grounding.exercised) if grounding.exercised is not None else None
m_vector = allocate_m(grounding.m, R, grounding.policy, exercised)
rng = np.random.default_rng(seed)
p0 = uniform_init(mcfg.K) if mcfg.init == "uniform" else p_star.copy()
model = _make_mnist_model(model_cfg, mcfg, oracle)
def real_images(counts: np.ndarray) -> np.ndarray:
return sampler.draw(counts, rng)
# Generation 0: train on n real images drawn from p0 (the loop builds gen-0, not the VAE).
X0 = real_images(rng.multinomial(n, p0))
model.fit(X0, rng)
rows: list[dict] = []
def record(t: int) -> None:
row = {"generation": t}
row.update(measure_metrics(model.mode_distribution(rng), p_star, tail_mask,
regions, R, metrics))
rows.append(row)
record(0)
for t in range(1, generations + 1):
X_syn = model.sample(n, rng) # drift: n images from the parent
if m_vector is not None: # immigration: m real images
counts_real = structured_multinomial(m_vector, p_star, regions, grounding.policy, rng)
pool = np.concatenate([X_syn, real_images(counts_real)], axis=0)
else:
pool = X_syn
pupil = _make_mnist_model(model_cfg, mcfg, oracle)
pupil.fit(pool, rng)
model = pupil
record(t)
return pd.DataFrame(rows)

159
src/neural/mnist_oracle.py Normal file
View file

@ -0,0 +1,159 @@
"""The MNIST oracle — a frozen classifier standing in for "reality's no".
For the synthetic sandbox the oracle is exact (decode the identity segment). Real images have
no lossless barcode, so the oracle is a **frozen CNN** that predicts the digit class, combined
with a **deterministic** stroke-thickness bin, giving the mode
``class * style_bins + thickness_bin``. The CNN is trained once to high class accuracy and
cached; its **confusion matrix** over held-out real images is recorded as the measurement-noise
floor, so every MNIST result is read as a *sign relative to that floor*, never as an exact
magnitude.
"""
from __future__ import annotations
import hashlib
from pathlib import Path
import numpy as np
from .config import MnistCfg, OracleCfg
from .mnist_data import MnistData, assign_modes, fit_thickness_thresholds, thickness_bin
from .train import resolve_device, seed_everything, set_determinism
def _make_cnn(n_classes: int):
"""Build a small 2-conv MNIST classifier (lazy torch import)."""
import torch.nn as nn
return nn.Sequential(
nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 28 -> 14
nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 14 -> 7
nn.Flatten(), nn.Linear(32 * 7 * 7, 128), nn.ReLU(), nn.Linear(128, n_classes),
)
def _predict_classes(net, images: np.ndarray, device, batch: int = 1000) -> np.ndarray:
"""Batched argmax class prediction for a stack of images."""
import torch
net.eval()
x = torch.as_tensor(np.asarray(images, dtype="float32"), device=device)
out = np.empty(x.shape[0], dtype="int64")
with torch.no_grad():
for i in range(0, x.shape[0], batch):
out[i:i + batch] = net(x[i:i + batch]).argmax(1).cpu().numpy()
return out
def train_or_load_classifier(cfg: MnistCfg, ocfg: OracleCfg, data: MnistData,
seed, device=None):
"""Train (or load a cached) frozen digit classifier; return ``(net, checkpoint_sha256)``.
Args:
cfg (MnistCfg): Mode-truth config (for ``n_classes``).
ocfg (OracleCfg): Training/caching hyperparameters (``epochs``, ``lr``, ``cache``).
data (MnistData): Loaded MNIST split.
seed: Study-stream seed (int or SeedSequence).
device: Optional torch device; resolved from ``"auto"`` if ``None``.
Returns:
tuple: ``(net, sha256)`` the frozen network (eval mode) and the hex digest of its
checkpoint file (recorded in the manifest for provenance).
"""
import torch
set_determinism()
device = device or resolve_device("auto")
net = _make_cnn(cfg.n_classes).to(device)
cache = Path(ocfg.cache)
if cache.exists():
net.load_state_dict(torch.load(cache, map_location=device))
else:
g = seed_everything(seed)
net.train()
opt = torch.optim.Adam(net.parameters(), lr=ocfg.lr)
loss_fn = torch.nn.CrossEntropyLoss()
x = torch.as_tensor(data.train_x, device=device)
y = torch.as_tensor(data.train_y, device=device)
n, bs = x.shape[0], ocfg.batch_size
for _ in range(ocfg.epochs):
perm = torch.randperm(n, generator=g).to(device)
for i in range(0, n, bs):
idx = perm[i:i + bs]
loss = loss_fn(net(x[idx]), y[idx])
opt.zero_grad(); loss.backward(); opt.step()
cache.parent.mkdir(parents=True, exist_ok=True)
torch.save(net.state_dict(), cache)
net.eval()
for p in net.parameters():
p.requires_grad_(False)
return net, hashlib.sha256(cache.read_bytes()).hexdigest()
class ClassifierOracle:
"""Frozen-classifier oracle: image -> mode = predicted_class * style_bins + thickness_bin.
Args:
net: The frozen digit classifier.
cuts (np.ndarray): Per-class thickness cut-points ``(n_classes, S-1)``.
cfg (MnistCfg): Mode-truth config (``style_bins``, ``K``).
device: Torch device the classifier lives on.
"""
def __init__(self, net, cuts: np.ndarray, cfg: MnistCfg, device) -> None:
self.net = net
self.cuts = cuts
self.cfg = cfg
self.device = device
self.K = cfg.K
def classify(self, X: np.ndarray) -> np.ndarray:
"""Return the length-``n`` mode index for each image in ``X`` (``(n,1,28,28)``)."""
classes = _predict_classes(self.net, X, self.device)
bins = thickness_bin(X, classes, self.cuts)
return classes * self.cfg.style_bins + bins
def build_oracle(cfg: MnistCfg, ocfg: OracleCfg, data: MnistData, seed, device=None):
"""Fit thickness thresholds + the classifier and assemble the oracle.
Returns:
tuple: ``(oracle, cuts, ckpt_hash)`` the :class:`ClassifierOracle`, the per-class
thickness cut-points (needed to label real images as ground truth), and the classifier
checkpoint hash.
"""
device = device or resolve_device("auto")
cuts = fit_thickness_thresholds(data.train_x, data.train_y, cfg)
net, ckpt_hash = train_or_load_classifier(cfg, ocfg, data, seed, device)
return ClassifierOracle(net, cuts, cfg, device), cuts, ckpt_hash
def confusion_summary(oracle: ClassifierOracle, data: MnistData, cuts: np.ndarray,
cfg: MnistCfg) -> dict:
"""Measure the oracle's mode-level accuracy and confusion on the real test set.
True mode uses the *true* label + true thickness bin; predicted mode uses the oracle. The
diagonal rate is the measurement-noise floor the collapse metrics are read against.
Args:
oracle (ClassifierOracle): The assembled oracle.
data (MnistData): Loaded MNIST split (uses the test images/labels).
cuts (np.ndarray): Per-class thickness cut-points (for ground-truth modes).
cfg (MnistCfg): Mode-truth config.
Returns:
dict: ``mode_accuracy``, ``class_accuracy``, and ``confusion`` (``K x K`` nested list).
"""
true_modes = assign_modes(data.test_x, data.test_y, cuts, cfg)
pred_modes = oracle.classify(data.test_x)
pred_classes = pred_modes // cfg.style_bins
conf = np.zeros((cfg.K, cfg.K), dtype="int64")
for t, p in zip(true_modes, pred_modes):
conf[t, p] += 1
return {
"mode_accuracy": float(np.mean(pred_modes == true_modes)),
"class_accuracy": float(np.mean(pred_classes == data.test_y)),
"confusion": conf.tolist(),
}

125
src/neural/mnist_vae.py Normal file
View file

@ -0,0 +1,125 @@
"""Convolutional VAE generative model over MNIST images (the secondary-confirmation tier).
A small conv encoder maps a 28x28 image to a Gaussian latent; a conv decoder reconstructs it;
trained by the ELBO (binary cross-entropy reconstruction + ``beta``*KL). The VAE is the
*canonical* model in which generative collapse was first observed, so a VAE lineage that loses
its rare modes under dry self-training is direct evidence the effect is not a synthetic-sandbox
artefact.
Implements the same ``GenerativeModel`` protocol as the synthetic tier
(``fit`` / ``sample`` / ``mode_distribution``), so the metric readout is identical: sample
``n_eval`` images, classify them with the frozen oracle, normalise the mode histogram. Gen-0 is
built by the MNIST loop from the real-image sampler, so ``initialise`` is intentionally unused.
"""
from __future__ import annotations
import numpy as np
from .config import MnistCfg, ModelCfg
from .oracle import Oracle, measure_distribution
from .train import device_generator, resolve_device, seed_everything, set_determinism
def _make_conv_vae(latent: int):
"""Build a small conv VAE (lazy torch import)."""
import torch
import torch.nn as nn
class ConvVAE(nn.Module):
def __init__(self) -> None:
super().__init__()
self.enc = nn.Sequential(
nn.Conv2d(1, 32, 4, 2, 1), nn.ReLU(), # 28 -> 14
nn.Conv2d(32, 64, 4, 2, 1), nn.ReLU(), # 14 -> 7
nn.Flatten())
self.to_mu = nn.Linear(64 * 7 * 7, latent)
self.to_lv = nn.Linear(64 * 7 * 7, latent)
self.dec_in = nn.Linear(latent, 64 * 7 * 7)
self.dec = nn.Sequential(
nn.ConvTranspose2d(64, 32, 4, 2, 1), nn.ReLU(), # 7 -> 14
nn.ConvTranspose2d(32, 1, 4, 2, 1)) # 14 -> 28 (logits)
def encode(self, x):
h = self.enc(x)
return self.to_mu(h), self.to_lv(h)
def decode(self, z):
h = self.dec_in(z).view(-1, 64, 7, 7)
return self.dec(h) # logits
def forward(self, x):
mu, lv = self.encode(x)
z = mu + torch.exp(0.5 * lv) * torch.randn_like(lv)
logits = self.decode(z)
kl = -0.5 * torch.sum(1 + lv - mu.pow(2) - lv.exp(), dim=1).mean()
return logits, kl
return ConvVAE()
class ConvVAEGenerator:
"""Convolutional VAE over MNIST images.
Args:
cfg (MnistCfg): Mode-truth config (for ``K``).
model_cfg (ModelCfg): Hyperparameters (``latent``, ``beta``, ``epochs``, ``lr``,
``batch_size``, ``n_eval``, ``device``).
oracle (Oracle): Frozen classifier used to read the model's mode distribution.
"""
def __init__(self, cfg: MnistCfg, model_cfg: ModelCfg, oracle: Oracle) -> None:
self.cfg = cfg
self.mcfg = model_cfg
self.oracle = oracle
self.device = resolve_device(model_cfg.device)
self.net = None
set_determinism()
def initialise(self, p0: np.ndarray, rng: np.random.Generator) -> None:
"""Unused: the MNIST loop builds generation 0 from the real-image sampler."""
raise NotImplementedError("MNIST gen-0 is built by run_mnist_lineage via the sampler")
def fit(self, X: np.ndarray, rng: np.random.Generator) -> None:
"""Train (from scratch) on a batch of images ``X`` of shape ``(n,1,28,28)``."""
import torch
import torch.nn.functional as F
g = seed_everything(int(rng.integers(2 ** 31)))
net = _make_conv_vae(self.mcfg.latent).to(self.device)
net.train()
opt = torch.optim.Adam(net.parameters(), lr=self.mcfg.lr)
data = torch.as_tensor(np.asarray(X, dtype="float32"), device=self.device)
n, bs, beta = data.shape[0], self.mcfg.batch_size, self.mcfg.beta
for _ in range(self.mcfg.epochs):
perm = torch.randperm(n, generator=g).to(self.device)
for i in range(0, n, bs):
batch = data[perm[i:i + bs]]
logits, kl = net(batch)
recon = F.binary_cross_entropy_with_logits(logits, batch, reduction="none")
recon = recon.sum(dim=(1, 2, 3)).mean()
loss = recon + beta * kl
opt.zero_grad(); loss.backward(); opt.step()
net.eval()
self.net = net
def sample(self, n: int, rng: np.random.Generator) -> np.ndarray:
"""Draw ``n`` fresh images ``(n,1,28,28)`` in [0,1] via ``z ~ N(0,I)`` -> decode."""
import torch
if self.net is None:
raise RuntimeError("ConvVAEGenerator.sample called before fit")
gen = device_generator(int(rng.integers(2 ** 31)), self.device)
out = np.empty((n, 1, 28, 28), dtype="float32")
bs = 2000
with torch.no_grad():
for i in range(0, n, bs):
b = min(bs, n - i)
z = torch.randn(b, self.mcfg.latent, generator=gen, device=self.device)
out[i:i + b] = torch.sigmoid(self.net.decode(z)).cpu().numpy()
return out
def mode_distribution(self, rng: np.random.Generator) -> np.ndarray:
"""Estimate ``p_t`` by generate-and-classify over ``n_eval`` samples."""
X = self.sample(self.mcfg.n_eval, rng)
return measure_distribution(X, self.oracle, self.cfg.K)

View file

@ -230,7 +230,22 @@ C3 vertical claim deferred.*
while mean-distill stays flat (analytic + trained + rho=1 control). **grounding**: the reframed while mean-distill stays flat (analytic + trained + rho=1 control). **grounding**: the reframed
4-panel (forward-KL phase boundary, recovery, metric-choice). 4-panel (forward-KL phase boundary, recovery, metric-choice).
## Remaining **2026-07-05 — real-MNIST external-validity tier (`mnist_collapse`).**
- New image tier plugged into the existing contract (metrics/grounding/output are data-agnostic and
reused verbatim): `mnist_data.py` (load, per-class thickness bins, mode = class×thickness bijection,
`MnistSampler`), `mnist_oracle.py` (frozen CNN + deterministic thickness = `ClassifierOracle`,
confusion matrix), `mnist_vae.py` (`ConvVAEGenerator`), `mnist_loop.py` (`run_mnist_lineage`), plus
`kind=mnist_lineage` dispatch in `experiment.py`, `configs/neural/mnist_collapse.yaml`,
`figures/plot_mnist.py`, `figures/mnist_montage.py` (eyeball diagnostic), `MnistCfg`/`OracleCfg`.
- **Gates:** CNN mode accuracy **98.5%** (30×30 confusion matrix in the manifest = noise floor);
VAE gen-0 recovers full 30/30 support (over-smooths freq, KL≈0.5, no prior hole).
- **Result (4 reps):** dry (g=0) VAE **collapses to a single mode** (KL 0.5→18, support 30→1, tail
1.0→0.06, H→0); **g=0.1 holds all 30 modes** (KL≈0.6, full tail, H≈0.9). Collapse + grounding-rescue
confirmed on real images. VAE needs ~10% grounding vs synthetic ~5% (cf. the `grounding` finding).
**99 tests green** (+5 torchvision-gated). `make mnist` / `make env-mnist` (kept out of `make neural`).
## Remaining (all optional)
- [ ] **`region_matched`** grounding (R>1), **`remint`** re-mint gate (optional). - [ ] **`region_matched`** grounding (R>1), **`remint`** re-mint gate (optional).
- [ ] **VAE fidelity:** fix the prior-hole mismatch (KL-annealing / free-bits / larger latent) so it - [ ] **VAE fidelity:** fix the prior-hole mismatch (KL-annealing / free-bits / larger latent) so it

106
tests/test_mnist.py Normal file
View file

@ -0,0 +1,106 @@
"""MNIST-tier tests (skipped when torchvision is absent).
Gate the real-image confirmation tier: the oracle measures modes accurately, thickness binning
and mode assignment are well-formed, the sampler draws the requested modes, the VAE clears a
floor-aware gen-0 fidelity check, and a short dry lineage collapses more than a grounded one.
These are sign checks (the tier is confirmation-only), and they reuse the cached classifier so
they stay fast.
"""
from __future__ import annotations
import numpy as np
import pytest
pytest.importorskip("torchvision")
from knowledge.metrics import forward_kl # noqa: E402
from neural.config import MnistCfg, ModelCfg, OracleCfg # noqa: E402
from neural.mnist_data import ( # noqa: E402
MnistSampler, assign_modes, fit_thickness_thresholds, make_mnist_truth, thickness_bin,
)
from neural.mnist_loop import run_mnist_lineage # noqa: E402
from neural.mnist_oracle import build_oracle, confusion_summary # noqa: E402
from neural.mnist_vae import ConvVAEGenerator # noqa: E402
@pytest.fixture(scope="module")
def world():
"""Load MNIST + build the oracle/sampler/truth once (cached CNN -> fast)."""
from neural.mnist_data import load_mnist
cfg = MnistCfg()
data = load_mnist(cfg.data_root)
td = make_mnist_truth(cfg)
oracle, cuts, _ = build_oracle(cfg, OracleCfg(), data, seed=1)
modes = assign_modes(data.train_x, data.train_y, cuts, cfg)
sampler = MnistSampler(data.train_x, modes, cfg.K)
return cfg, data, td, oracle, cuts, sampler
def test_thickness_bins_and_mode_bijection(world):
cfg, data, td, oracle, cuts, sampler = world
sub = data.train_x[:2000]
lab = data.train_y[:2000]
bins = thickness_bin(sub, lab, cuts)
assert bins.min() >= 0 and bins.max() < cfg.style_bins # bins in range
modes = assign_modes(sub, lab, cuts, cfg)
assert np.array_equal(modes, lab * cfg.style_bins + bins) # exact bijection
assert modes.min() >= 0 and modes.max() < cfg.K
def test_sampler_draws_requested_modes(world):
cfg, data, td, oracle, cuts, sampler = world
counts = np.zeros(cfg.K, dtype=int)
counts[5] = 30 # mode 5 -> class 5//style_bins
counts[17] = 10 # mode 17 -> class 17//style_bins
imgs = sampler.draw(counts, np.random.default_rng(0))
assert imgs.shape == (40, 1, 28, 28)
# The sampler draws from the true per-mode pools, so the oracle should recover the two
# requested classes for the large majority of the draw (allowing the ~2% oracle error).
pred_classes = oracle.classify(imgs) // cfg.style_bins
want = {5 // cfg.style_bins, 17 // cfg.style_bins}
assert np.mean(np.isin(pred_classes, list(want))) > 0.9
def test_oracle_mode_accuracy_high(world):
cfg, data, td, oracle, cuts, sampler = world
conf = confusion_summary(oracle, data, cuts, cfg)
assert conf["mode_accuracy"] > 0.95 # solid measurement-noise floor
assert conf["class_accuracy"] > 0.97
def test_vae_gen0_recovers_full_support(world):
# gen-0 fidelity gate (floor-aware): the VAE must represent every mode (incl. the tail),
# else later tail loss would be underfitting, not collapse. Frequencies over-smooth (KL~0.5).
cfg, data, td, oracle, cuts, sampler = world
rng = np.random.default_rng(0)
mcfg = ModelCfg(kind="convvae", latent=32, epochs=15, n_eval=8000)
X0 = sampler.draw(rng.multinomial(6000, td.p_star), rng)
vae = ConvVAEGenerator(cfg, mcfg, oracle)
vae.fit(X0, rng)
p_hat = vae.mode_distribution(rng)
assert float(np.mean(p_hat[td.tail_mask] > 1e-9)) > 0.8 # tail represented
assert forward_kl(td.p_star, p_hat, 1e-9) < 1.2 # loose, floor-aware
def test_dry_collapses_more_than_grounded(world):
# The N1+N2 signs on real images: dry end-KL > grounded end-KL. Tiny config for speed.
cfg, data, td, oracle, cuts, sampler = world
base = {
"mnist": cfg.__dict__,
"model": ModelCfg(kind="convvae", latent=32, epochs=20, n_eval=6000).__dict__,
"generations": 8,
"metrics": {"kl_floor": 1e-9, "support_eps": 1e-9},
}
def run(g):
n = 6000
m = 0 if g == 0 else round(n * g / (1 - g))
c = {**base, "dynamics": {"n": n, "grounding": {"m": m, "policy": "proportional"}}}
return run_mnist_lineage(c, seed=0, oracle=oracle, sampler=sampler, td=td)
dry, grd = run(0.0), run(0.1)
assert dry["forward_kl"].iloc[-1] > dry["forward_kl"].iloc[0] + 1.0 # dry collapses
assert grd["forward_kl"].iloc[-1] < dry["forward_kl"].iloc[-1] # grounding arrests it
assert grd["support_size"].iloc[-1] > dry["support_size"].iloc[-1] # keeps more modes