diff --git a/README.md b/README.md index 7d53e6f..eb630a6 100644 --- a/README.md +++ b/README.md @@ -115,6 +115,28 @@ The optional `ShiftLinear` is a zero-padded, one-token causal reference. Its streaming-cache interface is intentionally unsupported until cache correctness is tested. +## Matched-backbone control lab + +Run a seeded geometric-versus-standard-additive smoke experiment from identical +initial weights and an identical pre-generated token stream: + +```bash +python experiments/paired_control.py --steps 20 --output paired-control.json +``` + +The machine-readable receipt records the complete experiment specification, +initialization and dataset SHA-256 digests, environment, per-step losses and +gradient norms, active-gradient parameter counts, throughput, and its own +integrity digest. The additive arm keeps the geometric controller parameters so +the serialized parameter count is identical, but explicitly reports that those +parameters are dormant. + +This follows the reproducibility patterns used by small training baselines and +evaluation harnesses while keeping the implementation original. Synthetic-token +loss is only a plumbing check. Publishable claims still require real data, +multiple seeds, controlled compute, held-out evaluation, and target-hardware +measurements. + ## License Apache License 2.0. See [LICENSE](LICENSE). diff --git a/experiments/paired_control.py b/experiments/paired_control.py new file mode 100644 index 0000000..d3cd4b4 --- /dev/null +++ b/experiments/paired_control.py @@ -0,0 +1,37 @@ +"""Run a matched-backbone geometric/additive control and emit a JSON receipt.""" + +import argparse +import json +from pathlib import Path + +from model.control_lab import PairedControlSpec, run_paired_control + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output", type=Path, default=Path("paired-control.json")) + parser.add_argument("--seed", type=int, default=1337) + parser.add_argument("--steps", type=int, default=20) + parser.add_argument("--batch-size", type=int, default=4) + parser.add_argument("--sequence-length", type=int, default=32) + parser.add_argument("--learning-rate", type=float, default=3e-4) + parser.add_argument("--device", default="cpu") + args = parser.parse_args() + + receipt = run_paired_control( + PairedControlSpec( + seed=args.seed, + steps=args.steps, + batch_size=args.batch_size, + sequence_length=args.sequence_length, + learning_rate=args.learning_rate, + device=args.device, + ) + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(receipt, indent=2, sort_keys=True) + "\n") + print(args.output) + + +if __name__ == "__main__": + main() diff --git a/model/__init__.py b/model/__init__.py index 4feefcc..1615365 100644 --- a/model/__init__.py +++ b/model/__init__.py @@ -1,5 +1,13 @@ """Public OUROBOROS model API.""" +from .control_lab import PairedControlSpec, run_paired_control, verify_control_receipt from .ouroboros import OuroborosConfig, OuroborosModel, OuroborosResidual -__all__ = ["OuroborosConfig", "OuroborosModel", "OuroborosResidual"] +__all__ = [ + "OuroborosConfig", + "OuroborosModel", + "OuroborosResidual", + "PairedControlSpec", + "run_paired_control", + "verify_control_receipt", +] diff --git a/model/control_lab.py b/model/control_lab.py new file mode 100644 index 0000000..515ed6c --- /dev/null +++ b/model/control_lab.py @@ -0,0 +1,182 @@ +"""Small, auditable paired-control experiments for geometric residual research.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass +import hashlib +import json +import platform +import time +from typing import Any + +import torch + +from .ouroboros import OuroborosConfig, OuroborosModel + + +SCHEMA_VERSION = 1 + + +def _canonical_json(value: Any) -> bytes: + return json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8") + + +def _state_digest(model: OuroborosModel) -> str: + digest = hashlib.sha256() + for name, tensor in sorted(model.state_dict().items()): + digest.update(name.encode("utf-8")) + digest.update(tensor.detach().cpu().contiguous().numpy().tobytes()) + return digest.hexdigest() + + +@dataclass(frozen=True) +class PairedControlSpec: + seed: int = 1337 + steps: int = 20 + batch_size: int = 4 + sequence_length: int = 32 + learning_rate: float = 3e-4 + device: str = "cpu" + deterministic_algorithms: bool = True + model_config: dict[str, Any] | None = None + + def validate(self) -> None: + if self.steps <= 0 or self.batch_size <= 0 or self.sequence_length <= 1: + raise ValueError("steps/batch_size must be positive and sequence_length > 1") + if self.learning_rate <= 0: + raise ValueError("learning_rate must be positive") + + +def _default_model_config() -> dict[str, Any]: + return { + "vocab_size": 256, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "hidden_size": 128, + "head_dim": 32, + "block_size": 128, + "using_groupnorm": False, + } + + +def _make_batches(spec: PairedControlSpec, vocab_size: int) -> list[torch.Tensor]: + generator = torch.Generator(device="cpu").manual_seed(spec.seed + 1) + return [ + torch.randint( + 0, + vocab_size, + (spec.batch_size, spec.sequence_length + 1), + generator=generator, + ) + for _ in range(spec.steps) + ] + + +def _run_arm(model, batches, learning_rate, device): + model.to(device).train() + optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate) + losses = [] + gradient_norms = [] + active_gradient_parameters = set() + started = time.perf_counter() + for batch in batches: + batch = batch.to(device) + optimizer.zero_grad(set_to_none=True) + _, loss = model(batch[:, :-1], targets=batch[:, 1:]) + loss.backward() + squared_norm = torch.zeros((), device=device) + for name, parameter in model.named_parameters(): + if parameter.grad is not None: + active_gradient_parameters.add(name) + squared_norm += parameter.grad.float().pow(2).sum() + optimizer.step() + losses.append(float(loss.detach().cpu())) + gradient_norms.append(float(squared_norm.sqrt().detach().cpu())) + duration = time.perf_counter() - started + token_count = spec_token_count = sum(batch[:, :-1].numel() for batch in batches) + return { + "losses": losses, + "gradient_norms": gradient_norms, + "active_gradient_parameter_count": len(active_gradient_parameters), + "duration_seconds": duration, + "tokens_per_second": token_count / duration, + "trained_tokens": spec_token_count, + } + + +def run_paired_control(spec: PairedControlSpec) -> dict[str, Any]: + """Run geometric and additive arms from the same initialization and batches.""" + + spec.validate() + config_values = _default_model_config() + config_values.update(spec.model_config or {}) + if spec.sequence_length > int(config_values["block_size"]): + raise ValueError("sequence_length exceeds model block_size") + + previous_determinism = torch.are_deterministic_algorithms_enabled() + torch.use_deterministic_algorithms(spec.deterministic_algorithms) + try: + torch.manual_seed(spec.seed) + geometric_config = OuroborosConfig( + **config_values, ouroboros_residual_mode="geometric" + ) + geometric = OuroborosModel(geometric_config) + initial_state = { + name: tensor.detach().clone() for name, tensor in geometric.state_dict().items() + } + additive_config = OuroborosConfig( + **config_values, ouroboros_residual_mode="additive" + ) + additive = OuroborosModel(additive_config) + additive.load_state_dict(initial_state) + initialization_sha256 = _state_digest(geometric) + if _state_digest(additive) != initialization_sha256: + raise RuntimeError("control arms did not start from identical parameters") + + batches = _make_batches(spec, geometric_config.vocab_size) + dataset_sha256 = hashlib.sha256( + b"".join(batch.numpy().tobytes() for batch in batches) + ).hexdigest() + arms = { + "geometric": _run_arm( + geometric, batches, spec.learning_rate, torch.device(spec.device) + ), + "additive": _run_arm( + additive, batches, spec.learning_rate, torch.device(spec.device) + ), + } + receipt = { + "schema_version": SCHEMA_VERSION, + "spec": asdict(spec), + "model_config": config_values, + "initialization_sha256": initialization_sha256, + "dataset_sha256": dataset_sha256, + "parameter_count": sum(parameter.numel() for parameter in geometric.parameters()), + "environment": { + "python": platform.python_version(), + "torch": torch.__version__, + "device": str(spec.device), + "deterministic_algorithms": spec.deterministic_algorithms, + }, + "arms": arms, + "limitations": [ + "The additive arm keeps geometric-controller parameters for an identical serialized parameter count, but those controllers are dormant.", + "Wall-clock throughput is observational and is not deterministic across hardware.", + "Synthetic-token loss is a smoke metric, not evidence of language-model quality.", + ], + } + receipt["receipt_sha256"] = hashlib.sha256(_canonical_json(receipt)).hexdigest() + return receipt + finally: + torch.use_deterministic_algorithms(previous_determinism) + + +def verify_control_receipt(receipt: dict[str, Any]) -> bool: + if receipt.get("schema_version") != SCHEMA_VERSION: + return False + claimed = receipt.get("receipt_sha256") + unsigned = dict(receipt) + unsigned.pop("receipt_sha256", None) + return isinstance(claimed, str) and hashlib.sha256( + _canonical_json(unsigned) + ).hexdigest() == claimed diff --git a/model/ouroboros.py b/model/ouroboros.py index d451233..0f5ebad 100644 --- a/model/ouroboros.py +++ b/model/ouroboros.py @@ -115,6 +115,9 @@ class OuroborosResidual(nn.Module): def __init__(self, config): super().__init__() hidden_size = int(config.hidden_size) + self.residual_mode = str( + getattr(config, "ouroboros_residual_mode", "geometric") + ) self.k_eps = float(getattr(config, "ouroboros_k_eps", 1e-6)) self.v_sigmoid = bool(getattr(config, "ouroboros_v_sigmoid", True)) @@ -191,7 +194,10 @@ def forward( return_diagnostics: bool = False, ): k, beta, target, projection, delta = self._components(x, k_in, context) - output = geometric_update(x, k, delta) + if self.residual_mode == "additive": + output = x + k_in + else: + output = geometric_update(x, k, delta) if not return_diagnostics: return output return output, self._diagnostics(x, output, k, beta, target, projection) @@ -265,6 +271,8 @@ class OuroborosConfig(PretrainedConfig): ouroboros_v_constant: bool = False ouroboros_v_constant_value: float = 2.0 ouroboros_collect_diagnostics: bool = False + # ``additive`` is a matched-backbone experimental control, not the default. + ouroboros_residual_mode: str = "geometric" # Initialize beta; clamped to [0, 2]. Use 1.0 by default for baseline comparability. ouroboros_beta_init: float = 1.0 @@ -317,7 +325,11 @@ def forward(self, idx, targets=None, return_logits=True, output_all_seq=False): logits = self.lm_head(x) logits = logits.float() # use tf32/fp32 for logits logits = logits * logits_scale - loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1) + loss = F.cross_entropy( + logits.reshape(-1, logits.size(-1)), + targets.reshape(-1), + ignore_index=-1, + ) elif output_all_seq: logits = self.lm_head(x[:, :, :]) logits = logits * logits_scale diff --git a/model/pydantic_config.py b/model/pydantic_config.py index 028f316..aa4ef71 100644 --- a/model/pydantic_config.py +++ b/model/pydantic_config.py @@ -31,4 +31,8 @@ def validate_pretrained_config_kwargs(config_cls, kwargs): raise ValueError("hidden_size must equal num_attention_heads * head_dim") if float(values["ouroboros_k_eps"]) <= 0: raise ValueError("ouroboros_k_eps must be positive") + if values["ouroboros_residual_mode"] not in {"geometric", "additive"}: + raise ValueError( + "ouroboros_residual_mode must be either 'geometric' or 'additive'" + ) return values diff --git a/tests/test_control_lab.py b/tests/test_control_lab.py new file mode 100644 index 0000000..2cb60ea --- /dev/null +++ b/tests/test_control_lab.py @@ -0,0 +1,66 @@ +import copy + +import pytest +import torch + +from model import OuroborosConfig, OuroborosResidual +from model.control_lab import ( + PairedControlSpec, + run_paired_control, + verify_control_receipt, +) + + +def tiny_model_config(): + return { + "vocab_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "hidden_size": 16, + "head_dim": 8, + "block_size": 8, + "using_groupnorm": False, + } + + +def test_additive_control_is_standard_branch_addition(): + residual = OuroborosResidual( + OuroborosConfig(**tiny_model_config(), ouroboros_residual_mode="additive") + ) + x = torch.randn(2, 4, 16) + branch = torch.randn_like(x) + torch.testing.assert_close( + residual(x, k_in=branch, context=x), + x + branch, + ) + + +def test_paired_control_emits_verifiable_matched_receipt(): + receipt = run_paired_control( + PairedControlSpec( + seed=7, + steps=2, + batch_size=2, + sequence_length=4, + model_config=tiny_model_config(), + ) + ) + + assert verify_control_receipt(receipt) + assert receipt["parameter_count"] > 0 + assert receipt["arms"]["geometric"]["trained_tokens"] == 16 + assert receipt["arms"]["additive"]["trained_tokens"] == 16 + assert len(receipt["arms"]["geometric"]["losses"]) == 2 + + tampered = copy.deepcopy(receipt) + tampered["arms"]["geometric"]["losses"][0] += 1 + assert not verify_control_receipt(tampered) + + +def test_control_spec_and_residual_mode_fail_closed(): + with pytest.raises(ValueError, match="sequence_length"): + run_paired_control( + PairedControlSpec(sequence_length=9, model_config=tiny_model_config()) + ) + with pytest.raises(ValueError, match="ouroboros_residual_mode"): + OuroborosConfig(**tiny_model_config(), ouroboros_residual_mode="unknown")