From 042c43f36d7ebac22f4b3ff078b1de42445de0b3 Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Thu, 6 Aug 2026 14:31:22 -0500 Subject: experiment: implement layerwise causal feedback bootstrap --- sdil/conv.py | 187 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 187 insertions(+) (limited to 'sdil/conv.py') diff --git a/sdil/conv.py b/sdil/conv.py index 1034202..731cbab 100644 --- a/sdil/conv.py +++ b/sdil/conv.py @@ -1136,6 +1136,193 @@ def hierarchical_parameter_subspace_calibration( return calibration +@torch.no_grad() +def layerwise_causal_feedback_observation( + net, x, y, clean_forward, output_signal, edge_index=None, + sigma=1e-2, generator=None): + """Observe one hidden-credit target using per-example task perturbations. + + ``edge_index=None`` denotes the dense readout feedback. Otherwise only + the parent population receiving ``Q[edge_index]`` is perturbed. Causal + queries use evaluation-mode normalization so each example's loss + difference depends only on its own perturbation. The returned tensors are + a complete, detached local observation: the subsequent update need not and + does not read a forward parameter. + """ + if not isinstance(net, CIFARHierarchicalFAResNet): + raise TypeError("layerwise causal calibration requires a hierarchical net") + if sigma <= 0: + raise ValueError("causal perturbation scale must be positive") + if generator is None: + generator = torch.Generator(device=x.device).manual_seed(0) + if len(clean_forward.get("hiddens", [])) != net.n_hidden: + raise ValueError("clean forward does not match hierarchical populations") + caches = clean_forward.get("caches") + if caches is None: + raise ValueError("layerwise causal calibration requires cached states") + if net.normalization == "batchnorm" and any( + cache["normalization"]["training"] for cache in caches): + raise ValueError( + "per-example causal observations require evaluation-mode BatchNorm") + + teaching, contexts, recipients = net.hierarchical_teaching( + output_signal, clean_forward, return_edge_contexts=True) + perturbations = [torch.zeros_like(value) + for value in clean_forward["hiddens"]] + batch = x.shape[0] + if edge_index is None: + recipient = net.n_hidden - 1 + hidden = clean_forward["hiddens"][recipient] + random = torch.empty( + batch, hidden.shape[1], device=hidden.device, + dtype=hidden.dtype).bernoulli_( + 0.5, generator=generator).mul_(2).sub_(1) + direction = random[:, :, None, None].expand_as(hidden) + else: + if not 1 <= edge_index < len(net.Q): + raise ValueError("feedback edge index is out of range") + recipient = recipients[edge_index] + hidden = clean_forward["hiddens"][recipient] + random = torch.empty_like(hidden).bernoulli_( + 0.5, generator=generator).mul_(2).sub_(1) + direction = random + perturbations[recipient] = direction.mul(sigma) + plus = F.cross_entropy(net.forward( + x, perturbations=perturbations, training=False, + update_stats=False)["logits"], y, reduction="none") + perturbations[recipient] = direction.mul(-sigma) + minus = F.cross_entropy(net.forward( + x, perturbations=perturbations, training=False, + update_stats=False)["logits"], y, reduction="none") + directional = (plus - minus) / (2.0 * sigma) + + if edge_index is None: + # A channel-constant hidden intervention is conjugate to the pooled + # feature vector. The target therefore has the coefficient scale of + # output_signal @ R.T, before division by the number of spatial sites. + target = -directional[:, None] * random + prediction = output_signal @ net.R_out.t() + return { + "kind": "readout", "edge_index": None, + "recipient": recipient, + "target": target.clone(), "prediction": prediction.clone(), + "features": output_signal.clone(), + "direction": direction.clone(), + "directional": directional.clone(), + } + + expand = directional.reshape( + batch, *([1] * (direction.ndim - 1))) + target = -expand * direction + return { + "kind": "convolution", "edge_index": int(edge_index), + "recipient": int(recipient), "target": target.clone(), + "prediction": teaching[recipient].clone(), + "context": contexts[edge_index].clone(), + "direction": direction.clone(), + "stride": int(net.layer_specs[edge_index].stride), + "padding": int(net.layer_specs[edge_index].padding), + "directional": directional.clone(), + } + + +@torch.no_grad() +def layerwise_causal_feedback_update(net, observation, eta=0.1, eps=1e-12): + """Apply one normalized local feedback update from a stored observation.""" + if not isinstance(net, CIFARHierarchicalFAResNet): + raise TypeError("layerwise causal calibration requires a hierarchical net") + if not 0.0 < eta <= 1.0 or eps <= 0: + raise ValueError("invalid layerwise causal update hyperparameters") + target = observation["target"] + prediction = observation["prediction"] + error = target - prediction + kind = observation.get("kind") + if kind == "readout": + features = observation["features"] + normalizer = float(features.square().sum() / features.shape[0]) + normalizer = max(normalizer, eps) + update = (error.t() @ features) / (features.shape[0] * normalizer) + if update.shape != net.R_out.shape: + raise ValueError("readout causal update shape mismatch") + net.R_out.add_(update, alpha=eta) + parameters = update.numel() + edge = "readout" + elif kind == "convolution": + index = int(observation["edge_index"]) + context = observation["context"] + batch = error.shape[0] + recipient_spatial = error.shape[2] * error.shape[3] + normalizer = float(context.square().sum() / (batch * recipient_spatial)) + normalizer = max(normalizer, eps) + update = torch.nn.grad.conv2d_weight( + error, net.Q[index].shape, context, + stride=int(observation["stride"]), + padding=int(observation["padding"])) + update.div_(batch * recipient_spatial * normalizer) + net.Q[index].add_(update, alpha=eta) + parameters = update.numel() + edge = index + else: + raise ValueError("unknown layerwise causal observation kind") + + target_power = float(target.square().sum()) + prediction_power = float(prediction.square().sum()) + denominator = math.sqrt(target_power * prediction_power) + return { + "edge": edge, + "field_prediction_target_cosine": ( + float((target * prediction).sum()) / denominator + if denominator else 0.0), + "target_rms": math.sqrt(target_power / target.numel()), + "prediction_rms": math.sqrt(prediction_power / prediction.numel()), + "field_error_rms": math.sqrt(float(error.square().mean())), + "parameter_update_rms": math.sqrt( + float(update.square().sum()) / parameters), + "local_context_power": normalizer, + } + + +@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.""" + if generator is None: + generator = torch.Generator(device=x.device).manual_seed(0) + clean = net.forward( + x, return_cache=True, training=False, update_stats=False) + output_signal = (torch.softmax(clean["logits"], dim=1) + - F.one_hot(y, net.n_classes).to(clean["logits"].dtype)) + metrics = [] + observation = layerwise_causal_feedback_observation( + net, x, y, clean, output_signal, edge_index=None, + sigma=sigma, generator=generator) + metrics.append(layerwise_causal_feedback_update( + net, observation, eta=eta)) + for index in reversed(range(1, len(net.Q))): + # The forward cache is unchanged by feedback learning. The teaching + # pass inside the observation is recomputed after every Q update, so + # the newly calibrated child field bootstraps its immediate parent. + observation = layerwise_causal_feedback_observation( + net, x, y, clean, output_signal, edge_index=index, + sigma=sigma, generator=generator) + metrics.append(layerwise_causal_feedback_update( + net, observation, eta=eta)) + return { + "events": len(metrics), + "logical_batch_loss_queries": 2 * len(metrics), + "per_example_causal_observations": x.shape[0] * len(metrics), + "mean_field_prediction_target_cosine": sum( + value["field_prediction_target_cosine"] for value in metrics) + / len(metrics), + "mean_parameter_update_rms": sum( + value["parameter_update_rms"] for value in metrics) + / len(metrics), + "max_parameter_update_rms": max( + value["parameter_update_rms"] for value in metrics), + "edges": metrics, + } + + @torch.no_grad() def hierarchical_mirror_observations(net, batch_size=1, noise_std=1.0, generator=None): -- cgit v1.2.3