diff options
| -rw-r--r-- | experiments/conv_local_smoke.py | 129 | ||||
| -rw-r--r-- | experiments/conv_run.py | 32 | ||||
| -rw-r--r-- | sdil/conv.py | 218 |
3 files changed, 329 insertions, 50 deletions
diff --git a/experiments/conv_local_smoke.py b/experiments/conv_local_smoke.py index 4a0f47d..d2428eb 100644 --- a/experiments/conv_local_smoke.py +++ b/experiments/conv_local_smoke.py @@ -24,6 +24,11 @@ def architecture_checks(): assert net.n_forward_parameters == parameters assert len(net.blocks) == 3 * ((depth - 2) // 6) assert len(net.W) + 1 == depth + batchnorm_parameters = {8: 75290, 20: 269722, 32: 464154, 56: 853018} + for depth, parameters in batchnorm_parameters.items(): + net = CIFARLocalResNet( + depth=depth, normalization="batchnorm", residual_scale=1.0) + assert net.n_forward_parameters == parameters for depth in (7, 9, 21): try: CIFARLocalResNet(depth=depth) @@ -62,8 +67,10 @@ def exact_local_gradient_check(): teaching = [-batch * hidden.grad for hidden in forward["hiddens"]] output_error = (torch.softmax(forward["logits"].detach(), dim=1) - F.one_hot(y, local.n_classes)) - directions, out_direction, bias_direction = local.local_ascent_directions( + (directions, gamma_directions, beta_directions, + out_direction, bias_direction) = local.local_ascent_directions( teaching, output_error, forward) + assert gamma_directions == beta_directions == [] relative_errors = [] for direction, parameter in zip(directions, local.W): @@ -94,6 +101,78 @@ def exact_local_gradient_check(): } +def exact_batchnorm_local_gradient_check(): + torch.manual_seed(31) + batch = 4 + x = torch.randn(batch, 3, 32, 32) + y = torch.tensor([0, 1, 2, 3]) + common = dict( + depth=8, base_width=2, seed=29, + normalization="batchnorm", residual_scale=1.0) + local = CIFARLocalResNet(**common) + bp = CIFARLocalResNet(**common) + parameters = local.W + local.gamma + local.beta + [local.W_out, local.b_out] + for parameter in parameters: + parameter.requires_grad_(True) + forward = local.forward( + x, return_cache=True, training=True, update_stats=True) + for hidden in forward["hiddens"]: + hidden.retain_grad() + loss = F.cross_entropy(forward["logits"], y) + loss.backward() + teaching = [-batch * hidden.grad for hidden in forward["hiddens"]] + output_error = (torch.softmax(forward["logits"].detach(), dim=1) + - F.one_hot(y, 10)) + (directions, gamma_directions, beta_directions, + out_direction, bias_direction) = local.local_ascent_directions( + teaching, output_error, forward) + groups = ( + (directions, local.W), + (gamma_directions, local.gamma), + (beta_directions, local.beta), + ) + relative_errors = [] + for direction_group, parameter_group in groups: + for direction, parameter in zip(direction_group, parameter_group): + absolute = (direction + parameter.grad).abs().max() + relative_errors.append(float( + absolute / parameter.grad.abs().max().clamp_min(1e-12))) + assert max(relative_errors) < 3e-5 + for parameter in parameters: + parameter.requires_grad_(False) + + eta = 0.013 + local.apply_ascent( + directions, out_direction, bias_direction, eta, + gamma_directions=gamma_directions, beta_directions=beta_directions) + bp.bp_step(x, y, eta) + parameter_differences = [ + float((left - right).abs().max()) + for left, right in zip( + local.W + local.gamma + local.beta + [local.W_out, local.b_out], + bp.W + bp.gamma + bp.beta + [bp.W_out, bp.b_out])] + running_differences = [ + float((left - right).abs().max()) + for left, right in zip( + local.running_mean + local.running_var, + bp.running_mean + bp.running_var)] + assert max(parameter_differences) < 2e-7 + assert max(running_differences) == 0.0 + + running_before = [value.clone() for value in local.running_mean + local.running_var] + clean = local.forward(x, training=True, update_stats=False) + simultaneous_conv_node_perturbation( + local, x, y, clean, sigma=1e-3, n_directions=1, + generator=torch.Generator(device="cpu").manual_seed(9)) + assert all(torch.equal(before, after) for before, after in zip( + running_before, local.running_mean + local.running_var)) + assert torch.equal(local.logits(x), local.logits(x)) + return { + "batchnorm_max_relative_local_gradient_error": max(relative_errors), + "batchnorm_post_update_parameter_max_error": max(parameter_differences), + } + + def perturbation_checks(): net = CIFARLocalResNet(depth=8, base_width=4, seed=3) x = torch.randn(2, 3, 32, 32) @@ -132,7 +211,9 @@ def perturbation_estimator_check(): 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] + derivative = diagnostics["directional_derivatives"][0] + assert derivative["coupling"] == "per_example_objective" + finite_difference = derivative["scaled_directional"] exact = sum( (batch * hidden.grad * direction).flatten(1).sum(dim=1) for hidden, direction in zip(clean["hiddens"], directions)) @@ -143,7 +224,46 @@ def perturbation_estimator_check(): assert torch.equal(target, expected) for parameter in parameters: parameter.requires_grad_(False) - return float(relative.max()) + no_norm_relative = float(relative.max()) + + batch = 3 + batchnorm = CIFARSDILResNet( + depth=8, base_width=2, seed=14, dtype=torch.float64, + normalization="batchnorm", residual_scale=1.0) + xb = torch.randn(batch, 3, 32, 32, dtype=torch.float64) + yb = torch.tensor([1, 4, 9]) + parameters = (batchnorm.W + batchnorm.gamma + batchnorm.beta + + [batchnorm.W_out, batchnorm.b_out]) + for parameter in parameters: + parameter.requires_grad_(True) + clean = batchnorm.forward(xb, training=True, update_stats=False) + for hidden in clean["hiddens"]: + hidden.retain_grad() + F.cross_entropy(clean["logits"], yb).backward() + targets, diagnostics = simultaneous_conv_node_perturbation( + batchnorm, xb, yb, clean, sigma=1e-6, n_directions=1, + generator=torch.Generator(device="cpu").manual_seed(101), + return_diagnostics=True) + directions = diagnostics["directions"][0] + derivative = diagnostics["directional_derivatives"][0] + assert derivative["coupling"] == "batch_objective" + scaled = derivative["scaled_directional"] + exact_sum_directional = sum( + float((batch * hidden.grad * direction).sum()) + for hidden, direction in zip(clean["hiddens"], directions)) + batchnorm_relative = abs(float(scaled[0]) - exact_sum_directional) / max( + abs(exact_sum_directional), 1e-12) + assert batchnorm_relative < 2e-6 + assert torch.equal(scaled, scaled[:1].expand_as(scaled)) + for target, direction in zip(targets, directions): + expected = -scaled[:, None, None, None] * direction + assert torch.equal(target, expected) + for parameter in parameters: + parameter.requires_grad_(False) + return { + "perturbation_jvp_max_relative_error": no_norm_relative, + "batchnorm_batch_objective_jvp_relative_error": batchnorm_relative, + } def apical_learning_checks(): @@ -197,7 +317,8 @@ def main(): architecture_checks() perturbation_checks() report = exact_local_gradient_check() - report["perturbation_jvp_max_relative_error"] = perturbation_estimator_check() + report.update(exact_batchnorm_local_gradient_check()) + report.update(perturbation_estimator_check()) report.update(apical_learning_checks()) print(report) print("ALL CONVOLUTIONAL LOCAL-ELIGIBILITY CHECKS PASSED") diff --git a/experiments/conv_run.py b/experiments/conv_run.py index 48eebd1..dd3f6c5 100644 --- a/experiments/conv_run.py +++ b/experiments/conv_run.py @@ -91,10 +91,14 @@ def scheduled_lr(base, epoch, args): def build(args): + residual_scale = args.residual_scale + if residual_scale is None and args.normalization == "batchnorm": + residual_scale = 1.0 common = dict( depth=args.depth, base_width=args.width, n_classes=10, device=args.device, seed=args.seed, weight_scale=args.weight_scale, - residual_scale=args.residual_scale) + residual_scale=residual_scale, normalization=args.normalization, + bn_momentum=args.bn_momentum, bn_eps=args.bn_eps) if args.mode == "bp": return CIFARLocalResNet(**common), None net = CIFARSDILResNet( @@ -151,7 +155,8 @@ def work_report(net, mode, counters): + counters["apical_warmup_examples"] + counters["perturbation_forward_examples"]), "logical_batch_loss_queries": counters["logical_batch_loss_queries"], - "scalar_loss_evaluations": counters["scalar_loss_evaluations"], + "causal_scalar_observations": counters["causal_scalar_observations"], + "per_example_cross_entropy_terms": counters["per_example_loss_terms"], "definition": ( "multiply-accumulates in conv/linear maps; one local weight correlation " "equals one forward-weight MAC count; BP reverse is estimated as one " @@ -198,7 +203,8 @@ def run(args): "perturbation_forward_examples": 0, "calibration_event_examples": 0, "logical_batch_loss_queries": 0, - "scalar_loss_evaluations": 0, + "causal_scalar_observations": 0, + "per_example_loss_terms": 0, "perturbation_events": 0, } log = { @@ -208,7 +214,10 @@ def run(args): "provenance": provenance(), "split": split, "architecture": { - "family": "normalization-free CIFAR 6n+2 ResNet, option-A shortcuts", + "family": "CIFAR 6n+2 ResNet, option-A shortcuts", + "normalization": net.normalization, + "bn_momentum": net.bn_momentum if net.normalization == "batchnorm" else None, + "bn_eps": net.bn_eps if net.normalization == "batchnorm" else None, "depth": net.depth, "blocks_per_stage": net.blocks_per_stage, "base_width": net.base_width, @@ -231,7 +240,7 @@ def run(args): except StopIteration: iterator = iter(train) x, _ = next(iterator) - forward = net.forward(x) + forward = net.forward(x, training=True, update_stats=False) net.predictor_step( forward["hiddens"], config.eta_P, config.nuisance_scale) counters["predictor_warmup_examples"] += x.shape[0] @@ -255,7 +264,9 @@ def run(args): 2 * config.pert_directions * batch) counters["calibration_event_examples"] += batch counters["logical_batch_loss_queries"] += 2 * config.pert_directions - counters["scalar_loss_evaluations"] += 2 * config.pert_directions * batch + counters["causal_scalar_observations"] += 2 * config.pert_directions * ( + 1 if net.normalization == "batchnorm" else batch) + counters["per_example_loss_terms"] += 2 * config.pert_directions * batch counters["perturbation_events"] += 1 train.g.set_state(loader_state) log["apical_warmup"] = { @@ -304,7 +315,10 @@ def run(args): 2 * config.pert_directions * batch) counters["calibration_event_examples"] += batch counters["logical_batch_loss_queries"] += 2 * config.pert_directions - counters["scalar_loss_evaluations"] += ( + counters["causal_scalar_observations"] += ( + 2 * config.pert_directions + * (1 if net.normalization == "batchnorm" else batch)) + counters["per_example_loss_terms"] += ( 2 * config.pert_directions * batch) counters["perturbation_events"] += 1 step += 1 @@ -435,6 +449,10 @@ def parse_args(): parser.add_argument("--weight_decay", type=float, default=5e-4) parser.add_argument("--weight_scale", type=float, default=1.0) parser.add_argument("--residual_scale", type=float) + parser.add_argument("--normalization", choices=("batchnorm", "none"), + default="batchnorm") + parser.add_argument("--bn_momentum", type=float, default=0.1) + parser.add_argument("--bn_eps", type=float, default=1e-5) parser.add_argument("--a_scale", type=float, default=1.0) parser.add_argument("--eta_A", type=float, default=0.01) parser.add_argument("--eta_P", type=float, default=0.01) diff --git a/sdil/conv.py b/sdil/conv.py index 996817b..3b8fde1 100644 --- a/sdil/conv.py +++ b/sdil/conv.py @@ -1,11 +1,11 @@ """Convolutional local-learning primitives for CIFAR residual networks. The forward topology is the standard CIFAR ``6n+2`` basic-block family with -option-A identity shortcuts. Batch normalization is deliberately absent: -its cross-example Jacobian obscures what information a synapse needs. A -fixed ``1/sqrt(number of blocks)`` residual multiplier keeps the otherwise -normalization-free network stable and is included explicitly in every local -eligibility calculation. +option-A identity shortcuts. It supports canonical BatchNorm as well as a +normalization-free ablation. BatchNorm's cross-example Jacobian is evaluated +inside the current layer only; the synaptic update still never reads a +downstream weight. Normalization-free networks use an explicit residual +multiplier, which is included in every local eligibility calculation. Forward parameters are plain tensors. The local rule uses only the stored presynaptic activation, a postsynaptic ReLU gate, and a teaching vector at the @@ -33,7 +33,7 @@ class ConvLayerSpec: class CIFARLocalResNet: - """Normalization-free CIFAR ResNet with explicit local eligibilities. + """CIFAR ResNet with explicit local eligibilities. ``depth`` must satisfy ``depth = 6n + 2``. Hidden populations are defined after the stem ReLU, after every block's first ReLU, and after every block @@ -43,7 +43,8 @@ class CIFARLocalResNet: def __init__(self, depth=20, base_width=16, n_classes=10, device="cpu", dtype=torch.float32, seed=0, weight_scale=1.0, - residual_scale=None): + residual_scale=None, normalization="none", bn_momentum=0.1, + bn_eps=1e-5): if depth < 8 or (depth - 2) % 6: raise ValueError(f"CIFAR ResNet depth must be 6n+2 and >=8, got {depth}") if base_width <= 0: @@ -54,6 +55,13 @@ class CIFARLocalResNet: self.n_classes = int(n_classes) self.device = str(device) self.dtype = dtype + if normalization not in ("none", "batchnorm"): + raise ValueError(f"unknown normalization: {normalization}") + self.normalization = normalization + self.bn_momentum = float(bn_momentum) + self.bn_eps = float(bn_eps) + if not 0.0 < self.bn_momentum <= 1.0 or self.bn_eps <= 0: + raise ValueError("invalid BatchNorm momentum/epsilon") self.n_blocks = 3 * self.blocks_per_stage self.residual_scale = (1.0 / math.sqrt(self.n_blocks) if residual_scale is None else float(residual_scale)) @@ -64,6 +72,10 @@ class CIFARLocalResNet: self.W = [] self.layer_specs = [] self.blocks = [] + self.gamma = [] + self.beta = [] + self.running_mean = [] + self.running_var = [] def add_conv(name, in_channels, out_channels, stride, hidden_shape, branch_scale=1.0): @@ -72,6 +84,13 @@ class CIFARLocalResNet: out_channels, in_channels, 3, 3, generator=generator) * (weight_scale * math.sqrt(2.0 / fan_in))) self.W.append(weight.to(device=device, dtype=dtype)) + if normalization == "batchnorm": + self.gamma.append(torch.ones(out_channels, device=device, dtype=dtype)) + self.beta.append(torch.zeros(out_channels, device=device, dtype=dtype)) + self.running_mean.append(torch.zeros( + out_channels, device=device, dtype=dtype)) + self.running_var.append(torch.ones( + out_channels, device=device, dtype=dtype)) self.layer_specs.append(ConvLayerSpec( name=name, stride=stride, padding=1, hidden_shape=tuple(hidden_shape), branch_scale=float(branch_scale))) @@ -114,6 +133,8 @@ class CIFARLocalResNet: self.mW = [torch.zeros_like(weight) for weight in self.W] self.mW_out = torch.zeros_like(self.W_out) self.mb_out = torch.zeros_like(self.b_out) + self.mgamma = [torch.zeros_like(value) for value in self.gamma] + self.mbeta = [torch.zeros_like(value) for value in self.beta] @property def hidden_shapes(self): @@ -126,6 +147,8 @@ class CIFARLocalResNet: @property def n_forward_parameters(self): return (sum(weight.numel() for weight in self.W) + + sum(value.numel() for value in self.gamma) + + sum(value.numel() for value in self.beta) + self.W_out.numel() + self.b_out.numel()) @property @@ -170,7 +193,36 @@ class CIFARLocalResNet: f"does not match hidden value {tuple(value.shape)}") return value + perturbation - def forward(self, x, perturbations=None, return_cache=False): + def _normalize(self, index, value, training, update_stats): + if self.normalization == "none": + return value, None + axes = (0, 2, 3) + if training: + mean = value.mean(dim=axes) + variance = value.var(dim=axes, unbiased=False) + if update_stats: + with torch.no_grad(): + count = value.numel() // value.shape[1] + unbiased = variance * count / max(1, count - 1) + self.running_mean[index].lerp_(mean.detach(), self.bn_momentum) + self.running_var[index].lerp_(unbiased.detach(), self.bn_momentum) + else: + mean = self.running_mean[index] + variance = self.running_var[index] + inverse_std = torch.rsqrt(variance + self.bn_eps) + normalized = ((value - mean[None, :, None, None]) + * inverse_std[None, :, None, None]) + output = (self.gamma[index][None, :, None, None] * normalized + + self.beta[index][None, :, None, None]) + cache = { + "normalized": normalized, + "inverse_std": inverse_std, + "training": bool(training), + } + return output, cache + + def forward(self, x, perturbations=None, return_cache=False, training=False, + update_stats=False): if x.ndim != 4 or tuple(x.shape[1:]) != (3, 32, 32): raise ValueError(f"expected CIFAR NCHW input, got {tuple(x.shape)}") if perturbations is not None and len(perturbations) != self.n_hidden: @@ -181,10 +233,12 @@ class CIFARLocalResNet: pre = x u = F.conv2d(pre, self.W[0], stride=1, padding=1) - h_clean = F.relu(u) + normalized, norm_cache = self._normalize(0, u, training, update_stats) + h_clean = F.relu(normalized) hiddens.append(h_clean) if return_cache: - caches.append({"pre": pre, "gate": u > 0}) + caches.append({"pre": pre, "gate": normalized > 0, + "normalization": norm_cache}) h = self._inject(h_clean, perturbations, 0) for block in self.blocks: @@ -196,18 +250,24 @@ class CIFARLocalResNet: pre_first = h u_first = F.conv2d( pre_first, self.W[first], stride=block["stride"], padding=1) - first_clean = F.relu(u_first) + normalized_first, first_norm_cache = self._normalize( + first, u_first, training, update_stats) + first_clean = F.relu(normalized_first) hiddens.append(first_clean) if return_cache: - caches.append({"pre": pre_first, "gate": u_first > 0}) + caches.append({"pre": pre_first, "gate": normalized_first > 0, + "normalization": first_norm_cache}) first_value = self._inject(first_clean, perturbations, first) u_second = F.conv2d(first_value, self.W[second], stride=1, padding=1) - block_pre = shortcut + self.residual_scale * u_second + normalized_second, second_norm_cache = self._normalize( + second, u_second, training, update_stats) + block_pre = shortcut + self.residual_scale * normalized_second block_clean = F.relu(block_pre) hiddens.append(block_clean) if return_cache: - caches.append({"pre": first_value, "gate": block_pre > 0}) + caches.append({"pre": first_value, "gate": block_pre > 0, + "normalization": second_norm_cache}) h = self._inject(block_clean, perturbations, second) features = h.mean(dim=(2, 3)) @@ -222,6 +282,26 @@ class CIFARLocalResNet: def logits(self, x): return self.forward(x)["logits"] + def _normalization_backward(self, index, delta, cache): + """Local BatchNorm Jacobian-vector product and affine directions.""" + if self.normalization == "none": + return delta, None, None + normalized = cache["normalized"] + gamma_direction = (delta * normalized).sum(dim=(0, 2, 3)) + beta_direction = delta.sum(dim=(0, 2, 3)) + scaled = delta * self.gamma[index][None, :, None, None] + inverse_std = cache["inverse_std"][None, :, None, None] + if cache["training"]: + count = delta.shape[0] * delta.shape[2] * delta.shape[3] + summed = scaled.sum(dim=(0, 2, 3), keepdim=True) + projected = (scaled * normalized).sum( + dim=(0, 2, 3), keepdim=True) + input_delta = (inverse_std / count) * ( + count * scaled - summed - normalized * projected) + else: + input_delta = inverse_std * scaled + return input_delta, gamma_direction, beta_direction + def local_ascent_directions(self, teaching, output_error, forward): """Return forward-parameter descent directions from local signals. @@ -238,6 +318,8 @@ class CIFARLocalResNet: raise ValueError("local directions require a cached forward pass") batch = output_error.shape[0] directions = [] + gamma_directions = [] + beta_directions = [] with torch.no_grad(): for index, (signal, cache, spec, weight) in enumerate(zip( teaching, caches, self.layer_specs, self.W)): @@ -245,22 +327,34 @@ class CIFARLocalResNet: raise ValueError( f"teaching {index} has {tuple(signal.shape[1:])}, " f"expected {spec.hidden_shape}") - delta = (signal * cache["gate"].to(signal.dtype) - * spec.branch_scale) + post_norm_delta = (signal * cache["gate"].to(signal.dtype) + * spec.branch_scale) + delta, gamma_direction, beta_direction = self._normalization_backward( + index, post_norm_delta, cache["normalization"]) direction = torch.nn.grad.conv2d_weight( cache["pre"].detach(), weight.shape, delta.detach(), stride=spec.stride, padding=spec.padding) directions.append(direction / batch) + if gamma_direction is not None: + gamma_directions.append(gamma_direction / batch) + beta_directions.append(beta_direction / batch) output_weight = -(output_error.t() @ forward["features"].detach()) / batch output_bias = -output_error.mean(dim=0) - return directions, output_weight, output_bias + return (directions, gamma_directions, beta_directions, + output_weight, output_bias) def apply_ascent(self, directions, output_weight, output_bias, eta_hidden, - eta_output=None, momentum=0.0, weight_decay=0.0): + eta_output=None, momentum=0.0, weight_decay=0.0, + gamma_directions=None, beta_directions=None): """Apply simultaneously computed directions with optional momentum.""" if len(directions) != len(self.W): raise ValueError("one direction is required for every convolution") eta_output = eta_hidden if eta_output is None else eta_output + gamma_directions = [] if gamma_directions is None else gamma_directions + beta_directions = [] if beta_directions is None else beta_directions + if self.normalization == "batchnorm" and not ( + len(gamma_directions) == len(beta_directions) == len(self.W)): + raise ValueError("BatchNorm directions must cover every convolution") with torch.no_grad(): for index, (weight, direction) in enumerate(zip(self.W, directions)): update = direction - weight_decay * weight @@ -268,6 +362,15 @@ class CIFARLocalResNet: self.mW[index].mul_(momentum).add_(update) update = self.mW[index] weight.add_(update, alpha=eta_hidden) + for index, (gamma_direction, beta_direction) in enumerate(zip( + gamma_directions, beta_directions)): + if momentum: + self.mgamma[index].mul_(momentum).add_(gamma_direction) + self.mbeta[index].mul_(momentum).add_(beta_direction) + gamma_direction = self.mgamma[index] + beta_direction = self.mbeta[index] + self.gamma[index].add_(gamma_direction, alpha=eta_hidden) + self.beta[index].add_(beta_direction, alpha=eta_hidden) out_update = output_weight - weight_decay * self.W_out if momentum: self.mW_out.mul_(momentum).add_(out_update) @@ -279,18 +382,30 @@ class CIFARLocalResNet: def bp_step(self, x, y, eta, momentum=0.0, weight_decay=0.0): """Exact-backprop comparator on the identical forward architecture.""" - parameters = self.W + [self.W_out, self.b_out] + parameters = self.W + self.gamma + self.beta + [self.W_out, self.b_out] for parameter in parameters: parameter.requires_grad_(True) - loss = F.cross_entropy(self.logits(x), y) + loss = F.cross_entropy( + self.forward(x, training=True, update_stats=True)["logits"], y) gradients = torch.autograd.grad(loss, parameters) + n_conv = len(self.W) with torch.no_grad(): - conv_directions = [-gradient for gradient in gradients[:-2]] + conv_directions = [-gradient for gradient in gradients[:n_conv]] + if self.normalization == "batchnorm": + gamma_directions = [ + -gradient for gradient in gradients[n_conv:2 * n_conv]] + beta_directions = [ + -gradient for gradient in gradients[2 * n_conv:3 * n_conv]] + else: + gamma_directions = [] + beta_directions = [] output_weight = -gradients[-2] output_bias = -gradients[-1] self.apply_ascent( conv_directions, output_weight, output_bias, eta, - momentum=momentum, weight_decay=weight_decay) + momentum=momentum, weight_decay=weight_decay, + gamma_directions=gamma_directions, + beta_directions=beta_directions) for parameter in parameters: parameter.requires_grad_(False) return float(loss.detach()) @@ -425,8 +540,9 @@ def simultaneous_conv_node_perturbation(net, x, y, clean_forward, sigma=1e-2, 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. + is handled by the variance theorem in ``THEORY.md``. The antithetic trials + are evaluated as separate B-sized batches: concatenating them would couple + their BatchNorm statistics and change the intervention being estimated. """ if sigma <= 0: raise ValueError("perturbation sigma must be positive") @@ -439,28 +555,47 @@ def simultaneous_conv_node_perturbation(net, x, y, clean_forward, sigma=1e-2, 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 = [] + plus_perturbations = [] + minus_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) + plus_perturbations.append(sigma * direction) + minus_perturbations.append(-sigma * direction) + plus_forward = net.forward( + x, perturbations=plus_perturbations, + training=True, update_stats=False) + minus_forward = net.forward( + x, perturbations=minus_perturbations, + training=True, update_stats=False) + plus = F.cross_entropy(plus_forward["logits"], y, reduction="none") + minus = F.cross_entropy(minus_forward["logits"], y, reduction="none") + if net.normalization == "batchnorm": + # BN couples examples. The per-example loss difference is not a + # valid node-perturbation target because ell_i also responds to + # xi_j for j != i. The scalar batch objective is valid; multiplying + # its derivative by B recovers the derivative of the summed loss, + # matching the per-example signal convention of the local update. + batch_directional = (plus.mean() - minus.mean()) / (2.0 * sigma) + directional = batch_directional.mul(x.shape[0]).expand(x.shape[0]) + else: + batch_directional = None + 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) + diagnostic_derivatives.append({ + "scaled_directional": directional, + "batch_mean_directional": batch_directional, + "coupling": ("batch_objective" if batch_directional is not None + else "per_example_objective"), + }) if return_diagnostics: return targets, { "directions": diagnostic_directions, @@ -501,7 +636,8 @@ 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) + forward = net.forward( + x, return_cache=True, training=True, update_stats=True) logits = forward["logits"] loss = F.cross_entropy(logits, y) output_error = (torch.softmax(logits, dim=1) @@ -519,12 +655,15 @@ def conv_local_step(net, x, y, config, step, generator=None): 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( + (directions, gamma_directions, beta_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) + momentum=config.momentum, weight_decay=config.weight_decay, + gamma_directions=gamma_directions, + beta_directions=beta_directions) calibration = None if did_perturb and config.learn_A: calibration = net.calibrate_apical( @@ -554,7 +693,8 @@ def conv_apical_calibration_step(net, x, y, config, generator=None): config.validate() if not config.learn_A: raise ValueError("apical-only calibration requires learn_A=True") - forward = net.forward(x, return_cache=False) + forward = net.forward( + x, return_cache=False, training=True, update_stats=False) logits = forward["logits"] output_error = (torch.softmax(logits, dim=1) - F.one_hot(y, net.n_classes).to(logits.dtype)) @@ -571,7 +711,7 @@ def conv_apical_calibration_step(net, x, y, config, generator=None): def conv_alignment_report(net, x, y, config): """Measure apical alignment to exact hidden gradients; never used to learn.""" - parameters = net.W + [net.W_out, net.b_out] + parameters = net.W + net.gamma + net.beta + [net.W_out, net.b_out] for parameter in parameters: parameter.requires_grad_(True) forward = net.forward(x) |