Spaces:
Sleeping
Sleeping
| """The ablation runner: every variant builds, trains a step, is scored, and the collected table | |
| labels each model's no-adaptation intervention.""" | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import pytest | |
| import torch | |
| from plastic.eval.contract import ContractSpec | |
| from scripts.experiments.coordinate_ablation import VARIANTS, build, collect, outer_loss_value, run_variant | |
| def test_post_boundary_loss_excludes_the_first_chunk_for_every_model(): | |
| from plastic.data.mechanisms import mechanism_batch, split_combinations | |
| train, _ = split_combinations(k=2, n_heldout=5, seed=0) | |
| g = torch.Generator().manual_seed(0) | |
| b = mechanism_batch(2, seq_len=32, episodes_per_seq=1, combos=train, policy="gaussian", rng=g) | |
| for name in ("full", "delta_baseline"): | |
| m = build(name, d_model=32, n_heads=2, n_layers=1, chunk=16, seed=0) | |
| with torch.no_grad(): | |
| whole = outer_loss_value(m, b.inputs, b.target_delta, outer_loss="all", chunk=16) | |
| post = outer_loss_value(m, b.inputs, b.target_delta, outer_loss="post_boundary", chunk=16) | |
| pred, _, _ = m(b.inputs, target_delta=b.target_delta) if name == "full" else m(b.inputs) | |
| per_step = (pred - b.target_delta).pow(2).mean(-1) | |
| assert whole == pytest.approx(float(per_step.mean())) | |
| assert post == pytest.approx(float(per_step[:, 16:].mean())) | |
| with pytest.raises(ValueError, match="longer than one chunk"): | |
| outer_loss_value(build("full", d_model=32, n_heads=2, n_layers=1, chunk=16, seed=0), b.inputs[:, :16], b.target_delta[:, :16], outer_loss="post_boundary", chunk=16) | |
| def test_every_variant_builds_with_its_switches(): | |
| for name in VARIANTS: | |
| m = build(name, d_model=32, n_heads=2, n_layers=1, chunk=16, seed=0) | |
| assert sum(p.numel() for p in m.parameters()) > 0 | |
| if VARIANTS[name] is not None: | |
| for k, v in VARIANTS[name].items(): | |
| assert getattr(m.cfg, k) == v, (name, k) | |
| wide = build("full", d_model=32, n_heads=2, n_layers=1, chunk=16, seed=0, eta_max=10.0, eta_init=1.0) | |
| assert wide.cfg.eta_max == 10.0 and wide.cfg.eta_init == 1.0 | |
| assert build("delta_baseline", d_model=32, n_heads=2, n_layers=1, chunk=16, seed=0, eta_max=10.0) is not None | |
| def test_run_variant_and_collect_smoke(tmp_path): | |
| args = argparse.Namespace( | |
| steps=2, batch=2, seq_len=32, episodes=1, lr=1e-3, log_every=1, d_model=32, n_heads=2, n_layers=1, | |
| chunk=16, seed=0, device="cpu", outer_loss="post_boundary", | |
| ) | |
| # the speed horizon must reach past the coordinate variants' first boundary (chunk 16) | |
| spec = ContractSpec(seq_len=32, eval_batch=2, stream_episodes=4, probe_steps=20) | |
| for v in ("full", "decay_only", "delta_baseline"): | |
| r = run_variant(v, args, str(tmp_path), spec=spec) | |
| assert r["contract"]["revert"]["ok"] is True | |
| assert len(r["train"]["log"]) >= 2 and r["train"]["log"][-1]["step"] == 2 | |
| assert (r["train"]["log"][-1]["eta"] is None) == (v == "delta_baseline") | |
| assert (tmp_path / "full.pt").exists() and (tmp_path / "delta_baseline.pt").exists() | |
| full = json.load(open(tmp_path / "full.json")) | |
| assert full["fast_signals"]["stepped_fraction"] == 1.0 and full["checkpoint"] == "full.pt" | |
| assert json.load(open(tmp_path / "delta_baseline.json"))["fast_signals"] is None | |
| readme = collect(str(tmp_path)) | |
| assert "| full |" in readme and "| decay_only |" in readme and "| delta_baseline |" in readme | |
| assert "→" in readme # the inner-loss diagnostic column is filled for coordinate variants | |
| assert "| full | post_boundary |" in readme and full["train"]["outer_loss"] == "post_boundary" | |
| assert "fast parameters frozen (freeze=True)" in readme and "writes disabled (beta_scale=0)" in readme | |
| saved = json.load(open(tmp_path / "full.json")) | |
| assert saved["no_adapt_label"].startswith("fast parameters frozen") | |
| assert torch.isfinite(torch.tensor(saved["train"]["log"][-1]["loss"])) | |