diff options
| -rw-r--r-- | experiments/conv_local_smoke.py | 101 | ||||
| -rw-r--r-- | experiments/conv_run.py | 44 | ||||
| -rw-r--r-- | sdil/conv.py | 152 |
3 files changed, 273 insertions, 24 deletions
diff --git a/experiments/conv_local_smoke.py b/experiments/conv_local_smoke.py index 4e388a0..e8cf294 100644 --- a/experiments/conv_local_smoke.py +++ b/experiments/conv_local_smoke.py @@ -8,10 +8,11 @@ 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 (CIFARHierarchicalFAResNet, CIFARLocalResNet, - CIFARSDILResNet, ConvSDILConfig, +from sdil.conv import (CIFARHierarchicalFAResNet, CIFARKPResNet, + CIFARLocalResNet, CIFARSDILResNet, ConvSDILConfig, channel_subspace_apical_calibration, - conv_hierarchical_step, conv_local_step, + conv_hierarchical_step, conv_kolen_pollack_step, + conv_local_step, hierarchical_mirror_observations, hierarchical_parameter_subspace_calibration, normalized_residual_mirror_update, @@ -843,6 +844,99 @@ def normalized_response_mirror_checks(): } +def kolen_pollack_checks(): + """KP's reciprocal correlations are local and preserve exact symmetry.""" + torch.manual_seed(127) + common = dict( + depth=8, base_width=2, seed=53, dtype=torch.float64, + normalization="batchnorm", residual_scale=1.0) + net = CIFARKPResNet(**common) + for index in range(1, len(net.Q)): + net.Q[index].copy_(net.W[index]) + net.R_out.copy_(-net.W_out.t()) + x = torch.randn(3, 3, 32, 32, dtype=torch.float64) + y = torch.tensor([1, 4, 7]) + forward = net.forward( + x, return_cache=True, training=True, update_stats=False) + output_error = (torch.softmax(forward["logits"], dim=1) + - F.one_hot(y, 10).to(torch.float64)) + teaching = net.hierarchical_teaching(output_error, forward) + (forward_directions, gamma_directions, beta_directions, + output_weight, output_bias) = net.local_ascent_directions( + teaching, output_error, forward) + reciprocal_directions, reciprocal_readout = ( + net.reciprocal_feedback_directions( + teaching, output_error, forward)) + direction_error = max([ + float((left - right).abs().max()) + for left, right in zip(forward_directions[1:], reciprocal_directions[1:]) + ] + [float((reciprocal_readout + output_weight.t()).abs().max())]) + assert direction_error < 1e-14 + + # Once local activities have been observed, neither forward nor feedback + # parameter values may alter the independently formed reciprocal update. + before = [value.clone() for value in reciprocal_directions[1:]] + before_readout = reciprocal_readout.clone() + for value in net.W + net.Q: + value.add_(torch.randn_like(value)) + net.W_out.add_(torch.randn_like(net.W_out)) + net.R_out.add_(torch.randn_like(net.R_out)) + independent_directions, independent_readout = ( + net.reciprocal_feedback_directions( + teaching, output_error, forward)) + independence_error = max([ + float((left - right).abs().max()) + for left, right in zip(before, independent_directions[1:]) + ] + [float((before_readout - independent_readout).abs().max())]) + assert independence_error == 0.0 + + # A fresh symmetric state must remain symmetric under two momentum steps. + net = CIFARKPResNet(**common) + for index in range(1, len(net.Q)): + net.Q[index].copy_(net.W[index]) + net.R_out.copy_(-net.W_out.t()) + for _ in range(2): + forward = net.forward( + x, return_cache=True, training=True, update_stats=False) + output_error = (torch.softmax(forward["logits"], dim=1) + - F.one_hot(y, 10).to(torch.float64)) + teaching = net.hierarchical_teaching(output_error, forward) + (forward_directions, gamma_directions, beta_directions, + output_weight, output_bias) = net.local_ascent_directions( + teaching, output_error, forward) + reciprocal_directions, reciprocal_readout = ( + net.reciprocal_feedback_directions( + teaching, output_error, forward)) + net.apply_reciprocal_ascent( + reciprocal_directions, reciprocal_readout, + eta_hidden=0.013, eta_output=0.017, momentum=0.9, + weight_decay=1e-4) + net.apply_ascent( + forward_directions, output_weight, output_bias, + eta_hidden=0.013, eta_output=0.017, momentum=0.9, + weight_decay=1e-4, gamma_directions=gamma_directions, + beta_directions=beta_directions) + symmetry_error = max([ + float((net.Q[index] - net.W[index]).abs().max()) + for index in range(1, len(net.Q)) + ] + [float((net.R_out + net.W_out.t()).abs().max())]) + assert symmetry_error < 1e-14 + + # Exercise the public training step and ensure it remains graph-free. + result = conv_kolen_pollack_step( + net, x, y, ConvSDILConfig( + eta=1e-3, eta_output=1e-3, momentum=0.9, + weight_decay=1e-4, learn_A=False, learn_P=False)) + assert math.isfinite(result["loss"]) + assert all(not value.requires_grad for value in + net.W + net.Q + [net.W_out, net.R_out, net.b_out]) + return { + "kp_local_direction_absolute_error": direction_error, + "kp_forward_parameter_independence_error": independence_error, + "kp_symmetric_update_absolute_error": symmetry_error, + } + + def apical_learning_checks(): torch.manual_seed(11) net = CIFARSDILResNet(depth=8, base_width=2, seed=6) @@ -951,6 +1045,7 @@ def main(): report.update(hierarchical_feedback_checks()) report.update(hierarchical_parameter_calibration_checks()) report.update(normalized_response_mirror_checks()) + report.update(kolen_pollack_checks()) 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 3b91da9..738c86d 100644 --- a/experiments/conv_run.py +++ b/experiments/conv_run.py @@ -12,13 +12,14 @@ import time import torch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from sdil.conv import (CIFARHierarchicalFAResNet, CIFARLocalResNet, - CIFARSDILResNet, ConvSDILConfig, +from sdil.conv import (CIFARHierarchicalFAResNet, CIFARKPResNet, + CIFARLocalResNet, CIFARSDILResNet, ConvSDILConfig, conv_alignment_report, conv_apical_calibration_step, conv_hierarchical_alignment_report, - conv_hierarchical_step, + conv_hierarchical_step, conv_kolen_pollack_step, conv_learned_hierarchical_step, conv_local_step, - evaluate_conv, hierarchical_parameter_subspace_calibration) + evaluate_conv, hierarchical_feedback_tracking_report, + hierarchical_parameter_subspace_calibration) from sdil.conv import (normalized_residual_mirror_step, normalized_response_mirror_step) from sdil.data import DATA_DIR, get_cifar_image_splits @@ -107,8 +108,10 @@ def build(args): bn_momentum=args.bn_momentum, bn_eps=args.bn_eps) if args.mode == "bp": return CIFARLocalResNet(**common), None - if args.mode in ("hfa", "lhfa", "wm", "rrm"): - net = CIFARHierarchicalFAResNet( + if args.mode in ("hfa", "lhfa", "wm", "rrm", "kp"): + network_class = ( + CIFARKPResNet if args.mode == "kp" else CIFARHierarchicalFAResNet) + net = network_class( **common, feedback_seed=args.apical_seed, feedback_scale=args.a_scale) config = ConvSDILConfig( @@ -164,6 +167,8 @@ def work_report(net, mode, counters): + counters["mirror_readout_examples"] * mirror_readout_macs) mirror_prediction = mirror_forward if mode == "rrm" else 0 mirror_correlation = mirror_forward + kp_feedback_correlation = ( + counters["ordinary_examples"] * apical_macs if mode == "kp" else 0) components = { "ordinary_forward_macs": normal_forward, "warmup_clean_forward_macs": warmup_forward, @@ -175,6 +180,7 @@ def work_report(net, mode, counters): "mirror_response_macs": mirror_forward, "mirror_feedback_prediction_macs": mirror_prediction, "mirror_local_correlation_macs": mirror_correlation, + "kp_reciprocal_correlation_macs": kp_feedback_correlation, } return { "forward_macs_per_example": forward_macs, @@ -262,7 +268,8 @@ def run(args): "protocol_family": "oral_a_cifar_local_resnet_development", "calibration_metric_space": ( None if config is None or args.mode == "hfa" else - "hierarchical_feedback_parameters" if args.mode == "lhfa" else { + "hierarchical_feedback_parameters" if args.mode == "lhfa" else + "reciprocal_local_activity_products" if args.mode == "kp" else { "wm": "local_parent_child_response", "rrm": "local_parent_child_response_residual", "unit_targets": "full_hidden_field", @@ -295,7 +302,7 @@ def run(args): if args.mode == "hfa" else 0), "adaptive_feedback_parameters": ( getattr(net, "n_fixed_feedback_parameters", 0) - if args.mode in ("lhfa", "wm", "rrm") else 0), + if args.mode in ("lhfa", "wm", "rrm", "kp") else 0), }, "epochs": [], } @@ -465,6 +472,10 @@ def run(args): result = conv_hierarchical_step(net, x, y, config) loss = result["loss"] did_perturb = False + elif args.mode == "kp": + result = conv_kolen_pollack_step(net, x, y, config) + loss = result["loss"] + did_perturb = False else: result = conv_local_step( net, x, y, config, step, generator=perturb_generator) @@ -510,6 +521,18 @@ def run(args): key: sum(value[key] for value in mirror_metrics) / len(mirror_metrics) for key in mirror_metrics[0] } + if args.mode == "kp": + tracking = hierarchical_feedback_tracking_report(net) + record["feedback_tracking"] = { + "mean_feedback_forward_cosine": tracking[ + "mean_feedback_forward_cosine"], + "mean_feedback_forward_relative_error": tracking[ + "mean_feedback_forward_relative_error"], + "min_feedback_forward_cosine": min( + tracking["feedback_forward_cosine"]), + "max_feedback_forward_relative_error": max( + tracking["feedback_forward_relative_error"]), + } if args.eval_every and (epoch + 1) % args.eval_every == 0: sync(args.device) eval_start = time.time() @@ -548,7 +571,7 @@ def run(args): diagnostic_start = time.time() diagnostics = (conv_hierarchical_alignment_report( net, train.x[:probe], train.y[:probe]) - if args.mode in ("hfa", "lhfa", "wm", "rrm") else + if args.mode in ("hfa", "lhfa", "wm", "rrm", "kp") else conv_alignment_report( net, train.x[:probe], train.y[:probe], config)) sync(args.device) @@ -597,7 +620,8 @@ def parse_args(): parser = argparse.ArgumentParser() parser.add_argument( "--mode", choices=( - "bp", "dfa", "hfa", "lhfa", "wm", "rrm", "sdil", "nodepert"), + "bp", "dfa", "hfa", "lhfa", "wm", "rrm", "kp", "sdil", + "nodepert"), required=True) parser.add_argument("--out", required=True) parser.add_argument("--device", default="cpu") diff --git a/sdil/conv.py b/sdil/conv.py index b3d9235..8331433 100644 --- a/sdil/conv.py +++ b/sdil/conv.py @@ -529,6 +529,101 @@ class CIFARHierarchicalFAResNet(CIFARLocalResNet): return teaching +class CIFARKPResNet(CIFARHierarchicalFAResNet): + """Modified Kolen--Pollack reciprocal-feedback baseline. + + Forward and reciprocal synapses receive the same two locally available + activity factors and independently form equal-shaped correlation updates. + Matching optimizer and decay dynamics then suppress their initial + difference without reading or copying either synaptic weight. This is the + inherited Akrout et al. baseline, not an SDIL contribution. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.mQ = [torch.zeros_like(value) for value in self.Q] + self.mR_out = torch.zeros_like(self.R_out) + + @torch.no_grad() + def reciprocal_feedback_directions(self, teaching, output_error, forward): + """Recompute reciprocal updates from local activities only. + + Q is stored in forward-kernel orientation because ``conv_transpose2d`` + applies its transpose during feedback. Consequently its local KP + correlation has the same stored orientation as the corresponding W + correlation. R is stored as ``-W_out.T`` under this runner's teaching + sign convention. + """ + if len(teaching) != self.n_hidden: + raise ValueError("KP teaching must cover every hidden population") + caches = forward.get("caches") + if caches is None: + raise ValueError("KP feedback updates require local forward caches") + batch = output_error.shape[0] + directions = [None] + for index in range(1, len(self.Q)): + signal = teaching[index] + cache = caches[index] + spec = self.layer_specs[index] + post_norm_delta = (signal * cache["gate"].to(signal.dtype) + * spec.branch_scale) + delta, _, _ = self._normalization_backward( + index, post_norm_delta, cache["normalization"]) + direction = torch.nn.grad.conv2d_weight( + cache["pre"].detach(), self.Q[index].shape, delta.detach(), + stride=spec.stride, padding=spec.padding) + directions.append(direction / batch) + readout_direction = ( + forward["features"].detach().t() @ output_error) / batch + return directions, readout_direction + + @torch.no_grad() + def apply_reciprocal_ascent(self, directions, readout_direction, + eta_hidden, eta_output=None, momentum=0.0, + weight_decay=0.0): + """Apply the independently formed KP feedback-synapse updates.""" + if len(directions) != len(self.Q) or directions[0] is not None: + raise ValueError("KP directions must cover Q[1:] only") + eta_output = eta_hidden if eta_output is None else eta_output + for index in range(1, len(self.Q)): + update = directions[index] - weight_decay * self.Q[index] + if momentum: + self.mQ[index].mul_(momentum).add_(update) + update = self.mQ[index] + self.Q[index].add_(update, alpha=eta_hidden) + update = readout_direction - weight_decay * self.R_out + if momentum: + self.mR_out.mul_(momentum).add_(update) + update = self.mR_out + self.R_out.add_(update, alpha=eta_output) + + +@torch.no_grad() +def hierarchical_feedback_tracking_report(net): + """Cheap parameter-space tracking diagnostics; never used for learning.""" + if not isinstance(net, CIFARHierarchicalFAResNet): + raise TypeError("feedback tracking requires a hierarchical network") + pairs = list(zip(net.Q[1:], net.W[1:])) + [ + (net.R_out, -net.W_out.t())] + cosines = [float(F.cosine_similarity( + feedback.flatten(), target.flatten(), dim=0)) + for feedback, target in pairs] + norm_ratios = [float( + feedback.norm() / target.norm().clamp_min(1e-30)) + for feedback, target in pairs] + relative_errors = [float( + (feedback - target).norm() / target.norm().clamp_min(1e-30)) + for feedback, target in pairs] + return { + "feedback_forward_cosine": cosines, + "feedback_forward_norm_ratio": norm_ratios, + "feedback_forward_relative_error": relative_errors, + "mean_feedback_forward_cosine": sum(cosines) / len(cosines), + "mean_feedback_forward_relative_error": ( + sum(relative_errors) / len(relative_errors)), + } + + @torch.no_grad() def hierarchical_parameter_subspace_calibration( net, x, y, clean_forward, output_signal, sigma=1e-2, @@ -884,6 +979,49 @@ def conv_hierarchical_step(net, x, y, config): } +def conv_kolen_pollack_step(net, x, y, config): + """One modified-KP update with independently computed reciprocal changes.""" + if not isinstance(net, CIFARKPResNet): + raise TypeError("KP step requires CIFARKPResNet") + config.validate() + with torch.no_grad(): + 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) + - F.one_hot(y, net.n_classes).to(logits.dtype)) + teaching = net.hierarchical_teaching(output_error, forward) + total_units = sum(value.numel() for value in teaching) + teaching_rms = math.sqrt( + sum(float(value.square().sum()) for value in teaching) / total_units) + (directions, gamma_directions, beta_directions, + output_weight, output_bias) = net.local_ascent_directions( + teaching, output_error, forward) + reciprocal_directions, reciprocal_readout = ( + net.reciprocal_feedback_directions( + teaching, output_error, forward)) + # Both parameter sets consume their own local correlation calculation. + # Apply only after every direction has been formed, preserving a + # simultaneous update with no weight or weight-change read across paths. + net.apply_reciprocal_ascent( + reciprocal_directions, reciprocal_readout, + eta_hidden=config.eta, eta_output=config.eta_output, + momentum=config.momentum, weight_decay=config.weight_decay) + 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, + gamma_directions=gamma_directions, + beta_directions=beta_directions) + return { + "loss": float(loss), "did_perturb": False, "calibration": None, + "predictor_mse": None, "teaching_rms": teaching_rms, + "raw_apical_rms": teaching_rms, + "innovation_rms": teaching_rms, + } + + def conv_learned_hierarchical_step(net, x, y, config, step, generator=None): """One task update with optional causal calibration of hierarchical Q/R.""" config.validate() @@ -941,22 +1079,14 @@ def conv_hierarchical_alignment_report(net, x, y): for left, right in zip(teaching, negative_gradients)] for parameter in parameters: parameter.requires_grad_(False) - feedback_pairs = list(zip(net.Q[1:], net.W[1:])) + [ - (net.R_out, -net.W_out.t())] - feedback_forward_cosine = [float(F.cosine_similarity( - feedback.flatten(), target.flatten(), dim=0)) - for feedback, target in feedback_pairs] - feedback_forward_norm_ratio = [ - float(feedback.norm() / target.norm().clamp_min(1e-30)) - for feedback, target in feedback_pairs] - return { + report = { "normalization_state": "training_batch_stats_without_running_update", "teaching_negative_gradient_cosine": values, "raw_negative_gradient_cosine": values, "innovation_negative_gradient_cosine": values, - "feedback_forward_cosine": feedback_forward_cosine, - "feedback_forward_norm_ratio": feedback_forward_norm_ratio, } + report.update(hierarchical_feedback_tracking_report(net)) + return report class CIFARSDILResNet(CIFARLocalResNet): |