oral-a: add convolutional apical calibration

2 files changed, 328 insertions, 1 deletions

diff --git a/experiments/conv_local_smoke.py b/experiments/conv_local_smoke.py
index 35073a6..4a0f47d 100644
--- a/experiments/conv_local_smoke.py
+++ b/experiments/conv_local_smoke.py
@@ -7,7 +7,8 @@ import torch
 import torch.nn.functional as F
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from sdil.conv import CIFARLocalResNet
+from sdil.conv import (CIFARLocalResNet, CIFARSDILResNet, ConvSDILConfig,
+                       conv_local_step, simultaneous_conv_node_perturbation)
 
 
 def architecture_checks():
@@ -111,10 +112,93 @@ def perturbation_checks():
         raise AssertionError("short perturbation list was accepted")
 
 
+def perturbation_estimator_check():
+    """Antithetic finite differences equal the simultaneous hidden JVP."""
+    torch.manual_seed(3)
+    batch = 2
+    net = CIFARSDILResNet(
+        depth=8, base_width=2, seed=4, dtype=torch.float64)
+    x = torch.randn(batch, 3, 32, 32, dtype=torch.float64)
+    y = torch.tensor([2, 7])
+    parameters = net.W + [net.W_out, net.b_out]
+    for parameter in parameters:
+        parameter.requires_grad_(True)
+    clean = net.forward(x, return_cache=True)
+    for hidden in clean["hiddens"]:
+        hidden.retain_grad()
+    F.cross_entropy(clean["logits"], y).backward()
+    generator = torch.Generator(device="cpu").manual_seed(99)
+    targets, diagnostics = simultaneous_conv_node_perturbation(
+        net, x, y, clean, sigma=1e-6, n_directions=1,
+        generator=generator, return_diagnostics=True)
+    directions = diagnostics["directions"][0]
+    finite_difference = diagnostics["directional_derivatives"][0]
+    exact = sum(
+        (batch * hidden.grad * direction).flatten(1).sum(dim=1)
+        for hidden, direction in zip(clean["hiddens"], directions))
+    relative = (finite_difference - exact).abs() / exact.abs().clamp_min(1e-12)
+    assert float(relative.max()) < 2e-6
+    for target, direction in zip(targets, directions):
+        expected = -finite_difference[:, None, None, None] * direction
+        assert torch.equal(target, expected)
+    for parameter in parameters:
+        parameter.requires_grad_(False)
+    return float(relative.max())
+
+
+def apical_learning_checks():
+    torch.manual_seed(11)
+    net = CIFARSDILResNet(depth=8, base_width=2, seed=6)
+    x = torch.randn(8, 3, 32, 32)
+    y = torch.arange(8) % 10
+    clean = net.forward(x, return_cache=True)
+    output_signal = (torch.softmax(clean["logits"], dim=1)
+                     - F.one_hot(y, 10))
+    prediction, _, _ = net.apical_components(
+        output_signal, clean["hiddens"], use_residual=True)
+    targets = [torch.randn_like(value) * 0.01 for value in prediction]
+    before = sum(float((target - value).square().sum())
+                 for target, value in zip(targets, prediction))
+    net.calibrate_apical(output_signal, prediction, targets, eta=0.1)
+    after_prediction, _, _ = net.apical_components(
+        output_signal, clean["hiddens"], use_residual=True)
+    after = sum(float((target - value).square().sum())
+                for target, value in zip(targets, after_prediction))
+    assert after < before
+
+    predictor_net = CIFARSDILResNet(depth=8, base_width=2, seed=8)
+    hiddens = [torch.randn(64, *shape) for shape in predictor_net.hidden_shapes]
+    initial = predictor_net.predictor_step(hiddens, eta=0.1, nuisance_scale=0.5)
+    final = initial
+    for _ in range(60):
+        final = predictor_net.predictor_step(hiddens, eta=0.1, nuisance_scale=0.5)
+    assert final < initial * 1e-3
+
+    weights_before = [weight.clone() for weight in net.W]
+    result = conv_local_step(
+        net, x[:2], y[:2],
+        ConvSDILConfig(
+            eta=1e-3, eta_A=1e-3, momentum=0.0, weight_decay=0.0,
+            pert_every=1),
+        step=0, generator=torch.Generator(device="cpu").manual_seed(7))
+    assert result["did_perturb"] and result["calibration"] is not None
+    assert torch.isfinite(torch.tensor(list(
+        value for key, value in result.items()
+        if isinstance(value, float) and key != "predictor_mse"))).all()
+    assert any(not torch.equal(before_weight, after_weight)
+               for before_weight, after_weight in zip(weights_before, net.W))
+    assert all(not parameter.requires_grad
+               for parameter in net.W + [net.W_out, net.b_out])
+    return {"apical_mse_ratio": after / before,
+            "predictor_mse_ratio": final / initial}
+
+
 def main():
     architecture_checks()
     perturbation_checks()
     report = exact_local_gradient_check()
+    report["perturbation_jvp_max_relative_error"] = perturbation_estimator_check()
+    report.update(apical_learning_checks())
     print(report)
     print("ALL CONVOLUTIONAL LOCAL-ELIGIBILITY CHECKS PASSED")
 
diff --git a/sdil/conv.py b/sdil/conv.py
index 5971dea..7f805fa 100644
--- a/sdil/conv.py
+++ b/sdil/conv.py
@@ -285,6 +285,249 @@ class CIFARLocalResNet:
         return float(loss.detach())
 
 
+class CIFARSDILResNet(CIFARLocalResNet):
+    """CIFAR local ResNet with per-unit apical vectorizers and predictors.
+
+    Each spatial feature unit has a class-error vectorizer ``A_l``.  A full
+    spatial template is used rather than broadcasting one coefficient over a
+    channel, because the true credit assigned to an early convolution is
+    strongly position dependent.  The predictor remains Harnett-faithful and
+    diagonal: each unit fits its own affine soma--apical relation.
+    """
+
+    def __init__(self, *args, a_scale=1.0, apical_seed=None, **kwargs):
+        model_seed = kwargs.get("seed", 0)
+        super().__init__(*args, **kwargs)
+        generator = torch.Generator(device="cpu").manual_seed(
+            model_seed + 10007 if apical_seed is None else apical_seed)
+        nuisance_generator = torch.Generator(device="cpu").manual_seed(
+            model_seed + 20011 if apical_seed is None else apical_seed + 1)
+        self.A = []
+        self.P = []
+        self.P_bias = []
+        self.Bnuis = []
+        for channels, height, width in self.hidden_shapes:
+            units = channels * height * width
+            # A global-average readout makes early per-unit gradients shrink
+            # approximately as 1/(H*W).  This scale keeps fixed-DFA controls
+            # finite while learned A remains free to change its gain.
+            std = a_scale / (height * width * math.sqrt(self.n_classes))
+            self.A.append((torch.randn(units, self.n_classes, generator=generator)
+                           * std).to(device=self.device, dtype=self.dtype))
+            shape = (channels, height, width)
+            self.P.append(torch.zeros(shape, device=self.device, dtype=self.dtype))
+            self.P_bias.append(torch.zeros(shape, device=self.device, dtype=self.dtype))
+            self.Bnuis.append(torch.exp(
+                0.25 * torch.randn(shape, generator=nuisance_generator)
+            ).to(device=self.device, dtype=self.dtype))
+
+    @property
+    def n_apical_parameters(self):
+        return (sum(value.numel() for value in self.A)
+                + sum(value.numel() for value in self.P)
+                + sum(value.numel() for value in self.P_bias))
+
+    def instruction(self, index, output_signal):
+        shape = self.hidden_shapes[index]
+        return (output_signal @ self.A[index].t()).reshape(
+            output_signal.shape[0], *shape)
+
+    def apical_components(self, output_signal, hiddens, nuisance_scale=0.0,
+                          use_residual=True):
+        """Return teaching, raw apical, and innovation at every population."""
+        if len(hiddens) != self.n_hidden:
+            raise ValueError("one somatic state is required per apical population")
+        teaching = []
+        raw_apical = []
+        innovations = []
+        for index, hidden in enumerate(hiddens):
+            instruction = self.instruction(index, output_signal)
+            traffic = nuisance_scale * self.Bnuis[index] * hidden
+            raw = instruction + traffic
+            baseline = self.P[index] * hidden + self.P_bias[index]
+            innovation = raw - baseline
+            teaching.append(innovation if use_residual else raw)
+            raw_apical.append(raw)
+            innovations.append(innovation)
+        return teaching, raw_apical, innovations
+
+    @torch.no_grad()
+    def predictor_step(self, hiddens, eta, nuisance_scale):
+        """Neutral-period normalized LMS fit to soma-predictable traffic."""
+        squared_error = 0.0
+        units = 0
+        for index, hidden in enumerate(hiddens):
+            target = nuisance_scale * self.Bnuis[index] * hidden
+            residual = target - self.P[index] * hidden - self.P_bias[index]
+            centered_h = hidden - hidden.mean(dim=0)
+            centered_r = residual - residual.mean(dim=0)
+            variance = centered_h.square().mean(dim=0)
+            self.P[index].add_(
+                (centered_r * centered_h).mean(dim=0) / (variance + 1e-6),
+                alpha=eta)
+            self.P_bias[index].add_(residual.mean(dim=0), alpha=eta)
+            squared_error += float(residual.square().sum())
+            units += residual.numel()
+        return squared_error / units
+
+    @torch.no_grad()
+    def calibrate_apical(self, output_signal, predicted_teaching, targets, eta):
+        """Local delta rule fitting innovation to causal perturbation targets."""
+        if not (len(predicted_teaching) == len(targets) == self.n_hidden):
+            raise ValueError("calibration lists must cover every hidden population")
+        batch = output_signal.shape[0]
+        before_error = 0.0
+        target_power = 0.0
+        dot = 0.0
+        prediction_power = 0.0
+        for index, (prediction, target) in enumerate(zip(predicted_teaching, targets)):
+            error = target - prediction
+            flat_error = error.flatten(1)
+            self.A[index].add_(flat_error.t() @ output_signal / batch, alpha=eta)
+            before_error += float(error.square().sum())
+            target_power += float(target.square().sum())
+            prediction_power += float(prediction.square().sum())
+            dot += float((target * prediction).sum())
+        denominator = math.sqrt(target_power * prediction_power)
+        return {
+            "calibration_mse": before_error / sum(
+                target.numel() for target in targets),
+            "target_power": target_power / sum(target.numel() for target in targets),
+            "prediction_target_cosine": dot / denominator if denominator else 0.0,
+        }
+
+
+@torch.no_grad()
+def simultaneous_conv_node_perturbation(net, x, y, clean_forward, sigma=1e-2,
+                                        n_directions=1, generator=None,
+                                        return_diagnostics=False):
+    """Forward-only antithetic targets for all convolutional populations.
+
+    Independent Rademacher interventions are injected into every hidden map in
+    the same plus/minus evaluations.  Cross-layer interference is zero mean and
+    is handled by the variance theorem in ``THEORY.md``.  Plus and minus trials
+    are concatenated into one expanded batch for GPU efficiency.
+    """
+    if sigma <= 0:
+        raise ValueError("perturbation sigma must be positive")
+    if n_directions < 1:
+        raise ValueError("n_directions must be positive")
+    if len(clean_forward["hiddens"]) != net.n_hidden:
+        raise ValueError("clean forward does not match network hidden populations")
+    if generator is None:
+        generator = torch.Generator(device=x.device).manual_seed(0)
+    targets = [torch.zeros_like(hidden) for hidden in clean_forward["hiddens"]]
+    diagnostic_directions = []
+    diagnostic_derivatives = []
+    expanded_x = torch.cat((x, x), dim=0)
+    expanded_y = torch.cat((y, y), dim=0)
+    for _ in range(n_directions):
+        directions = []
+        perturbations = []
+        for hidden in clean_forward["hiddens"]:
+            direction = torch.empty_like(hidden).bernoulli_(
+                0.5, generator=generator).mul_(2).sub_(1)
+            directions.append(direction)
+            perturbations.append(torch.cat(
+                (sigma * direction, -sigma * direction), dim=0))
+        perturbed = net.forward(expanded_x, perturbations=perturbations)
+        losses = F.cross_entropy(perturbed["logits"], expanded_y, reduction="none")
+        plus, minus = losses.chunk(2)
+        directional = (plus - minus) / (2.0 * sigma)
+        for index, direction in enumerate(directions):
+            expand = directional.reshape(
+                directional.shape[0], *([1] * (direction.ndim - 1)))
+            targets[index].add_(-expand * direction / n_directions)
+        if return_diagnostics:
+            diagnostic_directions.append(directions)
+            diagnostic_derivatives.append(directional)
+    if return_diagnostics:
+        return targets, {
+            "directions": diagnostic_directions,
+            "directional_derivatives": diagnostic_derivatives,
+        }
+    return targets
+
+
+@dataclass
+class ConvSDILConfig:
+    eta: float = 0.01
+    eta_output: float = None
+    eta_A: float = 0.01
+    eta_P: float = 0.01
+    momentum: float = 0.9
+    weight_decay: float = 5e-4
+    learn_A: bool = True
+    learn_P: bool = False
+    use_residual: bool = True
+    nuisance_scale: float = 0.0
+    pert_sigma: float = 1e-2
+    pert_every: int = 4
+    pert_directions: int = 1
+    direct_node_perturbation: bool = False
+
+    def validate(self):
+        if self.eta <= 0 or (self.eta_output is not None and self.eta_output <= 0):
+            raise ValueError("forward learning rates must be positive")
+        if self.eta_A < 0 or self.eta_P < 0:
+            raise ValueError("apical learning rates must be nonnegative")
+        if self.pert_every < 1 or self.pert_directions < 1:
+            raise ValueError("perturbation cadence/directions must be positive")
+        if self.direct_node_perturbation and self.pert_every != 1:
+            raise ValueError("direct node perturbation requires a target every step")
+
+
+def conv_local_step(net, x, y, config, step, generator=None):
+    """One DFA/learned-feedback/direct-NP minibatch update without autograd."""
+    config.validate()
+    with torch.no_grad():
+        forward = net.forward(x, return_cache=True)
+        logits = forward["logits"]
+        loss = F.cross_entropy(logits, y)
+        output_error = (torch.softmax(logits, dim=1)
+                        - F.one_hot(y, net.n_classes).to(logits.dtype))
+        teaching, raw, innovations = net.apical_components(
+            output_error, forward["hiddens"], config.nuisance_scale,
+            config.use_residual)
+        did_perturb = ((config.learn_A or config.direct_node_perturbation)
+                       and step % config.pert_every == 0)
+        targets = None
+        if did_perturb:
+            targets = simultaneous_conv_node_perturbation(
+                net, x, y, forward, sigma=config.pert_sigma,
+                n_directions=config.pert_directions, generator=generator)
+        weight_teaching = targets if config.direct_node_perturbation else teaching
+        if weight_teaching is None:
+            raise RuntimeError("direct perturbation target is unavailable")
+        directions, output_weight, output_bias = net.local_ascent_directions(
+            weight_teaching, output_error, forward)
+        net.apply_ascent(
+            directions, output_weight, output_bias,
+            eta_hidden=config.eta, eta_output=config.eta_output,
+            momentum=config.momentum, weight_decay=config.weight_decay)
+        calibration = None
+        if did_perturb and config.learn_A:
+            calibration = net.calibrate_apical(
+                output_error, teaching, targets, config.eta_A)
+        predictor_mse = None
+        if config.learn_P:
+            predictor_mse = net.predictor_step(
+                forward["hiddens"], config.eta_P, config.nuisance_scale)
+        return {
+            "loss": float(loss),
+            "did_perturb": did_perturb,
+            "calibration": calibration,
+            "predictor_mse": predictor_mse,
+            "teaching_rms": math.sqrt(sum(float(value.square().sum()) for value in teaching)
+                                      / sum(value.numel() for value in teaching)),
+            "raw_apical_rms": math.sqrt(sum(float(value.square().sum()) for value in raw)
+                                        / sum(value.numel() for value in raw)),
+            "innovation_rms": math.sqrt(
+                sum(float(value.square().sum()) for value in innovations)
+                / sum(value.numel() for value in innovations)),
+        }
+
+
 @torch.no_grad()
 def evaluate_conv(net, loader):
     correct = 0