diff options
Diffstat (limited to 'sdil')
| -rw-r--r-- | sdil/conv.py | 152 |
1 files changed, 141 insertions, 11 deletions
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): |