diff --git a/.gitignore b/.gitignore index 03eb840..a9854f8 100644 --- a/.gitignore +++ b/.gitignore @@ -218,4 +218,6 @@ __marimo__/ .streamlit/secrets.toml # Experimental documents for context, not to be committed -documents/* \ No newline at end of file +documents/* +dnc-context/* +tmp/* \ No newline at end of file diff --git a/README.md b/README.md index a624150..f009511 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,25 @@ The training core delegates task construction and task-specific behavior to a da The DNC supplies a memory matrix and differentiable read/write interface. The memory can retain information beyond a controller's immediate hidden state. Content-based addressing, read heads, and temporal-link structures are part of the underlying DNC implementation. The temporal link matrix has quadratic cost in the number of memory cells, which motivates the project's link-matrix ablations and memory-size investigations. +### Read/write mechanism roadmap + +The split-graph memory exposes independent read and write variants while retaining the DNC usage, allocation, write addressing, read modes, and temporal-link machinery: + +- `--dnc-read-variant cosine` is the DNC baseline. `relational-mlp` replaces cosine similarity with a learned score for each `(query, memory row)` pair. `relational-residual` adds a scaled learned pairwise score to cosine and zero-initializes its final layer, so it starts exactly at the DNC baseline. `--dnc-read-residual-scale` controls that correction's strength (default `1.0`). Track A logs the correction-to-cosine RMS ratio to show whether it contributes meaningfully. Write content addressing remains cosine-based for both learned readers. +- `--dnc-write-variant learned` is the DNC baseline. `kanerva-closed-form` turns the existing write-value projection into an observation and applies a per-row online isotropic Gaussian posterior update. It does not import the Kanerva Machine's exchangeable episode objective or claim to reproduce its full generative model. +- The variants compose, so the four baseline / Track A / Track B / combined cells are expressible through the same training and checkpoint pipeline. The closed-form writer uses a deterministic observation projection; the stochastic write head is not installed for that variant. +- The graph curriculum emits ID, OOD, and OOD-offset metrics as before. Graph runs also write `run__relational_probe.csv`, with cosine top-1, selected read-head top-1, and held-out-composition accuracy. `data/graph_traversal/order_probe.py` supplies paired prefix-permutation helpers for the Q99 order-sensitivity evaluation. +- Split-graph graph runs also write `run__track_a.csv` and `run__track_b.csv`. Track A records ID/OOD accuracy, relational stress and held-out-composition accuracy, ordinary-input read-score/cosine correlation, and read-content entropy/maximum weight. Track B records the Q99 order-sensitivity gate, ID/OOD posterior variance and write-update diagnostics, and training-window loss, gradient norm, AMP scale, and learning rate. These are experiment logs; no training or evaluation results are checked in. + +Use the selected split-graph baseline with either supported controller combiner: + +```bash +python -m initium.core_training 0.0 --split-graph --split-graph-variant mamba1 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc +python -m initium.core_training 0.0 --split-graph --split-graph-variant mamba1 --split-graph-combiner-mode controller --split-graph-combiner-variant mamba+cfc +``` + +Add `--dnc-read-variant relational-mlp` to reproduce the original Track A reader, or `--dnc-read-variant relational-residual --dnc-read-residual-scale 1.0` to use the cosine-anchored Track A reader. Lower scales such as `0.5` are available for follow-up comparisons. Use `--dnc-write-variant kanerva-closed-form` for Track B. Add a reader and writer variant together only for the combined cell after each track is characterized. Dense temporal links remain the default. The short-run baseline reproduction gate and the Q99 paired-order evaluation are deliberate user-run experiment steps; no run results are assumed by this repository change. + ### Controllers The training core exposes controller choices through its command-line interface and configuration: diff --git a/src/__init__.py b/src/__init__.py index cab731d..8d79c44 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -1,3 +1,3 @@ """initium - modular Differentiable Neural Computer research framework.""" -__version__ = "0.2.0" +__version__ = "0.2.1" diff --git a/src/initium/config/controller_config.py b/src/initium/config/controller_config.py index b3a4e12..5da51c8 100644 --- a/src/initium/config/controller_config.py +++ b/src/initium/config/controller_config.py @@ -67,6 +67,13 @@ SPLIT_GRAPH_COMBINER_VARIANT = "mamba1" SPLIT_GRAPH_COMBINER_NUM_BLOCKS = 1 +# DNC memory mechanism variants (only applied by the split-graph assembly). +DNC_READ_VARIANT = "cosine" # "cosine" | "relational-mlp" | "relational-residual" +DNC_READ_RESIDUAL_SCALE = 1.0 # Gain applied before tanh for relational-residual. +DNC_READ_RESIDUAL_MAX_RATIO = 0.5 # Maximum residual magnitude relative to cosine RMS. +DNC_WRITE_VARIANT = "learned" # "learned" | "kanerva-closed-form" +DNC_WRITE_OBSERVATION_VARIANCE = 1.0 # fixed isotropic likelihood variance + # ---- Option 1: link-matrix ablation/sparsification at fixed N ----------------- LINK_MATRIX_MODE = "dense" # "dense" | "ablated" | "sparse_topk" LINK_MATRIX_TOPK: int | None = None # int, required only for "sparse_topk" diff --git a/src/initium/core_training.py b/src/initium/core_training.py index bff8630..e25209a 100644 --- a/src/initium/core_training.py +++ b/src/initium/core_training.py @@ -61,6 +61,11 @@ CFC_MODE, CFC_RESIDUAL, CONTROLLER_TYPE, + DNC_READ_RESIDUAL_MAX_RATIO, + DNC_READ_RESIDUAL_SCALE, + DNC_READ_VARIANT, + DNC_WRITE_OBSERVATION_VARIANCE, + DNC_WRITE_VARIANT, DYNAMIC_N_CEILING, DYNAMIC_N_COOLDOWN_STEPS, DYNAMIC_N_EMA_DECAY, @@ -139,6 +144,15 @@ WARMUP_STEPS, ) from initium.data.dataset_registry import get_dataset +from initium.data.graph_traversal.order_probe import ( + build_q99_probe_batch, + evaluate_q99_order_sensitivity, +) +from initium.data.graph_traversal.relational_probe import ( + build_relational_probe, + cosine_top1_accuracy, + score_probe_cases, +) from initium.data.prepared_dataset import load_prepared_dataset, save_prepared_dataset from initium.mamba_controller.mamba_controller import MambaDNC from initium.mamba_controller.split_graph_dnc import SplitGraphDNC @@ -566,6 +580,11 @@ def run( split_graph_combiner_mode: str = SPLIT_GRAPH_COMBINER_MODE, split_graph_combiner_variant: str = SPLIT_GRAPH_COMBINER_VARIANT, split_graph_combiner_num_blocks: int = SPLIT_GRAPH_COMBINER_NUM_BLOCKS, + dnc_read_variant: str = DNC_READ_VARIANT, + dnc_read_residual_scale: float = DNC_READ_RESIDUAL_SCALE, + dnc_read_residual_max_ratio: float = DNC_READ_RESIDUAL_MAX_RATIO, + dnc_write_variant: str = DNC_WRITE_VARIANT, + dnc_write_observation_variance: float = DNC_WRITE_OBSERVATION_VARIANCE, dataset_type: str = DATASET_TYPE, dataset_link: str | None = DATASET_LINK, test_dataset_link: str | None = None, @@ -646,6 +665,114 @@ def run( ] ) + relational_probe_cases = ( + build_relational_probe(cell_size=MODEL_CELL_SIZE) if dataset_type == "graph" else None + ) + relational_probe_file = None + relational_probe_writer = None + if relational_probe_cases is not None: + relational_probe_file = open( + os.path.join(LOG_DIR, f"run_{run_id}_relational_probe.csv"), + "a" if resuming else "w", + newline="", + ) + relational_probe_writer = csv.writer(relational_probe_file) + if not resuming: + relational_probe_writer.writerow( + [ + "step", + "lesson", + "cosine_top1_accuracy", + "selected_read_top1_accuracy", + "selected_read_held_out_accuracy", + "cases", + ] + ) + + track_a_file = None + track_a_writer = None + track_b_file = None + track_b_writer = None + if split_graph_enabled and dataset_type == "graph": + track_a_file = open( + os.path.join(LOG_DIR, f"run_{run_id}_track_a.csv"), + "a" if resuming else "w", + newline="", + ) + track_a_writer = csv.writer(track_a_file) + if not resuming: + track_a_writer.writerow( + [ + "step", + "lesson", + "read_variant", + "id_triple_acc", + "ood_triple_acc", + "ood_offset_triple", + "stress_cosine_top1_acc", + "stress_selected_read_top1_acc", + "stress_held_out_acc", + "ordinary_read_score_cosine_rho_mean", + "ordinary_read_diagnostic_count", + "read_content_entropy_mean", + "read_content_max_weight_mean", + "relational_read_residual_rms_mean", + "relational_read_residual_to_cosine_rms_mean", + ] + ) + track_b_file = open( + os.path.join(LOG_DIR, f"run_{run_id}_track_b.csv"), + "a" if resuming else "w", + newline="", + ) + track_b_writer = csv.writer(track_b_file) + if not resuming: + track_b_writer.writerow( + [ + "event", + "step", + "lesson", + "write_variant", + "id_triple_acc", + "ood_triple_acc", + "ood_offset_triple", + "q99_gate", + "q99_tested_episodes", + "q99_tested_fact_episodes", + "q99_tested_query_episodes", + "q99_ordered_query_accuracy", + "q99_fact_permuted_accuracy", + "q99_fact_accuracy_delta", + "q99_fact_prediction_retention", + "q99_swapped_query_accuracy", + "q99_query_order_accuracy_drop", + "q99_query_prediction_flip_rate", + "q99_query_memory_incremental_order_effect", + "id_posterior_mean_variance", + "id_posterior_variance_min", + "id_posterior_variance_max", + "id_posterior_variance_reduction", + "id_posterior_mean_update_abs", + "id_active_write_cell_fraction", + "id_write_weight_mass", + "id_diagnostic_steps", + "ood_posterior_mean_variance", + "ood_posterior_variance_min", + "ood_posterior_variance_max", + "ood_posterior_variance_reduction", + "ood_posterior_mean_update_abs", + "ood_active_write_cell_fraction", + "ood_write_weight_mass", + "ood_diagnostic_steps", + "train_task_loss", + "train_kl_loss", + "train_total_loss", + "train_grad_norm", + "amp_scale", + "learning_rate", + ] + ) + ogs_log_path = os.path.join(LOG_DIR, f"run_{run_id}_ogs.csv") ogs_log_file = open(ogs_log_path, "a" if resuming else "w", newline="") ogs_log_writer = csv.writer(ogs_log_file) @@ -1015,6 +1142,11 @@ def run( moe_capacity_factor=moe_capacity_factor, moe_load_balance_alpha=moe_load_balance_alpha, moe_cfc_multi_source=split_graph_combiner_cfc_multi_source_moe, + dnc_read_variant=dnc_read_variant, + dnc_read_residual_scale=dnc_read_residual_scale, + dnc_read_residual_max_ratio=dnc_read_residual_max_ratio, + dnc_write_variant=dnc_write_variant, + dnc_write_observation_variance=dnc_write_observation_variance, device=device, ).to(device) else: @@ -1076,6 +1208,11 @@ def run( else None ), "split_graph_combiner_cfc_multi_source_moe": split_graph_combiner_cfc_multi_source_moe, + "dnc_read_variant": dnc_read_variant, + "dnc_read_residual_scale": dnc_read_residual_scale, + "dnc_read_residual_max_ratio": dnc_read_residual_max_ratio, + "dnc_write_variant": dnc_write_variant, + "dnc_write_observation_variance": dnc_write_observation_variance, } if link_matrix_mode != "dense": @@ -1096,9 +1233,15 @@ def run( # step 0 the sampled mean exactly matches Phase 0's write vector. # sample=True for all beta in BETAS_TO_SWEEP now that beta=0 is dropped # from the sweep (see header note) -- beta_target is always > 0.0 here. - stochastic_heads = install_stochastic_write_heads( - rnn, device=device, sample=(beta_target > 0.0) - ) + if split_graph_enabled and dnc_write_variant == "kanerva-closed-form": + # The Bayesian update consumes a deterministic observation from the + # DNC write-value projection. Sampling a stochastic write head here + # would make uncertainty sampling a second, confounded write change. + stochastic_heads = [] + else: + stochastic_heads = install_stochastic_write_heads( + rnn, device=device, sample=(beta_target > 0.0) + ) optimizer = torch.optim.Adam(list(rnn.parameters()) + list(output_proj.parameters()), lr=LR) # Option 3 (Concept 24/LB-17): multi-axis efficiency accounting. @@ -1110,7 +1253,7 @@ def run( p.numel() for p in output_proj.parameters() ) amp_enabled = USE_AMP and device.type == "cuda" - kl_on = stochastic_heads[0].sample # False when beta==0 -> no KL/prior terms in console + kl_on = bool(stochastic_heads) and stochastic_heads[0].sample print( f"[{run_id}] params {param_count} | kl_terms={'on' if kl_on else 'off'} | amp={'on' if amp_enabled else 'off'}" ) @@ -1551,6 +1694,38 @@ def set_lr(step_now): parts.append(dyn_n_str.strip(" |")) print(" | ".join(parts)) + if ( + track_b_writer is not None + and track_b_file is not None + and dnc_write_variant == "kanerva-closed-form" + ): + train_total_loss = avg_task + kl_contrib + track_b_writer.writerow( + [ + "training_window", + step, + curriculum.lesson + 1, + dnc_write_variant, + None, + None, + None, + *([None] * 28), + avg_task, + avg_kl, + train_total_loss, + avg_grad_norm, + amp_scale, + current_lr, + ] + ) + track_b_file.flush() + print( + f"[{run_id}] [Track B / train] task/total loss=" + f"{avg_task:.5f}/{train_total_loss:.5f} | " + f"grad_norm={avg_grad_norm:.4f} | amp_scale={amp_scale:.1f} | " + f"lr={current_lr:.6g}" + ) + log_writer.writerow( [ step, @@ -1629,9 +1804,34 @@ def set_lr(step_now): if step % EVAL_EVERY == 0: pre_advance_lesson = curriculum.lesson # capture before maybe_advance can bump it + id_stage_memory = rnn.memories[0] + if hasattr(id_stage_memory, "begin_stage_diagnostics"): + id_stage_memory.begin_stage_diagnostics() _, id_triple_acc, id_perfect_frac = curriculum.maybe_advance( rnn, device, step=step, optimizer=optimizer ) + id_stage_diagnostics = ( + id_stage_memory.finish_stage_diagnostics() + if hasattr(id_stage_memory, "finish_stage_diagnostics") + else {} + ) + if ( + relational_probe_writer is not None + and relational_probe_file is not None + and relational_probe_cases is not None + ): + probe_metrics = score_probe_cases(rnn.memories[0], relational_probe_cases) + relational_probe_writer.writerow( + [ + step, + curriculum.lesson + 1, + cosine_top1_accuracy(relational_probe_cases), + probe_metrics["accuracy"], + probe_metrics["held_out_accuracy"], + int(probe_metrics["cases"]), + ] + ) + relational_probe_file.flush() if curriculum.lesson != pre_advance_lesson: lesson_dip_start_step = step @@ -1693,6 +1893,9 @@ def set_lr(step_now): ood_field_log: dict[tuple[int, int], list[int]] = {} ood_rng_state = ood_rng.getstate() ood_torch_rng = _capture_torch_rng() + ood_stage_memory = rnn.memories[0] + if hasattr(ood_stage_memory, "begin_stage_diagnostics"): + ood_stage_memory.begin_stage_diagnostics() ood_triple_acc, ood_perfect_frac, ood_hop_breakdown = dataset.evaluate_ood( rnn, device, @@ -1700,6 +1903,11 @@ def set_lr(step_now): rng=ood_rng, field_log=ood_field_log, ) + ood_stage_diagnostics = ( + ood_stage_memory.finish_stage_diagnostics() + if hasattr(ood_stage_memory, "finish_stage_diagnostics") + else {} + ) ood_ablated_acc = None if hasattr(dataset, "evaluate_ood_ablated"): ood_ablated_rng = random.Random() @@ -1722,6 +1930,104 @@ def set_lr(step_now): ] ) ood_log_file.flush() + if ( + track_a_writer is not None + and track_a_file is not None + and relational_probe_cases is not None + ): + rho_key = "ordinary_read_score_cosine_correlation_mean" + rho = id_stage_diagnostics.get( + rho_key, + 1.0 if dnc_read_variant == "cosine" else None, + ) + rho_count = id_stage_diagnostics.get("diagnostic_steps", 0.0) + stress_metrics = score_probe_cases(rnn.memories[0], relational_probe_cases) + track_a_writer.writerow( + [ + step, + pre_advance_lesson + 1, + dnc_read_variant, + id_triple_acc, + ood_triple_acc, + id_triple_acc - ood_triple_acc, + cosine_top1_accuracy(relational_probe_cases), + stress_metrics["accuracy"], + stress_metrics["held_out_accuracy"], + rho, + rho_count, + id_stage_diagnostics.get("read_content_entropy_mean"), + id_stage_diagnostics.get("read_content_max_weight_mean"), + id_stage_diagnostics.get("relational_read_residual_rms_mean"), + id_stage_diagnostics.get("relational_read_residual_to_cosine_rms_mean"), + ] + ) + track_a_file.flush() + print( + f"[{run_id}] [Track A] read={dnc_read_variant} | " + f"ID/OOD/offset={id_triple_acc:.2f}/{ood_triple_acc:.2f}/" + f"{id_triple_acc - ood_triple_acc:.2f} | " + f"stress cosine/read/held-out=" + f"{cosine_top1_accuracy(relational_probe_cases):.3f}/" + f"{stress_metrics['accuracy']:.3f}/" + f"{stress_metrics['held_out_accuracy']:.3f} | " + f"ordinary MLP-cosine rho=" + f"{rho if rho is not None else float('nan'):.4f} | " + f"read entropy/max=" + f"{id_stage_diagnostics.get('read_content_entropy_mean', float('nan')):.3f}/" + f"{id_stage_diagnostics.get('read_content_max_weight_mean', float('nan')):.3f} | " + f"residual/cosine rms=" + f"{id_stage_diagnostics.get('relational_read_residual_to_cosine_rms_mean', float('nan')):.3f}" + ) + if track_b_writer is not None and track_b_file is not None: + id_write_values = [ + id_stage_diagnostics.get("posterior_variance_mean"), + id_stage_diagnostics.get("posterior_variance_min"), + id_stage_diagnostics.get("posterior_variance_max"), + id_stage_diagnostics.get("posterior_variance_reduction_mean"), + id_stage_diagnostics.get("posterior_mean_update_abs_mean"), + id_stage_diagnostics.get("active_write_cell_fraction_mean"), + id_stage_diagnostics.get("write_weight_mass_mean"), + id_stage_diagnostics.get("diagnostic_steps"), + ] + ood_write_values = [ + ood_stage_diagnostics.get("posterior_variance_mean"), + ood_stage_diagnostics.get("posterior_variance_min"), + ood_stage_diagnostics.get("posterior_variance_max"), + ood_stage_diagnostics.get("posterior_variance_reduction_mean"), + ood_stage_diagnostics.get("posterior_mean_update_abs_mean"), + ood_stage_diagnostics.get("active_write_cell_fraction_mean"), + ood_stage_diagnostics.get("write_weight_mass_mean"), + ood_stage_diagnostics.get("diagnostic_steps"), + ] + track_b_writer.writerow( + [ + "evaluation", + step, + pre_advance_lesson + 1, + dnc_write_variant, + id_triple_acc, + ood_triple_acc, + id_triple_acc - ood_triple_acc, + *([None] * 12), + *id_write_values, + *ood_write_values, + *([None] * 6), + ] + ) + track_b_file.flush() + if dnc_write_variant == "kanerva-closed-form": + print( + f"[{run_id}] [Track B] write={dnc_write_variant} | " + f"posterior variance ID/OOD=" + f"{id_stage_diagnostics.get('posterior_variance_mean', float('nan')):.5f}/" + f"{ood_stage_diagnostics.get('posterior_variance_mean', float('nan')):.5f} | " + f"mean update ID/OOD=" + f"{id_stage_diagnostics.get('posterior_mean_update_abs_mean', float('nan')):.6g}/" + f"{ood_stage_diagnostics.get('posterior_mean_update_abs_mean', float('nan')):.6g} | " + f"active cells ID/OOD=" + f"{id_stage_diagnostics.get('active_write_cell_fraction_mean', float('nan')):.4f}/" + f"{ood_stage_diagnostics.get('active_write_cell_fraction_mean', float('nan')):.4f}" + ) _write_ogs_row( ogs_log_writer, step, @@ -1879,15 +2185,43 @@ def set_lr(step_now): shutil.copyfile(ckpt_path, latest_path) print(f"\n[{run_id}] Training complete. Final evaluation on training-distribution lesson:") + final_id_stage_memory = rnn.memories[0] + if hasattr(final_id_stage_memory, "begin_stage_diagnostics"): + final_id_stage_memory.begin_stage_diagnostics() _, id_triple_acc, id_perfect_frac = curriculum.maybe_advance( rnn, device, step=step, optimizer=optimizer ) + final_id_stage_diagnostics = ( + final_id_stage_memory.finish_stage_diagnostics() + if hasattr(final_id_stage_memory, "finish_stage_diagnostics") + else {} + ) + if ( + relational_probe_writer is not None + and relational_probe_file is not None + and relational_probe_cases is not None + ): + probe_metrics = score_probe_cases(rnn.memories[0], relational_probe_cases) + relational_probe_writer.writerow( + [ + step, + curriculum.lesson + 1, + cosine_top1_accuracy(relational_probe_cases), + probe_metrics["accuracy"], + probe_metrics["held_out_accuracy"], + int(probe_metrics["cases"]), + ] + ) + relational_probe_file.flush() print(f"\n[{run_id}] Generalization test:") final_ood_rng_state = ood_rng.getstate() final_ood_torch_rng = _capture_torch_rng() final_ood_field_log: dict[tuple[int, int], list[int]] = {} + final_ood_stage_memory = rnn.memories[0] + if hasattr(final_ood_stage_memory, "begin_stage_diagnostics"): + final_ood_stage_memory.begin_stage_diagnostics() ood_triple_acc, ood_perfect_frac, ood_hop_breakdown = dataset.evaluate_ood( rnn, device, @@ -1896,10 +2230,98 @@ def set_lr(step_now): verbose_n=10, field_log=final_ood_field_log, ) + final_ood_stage_diagnostics = ( + final_ood_stage_memory.finish_stage_diagnostics() + if hasattr(final_ood_stage_memory, "finish_stage_diagnostics") + else {} + ) ood_offset_triple = id_triple_acc - ood_triple_acc ood_offset_perfect = id_perfect_frac - ood_perfect_frac + if ( + track_a_writer is not None + and track_a_file is not None + and relational_probe_cases is not None + ): + final_stress_metrics = score_probe_cases(rnn.memories[0], relational_probe_cases) + rho = final_id_stage_diagnostics.get( + "ordinary_read_score_cosine_correlation_mean", + 1.0 if dnc_read_variant == "cosine" else None, + ) + track_a_writer.writerow( + [ + step, + curriculum.lesson + 1, + dnc_read_variant, + id_triple_acc, + ood_triple_acc, + ood_offset_triple, + cosine_top1_accuracy(relational_probe_cases), + final_stress_metrics["accuracy"], + final_stress_metrics["held_out_accuracy"], + rho, + final_id_stage_diagnostics.get("diagnostic_steps", 0.0), + final_id_stage_diagnostics.get("read_content_entropy_mean"), + final_id_stage_diagnostics.get("read_content_max_weight_mean"), + final_id_stage_diagnostics.get("relational_read_residual_rms_mean"), + final_id_stage_diagnostics.get("relational_read_residual_to_cosine_rms_mean"), + ] + ) + track_a_file.flush() + print( + f"[{run_id}] [Track A / final] stress cosine/read/held-out=" + f"{cosine_top1_accuracy(relational_probe_cases):.3f}/" + f"{final_stress_metrics['accuracy']:.3f}/" + f"{final_stress_metrics['held_out_accuracy']:.3f} | " + f"ordinary MLP-cosine rho=" + f"{rho if rho is not None else float('nan'):.4f} | " + f"read entropy/max=" + f"{final_id_stage_diagnostics.get('read_content_entropy_mean', float('nan')):.3f}/" + f"{final_id_stage_diagnostics.get('read_content_max_weight_mean', float('nan')):.3f} | " + f"residual/cosine rms=" + f"{final_id_stage_diagnostics.get('relational_read_residual_to_cosine_rms_mean', float('nan')):.3f}" + ) + if track_b_writer is not None and track_b_file is not None: + + def _write_diagnostic_values(diagnostics): + return [ + diagnostics.get("posterior_variance_mean"), + diagnostics.get("posterior_variance_min"), + diagnostics.get("posterior_variance_max"), + diagnostics.get("posterior_variance_reduction_mean"), + diagnostics.get("posterior_mean_update_abs_mean"), + diagnostics.get("active_write_cell_fraction_mean"), + diagnostics.get("write_weight_mass_mean"), + diagnostics.get("diagnostic_steps"), + ] + + track_b_writer.writerow( + [ + "final_evaluation", + step, + curriculum.lesson + 1, + dnc_write_variant, + id_triple_acc, + ood_triple_acc, + ood_offset_triple, + *([None] * 12), + *_write_diagnostic_values(final_id_stage_diagnostics), + *_write_diagnostic_values(final_ood_stage_diagnostics), + *([None] * 6), + ] + ) + track_b_file.flush() + if dnc_write_variant == "kanerva-closed-form": + print( + f"[{run_id}] [Track B / final] posterior variance ID/OOD=" + f"{final_id_stage_diagnostics.get('posterior_variance_mean', float('nan')):.5f}/" + f"{final_ood_stage_diagnostics.get('posterior_variance_mean', float('nan')):.5f} | " + f"mean update ID/OOD=" + f"{final_id_stage_diagnostics.get('posterior_mean_update_abs_mean', float('nan')):.6g}/" + f"{final_ood_stage_diagnostics.get('posterior_mean_update_abs_mean', float('nan')):.6g}" + ) + ood_log_writer.writerow( [ step, @@ -1921,6 +2343,58 @@ def set_lr(step_now): final_ood_ablated_acc, _final_ood_ablated_perfect = dataset.evaluate_ood_ablated( rnn, device, final_ood_episodes, final_ood_ablated_rng ) + + q99_metrics = None + if split_graph_enabled and dnc_write_variant == "kanerva-closed-form": + q99_inputs, q99_targets, q99_answer_mask = build_q99_probe_batch() + q99_metrics = evaluate_q99_order_sensitivity( + rnn, + output_proj, + q99_inputs, + q99_targets, + q99_answer_mask, + ) + print( + f"[{run_id}] [Track B / Q99] post-training gate={q99_metrics['gate']} | " + f"episodes={q99_metrics['tested_episodes']} | " + f"ordered/fact-permuted accuracy=" + f"{q99_metrics['ordered_query_accuracy']:.2f}/" + f"{q99_metrics['fact_permuted_accuracy']:.2f}% " + f"(delta {q99_metrics['fact_accuracy_delta']:+.2f} pp) | " + f"ordered/swapped query accuracy=" + f"{q99_metrics['ordered_query_accuracy']:.2f}/" + f"{q99_metrics['swapped_query_accuracy']:.2f}% " + f"(drop {q99_metrics['query_order_accuracy_drop']:+.2f} pp) | " + f"query memory effect=" + f"{q99_metrics['query_memory_incremental_order_effect']:.6g}" + ) + if track_b_writer is not None and track_b_file is not None: + track_b_writer.writerow( + [ + "q99_semantic_order_check", + step, + curriculum.lesson + 1, + dnc_write_variant, + id_triple_acc, + ood_triple_acc, + ood_offset_triple, + q99_metrics["gate"], + q99_metrics["tested_episodes"], + q99_metrics["tested_fact_episodes"], + q99_metrics["tested_query_episodes"], + q99_metrics["ordered_query_accuracy"], + q99_metrics["fact_permuted_accuracy"], + q99_metrics["fact_accuracy_delta"], + q99_metrics["fact_prediction_retention"], + q99_metrics["swapped_query_accuracy"], + q99_metrics["query_order_accuracy_drop"], + q99_metrics["query_prediction_flip_rate"], + q99_metrics["query_memory_incremental_order_effect"], + *([None] * 16), + *([None] * 6), + ] + ) + track_b_file.flush() _write_ogs_row( ogs_log_writer, step, @@ -1962,6 +2436,10 @@ def set_lr(step_now): "ood_offset_perfect": ood_offset_perfect, "total_elapsed_sec": time.time() - t_run_start, } + if q99_metrics is not None: + summary["q99_gate"] = q99_metrics["gate"] + summary["q99_fact_permuted_accuracy"] = q99_metrics["fact_permuted_accuracy"] + summary["q99_query_order_accuracy_drop"] = q99_metrics["query_order_accuracy_drop"] print(f"\n[{run_id}] SUMMARY: {summary}") log_writer.writerow([]) @@ -1969,6 +2447,12 @@ def set_lr(step_now): log_writer.writerow([""] + list(summary.values())) log_file.close() ood_log_file.close() + if relational_probe_file is not None: + relational_probe_file.close() + if track_a_file is not None: + track_a_file.close() + if track_b_file is not None: + track_b_file.close() ogs_log_file.close() lesson_log_file.close() prior_log_file.close() @@ -2223,7 +2707,7 @@ def set_lr(step_now): type=str, default=SPLIT_GRAPH_COMBINER_MODE, choices=["linear", "controller"], - help="v11: 'linear' (default) is the original, already-efficient " + help="'linear' (default) is the original, already-efficient " "stateless combiner. 'controller' drives the sequential " "addressing step with a real interleaved Mamba controller " "cell instead (see --split-graph-combiner-variant).", @@ -2263,6 +2747,42 @@ def set_lr(step_now): "(--split-graph-combiner-mode=controller only). Kept small " "by default -- this step is meant to stay cheap.", ) + parser.add_argument( + "--dnc-read-variant", + choices=["cosine", "relational-mlp", "relational-residual"], + default=DNC_READ_VARIANT, + help="Read content score for --split-graph. relational-mlp replaces " + "cosine similarity; relational-residual adds a learned pairwise score " + "to cosine and starts exactly at the DNC baseline. DNC write addressing " + "and temporal links remain intact.", + ) + parser.add_argument( + "--dnc-read-residual-scale", + type=float, + default=DNC_READ_RESIDUAL_SCALE, + help="Gain applied before tanh to the normalized, centered relational correction; " + "ignored by other read variants. The correction is still bounded by the max-ratio.", + ) + parser.add_argument( + "--dnc-read-residual-max-ratio", + type=float, + default=DNC_READ_RESIDUAL_MAX_RATIO, + help="Maximum absolute relational correction as a fraction of per-query cosine RMS; " + "ignored by other read variants.", + ) + parser.add_argument( + "--dnc-write-variant", + choices=["learned", "kanerva-closed-form"], + default=DNC_WRITE_VARIANT, + help="Write update for --split-graph. kanerva-closed-form performs an " + "online Gaussian posterior update at each timestep.", + ) + parser.add_argument( + "--dnc-write-observation-variance", + type=float, + default=DNC_WRITE_OBSERVATION_VARIANCE, + help="Fixed positive observation variance for the Kanerva-style write.", + ) parser.add_argument( "--dataset-type", type=str, @@ -2328,6 +2848,16 @@ def set_lr(step_now): if args.prepare_only and args.load_prepared: parser.error("--prepare-only cannot be combined with --load-prepared") + if args.dnc_write_observation_variance <= 0: + parser.error("--dnc-write-observation-variance must be positive") + if args.dnc_read_residual_scale < 0: + parser.error("--dnc-read-residual-scale must be non-negative") + if args.dnc_read_residual_max_ratio < 0: + parser.error("--dnc-read-residual-max-ratio must be non-negative") + if not args.split_graph and ( + args.dnc_read_variant != "cosine" or args.dnc_write_variant != "learned" + ): + parser.error("--dnc-read/write-variant require --split-graph") if args.prepare_only and not args.save: parser.error("--prepare-only requires --save DIRECTORY") if args.save and args.load_prepared: @@ -2402,6 +2932,13 @@ def set_lr(step_now): run_id = f"{run_id}_moe{args.moe_num_experts}e" if args.split_graph: run_id = f"{run_id}_splitgraph_{args.split_graph_variant.replace('+', '')}" + if args.dnc_read_variant != "cosine" or args.dnc_write_variant != "learned": + run_id += f"_read-{args.dnc_read_variant}_write-{args.dnc_write_variant}" + if args.dnc_read_variant == "relational-residual": + run_id += f"_rscale{args.dnc_read_residual_scale:g}" + run_id += f"_rmax{args.dnc_read_residual_max_ratio:g}" + if args.dnc_write_variant == "kanerva-closed-form": + run_id += f"_obsvar{args.dnc_write_observation_variance:g}" if args.split_graph_combiner_mode == "controller": run_id = f"{run_id}_combctrl{args.split_graph_combiner_variant.replace('+', '')}" if args.dataset_type != "graph": @@ -2441,6 +2978,11 @@ def set_lr(step_now): split_graph_combiner_mode=args.split_graph_combiner_mode, split_graph_combiner_variant=args.split_graph_combiner_variant, split_graph_combiner_num_blocks=args.split_graph_combiner_num_blocks, + dnc_read_variant=args.dnc_read_variant, + dnc_read_residual_scale=args.dnc_read_residual_scale, + dnc_read_residual_max_ratio=args.dnc_read_residual_max_ratio, + dnc_write_variant=args.dnc_write_variant, + dnc_write_observation_variance=args.dnc_write_observation_variance, dataset_type=args.dataset_type, dataset_link=args.dataset_link, test_dataset_link=args.test_dataset_link, diff --git a/src/initium/data/graph_traversal/order_probe.py b/src/initium/data/graph_traversal/order_probe.py new file mode 100644 index 0000000..8962a30 --- /dev/null +++ b/src/initium/data/graph_traversal/order_probe.py @@ -0,0 +1,219 @@ +"""Semantic ordering probes for online memory writes (Q99).""" + +from __future__ import annotations + +import random + +import numpy as np +import torch + +Q99_PROBE_EPISODES = 128 +Q99_MIN_ORDERED_QUERY_ACCURACY = 50.0 +Q99_MAX_FACT_PERMUTATION_DROP = 5.0 +Q99_MIN_QUERY_ORDER_DROP = 1.0 + + +def build_q99_probe_batch(batch_size: int = Q99_PROBE_EPISODES, seed: int = 991): + """Build a reproducible paired batch without advancing training RNG state.""" + from initium.data.graph_traversal.graph_traversal import ( + build_traversal_episode, + collate_fn, + ) + + python_state = random.getstate() + numpy_state = np.random.get_state() + try: + random.seed(seed) + np.random.seed(seed) + episodes: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + while len(episodes) < batch_size: + episode = build_traversal_episode(10, (2, 4), (2, 4)) + if episode is not None: + episodes.append(episode) + return collate_fn(episodes) + finally: + random.setstate(python_state) + np.random.set_state(numpy_state) + + +def _fact_row_indices(episode: torch.Tensor, label_dim: int) -> torch.Tensor: + destination = episode[:, 2 * label_dim : 3 * label_dim].abs().sum(dim=-1) > 0 + prediction_required = episode[:, -1] > 0.5 + return torch.nonzero(destination & ~prediction_required, as_tuple=False).flatten() + + +def _path_query_indices(episode: torch.Tensor, label_dim: int) -> torch.Tensor: + source = episode[:, :label_dim].abs().sum(dim=-1) > 0 + edge = episode[:, label_dim : 2 * label_dim].abs().sum(dim=-1) > 0 + destination = episode[:, 2 * label_dim : 3 * label_dim].abs().sum(dim=-1) > 0 + prediction_required = episode[:, -1] > 0.5 + return torch.nonzero( + ~destination & ~prediction_required & (source | edge), as_tuple=False + ).flatten() + + +def permute_graph_fact_writes( + inputs: torch.Tensor, *, label_dim: int = 30, seed: int = 992 +) -> tuple[torch.Tensor, int]: + """Shuffle graph facts as a set, preserving the fact/query phase marker.""" + if inputs.ndim != 3 or inputs.size(-1) < 3 * label_dim + 2: + raise ValueError("expected graph inputs shaped (batch, time, 3*label_dim + 2)") + result = inputs.clone() + rng = random.Random(seed) + tested = 0 + for batch_index in range(inputs.size(0)): + indices = _fact_row_indices(inputs[batch_index], label_dim) + if indices.numel() < 2: + continue + order = list(range(indices.numel())) + rng.shuffle(order) + # Avoid treating an unchanged draw as a successful permutation. + if order == list(range(len(order))): + order = order[1:] + order[:1] + source_rows = indices[torch.tensor(order, device=indices.device)] + result[batch_index, indices] = inputs[batch_index, source_rows] + # The first-row marker denotes entry to the fact phase, not a property + # of an individual fact, so it stays at the first fact after shuffling. + result[batch_index, indices, -2] = 0.0 + result[batch_index, indices[0], -2] = 1.0 + tested += 1 + return result, tested + + +def swap_first_two_path_query_writes( + inputs: torch.Tensor, *, label_dim: int = 30 +) -> tuple[torch.Tensor, int]: + """Make an intentionally invalid order control for the dependent path queries.""" + if inputs.ndim != 3 or inputs.size(-1) < 3 * label_dim + 2: + raise ValueError("expected graph inputs shaped (batch, time, 3*label_dim + 2)") + result = inputs.clone() + tested = 0 + for batch_index in range(inputs.size(0)): + indices = _path_query_indices(inputs[batch_index], label_dim) + if indices.numel() < 2: + continue + first, second = indices[:2] + result[batch_index, first] = inputs[batch_index, second] + result[batch_index, second] = inputs[batch_index, first] + tested += 1 + return result, tested + + +def _triple_accuracy(logits: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor) -> float: + predicted = logits.argmax(dim=-1) + correct_triples = predicted.eq(targets).all(dim=-1) + selected = mask.to(device=logits.device, dtype=torch.bool) + if not selected.any(): + return 0.0 + return float(correct_triples[selected].float().mean().item() * 100.0) + + +def _answer_prediction_retention( + first: torch.Tensor, second: torch.Tensor, mask: torch.Tensor +) -> float: + selected = mask.to(device=first.device, dtype=torch.bool) + if not selected.any(): + return 0.0 + same = first.argmax(dim=-1).eq(second.argmax(dim=-1)).all(dim=-1) + return float(same[selected].float().mean().item() * 100.0) + + +@torch.no_grad() +def evaluate_q99_order_sensitivity( + model, + output_projection, + inputs: torch.Tensor, + target_digits: torch.Tensor, + answer_mask: torch.Tensor, + *, + atol: float = 1e-6, +) -> dict[str, float | bool | int | str]: + """Check fact-order invariance and ordered-query correctness separately. + + Graph facts are a set: shuffling them must preserve answer accuracy. + Path queries form a dependent sequence: swapping the first two query writes + is an intentionally invalid control and should reduce correctness. The + latter is compared with memory ablated to isolate memory's incremental role. + """ + fact_permuted, fact_episodes = permute_graph_fact_writes(inputs) + query_swapped, query_episodes = swap_first_two_path_query_writes(inputs) + if fact_episodes == 0 or query_episodes == 0: + return { + "gate": "not_applicable", + "tested_episodes": int(inputs.size(0)), + "tested_fact_episodes": fact_episodes, + "tested_query_episodes": query_episodes, + "ordered_query_accuracy": 0.0, + "fact_permuted_accuracy": 0.0, + "fact_accuracy_delta": 0.0, + "fact_prediction_retention": 0.0, + "swapped_query_accuracy": 0.0, + "query_order_accuracy_drop": 0.0, + "query_prediction_flip_rate": 0.0, + "query_memory_incremental_order_effect": 0.0, + } + + was_training = model.training + projection_was_training = output_projection.training + model.eval() + output_projection.eval() + + def forward(sequence: torch.Tensor, pass_through_memory: bool) -> torch.Tensor: + raw_output, _ = model( + sequence, + (None, None, None), + reset_experience=True, + pass_through_memory=pass_through_memory, + ) + return output_projection(raw_output.transpose(0, 1)).view( + sequence.size(0), sequence.size(1), 9, 10 + ) + + try: + device = next(model.parameters()).device + inputs = inputs.to(device) + fact_permuted = fact_permuted.to(device) + query_swapped = query_swapped.to(device) + target_digits = target_digits.to(device=device, dtype=torch.long) + answer_mask = answer_mask.to(device=device, dtype=torch.bool) + + ordered_full = forward(inputs, True) + fact_permuted_full = forward(fact_permuted, True) + query_swapped_full = forward(query_swapped, True) + ordered_ablate = forward(inputs, False) + query_swapped_ablate = forward(query_swapped, False) + finally: + model.train(was_training) + output_projection.train(projection_was_training) + + ordered_accuracy = _triple_accuracy(ordered_full, target_digits, answer_mask) + fact_accuracy = _triple_accuracy(fact_permuted_full, target_digits, answer_mask) + swapped_accuracy = _triple_accuracy(query_swapped_full, target_digits, answer_mask) + query_drop = ordered_accuracy - swapped_accuracy + query_flip = _answer_prediction_retention(ordered_full, query_swapped_full, answer_mask) + memory_delta = ( + (ordered_full - query_swapped_full) - (ordered_ablate - query_swapped_ablate) + ).abs()[answer_mask] + memory_effect = float(memory_delta.mean().item()) if memory_delta.numel() else 0.0 + fact_delta = fact_accuracy - ordered_accuracy + + fact_order_ok = fact_delta >= -Q99_MAX_FACT_PERMUTATION_DROP + ordered_queries_ok = ordered_accuracy >= Q99_MIN_ORDERED_QUERY_ACCURACY + query_order_ok = query_drop >= Q99_MIN_QUERY_ORDER_DROP and memory_effect > atol + gate_pass = fact_order_ok and ordered_queries_ok and query_order_ok + return { + "gate": "pass" if gate_pass else "fail", + "tested_episodes": int(inputs.size(0)), + "tested_fact_episodes": fact_episodes, + "tested_query_episodes": query_episodes, + "ordered_query_accuracy": ordered_accuracy, + "fact_permuted_accuracy": fact_accuracy, + "fact_accuracy_delta": fact_delta, + "fact_prediction_retention": _answer_prediction_retention( + ordered_full, fact_permuted_full, answer_mask + ), + "swapped_query_accuracy": swapped_accuracy, + "query_order_accuracy_drop": query_drop, + "query_prediction_flip_rate": 100.0 - query_flip, + "query_memory_incremental_order_effect": memory_effect, + } diff --git a/src/initium/data/graph_traversal/relational_probe.py b/src/initium/data/graph_traversal/relational_probe.py new file mode 100644 index 0000000..fbfd6ef --- /dev/null +++ b/src/initium/data/graph_traversal/relational_probe.py @@ -0,0 +1,123 @@ +"""Deterministic relational retrieval probes for DNC read mechanisms. + +The fixtures encode a small knowledge graph with role-separated node/relation +features and deliberately include high-overlap distractors. They are an eval +artifact for comparing cosine and relational read heads; they do not alter the +training curriculum or run automatically during training. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class RelationalProbeCase: + """Candidate memory rows, query key, and correct row index for one query.""" + + memory: torch.Tensor # (N, W) + query: torch.Tensor # (W,) + target_index: int + relation_id: int + held_out_composition: bool + + +def build_relational_probe( + *, cases: int = 8, cell_size: int = 64, seed: int = 17098 +) -> list[RelationalProbeCase]: + """Build fixed episodic fixtures with conjunctive relation cues. + + Each case uses role-separated source, relation, and binding channels. The + target row stores a pair-binding code (source multiplied by relation), + while a distractor repeats the query's source and relation as a bag of + features. Cosine therefore prefers the distractor; a relational scorer + can learn the pair interaction. Composition ids in the latter half are + held out so callers can report compositional generalization separately. + """ + if cases < 2: + raise ValueError("cases must be at least 2 to reserve held-out compositions") + # Three role channels are sufficient. Tiny smoke configurations use + # cell_size=8; requiring four dimensions per channel made the probe abort + # those runs before training, even though the fixture fits in fewer. + if cell_size < 6: + raise ValueError("cell_size must be at least 6") + + generator = torch.Generator(device="cpu").manual_seed(seed) + width = cell_size // 3 + codebook = torch.randn(3, width, width, generator=generator) + codebook = torch.nn.functional.normalize(codebook, dim=-1) + fixtures = [] + for index in range(cases): + source = codebook[0, index % width] + relation_id = (index * 3 + 1) % width + relation = codebook[1, relation_id] + binding = torch.nn.functional.normalize(source * relation, dim=0) + padding = torch.zeros(cell_size - 3 * width) + target = torch.cat((torch.zeros(width), torch.zeros(width), binding, padding)) + wrong_relation = codebook[1, (relation_id + 1) % width] + feature_lure = torch.cat((source, relation, torch.zeros(width), padding)) + source_only = torch.cat( + ( + torch.zeros(width), + torch.zeros(width), + torch.nn.functional.normalize(source * wrong_relation, dim=0), + padding, + ) + ) + # A fourth row is an unrelated distractor with a large norm to ensure + # the scorer responds to relational features rather than row magnitude. + noise = torch.randn(cell_size, generator=generator) * 0.1 + memory = torch.stack((feature_lure, source_only, target, noise)) + query = torch.cat((source, relation, torch.zeros(width), padding)) + fixtures.append( + RelationalProbeCase( + memory=memory, + query=query, + target_index=2, + relation_id=relation_id, + held_out_composition=index >= cases // 2, + ) + ) + return fixtures + + +def cosine_top1_accuracy(cases: list[RelationalProbeCase]) -> float: + """Return cosine top-1 accuracy for checking fixture construction.""" + if not cases: + return 0.0 + correct = 0 + for case in cases: + scores = torch.nn.functional.cosine_similarity(case.memory, case.query.unsqueeze(0), dim=-1) + correct += int(scores.argmax().item() == case.target_index) + return correct / len(cases) + + +def score_probe_cases(memory_module, cases: list[RelationalProbeCase]) -> dict[str, float]: + """Report relational-head top-1 accuracy overall and on held-out pairs.""" + if not cases: + return {"accuracy": 0.0, "held_out_accuracy": 0.0, "cases": 0.0} + device = next(memory_module.parameters()).device + correct, held_out_correct, held_out_count = 0, 0, 0 + with torch.no_grad(): + for case in cases: + memory = case.memory.to(device).unsqueeze(0) + query = case.query.to(device).view(1, 1, -1) + strengths = torch.ones(1, 1, device=device) + score_fn = getattr( + memory_module, + "read_content_weightings", + memory_module.content_weightings, + ) + weights = score_fn(memory, query, strengths) + is_correct = int(weights[0, 0].argmax().item() == case.target_index) + correct += is_correct + if case.held_out_composition: + held_out_correct += is_correct + held_out_count += 1 + return { + "accuracy": correct / len(cases), + "held_out_accuracy": held_out_correct / max(1, held_out_count), + "cases": float(len(cases)), + } diff --git a/src/initium/inference/model_loader.py b/src/initium/inference/model_loader.py index 931a1c4..0b68989 100644 --- a/src/initium/inference/model_loader.py +++ b/src/initium/inference/model_loader.py @@ -176,6 +176,17 @@ def _build_split_graph(c: _Cfg, input_dim, hidden, nr_cells, cell_size, read_hea combiner_num_blocks=c.get( "split_graph_combiner_num_blocks", cc.SPLIT_GRAPH_COMBINER_NUM_BLOCKS ), + dnc_read_variant=c.get("dnc_read_variant", cc.DNC_READ_VARIANT), + dnc_read_residual_scale=c.get( + # Older relational-residual checkpoints used gain 1. + "dnc_read_residual_scale", + 1.0, + ), + dnc_read_residual_max_ratio=c.get("dnc_read_residual_max_ratio", 0.5), + dnc_write_variant=c.get("dnc_write_variant", cc.DNC_WRITE_VARIANT), + dnc_write_observation_variance=c.get( + "dnc_write_observation_variance", cc.DNC_WRITE_OBSERVATION_VARIANCE + ), independent_linears=True, moe_enabled=_moe["moe_enabled"], moe_num_experts=_moe["moe_num_experts"], diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py index f31b077..11e3af8 100644 --- a/src/initium/mamba_controller/split_graph_dnc.py +++ b/src/initium/mamba_controller/split_graph_dnc.py @@ -114,7 +114,6 @@ import torch import torch.nn as nn -from dnc.memory import Memory from initium.LNN_controller.cfc_backbone_parallel import ( build_parallel_backbone, @@ -128,6 +127,7 @@ from initium.mamba_controller.mamba2_controller import Mamba2ControllerWrapper from initium.mamba_controller.mamba3_controller import Mamba3ControllerWrapper from initium.mamba_controller.mamba_controller import MambaControllerWrapper +from initium.memory_manipulation.dnc_mechanisms import build_memory class SplitGraphDNC(nn.Module): @@ -195,6 +195,11 @@ def __init__( # False to fall back to plain per-block MoE on the pre-concatenated # input, for direct A/B comparison. moe_cfc_multi_source: bool = True, + dnc_read_variant: str = "cosine", + dnc_read_residual_scale: float = 1.0, + dnc_read_residual_max_ratio: float = 0.5, + dnc_write_variant: str = "learned", + dnc_write_observation_variance: float = 1.0, ): super().__init__() if not independent_linears: @@ -218,6 +223,11 @@ def __init__( self.device = device self.mamba_variant = mamba_variant self.combine_reads = combine_reads + self.dnc_read_variant = dnc_read_variant + self.dnc_read_residual_scale = dnc_read_residual_scale + self.dnc_read_residual_max_ratio = dnc_read_residual_max_ratio + self.dnc_write_variant = dnc_write_variant + self.dnc_write_observation_variance = dnc_write_observation_variance # Mamba bookkeeping, for checkpoint self-description -- same # rationale/convention as MambaDNC (mamba_controller.py). @@ -437,13 +447,18 @@ def __init__( self.memories = [] self.memories.append( - Memory( + build_memory( input_size=self.output_size, nr_cells=self.nr_cells, cell_size=self.w, read_heads=self.r, device=self.device, independent_linears=independent_linears, + read_variant=dnc_read_variant, + write_variant=dnc_write_variant, + read_residual_scale=dnc_read_residual_scale, + read_residual_max_ratio=dnc_read_residual_max_ratio, + observation_variance=dnc_write_observation_variance, ) ) # See module docstring's API-compatibility section for why this diff --git a/src/initium/memory_manipulation/dnc_mechanisms.py b/src/initium/memory_manipulation/dnc_mechanisms.py new file mode 100644 index 0000000..5a7571a --- /dev/null +++ b/src/initium/memory_manipulation/dnc_mechanisms.py @@ -0,0 +1,355 @@ +"""Composable DNC read-addressing and write-value mechanisms. + +The stock pytorch-dnc Memory is kept as the baseline. This subclass changes +only the read content score and/or the update applied to the selected rows; +usage, allocation, erase/write addressing, read modes and temporal links keep +their upstream implementations and semantics. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from dnc.memory import Memory + +READ_VARIANTS = ("cosine", "relational-mlp", "relational-residual") +WRITE_VARIANTS = ("learned", "kanerva-closed-form") + + +class MechanismMemory(Memory): + """Stock DNC memory with independently selectable read and write rules.""" + + def __init__( + self, + *args, + read_variant: str = "cosine", + write_variant: str = "learned", + read_residual_scale: float = 1.0, + read_residual_max_ratio: float = 0.5, + relational_hidden_size: int | None = None, + observation_variance: float = 1.0, + **kwargs, + ): + if read_variant not in READ_VARIANTS: + raise ValueError(f"Unknown DNC read variant: {read_variant!r}") + if write_variant not in WRITE_VARIANTS: + raise ValueError(f"Unknown DNC write variant: {write_variant!r}") + if read_residual_scale < 0: + raise ValueError("read_residual_scale must be non-negative") + if read_residual_max_ratio < 0: + raise ValueError("read_residual_max_ratio must be non-negative") + if observation_variance <= 0: + raise ValueError("observation_variance must be positive") + super().__init__(*args, **kwargs) + self.read_variant = read_variant + self.write_variant = write_variant + self.read_residual_scale = float(read_residual_scale) + self.read_residual_max_ratio = float(read_residual_max_ratio) + self.observation_variance = float(observation_variance) + self._collect_stage_diagnostics = False + self._stage_diagnostic_sums: dict[str, torch.Tensor] = {} + self._stage_diagnostic_counts: dict[str, int] = {} + self._stage_diagnostic_min: dict[str, torch.Tensor] = {} + self._stage_diagnostic_max: dict[str, torch.Tensor] = {} + self.relational_score: nn.Module | None = None + if read_variant in {"relational-mlp", "relational-residual"}: + hidden = relational_hidden_size or max(32, self.cell_size) + self.relational_score = nn.Sequential( + nn.Linear(2 * self.cell_size, hidden), + nn.GELU(), + nn.Linear(hidden, 1), + ) + if read_variant == "relational-residual": + # Preserve the DNC content-addressing behavior at step zero; + # the learned relation function starts as a zero residual + # and is then optimized from task feedback. + output_layer = self.relational_score[-1] + if not isinstance(output_layer, nn.Linear): + raise TypeError("relational score output layer must be linear") + nn.init.zeros_(output_layer.weight) + if output_layer.bias is not None: + nn.init.zeros_(output_layer.bias) + if self.device is not None: + self.relational_score.to(self.device) + + def begin_stage_diagnostics(self) -> None: + """Collect lightweight read/write statistics until finish is called.""" + self._collect_stage_diagnostics = True + self._stage_diagnostic_sums = {} + self._stage_diagnostic_counts = {} + self._stage_diagnostic_min = {} + self._stage_diagnostic_max = {} + + def _record_stage_value(self, name: str, value: torch.Tensor) -> None: + values = value.detach().float() + value_mean = values.mean() + value_min = values.min() + value_max = values.max() + if name in self._stage_diagnostic_sums: + self._stage_diagnostic_sums[name] = self._stage_diagnostic_sums[name] + value_mean + self._stage_diagnostic_min[name] = torch.minimum( + self._stage_diagnostic_min[name], value_min + ) + self._stage_diagnostic_max[name] = torch.maximum( + self._stage_diagnostic_max[name], value_max + ) + self._stage_diagnostic_counts[name] += 1 + else: + self._stage_diagnostic_sums[name] = value_mean + self._stage_diagnostic_min[name] = value_min + self._stage_diagnostic_max[name] = value_max + self._stage_diagnostic_counts[name] = 1 + + def finish_stage_diagnostics(self) -> dict[str, float]: + """Stop collection and return per-write/read means and extrema.""" + self._collect_stage_diagnostics = False + result: dict[str, float] = {} + for name, total in self._stage_diagnostic_sums.items(): + result[f"{name}_mean"] = float((total / self._stage_diagnostic_counts[name]).item()) + result[f"{name}_min"] = float(self._stage_diagnostic_min[name].item()) + result[f"{name}_max"] = float(self._stage_diagnostic_max[name].item()) + result["diagnostic_steps"] = float(max(self._stage_diagnostic_counts.values(), default=0)) + return result + + def new(self, batch_size: int = 1): + hidden = super().new(batch_size) + if self.write_variant == "kanerva-closed-form": + # Isotropic Gaussian covariance per row, represented by its scalar + # variance. This is episode state, not a learned model parameter. + hidden["write_posterior_variance"] = torch.ones( + batch_size, + self.nr_cells, + 1, + device=hidden["memory"].device, + dtype=hidden["memory"].dtype, + ) + return hidden + + def clone(self, hidden): + cloned = super().clone(hidden) + if self.write_variant == "kanerva-closed-form": + cloned["write_posterior_variance"] = hidden["write_posterior_variance"].clone() + return cloned + + def erase(self, hidden): + hidden = super().erase(hidden) + if self.write_variant == "kanerva-closed-form": + hidden["write_posterior_variance"].fill_(1.0) + return hidden + + def _read_score_components(self, memory, keys): + """Return cosine and learned score components before key-strength scaling.""" + if self.read_variant == "cosine": + return self._dnc_cosine_scores(memory, keys), None + batch, rows, width = memory.shape + heads = keys.size(1) + candidates = memory[:, None, :, :].expand(batch, heads, rows, width) + queries = keys[:, :, None, :].expand(batch, heads, rows, width) + pair = torch.cat((queries, candidates), dim=-1) + relational_scores = self.relational_score(pair).squeeze(-1) + if self.read_variant == "relational-residual": + cosine_scores = self._dnc_cosine_scores(memory, keys) + return cosine_scores, relational_scores + return relational_scores, None + + def _read_similarity_scores(self, memory, keys): + """Return one unscaled read score for every query and memory row.""" + base_scores, residual_scores = self._read_score_components(memory, keys) + if residual_scores is not None: + residual_scores = residual_scores - residual_scores.mean(dim=-1, keepdim=True) + # Add epsilon before sqrt: the residual scorer is zero-initialized, + # so RMS can be exactly zero on the first forward pass. Clamping + # after sqrt leaves an infinite sqrt derivative in the graph. + residual_rms = (residual_scores.square().mean(dim=-1, keepdim=True) + 1e-12).sqrt() + normalized_residual = residual_scores / residual_rms + cosine_rms = (base_scores.square().mean(dim=-1, keepdim=True) + 1e-12).sqrt() + bounded_residual = ( + self.read_residual_max_ratio + * cosine_rms + * torch.tanh(self.read_residual_scale * normalized_residual) + ) + return base_scores + bounded_residual + return base_scores + + @staticmethod + def _dnc_cosine_scores(memory, keys): + """Match pytorch-dnc's cosine denominator, including its 1e-6 delta.""" + queries = keys.unsqueeze(2) + candidates = memory.unsqueeze(1) + dot = (queries * candidates).sum(dim=-1) + query_norm = torch.linalg.vector_norm(queries, dim=-1) + memory_norm = torch.linalg.vector_norm(candidates, dim=-1) + return dot / (query_norm * memory_norm + 1e-6) + + def read_content_weightings(self, memory, keys, strengths): + """Score rows, retaining DNC's key-strength softmax normalization.""" + scores = self._read_similarity_scores(memory, keys) + return F.softmax(scores * strengths.unsqueeze(-1), dim=-1) + + def read(self, read_keys, read_strengths, read_modes, hidden): + if self.read_variant == "cosine": + if self._collect_stage_diagnostics: + # Use the upstream content-weighting implementation so the + # baseline diagnostics match the actual DNC read path exactly. + with torch.no_grad(): + content_weights = self.content_weightings( + hidden["memory"], read_keys, read_strengths + ) + entropy = -( + content_weights.clamp_min(1e-12) * content_weights.clamp_min(1e-12).log() + ).sum(dim=-1) + self._record_stage_value("read_content_entropy", entropy) + self._record_stage_value( + "read_content_max_weight", content_weights.max(dim=-1).values + ) + self._record_stage_value( + "ordinary_read_score_cosine_correlation", + torch.ones_like(read_strengths), + ) + return super().read(read_keys, read_strengths, read_modes, hidden) + if self._collect_stage_diagnostics: + with torch.no_grad(): + memory = hidden["memory"] + cosine, residual = self._read_score_components(memory, read_keys) + scores = self._read_similarity_scores(memory, read_keys) + cosine_centered = cosine - cosine.mean(dim=-1, keepdim=True) + scores_centered = scores - scores.mean(dim=-1, keepdim=True) + numerator = (cosine_centered * scores_centered).sum(dim=-1) + denominator = torch.sqrt( + cosine_centered.square().sum(dim=-1) * scores_centered.square().sum(dim=-1) + ).clamp_min(1e-12) + self._record_stage_value( + "ordinary_read_score_cosine_correlation", + numerator / denominator, + ) + if residual is not None: + effective_residual = scores - cosine + cosine_rms = cosine.square().mean(dim=-1).sqrt().clamp_min(1e-12) + residual_rms = effective_residual.square().mean(dim=-1).sqrt() + self._record_stage_value("relational_read_residual_rms", residual_rms) + self._record_stage_value( + "relational_read_residual_to_cosine_rms", + self.read_residual_scale * residual_rms / cosine_rms, + ) + content_weights = self.read_content_weightings(hidden["memory"], read_keys, read_strengths) + if self._collect_stage_diagnostics: + with torch.no_grad(): + entropy = -( + content_weights.clamp_min(1e-12) * content_weights.clamp_min(1e-12).log() + ).sum(dim=-1) + self._record_stage_value("read_content_entropy", entropy) + self._record_stage_value( + "read_content_max_weight", content_weights.max(dim=-1).values + ) + hidden["read_weights"] = self.read_weightings( + hidden["memory"], + content_weights, + hidden["link_matrix"], + read_modes, + hidden["read_weights"], + ) + read_vectors = self.read_vectors(hidden["memory"], hidden["read_weights"]) + return read_vectors, hidden + + def write( + self, + write_key, + write_vector, + erase_vector, + free_gates, + read_strengths, + write_strength, + write_gate, + allocation_gate, + hidden, + ): + if self.write_variant == "learned": + return super().write( + write_key, + write_vector, + erase_vector, + free_gates, + read_strengths, + write_strength, + write_gate, + allocation_gate, + hidden, + ) + + # Keep DNC's usage, allocation and content-based write addressing. + hidden["usage_vector"] = self.get_usage_vector( + hidden["usage_vector"], + free_gates, + hidden["read_weights"], + hidden["write_weights"], + ) + write_content_weights = self.content_weightings(hidden["memory"], write_key, write_strength) + allocation, _ = self.allocate(hidden["usage_vector"], allocation_gate * write_gate) + hidden["write_weights"] = self.write_weighting( + hidden["memory"], + write_content_weights, + allocation, + write_gate, + allocation_gate, + ) + + # Exact online posterior mean/covariance update for each independently + # modeled row: prior N(memory_i, variance_i I), observation + # N(write_vector, observation_variance I), with DNC write weight as + # fractional observation precision. The learned projection supplies + # the observation; Bayesian fusion replaces erase-and-add blending. + weight = hidden["write_weights"].transpose(1, 2) + variance = hidden["write_posterior_variance"] + noise = self.observation_variance + gain = (weight * variance) / (noise + weight * variance) + observation = write_vector.expand_as(hidden["memory"]) + old_memory = hidden["memory"] + posterior_memory = old_memory + gain * (observation - hidden["memory"]) + hidden["memory"] = posterior_memory + hidden["write_posterior_variance"] = variance * noise / (noise + weight * variance) + if self._collect_stage_diagnostics: + posterior_variance = hidden["write_posterior_variance"] + self._record_stage_value( + "posterior_mean_update_abs", (posterior_memory - old_memory).abs() + ) + self._record_stage_value("posterior_variance", posterior_variance) + self._record_stage_value("posterior_variance_reduction", variance - posterior_variance) + self._record_stage_value("write_weight_mass", weight) + self._record_stage_value("active_write_cell_fraction", (weight > 1e-6).float()) + + hidden["link_matrix"] = self.get_link_matrix( + hidden["link_matrix"], hidden["write_weights"], hidden["precedence"] + ) + hidden["precedence"] = self.update_precedence(hidden["precedence"], hidden["write_weights"]) + return hidden + + +def build_memory( + *, + input_size: int, + nr_cells: int, + cell_size: int, + read_heads: int, + device=None, + independent_linears: bool = True, + read_variant: str = "cosine", + write_variant: str = "learned", + read_residual_scale: float = 1.0, + read_residual_max_ratio: float = 0.5, + observation_variance: float = 1.0, +): + """Construct a stock-compatible memory for the selected mechanism pair.""" + return MechanismMemory( + input_size=input_size, + nr_cells=nr_cells, + cell_size=cell_size, + read_heads=read_heads, + device=device, + independent_linears=independent_linears, + read_variant=read_variant, + write_variant=write_variant, + read_residual_scale=read_residual_scale, + read_residual_max_ratio=read_residual_max_ratio, + observation_variance=observation_variance, + ) diff --git a/src/initium/memory_manipulation/dynamic_memory_resize.py b/src/initium/memory_manipulation/dynamic_memory_resize.py index 521b45c..98256b3 100644 --- a/src/initium/memory_manipulation/dynamic_memory_resize.py +++ b/src/initium/memory_manipulation/dynamic_memory_resize.py @@ -50,6 +50,7 @@ import torch.nn as nn from dnc.memory import Memory +from initium.memory_manipulation.dnc_mechanisms import MechanismMemory, build_memory from initium.memory_manipulation.link_matrix_ablation import AblatableSparseLinkMemory @@ -95,14 +96,25 @@ def resize_memory(model, new_nr_cells: int, device=None, layer: int = 0, optimiz # used by the optimizer resync at the very end of this function. old_named_params = dict(old_memory.named_parameters()) - new_memory = Memory( + memory_args = dict( input_size=old_memory.input_size, - nr_cells=new_nr_cells, # VERIFY kwarg name + nr_cells=new_nr_cells, cell_size=old_memory.cell_size, read_heads=old_memory.read_heads, independent_linears=True, - device=device, # VERIFY kwarg name -- see note below - ).to(device) + device=device, + ) + if isinstance(old_memory, MechanismMemory): + new_memory = build_memory( + **memory_args, + read_variant=old_memory.read_variant, + write_variant=old_memory.write_variant, + read_residual_scale=old_memory.read_residual_scale, + read_residual_max_ratio=old_memory.read_residual_max_ratio, + observation_variance=old_memory.observation_variance, + ).to(device) + else: + new_memory = Memory(**memory_args).to(device) # Transplant every N-independent learned sublayer except # write_vector_transform (that one is a StochasticWriteHead, not a @@ -125,7 +137,13 @@ def resize_memory(model, new_nr_cells: int, device=None, layer: int = 0, optimiz # event (Option 2) can never silently undo an active link-matrix # ablation/sparsification (Option 1) mid-run. if isinstance(old_memory, AblatableSparseLinkMemory): - new_memory.__class__ = AblatableSparseLinkMemory + if not isinstance(new_memory, AblatableSparseLinkMemory): + original_class = new_memory.__class__ + new_memory.__class__ = type( + f"Ablatable{original_class.__name__}", + (AblatableSparseLinkMemory, original_class), + {}, + ) new_memory.link_matrix_mode = old_memory.link_matrix_mode if hasattr(old_memory, "link_matrix_topk"): new_memory.link_matrix_topk = old_memory.link_matrix_topk diff --git a/src/initium/memory_manipulation/link_matrix_ablation.py b/src/initium/memory_manipulation/link_matrix_ablation.py index a8de7df..50f5d10 100644 --- a/src/initium/memory_manipulation/link_matrix_ablation.py +++ b/src/initium/memory_manipulation/link_matrix_ablation.py @@ -123,7 +123,17 @@ def patch_link_matrix(model, mode: str, topk: int | None = None, layer: int | No targets = model.memories if layer is None else [model.memories[layer]] for memory in targets: - memory.__class__ = AblatableSparseLinkMemory + if not isinstance(memory, AblatableSparseLinkMemory): + # Preserve mechanism-specific overrides (for example relational + # read scoring and Bayesian writes) while layering this isolated + # link-matrix override on top of the memory's existing class. + original_class = memory.__class__ + combined_class = type( + f"Ablatable{original_class.__name__}", + (AblatableSparseLinkMemory, original_class), + {}, + ) + memory.__class__ = combined_class memory.link_matrix_mode = mode if topk is not None: memory.link_matrix_topk = topk