From a4858a04acc5adca65a437295b63b426b6358831 Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Wed, 22 Jul 2026 14:46:41 -0500 Subject: experiment: implement KP mixed-traffic innovation controls --- experiments/analyze_kp_innovation_short.py | 8 + experiments/conv_local_smoke.py | 129 ++++++++++++- experiments/conv_run.py | 174 ++++++++++++++--- sdil/conv.py | 296 +++++++++++++++++++++++++++++ 4 files changed, 579 insertions(+), 28 deletions(-) diff --git a/experiments/analyze_kp_innovation_short.py b/experiments/analyze_kp_innovation_short.py index 4c70575..5f21dea 100644 --- a/experiments/analyze_kp_innovation_short.py +++ b/experiments/analyze_kp_innovation_short.py @@ -82,12 +82,20 @@ def main(): if len(tracking[rule]) != 20 or any(value is None for value in tracking[rule]): raise ValueError(f"MT-1 {rule} tracking trajectory is incomplete") for row, values in zip(records[rule]["epochs"], tracking[rule]): + mixed = row.get("mixed_apical") + if mixed is None: + raise ValueError(f"MT-1 {rule} mixed-apical trajectory is incomplete") trajectory_values.extend([ float(row["train_loss"]), float(values["mean_feedback_forward_cosine"]), float(values["mean_feedback_forward_relative_error"]), float(values["min_feedback_forward_cosine"]), float(values["max_feedback_forward_relative_error"]), + float(mixed["teaching_rms"]), + float(mixed["instruction_rms"]), + float(mixed["raw_apical_rms"]), + float(mixed["innovation_rms"]), + float(mixed["traffic_rms"]), ]) all_finite = all(record["final"]["finite"] for record in records.values()) diff --git a/experiments/conv_local_smoke.py b/experiments/conv_local_smoke.py index e8cf294..887674e 100644 --- a/experiments/conv_local_smoke.py +++ b/experiments/conv_local_smoke.py @@ -8,10 +8,12 @@ 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, CIFARKPResNet, - CIFARLocalResNet, CIFARSDILResNet, ConvSDILConfig, +from sdil.conv import (CIFARHierarchicalFAResNet, CIFARKPMixedTrafficResNet, + CIFARKPResNet, CIFARLocalResNet, CIFARSDILResNet, + ConvSDILConfig, channel_subspace_apical_calibration, conv_hierarchical_step, conv_kolen_pollack_step, + conv_kp_mixed_traffic_step, conv_local_step, hierarchical_mirror_observations, hierarchical_parameter_subspace_calibration, @@ -937,6 +939,128 @@ def kolen_pollack_checks(): } +def kp_mixed_traffic_checks(): + """Mixed traffic isolates subtraction from norm and preserves KP locality.""" + torch.manual_seed(211) + common = dict( + depth=8, base_width=2, seed=67, dtype=torch.float64, + normalization="batchnorm", residual_scale=1.0, + traffic_seed=4000) + net = CIFARKPMixedTrafficResNet(**common) + x = torch.randn(5, 3, 32, 32, dtype=torch.float64) + y = torch.tensor([0, 2, 4, 6, 8]) + 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)) + instruction = net.hierarchical_teaching(output_error, forward) + + # Zero traffic and zero predictor collapse all three rules to clean KP. + zero_errors = [] + for rule in ("raw", "matched", "innovation"): + components = net.mixed_apical_components( + instruction, forward["hiddens"], rule) + zero_errors.extend(float((left - right).abs().max()) for left, right in + zip(components["used"], instruction)) + assert max(zero_errors) == 0.0 + + calibration = net.calibrate_traffic_gain( + instruction, forward["hiddens"], target_ratio=4.0) + ratio_error = max(abs(value - 4.0) for value in + calibration["realized_traffic_instruction_rms_ratio"]) + assert ratio_error < 1e-12 + + # An exact per-unit predictor removes all predictable traffic. + for slope, gain, coefficient in zip( + net.P_traffic, net.traffic_gain, net.B_traffic): + slope.copy_(gain * coefficient) + exact = net.mixed_apical_components( + instruction, forward["hiddens"], "innovation") + exact_predictor_error = max(float((left - right).abs().max()) + for left, right in zip( + exact["innovation"], instruction)) + assert exact_predictor_error < 1e-14 + + for slope, bias in zip(net.P_traffic, net.P_traffic_bias): + slope.zero_() + bias.zero_() + components = net.mixed_apical_components( + instruction, forward["hiddens"], "matched") + norm_errors = [] + direction_errors = [] + for raw, innovation, matched in zip( + components["raw"], components["innovation"], + components["matched"]): + raw_flat = raw.flatten(1) + innovation_flat = innovation.flatten(1) + matched_flat = matched.flatten(1) + norm_errors.append(float(( + (matched_flat.norm(dim=1) - innovation_flat.norm(dim=1)).abs() + / innovation_flat.norm(dim=1).clamp_min(1e-30)).max())) + direction_errors.append(float((F.cosine_similarity( + raw_flat, matched_flat, dim=1) - 1.0).abs().max())) + assert max(norm_errors) < 1e-12 + assert max(direction_errors) < 1e-12 + + # All used signals retain equal independently recomputed local KP products. + correlation_errors = [] + for rule in ("raw", "matched", "innovation"): + used = net.mixed_apical_components( + instruction, forward["hiddens"], rule)["used"] + forward_directions, _, _, output_weight, _ = ( + net.local_ascent_directions(used, output_error, forward)) + reciprocal, reciprocal_readout = net.reciprocal_feedback_directions( + used, output_error, forward) + correlation_errors.extend(float((left - right).abs().max()) + for left, right in zip( + forward_directions[1:], reciprocal[1:])) + correlation_errors.append(float( + (reciprocal_readout + output_weight.t()).abs().max())) + assert max(correlation_errors) < 1e-14 + + # Predictor plasticity consumes only the supplied soma/traffic pair: once + # those are fixed, changing every forward/feedback weight has no effect. + left = CIFARKPMixedTrafficResNet(**common) + right = CIFARKPMixedTrafficResNet(**common) + for left_gain, right_gain, source in zip( + left.traffic_gain, right.traffic_gain, net.traffic_gain): + left_gain.copy_(source) + right_gain.copy_(source) + fixed_hiddens = [value.detach().clone() for value in forward["hiddens"]] + for value in right.W + right.Q + [right.W_out, right.R_out]: + value.add_(torch.randn_like(value)) + left.predictor_step(fixed_hiddens, eta=0.1) + right.predictor_step(fixed_hiddens, eta=0.1) + predictor_independence_error = max(float((a - b).abs().max()) + for a, b in zip( + left.P_traffic + left.P_traffic_bias, + right.P_traffic + right.P_traffic_bias)) + assert predictor_independence_error == 0.0 + + net.traffic_rule = "innovation" + result = conv_kp_mixed_traffic_step( + net, x, y, ConvSDILConfig( + eta=1e-4, eta_output=1e-4, eta_P=0.1, momentum=0.0, + weight_decay=0.0, learn_A=False, learn_P=True), + step=0, rule="innovation", predictor_every=16) + assert math.isfinite(result["loss"]) and result["did_predictor_update"] + assert all(not value.requires_grad for value in + net.W + net.Q + net.P_traffic + net.P_traffic_bias + + [net.W_out, net.R_out, net.b_out]) + assert net.mixed_elementwise_ops_per_example("matched") > ( + net.mixed_elementwise_ops_per_example("raw")) + return { + "kp_traffic_zero_limit_error": max(zero_errors), + "kp_traffic_ratio_error": ratio_error, + "kp_traffic_exact_predictor_error": exact_predictor_error, + "kp_traffic_matched_norm_error": max(norm_errors), + "kp_traffic_matched_direction_error": max(direction_errors), + "kp_traffic_reciprocal_correlation_error": max(correlation_errors), + "kp_traffic_predictor_parameter_independence_error": ( + predictor_independence_error), + } + + def apical_learning_checks(): torch.manual_seed(11) net = CIFARSDILResNet(depth=8, base_width=2, seed=6) @@ -1046,6 +1170,7 @@ def main(): report.update(hierarchical_parameter_calibration_checks()) report.update(normalized_response_mirror_checks()) report.update(kolen_pollack_checks()) + report.update(kp_mixed_traffic_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 738c86d..d2c0653 100644 --- a/experiments/conv_run.py +++ b/experiments/conv_run.py @@ -12,11 +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, CIFARKPResNet, - CIFARLocalResNet, CIFARSDILResNet, ConvSDILConfig, +from sdil.conv import (CIFARHierarchicalFAResNet, CIFARKPMixedTrafficResNet, + CIFARKPResNet, CIFARLocalResNet, CIFARSDILResNet, + ConvSDILConfig, conv_alignment_report, conv_apical_calibration_step, conv_hierarchical_alignment_report, conv_hierarchical_step, conv_kolen_pollack_step, + conv_kp_mixed_traffic_alignment_report, + conv_kp_mixed_traffic_step, conv_learned_hierarchical_step, conv_local_step, evaluate_conv, hierarchical_feedback_tracking_report, hierarchical_parameter_subspace_calibration) @@ -108,17 +111,25 @@ 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", "kp"): - network_class = ( - CIFARKPResNet if args.mode == "kp" else CIFARHierarchicalFAResNet) - net = network_class( - **common, feedback_seed=args.apical_seed, - feedback_scale=args.a_scale) + if args.mode in ("hfa", "lhfa", "wm", "rrm", "kp", "kp_traffic"): + if args.mode == "kp_traffic": + net = CIFARKPMixedTrafficResNet( + **common, feedback_seed=args.apical_seed, + feedback_scale=args.a_scale, traffic_seed=args.traffic_seed) + net.traffic_rule = args.traffic_rule + else: + 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( eta=args.lr, eta_output=args.output_lr, eta_A=(args.eta_A if args.mode == "lhfa" else 0.0), + eta_P=args.eta_P, momentum=args.momentum, weight_decay=args.weight_decay, - learn_A=args.mode == "lhfa", learn_P=False, + learn_A=args.mode == "lhfa", learn_P=args.mode == "kp_traffic", pert_sigma=args.pert_sigma, pert_every=args.pert_every, pert_directions=args.pert_directions, apical_calibration_mode="hierarchical_parameter_subspace") @@ -145,7 +156,9 @@ def work_report(net, mode, counters): apical_macs = getattr(net, "apical_macs_per_example", 0) normal_forward = counters["ordinary_examples"] * forward_macs warmup_forward = ((counters["predictor_warmup_examples"] - + counters["apical_warmup_examples"]) * forward_macs) + + counters["apical_warmup_examples"] + + counters["traffic_calibration_examples"] + + counters["traffic_audit_examples"]) * forward_macs) calibration_forward = counters["perturbation_forward_examples"] * forward_macs if mode == "bp": bp_reverse = 2 * counters["ordinary_examples"] * forward_macs @@ -156,7 +169,9 @@ def work_report(net, mode, counters): bp_reverse = 0 local_correlation = counters["ordinary_examples"] * forward_macs apical_inference = ((counters["ordinary_examples"] - + counters["apical_warmup_examples"]) * apical_macs) + + counters["apical_warmup_examples"] + + counters["traffic_calibration_examples"]) + * apical_macs) regression_multiplier = 2 if mode == "lhfa" else 1 apical_regression = (regression_multiplier * counters["calibration_event_examples"] @@ -168,7 +183,17 @@ def work_report(net, mode, counters): 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) + counters["ordinary_examples"] * apical_macs + if mode in ("kp", "kp_traffic") else 0) + if mode == "kp_traffic": + elementwise_operations = ( + counters["ordinary_examples"] + * net.mixed_elementwise_ops_per_example() + + (counters["predictor_warmup_examples"] + + counters["predictor_update_examples"]) + * net.predictor_elementwise_ops_per_example) + else: + elementwise_operations = 0 components = { "ordinary_forward_macs": normal_forward, "warmup_clean_forward_macs": warmup_forward, @@ -187,12 +212,17 @@ def work_report(net, mode, counters): "apical_macs_per_example": apical_macs, "components": components, "total_macs_estimate": sum(components.values()), + "elementwise_operations_estimate": elementwise_operations, "total_clean_forward_examples": ( counters["ordinary_examples"] + counters["predictor_warmup_examples"] - + counters["apical_warmup_examples"]), + + counters["apical_warmup_examples"] + + counters["traffic_calibration_examples"] + + counters["traffic_audit_examples"]), "total_forward_equivalent_examples": ( counters["ordinary_examples"] + counters["predictor_warmup_examples"] + counters["apical_warmup_examples"] + + counters["traffic_calibration_examples"] + + counters["traffic_audit_examples"] + counters["perturbation_forward_examples"]), "logical_batch_loss_queries": counters["logical_batch_loss_queries"], "causal_scalar_observations": counters["causal_scalar_observations"], @@ -203,18 +233,28 @@ def work_report(net, mode, counters): "multiply-accumulates in conv/linear maps; one local weight correlation " "equals one forward-weight MAC count; BP reverse is estimated as one " "weight-gradient plus one activation-gradient convolution per forward " - "convolution; elementwise nonlinearities and optimizer arithmetic excluded"), + "convolution; mixed-traffic/predictor elementwise arithmetic is reported " + "separately and is not folded into MACs"), } def run(args): if args.eval_split == "test" and args.eval_every: raise ValueError("test protocols must use --eval_every 0 (one final evaluation)") - if args.mode not in ("sdil", "lhfa") and ( + if args.mode not in ("sdil", "lhfa", "kp_traffic") and ( args.a_warmup_steps or args.learn_P): raise ValueError("apical/predictor warmup is restricted to SDIL") if args.mode == "lhfa" and args.learn_P: raise ValueError("predictor learning is not defined for learned HFA") + if args.mode == "kp_traffic": + if not args.learn_P or args.predictor_warmup_steps < 1: + raise ValueError("mixed-traffic KP requires predictor warmup") + if args.predictor_every < 1: + raise ValueError("mixed-traffic KP requires a predictor cadence") + if args.traffic_calibration_examples < 1 or args.traffic_ratio <= 0: + raise ValueError("invalid mixed-traffic calibration") + elif args.predictor_every: + raise ValueError("predictor cadence is restricted to mixed-traffic KP") if args.mode not in ("wm", "rrm") and args.mirror_warmup_steps: raise ValueError("mirror warmup is restricted to weight mirror mode") if args.mirror_every < 1 or args.mirror_batch_size < 1: @@ -262,12 +302,17 @@ def run(args): "mirror_conv_examples": 0, "mirror_readout_examples": 0, "mirror_events": 0, + "traffic_calibration_examples": 0, + "traffic_audit_examples": 0, + "predictor_update_examples": 0, } log = { "schema_version": 1, "protocol_family": "oral_a_cifar_local_resnet_development", "calibration_metric_space": ( None if config is None or args.mode == "hfa" else + "reciprocal_local_activity_products_with_mixed_apical_traffic" + if args.mode == "kp_traffic" else "hierarchical_feedback_parameters" if args.mode == "lhfa" else "reciprocal_local_activity_products" if args.mode == "kp" else { "wm": "local_parent_child_response", @@ -302,7 +347,8 @@ 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", "kp") else 0), + if args.mode in ("lhfa", "wm", "rrm", "kp", "kp_traffic") + else 0), }, "epochs": [], } @@ -313,6 +359,33 @@ def run(args): apical_warmup_wall = 0.0 mirror_warmup_wall = 0.0 loader_state = train.g.get_state().clone() + traffic_calibration_batch = None + if args.mode == "kp_traffic": + count = min(args.traffic_calibration_examples, train.x.shape[0]) + if count != args.traffic_calibration_examples: + raise ValueError("traffic calibration prefix is unavailable") + calibration_x = train.x[:count] + calibration_y = train.y[:count] + traffic_calibration_batch = (calibration_x, calibration_y) + calibration_forward = net.forward( + calibration_x, return_cache=True, training=True, + update_stats=False) + calibration_output_error = ( + torch.softmax(calibration_forward["logits"], dim=1) + - torch.nn.functional.one_hot( + calibration_y, net.n_classes).to( + calibration_forward["logits"].dtype)) + calibration_instruction = net.hierarchical_teaching( + calibration_output_error, calibration_forward) + log["traffic_calibration"] = net.calibrate_traffic_gain( + calibration_instruction, calibration_forward["hiddens"], + args.traffic_ratio) + log["traffic_calibration"].update({ + "examples": count, + "data_source": "first unaugmented development-training examples", + "uses_validation_endpoint": False, + }) + counters["traffic_calibration_examples"] += count if args.mode in ("wm", "rrm") and args.mirror_warmup_steps: sync(args.device) mirror_start = time.time() @@ -342,6 +415,7 @@ def run(args): sync(args.device) warmup_start = time.time() iterator = iter(train) + predictor_metrics = [] for _ in range(args.predictor_warmup_steps): try: x, _ = next(iterator) @@ -349,10 +423,30 @@ def run(args): iterator = iter(train) x, _ = next(iterator) forward = net.forward(x, training=True, update_stats=False) - net.predictor_step( - forward["hiddens"], config.eta_P, config.nuisance_scale) + predictor_metric = (net.predictor_step( + forward["hiddens"], config.eta_P) + if args.mode == "kp_traffic" else net.predictor_step( + forward["hiddens"], config.eta_P, + config.nuisance_scale)) + predictor_metrics.append(predictor_metric) counters["predictor_warmup_examples"] += x.shape[0] train.g.set_state(loader_state) + log["predictor_warmup"] = { + "steps": args.predictor_warmup_steps, + "first_mse": predictor_metrics[0], + "mean_mse": sum(predictor_metrics) / len(predictor_metrics), + "last_mse": predictor_metrics[-1], + "instruction_present": False, + } + if args.mode == "kp_traffic": + audit_x, _ = traffic_calibration_batch + audit_forward = net.forward( + audit_x, training=True, update_stats=False) + log["predictor_warmup"][ + "post_warmup_traffic_residual_rms_ratio"] = ( + net.predictor_traffic_residual_rms_ratio( + audit_forward["hiddens"])) + counters["traffic_audit_examples"] += audit_x.shape[0] sync(args.device) predictor_warmup_wall = time.time() - warmup_start @@ -426,6 +520,7 @@ def run(args): examples = 0 calibration_metrics = [] mirror_metrics = [] + signal_metrics = [] for x, y in train: batch = x.shape[0] if args.mode == "bp": @@ -476,6 +571,17 @@ def run(args): result = conv_kolen_pollack_step(net, x, y, config) loss = result["loss"] did_perturb = False + elif args.mode == "kp_traffic": + result = conv_kp_mixed_traffic_step( + net, x, y, config, step, args.traffic_rule, + args.predictor_every) + loss = result["loss"] + did_perturb = False + signal_metrics.append({key: result[key] for key in ( + "teaching_rms", "instruction_rms", "raw_apical_rms", + "innovation_rms", "traffic_rms")}) + if result["did_predictor_update"]: + counters["predictor_update_examples"] += batch else: result = conv_local_step( net, x, y, config, step, generator=perturb_generator) @@ -521,7 +627,12 @@ 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": + if signal_metrics: + record["mixed_apical"] = { + key: sum(value[key] for value in signal_metrics) + / len(signal_metrics) for key in signal_metrics[0] + } + if args.mode in ("kp", "kp_traffic"): tracking = hierarchical_feedback_tracking_report(net) record["feedback_tracking"] = { "mean_feedback_forward_cosine": tracking[ @@ -569,11 +680,15 @@ def run(args): probe = min(args.alignment_probe, train.x.shape[0]) sync(args.device) diagnostic_start = time.time() - diagnostics = (conv_hierarchical_alignment_report( - net, train.x[:probe], train.y[:probe]) - if args.mode in ("hfa", "lhfa", "wm", "rrm", "kp") else - conv_alignment_report( - net, train.x[:probe], train.y[:probe], config)) + if args.mode == "kp_traffic": + diagnostics = conv_kp_mixed_traffic_alignment_report( + net, train.x[:probe], train.y[:probe], args.traffic_rule) + elif args.mode in ("hfa", "lhfa", "wm", "rrm", "kp"): + diagnostics = conv_hierarchical_alignment_report( + net, train.x[:probe], train.y[:probe]) + else: + diagnostics = conv_alignment_report( + net, train.x[:probe], train.y[:probe], config) sync(args.device) diagnostics["wall_s"] = time.time() - diagnostic_start values = diagnostics["teaching_negative_gradient_cosine"] @@ -620,8 +735,8 @@ def parse_args(): parser = argparse.ArgumentParser() parser.add_argument( "--mode", choices=( - "bp", "dfa", "hfa", "lhfa", "wm", "rrm", "kp", "sdil", - "nodepert"), + "bp", "dfa", "hfa", "lhfa", "wm", "rrm", "kp", + "kp_traffic", "sdil", "nodepert"), required=True) parser.add_argument("--out", required=True) parser.add_argument("--device", default="cpu") @@ -673,6 +788,13 @@ def parse_args(): choices=("unit_targets", "channel_subspace", "vectorizer_subspace"), default="unit_targets") parser.add_argument("--predictor_warmup_steps", type=int, default=0) + parser.add_argument("--predictor_every", type=int, default=0) + parser.add_argument("--traffic_rule", + choices=("raw", "matched", "innovation"), + default="innovation") + parser.add_argument("--traffic_seed", type=int, default=4000) + parser.add_argument("--traffic_ratio", type=float, default=4.0) + parser.add_argument("--traffic_calibration_examples", type=int, default=64) parser.add_argument("--a_warmup_steps", type=int, default=0) parser.add_argument("--mirror_warmup_steps", type=int, default=0) parser.add_argument("--mirror_every", type=int, default=16) diff --git a/sdil/conv.py b/sdil/conv.py index 8331433..0beb249 100644 --- a/sdil/conv.py +++ b/sdil/conv.py @@ -598,6 +598,177 @@ class CIFARKPResNet(CIFARHierarchicalFAResNet): self.R_out.add_(update, alpha=eta_output) +class CIFARKPMixedTrafficResNet(CIFARKPResNet): + """KP credit with a per-unit mixed apical compartment and neutral predictor. + + The reciprocal KP pathway supplies the instructional field. Fixed + soma-coupled traffic is added locally at every hidden population, while a + diagonal affine predictor is fitted only to instruction-off observations. + Which of raw, norm-matched raw, or innovation drives plasticity is selected + by the runner; every condition still computes and trains the predictor. + """ + + def __init__(self, *args, traffic_seed=4000, **kwargs): + super().__init__(*args, **kwargs) + generator = torch.Generator(device="cpu").manual_seed(traffic_seed) + self.B_traffic = [] + self.traffic_gain = [] + self.P_traffic = [] + self.P_traffic_bias = [] + for shape in self.hidden_shapes: + coefficient = torch.exp( + 0.25 * torch.randn(shape, generator=generator)) + self.B_traffic.append(coefficient.to( + device=self.device, dtype=self.dtype)) + self.traffic_gain.append(torch.zeros( + (), device=self.device, dtype=self.dtype)) + self.P_traffic.append(torch.zeros( + shape, device=self.device, dtype=self.dtype)) + self.P_traffic_bias.append(torch.zeros( + shape, device=self.device, dtype=self.dtype)) + self.traffic_rule = None + + @property + def n_predictor_parameters(self): + return (sum(value.numel() for value in self.P_traffic) + + sum(value.numel() for value in self.P_traffic_bias)) + + @property + def n_fixed_traffic_coefficients(self): + return sum(value.numel() for value in self.B_traffic) + + @property + def n_apical_parameters(self): + # Reciprocal Q/R parameters are logged separately as adaptive feedback. + return self.n_predictor_parameters + + @property + def mixed_units_per_example(self): + return sum(math.prod(shape) for shape in self.hidden_shapes) + + def mixed_elementwise_ops_per_example(self, rule=None): + """Transparent arithmetic count beyond convolution/linear MACs. + + Five operations per unit form traffic, affine prediction, raw, and + innovation. Norm matching additionally charges squares/reductions, + rescaling, and scalar norm arithmetic conservatively as five per unit. + """ + rule = self.traffic_rule if rule is None else rule + if rule not in ("raw", "matched", "innovation"): + raise ValueError(f"unknown mixed-traffic rule: {rule}") + multiplier = 10 if rule == "matched" else 5 + return multiplier * self.mixed_units_per_example + + @property + def predictor_elementwise_ops_per_example(self): + # Residual, centering, moments, normalization, and affine updates. + return 12 * self.mixed_units_per_example + + @torch.no_grad() + def traffic_fields(self, hiddens): + if len(hiddens) != self.n_hidden: + raise ValueError("traffic requires every somatic population") + return [gain * coefficient * hidden for gain, coefficient, hidden in zip( + self.traffic_gain, self.B_traffic, hiddens)] + + @torch.no_grad() + def calibrate_traffic_gain(self, instruction, hiddens, target_ratio): + """Fix one gain per layer from an initialization-only training prefix.""" + if target_ratio <= 0: + raise ValueError("traffic ratio must be positive") + if not (len(instruction) == len(hiddens) == self.n_hidden): + raise ValueError("traffic calibration must cover every population") + instruction_rms = [] + unscaled_traffic_rms = [] + gains = [] + realized = [] + for signal, hidden, coefficient, gain in zip( + instruction, hiddens, self.B_traffic, self.traffic_gain): + signal_scale = signal.square().mean().sqrt() + traffic_scale = (coefficient * hidden).square().mean().sqrt() + if float(signal_scale) <= 0 or float(traffic_scale) <= 0: + raise ValueError("traffic calibration encountered a zero RMS") + value = target_ratio * signal_scale / traffic_scale + gain.copy_(value) + achieved = (gain * coefficient * hidden).square().mean().sqrt() + instruction_rms.append(float(signal_scale)) + unscaled_traffic_rms.append(float(traffic_scale)) + gains.append(float(gain)) + realized.append(float(achieved / signal_scale)) + return { + "target_ratio": float(target_ratio), + "instruction_rms": instruction_rms, + "unscaled_traffic_rms": unscaled_traffic_rms, + "traffic_gain": gains, + "realized_traffic_instruction_rms_ratio": realized, + } + + @torch.no_grad() + def mixed_apical_components(self, instruction, hiddens, rule): + """Return used, raw, innovation, matched, and traffic fields.""" + if rule not in ("raw", "matched", "innovation"): + raise ValueError(f"unknown mixed-traffic rule: {rule}") + if not (len(instruction) == len(hiddens) == self.n_hidden): + raise ValueError("mixed apical inputs must cover every population") + traffic = self.traffic_fields(hiddens) + raw = [] + innovation = [] + matched = [] + for signal, hidden, ordinary, slope, bias in zip( + instruction, hiddens, traffic, + self.P_traffic, self.P_traffic_bias): + apical = signal + ordinary + residual = apical - (slope * hidden + bias) + raw_norm = apical.flatten(1).norm(dim=1).clamp_min(1e-30) + residual_norm = residual.flatten(1).norm(dim=1) + scale = (residual_norm / raw_norm).reshape( + residual.shape[0], *([1] * (residual.ndim - 1))) + raw.append(apical) + innovation.append(residual) + matched.append(scale * apical) + choices = {"raw": raw, "matched": matched, "innovation": innovation} + return { + "used": choices[rule], "raw": raw, "innovation": innovation, + "matched": matched, "traffic": traffic, + "instruction": instruction, + } + + @torch.no_grad() + def predictor_step(self, hiddens, eta): + """Instruction-off normalized-LMS update from local soma/traffic pairs.""" + if not 0.0 < eta <= 1.0: + raise ValueError("predictor learning rate must lie in (0, 1]") + traffic = self.traffic_fields(hiddens) + squared_error = 0.0 + units = 0 + for hidden, target, slope, bias in zip( + hiddens, traffic, self.P_traffic, self.P_traffic_bias): + residual = target - slope * hidden - bias + centered_h = hidden - hidden.mean(dim=0) + centered_r = residual - residual.mean(dim=0) + variance = centered_h.square().mean(dim=0) + slope.add_((centered_r * centered_h).mean(dim=0) + / (variance + 1e-12), alpha=eta) + bias.add_(residual.mean(dim=0), alpha=eta) + squared_error += float(residual.square().sum()) + units += residual.numel() + return squared_error / units + + @torch.no_grad() + def predictor_traffic_residual_rms_ratio(self, hiddens): + traffic = self.traffic_fields(hiddens) + residual_power = 0.0 + traffic_power = 0.0 + for hidden, target, slope, bias in zip( + hiddens, traffic, self.P_traffic, self.P_traffic_bias): + residual = target - slope * hidden - bias + residual_power += float(residual.square().sum()) + traffic_power += float(target.square().sum()) + if traffic_power <= 0: + raise ValueError("predictor audit requires nonzero traffic") + return math.sqrt(residual_power / traffic_power) + + @torch.no_grad() def hierarchical_feedback_tracking_report(net): """Cheap parameter-space tracking diagnostics; never used for learning.""" @@ -1022,6 +1193,64 @@ def conv_kolen_pollack_step(net, x, y, config): } +def conv_kp_mixed_traffic_step(net, x, y, config, step, rule, + predictor_every): + """One reciprocal-KP update using a selected mixed-apical signal.""" + if not isinstance(net, CIFARKPMixedTrafficResNet): + raise TypeError("mixed-traffic KP step requires CIFARKPMixedTrafficResNet") + if predictor_every < 1: + raise ValueError("predictor cadence must be positive") + 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)) + instruction = net.hierarchical_teaching(output_error, forward) + components = net.mixed_apical_components( + instruction, forward["hiddens"], rule) + used = components["used"] + total_units = sum(value.numel() for value in used) + + def rms(values): + return math.sqrt( + sum(float(value.square().sum()) for value in values) / total_units) + + (directions, gamma_directions, beta_directions, + output_weight, output_bias) = net.local_ascent_directions( + used, output_error, forward) + reciprocal_directions, reciprocal_readout = ( + net.reciprocal_feedback_directions(used, output_error, forward)) + # Form both local correlations before either parameter path changes. + 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) + did_predictor_update = step % predictor_every == 0 + predictor_mse = None + if did_predictor_update: + predictor_mse = net.predictor_step( + forward["hiddens"], config.eta_P) + return { + "loss": float(loss), "did_perturb": False, "calibration": None, + "did_predictor_update": did_predictor_update, + "predictor_mse": predictor_mse, + "teaching_rms": rms(used), + "instruction_rms": rms(components["instruction"]), + "raw_apical_rms": rms(components["raw"]), + "innovation_rms": rms(components["innovation"]), + "traffic_rms": rms(components["traffic"]), + } + + 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() @@ -1089,6 +1318,73 @@ def conv_hierarchical_alignment_report(net, x, y): return report +def conv_kp_mixed_traffic_alignment_report(net, x, y, rule): + """Same-state audit of instruction/raw/innovation/matched directions.""" + if not isinstance(net, CIFARKPMixedTrafficResNet): + raise TypeError("mixed-traffic audit requires CIFARKPMixedTrafficResNet") + 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, return_cache=True, training=True, update_stats=False) + gradients = torch.autograd.grad( + F.cross_entropy(forward["logits"], y), forward["hiddens"]) + batch = x.shape[0] + negative_gradients = [-batch * value.detach() for value in gradients] + with torch.no_grad(): + output_error = (torch.softmax(forward["logits"], dim=1) + - F.one_hot(y, net.n_classes).to(forward["logits"].dtype)) + instruction = net.hierarchical_teaching(output_error, forward) + components = net.mixed_apical_components( + instruction, forward["hiddens"], rule) + + def align(values): + return [float(F.cosine_similarity( + left.flatten(1), right.flatten(1), dim=1).mean()) + for left, right in zip(values, negative_gradients)] + + norm_errors = [] + direction_errors = [] + for raw, innovation, matched in zip( + components["raw"], components["innovation"], + components["matched"]): + raw_flat = raw.flatten(1) + innovation_flat = innovation.flatten(1) + matched_flat = matched.flatten(1) + target_norm = innovation_flat.norm(dim=1) + norm_errors.append(float(( + (matched_flat.norm(dim=1) - target_norm).abs() + / target_norm.clamp_min(1e-30)).max())) + raw_match_cosine = F.cosine_similarity( + raw_flat, matched_flat, dim=1) + direction_errors.append(float((raw_match_cosine - 1.0).abs().max())) + + instruction_power = sum(float(value.square().sum()) + for value in components["instruction"]) + traffic_power = sum(float(value.square().sum()) + for value in components["traffic"]) + report = { + "normalization_state": "training_batch_stats_without_running_update", + "teaching_negative_gradient_cosine": align(components["used"]), + "used_negative_gradient_cosine": align(components["used"]), + "instruction_negative_gradient_cosine": align( + components["instruction"]), + "raw_negative_gradient_cosine": align(components["raw"]), + "innovation_negative_gradient_cosine": align( + components["innovation"]), + "matched_negative_gradient_cosine": align(components["matched"]), + "max_norm_match_relative_error": max(norm_errors), + "max_norm_match_direction_error": max(direction_errors), + "traffic_instruction_rms_ratio": math.sqrt( + traffic_power / instruction_power), + "predictor_traffic_residual_rms_ratio": ( + net.predictor_traffic_residual_rms_ratio(forward["hiddens"])), + } + report.update(hierarchical_feedback_tracking_report(net)) + for parameter in parameters: + parameter.requires_grad_(False) + return report + + class CIFARSDILResNet(CIFARLocalResNet): """CIFAR local ResNet with per-unit apical vectorizers and predictors. -- cgit v1.2.3