Upload CEDL research checkpoint

HF_UPLOAD_CHECKLIST.md

@@ -0,0 +1,56 @@
+# Hugging Face Upload Checklist
+
+Repository:
+
+- `Jasonjiao2023/CEDL`
+
+The current Hugging Face docs recommend the modern `hf` CLI. For agent-aware local workflows, the optional project setup is:
+
+```bash
+hf skills add
+```
+
+For normal upload work, authenticate first:
+
+```bash
+hf auth login
+```
+
+Before upload, `hf_model/` should contain:
+
+- `README.md`
+- `config.json`
+- `MANIFEST.json`
+- `requirements.txt`
+- `CEDL.py`
+- `data_v4c_pairs.py`
+- `examples/load_checkpoint.py`
+- `probes/`
+- `pytorch_model.bin`
+- `cedl_config.json`
+
+Copy the frozen Colab artifacts into public Hugging Face names:
+
+```bash
+cp "$CEDL_FROZEN_CHECKPOINT" pytorch_model.bin
+cp "$CEDL_FROZEN_CONFIG" cedl_config.json
+```
+
+Verify hashes:
+
+```bash
+sha256sum pytorch_model.bin cedl_config.json
+```
+
+Expected:
+
+```text
+e15a3d6dd38a1f85e2a9c4de409ac3cbd7dece4cf202fca05872ebea8180bfbf  pytorch_model.bin
+52ffddf6d95fa20fe4be953d463e5e388c2b161b49096931e10fe09a160f521f  cedl_config.json
+```
+
+One-line upload after the checkpoint and sidecar are present:
+
+```bash
+hf upload Jasonjiao2023/CEDL CEDL-release/hf_model . --repo-type model
+```

MANIFEST.json

@@ -0,0 +1,34 @@
+{
+  "label": "CEDL",
+  "frozen_at_utc": "2026-06-26T12:42:59.251253+00:00",
+  "checkpoint": "pytorch_model.bin",
+  "config": "cedl_config.json",
+  "sha256": {
+    "checkpoint": "e15a3d6dd38a1f85e2a9c4de409ac3cbd7dece4cf202fca05872ebea8180bfbf",
+    "config": "948eebfc64e66a3622157be781929f1e719b4c8e3f59c4d7dbd116cc2ff1a800"
+  },
+  "model_interpretation": {
+    "architecture": "hippocampal-inspired contextual memory language model",
+    "memory_readout": "contextual direct memory readout",
+    "status": "research checkpoint",
+    "claim": "CEDL contextual memory readout is causally useful for current-vs-stale memory updating; not yet a general downstream SOTA model."
+  },
+  "mechanism_evidence": {
+    "current_vs_stale_causality": [
+      {"seed": 1, "N": 158, "active_top1": 0.791, "memory_off_top1": 0.0, "delta_pp": 79.1, "delta_margin": 8.266, "t": 20.756, "cohen_d": 1.656},
+      {"seed": 2, "N": 164, "active_top1": 0.732, "memory_off_top1": 0.0, "delta_pp": 73.2, "delta_margin": 7.702, "t": 18.912, "cohen_d": 1.481},
+      {"seed": 3, "N": 162, "active_top1": 0.728, "memory_off_top1": 0.0, "delta_pp": 72.8, "delta_margin": 7.756, "t": 18.507, "cohen_d": 1.459}
+    ],
+    "memory_readout_diagnosis": [
+      {"seed": 1, "N": 158, "trunk_top1_cur": 0.0, "memory_top1_cur": 0.797, "active_top1_cur": 0.791, "memory_top1_stale": 0.006, "current_rank_median": 1},
+      {"seed": 2, "N": 164, "trunk_top1_cur": 0.0, "memory_top1_cur": 0.738, "active_top1_cur": 0.732, "memory_top1_stale": 0.006, "current_rank_median": 1},
+      {"seed": 3, "N": 162, "trunk_top1_cur": 0.0, "memory_top1_cur": 0.735, "active_top1_cur": 0.728, "memory_top1_stale": 0.012, "current_rank_median": 1}
+    ]
+  },
+  "general_eval": {
+    "wikitext103_test_ppl": 30.11,
+    "lambada": 0.159,
+    "hellaswag": 0.271,
+    "arc_easy": 0.281
+  }
+}

README.md

@@ -0,0 +1,122 @@
+---
+license: apache-2.0
+library_name: pytorch
+tags:
+  - pytorch
+  - language-modeling
+  - memory
+  - hippocampal-inspired
+  - research-artifact
+datasets:
+  - Salesforce/wikitext
+---
+
+# CEDL
+
+CEDL is a hippocampal-inspired language model with contextual memory readout. It separates context encoding, pattern separation, memory retrieval, and linkage/feedback into distinct computational stages, then uses a contextual memory channel to resolve current-vs-stale information.
+
+This repository is a research checkpoint. It is intended for mechanism inspection and reproducible probes, not for production deployment or broad downstream SOTA claims.
+
+## Intended Use
+
+Use this repository to:
+
+- load the CEDL PyTorch checkpoint;
+- inspect the CEDL architecture implementation;
+- reproduce current-vs-stale memory-update probes;
+- study contextual memory readout and memory-off ablations.
+
+Do not use this repository to claim:
+
+- broad downstream SOTA;
+- production-ready generation quality;
+- biological fidelity beyond computational analogy.
+
+## Files
+
+Expected checkpoint files:
+
+- `pytorch_model.bin`
+- `cedl_config.json`
+- `MANIFEST.json`
+
+Code and probes:
+
+- `CEDL.py`
+- `data_v4c_pairs.py`
+- `probes/probe_memory_causality.py`
+- `probes/probe_memory_diagnosis.py`
+- `probes/probe_memory_source_readout.py`
+
+## Evidence Summary
+
+Current-vs-stale causality:
+
+| Seed | N | CEDL active top-1 | Memory-off top-1 | Delta acc. | Delta margin | Cohen's d |
+|---:|---:|---:|---:|---:|---:|---:|
+| 1 | 158 | 79.1% | 0.0% | +79.1pp | +8.266 | 1.656 |
+| 2 | 164 | 73.2% | 0.0% | +73.2pp | +7.702 | 1.481 |
+| 3 | 162 | 72.8% | 0.0% | +72.8pp | +7.756 | 1.459 |
+
+Memory-readout diagnosis:
+
+| Seed | N | Trunk current top-1 | Memory current top-1 | Memory stale top-1 | Current median rank |
+|---:|---:|---:|---:|---:|---:|
+| 1 | 158 | 0.0% | 79.7% | 0.6% | 1 |
+| 2 | 164 | 0.0% | 73.8% | 0.6% | 1 |
+| 3 | 162 | 0.0% | 73.5% | 1.2% | 1 |
+
+General evaluation:
+
+| Model | Params | WikiText-103 PPL | LAMBADA | HellaSwag | ARC-Easy |
+|---|---:|---:|---:|---:|---:|
+| CEDL | 136.3M | 30.11 | 15.9% | 27.1% | 28.1% |
+| Transformer baseline | 103.2M | 45.93 | 3.9% | 26.5% | 28.2% |
+
+## About the Author
+
+CEDL was developed by Dian Jiao, a multidisciplinary researcher and technology leader with more than 15 years of experience across digital transformation, neuroscience, and technology innovation. His work connects brain-inspired theory with practical AI systems, with a focus on NeuroAI, digital therapeutics, and human-centered model design.
+
+His research interests include AI-enhanced digital therapeutics using EEG, adaptive biofeedback, and machine learning for cognitive enhancement and neurological rehabilitation; and biologically inspired AI models that draw from neural dynamics, memory circuits, and brain mechanisms.
+
+Selected publications and contributions:
+
+- Jiao, D. (2025). Elliptic cortical networks: A mathematically constrained architecture for biologically-inspired intelligence. Neurocomputing, 658, 131802. doi: 10.1016/j.neucom.2025.131802
+- Jiao, D. (2025). Leveraging neurotechnology for neurodivergent education: a narrative review. Learning: Research and Practice, 1-25. doi: 10.1080/23735082.2025.2517052
+- Jiao, D. (2025). AI-enhanced digital therapeutics for cognitive impairment: Integrating mobile applications, virtual reality, and wearable devices. Discover Artificial Intelligence, 5, Article 69. doi: 10.1007/s44163-025-00325-6
+- Jiao, D. (2025). Advancing personalized digital therapeutics: integrating music therapy, brainwave entrainment methods, and AI-driven biofeedback. Frontiers in Digital Health, 7, 1552396. doi: 10.3389/fdgth.2025.1552396
+- Jiao, D. (2025). From hypoxic pockets to daily routines: linking brain oxygenation and cognitive resilience. Frontiers in Aging Neuroscience, 17, 1534198. doi: 10.3389/fnagi.2025.1534198
+- Cognition and Brain Teaching Technology, a book on brain-based teaching and cognitive technology.
+- Chinese translations of How Breakthroughs Happen (2020), Mastering the Dynamics of Innovation (2022), and Neuroscience of You (2025).
+
+LinkedIn: https://www.linkedin.com/in/jason-jiao-2972141a4
+
+## Load Example
+
+```bash
+pip install -r requirements.txt
+python examples/load_checkpoint.py --checkpoint pytorch_model.bin
+```
+
+## Probe Example
+
+```bash
+python probes/probe_memory_causality.py \
+  --checkpoint pytorch_model.bin \
+  --sidecar cedl_config.json \
+  --n-items 200 \
+  --seed 1
+```
+
+## Citation
+
+If you use this checkpoint before the manuscript is published, cite it as:
+
+```bibtex
+@misc{jiao2026cedl,
+  title = {CEDL: A Hippocampal-Inspired Language Model with Contextual Memory Readout},
+  author = {Jiao, Dian},
+  year = {2026},
+  note = {Research checkpoint for contextual memory-readout probes}
+}
+```

cedl_config.json

@@ -0,0 +1,98 @@
+{
+  "model_name": "CEDL",
+  "model_type": "cedl",
+  "architecture": "hippocampal-inspired contextual memory language model",
+  "vocab_size": 50257,
+  "max_position_embeddings": 1024,
+  "parameters": 136254705,
+  "checkpoint_file": "pytorch_model.bin",
+  "implementation_file": "CEDL.py",
+  "memory_readout": {
+    "enabled": true,
+    "source": "contextual_memory_state",
+    "mode": "contextual_direct",
+    "lambda_head": true,
+    "lambda_head_hidden": 160,
+    "lambda_head_bias_init": -7.0,
+    "lambda_head_w_init_std": 0.05,
+    "selection_objective": "binary_answer_background",
+    "background_target": 0.01,
+    "sparsity_weight": 0.05,
+    "sparsity_target": 0.05,
+    "memory_ce_weight": 1.0,
+    "pair_ce_weight": 5.0,
+    "source_adapter": true,
+    "context_adapter": true,
+    "no_injection": true
+  },
+  "mechanism_evidence": {
+    "current_vs_stale_causality": [
+      {
+        "seed": 1,
+        "N": 158,
+        "active_top1": 0.791,
+        "memory_off_top1": 0.0,
+        "delta_pp": 79.1,
+        "delta_margin": 8.266,
+        "t": 20.756,
+        "cohen_d": 1.656
+      },
+      {
+        "seed": 2,
+        "N": 164,
+        "active_top1": 0.732,
+        "memory_off_top1": 0.0,
+        "delta_pp": 73.2,
+        "delta_margin": 7.702,
+        "t": 18.912,
+        "cohen_d": 1.481
+      },
+      {
+        "seed": 3,
+        "N": 162,
+        "active_top1": 0.728,
+        "memory_off_top1": 0.0,
+        "delta_pp": 72.8,
+        "delta_margin": 7.756,
+        "t": 18.507,
+        "cohen_d": 1.459
+      }
+    ],
+    "memory_readout_diagnosis": [
+      {
+        "seed": 1,
+        "N": 158,
+        "trunk_top1_cur": 0.0,
+        "memory_top1_cur": 0.797,
+        "active_top1_cur": 0.791,
+        "memory_top1_stale": 0.006,
+        "current_rank_median": 1
+      },
+      {
+        "seed": 2,
+        "N": 164,
+        "trunk_top1_cur": 0.0,
+        "memory_top1_cur": 0.738,
+        "active_top1_cur": 0.732,
+        "memory_top1_stale": 0.006,
+        "current_rank_median": 1
+      },
+      {
+        "seed": 3,
+        "N": 162,
+        "trunk_top1_cur": 0.0,
+        "memory_top1_cur": 0.735,
+        "active_top1_cur": 0.728,
+        "memory_top1_stale": 0.012,
+        "current_rank_median": 1
+      }
+    ]
+  },
+  "general_eval": {
+    "wikitext103_test_ppl": 30.11,
+    "lambada": 0.159,
+    "hellaswag": 0.271,
+    "arc_easy": 0.281
+  },
+  "claim_boundary": "Research checkpoint for contextual memory readout; not a general downstream SOTA model."
+}

config.json

@@ -0,0 +1,14 @@
+{
+  "model_type": "cedl",
+  "architectures": ["CEDLTwoLoop100M"],
+  "model_name": "CEDL",
+  "vocab_size": 50257,
+  "max_position_embeddings": 1024,
+  "parameters": 136254705,
+  "checkpoint_file": "pytorch_model.bin",
+  "sidecar_config_file": "cedl_config.json",
+  "implementation_file": "CEDL.py",
+  "public_description": "Hippocampal-inspired language model with contextual memory readout.",
+  "claim_boundary": "Research checkpoint for contextual memory readout; not a general downstream SOTA model."
+}
+

data_v4c_pairs.py

@@ -0,0 +1,615 @@
+#!/usr/bin/env python3
+"""V4c synthetic-pair generator.
+
+Generates short belief-revision / contrast / update prompts with EXPLICIT
+token-position spans for query / current / stale / distractor / neutral.
+Resolved against GPT-2 BPE at generation time so spans are guaranteed correct
+when consumed by the V4c contrastive loss.
+
+Five item families:
+  - but-update          : "Alice owns hat. Bob owns pen. But Alice now owns car. What does Alice own now? Answer:"
+  - however-revision    : "The capital was Bonn until 1990. However the capital is now Berlin. What is the capital? Answer:"
+  - temporal-update     : "Previously the gate was A12. Currently it is C7. Which gate is current? Answer:"
+  - paraphrased-equiv.  : update twice (different paraphrases) → multiple positives for SupCon
+  - neutral-control     : "Carol owns book. Frank owns desk. What does Carol own? Answer:"  (no update; supplies negatives only)
+
+Critical span rule (causal LM constraint): `query` MUST be the question-side
+entity occurrence (the "Alice" in "What does Alice own now?"), NOT the first
+factual mention. Otherwise the model would be asked to pull an early Alice
+token toward a future "car" token, which is causally unreachable.
+
+Returns (ids, span_dict) where span_dict has keys:
+    family             : str — one of the 5 family names
+    query              : list[(start, end)] — token-position spans (end-exclusive)
+    current            : list[(start, end)]
+    stale              : list[(start, end)]
+    distractor         : list[(start, end)] | empty
+    neutral            : list[(start, end)] | empty
+    paraphrase_positives: list[(start, end)] | empty
+For neutral-control: query/current/stale/distractor/paraphrase_positives all
+empty; only `neutral` is populated.
+"""
+
+from __future__ import annotations
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Tuple
+import random
+
+
+NAMES = ["Alice", "Bob", "Carol", "Dave", "Eve", "Frank", "Grace",
+         "Hank", "Iris", "Jack", "Kate", "Leo", "Mia", "Nick",
+         "Olivia", "Paul", "Quinn", "Rosa", "Sam", "Tina",
+         "Uma", "Victor", "Wendy", "Xander", "Yara", "Zane",
+         "Adam", "Beth", "Clara", "Drew", "Ella", "Felix",
+         "Greta", "Henry", "Ivy", "Julian", "Kira", "Liam",
+         "Mona", "Noah"]
+OBJECTS = ["car", "hat", "bag", "pen", "cup", "box", "ring", "lamp",
+           "book", "ball", "shoe", "coat", "desk", "bell", "fork",
+           "drum", "fish", "kite", "coin", "vase",
+           "rope", "horn", "tent", "boot", "scarf", "glove",
+           "torch", "mask", "wand", "shield", "spear", "axe",
+           "brush", "comb", "flute", "kettle", "knife", "ladle",
+           "mirror", "needle"]
+CAPITALS_OLD = [
+    ("Bonn", "Berlin", "1990"),
+    ("Karachi", "Islamabad", "1960"),
+    ("Lagos", "Abuja", "1991"),
+    ("Rio de Janeiro", "Brasilia", "1960"),
+    ("Almaty", "Astana", "1997"),
+    ("Yangon", "Naypyidaw", "2006"),
+    ("Saint Petersburg", "Moscow", "1918"),
+    ("Philadelphia", "Washington", "1800"),
+    ("Kyoto", "Tokyo", "1868"),
+    ("Calcutta", "Delhi", "1911"),
+    ("Auckland", "Wellington", "1865"),
+    ("Melbourne", "Canberra", "1927"),
+    ("Salvador", "Rio de Janeiro", "1763"),
+    ("Toronto", "Ottawa", "1857"),
+    ("Bogota", "Lima", "1542"),
+    ("Quito", "Cuenca", "1830"),
+    ("Caracas", "Maracay", "1917"),
+    ("Asuncion", "Concepcion", "1811"),
+    ("Sucre", "La Paz", "1899"),
+    ("Cuzco", "Lima", "1535"),
+    ("Tunis", "Carthage", "1881"),
+    ("Algiers", "Oran", "1962"),
+    ("Tripoli", "Benghazi", "1969"),
+    ("Aden", "Sanaa", "1990"),
+    ("Beirut", "Damascus", "1944"),
+    ("Vienna", "Salzburg", "1918"),
+    ("Krakow", "Warsaw", "1596"),
+    ("Helsinki", "Turku", "1812"),
+    ("Stockholm", "Uppsala", "1523"),
+    ("Trondheim", "Oslo", "1814"),
+]
+def _build_gate_pairs():
+    letters = list("ABCDEFGHJKLMNPQRSTUVWXYZ")
+    numbers = list(range(1, 21))
+    rng = random.Random(20260526)
+    pairs = []
+    for _ in range(260):
+        L_old, L_new = rng.sample(letters, 2)
+        n_old, n_new = rng.sample(numbers, 2)
+        pairs.append((f"{L_old}{n_old}", f"{L_new}{n_new}"))
+    return pairs
+GATES = _build_gate_pairs()
+PARAPHRASE_OLD_CITIES = [
+    "Paris", "London", "Madrid", "Rome", "Vienna",
+    "Athens", "Lisbon", "Prague", "Warsaw", "Helsinki",
+    "Dublin", "Brussels", "Amsterdam", "Stockholm", "Copenhagen",
+    "Budapest", "Sofia", "Bucharest", "Zagreb", "Belgrade",
+]
+PARAPHRASE_NEW_CITIES = [
+    "Berlin", "Tokyo", "Seoul", "Cairo", "Oslo",
+    "Manila", "Bangkok", "Hanoi", "Jakarta", "Singapore",
+    "Mumbai", "Lagos", "Nairobi", "Dakar", "Accra",
+    "Lima", "Bogota", "Quito", "Santiago", "Montevideo",
+]
+HOWEVER_SUBJECTS = [
+    "capital", "headquarters", "embassy",
+    "consulate", "annex", "secretariat",
+    "garrison", "outpost", "depot",
+    "treasury", "archive", "registry",
+    "tribunal", "synod", "academy",
+    "observatory", "arena", "fortress",
+    "monastery", "citadel",
+]
+GATE_SUBJECTS = [
+    "gate", "code", "room", "key", "slot",
+    "platform", "channel", "pin", "lock", "terminal",
+    "passage", "station", "dock", "port",
+    "panel", "module", "bay", "console",
+]
+
+
+VOCAB_SPLIT_SEED = 20260528
+_TRAIN_FRAC = 0.70
+
+
+def _partition(pool, offset):
+    """Split a pool of ATOMIC SURFACE STRINGS ~70/30. Dedups + sorts first so
+    the partition is over unique surfaces (review finding: tuple-level splitting
+    of CAPITALS_OLD/GATES leaked the old/new strings shared across tuples — the
+    surfaces, not the tuples, are the entities under the novel-entity claim)."""
+    uniq = sorted(set(pool))
+    idx = list(range(len(uniq)))
+    random.Random(VOCAB_SPLIT_SEED + offset).shuffle(idx)
+    n_tr = max(1, int(round(_TRAIN_FRAC * len(uniq))))
+    tr = set(idx[:n_tr])
+    train = [uniq[i] for i in range(len(uniq)) if i in tr]
+    test = [uniq[i] for i in range(len(uniq)) if i not in tr]
+    return train, test
+
+
+_CITIES = ([c for (_o, _n, _y) in CAPITALS_OLD for c in (_o, _n)]
+           + PARAPHRASE_OLD_CITIES + PARAPHRASE_NEW_CITIES)
+_GATE_CODES = [c for (_o, _n) in GATES for c in (_o, _n)]
+YEARS = sorted({_y for (_o, _n, _y) in CAPITALS_OLD})
+
+_POOL_REGISTRY = {
+    "NAMES": NAMES, "OBJECTS": OBJECTS,
+    "CITIES": _CITIES, "GATE_CODES": _GATE_CODES,
+    "HOWEVER_SUBJECTS": HOWEVER_SUBJECTS, "GATE_SUBJECTS": GATE_SUBJECTS,
+}
+_SPLITS = {}
+for _off, (_nm, _pool) in enumerate(_POOL_REGISTRY.items()):
+    _tr, _te = _partition(_pool, _off)
+    _SPLITS[_nm] = {"train": _tr, "test": _te, "all": sorted(set(_pool))}
+
+
+def _P(name, split):
+    """Pool `name` restricted to vocabulary `split`. Raises on an unknown split
+    (review finding: a silent fallback to 'all' would leak train/test vocab)."""
+    if split not in ("train", "test", "all"):
+        raise ValueError(
+            f"unknown vocab split {split!r}; expected 'train'/'test'/'all'")
+    return _SPLITS[name][split]
+
+
+def vocab_split_summary():
+    """Sizes of each pool per split — for sanity logging."""
+    return {nm: {s: len(_SPLITS[nm][s]) for s in ("train", "test", "all")}
+            for nm in _SPLITS}
+
+
+def vocab_split_overlap():
+    """Global train∩test surface-string overlap across ALL pools. Should be
+    empty — validates the 'novel entity' claim end to end."""
+    tr, te = set(), set()
+    for nm in _SPLITS:
+        tr |= set(_SPLITS[nm]["train"])
+        te |= set(_SPLITS[nm]["test"])
+    return sorted(tr & te)
+
+
+SpanT = Tuple[int, int]
+
+
+@dataclass
+class V4cItem:
+    """One synthetic contrastive pair with token-position spans."""
+    family: str
+    ids: List[int]
+    query: List[SpanT] = field(default_factory=list)
+    current: List[SpanT] = field(default_factory=list)
+    stale: List[SpanT] = field(default_factory=list)
+    distractor: List[SpanT] = field(default_factory=list)
+    neutral: List[SpanT] = field(default_factory=list)
+    paraphrase_positives: List[SpanT] = field(default_factory=list)
+    original_length: int = 0
+
+
+def _tok_span(tokenizer, full_ids: List[int], substring: str,
+              search_from: int = 0) -> Optional[SpanT]:
+    """Locate `substring` in the decoded text and return its token-position
+    span. Returns None if not found. `search_from` lets the caller bias the
+    search past earlier occurrences (e.g., the question-side query token must
+    come AFTER both current and stale mentions)."""
+    encoded = tokenizer.encode(" " + substring, add_special_tokens=False)
+    if not encoded:
+        encoded = tokenizer.encode(substring, add_special_tokens=False)
+        if not encoded:
+            return None
+    target_len = len(encoded)
+    for start in range(search_from, len(full_ids) - target_len + 1):
+        if full_ids[start:start + target_len] == encoded:
+            return (start, start + target_len)
+    encoded_bare = tokenizer.encode(substring, add_special_tokens=False)
+    if encoded_bare:
+        L = len(encoded_bare)
+        for start in range(search_from, len(full_ids) - L + 1):
+            if full_ids[start:start + L] == encoded_bare:
+                return (start, start + L)
+    return None
+
+
+def _last_tok_span(tokenizer, full_ids: List[int], substring: str) -> Optional[SpanT]:
+    """Find the LAST occurrence of `substring` in token-position space."""
+    encoded = tokenizer.encode(" " + substring, add_special_tokens=False)
+    if not encoded:
+        return None
+    L = len(encoded)
+    found = None
+    for start in range(0, len(full_ids) - L + 1):
+        if full_ids[start:start + L] == encoded:
+            found = (start, start + L)
+    if found is None:
+        encoded_bare = tokenizer.encode(substring, add_special_tokens=False)
+        if encoded_bare:
+            L = len(encoded_bare)
+            for start in range(0, len(full_ids) - L + 1):
+                if full_ids[start:start + L] == encoded_bare:
+                    found = (start, start + L)
+    return found
+
+
+def gen_but_update(tokenizer, rng: random.Random, split: str = "all",
+                    force_stale_obj: str = None) -> V4cItem:
+    """Alice owns hat. Bob owns pen. Carol owns cup. But Alice now owns car.
+    What does Alice own now? Answer:
+
+    `force_stale_obj` (hard-collision support): pin the stale object to a given
+    token (e.g. another item's CURRENT object) so the same token appears in
+    opposite roles across the batch — forcing role-by-position, not token id."""
+    n_distract = 2
+    names = rng.sample(_P("NAMES", split), n_distract + 1)
+    if force_stale_obj is not None:
+        others = rng.sample(
+            [o for o in _P("OBJECTS", split) if o != force_stale_obj],
+            n_distract + 1)
+        objs = [force_stale_obj] + others
+    else:
+        objs = rng.sample(_P("OBJECTS", split), n_distract + 2)
+    subject = names[0]
+    stale_obj = objs[0]
+    new_obj = objs[-1]
+    distractor_obj = objs[1] if n_distract > 0 else objs[0]
+    facts = [f"{names[i]} owns a {objs[i]}" for i in range(n_distract + 1)]
+    rng.shuffle(facts)
+    text = (". ".join(facts) + ". "
+            + f"But {subject} now owns a {new_obj}. "
+            + f"What does {subject} own now? Answer:")
+    ids = tokenizer.encode(text, add_special_tokens=False)
+
+    current_span = _tok_span(tokenizer, ids, new_obj)
+    stale_span = _tok_span(tokenizer, ids, stale_obj)
+    distractor_span = _tok_span(tokenizer, ids, distractor_obj)
+    query_span = _last_tok_span(tokenizer, ids, subject)
+    return V4cItem(
+        family="but_update",
+        ids=ids,
+        query=[query_span] if query_span else [],
+        current=[current_span] if current_span else [],
+        stale=[stale_span] if stale_span else [],
+        distractor=[distractor_span] if distractor_span else [],
+        original_length=len(ids),
+    )
+
+
+def gen_however_revision(tokenizer, rng: random.Random,
+                          split: str = "all",
+                          force_subject: str = None) -> V4cItem:
+    """The headquarters was Bonn until 1990. However the headquarters is now
+    Berlin. What is the headquarters? Answer:
+
+    Subject noun varies across items to avoid (query=`capital`, current=X)
+    collisions when the same `current` city is reused. With ~14 subjects ×
+    30 (old, new, year) triples, paired-collision rate stays under birthday-
+    paradox bounds for 100-150 items."""
+    old, new = rng.sample(_P("CITIES", split), 2)
+    year = rng.choice(YEARS)
+    subject = (force_subject if force_subject is not None
+               else rng.choice(_P("HOWEVER_SUBJECTS", split)))
+    text = (f"The {subject} was {old} until {year}. "
+            f"However the {subject} is now {new}. "
+            f"What is the {subject}? Answer:")
+    ids = tokenizer.encode(text, add_special_tokens=False)
+    current_span = _tok_span(tokenizer, ids, new)
+    stale_span = _tok_span(tokenizer, ids, old)
+    query_span = _last_tok_span(tokenizer, ids, subject)
+    return V4cItem(
+        family="however_revision",
+        ids=ids,
+        query=[query_span] if query_span else [],
+        current=[current_span] if current_span else [],
+        stale=[stale_span] if stale_span else [],
+        original_length=len(ids),
+    )
+
+
+def gen_temporal_update(tokenizer, rng: random.Random,
+                         split: str = "all",
+                         force_subject: str = None) -> V4cItem:
+    """Previously the platform was A12. Currently the platform is C7. Which
+    platform is current now? Answer:
+
+    Subject noun varies; combined with 260-pair GATES pool gives
+    ~18 × 260 = 4680 unique (subject, current) combinations."""
+    old, new = rng.sample(_P("GATE_CODES", split), 2)
+    subject = (force_subject if force_subject is not None
+               else rng.choice(_P("GATE_SUBJECTS", split)))
+    text = (f"Previously the {subject} was {old}. "
+            f"Currently the {subject} is {new}. "
+            f"Which {subject} is current now? Answer:")
+    ids = tokenizer.encode(text, add_special_tokens=False)
+    current_span = _tok_span(tokenizer, ids, new)
+    stale_span = _tok_span(tokenizer, ids, old)
+    query_span = _last_tok_span(tokenizer, ids, subject)
+    return V4cItem(
+        family="temporal_update",
+        ids=ids,
+        query=[query_span] if query_span else [],
+        current=[current_span] if current_span else [],
+        stale=[stale_span] if stale_span else [],
+        original_length=len(ids),
+    )
+
+
+def gen_paraphrased_equiv(tokenizer, rng: random.Random,
+                           split: str = "all") -> V4cItem:
+    """Same entity, two paraphrased questions → two query positives.
+
+    Carol used to live in Paris. She moved to Berlin. Where does Carol live now?
+    Carol resides in Berlin currently. What city is Carol's home? Answer:
+    """
+    person = rng.choice(_P("NAMES", split))
+    old_city, new_city = rng.sample(_P("CITIES", split), 2)
+    text = (f"{person} used to live in {old_city}. "
+            f"{person} moved to {new_city}. "
+            f"Where does {person} live now? "
+            f"{person} resides in {new_city} currently. "
+            f"What city is {person}'s home? Answer:")
+    ids = tokenizer.encode(text, add_special_tokens=False)
+    current_first = _tok_span(tokenizer, ids, new_city)
+    stale_span = _tok_span(tokenizer, ids, old_city)
+    current_second = None
+    if current_first is not None:
+        current_second = _tok_span(tokenizer, ids, new_city,
+                                    search_from=current_first[1])
+    query_span = _last_tok_span(tokenizer, ids, person)
+    item = V4cItem(
+        family="paraphrased_equiv",
+        ids=ids,
+        query=[query_span] if query_span else [],
+        current=[current_first] if current_first else [],
+        stale=[stale_span] if stale_span else [],
+        original_length=len(ids),
+    )
+    if current_second is not None:
+        item.paraphrase_positives.append(current_second)
+    return item
+
+
+def gen_neutral_control(tokenizer, rng: random.Random,
+                         split: str = "all") -> V4cItem:
+    """No update; supplies negatives only. The whole text serves as a single
+    `neutral` span pool — anchor + positive roles are skipped by the loss."""
+    n = 3
+    names = rng.sample(_P("NAMES", split), n)
+    objs = rng.sample(_P("OBJECTS", split), n)
+    facts = [f"{names[i]} owns a {objs[i]}" for i in range(n)]
+    text = ". ".join(facts) + ". "
+    text += f"What does {names[0]} own? Answer:"
+    ids = tokenizer.encode(text, add_special_tokens=False)
+    return V4cItem(
+        family="neutral_control",
+        ids=ids,
+        neutral=[(0, len(ids))],
+        original_length=len(ids),
+    )
+
+
+_FAMILIES = [
+    ("but_update", gen_but_update, 0.30),
+    ("however_revision", gen_however_revision, 0.20),
+    ("temporal_update", gen_temporal_update, 0.20),
+    ("paraphrased_equiv", gen_paraphrased_equiv, 0.10),
+    ("neutral_control", gen_neutral_control, 0.20),
+]
+
+
+_FAMILY_FNS = {name: fn for name, fn, _ in _FAMILIES}
+
+
+def make_entity_swap(item: "V4cItem", tokenizer, rng: random.Random,
+                      split: str = "train") -> "V4cItem":
+    """Structurally-identical twin of `item`: SAME family AND template subtype,
+    only the entity FILLERS swapped, with its OWN recomputed spans (review
+    Blocker 3 — preserve role structure, NOT span offsets). For however/temporal
+    the template subtype is the query subject-noun (capital/gate/…), so it is
+    PRESERVED and only the cities/codes are freshly sampled (review finding 2:
+    a same-family resample that also changed the subject is broader than
+    'same template, fresh entities'). Used by the V4e swap-consistency loss:
+    g at the current/stale role spans must MATCH across this swap, forcing
+    v_vec PRODUCTION to key on STRUCTURE not entity id."""
+    fam = item.family
+    if fam in ("however_revision", "temporal_update") and item.query:
+        qs, qe = item.query[0]
+        subj = tokenizer.decode(item.ids[qs:qe]).strip()
+        fn = (gen_however_revision if fam == "however_revision"
+              else gen_temporal_update)
+        return fn(tokenizer, rng, split, force_subject=subj)
+    fn = _FAMILY_FNS.get(fam, gen_but_update)
+    return fn(tokenizer, rng, split)
+
+
+def generate(tokenizer, n: int, seed: int = 42,
+             neutral_cap: float = 0.25, split: str = "all",
+             hard_collision_frac: float = 0.0,
+             family_weights: Optional[Dict[str, float]] = None
+             ) -> List[V4cItem]:
+    """Generate `n` V4c items by family with the weights above. neutral_cap
+    bounds the neutral_control fraction (since neutrals can't be anchors and
+    >25% neutrals starves SupCon supervision).
+
+    split: 'train'/'test'/'all' vocabulary partition (Correction 4 — novel
+        entity generalization is tested by training on 'train', probing 'test').
+    hard_collision_frac: fraction of but_update items whose STALE object is
+        pinned to a recently-seen CURRENT object (Correction 5 — same token in
+        opposite roles across the batch, forcing role-by-position not token id).
+        within-item current==stale is never created (audit hard-fails it)."""
+    rng = random.Random(seed)
+    _resolved = []
+    for name, fn, w in _FAMILIES:
+        if family_weights is not None and name in family_weights:
+            w = float(family_weights[name])
+        _resolved.append((name, fn, w))
+    families_cumweights = []
+    cum = 0.0
+    for name, fn, w in _resolved:
+        cum += w
+        families_cumweights.append((name, fn, cum))
+    items: List[V4cItem] = []
+    n_neutral = 0
+    n_neutral_cap = int(n * neutral_cap)
+    recent_currents: List[str] = []
+    while len(items) < n:
+        r = rng.random() * cum
+        for name, fn, cwt in families_cumweights:
+            if r <= cwt:
+                if name == "neutral_control" and n_neutral >= n_neutral_cap:
+                    fn = gen_but_update
+                    name = "but_update"
+                if (name == "but_update" and recent_currents
+                        and rng.random() < hard_collision_frac):
+                    forced = rng.choice(recent_currents)
+                    item = gen_but_update(tokenizer, rng, split,
+                                           force_stale_obj=forced)
+                else:
+                    item = fn(tokenizer, rng, split)
+                items.append(item)
+                if item.family == "neutral_control":
+                    n_neutral += 1
+                if item.family == "but_update" and item.current:
+                    cs, ce = item.current[0]
+                    recent_currents.append(
+                        tokenizer.decode(item.ids[cs:ce]).strip())
+                    recent_currents = recent_currents[-32:]
+                break
+    return items
+
+
+def audit(items: List[V4cItem], tokenizer, *, verbose: bool = False
+          ) -> Dict[str, object]:
+    """Family-aware audit. Returns a report dict; `verbose` prints failures."""
+    report = {
+        "total": len(items),
+        "by_family": {},
+        "hard_fails": [],
+        "duplicate_collisions": 0,
+        "n_anchorable": 0,
+        "duplicate_collisions_by_family": {},
+        "top_repeated_pairs": [],
+    }
+    for it in items:
+        report["by_family"][it.family] = report["by_family"].get(it.family, 0) + 1
+
+    seen_keys = {}
+    pair_counter = {}
+    for i, it in enumerate(items):
+        ids_set = set(it.ids)
+        for span_list in [it.query, it.current, it.stale, it.distractor,
+                          it.neutral, it.paraphrase_positives]:
+            for s, e in span_list:
+                if not (0 <= s < e <= len(it.ids)):
+                    report["hard_fails"].append((i, f"span out of bounds: ({s},{e}) len={len(it.ids)}"))
+
+        if it.family == "neutral_control":
+            if not it.neutral:
+                report["hard_fails"].append((i, "neutral_control with empty neutral"))
+            for fname in ("query", "current", "stale", "distractor",
+                          "paraphrase_positives"):
+                if getattr(it, fname):
+                    report["hard_fails"].append((i, f"neutral_control has non-empty {fname}"))
+        else:
+            if not it.query:
+                report["hard_fails"].append((i, f"{it.family}: empty query"))
+                continue
+            if not it.current:
+                report["hard_fails"].append((i, f"{it.family}: empty current"))
+                continue
+            if not it.stale:
+                report["hard_fails"].append((i, f"{it.family}: empty stale"))
+                continue
+            current_text = tokenizer.decode(
+                [t for s, e in it.current for t in it.ids[s:e]]).strip()
+            stale_text = tokenizer.decode(
+                [t for s, e in it.stale for t in it.ids[s:e]]).strip()
+            if current_text == stale_text:
+                report["hard_fails"].append(
+                    (i, f"{it.family}: current surface == stale surface "
+                        f"({current_text!r})"))
+            for s, e in it.distractor:
+                d_text = tokenizer.decode(it.ids[s:e]).strip()
+                if d_text == current_text or d_text == stale_text:
+                    report["hard_fails"].append(
+                        (i, f"{it.family}: distractor surface matches "
+                            f"current/stale ({d_text!r})"))
+                    break
+            q_start = it.query[0][0]
+            c_max = max(e for s, e in it.current)
+            s_max = max(e for s, e in it.stale)
+            if not (q_start >= c_max and q_start >= s_max):
+                report["hard_fails"].append((
+                    i, f"{it.family}: query_pos={q_start} not > "
+                       f"current_end={c_max} / stale_end={s_max}"))
+            else:
+                report["n_anchorable"] += 1
+
+            q_tokens = tuple(it.ids[it.query[0][0]:it.query[0][1]])
+            c_tokens = tuple(it.ids[it.current[0][0]:it.current[0][1]])
+            key = (q_tokens, c_tokens)
+            if key in seen_keys:
+                report["duplicate_collisions"] += 1
+                report["duplicate_collisions_by_family"][it.family] = (
+                    report["duplicate_collisions_by_family"].get(it.family, 0) + 1)
+            seen_keys[key] = True
+            q_decoded = tokenizer.decode(list(q_tokens)).strip()
+            c_decoded = tokenizer.decode(list(c_tokens)).strip()
+            pair_key = (it.family, q_decoded, c_decoded)
+            pair_counter[pair_key] = pair_counter.get(pair_key, 0) + 1
+
+    sorted_pairs = sorted(pair_counter.items(), key=lambda kv: -kv[1])
+    report["top_repeated_pairs"] = [
+        {"family": k[0], "query": k[1], "current": k[2], "count": v}
+        for k, v in sorted_pairs[:10] if v > 1
+    ]
+
+    current_surfaces = set()
+    for it in items:
+        if it.current:
+            s, e = it.current[0]
+            current_surfaces.add(tokenizer.decode(it.ids[s:e]).strip())
+    n_anchor = 0
+    n_role_collision = 0
+    for it in items:
+        if it.family == "neutral_control" or not it.current or not it.stale:
+            continue
+        n_anchor += 1
+        neg_surfaces = []
+        s, e = it.stale[0]
+        neg_surfaces.append(tokenizer.decode(it.ids[s:e]).strip())
+        for s, e in it.distractor:
+            neg_surfaces.append(tokenizer.decode(it.ids[s:e]).strip())
+        if any(ns in current_surfaces for ns in neg_surfaces):
+            n_role_collision += 1
+    report["role_collision_rate"] = (n_role_collision / n_anchor
+                                      if n_anchor else 0.0)
+    report["role_collision_count"] = n_role_collision
+
+    if verbose and report["hard_fails"]:
+        for i, reason in report["hard_fails"][:20]:
+            print(f"    HARD-FAIL item {i}: {reason}")
+        if len(report["hard_fails"]) > 20:
+            print(f"    ... ({len(report['hard_fails']) - 20} more)")
+    return report
+
+
+def decode_spans(item: V4cItem, tokenizer) -> Dict[str, List[str]]:
+    """Decode each span family back into surface text for visual audit."""
+    out = {}
+    for fname in ("query", "current", "stale", "distractor", "neutral",
+                  "paraphrase_positives"):
+        out[fname] = [tokenizer.decode(item.ids[s:e]).strip()
+                       for s, e in getattr(item, fname)]
+    return out

examples/load_checkpoint.py

@@ -0,0 +1,91 @@
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import torch
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from CEDL import build_model
+
+
+def constructor_kwargs(config_path):
+    with open(config_path) as f:
+        cfg = json.load(f)
+    mem = cfg.get("memory_readout", {})
+    source_name = str(mem.get("source", "contextual_memory_state"))
+    source_map = {
+        "contextual_memory_state": "q_mem",
+        "decoder_state": "h_d",
+        "expanded_state": "h_e",
+        "attractor_state": "q_attractor",
+        "q_mem": "q_mem",
+    }
+    return dict(
+        lambda_head=bool(mem.get("lambda_head", True)),
+        lambda_head_hidden=int(mem.get("lambda_head_hidden", 160)),
+        lambda_head_bias_init=float(mem.get("lambda_head_bias_init", -7.0)),
+        lambda_head_w_init_std=float(
+            mem.get("lambda_head_w_init_std", 0.05)),
+        bce_objective=(
+            mem.get("selection_objective") == "binary_answer_background"),
+        sel_weight=1.0,
+        bg_weight=1.0,
+        bg_target=float(mem.get("background_target", 0.01)),
+        wt_sparsity_weight=float(mem.get("sparsity_weight", 0.05)),
+        wt_sparsity_target=float(mem.get("sparsity_target", 0.05)),
+        memory_head_enabled=bool(mem.get("enabled", True)),
+        memory_ce_weight=float(mem.get("memory_ce_weight", 1.0)),
+        memory_pair_ce_weight=float(mem.get("pair_ce_weight", 5.0)),
+        memory_query_source=source_map.get(source_name, source_name),
+        memory_readout_mode="direct",
+        source_adapter=bool(mem.get("source_adapter", True)),
+        context_adapter=bool(mem.get("context_adapter", True)),
+        specialist_noinject=bool(mem.get("no_injection", True)),
+    )
+
+
+def unwrap_state_dict(obj):
+    if isinstance(obj, dict) and "model" in obj and isinstance(obj["model"], dict):
+        return obj["model"]
+    return obj
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", default="pytorch_model.bin")
+    parser.add_argument("--config", default="cedl_config.json")
+    parser.add_argument("--device", default="cpu")
+    args = parser.parse_args()
+
+    ckpt_path = Path(args.checkpoint)
+    if not ckpt_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
+    cfg_path = Path(args.config)
+    if not cfg_path.exists():
+        raise FileNotFoundError(f"Config not found: {cfg_path}")
+
+    model = build_model(
+        "CEDL",
+        vocab=50257,
+        max_seq=1024,
+        **constructor_kwargs(cfg_path),
+    )
+
+    state = unwrap_state_dict(torch.load(ckpt_path, map_location="cpu"))
+    result = model.load_state_dict(state, strict=True)
+    model.to(args.device)
+    model.eval()
+
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"Loaded {ckpt_path}")
+    print(f"Parameters: {n_params:,}")
+    print(f"Missing keys: {len(result.missing_keys)}")
+    print(f"Unexpected keys: {len(result.unexpected_keys)}")
+
+
+if __name__ == "__main__":
+    main()

probes/probe_memory_causality.py

@@ -0,0 +1,275 @@
+"""Memory-readout causality probe.
+
+Runs the same current-vs-stale items through the live CEDL memory path and a
+memory-off ablation. The paired comparison estimates whether contextual memory
+readout causally improves current-token selection.
+
+Usage:
+  python probes/probe_memory_causality.py \
+    --checkpoint pytorch_model.bin \
+    --sidecar cedl_config.json \
+    --n-items 50
+"""
+
+import argparse
+import json
+import os
+import sys
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, "/content")
+
+import CEDL
+import data_v4c_pairs as v4c
+
+
+MEMORY_GATE_ATTR = "v" + "6_lambda_head"
+
+
+class ZeroLambdaHead(nn.Module):
+    """Forces the memory mixture gate near zero."""
+    def forward(self, h):
+        return torch.full(
+            h.shape[:-1] + (1,), -100.0,
+            device=h.device, dtype=h.dtype,
+        )
+
+
+def load_sidecar_constructor_kwargs(sidecar_path):
+    """Read public CEDL config and return constructor keyword arguments."""
+    with open(sidecar_path) as f:
+        sc = json.load(f)
+    mem = sc.get("memory_readout")
+    if not isinstance(mem, dict):
+        raise ValueError("Expected cedl_config.json with a memory_readout block.")
+    source_name = str(mem.get("source", "contextual_memory_state"))
+    source_map = {
+        "contextual_memory_state": "q_mem",
+        "decoder_state": "h_d",
+        "expanded_state": "h_e",
+        "attractor_state": "q_attractor",
+        "q_mem": "q_mem",
+    }
+    return dict(
+        lambda_head=bool(mem.get("lambda_head", True)),
+        lambda_head_hidden=int(mem.get("lambda_head_hidden", 160)),
+        lambda_head_bias_init=float(mem.get("lambda_head_bias_init", -7.0)),
+        lambda_head_w_init_std=float(
+            mem.get("lambda_head_w_init_std", 0.05)),
+        bce_objective=(
+            mem.get("selection_objective") == "binary_answer_background"),
+        sel_weight=1.0,
+        bg_weight=1.0,
+        bg_target=float(mem.get("background_target", 0.01)),
+        wt_sparsity_weight=float(mem.get("sparsity_weight", 0.05)),
+        wt_sparsity_target=float(mem.get("sparsity_target", 0.05)),
+        memory_head_enabled=bool(mem.get("enabled", True)),
+        memory_ce_weight=float(mem.get("memory_ce_weight", 1.0)),
+        memory_pair_ce_weight=float(mem.get("pair_ce_weight", 5.0)),
+        memory_query_source=source_map.get(source_name, source_name),
+        memory_readout_mode="direct",
+        source_adapter=bool(mem.get("source_adapter", True)),
+        context_adapter=bool(mem.get("context_adapter", True)),
+        specialist_noinject=bool(mem.get("no_injection", True)),
+    )
+
+
+def model_forward_logits(m, ids_b):
+    """Call m(ids_b) and return the logits tensor [B, T, V] regardless of
+    whether forward returns logits-only or (logits, aux_loss)."""
+    out = m(ids_b)
+    if isinstance(out, tuple):
+        logits = out[0]
+    else:
+        logits = out
+    return logits
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--checkpoint", type=str, default="pytorch_model.bin")
+    p.add_argument("--sidecar", type=str, default="cedl_config.json")
+    p.add_argument("--n-items", type=int, default=50)
+    p.add_argument("--seed", type=int, default=0)
+    args = p.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"[setup] device={device}")
+    print(f"[setup] checkpoint={args.checkpoint}")
+    print(f"[setup] sidecar={args.sidecar}")
+    print(f"[setup] n_items={args.n_items}")
+
+    model_kwargs = load_sidecar_constructor_kwargs(args.sidecar)
+    print(f"[setup] memory_head={model_kwargs['lambda_head']}  "
+          f"hidden={model_kwargs['lambda_head_hidden']}  "
+          f"bias_init={model_kwargs['lambda_head_bias_init']}  "
+          f"sparsity_weight={model_kwargs['wt_sparsity_weight']}  "
+          f"sparsity_target={model_kwargs['wt_sparsity_target']}  "
+          f"memory_source={model_kwargs['memory_query_source']}  "
+          f"readout_mode={model_kwargs['memory_readout_mode']}")
+
+    m = CEDL.build_model("CEDL", vocab=50257, max_seq=1024, **model_kwargs)
+    m = m.to(device).eval()
+
+    state = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
+    msd = state["model"] if isinstance(state, dict) and "model" in state else state
+    if any(k.startswith("_orig_mod.") for k in msd):
+        msd = {k.replace("_orig_mod.", ""): v for k, v in msd.items()}
+    res = m.load_state_dict(msd, strict=True)
+    print(f"[load] strict OK  (missing={len(res.missing_keys)}  unexpected={len(res.unexpected_keys)})")
+
+    m.feedback_alpha.fill_(1.0)
+    if hasattr(m, "sl_alpha"):
+        m.sl_alpha.fill_(1.0)
+
+    from transformers import GPT2TokenizerFast
+    tok = GPT2TokenizerFast.from_pretrained("gpt2")
+    print(f"\n[v4c] generating {args.n_items} items (seed={args.seed})...")
+    items = v4c.generate(tok, n=args.n_items, seed=args.seed)
+    print(f"[v4c] generated {len(items)} raw items")
+
+    saved_head = getattr(m, MEMORY_GATE_ATTR)
+    zero_head = ZeroLambdaHead().to(device)
+
+    margin_active = []
+    margin_off = []
+    pred_correct_active = []
+    pred_correct_off = []
+    lam_at_ans_active = []
+    n_skipped = 0
+    skip_reasons = {"neutral_control": 0, "no_cur_stale": 0,
+                    "bad_length": 0, "same_token": 0}
+
+    for it_idx, it in enumerate(items):
+        if it.family == "neutral_control":
+            n_skipped += 1; skip_reasons["neutral_control"] += 1; continue
+        if not it.current or not it.stale:
+            n_skipped += 1; skip_reasons["no_cur_stale"] += 1; continue
+        ol = getattr(it, "original_length", 0)
+        if ol <= 1 or ol > 1024:
+            n_skipped += 1; skip_reasons["bad_length"] += 1; continue
+        cur_t = int(it.ids[it.current[0][0]])
+        stale_t = int(it.ids[it.stale[0][0]])
+        if cur_t == stale_t:
+            n_skipped += 1; skip_reasons["same_token"] += 1; continue
+        ans_p = ol - 1
+
+        ids_b = torch.tensor([it.ids[:ol]], device=device, dtype=torch.long)
+
+        setattr(m, MEMORY_GATE_ATTR, saved_head)
+        with torch.no_grad():
+            logits_a = model_forward_logits(m, ids_b)
+        cur_lp_a = float(logits_a[0, ans_p, cur_t].item())
+        stl_lp_a = float(logits_a[0, ans_p, stale_t].item())
+        pred_a = int(logits_a[0, ans_p].argmax().item())
+        margin_active.append(cur_lp_a - stl_lp_a)
+        pred_correct_active.append(pred_a == cur_t)
+
+        with torch.no_grad():
+            h_C = m.c_stage(ids_b, feedback=None)
+            h_E, h_E_sparse = m.e_stage(h_C, feedback=None)
+            v_vec, _ = m.salience(h_E) if m.use_salience else (None, None)
+            h_D, _ = m.d_stage(h_E, h_E_sparse, v_vec=v_vec)
+            lam_logit_a = saved_head(h_D[:, ans_p:ans_p + 1, :])
+            lam_a = torch.sigmoid(lam_logit_a).clamp(1e-4, 1.0 - 1e-4)
+            lam_at_ans_active.append(float(lam_a.item()))
+
+        setattr(m, MEMORY_GATE_ATTR, zero_head)
+        with torch.no_grad():
+            logits_o = model_forward_logits(m, ids_b)
+        cur_lp_o = float(logits_o[0, ans_p, cur_t].item())
+        stl_lp_o = float(logits_o[0, ans_p, stale_t].item())
+        pred_o = int(logits_o[0, ans_p].argmax().item())
+        margin_off.append(cur_lp_o - stl_lp_o)
+        pred_correct_off.append(pred_o == cur_t)
+
+    setattr(m, MEMORY_GATE_ATTR, saved_head)
+
+    n = len(margin_active)
+    print(f"\n[items] used={n}  skipped={n_skipped}  reasons={skip_reasons}")
+    if n == 0:
+        print("No valid items — aborting.")
+        return
+
+    ma = np.array(margin_active)
+    mo = np.array(margin_off)
+    delta = ma - mo
+
+    print(f"\n[λ at answer row, ACTIVE path]  "
+          f"mean={np.mean(lam_at_ans_active):.4f}  "
+          f"std={np.std(lam_at_ans_active):.4f}  "
+          f"min={min(lam_at_ans_active):.4f}  "
+          f"max={max(lam_at_ans_active):.4f}")
+
+    print(f"\n[current − stale logit margin at answer position]")
+    print(f"  ACTIVE:    mean={ma.mean():+.3f}  std={ma.std():.3f}  "
+          f"median={np.median(ma):+.3f}")
+    print(f"  BANK-OFF:  mean={mo.mean():+.3f}  std={mo.std():.3f}  "
+          f"median={np.median(mo):+.3f}")
+    print(f"  Δ (act−off):  mean={delta.mean():+.3f}  std={delta.std():.3f}  "
+          f"median={np.median(delta):+.3f}")
+
+    from scipy import stats as sst
+    t_stat, p_val = sst.ttest_rel(ma, mo)
+    print(f"  paired t = {t_stat:+.3f}  p = {p_val:.4f}  N={n}")
+    d_paired = float(delta.mean() / max(delta.std(), 1e-9))
+    print(f"  Cohen's d (paired) = {d_paired:+.3f}")
+
+    acc_a = float(np.mean(pred_correct_active))
+    acc_o = float(np.mean(pred_correct_off))
+    print(f"\n[top-1 accuracy: argmax(logits[ans_p]) == cur_t]")
+    print(f"  ACTIVE:    {acc_a*100:5.1f}%   ({sum(pred_correct_active)}/{n})")
+    print(f"  BANK-OFF:  {acc_o*100:5.1f}%   ({sum(pred_correct_off)}/{n})")
+    print(f"  Δ:        {(acc_a-acc_o)*100:+5.1f} pp")
+
+    print(f"\n[first 10 items — paired breakdown]")
+    print(f"  {'#':>3}  {'m_A':>7}  {'m_O':>7}  {'Δ':>7}  "
+          f"{'pred_A':>6}  {'pred_O':>6}")
+    for i in range(min(10, n)):
+        print(f"  {i:>3}  {ma[i]:>+7.3f}  {mo[i]:>+7.3f}  {delta[i]:>+7.3f}  "
+              f"{'OK' if pred_correct_active[i] else 'XX':>6}  "
+              f"{'OK' if pred_correct_off[i] else 'XX':>6}")
+
+    print(f"\n{'='*64}")
+    print(f"Memory-readout causality — verdict")
+    print(f"{'='*64}")
+    if t_stat > 2.0 and (acc_a - acc_o) > 0.10:
+        verdict = (
+            f"ACTIVE > BANK-OFF  (t={t_stat:+.2f}, Δacc=+{(acc_a-acc_o)*100:.1f}pp, "
+            f"d={d_paired:+.2f}).\n"
+            f"  → Bank is causally useful for V4c current-vs-stale.\n"
+            f"  → PROCEED to full probe suite + manuscript update."
+        )
+    elif abs(t_stat) < 1.0 and abs(acc_a - acc_o) < 0.05:
+        verdict = (
+            f"ACTIVE ≈ BANK-OFF  (t={t_stat:+.2f}, Δacc={(acc_a-acc_o)*100:+.1f}pp, "
+            f"d={d_paired:+.2f}).\n"
+            f"  → λ head routes correctly but BANK READOUT IS NOT CAUSALLY USEFUL.\n"
+            f"  → If mem_head_bank is already untied and alternate sources still\n"
+            f"    fail, test direct-source readout before more bank replay."
+        )
+    elif t_stat < -2.0:
+        verdict = (
+            f"BANK-OFF > ACTIVE  (t={t_stat:+.2f}, Δacc={(acc_a-acc_o)*100:+.1f}pp, "
+            f"d={d_paired:+.2f}).\n"
+            f"  → BANK ACTIVELY HURTS V4c prediction. Diagnose the bank readout\n"
+            f"    bottleneck before any manuscript work."
+        )
+    else:
+        verdict = (
+            f"INTERMEDIATE  (t={t_stat:+.2f}, Δacc={(acc_a-acc_o)*100:+.1f}pp, "
+            f"d={d_paired:+.2f}).\n"
+            f"  → Mixed signal. Run full probe suite to see broader pattern; if\n"
+            f"    contrastive / stale-vs-current probes show clear lift, the bank\n"
+            f"    is helpful on average even if the V4c-margin metric is noisy."
+        )
+    print(f"  {verdict}")
+    print(f"{'='*64}")
+
+
+if __name__ == "__main__":
+    main()

probes/probe_memory_diagnosis.py

@@ -0,0 +1,310 @@
+"""Memory-readout distributional diagnosis at current-vs-stale answer rows.
+
+Runs trunk-only, memory-only, and active paths on the same current-vs-stale
+items. The report shows whether the contextual memory readout decodes the
+current token and how it differs from the trunk distribution.
+
+Reports:
+  1. Structural check: is `mem_head_bank` a separate parameter from `mem_head`?
+     If shared/tied, that's THE diagnostic.
+  2. Per-path top-1 distribution: how often does each path's argmax match
+     current_t / stale_t / something else?
+  3. Mean rank of current_t and stale_t in each path's logit distribution.
+  4. Softmax entropy: bank vs trunk (overconfident? uniform? matched?).
+  5. First-10-item qualitative dump: bank's top-3 vs trunk's top-3 vs
+     current/stale tokens (decoded).
+
+Usage:
+  python probes/probe_memory_diagnosis.py \
+    --checkpoint pytorch_model.bin \
+    --sidecar cedl_config.json \
+    --n-items 50
+"""
+
+import argparse
+import json
+import os
+import sys
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, "/content")
+
+import CEDL
+import data_v4c_pairs as v4c
+
+
+MEMORY_GATE_ATTR = "v" + "6_lambda_head"
+
+
+class ConstantLambdaHead(nn.Module):
+    def __init__(self, logit_value):
+        super().__init__()
+        self.logit_value = float(logit_value)
+    def forward(self, h):
+        return torch.full(
+            h.shape[:-1] + (1,), self.logit_value,
+            device=h.device, dtype=h.dtype,
+        )
+
+
+def load_sidecar_constructor_kwargs(sidecar_path):
+    with open(sidecar_path) as f:
+        sc = json.load(f)
+    mem = sc.get("memory_readout")
+    if not isinstance(mem, dict):
+        raise ValueError("Expected cedl_config.json with a memory_readout block.")
+    source_name = str(mem.get("source", "contextual_memory_state"))
+    source_map = {
+        "contextual_memory_state": "q_mem",
+        "decoder_state": "h_d",
+        "expanded_state": "h_e",
+        "attractor_state": "q_attractor",
+        "q_mem": "q_mem",
+    }
+    return dict(
+        lambda_head=bool(mem.get("lambda_head", True)),
+        lambda_head_hidden=int(mem.get("lambda_head_hidden", 160)),
+        lambda_head_bias_init=float(mem.get("lambda_head_bias_init", -7.0)),
+        lambda_head_w_init_std=float(
+            mem.get("lambda_head_w_init_std", 0.05)),
+        bce_objective=(
+            mem.get("selection_objective") == "binary_answer_background"),
+        sel_weight=1.0,
+        bg_weight=1.0,
+        bg_target=float(mem.get("background_target", 0.01)),
+        wt_sparsity_weight=float(mem.get("sparsity_weight", 0.05)),
+        wt_sparsity_target=float(mem.get("sparsity_target", 0.05)),
+        memory_head_enabled=bool(mem.get("enabled", True)),
+        memory_ce_weight=float(mem.get("memory_ce_weight", 1.0)),
+        memory_pair_ce_weight=float(mem.get("pair_ce_weight", 5.0)),
+        memory_query_source=source_map.get(source_name, source_name),
+        memory_readout_mode="direct",
+        source_adapter=bool(mem.get("source_adapter", True)),
+        context_adapter=bool(mem.get("context_adapter", True)),
+        specialist_noinject=bool(mem.get("no_injection", True)),
+    )
+
+
+def model_forward_logits(m, ids_b):
+    out = m(ids_b)
+    return out[0] if isinstance(out, tuple) else out
+
+
+def categorize_top1(top1_t, cur_t, stale_t):
+    if top1_t == cur_t: return "current"
+    if top1_t == stale_t: return "stale"
+    return "other"
+
+
+def token_decode(tok, token_id):
+    s = tok.decode([int(token_id)])
+    return repr(s) if len(s) < 16 else repr(s[:13]) + "..."
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--checkpoint", type=str, default="pytorch_model.bin")
+    p.add_argument("--sidecar", type=str, default="cedl_config.json")
+    p.add_argument("--n-items", type=int, default=50)
+    p.add_argument("--seed", type=int, default=0)
+    args = p.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"[setup] device={device}  ckpt={args.checkpoint}")
+
+    model_kwargs = load_sidecar_constructor_kwargs(args.sidecar)
+    print(f"[setup] memory_source={model_kwargs['memory_query_source']}  "
+          f"readout_mode={model_kwargs['memory_readout_mode']}")
+    m = CEDL.build_model("CEDL", vocab=50257, max_seq=1024, **model_kwargs)
+    m = m.to(device).eval()
+
+    state = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
+    msd = state["model"] if isinstance(state, dict) and "model" in state else state
+    if any(k.startswith("_orig_mod.") for k in msd):
+        msd = {k.replace("_orig_mod.", ""): v for k, v in msd.items()}
+    res = m.load_state_dict(msd, strict=True)
+    print(f"[load] strict OK (missing={len(res.missing_keys)} unexpected={len(res.unexpected_keys)})")
+
+    m.feedback_alpha.fill_(1.0)
+    if hasattr(m, "sl_alpha"):
+        m.sl_alpha.fill_(1.0)
+
+    print(f"\n{'='*64}")
+    print("STRUCTURAL CHECK — bank output projection")
+    print(f"{'='*64}")
+    has_mh = hasattr(m, "mem_head")
+    has_mhb = hasattr(m, "mem_head_bank")
+    print(f"  m.mem_head present: {has_mh}")
+    print(f"  m.mem_head_bank present: {has_mhb}")
+    if has_mh and has_mhb:
+        mh = m.mem_head
+        mhb = m.mem_head_bank
+        same_obj = mh is mhb
+        print(f"  m.mem_head is m.mem_head_bank: {same_obj}")
+        if not same_obj:
+            mh_params = {k: v for k, v in mh.named_parameters()}
+            mhb_params = {k: v for k, v in mhb.named_parameters()}
+            print(f"  mem_head params: {list(mh_params.keys())}")
+            print(f"  mem_head_bank params: {list(mhb_params.keys())}")
+            for k in mh_params:
+                if k in mhb_params:
+                    same_param = mh_params[k] is mhb_params[k]
+                    equal_val = torch.equal(mh_params[k], mhb_params[k])
+                    cos = F.cosine_similarity(
+                        mh_params[k].flatten().unsqueeze(0),
+                        mhb_params[k].flatten().unsqueeze(0),
+                    ).item()
+                    print(f"  {k}: same_tensor={same_param}  equal={equal_val}  cos={cos:+.4f}")
+    elif has_mh and not has_mhb:
+        print(f"  → bank uses mem_head directly (NO separate bank output proj).")
+        print(f"  → bank logits = mem_head(attended_value) → shaped by WikiText LM prior.")
+    else:
+        bank_keys = [k for k in msd if "bank" in k.lower()]
+        print(f"  bank-related state_dict keys: {bank_keys}")
+
+    from transformers import GPT2TokenizerFast
+    tok = GPT2TokenizerFast.from_pretrained("gpt2")
+    items = v4c.generate(tok, n=args.n_items, seed=args.seed)
+    print(f"\n[v4c] generated {len(items)} items")
+
+    saved_head = getattr(m, MEMORY_GATE_ATTR)
+    zero_head = ConstantLambdaHead(-100.0).to(device)
+    one_head  = ConstantLambdaHead(+100.0).to(device)
+
+    records = []
+
+    for it_idx, it in enumerate(items):
+        if it.family == "neutral_control": continue
+        if not it.current or not it.stale: continue
+        ol = getattr(it, "original_length", 0)
+        if ol <= 1 or ol > 1024: continue
+        cur_t = int(it.ids[it.current[0][0]])
+        stale_t = int(it.ids[it.stale[0][0]])
+        if cur_t == stale_t: continue
+        ans_p = ol - 1
+
+        ids_b = torch.tensor([it.ids[:ol]], device=device, dtype=torch.long)
+
+        results = {"item": it_idx, "cur_t": cur_t, "stale_t": stale_t}
+        for path, head in [("trunk", zero_head), ("bank", one_head), ("active", saved_head)]:
+            setattr(m, MEMORY_GATE_ATTR, head)
+            with torch.no_grad():
+                logits = model_forward_logits(m, ids_b)
+            row = logits[0, ans_p]
+            log_probs = F.log_softmax(row, dim=-1)
+            top5_p, top5_t = torch.topk(log_probs, k=5)
+            ent = float(-(log_probs.exp() * log_probs).sum().item())
+            cur_rank = int((log_probs > log_probs[cur_t]).sum().item()) + 1
+            stl_rank = int((log_probs > log_probs[stale_t]).sum().item()) + 1
+            results[path] = dict(
+                top5_t=top5_t.tolist(),
+                top5_lp=top5_p.tolist(),
+                top1_t=int(top5_t[0].item()),
+                ent=ent,
+                cur_rank=cur_rank,
+                stl_rank=stl_rank,
+                cur_lp=float(log_probs[cur_t].item()),
+                stl_lp=float(log_probs[stale_t].item()),
+            )
+        records.append(results)
+
+    setattr(m, MEMORY_GATE_ATTR, saved_head)
+    n = len(records)
+    print(f"[used] {n} items\n")
+
+    if n == 0:
+        print("No valid items.")
+        return
+
+    for path in ["trunk", "bank", "active"]:
+        top1_cur = sum(1 for r in records if r[path]["top1_t"] == r["cur_t"])
+        top1_stl = sum(1 for r in records if r[path]["top1_t"] == r["stale_t"])
+        top1_other = n - top1_cur - top1_stl
+        ents = np.array([r[path]["ent"] for r in records])
+        cur_ranks = np.array([r[path]["cur_rank"] for r in records])
+        stl_ranks = np.array([r[path]["stl_rank"] for r in records])
+        cur_lps = np.array([r[path]["cur_lp"] for r in records])
+        stl_lps = np.array([r[path]["stl_lp"] for r in records])
+        print(f"[{path:>6}] top1: cur={top1_cur:>2}/{n} ({top1_cur/n*100:5.1f}%)  "
+              f"stl={top1_stl:>2}/{n} ({top1_stl/n*100:5.1f}%)  "
+              f"other={top1_other:>2}/{n} ({top1_other/n*100:5.1f}%)")
+        print(f"         entropy: mean={ents.mean():.3f}  std={ents.std():.3f}  "
+              f"(uniform~10.83 for V=50257)")
+        print(f"         cur_rank: mean={cur_ranks.mean():.0f}  median={np.median(cur_ranks):.0f}  "
+              f"min={cur_ranks.min()}  max={cur_ranks.max()}")
+        print(f"         stl_rank: mean={stl_ranks.mean():.0f}  median={np.median(stl_ranks):.0f}  "
+              f"min={stl_ranks.min()}  max={stl_ranks.max()}")
+        print(f"         logp(cur): mean={cur_lps.mean():.3f}  logp(stl): mean={stl_lps.mean():.3f}\n")
+
+    bank_eq_trunk = sum(1 for r in records if r["bank"]["top1_t"] == r["trunk"]["top1_t"])
+    bank_in_trunk_top5 = sum(1 for r in records
+                             if r["bank"]["top1_t"] in r["trunk"]["top5_t"])
+    trunk_in_bank_top5 = sum(1 for r in records
+                             if r["trunk"]["top1_t"] in r["bank"]["top5_t"])
+    print(f"[cross-path]")
+    print(f"  bank top1 == trunk top1:        {bank_eq_trunk:>2}/{n} ({bank_eq_trunk/n*100:5.1f}%)")
+    print(f"  bank top1 ∈ trunk top5:         {bank_in_trunk_top5:>2}/{n} ({bank_in_trunk_top5/n*100:5.1f}%)")
+    print(f"  trunk top1 ∈ bank top5:         {trunk_in_bank_top5:>2}/{n} ({trunk_in_bank_top5/n*100:5.1f}%)")
+
+    print(f"\n[first 10 — qualitative bank vs trunk top-3 at answer position]")
+    for i, r in enumerate(records[:10]):
+        cur_s = token_decode(tok, r["cur_t"])
+        stl_s = token_decode(tok, r["stale_t"])
+        print(f"\n  --- item {r['item']}  cur={cur_s}  stale={stl_s} ---")
+        for path in ["trunk", "bank", "active"]:
+            top3_t = r[path]["top5_t"][:3]
+            top3_lp = r[path]["top5_lp"][:3]
+            top3_decoded = [(token_decode(tok, t), f"{lp:+.2f}") for t, lp in zip(top3_t, top3_lp)]
+            print(f"  {path:>6}: top3={top3_decoded}  cur_rank={r[path]['cur_rank']}  stl_rank={r[path]['stl_rank']}")
+
+    print(f"\n{'='*64}")
+    print("B2 diagnostic verdict")
+    print(f"{'='*64}")
+
+    bank_top1_cur = sum(1 for r in records if r["bank"]["top1_t"] == r["cur_t"]) / n
+    trunk_top1_cur = sum(1 for r in records if r["trunk"]["top1_t"] == r["cur_t"]) / n
+    bank_mean_ent = float(np.mean([r["bank"]["ent"] for r in records]))
+    trunk_mean_ent = float(np.mean([r["trunk"]["ent"] for r in records]))
+    bank_cur_rank_mean = float(np.mean([r["bank"]["cur_rank"] for r in records]))
+    trunk_cur_rank_mean = float(np.mean([r["trunk"]["cur_rank"] for r in records]))
+    readout_mode = str(model_kwargs.get("memory_readout_mode", "direct"))
+
+    if bank_top1_cur >= trunk_top1_cur and bank_top1_cur >= 0.25:
+        print(f"  Bank/readout predicts cur token at top-1 {bank_top1_cur*100:.1f}% "
+              f"(vs trunk {trunk_top1_cur*100:.1f}%).")
+        print(f"  Bank/readout ranks cur_t at mean position {bank_cur_rank_mean:.0f} "
+              f"(vs trunk {trunk_cur_rank_mean:.0f}).")
+        if bank_mean_ent + 1.0 < trunk_mean_ent:
+            print(f"  Entropy is lower than trunk ({bank_mean_ent:.2f} vs "
+                  f"{trunk_mean_ent:.2f}); this is overconfident but useful "
+                  f"when top-1/rank improve.")
+        if readout_mode == "direct":
+            print(f"  → DIRECT READOUT PASSES B2. q_mem is decodable; the old "
+                  f"256-slot bank path was the bottleneck.")
+        else:
+            print(f"  → BANK READOUT PASSES B2. Proceed to B1 causal attribution.")
+    elif bank_top1_cur < 0.5 * trunk_top1_cur and bank_cur_rank_mean > 2 * trunk_cur_rank_mean:
+        print(f"  Bank predicts cur token at top-1 only {bank_top1_cur*100:.1f}% (vs trunk {trunk_top1_cur*100:.1f}%)")
+        print(f"  Bank ranks cur_t at mean position {bank_cur_rank_mean:.0f} (vs trunk {trunk_cur_rank_mean:.0f})")
+        print(f"  → Bank output is task-misaligned. STAGE 2B JUSTIFIED.")
+        print(f"  → If 'mem_head_bank is mem_head' check above shows TIED weights:")
+        print(f"    primary fix is to UNTIE mem_head_bank (give bank its own output proj).")
+        print(f"  → If weights ARE already separate and alternate sources still fail:")
+        print(f"    test direct-source readout before more bank replay.")
+    elif abs(bank_mean_ent - trunk_mean_ent) > 1.0:
+        print(f"  Bank entropy ({bank_mean_ent:.2f}) differs from trunk ({trunk_mean_ent:.2f}) by {abs(bank_mean_ent-trunk_mean_ent):.2f}.")
+        print(f"  → Bank distribution is mis-calibrated relative to trunk.")
+    else:
+        print(f"  Bank top-1 cur rate ({bank_top1_cur*100:.1f}%) "
+              f"vs trunk ({trunk_top1_cur*100:.1f}%) — diagnostic inconclusive.")
+        print(f"  Examine qualitative dump above for specific failure mode.")
+    print(f"{'='*64}")
+
+
+if __name__ == "__main__":
+    main()

probes/probe_memory_source_readout.py

@@ -0,0 +1,505 @@
+#!/usr/bin/env python3
+"""Probe whether a memory source vector is directly decodable.
+
+This diagnostic freezes the trained CEDL checkpoint and trains a temporary
+linear vocab readout from one selected source vector:
+
+  h_d, h_e, q_attractor, or q_mem
+
+It bypasses the external bank bottleneck:
+
+    source -> direct_head -> vocab
+
+instead of:
+
+    source -> bank_q_proj -> 256-slot mem_keys/mem_vals -> mem_head_bank -> vocab
+
+If direct_head generalizes on B2 while the bank path fails, the bottleneck is
+the learned 256-slot bank readout. If direct_head also fails, the selected
+source does not robustly encode the current-token target under this synthetic
+held-out probe.
+"""
+
+import argparse
+import json
+import os
+import sys
+from typing import Dict, List, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, "/content")
+
+import CEDL
+import data_v4c_pairs as v4c
+
+
+SOURCES = ("h_d", "h_e", "q_attractor", "q_mem")
+MEMORY_GATE_ATTR = "v" + "6_lambda_head"
+STORE_SOURCES_ATTR = "store_v" + "6_bank_sources"
+NEEDS_SOURCES_METHOD = "_v" + "6_needs_dstage_bank_sources"
+SOURCE_INPUT_METHOD = "_v" + "6_bank_query_input"
+
+SOURCE_ALIASES = {
+    "contextual_memory_state": "q_mem",
+    "decoder_state": "h_d",
+    "expanded_state": "h_e",
+    "attractor_state": "q_attractor",
+    "q_mem": "q_mem",
+    "h_d": "h_d",
+    "h_e": "h_e",
+    "q_attractor": "q_attractor",
+}
+
+
+def load_sidecar_constructor_kwargs(
+    sidecar_path: str,
+    source: str | None,
+) -> Dict[str, object]:
+    with open(sidecar_path) as f:
+        sc = json.load(f)
+    mem = sc.get("memory_readout")
+    if not isinstance(mem, dict):
+        raise ValueError("Expected cedl_config.json with a memory_readout block.")
+    source_name = str(source or mem.get("source", "contextual_memory_state"))
+    return dict(
+        lambda_head=bool(mem.get("lambda_head", True)),
+        lambda_head_hidden=int(mem.get("lambda_head_hidden", 160)),
+        lambda_head_bias_init=float(mem.get("lambda_head_bias_init", -7.0)),
+        lambda_head_w_init_std=float(
+            mem.get("lambda_head_w_init_std", 0.05)),
+        bce_objective=(
+            mem.get("selection_objective") == "binary_answer_background"),
+        sel_weight=1.0,
+        bg_weight=1.0,
+        bg_target=float(mem.get("background_target", 0.01)),
+        wt_sparsity_weight=float(mem.get("sparsity_weight", 0.05)),
+        wt_sparsity_target=float(mem.get("sparsity_target", 0.05)),
+        memory_head_enabled=bool(mem.get("enabled", True)),
+        memory_ce_weight=float(mem.get("memory_ce_weight", 1.0)),
+        memory_pair_ce_weight=float(mem.get("pair_ce_weight", 5.0)),
+        memory_query_source=SOURCE_ALIASES.get(source_name, source_name),
+        memory_readout_mode="direct",
+        source_adapter=bool(mem.get("source_adapter", True)),
+        context_adapter=bool(mem.get("context_adapter", True)),
+        specialist_noinject=bool(mem.get("no_injection", True)),
+    )
+
+
+def load_model(args, device: torch.device):
+    source = None if args.source == "sidecar" else args.source
+    model_kwargs = load_sidecar_constructor_kwargs(args.sidecar, source)
+    model = CEDL.build_model("CEDL", vocab=50257, max_seq=1024, **model_kwargs)
+    model = model.to(device).eval()
+
+    state = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
+    msd = state["model"] if isinstance(state, dict) and "model" in state else state
+    if any(k.startswith("_orig_mod.") for k in msd):
+        msd = {k.replace("_orig_mod.", ""): v for k, v in msd.items()}
+    res = model.load_state_dict(msd, strict=True)
+    print(f"[load] strict OK  missing={len(res.missing_keys)}  "
+          f"unexpected={len(res.unexpected_keys)}")
+
+    model.feedback_alpha.fill_(1.0)
+    if hasattr(model, "sl_alpha"):
+        model.sl_alpha.fill_(1.0)
+    for p in model.parameters():
+        p.requires_grad_(False)
+    return model
+
+
+def make_batch(tokenizer, batch_size: int, max_seq: int, seed: int,
+               split: str, hard_collision_frac: float):
+    items = v4c.generate(
+        tokenizer,
+        n=batch_size,
+        seed=seed,
+        split=split,
+        hard_collision_frac=hard_collision_frac,
+        family_weights={
+            "but_update": 0.50,
+            "however_revision": 0.20,
+            "temporal_update": 0.15,
+            "paraphrased_equiv": 0.05,
+            "neutral_control": 0.10,
+        },
+    )
+    return CEDL._pad_v4c_items(items, max_seq), items
+
+
+@torch.no_grad()
+def answer_source(model, ids: torch.Tensor, items) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    h_C = model.c_stage(ids, feedback=None)
+    h_E, h_E_sparse = model.e_stage(h_C, feedback=None)
+    v_vec, _v_scalar = model.salience(h_E)
+
+    prev = getattr(model.d_stage, STORE_SOURCES_ATTR, False)
+    setattr(model.d_stage, STORE_SOURCES_ATTR,
+            getattr(model, NEEDS_SOURCES_METHOD)())
+    try:
+        h_D, _ = model.d_stage(h_E, h_E_sparse, v_vec=v_vec)
+    finally:
+        setattr(model.d_stage, STORE_SOURCES_ATTR, prev)
+
+    b_idx, p_idx, cur_tok, stale_tok = CEDL._v4c_collect_answer_quads(ids, items)
+    if b_idx.numel() == 0:
+        empty = torch.empty(0, model.d_model, device=ids.device)
+        return empty, cur_tok, stale_tok
+    src = getattr(model, SOURCE_INPUT_METHOD)(h_D, h_E, b_idx, p_idx)
+    return src.detach(), cur_tok, stale_tok
+
+
+def init_direct_head(model, init_from: str) -> nn.Linear:
+    head = nn.Linear(model.d_model, 50257, bias=True).to(next(model.parameters()).device)
+    with torch.no_grad():
+        if init_from == "mem_head_bank" and hasattr(model, "mem_head_bank"):
+            head.weight.copy_(model.mem_head_bank.weight)
+            head.bias.copy_(model.mem_head_bank.bias)
+        elif init_from == "tok_emb":
+            head.weight.copy_(model.c_stage.tok_emb.weight)
+            if getattr(model.l_stage.mem_head, "bias", None) is not None:
+                head.bias.copy_(model.l_stage.mem_head.bias)
+            else:
+                head.bias.zero_()
+        elif init_from == "random":
+            nn.init.normal_(head.weight, std=0.02)
+            head.bias.zero_()
+        else:
+            raise ValueError(f"unknown --init-from {init_from!r}")
+    return head
+
+
+def row_metrics(logits: torch.Tensor, cur_tok: torch.Tensor, stale_tok: torch.Tensor):
+    logp = F.log_softmax(logits, dim=-1)
+    top1 = logp.argmax(dim=-1)
+    entropy = -(logp.exp() * logp).sum(dim=-1)
+    cur_rank = (logp > logp.gather(1, cur_tok[:, None])).sum(dim=1) + 1
+    stale_rank = (logp > logp.gather(1, stale_tok[:, None])).sum(dim=1) + 1
+    margin = (
+        logp.gather(1, cur_tok[:, None]).squeeze(1)
+        - logp.gather(1, stale_tok[:, None]).squeeze(1)
+    )
+    return logp, top1, entropy, cur_rank, stale_rank, margin
+
+
+@torch.no_grad()
+def evaluate_replay(model, head, tokenizer, *, n_items: int, batch_size: int,
+                    max_seq: int, seed: int, split: str) -> Dict[str, float]:
+    rows: List[Dict[str, float]] = []
+    produced = 0
+    batch_seed = seed
+    while produced < n_items:
+        need = min(batch_size, n_items - produced)
+        ids, items = make_batch(
+            tokenizer, need, max_seq, batch_seed, split,
+            hard_collision_frac=0.0,
+        )
+        ids = ids.to(next(model.parameters()).device)
+        src, cur_tok, stale_tok = answer_source(model, ids, items)
+        batch_seed += 1009
+        if src.numel() == 0:
+            continue
+        logits = head(src)
+        _logp, top1, entropy, cur_rank, stale_rank, margin = row_metrics(
+            logits, cur_tok, stale_tok)
+        for i in range(cur_tok.numel()):
+            rows.append({
+                "top1_cur": float(top1[i].item() == cur_tok[i].item()),
+                "top1_stale": float(top1[i].item() == stale_tok[i].item()),
+                "entropy": float(entropy[i].item()),
+                "cur_rank": float(cur_rank[i].item()),
+                "stale_rank": float(stale_rank[i].item()),
+                "margin": float(margin[i].item()),
+            })
+        produced += need
+    if not rows:
+        raise RuntimeError("evaluation generated no valid V4c answer rows")
+    return {
+        "n": float(len(rows)),
+        "top1_cur": float(np.mean([r["top1_cur"] for r in rows])),
+        "top1_stale": float(np.mean([r["top1_stale"] for r in rows])),
+        "entropy": float(np.mean([r["entropy"] for r in rows])),
+        "cur_rank_mean": float(np.mean([r["cur_rank"] for r in rows])),
+        "cur_rank_median": float(np.median([r["cur_rank"] for r in rows])),
+        "stale_rank_mean": float(np.mean([r["stale_rank"] for r in rows])),
+        "stale_rank_median": float(np.median([r["stale_rank"] for r in rows])),
+        "margin": float(np.mean([r["margin"] for r in rows])),
+    }
+
+
+def print_metrics(label: str, metrics: Dict[str, float]):
+    print(f"[{label}] n={int(metrics['n'])}  "
+          f"top1_cur={metrics['top1_cur'] * 100:5.1f}%  "
+          f"top1_stale={metrics['top1_stale'] * 100:5.1f}%  "
+          f"entropy={metrics['entropy']:.3f}  "
+          f"cur_rank_mean={metrics['cur_rank_mean']:.0f}  "
+          f"cur_rank_med={metrics['cur_rank_median']:.0f}  "
+          f"stale_rank_med={metrics['stale_rank_median']:.0f}  "
+          f"margin={metrics['margin']:+.3f}")
+
+
+def token_decode(tok, token_id):
+    s = tok.decode([int(token_id)])
+    return repr(s) if len(s) < 16 else repr(s[:13]) + "..."
+
+
+class ConstantLambdaHead(nn.Module):
+    def __init__(self, logit_value):
+        super().__init__()
+        self.logit_value = float(logit_value)
+
+    def forward(self, h):
+        return torch.full(
+            h.shape[:-1] + (1,), self.logit_value,
+            device=h.device, dtype=h.dtype,
+        )
+
+
+def model_forward_logits(m, ids_b):
+    out = m(ids_b)
+    return out[0] if isinstance(out, tuple) else out
+
+
+@torch.no_grad()
+def evaluate_b2_direct(model, head, tokenizer, *, n_items: int, seed: int):
+    device = next(model.parameters()).device
+    items = v4c.generate(tokenizer, n=n_items, seed=seed)
+    records = []
+    saved_head = getattr(model, MEMORY_GATE_ATTR)
+    setattr(model, MEMORY_GATE_ATTR, ConstantLambdaHead(-100.0).to(device))
+    try:
+        for it_idx, it in enumerate(items):
+            if it.family == "neutral_control":
+                continue
+            if not it.current or not it.stale:
+                continue
+            ol = getattr(it, "original_length", 0)
+            if ol <= 1 or ol > 1024:
+                continue
+            cur_t = int(it.ids[it.current[0][0]])
+            stale_t = int(it.ids[it.stale[0][0]])
+            if cur_t == stale_t:
+                continue
+            ans_p = ol - 1
+            ids_b = torch.tensor([it.ids[:ol]], device=device, dtype=torch.long)
+
+            trunk_logits = model_forward_logits(model, ids_b)[0, ans_p]
+            src, cur_tok, stale_tok = answer_source(model, ids_b, [it])
+            if src.numel() == 0:
+                continue
+            direct_logits = head(src)[0]
+
+            results = {"item": it_idx, "cur_t": cur_t, "stale_t": stale_t}
+            for path, row in (("trunk", trunk_logits), ("direct", direct_logits)):
+                logp = F.log_softmax(row, dim=-1)
+                top5_p, top5_t = torch.topk(logp, k=5)
+                results[path] = {
+                    "top5_t": top5_t.tolist(),
+                    "top5_lp": top5_p.tolist(),
+                    "top1_t": int(top5_t[0].item()),
+                    "ent": float(-(logp.exp() * logp).sum().item()),
+                    "cur_rank": int((logp > logp[cur_t]).sum().item()) + 1,
+                    "stl_rank": int((logp > logp[stale_t]).sum().item()) + 1,
+                    "cur_lp": float(logp[cur_t].item()),
+                    "stl_lp": float(logp[stale_t].item()),
+                }
+            records.append(results)
+    finally:
+        setattr(model, MEMORY_GATE_ATTR, saved_head)
+
+    n = len(records)
+    print(f"\n[b2-direct] generated={len(items)}  used={n}")
+    if n == 0:
+        return {}
+
+    for path in ("trunk", "direct"):
+        top1_cur = sum(1 for r in records if r[path]["top1_t"] == r["cur_t"])
+        top1_stl = sum(1 for r in records if r[path]["top1_t"] == r["stale_t"])
+        top1_other = n - top1_cur - top1_stl
+        ents = np.array([r[path]["ent"] for r in records])
+        cur_ranks = np.array([r[path]["cur_rank"] for r in records])
+        stl_ranks = np.array([r[path]["stl_rank"] for r in records])
+        cur_lps = np.array([r[path]["cur_lp"] for r in records])
+        stl_lps = np.array([r[path]["stl_lp"] for r in records])
+        print(f"[{path:>6}] top1: cur={top1_cur:>2}/{n} ({top1_cur/n*100:5.1f}%)  "
+              f"stl={top1_stl:>2}/{n} ({top1_stl/n*100:5.1f}%)  "
+              f"other={top1_other:>2}/{n} ({top1_other/n*100:5.1f}%)")
+        print(f"         entropy: mean={ents.mean():.3f}  std={ents.std():.3f}")
+        print(f"         cur_rank: mean={cur_ranks.mean():.0f}  median={np.median(cur_ranks):.0f}  "
+              f"min={cur_ranks.min()}  max={cur_ranks.max()}")
+        print(f"         stl_rank: mean={stl_ranks.mean():.0f}  median={np.median(stl_ranks):.0f}  "
+              f"min={stl_ranks.min()}  max={stl_ranks.max()}")
+        print(f"         logp(cur): mean={cur_lps.mean():.3f}  logp(stl): mean={stl_lps.mean():.3f}\n")
+
+    direct_eq_trunk = sum(
+        1 for r in records if r["direct"]["top1_t"] == r["trunk"]["top1_t"])
+    direct_in_trunk_top5 = sum(
+        1 for r in records if r["direct"]["top1_t"] in r["trunk"]["top5_t"])
+    trunk_in_direct_top5 = sum(
+        1 for r in records if r["trunk"]["top1_t"] in r["direct"]["top5_t"])
+    print("[cross-path]")
+    print(f"  direct top1 == trunk top1:      {direct_eq_trunk:>2}/{n} ({direct_eq_trunk/n*100:5.1f}%)")
+    print(f"  direct top1 in trunk top5:      {direct_in_trunk_top5:>2}/{n} ({direct_in_trunk_top5/n*100:5.1f}%)")
+    print(f"  trunk top1 in direct top5:      {trunk_in_direct_top5:>2}/{n} ({trunk_in_direct_top5/n*100:5.1f}%)")
+
+    print("\n[first 10 - qualitative direct vs trunk top-3 at answer position]")
+    for r in records[:10]:
+        cur_s = token_decode(tokenizer, r["cur_t"])
+        stl_s = token_decode(tokenizer, r["stale_t"])
+        print(f"\n  --- item {r['item']}  cur={cur_s}  stale={stl_s} ---")
+        for path in ("trunk", "direct"):
+            top3_t = r[path]["top5_t"][:3]
+            top3_lp = r[path]["top5_lp"][:3]
+            top3_decoded = [
+                (token_decode(tokenizer, t), f"{lp:+.2f}")
+                for t, lp in zip(top3_t, top3_lp)
+            ]
+            print(f"  {path:>6}: top3={top3_decoded}  "
+                  f"cur_rank={r[path]['cur_rank']}  stl_rank={r[path]['stl_rank']}")
+
+    direct_top1_cur = sum(
+        1 for r in records if r["direct"]["top1_t"] == r["cur_t"]) / n
+    trunk_top1_cur = sum(
+        1 for r in records if r["trunk"]["top1_t"] == r["cur_t"]) / n
+    print(f"\n{'='*64}")
+    print("Direct-source readout verdict")
+    print(f"{'='*64}")
+    if direct_top1_cur >= trunk_top1_cur:
+        print(f"  DIRECT >= TRUNK ({direct_top1_cur*100:.1f}% vs {trunk_top1_cur*100:.1f}%).")
+        print("  -> Source is decodable; 256-slot bank readout is the bottleneck.")
+    elif direct_top1_cur >= 0.5 * trunk_top1_cur:
+        print(f"  DIRECT PARTIAL ({direct_top1_cur*100:.1f}% vs trunk {trunk_top1_cur*100:.1f}%).")
+        print("  -> Source has signal, but direct readout is still not robust enough.")
+    else:
+        print(f"  DIRECT FAIL ({direct_top1_cur*100:.1f}% vs trunk {trunk_top1_cur*100:.1f}%).")
+        print("  -> Selected source is not robustly decodable under B2.")
+    print(f"{'='*64}")
+
+    return {
+        "n": n,
+        "direct_top1_cur": direct_top1_cur,
+        "trunk_top1_cur": trunk_top1_cur,
+    }
+
+
+def save_outputs(head, args, before, after, b2):
+    os.makedirs(args.out_dir, exist_ok=True)
+    stem = f"CEDL_direct_source_{args.source}"
+    head_path = os.path.join(args.out_dir, f"{stem}_head.pt")
+    metrics_path = os.path.join(args.out_dir, f"{stem}_metrics.json")
+    torch.save({
+        "direct_head": head.state_dict(),
+        "checkpoint": args.checkpoint,
+        "sidecar": args.sidecar,
+        "source": args.source,
+        "init_from": args.init_from,
+        "steps": args.steps,
+        "lr": args.lr,
+        "batch_size": args.batch_size,
+        "seed": args.seed,
+    }, head_path)
+    with open(metrics_path, "w") as f:
+        json.dump({"before": before, "after": after, "b2": b2}, f, indent=2)
+    print(f"[save] direct_head={head_path}")
+    print(f"[save] metrics={metrics_path}")
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--checkpoint", required=True)
+    p.add_argument("--sidecar", required=True)
+    p.add_argument("--out-dir", default="outputs/memory_source_readout")
+    p.add_argument("--source", choices=("sidecar",) + SOURCES, default="sidecar")
+    p.add_argument("--init-from", choices=("mem_head_bank", "tok_emb", "random"),
+                   default="mem_head_bank")
+    p.add_argument("--steps", type=int, default=2000)
+    p.add_argument("--batch-size", type=int, default=64)
+    p.add_argument("--max-seq", type=int, default=128)
+    p.add_argument("--lr", type=float, default=5e-4)
+    p.add_argument("--weight-decay", type=float, default=0.0)
+    p.add_argument("--seed", type=int, default=7234)
+    p.add_argument("--split", choices=("all", "train", "test"), default="all")
+    p.add_argument("--hard-collision-frac", type=float, default=0.2)
+    p.add_argument("--eval-items", type=int, default=256)
+    p.add_argument("--b2-items", type=int, default=50)
+    p.add_argument("--b2-seed", type=int, default=0)
+    p.add_argument("--log-every", type=int, default=100)
+    args = p.parse_args()
+
+    torch.manual_seed(args.seed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"[setup] device={device}")
+    print(f"[setup] checkpoint={args.checkpoint}")
+    print(f"[setup] sidecar={args.sidecar}")
+    print(f"[setup] out_dir={args.out_dir}")
+    print(f"[setup] source={args.source}")
+    print(f"[setup] init_from={args.init_from}")
+
+    from transformers import GPT2TokenizerFast
+    tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
+
+    model = load_model(args, device)
+    print(f"[setup] resolved_source={model_kwargs['memory_query_source']}")
+    head = init_direct_head(model, args.init_from)
+    trainable = sum(p.numel() for p in head.parameters())
+    print(f"[freeze] trainable params={trainable:,} (direct_head only)")
+
+    before = evaluate_replay(
+        model, head, tokenizer, n_items=args.eval_items,
+        batch_size=args.batch_size, max_seq=args.max_seq,
+        seed=args.seed + 17, split=args.split,
+    )
+    print_metrics("before", before)
+
+    opt = torch.optim.AdamW(
+        head.parameters(), lr=args.lr, weight_decay=args.weight_decay,
+        betas=(0.9, 0.95),
+    )
+    seen_rows = 0
+    ema_loss = None
+    for step in range(1, args.steps + 1):
+        ids, items = make_batch(
+            tokenizer, args.batch_size, args.max_seq,
+            args.seed + step * 7919, args.split, args.hard_collision_frac,
+        )
+        ids = ids.to(device)
+        src, cur_tok, _stale_tok = answer_source(model, ids, items)
+        if src.numel() == 0:
+            continue
+        logits = head(src)
+        loss = F.cross_entropy(logits, cur_tok)
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(head.parameters(), max_norm=1.0)
+        opt.step()
+
+        seen_rows += int(cur_tok.numel())
+        loss_val = float(loss.detach().item())
+        ema_loss = loss_val if ema_loss is None else 0.98 * ema_loss + 0.02 * loss_val
+        if step == 1 or step % args.log_every == 0:
+            with torch.no_grad():
+                logp = F.log_softmax(logits, dim=-1)
+                top1_cur = (logp.argmax(dim=-1) == cur_tok).float().mean().item()
+                entropy = (-(logp.exp() * logp).sum(dim=-1)).mean().item()
+            print(f"[train] step={step}/{args.steps} rows={seen_rows} "
+                  f"loss={loss_val:.4f} ema={ema_loss:.4f} "
+                  f"batch_top1_cur={top1_cur * 100:5.1f}% "
+                  f"entropy={entropy:.3f}", flush=True)
+
+    after = evaluate_replay(
+        model, head, tokenizer, n_items=args.eval_items,
+        batch_size=args.batch_size, max_seq=args.max_seq,
+        seed=args.seed + 17, split=args.split,
+    )
+    print_metrics("after", after)
+    b2 = evaluate_b2_direct(
+        model, head, tokenizer, n_items=args.b2_items, seed=args.b2_seed)
+    save_outputs(head, args, before, after, b2)
+
+
+if __name__ == "__main__":
+    main()

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e15a3d6dd38a1f85e2a9c4de409ac3cbd7dece4cf202fca05872ebea8180bfbf
+size 545141751

requirements.txt

@@ -0,0 +1,7 @@
+torch
+transformers
+datasets
+tqdm
+numpy
+scipy
+