diff --git a/docs/design/quality-experiment-matrix.md b/docs/design/quality-experiment-matrix.md index f613615e..a5eda6bf 100644 --- a/docs/design/quality-experiment-matrix.md +++ b/docs/design/quality-experiment-matrix.md @@ -1693,3 +1693,4 @@ self-distillation, trajectory RL) are in | E63 | Gate calibration | ECE / selective accuracy / abstention on `FastPathGate` | proposed | | E64 | Trajectory-aligned RL | MDPO/d1-style on intermediate MaskGIT states | proposed | | E65 | Schema generalization | Held-out schemas / rename / `toy-layout` transfer | proposed | +| VSS3-02 | Cost-to-go energy ranker | `cost_to_go_loss_weight` + `cost_to_go_hidden_dim` | wired | diff --git a/docs/design/research-lineage.md b/docs/design/research-lineage.md index 5c726fb6..f0f88eb2 100644 --- a/docs/design/research-lineage.md +++ b/docs/design/research-lineage.md @@ -527,7 +527,7 @@ batch here avoids treating adjacent papers as independent evidence for one lever | [DeepCoder](https://arxiv.org/abs/1611.01989) (Balog et al., ICLR 2017) | **Adjacent** | Learned search-guidance / candidate ranking; the contract keeps learned scores soft (`rank_forest`) and does not reimplement its attribute predictor or DSL | | Counterexample-guided (neural) synthesis / CEGIS — see the *LLM-Modulo / CEGIS planning* anchor above and R19/R23 in [lattice-recursive-search.md](lattice-recursive-search.md) | **Adapted boundary** | Deduction-vs-decision split and local nogoods adopt the counterexample → refinement principle; neural synthesis training is not reimplemented | | [egg / e-graphs](https://arxiv.org/abs/2004.03082) (Willsey et al., POPL 2021) | **Adjacent** | Post-realization canonicalization is motivated by equality saturation, but the contract is not an e-graph / equality-saturation engine (cf. `iter-e252-canonicalizer-20260716.md`) | -| EDLM — energy-based diffusion language model (2024) | **Adjacent** | Energy/score-based candidate ranking is analogous to the soft-scoring layer; no EDLM training or energy head is implemented or assumed | +| EDLM — energy-based diffusion language model (2024) | **Adapted boundary** | VSS3-02 wires a small regression+pairwise cost-to-go head on TwoTower that ranks live solver candidates (not a full EDLM); fixture-level head-only training, no production energy model or ship claim | | [TreeDiff](https://arxiv.org/abs/2508.01473) (Zeng et al., 2025) | **Adapted boundary** | Tree-edit diffusion informs late realization against AST/choice IR; architecture and training remain future empirical work (as R13) | | [Lattice Deduction Transformer](https://arxiv.org/abs/2605.08605) (Davis et al., 2026) | **Adapted** | Monotone lattice projection plus rollback are the closure/deduction model; LDT architecture, alpha supervision, and training remain future work (as R7) | diff --git a/docs/design/verified-scope-solver.md b/docs/design/verified-scope-solver.md index 3e371458..af5856e2 100644 --- a/docs/design/verified-scope-solver.md +++ b/docs/design/verified-scope-solver.md @@ -286,6 +286,35 @@ controller validates the returned sequence is a permutation and rejects any missing/extra/duplicate); `CERTIFIED_UNSAT` is impossible when any required branch was UNKNOWN or budget-truncated; every `SOLVED` carries a final verifier report. +## Implemented cost-to-go energy ranker (VSS3-02 / SLM-70) + +[`dsl/solver/energy_ranker.py`](../../src/slm_training/dsl/solver/energy_ranker.py) +adds `CandidateEnergyInput`, `CandidateEnergyOutput`, and `EnergyCandidateRanker`. +The ranker wraps a `CandidateEnergyScorer` and returns a **permutation** of the +exact live candidates supplied by the solver. It fail-closed falls back to the +canonical value order when the scorer returns the wrong length, non-finite +values, a membership change, or raises an exception. + +[`models/twotower.py`](../../src/slm_training/models/twotower.py) adds an optional +`cost_to_go_head` auxiliary MLP, plus `score_candidates(...)` and +`cost_to_go_loss(...)`. The loss consumes `CandidateCostRow` records from the +VSS3-01 replay-verified supervision corpus and: + +- masks rows where `cost_observed=False` from the regression loss, +- adds a pairwise ranking loss over candidates sharing the same `(state_fingerprint, hole_id)`, +- reports `cost_to_go_regression_loss`, `cost_to_go_pairwise_loss`, and `cost_to_go_observed_rows`. + +[`harnesses/model_build/cost_to_go_train.py`](../../src/slm_training/harnesses/model_build/cost_to_go_train.py) +provides a head-only trainer that freezes the TwoTower backbone and optimizes only +the new head, analogous to `trust_train.py`. + +Runtime contract: + +- Exact closure computes the hard live set **before** energy scoring. +- Energy output may change candidate order, expanded nodes, backtracks, and latency only. +- Invalid scorer output triggers the configured deterministic fallback and records the fallback; membership is unchanged. +- The ranker never certifies support, suppresses `UNKNOWN`, or bypasses the final verifier. + ## Reference support semantics | Verdict | Requirement | Removal permitted? | @@ -401,6 +430,7 @@ the closure. | `HoleId`, `SolverBounds`, `DomainValue`, `HoleDomain`, `FiniteDomainState` (`dsl/solver/state.py`) | Immutable JSON-safe finite-domain carrier with canonical hard-state identity and monotone operations. | **Implemented by VSS0-03.** Torch-free and not invoked by default; soft scores and proof artifacts remain outside the state. | | `completion_forest_state` (`dsl/solver/adapters.py`) | One-hole projection of full compiler completion paths plus coverage and `UNKNOWN` provenance. | **Implemented by VSS0-03.** Retains the compiler as owner; a singleton solves only the projection. | | `verification capsule`, `proof certificate` | Not implemented yet. | **Future VSS issues.** Capsule solving and replayable proof ownership are not duplicated here. | +| `EnergyCandidateRanker` / `cost_to_go_head` (`dsl/solver/energy_ranker.py`, `models/twotower.py`) | Optional learned cost-to-go ranking over the exact live candidate set. | **Implemented by VSS3-02.** Permutation-only; fail-closed fallback on invalid output. No support certification or verifier bypass. | ## End-to-end example (partial coverage → live `UNKNOWN`) @@ -527,3 +557,16 @@ integration, no learned energy, no persistent clause DB, and **no ship claim**. Verified: `python -m pytest tests/test_dsl/test_solver_controller.py tests/test_dsl/test_lattice_search.py -q` (94 in the full solver+compat suite) and `python -m scripts.repo_policy`. + +2026-07-18 — VSS3-02 / SLM-70 implements the cost-to-go energy scorer wiring: +`EnergyCandidateRanker` enforces permutation-only ordering with fail-closed +fallback, `TwoTowerModel` grows an optional `cost_to_go_head` trained from +`CandidateCostRow` records, and `cost_to_go_train.py` provides a head-only trainer. +Pinned by tiny closed fixtures (`tests/test_dsl/test_energy_ranker.py`, +`tests/test_models/test_cost_to_go.py`): ranker permutation/fallback, lower-cost +candidates receiving lower energy after overfit, `cost_observed=False` rows masked +from regression, and old checkpoints loading with the new head allowed as missing. +No full train/eval/benchmark/checkpoint was run, and **no ship or speed claim** is +made. Verified: `python -m pytest tests/test_dsl/test_energy_ranker.py +tests/test_models/test_cost_to_go.py tests/test_dsl/test_solver_controller.py -q` +(45 passed) and `python -m scripts.repo_policy`. diff --git a/docs/design/vss3-02-cost-to-go.md b/docs/design/vss3-02-cost-to-go.md new file mode 100644 index 00000000..5931ecc0 --- /dev/null +++ b/docs/design/vss3-02-cost-to-go.md @@ -0,0 +1,52 @@ +# VSS3-02: Train a cost-to-go energy scorer that ranks live solver candidates only + +**Issue:** SLM-70 +**Status:** wiring / fixture evidence. No train, eval, benchmark, model, checkpoint, or ship claim. + +## What was added + +A cost-to-go energy scorer that plugs into the verified-scope solver as a permutation-only ranker. + +Solver-side contract in `src/slm_training/dsl/solver/energy_ranker.py`: + +- `CandidateEnergyInput` / `CandidateEnergyOutput` — typed scoring envelope. +- `CandidateEnergyScorer` protocol and `EnergyCandidateRanker` implementing `CandidateRanker`. +- Fail-closed fallback: wrong-length, non-finite, membership-changing, or exception-throwing scorer output returns the canonical value order and records the fallback. + +Model-side head in `src/slm_training/models/twotower.py`: + +- `TwoTowerConfig` fields `cost_to_go_loss_weight` and `cost_to_go_hidden_dim`. +- Optional `cost_to_go_head` MLP built with `isolated_aux_init` when enabled. +- `score_candidates(state, hole_id, values, context_prompt=None)` returns `CandidateEnergyOutput`. +- `cost_to_go_loss(rows, prompts=None)` regresses `log1p(nodes + verifier_calls + backtracks + depth)` for rows with `cost_observed=True` and adds a pairwise ranking loss over candidates sharing the same `(state_fingerprint, hole_id)`. +- `UNKNOWN`/censored rows (`cost_observed=False`) are masked from the regression loss. +- Prefix `"cost_to_go_head."` added to optimizer groups and `allowed_missing` checkpoint set. + +Harness wiring: + +- `src/slm_training/harnesses/model_build/config.py` and `factory.py` propagate the new config fields. +- `scripts/train_model.py` accepts `--cost-to-go-loss-weight` and `--cost-to-go-hidden-dim`. +- `src/slm_training/harnesses/model_build/cost_to_go_train.py` provides head-only training (freeze backbone, optimize head) and a checkpoint-save path. +- `scripts/train_cost_to_go.py` thin CLI over the head-only trainer. + +## Verified + +- `ruff check` passes on touched files. +- `python -m compileall` passes. +- `pytest tests/test_dsl/test_energy_ranker.py tests/test_models/test_cost_to_go.py tests/test_dsl/test_solver_controller.py -q` → 45 passed. +- `python -m scripts.repo_policy` ok. +- `git diff --check` clean. +- `.githooks/check-changed` → all checks passed (590 passed across targeted suites). + +## Design boundaries preserved + +- `EnergyCandidateRanker` validates that output is a permutation of the exact live values; it cannot add or drop candidates. +- The energy head is trained only from replay-verified `CandidateCostRow` records; `cost_observed=False` rows are not treated as negatives. +- No support certification, verifier bypass, or unconstrained fallback is introduced. +- Old checkpoints load cleanly with the new head allowed as missing. + +## Caveats + +- This is fixture wiring only: the head trains on tiny synthetic rows in tests and has not been trained on a real VSS3-01 corpus or evaluated on a solver benchmark. +- Candidate features are currently the raw `token_ids` payload from `DomainValue`; richer state/hole features are future work. +- No checkpoint ship, model-card, or readiness claim is made. diff --git a/scripts/train_cost_to_go.py b/scripts/train_cost_to_go.py new file mode 100644 index 00000000..8d9ea1f8 --- /dev/null +++ b/scripts/train_cost_to_go.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python3 +"""Train only the VSS3-02 cost-to-go head from a solver-supervision corpus.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from slm_training.harnesses.model_build.cost_to_go_train import ( + train_cost_to_go_from_paths, +) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--rows", type=Path, required=True) + parser.add_argument("--out-dir", type=Path, required=True) + parser.add_argument("--steps", type=int, default=40) + parser.add_argument("--batch-size", type=int, default=8) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--device", default="cpu") + args = parser.parse_args(argv) + summary = train_cost_to_go_from_paths( + checkpoint=args.checkpoint, + rows_path=args.rows, + out_dir=args.out_dir, + steps=args.steps, + batch_size=args.batch_size, + lr=args.lr, + device=args.device, + ) + print(summary) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/train_model.py b/scripts/train_model.py index 766ebe18..b7bd0c41 100644 --- a/scripts/train_model.py +++ b/scripts/train_model.py @@ -571,6 +571,18 @@ def main(argv: list[str] | None = None) -> int: default=0.0, help="Cheap structural force-align aux loss weight (0 disables).", ) + parser.add_argument( + "--cost-to-go-loss-weight", + type=float, + default=0.0, + help="VSS3-02: weight for solver cost-to-go regression+pairwise loss (0 disables).", + ) + parser.add_argument( + "--cost-to-go-hidden-dim", + type=int, + default=0, + help="VSS3-02: hidden dimension for the cost-to-go MLP head (0 disables).", + ) parser.add_argument( "--grad-accum", type=int, @@ -798,6 +810,8 @@ def main(argv: list[str] | None = None) -> int: fuse_ltr_loss=fuse_ltr, grammar_fastpath=True, fastpath_aux_weight=args.fastpath_aux_weight, + cost_to_go_loss_weight=args.cost_to_go_loss_weight, + cost_to_go_hidden_dim=args.cost_to_go_hidden_dim, grammar_trust_model=args.grammar_trust_model, noise_rate=args.noise_rate, eval_every=args.eval_every, diff --git a/src/slm_training/dsl/solver/__init__.py b/src/slm_training/dsl/solver/__init__.py index cb91a09b..9f60204e 100644 --- a/src/slm_training/dsl/solver/__init__.py +++ b/src/slm_training/dsl/solver/__init__.py @@ -25,6 +25,13 @@ default_hole_selector, search, ) +from slm_training.dsl.solver.energy_ranker import ( + CandidateEnergyInput, + CandidateEnergyOutput, + CandidateEnergyScorer, + EnergyCandidateRanker, + make_stub_energy_scorer, +) from slm_training.dsl.solver.closure import ( CertifiedDeduction, ClosureCounters, @@ -54,8 +61,12 @@ __all__ = [ "BaselineRanker", + "CandidateEnergyInput", + "CandidateEnergyOutput", + "CandidateEnergyScorer", "CandidateRanker", "CertifiedDeduction", + "EnergyCandidateRanker", "ClosureCounters", "ClosureResult", "DomainValue", @@ -93,5 +104,6 @@ "completion_forest_state", "default_query_order", "exact_closure", + "make_stub_energy_scorer", "replay_support_certificate", ] diff --git a/src/slm_training/dsl/solver/energy_ranker.py b/src/slm_training/dsl/solver/energy_ranker.py new file mode 100644 index 00000000..b2b4e285 --- /dev/null +++ b/src/slm_training/dsl/solver/energy_ranker.py @@ -0,0 +1,156 @@ +"""Cost-to-go energy ranker for the verified solver (VSS3-02 / SLM-70). + +The ranker has no authority over candidate membership, support, or final +verification. It only orders the exact live candidates supplied by the solver. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +from slm_training.dsl.solver.controller import CandidateRanker +from slm_training.dsl.solver.state import DomainValue, FiniteDomainState, HoleId + +if TYPE_CHECKING: + import torch + + +@dataclass(frozen=True) +class CandidateEnergyInput: + """Features for one scoring call.""" + + state_fingerprint: str + state_features: torch.Tensor + hole_features: torch.Tensor + candidate_ids: tuple[DomainValue, ...] + candidate_features: torch.Tensor + + +@dataclass(frozen=True) +class CandidateEnergyOutput: + """Energy scores for the supplied candidates.""" + + energies: torch.Tensor # shape [num_live_candidates] + candidate_ids: tuple[DomainValue, ...] + scorer_id: str + + +@runtime_checkable +class CandidateEnergyScorer(Protocol): + """Callable that scores live candidates without altering membership.""" + + def __call__( + self, + state: FiniteDomainState, + hole_id: HoleId, + values: tuple[DomainValue, ...], + ) -> CandidateEnergyOutput: ... + + +class EnergyCandidateRanker(CandidateRanker): + """Permutation-only ranker backed by an energy scorer. + + Falls back to the canonical value order when the scorer returns an invalid + result (wrong length, non-finite values, or membership change). + """ + + def __init__( + self, + scorer: CandidateEnergyScorer, + *, + ranker_id: str | None = None, + fallback: str = "deterministic", + ) -> None: + self._scorer = scorer + self._ranker_id = ranker_id or f"energy-{getattr(scorer, 'scorer_id', id(scorer))}" + self._fallback = fallback + self._fallback_count = 0 + + @property + def ranker_id(self) -> str: + return self._ranker_id + + @property + def fallback_count(self) -> int: + return self._fallback_count + + def rank( + self, + state: FiniteDomainState, + hole_id: HoleId, + values: tuple[DomainValue, ...], + ) -> tuple[DomainValue, ...]: + import torch + + if not values: + return tuple(values) + try: + output = self._scorer(state, hole_id, values) + except Exception: # noqa: BLE001 + self._fallback_count += 1 + return tuple(values) + + if not isinstance(output, CandidateEnergyOutput): + self._fallback_count += 1 + return tuple(values) + + if len(output.candidate_ids) != len(values) or len(output.energies) != len( + values + ): + self._fallback_count += 1 + return tuple(values) + + if set(output.candidate_ids) != set(values): + self._fallback_count += 1 + return tuple(values) + + if not torch.isfinite(output.energies).all(): + self._fallback_count += 1 + return tuple(values) + + # Lower energy ranks earlier. + order = torch.argsort(output.energies, stable=True) + ranked = tuple(output.candidate_ids[int(i)] for i in order.tolist()) + if set(ranked) != set(values) or len(ranked) != len(values): + self._fallback_count += 1 + return tuple(values) + return ranked + + +def make_stub_energy_scorer( + scorer_id: str = "stub-v1", +) -> CandidateEnergyScorer: + """Return a deterministic scorer that assigns lower energy to shorter payloads. + + Useful for fixture tests: it never touches a model and always returns finite + energies for the exact supplied values. + """ + + def scorer( + state: FiniteDomainState, + hole_id: HoleId, + values: tuple[DomainValue, ...], + ) -> CandidateEnergyOutput: + import torch + + energies = torch.tensor( + [float(len(str(value.payload_json))) for value in values], + dtype=torch.float32, + ) + return CandidateEnergyOutput( + energies=energies, + candidate_ids=values, + scorer_id=scorer_id, + ) + + return scorer + + +__all__ = [ + "CandidateEnergyInput", + "CandidateEnergyOutput", + "CandidateEnergyScorer", + "EnergyCandidateRanker", + "make_stub_energy_scorer", +] diff --git a/src/slm_training/harnesses/model_build/config.py b/src/slm_training/harnesses/model_build/config.py index 63bb9a6e..576d5bb5 100644 --- a/src/slm_training/harnesses/model_build/config.py +++ b/src/slm_training/harnesses/model_build/config.py @@ -253,6 +253,9 @@ class ModelBuildConfig: speculative_successor: bool = False # E74 successor-state cache speculative_fanout: int = 2 speculative_overlap: bool = False + # VSS3-02 (SLM-70): cost-to-go energy head for solver candidate ranking. + cost_to_go_loss_weight: float = 0.0 + cost_to_go_hidden_dim: int = 0 # Hugging Face Bucket for durable checkpoints (full HF-context trains). # None → default hf://buckets/TKendrick/OpenUI when sync is enabled. # Empty string → disable auto bucket selection. diff --git a/src/slm_training/harnesses/model_build/cost_to_go_train.py b/src/slm_training/harnesses/model_build/cost_to_go_train.py new file mode 100644 index 00000000..ca090b31 --- /dev/null +++ b/src/slm_training/harnesses/model_build/cost_to_go_train.py @@ -0,0 +1,121 @@ +"""Head-only trainer for the VSS3-02 cost-to-go energy scorer (SLM-70). + +Freezes the TwoTower backbone and trains only ``cost_to_go_head`` from +replay-verified solver supervision rows. +""" + +from __future__ import annotations + +import json +from dataclasses import asdict +from pathlib import Path +from typing import Any + +import torch + +from slm_training.harnesses.distill.solver_supervision import CandidateCostRow +from slm_training.models.twotower import TwoTowerModel + + +def _load_rows(path: Path) -> list[CandidateCostRow]: + rows: list[CandidateCostRow] = [] + with path.open("r", encoding="utf-8") as fh: + for line in fh: + line = line.strip() + if not line: + continue + data = json.loads(line) + if data.get("row_kind") == "candidate_cost": + rows.append(CandidateCostRow.from_dict(data)) + return rows + + +def train_cost_to_go( + model: TwoTowerModel, + rows: list[CandidateCostRow], + *, + steps: int = 40, + batch_size: int = 8, + lr: float = 1e-3, + device: str | None = None, +) -> dict[str, Any]: + """Freeze context/denoiser/base; optimize only ``cost_to_go_head``.""" + if device: + model.to(device) + for p in model.parameters(): + p.requires_grad_(False) + if model.cost_to_go_head is None: + raise RuntimeError("model has no cost_to_go_head; set cost_to_go_hidden_dim > 0") + for p in model.cost_to_go_head.parameters(): + p.requires_grad_(True) + + opt = torch.optim.AdamW(model.cost_to_go_head.parameters(), lr=lr) + losses: list[float] = [] + n = max(1, len(rows)) + rng = torch.Generator() + rng.manual_seed(int(model.config.seed)) + + for step in range(max(1, steps)): + # Deterministic cyclic mini-batch. + start = (step * batch_size) % n + batch_rows = rows[start : start + batch_size] + if len(batch_rows) < batch_size: + batch_rows = batch_rows + rows[: batch_size - len(batch_rows)] + loss = model.cost_to_go_loss(batch_rows) + opt.zero_grad(set_to_none=True) + loss.backward() + opt.step() + losses.append(float(loss.detach().cpu())) + + # Restore trainability for non-head params for subsequent stages. + for p in model.parameters(): + p.requires_grad_(True) + model.config.cost_to_go_loss_weight = 1.0 + return { + "steps": steps, + "batch_size": batch_size, + "lr": lr, + "last_loss": losses[-1] if losses else None, + "mean_loss": sum(losses) / max(1, len(losses)), + "row_count": n, + } + + +def train_cost_to_go_from_paths( + checkpoint: Path | str, + rows_path: Path | str, + *, + out_dir: Path | str, + steps: int = 40, + batch_size: int = 8, + lr: float = 1e-3, + device: str = "cpu", +) -> dict[str, Any]: + """Load a checkpoint and a solver-supervision JSONL, train the head, save result.""" + ckpt = Path(checkpoint) + out = Path(out_dir) + out.mkdir(parents=True, exist_ok=True) + model = TwoTowerModel.from_checkpoint(ckpt, device=device) + rows = _load_rows(Path(rows_path)) + if not rows: + raise ValueError(f"no candidate_cost rows found in {rows_path}") + summary = train_cost_to_go( + model, rows, steps=steps, batch_size=batch_size, lr=lr, device=device + ) + dest = out / "checkpoints" / "last.pt" + dest.parent.mkdir(parents=True, exist_ok=True) + model.save(dest) + meta = { + "scorer_id": "twotower-cost-to-go-v1", + "source_checkpoint": str(ckpt), + "source_rows": str(rows_path), + "training": summary, + "config": {k: v for k, v in asdict(model.config).items() if k.startswith("cost_to_go")}, + } + (out / "cost_to_go_train_summary.json").write_text( + json.dumps(meta, indent=2) + "\n", encoding="utf-8" + ) + return {**summary, "checkpoint": str(dest), "summary_path": str(out / "cost_to_go_train_summary.json")} + + +__all__ = ["train_cost_to_go", "train_cost_to_go_from_paths"] diff --git a/src/slm_training/harnesses/model_build/factory.py b/src/slm_training/harnesses/model_build/factory.py index 8f4573d6..b120b2b1 100644 --- a/src/slm_training/harnesses/model_build/factory.py +++ b/src/slm_training/harnesses/model_build/factory.py @@ -150,6 +150,8 @@ def apply_runtime_overrides(model: Any, config: ModelBuildConfig) -> Any: "speculative_successor", "speculative_fanout", "speculative_overlap", + "cost_to_go_loss_weight", + "cost_to_go_hidden_dim", ): if allowed is not None and key not in allowed: continue @@ -437,6 +439,10 @@ def _twotower_config_from_build(config: ModelBuildConfig) -> "TwoTowerConfig": speculative_successor=bool(getattr(config, "speculative_successor", False)), speculative_fanout=int(getattr(config, "speculative_fanout", 2) or 2), speculative_overlap=bool(getattr(config, "speculative_overlap", False)), + cost_to_go_loss_weight=float( + getattr(config, "cost_to_go_loss_weight", 0.0) or 0.0 + ), + cost_to_go_hidden_dim=int(getattr(config, "cost_to_go_hidden_dim", 0) or 0), seed=config.seed, ) diff --git a/src/slm_training/models/twotower.py b/src/slm_training/models/twotower.py index 982b94dd..eb212d8a 100644 --- a/src/slm_training/models/twotower.py +++ b/src/slm_training/models/twotower.py @@ -386,6 +386,9 @@ class TwoTowerConfig: speculative_successor: bool = False speculative_fanout: int = 2 speculative_overlap: bool = False + # VSS3-02 (SLM-70): cost-to-go energy head for solver candidate ranking. + cost_to_go_loss_weight: float = 0.0 + cost_to_go_hidden_dim: int = 0 # 0 disables the head def _pad_batch( @@ -432,6 +435,8 @@ def _load_checkpoint_state( allowed_missing |= {key for key in missing if key.startswith("trust_gate.")} # V7 survival head is likewise a plug-in (trained via survival_train, E73). allowed_missing |= {key for key in missing if key.startswith("survival_head.")} + # VSS3-02 (SLM-70): cost-to-go head is a plug-in trained from solver supervision. + allowed_missing |= {key for key in missing if key.startswith("cost_to_go_head.")} bad_missing = sorted(set(missing) - allowed_missing) # V5 may have checkpointed a zero kind_lookup even when unused; the # non-factorized path now uses a non-persistent stub, so treat that legacy @@ -679,6 +684,23 @@ def isolated_aux_init(factory, offset: int): self.survival_head = isolated_aux_init( lambda: FastPathGate(self.config.d_model), 109 ) + # VSS3-02 (SLM-70): cost-to-go energy head for solver candidate ranking. + cost_to_go_dim = int(getattr(self.config, "cost_to_go_hidden_dim", 0) or 0) + cost_to_go_weight = float( + getattr(self.config, "cost_to_go_loss_weight", 0.0) or 0.0 + ) + self.cost_to_go_head = ( + isolated_aux_init( + lambda: nn.Sequential( + nn.Linear(self.config.d_model, max(1, cost_to_go_dim)), + nn.ReLU(), + nn.Linear(max(1, cost_to_go_dim), 1), + ), + 110, + ) + if cost_to_go_dim > 0 or cost_to_go_weight > 0.0 + else None + ) # V7 decode telemetry (MaskGIT path): forwards, successor hits/misses. self.speculative_stats = SpeculativeStats() self._rng = random.Random(self.config.seed) @@ -841,6 +863,7 @@ def optimizer_parameter_groups(self) -> list[dict[str, list[nn.Parameter]]]: "binder_arity_head.", "trust_gate.", "survival_head.", + "cost_to_go_head.", ) grouped: dict[str, list[nn.Parameter]] = {"base": []} for name, parameter in self.named_parameters(): @@ -2372,6 +2395,161 @@ def training_loss(self, batch: list[ExampleRecord]) -> torch.Tensor: return mask_loss + def score_candidates( + self, + state: Any, + hole_id: Any, + values: tuple[Any, ...], + context_prompt: str | None = None, + ) -> Any: + """Score live solver candidates with the cost-to-go head. + + Returns a ``CandidateEnergyOutput``. The state/hole_id are accepted as + opaque identities; the model only needs the candidate values and an + optional user prompt for context encoding. + """ + from slm_training.dsl.solver.energy_ranker import CandidateEnergyOutput + + if self.cost_to_go_head is None: + raise RuntimeError("cost_to_go_head is disabled in config") + if not values: + return CandidateEnergyOutput( + energies=torch.empty(0), + candidate_ids=values, + scorer_id="twotower-cost-to-go-v1", + ) + + device = self.device_name + prompts = [context_prompt or ""] * len(values) + ctx, ctx_pad = self._encode_context(prompts) + canvases: list[list[int]] = [] + for value in values: + payload = getattr(value, "payload", None) + token_ids = ( + list(payload.get("token_ids", [])) + if isinstance(payload, dict) + else [] + ) + canvases.append([self.tokenizer.bos_id] + token_ids) + ids = _pad_batch(canvases, self.tokenizer.pad_id, device=device) + hidden = self._denoiser_hidden(ids, ctx, ctx_pad) + lengths = (ids != self.tokenizer.pad_id).sum(dim=1) + last_idx = (lengths - 1).clamp_min(0) + pooled = hidden[torch.arange(hidden.size(0), device=hidden.device), last_idx] + energies = self.cost_to_go_head(pooled).squeeze(-1) + return CandidateEnergyOutput( + energies=energies, + candidate_ids=values, + scorer_id="twotower-cost-to-go-v1", + ) + + def cost_to_go_loss( + self, + rows: list[Any], + prompts: list[str] | None = None, + ) -> torch.Tensor: + """Regression + pairwise ranking loss over solver-supervision rows. + + Rows are expected to expose the same fields as + ``CandidateCostRow`` from the VSS3-01 supervision corpus. + """ + if self.cost_to_go_head is None: + return torch.tensor(0.0, device=self.device_name) + + from slm_training.harnesses.distill.solver_supervision import CandidateCostRow + + valid_rows: list[CandidateCostRow] = [] + for row in rows: + if isinstance(row, CandidateCostRow): + valid_rows.append(row) + elif isinstance(row, dict): + valid_rows.append(CandidateCostRow.from_dict(row)) + else: + raise TypeError(f"unexpected cost row type: {type(row)}") + + if not valid_rows: + return torch.tensor(0.0, device=self.device_name) + + device = self.device_name + if prompts is None: + prompts = [""] * len(valid_rows) + + # Build one candidate canvas per row using the row's chosen candidate. + canvases: list[list[int]] = [] + for row, prompt in zip(valid_rows, prompts, strict=True): + token_ids = list(row.candidate.payload.get("token_ids", [])) if isinstance( + row.candidate.payload, dict + ) else [] + canvases.append([self.tokenizer.bos_id] + token_ids) + ids = _pad_batch(canvases, self.tokenizer.pad_id, device=device) + ctx, ctx_pad = self._encode_context(prompts) + hidden = self._denoiser_hidden(ids, ctx, ctx_pad) + lengths = (ids != self.tokenizer.pad_id).sum(dim=1) + last_idx = (lengths - 1).clamp_min(0) + pooled = hidden[torch.arange(hidden.size(0), device=hidden.device), last_idx] + pred = self.cost_to_go_head(pooled).squeeze(-1) + + # Target: log1p(work). Work = nodes + verifier_calls + backtracks + depth. + work = torch.tensor( + [ + float(row.nodes) + + float(row.verifier_calls) + + float(row.backtracks) + + float(row.depth) + for row in valid_rows + ], + device=device, + dtype=pred.dtype, + ) + target = torch.log1p(work) + observed = torch.tensor( + [bool(row.cost_observed) for row in valid_rows], + device=device, + dtype=torch.bool, + ) + + regression_loss = torch.tensor(0.0, device=device, dtype=pred.dtype) + if observed.any(): + delta = pred[observed] - target[observed] + regression_loss = torch.where( + delta.abs() < 1.0, + 0.5 * delta * delta, + delta.abs() - 0.5, + ).mean() + + pairwise_loss = torch.tensor(0.0, device=device, dtype=pred.dtype) + groups: dict[tuple[str, str], list[int]] = {} + for idx, row in enumerate(valid_rows): + if not row.cost_observed: + continue + key = (row.state_fingerprint, str(row.hole_id)) + groups.setdefault(key, []).append(idx) + margins: list[torch.Tensor] = [] + for indices in groups.values(): + if len(indices) < 2: + continue + p = pred[indices] + t = target[indices] + # For each ordered pair where target_i < target_j, require energy_i < energy_j. + for i in range(len(indices)): + for j in range(len(indices)): + if i == j: + continue + if t[i] < t[j]: + margins.append(F.relu(1.0 + p[i] - p[j])) + if margins: + pairwise_loss = torch.stack(margins).mean() + + weight = float(self.config.cost_to_go_loss_weight or 0.0) + total = weight * (regression_loss + pairwise_loss) + self.last_training_metrics = { + "cost_to_go_regression_loss": float(regression_loss.detach().cpu()), + "cost_to_go_pairwise_loss": float(pairwise_loss.detach().cpu()), + "cost_to_go_total_loss": float(total.detach().cpu()), + "cost_to_go_observed_rows": int(observed.sum().cpu()), + } + return total + def _format_one_context( self, prompt: str, diff --git a/tests/test_dsl/test_energy_ranker.py b/tests/test_dsl/test_energy_ranker.py new file mode 100644 index 00000000..3d3ff77c --- /dev/null +++ b/tests/test_dsl/test_energy_ranker.py @@ -0,0 +1,110 @@ +"""Tests for the VSS3-02 energy ranker (SLM-70).""" + +from __future__ import annotations + +import torch + +from slm_training.dsl.solver.energy_ranker import ( + CandidateEnergyOutput, + EnergyCandidateRanker, + make_stub_energy_scorer, +) +from slm_training.dsl.solver.state import DomainValue, FiniteDomainState, HoleId, SolverBounds + + +def _hole_id(name: str) -> HoleId: + return HoleId(namespace=name, path=(), kind="test") + + +def _state(values: tuple[DomainValue, ...]) -> FiniteDomainState: + return FiniteDomainState( + problem_id="p1", + pack_id="openui", + constraint_version="v1", + bounds=SolverBounds( + max_tokens=128, + max_nodes=64, + max_depth=8, + max_backtracks=4, + max_verifier_calls=16, + ), + holes=(), + decision_level=0, + ) + + +def _value(tag: str, token_ids: list[int]) -> DomainValue: + return DomainValue.create(tag, {"token_ids": token_ids}) + + +def test_energy_ranker_returns_permutation() -> None: + values = (_value("a", [1, 2]), _value("b", [3]), _value("c", [4, 5, 6])) + ranker = EnergyCandidateRanker(make_stub_energy_scorer()) + ranked = ranker.rank(_state(values), _hole_id("h1"), values) + assert set(ranked) == set(values) + assert len(ranked) == len(values) + assert ranker.fallback_count == 0 + + +def test_energy_ranker_sorts_by_ascending_energy() -> None: + values = (_value("long", [1] * 10), _value("short", [1]), _value("mid", [1] * 5)) + ranker = EnergyCandidateRanker(make_stub_energy_scorer()) + ranked = ranker.rank(_state(values), _hole_id("h1"), values) + # Shorter payload -> lower energy. + assert ranked[0].tag == "short" + assert ranked[1].tag == "mid" + assert ranked[2].tag == "long" + + +def test_energy_ranker_falls_back_on_wrong_length() -> None: + values = (_value("a", [1]), _value("b", [2])) + + def bad_scorer(state, hole_id, values): + return CandidateEnergyOutput( + energies=torch.tensor([0.0]), + candidate_ids=values[:1], + scorer_id="bad", + ) + + ranker = EnergyCandidateRanker(bad_scorer) + ranked = ranker.rank(_state(values), _hole_id("h1"), values) + assert ranked == values + assert ranker.fallback_count == 1 + + +def test_energy_ranker_falls_back_on_nan() -> None: + values = (_value("a", [1]), _value("b", [2])) + + def nan_scorer(state, hole_id, values): + return CandidateEnergyOutput( + energies=torch.tensor([0.0, float("nan")]), + candidate_ids=values, + scorer_id="nan", + ) + + ranker = EnergyCandidateRanker(nan_scorer) + ranked = ranker.rank(_state(values), _hole_id("h1"), values) + assert ranked == values + assert ranker.fallback_count == 1 + + +def test_energy_ranker_falls_back_on_membership_change() -> None: + values = (_value("a", [1]), _value("b", [2])) + other = _value("c", [3]) + + def swap_scorer(state, hole_id, values): + return CandidateEnergyOutput( + energies=torch.tensor([0.0, 1.0]), + candidate_ids=(other, values[1]), + scorer_id="swap", + ) + + ranker = EnergyCandidateRanker(swap_scorer) + ranked = ranker.rank(_state(values), _hole_id("h1"), values) + assert ranked == values + assert ranker.fallback_count == 1 + + +def test_energy_ranker_empty_values() -> None: + ranker = EnergyCandidateRanker(make_stub_energy_scorer()) + assert ranker.rank(_state(()), _hole_id("h1"), ()) == () diff --git a/tests/test_models/test_cost_to_go.py b/tests/test_models/test_cost_to_go.py new file mode 100644 index 00000000..3592412a --- /dev/null +++ b/tests/test_models/test_cost_to_go.py @@ -0,0 +1,235 @@ +"""Tests for the VSS3-02 cost-to-go head on TwoTower (SLM-70).""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") + +from slm_training.dsl.schema import ExampleRecord +from slm_training.dsl.solver.state import DomainValue, HoleId +from slm_training.harnesses.distill.solver_supervision import CandidateCostRow, SearchCounters +from slm_training.models.twotower import ( + TwoTowerConfig, + TwoTowerModel, + _load_checkpoint_state, +) + +HERO = 'root = Stack([hero], "column")\nhero_title = TextContent(":hero.title")\nhero_body = TextContent(":hero.body")\nhero = Card([hero_title, hero_body])' + + +def _hole_id(name: str) -> HoleId: + return HoleId(namespace=name, path=(), kind="test") + + +def _make_rows() -> list[CandidateCostRow]: + """Return candidate_cost rows where lower token-count candidates are cheaper.""" + counters = SearchCounters( + nodes=1, tokens=1, depth=1, backtracks=0, verifier_calls=0 + ) + base = { + "schema_version": 1, + "row_kind": "candidate_cost", + "state_fingerprint": "fp1", + "parent_fingerprint": None, + "problem_id": "p1", + "pack_id": "openui", + "constraint_version": "v1", + "program_family_id": "fam1", + "lineage_id": "lin1", + "split_group_id": "sg1", + "split": "train", + "capsule_id": None, + "hole_id": _hole_id("h1"), + "ranker_id": "baseline", + "chosen": False, + "support_verdict": "unknown", + "terminal_success": False, + "cost_observed": True, + "censor_reason": None, + "conflict_reason": None, + "counters": counters, + "final_trajectory_status": "unknown", + "trace_id": None, + "trajectory_id": None, + } + + def row(candidate: DomainValue, *, nodes: int, tokens: int, observed: bool = True) -> CandidateCostRow: + return CandidateCostRow( + **{ + **base, + "candidate": candidate, + "nodes": nodes, + "tokens": tokens, + "depth": nodes, + "backtracks": 0, + "verifier_calls": 0, + "cost_observed": observed, + } + ) + + # Two candidates in the same (state, hole): short is cheaper. + a = DomainValue.create("short", {"token_ids": [10, 11]}) + b = DomainValue.create("long", {"token_ids": [10, 11, 12, 13]}) + return [ + row(a, nodes=2, tokens=2), + row(b, nodes=8, tokens=8), + # Another state with a censored row. + CandidateCostRow( + **{ + **base, + "state_fingerprint": "fp2", + "candidate": DomainValue.create("c", {"token_ids": [20]}), + "nodes": 0, + "tokens": 0, + "depth": 0, + "backtracks": 0, + "verifier_calls": 0, + "cost_observed": False, + } + ), + ] + + +def test_head_created_when_enabled() -> None: + records = [ExampleRecord(id="a", prompt="Hero", openui=HERO, split="train")] + model = TwoTowerModel.from_records( + records, + config=TwoTowerConfig( + d_model=32, + n_heads=4, + context_layers=1, + denoiser_layers=1, + cost_to_go_hidden_dim=16, + ), + device="cpu", + ) + assert model.cost_to_go_head is not None + + +def test_head_absent_when_disabled() -> None: + records = [ExampleRecord(id="a", prompt="Hero", openui=HERO, split="train")] + model = TwoTowerModel.from_records( + records, + config=TwoTowerConfig( + d_model=32, + n_heads=4, + context_layers=1, + denoiser_layers=1, + ), + device="cpu", + ) + assert model.cost_to_go_head is None + + +def test_loss_decreases_on_synthetic_rows() -> None: + records = [ExampleRecord(id="a", prompt="Hero", openui=HERO, split="train")] + model = TwoTowerModel.from_records( + records, + config=TwoTowerConfig( + d_model=32, + n_heads=4, + context_layers=1, + denoiser_layers=1, + cost_to_go_hidden_dim=16, + cost_to_go_loss_weight=1.0, + ), + device="cpu", + ) + rows = _make_rows() + opt = torch.optim.AdamW(model.cost_to_go_head.parameters(), lr=1e-2) + losses: list[float] = [] + for _ in range(20): + loss = model.cost_to_go_loss(rows) + opt.zero_grad(set_to_none=True) + loss.backward() + opt.step() + losses.append(float(loss.detach().cpu())) + assert losses[-1] < losses[0] + assert model.last_training_metrics["cost_to_go_observed_rows"] == 2 + + +def test_censored_rows_are_masked_from_regression() -> None: + records = [ExampleRecord(id="a", prompt="Hero", openui=HERO, split="train")] + model = TwoTowerModel.from_records( + records, + config=TwoTowerConfig( + d_model=32, + n_heads=4, + context_layers=1, + denoiser_layers=1, + cost_to_go_hidden_dim=16, + cost_to_go_loss_weight=1.0, + ), + device="cpu", + ) + rows = _make_rows() + loss = model.cost_to_go_loss(rows) + loss.backward() + # Censored row is at index 2; observed mask should only include 0,1. + assert model.last_training_metrics["cost_to_go_observed_rows"] == 2 + + +def test_old_checkpoint_loads_with_missing_head(tmp_path: Path) -> None: + records = [ExampleRecord(id="a", prompt="Hero", openui=HERO, split="train")] + old_model = TwoTowerModel.from_records( + records, + config=TwoTowerConfig( + d_model=32, + n_heads=4, + context_layers=1, + denoiser_layers=1, + ), + device="cpu", + ) + ckpt = tmp_path / "old.pt" + old_model.save(ckpt) + + new_model = TwoTowerModel.from_records( + records, + config=TwoTowerConfig( + d_model=32, + n_heads=4, + context_layers=1, + denoiser_layers=1, + cost_to_go_hidden_dim=16, + ), + device="cpu", + ) + payload = torch.load(ckpt, map_location="cpu", weights_only=True) + _load_checkpoint_state(new_model, payload["state_dict"]) + assert new_model.cost_to_go_head is not None + + +def test_pairwise_ordering_follows_cost() -> None: + records = [ExampleRecord(id="a", prompt="Hero", openui=HERO, split="train")] + model = TwoTowerModel.from_records( + records, + config=TwoTowerConfig( + d_model=32, + n_heads=4, + context_layers=1, + denoiser_layers=1, + cost_to_go_hidden_dim=16, + cost_to_go_loss_weight=1.0, + ), + device="cpu", + ) + rows = _make_rows() + opt = torch.optim.AdamW(model.cost_to_go_head.parameters(), lr=1e-2) + for _ in range(30): + loss = model.cost_to_go_loss(rows) + opt.zero_grad(set_to_none=True) + loss.backward() + opt.step() + + short = DomainValue.create("short", {"token_ids": [10, 11]}) + long = DomainValue.create("long", {"token_ids": [10, 11, 12, 13]}) + state = None + output = model.score_candidates( + state, _hole_id("h1"), (short, long), context_prompt="Hero" + ) + energies = output.energies.detach() + assert energies[0] < energies[1]