#!/usr/bin/env python3 """ Does inference-depth degradation come from step-embedding extrapolation (B) or from genuine recurrent saturation (A)? HYPOTHESIS The core applies `step_emb[min(step, max_steps)]`. Beyond `max_steps` every additional step receives the SAME embedding as the last trained step, so the core gets an out-of-distribution "which step am I on" signal. If that is what breaks deep inference, then removing the step embedding entirely (making the core genuinely step-invariant) should change the shape of the depth curve. DESIGN (weights are FIXED and IDENTICAL in every condition) condition `clamped` : stock behaviour, step_emb[min(step, max_steps)] condition `zeroed` : step_emb forced to 0 at every step -> step-invariant condition `frozen_k` : every step uses step_emb[k] for a fixed in-distribution k, isolating "OOD signal" from "no signal at all" Each condition is evaluated at depths 1/2/4/8/16/32 against a FIXED opponent pool (random-legal and one-ply material), with identical games, temperature, seed and adjudication. Only the depth and the embedding policy vary. DISCRIMINATION A (saturation) -> all conditions degrade at the same depth B (OOD step embeddings) -> `zeroed`/`frozen_k` keep improving where `clamped` degrades Anything else -> report as ambiguous Also records a weight-free signal: policy agreement with the depth-1 policy and mean entropy per depth, to show whether deep inference is drifting or collapsing. """ from __future__ import annotations import argparse import contextlib import json import os import sys import time import torch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from tinychess.arena import (MaterialAgent, RandomAgent, elo_from_score, head_to_head, wilson_interval) from tinychess.arena import policy_diagnostics from tinychess.batching import encode_batch from tinychess.checkpoint import load_model_only @contextlib.contextmanager def step_embedding_mode(model, mode: str, k: int = 0): """Temporarily replace the core's step embedding policy. Weights unchanged.""" core = model.core original = core.step_emb.data.clone() try: if mode == "clamped": pass # stock behaviour elif mode == "zeroed": core.step_emb.data.zero_() # step-invariant core elif mode == "frozen": row = original[min(k, original.shape[0] - 1)].clone() core.step_emb.data[:] = row # same in-distribution signal else: raise ValueError(mode) yield finally: core.step_emb.data.copy_(original) def make_boards(n=48, seed=0): """Fixed, reproducible mid-game position set (same recipe as exp_compute_scaling).""" import random import chess rng = random.Random(seed) boards = [] while len(boards) < n: b = chess.Board() for _ in range(rng.randint(0, 40)): ms = list(b.legal_moves) if not ms: break b.push(rng.choice(ms)) if not b.is_game_over(): boards.append(b) return boards def policy_stats(model, boards, steps): """Top-1 move ids, entropy, top-1 prob and illegal mass at a given depth.""" pb = encode_batch(boards) with torch.no_grad(): out = model(**pb.model_args(), steps=steps) logits = out["logits"] d = policy_diagnostics(model, boards, steps=steps) return logits.argmax(-1), d["mean_entropy"], d["mean_top1_prob"], d["illegal_prob_mass"] def main(): ap = argparse.ArgumentParser() ap.add_argument("--ckpt", required=True) ap.add_argument("--depths", default="1,2,4,8,16,32") ap.add_argument("--modes", default="clamped,zeroed,frozen") ap.add_argument("--frozen-k", type=int, default=-1, help="step index for 'frozen' mode; -1 = last trained step") ap.add_argument("--games", type=int, default=40) ap.add_argument("--max-plies", type=int, default=140) ap.add_argument("--temperature", type=float, default=0.3) ap.add_argument("--adjudicate", default="material", choices=["material", "none"]) ap.add_argument("--seed", type=int, default=0) ap.add_argument("--out", default="runs/step_embedding.json") a = ap.parse_args() depths = [int(x) for x in a.depths.split(",")] modes = [m.strip() for m in a.modes.split(",")] model, cfg = load_model_only(a.ckpt) rep = model.param_report() # `cfg.max_steps` is the TRAINED depth. Capture it before any mutation: # model.forward clamps n to cfg.max_steps, so without lifting the cap every # depth > max_steps silently runs only max_steps iterations. trained_max_steps = int(cfg.max_steps) max_steps = trained_max_steps if max(depths) > trained_max_steps: import torch.nn as nn old = model.core.step_emb.data new = torch.zeros(max(depths) + 1, old.shape[1]) new[:old.shape[0]] = old new[old.shape[0]:] = old[-1] # stock behaviour: repeat last row model.core.step_emb = nn.Parameter(new) model.cfg.max_steps = max(depths) model.core.cfg.max_steps = max(depths) print(f"[lifted] step cap {trained_max_steps} -> {max(depths)} " f"(step_emb rows {old.shape[0]} -> {new.shape[0]})") k = a.frozen_k if a.frozen_k >= 0 else trained_max_steps print(f"ckpt={a.ckpt}") print(f"params={rep['total']:,} trained max_steps={trained_max_steps} frozen_k={k}") print(f"depths={depths} modes={modes} games={a.games} adj={a.adjudicate}\n") boards = make_boards(48, seed=a.seed) opponents = {"random_legal": RandomAgent(), "material_greedy": MaterialAgent()} res = {"checkpoint": a.ckpt, "params": rep["total"], "trained_max_steps": trained_max_steps, "frozen_k": k, "args": vars(a), "conditions": {}} ref_top1 = {} for mode in modes: res["conditions"][mode] = {} with step_embedding_mode(model, mode, k): for d in depths: entry = {} t0 = time.time() top1, ent, top1p, illegal = policy_stats(model, boards, d) entry["mean_entropy"] = round(ent, 4) entry["mean_top1_prob"] = round(top1p, 4) entry["illegal_prob_mass"] = float(illegal) entry["ms_per_position"] = round( 1000 * (time.time() - t0) / len(boards), 3) if d == depths[0]: ref_top1[mode] = top1 entry["top1_agreement_vs_depth1"] = round( float((top1 == ref_top1[mode]).float().mean()), 4) for oname, opp in opponents.items(): h = head_to_head(model, opp, n_games=a.games, steps=d, max_plies=a.max_plies, temperature=a.temperature, adjudicate=a.adjudicate, seed=100 + a.seed) ci = wilson_interval(h["wins"], h["draws"], h["games"]) entry[oname] = { "score": round(h["score"], 4), "ci95": [round(ci[0], 3), round(ci[1], 3)], "elo": round(elo_from_score(h["score"]), 1), "W": h["wins"], "D": h["draws"], "L": h["losses"], "mean_plies": round(h["mean_plies"], 1), } res["conditions"][mode][str(d)] = entry r = entry["random_legal"]; mt = entry["material_greedy"] print(f" [{mode:8s}] d={d:3d} rand={r['score']:.3f}" f"{str(r['ci95']):16s} mat={mt['score']:.3f}" f"{str(mt['ci95']):16s} ent={ent:.2f} " f"agree_d1={entry['top1_agreement_vs_depth1']:.2f} " f"{entry['ms_per_position']:.1f}ms") # ---- verdict ---- def curve(mode, opp): return [(d, res["conditions"][mode][str(d)][opp]["score"]) for d in depths] verdict = {} for opp in opponents: v = {} for mode in modes: c = curve(mode, opp) best_d, best_s = max(c, key=lambda t: t[1]) deep = [s for d, s in c if d > trained_max_steps] shallow = [s for d, s in c if d <= trained_max_steps] ms = (sum(shallow) / len(shallow)) if shallow else None md = (sum(deep) / len(deep)) if deep else None v[mode] = {"curve": c, "best_depth": best_d, "best_score": best_s, "n_shallow": len(shallow), "n_deep": len(deep), "mean_shallow": round(ms, 4) if ms is not None else None, "mean_deep": round(md, 4) if md is not None else None, "deep_minus_shallow": (round(md - ms, 4) if (ms is not None and md is not None) else None)} verdict[opp] = v res["verdict"] = verdict print("\n=== deep-vs-shallow (mean score beyond trained max_steps minus within) ===") for opp, v in verdict.items(): print(f" {opp}:") for mode, d in v.items(): if d["deep_minus_shallow"] is None: print(f" {mode:8s} best_depth={d['best_depth']:3d} " f"shallow={d['mean_shallow']} deep=n/a " f"(no depths beyond trained max_steps)") else: print(f" {mode:8s} best_depth={d['best_depth']:3d} " f"shallow={d['mean_shallow']:.3f} deep={d['mean_deep']:.3f} " f"delta={d['deep_minus_shallow']:+.3f}") os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True) json.dump(res, open(a.out, "w"), indent=2) print(f"\n[saved] {a.out}") if __name__ == "__main__": main()