experiment: implement shared-feedback feasibility gate

3 files changed, 501 insertions, 0 deletions

diff --git a/experiments/shared_feedback_s0.py b/experiments/shared_feedback_s0.py
new file mode 100644
index 0000000..719a333
--- /dev/null
+++ b/experiments/shared_feedback_s0.py
@@ -0,0 +1,162 @@
+#!/usr/bin/env python3
+"""Frozen single-run S0 screen from SHARED_FEEDBACK.md."""
+
+import argparse
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+import time
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from sdil.shared_feedback import (
+    CONDITIONS, SharedFeedbackConfig, SharedFeedbackNet,
+    conditional_selector_data, evaluate_shared_feedback, shared_feedback_step,
+)
+
+
+ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_OUT = ROOT / "results" / "shared_feedback" / "s0.json"
+
+
+def git_output(*args):
+    return subprocess.run(
+        ["git", *args], cwd=ROOT, check=True, capture_output=True,
+        text=True).stdout.strip()
+
+
+def train_condition(condition, base, train, validation, neutral, device):
+    net = base.clone()
+    x_train, z_train, y_train = train
+    x_val, z_val, y_val = validation
+    x_neutral, z_neutral = neutral
+    shuffle = torch.Generator(device="cpu").manual_seed(3101)
+    epoch_losses = []
+    predictor_reports = []
+    first_nonfinite_epoch = None
+    started = time.time()
+    for epoch in range(40):
+        if condition in ("innovation", "matched_raw"):
+            predictor_reports = net.fit_neutral_predictor(x_neutral, z_neutral)
+        permutation = torch.randperm(x_train.shape[0], generator=shuffle)
+        losses = []
+        for start in range(0, x_train.shape[0], 128):
+            indices = permutation[start:start + 128].to(device)
+            loss, _ = shared_feedback_step(
+                net, x_train[indices], z_train[indices], y_train[indices], condition)
+            losses.append(loss)
+        mean_loss = sum(losses) / len(losses)
+        epoch_losses.append(mean_loss)
+        if not torch.isfinite(torch.tensor(mean_loss)):
+            first_nonfinite_epoch = epoch
+            break
+    if condition not in ("innovation", "matched_raw"):
+        predictor_reports = net.fit_neutral_predictor(x_neutral, z_neutral)
+    endpoint = evaluate_shared_feedback(net, x_val, z_val, y_val)
+    lesion = evaluate_shared_feedback(
+        net, x_val, z_val, y_val, context_enabled=False)
+    return {
+        "condition": condition,
+        "epochs_completed": len(epoch_losses),
+        "epoch_train_loss": epoch_losses,
+        "first_nonfinite_epoch": first_nonfinite_epoch,
+        "finite": first_nonfinite_epoch is None,
+        "validation": endpoint,
+        "context_lesion_validation": lesion,
+        "predictor": predictor_reports,
+        "wall_seconds": time.time() - started,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--out", type=Path, default=DEFAULT_OUT)
+    args = parser.parse_args()
+    if git_output("status", "--porcelain", "--untracked-files=no"):
+        raise RuntimeError("S0 requires clean tracked source")
+    device = torch.device(args.device)
+    torch.manual_seed(3101)
+    if device.type == "cpu":
+        torch.set_num_threads(1)
+
+    config = SharedFeedbackConfig()
+    base = SharedFeedbackNet(config, seed=3101, device=device)
+    train = conditional_selector_data(8192, 3101, device)
+    validation = conditional_selector_data(2048, 3102, device)
+    neutral = (train[0][:512], train[1][:512])
+    records = [train_condition(
+        condition, base, train, validation, neutral, device)
+               for condition in CONDITIONS]
+    by_name = {row["condition"]: row for row in records}
+    oracle = 100.0 * by_name["oracle"]["validation"]["accuracy"]
+    raw = 100.0 * by_name["raw_shared"]["validation"]["accuracy"]
+    innovation = 100.0 * by_name["innovation"]["validation"]["accuracy"]
+    matched = 100.0 * by_name["matched_raw"]["validation"]["accuracy"]
+    lesion_drop = 100.0 * (
+        by_name["oracle"]["validation"]["accuracy"]
+        - by_name["oracle"]["context_lesion_validation"]["accuracy"])
+    predictor = by_name["innovation"]["predictor"]
+    checks = {
+        "oracle_at_least_90": oracle >= 90.0,
+        "oracle_context_lesion_drop_at_least_10": lesion_drop >= 10.0,
+        "nonzero_context_every_layer": all(
+            row["context_rms"] > 0 for row in predictor),
+        "raw_below_oracle_by_5_or_nonfinite": (
+            oracle - raw >= 5.0 or not by_name["raw_shared"]["finite"]),
+        "innovation_above_raw_by_5": innovation - raw >= 5.0,
+        "innovation_within_3_of_oracle": oracle - innovation <= 3.0,
+        "innovation_within_2_of_exact_subtraction": abs(
+            innovation - oracle) <= 2.0,
+        "matched_raw_below_innovation_by_3": innovation - matched >= 3.0,
+        "predictor_mean_r2_at_least_0p8": (
+            sum(row["mean_per_cell_r2"] for row in predictor)
+            / len(predictor) >= 0.8),
+        "predictor_residual_ratio_at_most_0p25": max(
+            row["residual_context_rms_ratio"] for row in predictor) <= 0.25,
+        "zero_instruction_observations": max(
+            row["instruction_observations"] for record in records
+            for row in record["predictor"]) == 0,
+    }
+    report = {
+        "stage": "shared_feedback_s0",
+        "gate": "pass" if all(checks.values()) else "fail",
+        "checks": checks,
+        "config": config.__dict__,
+        "data": {
+            "train_examples": 8192, "validation_examples": 2048,
+            "neutral_examples_per_epoch": 512, "batch_size": 128,
+            "epochs": 40, "data_seed": 3101,
+            "validation_seed": 3102, "test_generated": False,
+        },
+        "records": records,
+        "summary": {
+            "oracle_validation_accuracy_percent": oracle,
+            "raw_validation_accuracy_percent": raw,
+            "innovation_validation_accuracy_percent": innovation,
+            "matched_raw_validation_accuracy_percent": matched,
+            "oracle_context_lesion_drop_points": lesion_drop,
+        },
+        "provenance": {
+            "git_commit": git_output("rev-parse", "HEAD"),
+            "git_dirty_tracked": False,
+            "device": str(device), "torch_version": torch.__version__,
+            "cuda_device_name": (
+                torch.cuda.get_device_name(device) if device.type == "cuda" else None),
+            "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
+        },
+    }
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    with open(args.out, "w", encoding="utf-8") as handle:
+        json.dump(report, handle, indent=2, sort_keys=True)
+        handle.write("\n")
+    print(json.dumps({"gate": report["gate"], **report["summary"],
+                      "checks": checks}, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()
+
diff --git a/experiments/shared_feedback_smoke.py b/experiments/shared_feedback_smoke.py
new file mode 100644
index 0000000..672d123
--- /dev/null
+++ b/experiments/shared_feedback_smoke.py
@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+"""Deterministic mechanics checks for SHARED_FEEDBACK.md S0."""
+
+import os
+import sys
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from sdil.shared_feedback import (
+    CONDITIONS, SharedFeedbackConfig, SharedFeedbackNet,
+    conditional_selector_data, select_shared_signal, shared_feedback_step,
+)
+
+
+def maximum_difference(left, right):
+    return max(float((a - b).abs().max()) for a, b in zip(left, right))
+
+
+def main():
+    torch.set_num_threads(1)
+    config = SharedFeedbackConfig(width=8, hidden_layers=2)
+    base = SharedFeedbackNet(config, seed=3101, dtype=torch.float64)
+    x, z, y = conditional_selector_data(32, 3101)
+    x = x.to(torch.float64)
+
+    clones = {condition: base.clone() for condition in CONDITIONS}
+    forward_error = max(float((
+        clones["oracle"].forward(x, z)["logits"]
+        - clones[condition].forward(x, z)["logits"]).abs().max())
+                        for condition in CONDITIONS)
+    assert forward_error == 0.0
+
+    enabled = base.forward(x, z)
+    lesioned = base.forward(x, z, context_enabled=False)
+    lesion_context_error = max(float(value.abs().max())
+                               for value in lesioned["context"])
+    assert lesion_context_error == 0.0
+    assert any(not torch.equal(a, b) for a, b in zip(
+        enabled["h"][1:-1], lesioned["h"][1:-1]))
+
+    reports = base.fit_neutral_predictor(x, z)
+    assert all(row["instruction_observations"] == 0 for row in reports)
+    layer = 0
+    soma = base.forward(x, z)["h"][1]
+    context = base.context_fields(z)[layer]
+    instruction = torch.randn_like(context)
+    exact, raw, innovation = select_shared_signal(
+        base, layer, instruction, context, soma, "exact_subtraction")
+    assert torch.equal(raw, instruction + context)
+    assert torch.allclose(exact, instruction, atol=1e-15, rtol=1e-15)
+    matched, _, _ = select_shared_signal(
+        base, layer, instruction, context, soma, "matched_raw")
+    assert torch.allclose(matched.norm(dim=1), innovation.norm(dim=1),
+                          atol=1e-12, rtol=1e-12)
+
+    # With context and predictor exactly zero, all four signal rules coincide.
+    zero = base.clone()
+    for value in zero.C + zero.P + zero.P_bias:
+        value.zero_()
+    zero_clones = {condition: zero.clone() for condition in CONDITIONS}
+    for condition, net in zero_clones.items():
+        shared_feedback_step(net, x, z, y, condition)
+    reference = zero_clones["oracle"]
+    zero_context_update_error = max(
+        maximum_difference(reference.W, net.W)
+        + maximum_difference(reference.Q[1:], net.Q[1:])
+        for condition, net in zero_clones.items() if condition != "oracle")
+    assert zero_context_update_error < 1e-14
+
+    # The reciprocal change is independently applied from the same local
+    # direction; it is not copied from the updated forward tensor.
+    kp = base.clone()
+    before_w = [value.clone() for value in kp.W]
+    before_q = [None] + [value.clone() for value in kp.Q[1:]]
+    _, auxiliary = shared_feedback_step(kp, x, z, y, "oracle")
+    reciprocal_direction_error = 0.0
+    for layer in range(1, len(kp.W)):
+        expected_w = ((kp.W[layer] - before_w[layer])
+                      / kp.config.learning_rate)
+        expected_q = ((kp.Q[layer] - before_q[layer])
+                      / kp.config.reciprocal_learning_rate)
+        reciprocal_direction_error = max(
+            reciprocal_direction_error,
+            float((expected_w - expected_q
+                   - kp.config.weight_decay
+                   * (before_q[layer] - before_w[layer])).abs().max()))
+    assert reciprocal_direction_error < 1e-12
+    assert all(not value.requires_grad for value in kp.W + kp.Q[1:])
+    assert all(direction.grad_fn is None for direction in auxiliary["directions"])
+
+    print({
+        "forward_identity_error": forward_error,
+        "lesion_context_error": lesion_context_error,
+        "zero_context_update_error": zero_context_update_error,
+        "reciprocal_direction_error": reciprocal_direction_error,
+        "neutral_instruction_observations": max(
+            row["instruction_observations"] for row in reports),
+    })
+
+
+if __name__ == "__main__":
+    main()
diff --git a/sdil/shared_feedback.py b/sdil/shared_feedback.py
new file mode 100644
index 0000000..9a38e0f
--- /dev/null
+++ b/sdil/shared_feedback.py
@@ -0,0 +1,236 @@
+"""Minimal endogenous shared-apical-path feasibility model.
+
+The context field in this module is part of the forward computation and is
+required to solve the conditional task.  During learning, the same apical
+measurement contains that ordinary field plus reciprocal KP instruction.
+There is no generated nuisance or bias term.
+"""
+
+from dataclasses import dataclass
+import math
+
+import torch
+import torch.nn.functional as F
+
+
+CONDITIONS = ("oracle", "raw_shared", "innovation", "matched_raw")
+
+
+def conditional_selector_data(n, seed, device="cpu"):
+    """Balanced contextual selector task with no context on the basal input."""
+    if n % 2:
+        raise ValueError("conditional selector data size must be even")
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+    x = torch.randn(n, 2, generator=generator)
+    z = torch.arange(n, dtype=torch.long).remainder(2)
+    permutation = torch.randperm(n, generator=generator)
+    x, z = x[permutation], z[permutation]
+    selected = x.gather(1, z[:, None]).squeeze(1)
+    y = (selected > 0).long()
+    return x.to(device), z.to(device), y.to(device)
+
+
+@dataclass(frozen=True)
+class SharedFeedbackConfig:
+    width: int = 64
+    hidden_layers: int = 2
+    learning_rate: float = 0.03
+    reciprocal_learning_rate: float = 0.03
+    momentum: float = 0.9
+    weight_decay: float = 1e-4
+    context_scale: float = 1.0
+
+
+class SharedFeedbackNet:
+    """Context-conditioned MLP with independently stored reciprocal weights."""
+
+    def __init__(self, config=SharedFeedbackConfig(), seed=3101,
+                 device="cpu", dtype=torch.float32):
+        if config.hidden_layers < 1:
+            raise ValueError("shared-feedback model needs a hidden population")
+        self.config = config
+        self.device = torch.device(device)
+        self.dtype = dtype
+        sizes = [2] + [config.width] * config.hidden_layers + [2]
+        generator = torch.Generator(device="cpu").manual_seed(seed)
+        reciprocal_generator = torch.Generator(device="cpu").manual_seed(seed + 1)
+        context_generator = torch.Generator(device="cpu").manual_seed(seed + 2)
+
+        def normal(shape, scale, source):
+            return (torch.randn(*shape, generator=source) * scale).to(
+                device=self.device, dtype=dtype)
+
+        self.W = [normal((sizes[i + 1], sizes[i]),
+                         1.0 / math.sqrt(sizes[i]), generator)
+                  for i in range(len(sizes) - 1)]
+        # Q[i] corresponds to W[i] and has the same storage orientation.  Q[0]
+        # is absent because the basal input does not need a transported field.
+        self.Q = [None] + [normal(tuple(self.W[i].shape),
+                                  1.0 / math.sqrt(sizes[i]),
+                                  reciprocal_generator)
+                           for i in range(1, len(self.W))]
+        self.C = [normal((config.width, 2),
+                         config.context_scale / math.sqrt(2.0),
+                         context_generator)
+                  for _ in range(config.hidden_layers)]
+        self.P = [torch.zeros(config.width, device=self.device, dtype=dtype)
+                  for _ in range(config.hidden_layers)]
+        self.P_bias = [torch.zeros_like(value) for value in self.P]
+
+        self.mW = [torch.zeros_like(value) for value in self.W]
+        self.mQ = [None] + [torch.zeros_like(value) for value in self.Q[1:]]
+
+    def clone(self):
+        copied = SharedFeedbackNet(
+            self.config, seed=0, device=self.device, dtype=self.dtype)
+        for name in ("W", "Q", "C", "P", "P_bias", "mW", "mQ"):
+            source = getattr(self, name)
+            target = []
+            for value in source:
+                target.append(None if value is None else value.clone())
+            setattr(copied, name, target)
+        return copied
+
+    def context_fields(self, z, enabled=True):
+        onehot = F.one_hot(z, num_classes=2).to(self.dtype)
+        if not enabled:
+            return [torch.zeros((z.shape[0], self.config.width),
+                                device=self.device, dtype=self.dtype)
+                    for _ in self.C]
+        return [onehot @ projection.t() for projection in self.C]
+
+    def forward(self, x, z, context_enabled=True):
+        fields = self.context_fields(z, enabled=context_enabled)
+        h = [x]
+        u = []
+        for layer in range(self.config.hidden_layers):
+            value = h[-1] @ self.W[layer].t() + fields[layer]
+            u.append(value)
+            h.append(torch.tanh(value))
+        logits = h[-1] @ self.W[-1].t()
+        h.append(logits)
+        return {"h": h, "u": u, "context": fields, "logits": logits}
+
+    def predictor(self, layer, soma):
+        return self.P[layer] * soma + self.P_bias[layer]
+
+    @torch.no_grad()
+    def fit_neutral_predictor(self, x, z):
+        """Per-cell affine neutral fit; labels and instruction are not inputs."""
+        state = self.forward(x, z)
+        reports = []
+        for layer, target in enumerate(state["context"]):
+            soma = state["h"][layer + 1]
+            soma_centered = soma - soma.mean(0)
+            target_centered = target - target.mean(0)
+            variance = soma_centered.square().mean(0)
+            covariance = (soma_centered * target_centered).mean(0)
+            slope = covariance / variance.clamp_min(1e-8)
+            intercept = target.mean(0) - slope * soma.mean(0)
+            self.P[layer].copy_(slope)
+            self.P_bias[layer].copy_(intercept)
+            prediction = slope * soma + intercept
+            residual = target - prediction
+            target_ss = target_centered.square().sum(0)
+            residual_ss = residual.square().sum(0)
+            valid = target_ss > 1e-12
+            r2 = 1.0 - residual_ss[valid] / target_ss[valid]
+            reports.append({
+                "mean_per_cell_r2": float(r2.mean()),
+                "context_rms": float(target.square().mean().sqrt()),
+                "residual_context_rms_ratio": float(
+                    residual.square().mean().sqrt()
+                    / target.square().mean().sqrt().clamp_min(1e-12)),
+                "neutral_observations": int(x.shape[0]),
+                "instruction_observations": 0,
+            })
+        return reports
+
+
+def select_shared_signal(net, layer, instruction, context, soma, condition):
+    raw = instruction + context
+    innovation = raw - net.predictor(layer, soma)
+    if condition == "oracle":
+        used = instruction
+    elif condition == "raw_shared":
+        used = raw
+    elif condition == "innovation":
+        used = innovation
+    elif condition == "matched_raw":
+        scale = (innovation.norm(dim=1, keepdim=True)
+                 / raw.norm(dim=1, keepdim=True).clamp_min(1e-12))
+        used = raw * scale
+    elif condition == "exact_subtraction":
+        used = raw - context
+    else:
+        raise ValueError(f"unknown shared-feedback condition: {condition}")
+    return used, raw, innovation
+
+
+@torch.no_grad()
+def shared_feedback_step(net, x, z, y, condition):
+    """One manual modified-KP step with a shared apical measurement."""
+    state = net.forward(x, z)
+    h, u, context = state["h"], state["u"], state["context"]
+    probabilities = torch.softmax(state["logits"], dim=1)
+    output_instruction = F.one_hot(y, num_classes=2).to(net.dtype) - probabilities
+    batch = x.shape[0]
+
+    hidden_deltas = [None] * net.config.hidden_layers
+    used_fields = [None] * net.config.hidden_layers
+    raw_fields = [None] * net.config.hidden_layers
+    innovation_fields = [None] * net.config.hidden_layers
+    child_delta = output_instruction
+    for layer in reversed(range(net.config.hidden_layers)):
+        instruction = child_delta @ net.Q[layer + 1]
+        used, raw, innovation = select_shared_signal(
+            net, layer, instruction, context[layer], h[layer + 1], condition)
+        delta = used * (1.0 - torch.tanh(u[layer]).square())
+        hidden_deltas[layer] = delta
+        used_fields[layer] = used
+        raw_fields[layer] = raw
+        innovation_fields[layer] = innovation
+        child_delta = delta
+
+    directions = [hidden_deltas[0].t() @ h[0] / batch]
+    for layer in range(1, net.config.hidden_layers):
+        directions.append(hidden_deltas[layer].t() @ h[layer] / batch)
+    directions.append(output_instruction.t() @ h[-2] / batch)
+
+    for layer, direction in enumerate(directions):
+        net.mW[layer].mul_(net.config.momentum).add_(
+            direction - net.config.weight_decay * net.W[layer])
+        net.W[layer].add_(net.mW[layer], alpha=net.config.learning_rate)
+        if layer > 0:
+            net.mQ[layer].mul_(net.config.momentum).add_(
+                direction - net.config.weight_decay * net.Q[layer])
+            net.Q[layer].add_(
+                net.mQ[layer], alpha=net.config.reciprocal_learning_rate)
+
+    loss = F.cross_entropy(state["logits"], y)
+    return float(loss), {
+        "directions": directions,
+        "used": used_fields,
+        "raw": raw_fields,
+        "innovation": innovation_fields,
+        "context": context,
+    }
+
+
+@torch.no_grad()
+def evaluate_shared_feedback(net, x, z, y, context_enabled=True,
+                             batch_size=512):
+    correct = 0
+    total_loss = 0.0
+    for start in range(0, x.shape[0], batch_size):
+        stop = min(start + batch_size, x.shape[0])
+        logits = net.forward(
+            x[start:stop], z[start:stop], context_enabled=context_enabled)["logits"]
+        total_loss += float(F.cross_entropy(
+            logits, y[start:stop], reduction="sum"))
+        correct += int((logits.argmax(1) == y[start:stop]).sum())
+    return {
+        "accuracy": correct / x.shape[0],
+        "loss": total_loss / x.shape[0],
+    }
+