Spaces:
Build error
Build error
File size: 3,999 Bytes
c10c824 ecd8f5c c10c824 ecd8f5c c10c824 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | """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"]))
|