diff options
Diffstat (limited to 'sdil')
| -rw-r--r-- | sdil/conv.py | 89 |
1 files changed, 89 insertions, 0 deletions
diff --git a/sdil/conv.py b/sdil/conv.py index 731cbab..404bca6 100644 --- a/sdil/conv.py +++ b/sdil/conv.py @@ -1283,6 +1283,95 @@ def layerwise_causal_feedback_update(net, observation, eta=0.1, eps=1e-12): @torch.no_grad() +def causal_readout_least_squares_fit( + net, observations, relative_ridge=1e-6): + """Fit readout feedback from stored local causal observations.""" + if not isinstance(net, CIFARHierarchicalFAResNet): + raise TypeError("causal readout fit requires a hierarchical net") + if not observations or relative_ridge <= 0: + raise ValueError("invalid causal readout fit inputs") + if any(value.get("kind") != "readout" for value in observations): + raise ValueError("readout fit received a convolutional observation") + features = torch.cat([value["features"] for value in observations], dim=0) + targets = torch.cat([value["target"] for value in observations], dim=0) + gram = features.t() @ features + ridge = relative_ridge * float(torch.trace(gram)) / gram.shape[0] + ridge = max(ridge, torch.finfo(gram.dtype).tiny) + regularized = gram + ridge * torch.eye( + gram.shape[0], device=gram.device, dtype=gram.dtype) + solution = torch.linalg.solve(regularized, features.t() @ targets) + fitted = features @ solution + before = features @ net.R_out.t() + net.R_out.copy_(solution.t()) + target_power = float(targets.square().sum()) + fitted_power = float(fitted.square().sum()) + denominator = math.sqrt(target_power * fitted_power) + return { + "examples": int(features.shape[0]), + "relative_ridge": float(relative_ridge), + "absolute_ridge": ridge, + "before_mse": float((targets - before).square().mean()), + "after_mse": float((targets - fitted).square().mean()), + "fit_target_cosine": ( + float((targets * fitted).sum()) / denominator + if denominator else 0.0), + "parameter_rms": math.sqrt(float(net.R_out.square().mean())), + } + + +@torch.no_grad() +def causal_conv_diagonal_least_squares_fit( + net, observations, relative_ridge=1e-3): + """Apply one diagonal-whitened local fit to a selected feedback edge.""" + if not isinstance(net, CIFARHierarchicalFAResNet): + raise TypeError("causal convolutional fit requires a hierarchical net") + if not observations or relative_ridge <= 0: + raise ValueError("invalid causal convolutional fit inputs") + if any(value.get("kind") != "convolution" for value in observations): + raise ValueError("convolutional fit received a readout observation") + indices = {int(value["edge_index"]) for value in observations} + if len(indices) != 1: + raise ValueError("one causal fit must contain exactly one feedback edge") + index = indices.pop() + numerator = torch.zeros_like(net.Q[index]) + diagonal = torch.zeros_like(net.Q[index]) + before_error = 0.0 + target_power = 0.0 + units = 0 + for value in observations: + error = value["target"] - value["prediction"] + context = value["context"] + stride = int(value["stride"]) + padding = int(value["padding"]) + numerator.add_(torch.nn.grad.conv2d_weight( + error, net.Q[index].shape, context, + stride=stride, padding=padding)) + diagonal.add_(torch.nn.grad.conv2d_weight( + torch.ones_like(error), net.Q[index].shape, context.square(), + stride=stride, padding=padding)) + before_error += float(error.square().sum()) + target_power += float(value["target"].square().sum()) + units += error.numel() + ridge = relative_ridge * float(diagonal.mean()) + ridge = max(ridge, torch.finfo(diagonal.dtype).tiny) + update = numerator / (diagonal + ridge) + net.Q[index].add_(update) + return { + "edge": index, + "observations": len(observations), + "examples": sum(value["target"].shape[0] for value in observations), + "relative_ridge": float(relative_ridge), + "absolute_ridge": ridge, + "before_field_mse": before_error / units, + "target_rms": math.sqrt(target_power / units), + "parameter_update_rms": math.sqrt(float(update.square().mean())), + "parameter_rms": math.sqrt(float(net.Q[index].square().mean())), + "minimum_diagonal": float(diagonal.min()), + "mean_diagonal": float(diagonal.mean()), + } + + +@torch.no_grad() def layerwise_causal_bootstrap_sweep( net, x, y, sigma=1e-2, eta=0.1, generator=None): """Calibrate readout then convolutional feedback in reverse DAG order.""" |