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 ++++++++++++++++++++++++----- 3 files changed, 283 insertions(+), 28 deletions(-) (limited to 'experiments') 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) -- cgit v1.2.3