feat: add reproducible matched-backbone control receipts

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).

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()

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",
+]

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

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

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

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")