From 0d4af59311e85dc678f14697994e34a78e0d1c76 Mon Sep 17 00:00:00 2001
From: Samuel <samueldanielyan@gmail.com>
Date: Fri, 2 Oct 2026 20:03:24 +0200
Subject: [PATCH 1/8] Add split-graph DNC mechanism variants

This change adds configurable split-graph DNC read/write mechanisms: relational-mlp read scoring and a Kanerva-style closed-form write posterior, wired through controller config, CLI arguments, and model loading. It also adds graph traversal probes for relational and Q99 order-sensitivity checks, updates experiment CSV logging for Track A/B diagnostics, and preserves mechanism-aware memory resizing and link-matrix ablations.
---
 .gitignore                                    |   4 +-
 README.md                                     |  19 +
 src/initium/config/controller_config.py       |   5 +
 src/initium/core_training.py                  | 468 +++++++++++++++++-
 .../data/graph_traversal/order_probe.py       | 185 +++++++
 .../data/graph_traversal/relational_probe.py  | 122 +++++
 src/initium/inference/model_loader.py         |   5 +
 .../mamba_controller/split_graph_dnc.py       |  13 +-
 .../memory_manipulation/dnc_mechanisms.py     | 283 +++++++++++
 .../dynamic_memory_resize.py                  |  26 +-
 .../link_matrix_ablation.py                   |  12 +-
 11 files changed, 1129 insertions(+), 13 deletions(-)
 create mode 100644 src/initium/data/graph_traversal/order_probe.py
 create mode 100644 src/initium/data/graph_traversal/relational_probe.py
 create mode 100644 src/initium/memory_manipulation/dnc_mechanisms.py

diff --git a/.gitignore b/.gitignore
index 03eb840..f185c83 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/*
\ No newline at end of file
diff --git a/README.md b/README.md
index a624150..40d5073 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` learns a score for each `(query, memory row)` pair; write content addressing remains cosine-based.
+- `--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_<id>_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_<id>_track_a.csv` and `run_<id>_track_b.csv`. Track A records ID/OOD accuracy, relational stress and held-out-composition accuracy, plus ordinary-input read-score/cosine correlation. 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` for Track A or `--dnc-write-variant kanerva-closed-form` for Track B. Add both 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/initium/config/controller_config.py b/src/initium/config/controller_config.py
index b3a4e12..f0404d8 100644
--- a/src/initium/config/controller_config.py
+++ b/src/initium/config/controller_config.py
@@ -67,6 +67,11 @@
 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"
+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..a54b0b8 100644
--- a/src/initium/core_training.py
+++ b/src/initium/core_training.py
@@ -69,6 +69,9 @@
     DYNAMIC_N_MODE,
     DYNAMIC_N_TRIGGER_FRAC,
     DYNAMIC_N_USAGE_HIGH,
+    DNC_READ_VARIANT,
+    DNC_WRITE_OBSERVATION_VARIANCE,
+    DNC_WRITE_VARIANT,
     HYBRID_CFC_NUM_BLOCKS,
     ISOLATE_LINK_ABLATION,
     LINK_MATRIX_MODE,
@@ -139,6 +142,15 @@
     WARMUP_STEPS,
 )
 from initium.data.dataset_registry import get_dataset
+from initium.data.graph_traversal.relational_probe import (
+    build_relational_probe,
+    cosine_top1_accuracy,
+    score_probe_cases,
+)
+from initium.data.graph_traversal.order_probe import (
+    build_q99_probe_batch,
+    evaluate_q99_order_sensitivity,
+)
 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 +578,9 @@ 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_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 +661,108 @@ 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",
+                ]
+            )
+        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_outputs_change_with_memory",
+                    "q99_outputs_change_without_memory",
+                    "q99_max_abs_output_delta",
+                    "q99_mean_abs_output_delta",
+                    "q99_prediction_flip_rate",
+                    "q99_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 +1132,9 @@ 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_write_variant=dnc_write_variant,
+            dnc_write_observation_variance=dnc_write_observation_variance,
             device=device,
         ).to(device)
     else:
@@ -1076,6 +1196,9 @@ 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_write_variant": dnc_write_variant,
+        "dnc_write_observation_variance": dnc_write_observation_variance,
     }
 
     if link_matrix_mode != "dense":
@@ -1096,9 +1219,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 +1239,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'}"
     )
@@ -1328,6 +1457,58 @@ def set_lr(step_now):
                 f"want to extend further."
             )
 
+    q99_metrics = None
+    if split_graph_enabled and dnc_write_variant == "kanerva-closed-form":
+        python_rng_state = random.getstate()
+        numpy_rng_state = np.random.get_state()
+        torch_rng_state = _capture_torch_rng()
+        try:
+            with _replay_torch_rng(torch_rng_state):
+                q99_inputs, _, q99_answer_mask = build_q99_probe_batch()
+                q99_metrics = evaluate_q99_order_sensitivity(
+                    rnn, output_proj, q99_inputs, q99_answer_mask
+                )
+        finally:
+            random.setstate(python_rng_state)
+            np.random.set_state(numpy_rng_state)
+
+        print(
+            f"[{run_id}] [Track B / Q99] gate={q99_metrics['gate']} | "
+            f"episodes={q99_metrics['tested_episodes']} | "
+            f"output_delta={q99_metrics.get('mean_abs_output_delta', 0.0):.6g} | "
+            f"prediction_flip={q99_metrics.get('prediction_flip_rate', 0.0):.4f} | "
+            f"memory_order_effect="
+            f"{q99_metrics.get('memory_incremental_order_effect', 0.0):.6g}"
+        )
+        if track_b_writer is not None:
+            track_b_writer.writerow(
+                [
+                    "q99_gate",
+                    step,
+                    curriculum.lesson + 1,
+                    dnc_write_variant,
+                    None,
+                    None,
+                    None,
+                    q99_metrics["gate"],
+                    q99_metrics["tested_episodes"],
+                    q99_metrics["outputs_change_with_memory"],
+                    q99_metrics["outputs_change_without_memory"],
+                    q99_metrics.get("max_abs_output_delta"),
+                    q99_metrics.get("mean_abs_output_delta"),
+                    q99_metrics.get("prediction_flip_rate"),
+                    q99_metrics.get("memory_incremental_order_effect"),
+                    *([None] * 16),
+                    *([None] * 6),
+                ]
+            )
+            track_b_file.flush()
+        if q99_metrics["gate"] != "pass":
+            raise RuntimeError(
+                "Track B Q99 order-sensitivity gate failed before training; "
+                "see the run's _track_b.csv q99_gate row."
+            )
+
     print(
         f"\n=== [{run_id}] Training (beta_target={beta_target}, beta_mode={beta_mode}) "
         f"{'[resumed]' if resuming else ''} ==="
@@ -1551,6 +1732,34 @@ def set_lr(step_now):
                 parts.append(dyn_n_str.strip(" |"))
             print(" | ".join(parts))
 
+            if track_b_writer 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] * 24),
+                        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 +1838,30 @@ 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:
+                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 +1923,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 +1933,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 +1960,91 @@ def set_lr(step_now):
                 ]
             )
             ood_log_file.flush()
+            if track_a_writer 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,
+                    ]
+                )
+                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}"
+                )
+            if track_b_writer 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] * 8),
+                        *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 +2202,39 @@ 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:
+        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 +2243,84 @@ 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:
+        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),
+            ]
+        )
+        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}"
+        )
+    if track_b_writer 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] * 8),
+                *_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,
@@ -1969,6 +2390,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()
@@ -2263,6 +2690,26 @@ 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"],
+        default=DNC_READ_VARIANT,
+        help="Read content score for --split-graph. relational-mlp replaces "
+        "only read similarity; DNC write addressing and temporal links remain intact.",
+    )
+    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 +2775,12 @@ 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 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 +2855,10 @@ 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_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 +2898,9 @@ 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_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..abd3971
--- /dev/null
+++ b/src/initium/data/graph_traversal/order_probe.py
@@ -0,0 +1,185 @@
+"""Utilities for the Q99 online-write order-sensitivity check."""
+
+from __future__ import annotations
+
+import torch
+
+
+def build_q99_probe_batch(batch_size: int = 8):
+    """Build graph episodes with at least two path writes for Q99 pairing."""
+    from initium.data.graph_traversal.graph_traversal import (
+        build_traversal_episode,
+        collate_fn,
+    )
+
+    episodes = []
+    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)
+
+
+def permute_early_writes(
+    inputs: torch.Tensor,
+    *,
+    prefix_length: int,
+    permutation: tuple[int, ...],
+) -> torch.Tensor:
+    """Return a copy with only the first ``prefix_length`` events reordered.
+
+    Inputs use the dataset's batch-first ``(batch, time, features)`` layout.
+    Queries and later sequence steps remain in their original positions.
+    ``permutation`` contains indices relative to the prefix and must be a
+    bijection, so callers can provide a deterministic paired episode.
+    """
+    if inputs.ndim != 3:
+        raise ValueError("inputs must have shape (batch, time, features)")
+    if not 0 < prefix_length <= inputs.size(1):
+        raise ValueError("prefix_length must be within the input sequence")
+    if len(permutation) != prefix_length or sorted(permutation) != list(range(prefix_length)):
+        raise ValueError("permutation must be a bijection over the write prefix")
+    result = inputs.clone()
+    indices = torch.tensor(permutation, device=inputs.device)
+    result[:, :prefix_length] = inputs[:, :prefix_length].index_select(1, indices)
+    return result
+
+
+def permute_early_graph_query_writes(
+    inputs: torch.Tensor, *, label_dim: int = 30
+) -> tuple[torch.Tensor, int]:
+    """Swap the first two path-query writes in each graph episode.
+
+    Graph facts have a populated destination field; path-query writes have a
+    source and/or edge field, no destination, and prediction_required=0. The
+    answer prompts and graph-fact write order are left untouched.
+    """
+    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)):
+        episode = inputs[batch_index]
+        source_present = episode[:, :label_dim].abs().sum(dim=-1) > 0
+        edge_present = episode[:, label_dim : 2 * label_dim].abs().sum(dim=-1) > 0
+        destination_present = episode[:, 2 * label_dim : 3 * label_dim].abs().sum(dim=-1) > 0
+        prediction_required = episode[:, -1] > 0.5
+        query_rows = torch.nonzero(
+            ~destination_present
+            & ~prediction_required
+            & (source_present | edge_present),
+            as_tuple=False,
+        ).flatten()
+        if query_rows.numel() < 2:
+            continue
+        first, second = query_rows[:2]
+        result[batch_index, first] = inputs[batch_index, second]
+        result[batch_index, second] = inputs[batch_index, first]
+        tested += 1
+    return result, tested
+
+
+def order_sensitivity_metrics(
+    ordered_output: torch.Tensor,
+    permuted_output: torch.Tensor,
+    *,
+    atol: float = 1e-6,
+    mask: torch.Tensor | None = None,
+) -> dict[str, float | bool]:
+    """Summarize whether task outputs respond to a prefix-order change."""
+    if ordered_output.shape != permuted_output.shape:
+        raise ValueError("paired outputs must have identical shapes")
+    if mask is not None:
+        if mask.shape != ordered_output.shape[: mask.ndim]:
+            raise ValueError("mask dimensions must match leading output dimensions")
+        selected = mask.to(device=ordered_output.device, dtype=torch.bool)
+        ordered_output = ordered_output[selected]
+        permuted_output = permuted_output[selected]
+        if ordered_output.numel() == 0:
+            raise ValueError("mask selects no output positions")
+    difference = (ordered_output - permuted_output).abs()
+    if ordered_output.ndim >= 2:
+        changed = ordered_output.argmax(dim=-1) != permuted_output.argmax(dim=-1)
+        prediction_flip_rate = changed.float().mean().item()
+    else:
+        prediction_flip_rate = float(
+            (ordered_output.round() != permuted_output.round()).float().mean().item()
+        )
+    return {
+        "outputs_change": bool((difference > atol).any().item()),
+        "max_absolute_difference": float(difference.max().item()),
+        "mean_absolute_difference": float(difference.mean().item()),
+        "prediction_flip_rate": prediction_flip_rate,
+    }
+
+
+@torch.no_grad()
+def evaluate_q99_order_sensitivity(
+    model,
+    output_projection,
+    inputs: torch.Tensor,
+    answer_mask: torch.Tensor,
+    *,
+    atol: float = 1e-6,
+) -> dict[str, float | bool | int | str]:
+    """Compare paired-order task outputs with and without external memory."""
+    permuted, tested_episodes = permute_early_graph_query_writes(inputs)
+    if tested_episodes == 0:
+        return {
+            "gate": "not_applicable",
+            "tested_episodes": 0,
+            "outputs_change_with_memory": False,
+            "outputs_change_without_memory": False,
+            "memory_incremental_order_effect": 0.0,
+            "prediction_flip_rate": 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:
+        inputs = inputs.to(next(model.parameters()).device)
+        permuted = permuted.to(inputs.device)
+        answer_mask = answer_mask.to(device=inputs.device, dtype=torch.bool)
+        ordered_full = forward(inputs, True)
+        permuted_full = forward(permuted, True)
+        ordered_ablate = forward(inputs, False)
+        permuted_ablate = forward(permuted, False)
+    finally:
+        model.train(was_training)
+        output_projection.train(projection_was_training)
+
+    full_metrics = order_sensitivity_metrics(
+        ordered_full, permuted_full, atol=atol, mask=answer_mask
+    )
+    ablated_metrics = order_sensitivity_metrics(
+        ordered_ablate, permuted_ablate, atol=atol, mask=answer_mask
+    )
+    delta_difference = (
+        (ordered_full - permuted_full) - (ordered_ablate - permuted_ablate)
+    ).abs()[answer_mask]
+    memory_effect = float(delta_difference.mean().item()) if delta_difference.numel() else 0.0
+    gate_pass = bool(full_metrics["outputs_change"]) and memory_effect > atol
+    return {
+        "gate": "pass" if gate_pass else "fail",
+        "tested_episodes": tested_episodes,
+        "outputs_change_with_memory": bool(full_metrics["outputs_change"]),
+        "outputs_change_without_memory": bool(ablated_metrics["outputs_change"]),
+        "max_abs_output_delta": float(full_metrics["max_absolute_difference"]),
+        "mean_abs_output_delta": float(full_metrics["mean_absolute_difference"]),
+        "prediction_flip_rate": float(full_metrics["prediction_flip_rate"]),
+        "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..3cdb866
--- /dev/null
+++ b/src/initium/data/graph_traversal/relational_probe.py
@@ -0,0 +1,122 @@
+"""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")
+    if cell_size < 12:
+        raise ValueError("cell_size must be at least 12")
+
+    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..7d91631 100644
--- a/src/initium/inference/model_loader.py
+++ b/src/initium/inference/model_loader.py
@@ -176,6 +176,11 @@ 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_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..5238fdc 100644
--- a/src/initium/mamba_controller/split_graph_dnc.py
+++ b/src/initium/mamba_controller/split_graph_dnc.py
@@ -114,7 +114,7 @@
 
 import torch
 import torch.nn as nn
-from dnc.memory import Memory
+from initium.memory_manipulation.dnc_mechanisms import build_memory
 
 from initium.LNN_controller.cfc_backbone_parallel import (
     build_parallel_backbone,
@@ -195,6 +195,9 @@ 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_write_variant: str = "learned",
+        dnc_write_observation_variance: float = 1.0,
     ):
         super().__init__()
         if not independent_linears:
@@ -218,6 +221,9 @@ def __init__(
         self.device = device
         self.mamba_variant = mamba_variant
         self.combine_reads = combine_reads
+        self.dnc_read_variant = dnc_read_variant
+        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 +443,16 @@ 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,
+                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..6fabc1d
--- /dev/null
+++ b/src/initium/memory_manipulation/dnc_mechanisms.py
@@ -0,0 +1,283 @@
+"""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")
+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",
+        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 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.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 == "relational-mlp":
+            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 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_content_weightings(self, memory, keys, strengths):
+        """Score candidate rows for reads; writes continue using DNC cosine."""
+        if self.read_variant == "cosine":
+            return self.content_weightings(memory, keys, strengths)
+        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)
+        logits = self.relational_score(pair).squeeze(-1)
+        return F.softmax(logits * strengths.unsqueeze(-1), dim=-1)
+
+    def read(self, read_keys, read_strengths, read_modes, hidden):
+        if self.read_variant == "cosine":
+            return super().read(read_keys, read_strengths, read_modes, hidden)
+        if self._collect_stage_diagnostics:
+            with torch.no_grad():
+                batch, rows, width = hidden["memory"].shape
+                heads = read_keys.size(1)
+                candidates = hidden["memory"][:, None, :, :].expand(
+                    batch, heads, rows, width
+                )
+                queries = read_keys[:, :, None, :].expand(batch, heads, rows, width)
+                cosine = F.cosine_similarity(queries, candidates, dim=-1)
+                relational = self.relational_score(
+                    torch.cat((queries, candidates), dim=-1)
+                ).squeeze(-1)
+                cosine_centered = cosine - cosine.mean(dim=-1, keepdim=True)
+                relational_centered = relational - relational.mean(dim=-1, keepdim=True)
+                numerator = (cosine_centered * relational_centered).sum(dim=-1)
+                denominator = torch.sqrt(
+                    cosine_centered.square().sum(dim=-1)
+                    * relational_centered.square().sum(dim=-1)
+                ).clamp_min(1e-12)
+                self._record_stage_value(
+                    "ordinary_read_score_cosine_correlation", numerator / denominator
+                )
+        content_weights = self.read_content_weightings(
+            hidden["memory"], read_keys, read_strengths
+        )
+        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",
+    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,
+        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..63a9f2f 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,23 @@ 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,
+            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 +135,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

From b46574a1beed4778963edc3a040f937a25b8b76d Mon Sep 17 00:00:00 2001
From: Samuel <samueldanielyan@gmail.com>
Date: Fri, 2 Oct 2026 20:08:41 +0200
Subject: [PATCH 2/8] Relax minimum relational probe cell size

Lower the minimum required `cell_size` for relational probes from 12 to 6 so tiny smoke-test configurations can run without aborting before training. The probe only needs three role channels, and the previous 4-dimension-per-channel requirement was unnecessarily restrictive for small fixtures.
---
 src/initium/data/graph_traversal/relational_probe.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/initium/data/graph_traversal/relational_probe.py b/src/initium/data/graph_traversal/relational_probe.py
index 3cdb866..b709700 100644
--- a/src/initium/data/graph_traversal/relational_probe.py
+++ b/src/initium/data/graph_traversal/relational_probe.py
@@ -38,8 +38,11 @@ def build_relational_probe(
     """
     if cases < 2:
         raise ValueError("cases must be at least 2 to reserve held-out compositions")
-    if cell_size < 12:
-        raise ValueError("cell_size must be at least 12")
+    # 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

From 0f52e6b5ba7136edf1ea96e3598d5ff2b7795ff6 Mon Sep 17 00:00:00 2001
From: Samuel <samueldanielyan@gmail.com>
Date: Sat, 3 Oct 2026 12:38:01 +0200
Subject: [PATCH 3/8] Add relational residual DNC reader

Add a new `relational-residual` read variant that keeps the DNC cosine baseline at initialization by zero-initializing the learned residual output, and extend split-graph diagnostics with read-content entropy and max-weight metrics. Update the config and CLI/docs to expose the variant and document the Track A/B experiment logs and reader choices.
---
 README.md                                     |  6 +-
 src/initium/config/controller_config.py       |  2 +-
 src/initium/core_training.py                  | 22 ++++-
 .../memory_manipulation/dnc_mechanisms.py     | 89 ++++++++++++++-----
 4 files changed, 90 insertions(+), 29 deletions(-)

diff --git a/README.md b/README.md
index 40d5073..d633aa5 100644
--- a/README.md
+++ b/README.md
@@ -91,11 +91,11 @@ The DNC supplies a memory matrix and differentiable read/write interface. The me
 
 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` learns a score for each `(query, memory row)` pair; write content addressing remains cosine-based.
+- `--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 learned pairwise score to cosine and zero-initializes its final layer, so it starts exactly at the DNC baseline. 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_<id>_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_<id>_track_a.csv` and `run_<id>_track_b.csv`. Track A records ID/OOD accuracy, relational stress and held-out-composition accuracy, plus ordinary-input read-score/cosine correlation. 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.
+- Split-graph graph runs also write `run_<id>_track_a.csv` and `run_<id>_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:
 
@@ -104,7 +104,7 @@ python -m initium.core_training 0.0 --split-graph --split-graph-variant mamba1 -
 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` for Track A or `--dnc-write-variant kanerva-closed-form` for Track B. Add both 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.
+Add `--dnc-read-variant relational-mlp` to reproduce the original Track A reader, or `--dnc-read-variant relational-residual` to evaluate the cosine-anchored Track A reader. 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
 
diff --git a/src/initium/config/controller_config.py b/src/initium/config/controller_config.py
index f0404d8..4d5c0a1 100644
--- a/src/initium/config/controller_config.py
+++ b/src/initium/config/controller_config.py
@@ -68,7 +68,7 @@
 SPLIT_GRAPH_COMBINER_NUM_BLOCKS = 1
 
 # DNC memory mechanism variants (only applied by the split-graph assembly).
-DNC_READ_VARIANT = "cosine"  # "cosine" | "relational-mlp"
+DNC_READ_VARIANT = "cosine"  # "cosine" | "relational-mlp" | "relational-residual"
 DNC_WRITE_VARIANT = "learned"  # "learned" | "kanerva-closed-form"
 DNC_WRITE_OBSERVATION_VARIANCE = 1.0  # fixed isotropic likelihood variance
 
diff --git a/src/initium/core_training.py b/src/initium/core_training.py
index a54b0b8..3be79e2 100644
--- a/src/initium/core_training.py
+++ b/src/initium/core_training.py
@@ -712,6 +712,8 @@ def run(
                     "stress_held_out_acc",
                     "ordinary_read_score_cosine_rho_mean",
                     "ordinary_read_diagnostic_count",
+                    "read_content_entropy_mean",
+                    "read_content_max_weight_mean",
                 ]
             )
         track_b_file = open(
@@ -1981,6 +1983,8 @@ def set_lr(step_now):
                         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"),
                     ]
                 )
                 track_a_file.flush()
@@ -1993,7 +1997,10 @@ def set_lr(step_now):
                     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"{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}"
                 )
             if track_b_writer is not None:
                 id_write_values = [
@@ -2271,6 +2278,8 @@ def set_lr(step_now):
                 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"),
             ]
         )
         track_a_file.flush()
@@ -2280,7 +2289,10 @@ def set_lr(step_now):
             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"{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}"
         )
     if track_b_writer is not None:
         def _write_diagnostic_values(diagnostics):
@@ -2692,10 +2704,12 @@ def _write_diagnostic_values(diagnostics):
     )
     parser.add_argument(
         "--dnc-read-variant",
-        choices=["cosine", "relational-mlp"],
+        choices=["cosine", "relational-mlp", "relational-residual"],
         default=DNC_READ_VARIANT,
         help="Read content score for --split-graph. relational-mlp replaces "
-        "only read similarity; DNC write addressing and temporal links remain intact.",
+        "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-write-variant",
diff --git a/src/initium/memory_manipulation/dnc_mechanisms.py b/src/initium/memory_manipulation/dnc_mechanisms.py
index 6fabc1d..6c20844 100644
--- a/src/initium/memory_manipulation/dnc_mechanisms.py
+++ b/src/initium/memory_manipulation/dnc_mechanisms.py
@@ -13,7 +13,7 @@
 import torch.nn.functional as F
 from dnc.memory import Memory
 
-READ_VARIANTS = ("cosine", "relational-mlp")
+READ_VARIANTS = ("cosine", "relational-mlp", "relational-residual")
 WRITE_VARIANTS = ("learned", "kanerva-closed-form")
 
 
@@ -45,13 +45,19 @@ def __init__(
         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 == "relational-mlp":
+        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.
+                nn.init.zeros_(self.relational_score[-1].weight)
+                nn.init.zeros_(self.relational_score[-1].bias)
             if self.device is not None:
                 self.relational_score.to(self.device)
 
@@ -124,46 +130,87 @@ def erase(self, hidden):
             hidden["write_posterior_variance"].fill_(1.0)
         return hidden
 
-    def read_content_weightings(self, memory, keys, strengths):
-        """Score candidate rows for reads; writes continue using DNC cosine."""
+    def _read_similarity_scores(self, memory, keys):
+        """Return one unscaled read score for every query and memory row."""
         if self.read_variant == "cosine":
-            return self.content_weightings(memory, keys, strengths)
+            return self._dnc_cosine_scores(memory, keys)
         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)
-        logits = self.relational_score(pair).squeeze(-1)
-        return F.softmax(logits * strengths.unsqueeze(-1), 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
+
+    @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():
-                batch, rows, width = hidden["memory"].shape
-                heads = read_keys.size(1)
-                candidates = hidden["memory"][:, None, :, :].expand(
-                    batch, heads, rows, width
-                )
-                queries = read_keys[:, :, None, :].expand(batch, heads, rows, width)
-                cosine = F.cosine_similarity(queries, candidates, dim=-1)
-                relational = self.relational_score(
-                    torch.cat((queries, candidates), dim=-1)
-                ).squeeze(-1)
+                memory = hidden["memory"]
+                cosine = self._dnc_cosine_scores(memory, read_keys)
+                scores = self._read_similarity_scores(memory, read_keys)
                 cosine_centered = cosine - cosine.mean(dim=-1, keepdim=True)
-                relational_centered = relational - relational.mean(dim=-1, keepdim=True)
-                numerator = (cosine_centered * relational_centered).sum(dim=-1)
+                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)
-                    * relational_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
+                    "ordinary_read_score_cosine_correlation",
+                    numerator / denominator,
                 )
         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,

From 41904ea320e55683112c771866f78ab466074db3 Mon Sep 17 00:00:00 2001
From: Samuel <samueldanielyan@gmail.com>
Date: Sat, 3 Oct 2026 16:02:53 +0200
Subject: [PATCH 4/8] Add configurable residual read scaling

This change introduces a configurable `--dnc-read-residual-scale` for the relational-residual DNC read variant, preserving the cosine baseline at initialization while allowing the learned correction to be strengthened or weakened. The setting is threaded through config, training, checkpoint loading, and memory resizing, and diagnostics now report the residual RMS and residual-to-cosine RMS ratio. The README is updated with the new flag and experiment guidance.
---
 README.md                                     |  4 +--
 src/initium/config/controller_config.py       |  1 +
 src/initium/core_training.py                  | 33 +++++++++++++++--
 src/initium/inference/model_loader.py         |  4 +++
 .../mamba_controller/split_graph_dnc.py       |  3 ++
 .../memory_manipulation/dnc_mechanisms.py     | 35 +++++++++++++++----
 .../dynamic_memory_resize.py                  |  1 +
 7 files changed, 71 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index d633aa5..305f273 100644
--- a/README.md
+++ b/README.md
@@ -91,7 +91,7 @@ The DNC supplies a memory matrix and differentiable read/write interface. The me
 
 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 learned pairwise score to cosine and zero-initializes its final layer, so it starts exactly at the DNC baseline. Write content addressing remains cosine-based for both learned readers.
+- `--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 `2.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_<id>_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.
@@ -104,7 +104,7 @@ python -m initium.core_training 0.0 --split-graph --split-graph-variant mamba1 -
 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` to evaluate the cosine-anchored Track A reader. 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.
+Add `--dnc-read-variant relational-mlp` to reproduce the original Track A reader, or `--dnc-read-variant relational-residual --dnc-read-residual-scale 2.0` to evaluate the strengthened cosine-anchored Track A reader. 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
 
diff --git a/src/initium/config/controller_config.py b/src/initium/config/controller_config.py
index 4d5c0a1..409629f 100644
--- a/src/initium/config/controller_config.py
+++ b/src/initium/config/controller_config.py
@@ -69,6 +69,7 @@
 
 # DNC memory mechanism variants (only applied by the split-graph assembly).
 DNC_READ_VARIANT = "cosine"  # "cosine" | "relational-mlp" | "relational-residual"
+DNC_READ_RESIDUAL_SCALE = 2.0  # Multiplier on the learned score correction for relational-residual.
 DNC_WRITE_VARIANT = "learned"  # "learned" | "kanerva-closed-form"
 DNC_WRITE_OBSERVATION_VARIANCE = 1.0  # fixed isotropic likelihood variance
 
diff --git a/src/initium/core_training.py b/src/initium/core_training.py
index 3be79e2..1e1de81 100644
--- a/src/initium/core_training.py
+++ b/src/initium/core_training.py
@@ -70,6 +70,7 @@
     DYNAMIC_N_TRIGGER_FRAC,
     DYNAMIC_N_USAGE_HIGH,
     DNC_READ_VARIANT,
+    DNC_READ_RESIDUAL_SCALE,
     DNC_WRITE_OBSERVATION_VARIANCE,
     DNC_WRITE_VARIANT,
     HYBRID_CFC_NUM_BLOCKS,
@@ -579,6 +580,7 @@ def run(
     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_write_variant: str = DNC_WRITE_VARIANT,
     dnc_write_observation_variance: float = DNC_WRITE_OBSERVATION_VARIANCE,
     dataset_type: str = DATASET_TYPE,
@@ -714,6 +716,8 @@ def run(
                     "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(
@@ -1135,6 +1139,7 @@ def run(
             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_write_variant=dnc_write_variant,
             dnc_write_observation_variance=dnc_write_observation_variance,
             device=device,
@@ -1199,6 +1204,7 @@ def run(
         ),
         "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_write_variant": dnc_write_variant,
         "dnc_write_observation_variance": dnc_write_observation_variance,
     }
@@ -1985,6 +1991,10 @@ def set_lr(step_now):
                         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()
@@ -2000,7 +2010,9 @@ def set_lr(step_now):
                     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"{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:
                 id_write_values = [
@@ -2280,6 +2292,10 @@ def set_lr(step_now):
                 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()
@@ -2292,7 +2308,9 @@ def set_lr(step_now):
             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"{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:
         def _write_diagnostic_values(diagnostics):
@@ -2711,6 +2729,14 @@ def _write_diagnostic_values(diagnostics):
         "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="Multiplier on the learned pairwise correction for relational-residual; "
+        "ignored by other read variants. Default 2.0 strengthens relation evidence "
+        "while preserving the cosine baseline at initialization.",
+    )
     parser.add_argument(
         "--dnc-write-variant",
         choices=["learned", "kanerva-closed-form"],
@@ -2871,6 +2897,8 @@ def _write_diagnostic_values(diagnostics):
             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}"
             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":
@@ -2913,6 +2941,7 @@ def _write_diagnostic_values(diagnostics):
             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_write_variant=args.dnc_write_variant,
             dnc_write_observation_variance=args.dnc_write_observation_variance,
             dataset_type=args.dataset_type,
diff --git a/src/initium/inference/model_loader.py b/src/initium/inference/model_loader.py
index 7d91631..16b7b34 100644
--- a/src/initium/inference/model_loader.py
+++ b/src/initium/inference/model_loader.py
@@ -177,6 +177,10 @@ def _build_split_graph(c: _Cfg, input_dim, hidden, nr_cells, cell_size, read_hea
             "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 were trained at scale 1.
+            "dnc_read_residual_scale", 1.0
+        ),
         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
diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py
index 5238fdc..d86e9ac 100644
--- a/src/initium/mamba_controller/split_graph_dnc.py
+++ b/src/initium/mamba_controller/split_graph_dnc.py
@@ -196,6 +196,7 @@ def __init__(
         # input, for direct A/B comparison.
         moe_cfc_multi_source: bool = True,
         dnc_read_variant: str = "cosine",
+        dnc_read_residual_scale: float = 2.0,
         dnc_write_variant: str = "learned",
         dnc_write_observation_variance: float = 1.0,
     ):
@@ -222,6 +223,7 @@ def __init__(
         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_write_variant = dnc_write_variant
         self.dnc_write_observation_variance = dnc_write_observation_variance
 
@@ -452,6 +454,7 @@ def __init__(
                 independent_linears=independent_linears,
                 read_variant=dnc_read_variant,
                 write_variant=dnc_write_variant,
+                read_residual_scale=dnc_read_residual_scale,
                 observation_variance=dnc_write_observation_variance,
             )
         )
diff --git a/src/initium/memory_manipulation/dnc_mechanisms.py b/src/initium/memory_manipulation/dnc_mechanisms.py
index 6c20844..5d42e8a 100644
--- a/src/initium/memory_manipulation/dnc_mechanisms.py
+++ b/src/initium/memory_manipulation/dnc_mechanisms.py
@@ -25,6 +25,7 @@ def __init__(
         *args,
         read_variant: str = "cosine",
         write_variant: str = "learned",
+        read_residual_scale: float = 2.0,
         relational_hidden_size: int | None = None,
         observation_variance: float = 1.0,
         **kwargs,
@@ -33,11 +34,14 @@ def __init__(
             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 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.observation_variance = float(observation_variance)
         self._collect_stage_diagnostics = False
         self._stage_diagnostic_sums: dict[str, torch.Tensor] = {}
@@ -130,10 +134,10 @@ def erase(self, hidden):
             hidden["write_posterior_variance"].fill_(1.0)
         return hidden
 
-    def _read_similarity_scores(self, memory, keys):
-        """Return one unscaled read score for every query and memory row."""
+    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)
+            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)
@@ -142,8 +146,15 @@ def _read_similarity_scores(self, memory, keys):
         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
+            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:
+            return base_scores + self.read_residual_scale * residual_scores
+        return base_scores
 
     @staticmethod
     def _dnc_cosine_scores(memory, keys):
@@ -185,7 +196,7 @@ def read(self, read_keys, read_strengths, read_modes, hidden):
         if self._collect_stage_diagnostics:
             with torch.no_grad():
                 memory = hidden["memory"]
-                cosine = self._dnc_cosine_scores(memory, read_keys)
+                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)
@@ -198,6 +209,16 @@ def read(self, read_keys, read_strengths, read_modes, hidden):
                     "ordinary_read_score_cosine_correlation",
                     numerator / denominator,
                 )
+                if residual is not None:
+                    cosine_rms = cosine.square().mean(dim=-1).sqrt().clamp_min(1e-12)
+                    residual_rms = 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
         )
@@ -314,6 +335,7 @@ def build_memory(
     independent_linears: bool = True,
     read_variant: str = "cosine",
     write_variant: str = "learned",
+    read_residual_scale: float = 2.0,
     observation_variance: float = 1.0,
 ):
     """Construct a stock-compatible memory for the selected mechanism pair."""
@@ -326,5 +348,6 @@ def build_memory(
         independent_linears=independent_linears,
         read_variant=read_variant,
         write_variant=write_variant,
+        read_residual_scale=read_residual_scale,
         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 63a9f2f..9215398 100644
--- a/src/initium/memory_manipulation/dynamic_memory_resize.py
+++ b/src/initium/memory_manipulation/dynamic_memory_resize.py
@@ -109,6 +109,7 @@ def resize_memory(model, new_nr_cells: int, device=None, layer: int = 0, optimiz
             **memory_args,
             read_variant=old_memory.read_variant,
             write_variant=old_memory.write_variant,
+            read_residual_scale=old_memory.read_residual_scale,
             observation_variance=old_memory.observation_variance,
         ).to(device)
     else:

From 63beadc4799ce5ccc9bfc87489d0604a9708f62c Mon Sep 17 00:00:00 2001
From: Samuel <samueldanielyan@gmail.com>
Date: Sat, 3 Oct 2026 17:15:21 +0200
Subject: [PATCH 5/8] Lower DNC residual read scale to 1.0

This updates the default relational-residual read correction from 2.0 to 1.0 across config and runtime defaults, and clarifies the docs that the residual remains cosine-anchored and exact at initialization. It also notes that lower scales such as 0.5 are available for follow-up comparisons while keeping the baseline and Track A setup consistent.
---
 README.md                                         | 4 ++--
 src/initium/config/controller_config.py           | 2 +-
 src/initium/core_training.py                      | 4 ++--
 src/initium/mamba_controller/split_graph_dnc.py   | 2 +-
 src/initium/memory_manipulation/dnc_mechanisms.py | 4 ++--
 5 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 305f273..f009511 100644
--- a/README.md
+++ b/README.md
@@ -91,7 +91,7 @@ The DNC supplies a memory matrix and differentiable read/write interface. The me
 
 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 `2.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-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_<id>_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.
@@ -104,7 +104,7 @@ python -m initium.core_training 0.0 --split-graph --split-graph-variant mamba1 -
 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 2.0` to evaluate the strengthened cosine-anchored Track A reader. 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.
+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
 
diff --git a/src/initium/config/controller_config.py b/src/initium/config/controller_config.py
index 409629f..7c68522 100644
--- a/src/initium/config/controller_config.py
+++ b/src/initium/config/controller_config.py
@@ -69,7 +69,7 @@
 
 # DNC memory mechanism variants (only applied by the split-graph assembly).
 DNC_READ_VARIANT = "cosine"  # "cosine" | "relational-mlp" | "relational-residual"
-DNC_READ_RESIDUAL_SCALE = 2.0  # Multiplier on the learned score correction for relational-residual.
+DNC_READ_RESIDUAL_SCALE = 1.0  # Multiplier on the learned score correction for relational-residual.
 DNC_WRITE_VARIANT = "learned"  # "learned" | "kanerva-closed-form"
 DNC_WRITE_OBSERVATION_VARIANCE = 1.0  # fixed isotropic likelihood variance
 
diff --git a/src/initium/core_training.py b/src/initium/core_training.py
index 1e1de81..8c877cb 100644
--- a/src/initium/core_training.py
+++ b/src/initium/core_training.py
@@ -2734,8 +2734,8 @@ def _write_diagnostic_values(diagnostics):
         type=float,
         default=DNC_READ_RESIDUAL_SCALE,
         help="Multiplier on the learned pairwise correction for relational-residual; "
-        "ignored by other read variants. Default 2.0 strengthens relation evidence "
-        "while preserving the cosine baseline at initialization.",
+        "ignored by other read variants. Default 1.0 preserves the tested scale; "
+        "the cosine baseline is exact at initialization.",
     )
     parser.add_argument(
         "--dnc-write-variant",
diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py
index d86e9ac..a07844a 100644
--- a/src/initium/mamba_controller/split_graph_dnc.py
+++ b/src/initium/mamba_controller/split_graph_dnc.py
@@ -196,7 +196,7 @@ def __init__(
         # input, for direct A/B comparison.
         moe_cfc_multi_source: bool = True,
         dnc_read_variant: str = "cosine",
-        dnc_read_residual_scale: float = 2.0,
+        dnc_read_residual_scale: float = 1.0,
         dnc_write_variant: str = "learned",
         dnc_write_observation_variance: float = 1.0,
     ):
diff --git a/src/initium/memory_manipulation/dnc_mechanisms.py b/src/initium/memory_manipulation/dnc_mechanisms.py
index 5d42e8a..67b4324 100644
--- a/src/initium/memory_manipulation/dnc_mechanisms.py
+++ b/src/initium/memory_manipulation/dnc_mechanisms.py
@@ -25,7 +25,7 @@ def __init__(
         *args,
         read_variant: str = "cosine",
         write_variant: str = "learned",
-        read_residual_scale: float = 2.0,
+        read_residual_scale: float = 1.0,
         relational_hidden_size: int | None = None,
         observation_variance: float = 1.0,
         **kwargs,
@@ -335,7 +335,7 @@ def build_memory(
     independent_linears: bool = True,
     read_variant: str = "cosine",
     write_variant: str = "learned",
-    read_residual_scale: float = 2.0,
+    read_residual_scale: float = 1.0,
     observation_variance: float = 1.0,
 ):
     """Construct a stock-compatible memory for the selected mechanism pair."""

From f50b0be98ec2bcf151facbe7f4bc5cf4bac762fe Mon Sep 17 00:00:00 2001
From: Samuel <samueldanielyan@gmail.com>
Date: Sun, 4 Oct 2026 12:45:52 +0200
Subject: [PATCH 6/8] Tune DNC residual read gating and Q99 probes

Adds a configurable cap for relational DNC residual read scores, exposed through config and CLI, and threads it through model loading and memory resizing. The Q99 order-sensitivity check is also refactored into separate fact-order and query-order probes with clearer semantic metrics and post-training gating. This keeps the residual correction bounded while making the order-sensitivity diagnostics more interpretable and robust. Includes a small .gitignore update for tmp files.
---
 .gitignore                                    |   4 +-
 src/initium/config/controller_config.py       |   3 +-
 src/initium/core_training.py                  | 156 ++++++-----
 .../data/graph_traversal/order_probe.py       | 264 ++++++++++--------
 src/initium/inference/model_loader.py         |   3 +-
 .../mamba_controller/split_graph_dnc.py       |   3 +
 .../memory_manipulation/dnc_mechanisms.py     |  24 +-
 .../dynamic_memory_resize.py                  |   1 +
 8 files changed, 276 insertions(+), 182 deletions(-)

diff --git a/.gitignore b/.gitignore
index f185c83..a9854f8 100644
--- a/.gitignore
+++ b/.gitignore
@@ -219,5 +219,5 @@ __marimo__/
 
 # Experimental documents for context, not to be committed
 documents/*
-
-dnc-context/*
\ No newline at end of file
+dnc-context/*
+tmp/*
\ No newline at end of file
diff --git a/src/initium/config/controller_config.py b/src/initium/config/controller_config.py
index 7c68522..5da51c8 100644
--- a/src/initium/config/controller_config.py
+++ b/src/initium/config/controller_config.py
@@ -69,7 +69,8 @@
 
 # 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  # Multiplier on the learned score correction for 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
 
diff --git a/src/initium/core_training.py b/src/initium/core_training.py
index 8c877cb..907171d 100644
--- a/src/initium/core_training.py
+++ b/src/initium/core_training.py
@@ -71,6 +71,7 @@
     DYNAMIC_N_USAGE_HIGH,
     DNC_READ_VARIANT,
     DNC_READ_RESIDUAL_SCALE,
+    DNC_READ_RESIDUAL_MAX_RATIO,
     DNC_WRITE_OBSERVATION_VARIANCE,
     DNC_WRITE_VARIANT,
     HYBRID_CFC_NUM_BLOCKS,
@@ -581,6 +582,7 @@ def run(
     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,
@@ -738,12 +740,16 @@ def run(
                     "ood_offset_triple",
                     "q99_gate",
                     "q99_tested_episodes",
-                    "q99_outputs_change_with_memory",
-                    "q99_outputs_change_without_memory",
-                    "q99_max_abs_output_delta",
-                    "q99_mean_abs_output_delta",
-                    "q99_prediction_flip_rate",
-                    "q99_memory_incremental_order_effect",
+                    "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",
@@ -1140,6 +1146,7 @@ def run(
             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,
@@ -1205,6 +1212,7 @@ def run(
         "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,
     }
@@ -1465,58 +1473,6 @@ def set_lr(step_now):
                 f"want to extend further."
             )
 
-    q99_metrics = None
-    if split_graph_enabled and dnc_write_variant == "kanerva-closed-form":
-        python_rng_state = random.getstate()
-        numpy_rng_state = np.random.get_state()
-        torch_rng_state = _capture_torch_rng()
-        try:
-            with _replay_torch_rng(torch_rng_state):
-                q99_inputs, _, q99_answer_mask = build_q99_probe_batch()
-                q99_metrics = evaluate_q99_order_sensitivity(
-                    rnn, output_proj, q99_inputs, q99_answer_mask
-                )
-        finally:
-            random.setstate(python_rng_state)
-            np.random.set_state(numpy_rng_state)
-
-        print(
-            f"[{run_id}] [Track B / Q99] gate={q99_metrics['gate']} | "
-            f"episodes={q99_metrics['tested_episodes']} | "
-            f"output_delta={q99_metrics.get('mean_abs_output_delta', 0.0):.6g} | "
-            f"prediction_flip={q99_metrics.get('prediction_flip_rate', 0.0):.4f} | "
-            f"memory_order_effect="
-            f"{q99_metrics.get('memory_incremental_order_effect', 0.0):.6g}"
-        )
-        if track_b_writer is not None:
-            track_b_writer.writerow(
-                [
-                    "q99_gate",
-                    step,
-                    curriculum.lesson + 1,
-                    dnc_write_variant,
-                    None,
-                    None,
-                    None,
-                    q99_metrics["gate"],
-                    q99_metrics["tested_episodes"],
-                    q99_metrics["outputs_change_with_memory"],
-                    q99_metrics["outputs_change_without_memory"],
-                    q99_metrics.get("max_abs_output_delta"),
-                    q99_metrics.get("mean_abs_output_delta"),
-                    q99_metrics.get("prediction_flip_rate"),
-                    q99_metrics.get("memory_incremental_order_effect"),
-                    *([None] * 16),
-                    *([None] * 6),
-                ]
-            )
-            track_b_file.flush()
-        if q99_metrics["gate"] != "pass":
-            raise RuntimeError(
-                "Track B Q99 order-sensitivity gate failed before training; "
-                "see the run's _track_b.csv q99_gate row."
-            )
-
     print(
         f"\n=== [{run_id}] Training (beta_target={beta_target}, beta_mode={beta_mode}) "
         f"{'[resumed]' if resuming else ''} ==="
@@ -1751,7 +1707,7 @@ def set_lr(step_now):
                         None,
                         None,
                         None,
-                        *([None] * 24),
+                        *([None] * 28),
                         avg_task,
                         avg_kl,
                         train_total_loss,
@@ -2044,7 +2000,7 @@ def set_lr(step_now):
                         id_triple_acc,
                         ood_triple_acc,
                         id_triple_acc - ood_triple_acc,
-                        *([None] * 8),
+                        *([None] * 12),
                         *id_write_values,
                         *ood_write_values,
                         *([None] * 6),
@@ -2334,7 +2290,7 @@ def _write_diagnostic_values(diagnostics):
                 id_triple_acc,
                 ood_triple_acc,
                 ood_offset_triple,
-                *([None] * 8),
+                *([None] * 12),
                 *_write_diagnostic_values(final_id_stage_diagnostics),
                 *_write_diagnostic_values(final_ood_stage_diagnostics),
                 *([None] * 6),
@@ -2372,6 +2328,58 @@ def _write_diagnostic_values(diagnostics):
             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:
+            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,
@@ -2413,6 +2421,14 @@ def _write_diagnostic_values(diagnostics):
         "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([])
@@ -2733,9 +2749,15 @@ def _write_diagnostic_values(diagnostics):
         "--dnc-read-residual-scale",
         type=float,
         default=DNC_READ_RESIDUAL_SCALE,
-        help="Multiplier on the learned pairwise correction for relational-residual; "
-        "ignored by other read variants. Default 1.0 preserves the tested scale; "
-        "the cosine baseline is exact at initialization.",
+        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",
@@ -2817,6 +2839,10 @@ def _write_diagnostic_values(diagnostics):
         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"
     ):
@@ -2899,6 +2925,7 @@ def _write_diagnostic_values(diagnostics):
                 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":
@@ -2942,6 +2969,7 @@ def _write_diagnostic_values(diagnostics):
             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,
diff --git a/src/initium/data/graph_traversal/order_probe.py b/src/initium/data/graph_traversal/order_probe.py
index abd3971..479c57d 100644
--- a/src/initium/data/graph_traversal/order_probe.py
+++ b/src/initium/data/graph_traversal/order_probe.py
@@ -1,116 +1,122 @@
-"""Utilities for the Q99 online-write order-sensitivity check."""
+"""Semantic ordering probes for online memory writes (Q99)."""
 
 from __future__ import annotations
 
+import random
+
+import numpy as np
 import torch
 
 
-def build_q99_probe_batch(batch_size: int = 8):
-    """Build graph episodes with at least two path writes for Q99 pairing."""
+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,
     )
 
-    episodes = []
-    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)
+    python_state = random.getstate()
+    numpy_state = np.random.get_state()
+    try:
+        random.seed(seed)
+        np.random.seed(seed)
+        episodes = []
+        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 permute_early_writes(
-    inputs: torch.Tensor,
-    *,
-    prefix_length: int,
-    permutation: tuple[int, ...],
-) -> torch.Tensor:
-    """Return a copy with only the first ``prefix_length`` events reordered.
-
-    Inputs use the dataset's batch-first ``(batch, time, features)`` layout.
-    Queries and later sequence steps remain in their original positions.
-    ``permutation`` contains indices relative to the prefix and must be a
-    bijection, so callers can provide a deterministic paired episode.
-    """
-    if inputs.ndim != 3:
-        raise ValueError("inputs must have shape (batch, time, features)")
-    if not 0 < prefix_length <= inputs.size(1):
-        raise ValueError("prefix_length must be within the input sequence")
-    if len(permutation) != prefix_length or sorted(permutation) != list(range(prefix_length)):
-        raise ValueError("permutation must be a bijection over the write prefix")
+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()
-    indices = torch.tensor(permutation, device=inputs.device)
-    result[:, :prefix_length] = inputs[:, :prefix_length].index_select(1, indices)
-    return result
+    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 permute_early_graph_query_writes(
+def swap_first_two_path_query_writes(
     inputs: torch.Tensor, *, label_dim: int = 30
 ) -> tuple[torch.Tensor, int]:
-    """Swap the first two path-query writes in each graph episode.
-
-    Graph facts have a populated destination field; path-query writes have a
-    source and/or edge field, no destination, and prediction_required=0. The
-    answer prompts and graph-fact write order are left untouched.
-    """
+    """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)):
-        episode = inputs[batch_index]
-        source_present = episode[:, :label_dim].abs().sum(dim=-1) > 0
-        edge_present = episode[:, label_dim : 2 * label_dim].abs().sum(dim=-1) > 0
-        destination_present = episode[:, 2 * label_dim : 3 * label_dim].abs().sum(dim=-1) > 0
-        prediction_required = episode[:, -1] > 0.5
-        query_rows = torch.nonzero(
-            ~destination_present
-            & ~prediction_required
-            & (source_present | edge_present),
-            as_tuple=False,
-        ).flatten()
-        if query_rows.numel() < 2:
+        indices = _path_query_indices(inputs[batch_index], label_dim)
+        if indices.numel() < 2:
             continue
-        first, second = query_rows[:2]
+        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 order_sensitivity_metrics(
-    ordered_output: torch.Tensor,
-    permuted_output: torch.Tensor,
-    *,
-    atol: float = 1e-6,
-    mask: torch.Tensor | None = None,
-) -> dict[str, float | bool]:
-    """Summarize whether task outputs respond to a prefix-order change."""
-    if ordered_output.shape != permuted_output.shape:
-        raise ValueError("paired outputs must have identical shapes")
-    if mask is not None:
-        if mask.shape != ordered_output.shape[: mask.ndim]:
-            raise ValueError("mask dimensions must match leading output dimensions")
-        selected = mask.to(device=ordered_output.device, dtype=torch.bool)
-        ordered_output = ordered_output[selected]
-        permuted_output = permuted_output[selected]
-        if ordered_output.numel() == 0:
-            raise ValueError("mask selects no output positions")
-    difference = (ordered_output - permuted_output).abs()
-    if ordered_output.ndim >= 2:
-        changed = ordered_output.argmax(dim=-1) != permuted_output.argmax(dim=-1)
-        prediction_flip_rate = changed.float().mean().item()
-    else:
-        prediction_flip_rate = float(
-            (ordered_output.round() != permuted_output.round()).float().mean().item()
-        )
-    return {
-        "outputs_change": bool((difference > atol).any().item()),
-        "max_absolute_difference": float(difference.max().item()),
-        "mean_absolute_difference": float(difference.mean().item()),
-        "prediction_flip_rate": prediction_flip_rate,
-    }
+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()
@@ -118,20 +124,34 @@ 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]:
-    """Compare paired-order task outputs with and without external memory."""
-    permuted, tested_episodes = permute_early_graph_query_writes(inputs)
-    if tested_episodes == 0:
+    """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": 0,
-            "outputs_change_with_memory": False,
-            "outputs_change_without_memory": False,
-            "memory_incremental_order_effect": 0.0,
-            "prediction_flip_rate": 0.0,
+            "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
@@ -151,35 +171,55 @@ def forward(sequence: torch.Tensor, pass_through_memory: bool) -> torch.Tensor:
         )
 
     try:
-        inputs = inputs.to(next(model.parameters()).device)
-        permuted = permuted.to(inputs.device)
-        answer_mask = answer_mask.to(device=inputs.device, dtype=torch.bool)
+        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)
-        permuted_full = forward(permuted, True)
+        fact_permuted_full = forward(fact_permuted, True)
+        query_swapped_full = forward(query_swapped, True)
         ordered_ablate = forward(inputs, False)
-        permuted_ablate = forward(permuted, False)
+        query_swapped_ablate = forward(query_swapped, False)
     finally:
         model.train(was_training)
         output_projection.train(projection_was_training)
 
-    full_metrics = order_sensitivity_metrics(
-        ordered_full, permuted_full, atol=atol, mask=answer_mask
-    )
-    ablated_metrics = order_sensitivity_metrics(
-        ordered_ablate, permuted_ablate, atol=atol, mask=answer_mask
+    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
     )
-    delta_difference = (
-        (ordered_full - permuted_full) - (ordered_ablate - permuted_ablate)
+    memory_delta = (
+        (ordered_full - query_swapped_full)
+        - (ordered_ablate - query_swapped_ablate)
     ).abs()[answer_mask]
-    memory_effect = float(delta_difference.mean().item()) if delta_difference.numel() else 0.0
-    gate_pass = bool(full_metrics["outputs_change"]) and memory_effect > atol
+    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": tested_episodes,
-        "outputs_change_with_memory": bool(full_metrics["outputs_change"]),
-        "outputs_change_without_memory": bool(ablated_metrics["outputs_change"]),
-        "max_abs_output_delta": float(full_metrics["max_absolute_difference"]),
-        "mean_abs_output_delta": float(full_metrics["mean_absolute_difference"]),
-        "prediction_flip_rate": float(full_metrics["prediction_flip_rate"]),
-        "memory_incremental_order_effect": memory_effect,
+        "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/inference/model_loader.py b/src/initium/inference/model_loader.py
index 16b7b34..f1a18e6 100644
--- a/src/initium/inference/model_loader.py
+++ b/src/initium/inference/model_loader.py
@@ -178,9 +178,10 @@ def _build_split_graph(c: _Cfg, input_dim, hidden, nr_cells, cell_size, read_hea
         ),
         dnc_read_variant=c.get("dnc_read_variant", cc.DNC_READ_VARIANT),
         dnc_read_residual_scale=c.get(
-            # Older relational-residual checkpoints were trained at scale 1.
+            # 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
diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py
index a07844a..8e218ac 100644
--- a/src/initium/mamba_controller/split_graph_dnc.py
+++ b/src/initium/mamba_controller/split_graph_dnc.py
@@ -197,6 +197,7 @@ def __init__(
         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,
     ):
@@ -224,6 +225,7 @@ def __init__(
         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
 
@@ -455,6 +457,7 @@ def __init__(
                 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,
             )
         )
diff --git a/src/initium/memory_manipulation/dnc_mechanisms.py b/src/initium/memory_manipulation/dnc_mechanisms.py
index 67b4324..a54d8cd 100644
--- a/src/initium/memory_manipulation/dnc_mechanisms.py
+++ b/src/initium/memory_manipulation/dnc_mechanisms.py
@@ -26,6 +26,7 @@ def __init__(
         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,
@@ -36,12 +37,15 @@ def __init__(
             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] = {}
@@ -153,7 +157,20 @@ 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:
-            return base_scores + self.read_residual_scale * residual_scores
+            residual_scores = residual_scores - residual_scores.mean(
+                dim=-1, keepdim=True
+            )
+            residual_rms = residual_scores.square().mean(
+                dim=-1, keepdim=True
+            ).sqrt()
+            normalized_residual = residual_scores / residual_rms.clamp_min(1e-6)
+            cosine_rms = base_scores.square().mean(dim=-1, keepdim=True).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
@@ -210,8 +227,9 @@ def read(self, read_keys, read_strengths, read_modes, hidden):
                     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 = residual.square().mean(dim=-1).sqrt()
+                    residual_rms = effective_residual.square().mean(dim=-1).sqrt()
                     self._record_stage_value(
                         "relational_read_residual_rms", residual_rms
                     )
@@ -336,6 +354,7 @@ def build_memory(
     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."""
@@ -349,5 +368,6 @@ def build_memory(
         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 9215398..98256b3 100644
--- a/src/initium/memory_manipulation/dynamic_memory_resize.py
+++ b/src/initium/memory_manipulation/dynamic_memory_resize.py
@@ -110,6 +110,7 @@ def resize_memory(model, new_nr_cells: int, device=None, layer: int = 0, optimiz
             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:

From bf800b33094e6f740b0f25fcc4ef3061e3db9a36 Mon Sep 17 00:00:00 2001
From: Samuel <samueldanielyan@gmail.com>
Date: Sun, 4 Oct 2026 13:28:22 +0200
Subject: [PATCH 7/8] Stabilize DNC residual normalization

Avoid zero-valued RMS denominators in residual score normalization during the first forward pass. This change adds a small epsilon before taking the square root for both residual and cosine RMS values, preventing NaNs and infinite gradients when the scorer is initialized to zero.
---
 src/initium/memory_manipulation/dnc_mechanisms.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/initium/memory_manipulation/dnc_mechanisms.py b/src/initium/memory_manipulation/dnc_mechanisms.py
index a54d8cd..d789f11 100644
--- a/src/initium/memory_manipulation/dnc_mechanisms.py
+++ b/src/initium/memory_manipulation/dnc_mechanisms.py
@@ -160,11 +160,16 @@ def _read_similarity_scores(self, memory, keys):
             residual_scores = residual_scores - residual_scores.mean(
                 dim=-1, keepdim=True
             )
-            residual_rms = residual_scores.square().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()
-            normalized_residual = residual_scores / residual_rms.clamp_min(1e-6)
-            cosine_rms = base_scores.square().mean(dim=-1, keepdim=True).sqrt()
             bounded_residual = (
                 self.read_residual_max_ratio
                 * cosine_rms

From 638431af3470fb13c028e0099fbcd64745d8709e Mon Sep 17 00:00:00 2001
From: Samuel <samueldanielyan@gmail.com>
Date: Tue, 6 Oct 2026 18:42:36 +0200
Subject: [PATCH 8/8] formatted and linted

---
 src/__init__.py                               |  2 +-
 src/initium/core_training.py                  | 77 +++++++++++--------
 .../data/graph_traversal/order_probe.py       | 14 +---
 .../data/graph_traversal/relational_probe.py  |  4 +-
 src/initium/inference/model_loader.py         |  3 +-
 .../mamba_controller/split_graph_dnc.py       |  2 +-
 .../memory_manipulation/dnc_mechanisms.py     | 73 ++++++------------
 7 files changed, 78 insertions(+), 97 deletions(-)

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/core_training.py b/src/initium/core_training.py
index 907171d..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,
@@ -69,11 +74,6 @@
     DYNAMIC_N_MODE,
     DYNAMIC_N_TRIGGER_FRAC,
     DYNAMIC_N_USAGE_HIGH,
-    DNC_READ_VARIANT,
-    DNC_READ_RESIDUAL_SCALE,
-    DNC_READ_RESIDUAL_MAX_RATIO,
-    DNC_WRITE_OBSERVATION_VARIANCE,
-    DNC_WRITE_VARIANT,
     HYBRID_CFC_NUM_BLOCKS,
     ISOLATE_LINK_ABLATION,
     LINK_MATRIX_MODE,
@@ -144,15 +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.graph_traversal.order_probe import (
-    build_q99_probe_batch,
-    evaluate_q99_order_sensitivity,
-)
 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
@@ -666,9 +666,7 @@ def run(
         )
 
     relational_probe_cases = (
-        build_relational_probe(cell_size=MODEL_CELL_SIZE)
-        if dataset_type == "graph"
-        else None
+        build_relational_probe(cell_size=MODEL_CELL_SIZE) if dataset_type == "graph" else None
     )
     relational_probe_file = None
     relational_probe_writer = None
@@ -1696,7 +1694,11 @@ def set_lr(step_now):
                 parts.append(dyn_n_str.strip(" |"))
             print(" | ".join(parts))
 
-            if track_b_writer is not None and dnc_write_variant == "kanerva-closed-form":
+            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(
                     [
@@ -1813,7 +1815,11 @@ def set_lr(step_now):
                 if hasattr(id_stage_memory, "finish_stage_diagnostics")
                 else {}
             )
-            if relational_probe_writer is not None:
+            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(
                     [
@@ -1924,7 +1930,11 @@ def set_lr(step_now):
                 ]
             )
             ood_log_file.flush()
-            if track_a_writer is not None:
+            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,
@@ -1948,9 +1958,7 @@ def set_lr(step_now):
                         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"
-                        ),
+                        id_stage_diagnostics.get("relational_read_residual_to_cosine_rms_mean"),
                     ]
                 )
                 track_a_file.flush()
@@ -1970,7 +1978,7 @@ def set_lr(step_now):
                     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:
+            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"),
@@ -2188,7 +2196,11 @@ def set_lr(step_now):
         if hasattr(final_id_stage_memory, "finish_stage_diagnostics")
         else {}
     )
-    if relational_probe_writer is not None:
+    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(
             [
@@ -2227,7 +2239,11 @@ def set_lr(step_now):
     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:
+    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",
@@ -2249,9 +2265,7 @@ def set_lr(step_now):
                 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"
-                ),
+                final_id_stage_diagnostics.get("relational_read_residual_to_cosine_rms_mean"),
             ]
         )
         track_a_file.flush()
@@ -2268,7 +2282,8 @@ def set_lr(step_now):
             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:
+    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"),
@@ -2353,7 +2368,7 @@ def _write_diagnostic_values(diagnostics):
             f"query memory effect="
             f"{q99_metrics['query_memory_incremental_order_effect']:.6g}"
         )
-        if track_b_writer is not None:
+        if track_b_writer is not None and track_b_file is not None:
             track_b_writer.writerow(
                 [
                     "q99_semantic_order_check",
@@ -2423,12 +2438,8 @@ def _write_diagnostic_values(diagnostics):
     }
     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"
-        ]
+        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([])
@@ -2696,7 +2707,7 @@ def _write_diagnostic_values(diagnostics):
         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).",
diff --git a/src/initium/data/graph_traversal/order_probe.py b/src/initium/data/graph_traversal/order_probe.py
index 479c57d..8962a30 100644
--- a/src/initium/data/graph_traversal/order_probe.py
+++ b/src/initium/data/graph_traversal/order_probe.py
@@ -7,7 +7,6 @@
 import numpy as np
 import torch
 
-
 Q99_PROBE_EPISODES = 128
 Q99_MIN_ORDERED_QUERY_ACCURACY = 50.0
 Q99_MAX_FACT_PERMUTATION_DROP = 5.0
@@ -26,7 +25,7 @@ def build_q99_probe_batch(batch_size: int = Q99_PROBE_EPISODES, seed: int = 991)
     try:
         random.seed(seed)
         np.random.seed(seed)
-        episodes = []
+        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:
@@ -191,21 +190,16 @@ def forward(sequence: torch.Tensor, pass_through_memory: bool) -> torch.Tensor:
     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
-    )
+    query_flip = _answer_prediction_retention(ordered_full, query_swapped_full, answer_mask)
     memory_delta = (
-        (ordered_full - query_swapped_full)
-        - (ordered_ablate - query_swapped_ablate)
+        (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
-    )
+    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",
diff --git a/src/initium/data/graph_traversal/relational_probe.py b/src/initium/data/graph_traversal/relational_probe.py
index b709700..fbfd6ef 100644
--- a/src/initium/data/graph_traversal/relational_probe.py
+++ b/src/initium/data/graph_traversal/relational_probe.py
@@ -89,9 +89,7 @@ def cosine_top1_accuracy(cases: list[RelationalProbeCase]) -> float:
         return 0.0
     correct = 0
     for case in cases:
-        scores = torch.nn.functional.cosine_similarity(
-            case.memory, case.query.unsqueeze(0), dim=-1
-        )
+        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)
 
diff --git a/src/initium/inference/model_loader.py b/src/initium/inference/model_loader.py
index f1a18e6..0b68989 100644
--- a/src/initium/inference/model_loader.py
+++ b/src/initium/inference/model_loader.py
@@ -179,7 +179,8 @@ def _build_split_graph(c: _Cfg, input_dim, hidden, nr_cells, cell_size, read_hea
         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_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),
diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py
index 8e218ac..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 initium.memory_manipulation.dnc_mechanisms import build_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):
diff --git a/src/initium/memory_manipulation/dnc_mechanisms.py b/src/initium/memory_manipulation/dnc_mechanisms.py
index d789f11..5a7571a 100644
--- a/src/initium/memory_manipulation/dnc_mechanisms.py
+++ b/src/initium/memory_manipulation/dnc_mechanisms.py
@@ -64,8 +64,12 @@ def __init__(
                 # 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.
-                nn.init.zeros_(self.relational_score[-1].weight)
-                nn.init.zeros_(self.relational_score[-1].bias)
+                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)
 
@@ -102,14 +106,10 @@ def finish_stage_diagnostics(self) -> dict[str, float]:
         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}_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)
-        )
+        result["diagnostic_steps"] = float(max(self._stage_diagnostic_counts.values(), default=0))
         return result
 
     def new(self, batch_size: int = 1):
@@ -118,7 +118,9 @@ def new(self, batch_size: int = 1):
             # 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,
+                batch_size,
+                self.nr_cells,
+                1,
                 device=hidden["memory"].device,
                 dtype=hidden["memory"].dtype,
             )
@@ -127,9 +129,7 @@ def new(self, batch_size: int = 1):
     def clone(self, hidden):
         cloned = super().clone(hidden)
         if self.write_variant == "kanerva-closed-form":
-            cloned["write_posterior_variance"] = hidden[
-                "write_posterior_variance"
-            ].clone()
+            cloned["write_posterior_variance"] = hidden["write_posterior_variance"].clone()
         return cloned
 
     def erase(self, hidden):
@@ -157,19 +157,13 @@ 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
-            )
+            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()
+            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()
+            cosine_rms = (base_scores.square().mean(dim=-1, keepdim=True) + 1e-12).sqrt()
             bounded_residual = (
                 self.read_residual_max_ratio
                 * cosine_rms
@@ -203,8 +197,7 @@ def read(self, read_keys, read_strengths, read_modes, hidden):
                         hidden["memory"], read_keys, read_strengths
                     )
                     entropy = -(
-                        content_weights.clamp_min(1e-12)
-                        * content_weights.clamp_min(1e-12).log()
+                        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(
@@ -224,8 +217,7 @@ def read(self, read_keys, read_strengths, read_modes, hidden):
                 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)
+                    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",
@@ -235,21 +227,16 @@ def read(self, read_keys, read_strengths, read_modes, hidden):
                     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_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
-        )
+        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()
+                    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(
@@ -297,12 +284,8 @@ def write(
             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
-        )
+        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,
@@ -322,13 +305,9 @@ def write(
         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"]
-        )
+        posterior_memory = old_memory + gain * (observation - hidden["memory"])
         hidden["memory"] = posterior_memory
-        hidden["write_posterior_variance"] = variance * noise / (
-            noise + weight * variance
-        )
+        hidden["write_posterior_variance"] = variance * noise / (noise + weight * variance)
         if self._collect_stage_diagnostics:
             posterior_variance = hidden["write_posterior_variance"]
             self._record_stage_value(
@@ -342,9 +321,7 @@ def write(
         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"]
-        )
+        hidden["precedence"] = self.update_precedence(hidden["precedence"], hidden["write_weights"])
         return hidden