diff options
Diffstat (limited to 'experiments/conv_local_smoke.py')
| -rw-r--r-- | experiments/conv_local_smoke.py | 149 |
1 files changed, 148 insertions, 1 deletions
diff --git a/experiments/conv_local_smoke.py b/experiments/conv_local_smoke.py index 99ca3b6..1ea0cea 100644 --- a/experiments/conv_local_smoke.py +++ b/experiments/conv_local_smoke.py @@ -8,7 +8,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, - conv_local_step, simultaneous_conv_node_perturbation) + channel_subspace_apical_calibration, conv_local_step, + simultaneous_conv_node_perturbation) def architecture_checks(): @@ -266,6 +267,138 @@ def perturbation_estimator_check(): } +def channel_subspace_estimator_check(): + """The structured estimator targets representable base/gate moments.""" + torch.manual_seed(71) + batch = 3 + 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] + estimated_base = [torch.zeros( + batch, value.shape[1], dtype=value.dtype) for value in hiddens] + estimated_gate = [torch.zeros_like(value) for value in estimated_base] + generator = torch.Generator(device="cpu").manual_seed(211) + directions = 4096 + inverse_sqrt_two = 1.0 / (2.0 ** 0.5) + for _ in range(directions): + hidden_directions = [] + base_random = [] + gate_random = [] + 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) + hidden_directions.append(( + base[:, :, None, None] + + torch.tanh(hidden) * gate[:, :, None, None]) + * inverse_sqrt_two) + base_random.append(base) + gate_random.append(gate) + # The exact loss derivative uses g=-negative_gradient and remains + # per-example without BatchNorm coupling. + directional = -sum( + (gradient * direction).flatten(1).sum(dim=1) + for gradient, direction in zip( + negative_gradients, hidden_directions)) + for index, (hidden, base, gate) in enumerate(zip( + hiddens, base_random, gate_random)): + spatial = hidden.shape[2] * hidden.shape[3] + scale = -(2.0 ** 0.5) / (spatial * directions) + estimated_base[index].add_(directional[:, None] * base, alpha=scale) + estimated_gate[index].add_(directional[:, None] * gate, alpha=scale) + exact_base = [value.mean(dim=(2, 3)) for value in negative_gradients] + exact_gate = [(value * torch.tanh(hidden)).mean(dim=(2, 3)) + for value, hidden in zip(negative_gradients, hiddens)] + estimated = torch.cat([ + value.flatten() for pair in zip(estimated_base, estimated_gate) + for value in pair]) + exact = torch.cat([ + value.flatten() for pair in zip(exact_base, exact_gate) + for value in pair]) + cosine = float(F.cosine_similarity(estimated, exact, dim=0)) + norm_ratio = float(estimated.norm() / exact.norm()) + assert cosine > 0.985 + assert 0.90 < norm_ratio < 1.10 + + # The executable antithetic implementation must match the same structured + # directional derivative, not an autograd surrogate. + # Use a fixed nondegenerate point. Width-one, zero-bias ReLU networks can + # contain structurally exact-zero preactivations, where central differences + # and PyTorch's chosen subgradient need not agree even as sigma -> 0. + torch.manual_seed(3) + net = CIFARSDILResNet( + depth=8, base_width=2, seed=72, dtype=torch.float64, + vectorizer_mode="channel_gated") + x = torch.randn(2, 3, 32, 32, dtype=torch.float64) + y = torch.tensor([2, 8]) + 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() + loss = F.cross_entropy(clean["logits"], y) + loss.backward() + output_signal = (torch.softmax(clean["logits"].detach(), dim=1) + - F.one_hot(y, 10).to(torch.float64)) + _, diagnostics = channel_subspace_apical_calibration( + net, x, y, clean, output_signal, sigma=1e-6, + n_directions=1, eta=0.0, + generator=torch.Generator(device="cpu").manual_seed(307), + return_diagnostics=True) + hidden_direction = diagnostics["directions"][0]["hidden"] + finite_difference = diagnostics["directional_derivatives"][0][ + "scaled_directional"] + exact_directional = sum( + (x.shape[0] * hidden.grad * direction).flatten(1).sum(dim=1) + for hidden, direction in zip(clean["hiddens"], hidden_direction)) + relative = ((finite_difference - exact_directional).abs() + / exact_directional.abs().clamp_min(1e-12)) + assert float(relative.max()) < 2e-6 + for parameter in parameters: + parameter.requires_grad_(False) + + # The local A/G update must equal the full-field delta rule after replacing + # only its two target moments with the structured causal estimates. + eta = 0.0023 + 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_signal @ before_a[index].t() + gate_coefficient = output_signal @ before_g[index].t() + gate = torch.tanh(hidden.detach()) + mean = gate.mean(dim=(2, 3)) + second = gate.square().mean(dim=(2, 3)) + base_error = diagnostics["target_base"][index] - ( + base + mean * gate_coefficient) + gate_error = diagnostics["target_gate"][index] - ( + mean * base + second * gate_coefficient) + expected_a.append(before_a[index] + eta * ( + base_error.t() @ output_signal / x.shape[0])) + expected_g.append(before_g[index] + eta * ( + gate_error.t() @ output_signal / x.shape[0])) + channel_subspace_apical_calibration( + net, x, y, clean, output_signal, sigma=1e-6, + n_directions=1, eta=eta, + generator=torch.Generator(device="cpu").manual_seed(307)) + 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 + return { + "channel_subspace_moment_cosine": cosine, + "channel_subspace_moment_norm_ratio": norm_ratio, + "channel_subspace_jvp_relative_error": float(relative.max()), + "channel_subspace_delta_rule_absolute_error": update_error, + } + + def apical_learning_checks(): torch.manual_seed(11) net = CIFARSDILResNet(depth=8, base_width=2, seed=6) @@ -342,6 +475,19 @@ def apical_learning_checks(): 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]) + + gated_weights_before = [value.clone() for value in gated.A + gated.A_gate] + gated_result = conv_local_step( + gated, x[:2], y[:2], + ConvSDILConfig( + eta=1e-3, eta_A=1e-3, momentum=0.0, weight_decay=0.0, + pert_every=1, apical_calibration_mode="channel_subspace"), + step=0, generator=torch.Generator(device="cpu").manual_seed(17)) + assert gated_result["did_perturb"] + assert all(torch.isfinite(torch.tensor(value)) for value in + gated_result["calibration"].values()) + assert any(not torch.equal(before, after) for before, after in zip( + gated_weights_before, gated.A + gated.A_gate)) return {"apical_mse_ratio": after / before, "gated_apical_mse_ratio": gated_after / gated_before, "gated_vectorizer_parameter_reduction": ( @@ -356,6 +502,7 @@ def main(): report = exact_local_gradient_check() report.update(exact_batchnorm_local_gradient_check()) report.update(perturbation_estimator_check()) + report.update(channel_subspace_estimator_check()) report.update(apical_learning_checks()) print(report) print("ALL CONVOLUTIONAL LOCAL-ELIGIBILITY CHECKS PASSED") |