algorithm: estimate causal credit in vectorizer space

1 files changed, 200 insertions, 1 deletions

diff --git a/experiments/conv_local_smoke.py b/experiments/conv_local_smoke.py
index 1ea0cea..8070d9f 100644
--- a/experiments/conv_local_smoke.py
+++ b/experiments/conv_local_smoke.py
@@ -9,7 +9,8 @@ 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, CIFARSDILResNet, ConvSDILConfig,
                        channel_subspace_apical_calibration, conv_local_step,
-                       simultaneous_conv_node_perturbation)
+                       simultaneous_conv_node_perturbation,
+                       vectorizer_subspace_apical_calibration)
 
 
 def architecture_checks():
@@ -391,6 +392,7 @@ def channel_subspace_estimator_check():
                        for actual, expected in zip(
                            net.A + net.A_gate, expected_a + expected_g))
     assert update_error < 1e-14
+
     return {
         "channel_subspace_moment_cosine": cosine,
         "channel_subspace_moment_norm_ratio": norm_ratio,
@@ -399,6 +401,202 @@ def channel_subspace_estimator_check():
     }
 
 
+def vectorizer_subspace_estimator_check():
+    """Direct A/G perturbations are unbiased and lower variance at batch 128."""
+    torch.manual_seed(81)
+    batch = 128
+    output_dim = 5
+    hiddens = [
+        torch.randn(batch, 2, 4, 4, dtype=torch.float64),
+        torch.randn(batch, 3, 2, 2, dtype=torch.float64),
+    ]
+    negative_gradients = [torch.randn_like(value) for value in hiddens]
+    output_signal = torch.randn(batch, output_dim, dtype=torch.float64)
+    exact_pairs = []
+    for hidden, target in zip(hiddens, negative_gradients):
+        exact_pairs.extend([
+            target.mean(dim=(2, 3)).t() @ output_signal / batch,
+            (target * torch.tanh(hidden)).mean(dim=(2, 3)).t()
+            @ output_signal / batch,
+        ])
+    exact = torch.cat([value.flatten() for value in exact_pairs])
+    coefficient_sum = torch.zeros_like(exact)
+    vectorizer_sum = torch.zeros_like(exact)
+    coefficient_mse = 0.0
+    vectorizer_mse = 0.0
+    directions = 2048
+    generator = torch.Generator(device="cpu").manual_seed(401)
+    inverse_sqrt_two = 1.0 / (2.0 ** 0.5)
+    for _ in range(directions):
+        random_coefficients = []
+        hidden_directions = []
+        for hidden in hiddens:
+            shape = (batch, hidden.shape[1])
+            base = torch.empty(shape, dtype=hidden.dtype).bernoulli_(
+                0.5, generator=generator).mul_(2).sub_(1)
+            gate = torch.empty_like(base).bernoulli_(
+                0.5, generator=generator).mul_(2).sub_(1)
+            random_coefficients.append((base, gate))
+            hidden_directions.append((
+                base[:, :, None, None]
+                + torch.tanh(hidden) * gate[:, :, None, None])
+                * inverse_sqrt_two)
+        directional = -sum((target * direction).sum()
+                           for target, direction in zip(
+                               negative_gradients, hidden_directions))
+        coefficient_values = []
+        for hidden, (base, gate) in zip(hiddens, random_coefficients):
+            spatial = hidden.shape[2] * hidden.shape[3]
+            base_target = -(2.0 ** 0.5) * directional * base / spatial
+            gate_target = -(2.0 ** 0.5) * directional * gate / spatial
+            coefficient_values.extend([
+                base_target.t() @ output_signal / batch,
+                gate_target.t() @ output_signal / batch,
+            ])
+        coefficient_sample = torch.cat(
+            [value.flatten() for value in coefficient_values])
+
+        random_matrices = []
+        hidden_directions = []
+        for hidden in hiddens:
+            shape = (hidden.shape[1], output_dim)
+            base = torch.empty(shape, dtype=hidden.dtype).bernoulli_(
+                0.5, generator=generator).mul_(2).sub_(1)
+            gate = torch.empty_like(base).bernoulli_(
+                0.5, generator=generator).mul_(2).sub_(1)
+            base_field = output_signal @ base.t()
+            gate_field = output_signal @ gate.t()
+            random_matrices.append((base, gate))
+            hidden_directions.append((
+                base_field[:, :, None, None]
+                + torch.tanh(hidden) * gate_field[:, :, None, None])
+                * inverse_sqrt_two)
+        directional = -sum((target * direction).sum()
+                           for target, direction in zip(
+                               negative_gradients, hidden_directions))
+        vectorizer_values = []
+        for hidden, (base, gate) in zip(hiddens, random_matrices):
+            spatial = hidden.shape[2] * hidden.shape[3]
+            scale = -(2.0 ** 0.5) * directional / (batch * spatial)
+            vectorizer_values.extend([scale * base, scale * gate])
+        vectorizer_sample = torch.cat(
+            [value.flatten() for value in vectorizer_values])
+        coefficient_sum.add_(coefficient_sample)
+        vectorizer_sum.add_(vectorizer_sample)
+        coefficient_mse += float((coefficient_sample - exact).square().mean())
+        vectorizer_mse += float((vectorizer_sample - exact).square().mean())
+    coefficient_mean = coefficient_sum / directions
+    vectorizer_mean = vectorizer_sum / directions
+    vectorizer_cosine = float(F.cosine_similarity(
+        vectorizer_mean, exact, dim=0))
+    vectorizer_norm_ratio = float(vectorizer_mean.norm() / exact.norm())
+    variance_ratio = vectorizer_mse / coefficient_mse
+    assert vectorizer_cosine > 0.95
+    assert 0.90 < vectorizer_norm_ratio < 1.10
+    assert variance_ratio < 0.25
+
+    # Match the executable forward-only derivative and its exact A/G update.
+    torch.manual_seed(3)
+    net = CIFARSDILResNet(
+        depth=8, base_width=2, seed=82, dtype=torch.float64,
+        vectorizer_mode="channel_gated")
+    x = torch.randn(2, 3, 32, 32, dtype=torch.float64)
+    y = torch.tensor([1, 6])
+    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()
+    output_error = (torch.softmax(clean["logits"].detach(), dim=1)
+                    - F.one_hot(y, 10).to(torch.float64))
+    _, diagnostics = vectorizer_subspace_apical_calibration(
+        net, x, y, clean, output_error, sigma=1e-6,
+        n_directions=1, eta=0.0,
+        generator=torch.Generator(device="cpu").manual_seed(503),
+        return_diagnostics=True)
+    finite_difference = diagnostics["directional_derivatives"][0][
+        "scaled_directional"]
+    exact_directional = sum(
+        (x.shape[0] * hidden.grad * direction).sum()
+        for hidden, direction in zip(
+            clean["hiddens"], diagnostics["directions"][0]["hidden"]))
+    jvp_relative = float((finite_difference - exact_directional).abs()
+                         / exact_directional.abs().clamp_min(1e-12))
+    assert jvp_relative < 2e-6
+    for parameter in parameters:
+        parameter.requires_grad_(False)
+
+    eta = 0.0017
+    before_a = [value.clone() for value in net.A]
+    before_g = [value.clone() for value in net.A_gate]
+    expected_a = []
+    expected_g = []
+    for index, hidden in enumerate(clean["hiddens"]):
+        base = output_error @ before_a[index].t()
+        gate_coefficient = output_error @ before_g[index].t()
+        gate = torch.tanh(hidden.detach())
+        mean = gate.mean(dim=(2, 3))
+        second = gate.square().mean(dim=(2, 3))
+        base_prediction = (base + mean * gate_coefficient).t() @ output_error / 2
+        gate_prediction = (
+            mean * base + second * gate_coefficient).t() @ output_error / 2
+        expected_a.append(before_a[index] + eta * (
+            diagnostics["target_base"][index] - base_prediction))
+        expected_g.append(before_g[index] + eta * (
+            diagnostics["target_gate"][index] - gate_prediction))
+    vectorizer_subspace_apical_calibration(
+        net, x, y, clean, output_error, sigma=1e-6,
+        n_directions=1, eta=eta,
+        generator=torch.Generator(device="cpu").manual_seed(503))
+    update_error = max(float((actual - expected).abs().max())
+                       for actual, expected in zip(
+                           net.A + net.A_gate, expected_a + expected_g))
+    assert update_error < 1e-14
+
+    batchnorm = CIFARSDILResNet(
+        depth=8, base_width=2, seed=83, dtype=torch.float64,
+        normalization="batchnorm", residual_scale=1.0,
+        vectorizer_mode="channel_gated")
+    xb = torch.randn(3, 3, 32, 32, dtype=torch.float64)
+    yb = torch.tensor([0, 4, 9])
+    parameters = (batchnorm.W + batchnorm.gamma + batchnorm.beta
+                  + [batchnorm.W_out, batchnorm.b_out])
+    for parameter in parameters:
+        parameter.requires_grad_(True)
+    clean_b = batchnorm.forward(xb, training=True, update_stats=False)
+    for hidden in clean_b["hiddens"]:
+        hidden.retain_grad()
+    F.cross_entropy(clean_b["logits"], yb).backward()
+    output_b = (torch.softmax(clean_b["logits"].detach(), dim=1)
+                - F.one_hot(yb, 10).to(torch.float64))
+    _, diagnostics_b = vectorizer_subspace_apical_calibration(
+        batchnorm, xb, yb, clean_b, output_b, sigma=1e-6,
+        n_directions=1, eta=0.0,
+        generator=torch.Generator(device="cpu").manual_seed(509),
+        return_diagnostics=True)
+    finite_b = diagnostics_b["directional_derivatives"][0][
+        "scaled_directional"]
+    exact_b = sum(
+        (xb.shape[0] * hidden.grad * direction).sum()
+        for hidden, direction in zip(
+            clean_b["hiddens"], diagnostics_b["directions"][0]["hidden"]))
+    batchnorm_relative = float(
+        (finite_b - exact_b).abs() / exact_b.abs().clamp_min(1e-12))
+    assert batchnorm_relative < 2e-6
+    for parameter in parameters:
+        parameter.requires_grad_(False)
+    return {
+        "vectorizer_subspace_mean_cosine": vectorizer_cosine,
+        "vectorizer_subspace_mean_norm_ratio": vectorizer_norm_ratio,
+        "vectorizer_vs_coefficient_mse_ratio": variance_ratio,
+        "vectorizer_subspace_jvp_relative_error": jvp_relative,
+        "vectorizer_subspace_batchnorm_jvp_relative_error": batchnorm_relative,
+        "vectorizer_subspace_delta_rule_absolute_error": update_error,
+    }
+
+
 def apical_learning_checks():
     torch.manual_seed(11)
     net = CIFARSDILResNet(depth=8, base_width=2, seed=6)
@@ -503,6 +701,7 @@ def main():
     report.update(exact_batchnorm_local_gradient_check())
     report.update(perturbation_estimator_check())
     report.update(channel_subspace_estimator_check())
+    report.update(vectorizer_subspace_estimator_check())
     report.update(apical_learning_checks())
     print(report)
     print("ALL CONVOLUTIONAL LOCAL-ELIGIBILITY CHECKS PASSED")