Spaces:
Build error
Build error
File size: 5,580 Bytes
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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | """The coordinate candidate under the contract: fresh state per measurement, freeze as the
no-adaptation control, revert exact, and both modes producing a serializable report."""
from __future__ import annotations
import json
import pytest
import torch
from plastic.data.mechanisms import mechanism_batch, split_combinations
from plastic.eval.contract import ContractSpec, run_contract
from plastic.eval.coordinate_learner import CoordinateLearner
from plastic.model.coordinate import CoordinateConfig, CoordinateDynamics
def _tiny(**kw) -> CoordinateDynamics:
torch.manual_seed(0)
chunk = kw.pop("chunk", 16)
cfg = CoordinateConfig(d_model=32, n_heads=2, n_layers=2, chunk=chunk, **kw)
return CoordinateDynamics(cfg)
def _spec() -> ContractSpec:
# one 32-step episode per scored row; chunk 16 puts one update boundary inside it
return ContractSpec(seq_len=32, eval_batch=2, stream_episodes=4, probe_steps=20)
def test_contract_refuses_a_chunk_as_long_as_the_episode():
learner = CoordinateLearner(_tiny(chunk=32), mode="frozen")
with pytest.raises(ValueError, match="no fast-update boundary"):
run_contract(learner, _spec(), seed=1)
def test_rows_are_measured_independently():
learner = CoordinateLearner(_tiny(), mode="frozen")
b = _batch(4)
both = learner.step_mse(b, adapt=True)
from dataclasses import replace
alone = [
learner.step_mse(replace(b, inputs=b.inputs[i : i + 1], target_delta=b.target_delta[i : i + 1], worlds=b.worlds[i : i + 1], reset_flag=b.reset_flag[i : i + 1]), adapt=True)
for i in range(2)
]
assert torch.allclose(both[0], alone[0][0], atol=1e-6) and torch.allclose(both[1], alone[1][0], atol=1e-6)
def _batch(seed: int = 0):
train, heldout = split_combinations(k=2, n_heldout=5, seed=0)
g = torch.Generator().manual_seed(seed)
return mechanism_batch(2, seq_len=32, episodes_per_seq=1, combos=heldout, policy="gaussian", rng=g)
def test_step_mse_starts_fresh_every_call_and_freeze_differs_from_adapt():
learner = CoordinateLearner(_tiny(), mode="frozen")
b = _batch()
a1 = learner.step_mse(b, adapt=True)
a2 = learner.step_mse(b, adapt=True)
f1 = learner.step_mse(b, adapt=False)
assert a1.shape == (2, 32) and torch.equal(a1, a2), "a measurement must not depend on earlier calls"
assert not torch.equal(a1, f1), "fast updates must change predictions after the first chunk"
# within the first chunk nothing has been proposed yet, so adapt and freeze agree there
assert torch.allclose(a1[:, :16], f1[:, :16])
def test_fast_signals_summary_reports_proposals_only_when_stepped():
learner = CoordinateLearner(_tiny(), mode="frozen")
b = _batch(3)
assert learner.fast_signals_summary() is None
learner.step_mse(b, adapt=True)
s = learner.fast_signals_summary()
# a 32-step episode with chunk 16 crosses two boundaries, at steps 16 and 32
assert s["boundaries"] == 2 and s["stepped_fraction"] == 1.0
assert s["inner_loss_before"] is not None and s["inner_loss_after"] is not None
assert len(s["dW_norm_by_layer"]) == 2 == len(s["dtheta_norm_by_layer"]) == len(s["eta_by_layer"])
assert s["r_peak_inf"] <= s["bound_peak"] + 1e-6
learner.step_mse(b, adapt=False)
f = learner.fast_signals_summary()
assert f["stepped_fraction"] == 0.0 and f["inner_loss_before"] is None and f["dW_norm_by_layer"] is None
def test_fast_updates_off_makes_adapt_equal_freeze():
learner = CoordinateLearner(_tiny(fast_updates=False), mode="frozen")
b = _batch(1)
assert torch.allclose(learner.step_mse(b, adapt=True), learner.step_mse(b, adapt=False))
@pytest.mark.parametrize("mode", ["frozen", "continued"])
def test_contract_runs_and_reverts_exactly(mode):
model = _tiny()
learner = CoordinateLearner(model, mode=mode, lr=1e-3, steps=2)
report = run_contract(learner, _spec(), seed=3)
json.dumps(report)
assert report["revert"]["ok"] is True, report["revert"]
assert report["compute"]["parameters"] == model.num_params()
if mode == "frozen":
assert report["acceptance"]["n_good"] == 0
assert all(row["delta_mse"] == 0.0 for row in report["transfer"].values())
else:
assert report["acceptance"] == {"accepted_good": 1.0, "refused_bad": 0.0, "n_good": 2, "n_bad": 1}
assert report["decisions"][0]["accepted"] is True
def _moved_keys(model: CoordinateDynamics, batch) -> set[str]:
"""Which fast-parameter keys changed after one adapting pass from the initial state."""
from plastic.model.coordinate import THETA_KEY, W_KEYS
x, y = batch.inputs, batch.target_delta
state0 = model.init_state(int(x.shape[0]))
with torch.no_grad():
_, state1, _ = model(x, state0, mode="chunk", target_delta=y, commit=True)
moved: set[str] = set()
for layer0, layer1 in zip(state0.omega, state1.omega):
for k in tuple(W_KEYS) + (THETA_KEY,):
if not torch.equal(layer0[k], layer1[k]):
moved.add("theta" if k == THETA_KEY else "W")
return moved
def test_decay_only_and_coordinates_only_switches_move_only_their_parameters():
# at initialization the fast step barely changes predictions (about 1e-5), so distinctness
# is checked on which fast parameters move, not on the outputs
b = _batch(2)
assert _moved_keys(_tiny(), b) == {"W", "theta"}
assert _moved_keys(_tiny(adapt_W=False), b) == {"theta"}
assert _moved_keys(_tiny(adapt_theta=False), b) == {"W"}
assert _moved_keys(_tiny(fast_updates=False), b) == set()
|