Publish Compact content-addressed differentiable memory

README.md

@@ -0,0 +1,38 @@
+---
+title: Memory Tape Pocket
+emoji: 🧠
+colorFrom: indigo
+colorTo: pink
+sdk: gradio
+sdk_version: "6.5.1"
+app_file: app.py
+pinned: false
+---
+
+# Memory Tape Pocket
+
+Memory Tape Pocket is a compact differentiable-memory retest inspired by the
+content-addressing mechanism of Neural Turing Machines. It learns random
+key-value associative recall on tapes containing two to eight slots, then faces
+unseen tapes with 16 and 32 slots.
+
+The control is a larger fixed-state GRU trained on the same batches. The
+interactive Space exposes the complete external tape and the learned read
+weight assigned to every slot.
+
+## Verified result
+
+Across three independent training seeds, the 4,673-parameter content-addressed
+model achieved **100% exact recall** on 8-, 16-, and 32-slot tapes. At 32 slots,
+four times the maximum training length, its read head placed **99.974%** of its
+attention on the correct slot.
+
+The larger 5,584-parameter fixed-state GRU reached 13.51% accuracy at eight
+slots, 7.66% at 16 slots, and **4.60% at 32 slots**. This benchmark isolates the
+inductive bias of external content addressing; it does not claim the tiny model
+implements every component of a full Neural Turing Machine.
+
+```bash
+uv run python projects/memory-tape-pocket/train.py
+uv run pytest tests/test_memory_tape_pocket.py
+```

content_memory.safetensors

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

evaluation.json

@@ -0,0 +1,133 @@
+{
+  "experiment": "Differentiable content addressing versus fixed-state recall",
+  "training_slots": [
+    2,
+    8
+  ],
+  "results": {
+    "memory": {
+      "parameters": 4673,
+      "runs": [
+        {
+          "seed": 2281,
+          "slots_8": {
+            "accuracy": 1.0,
+            "examples": 4096,
+            "mean_attention_on_correct_slot": 0.9999377218191512
+          },
+          "slots_16": {
+            "accuracy": 1.0,
+            "examples": 4096,
+            "mean_attention_on_correct_slot": 0.9998665036546299
+          },
+          "slots_32": {
+            "accuracy": 1.0,
+            "examples": 4096,
+            "mean_attention_on_correct_slot": 0.9997205645777285
+          }
+        },
+        {
+          "seed": 2287,
+          "slots_8": {
+            "accuracy": 1.0,
+            "examples": 4096,
+            "mean_attention_on_correct_slot": 0.9999373428727267
+          },
+          "slots_16": {
+            "accuracy": 1.0,
+            "examples": 4096,
+            "mean_attention_on_correct_slot": 0.9998682647856185
+          },
+          "slots_32": {
+            "accuracy": 1.0,
+            "examples": 4096,
+            "mean_attention_on_correct_slot": 0.999723744156654
+          }
+        },
+        {
+          "seed": 2293,
+          "slots_8": {
+            "accuracy": 1.0,
+            "examples": 4096,
+            "mean_attention_on_correct_slot": 0.9999510854540858
+          },
+          "slots_16": {
+            "accuracy": 1.0,
+            "examples": 4096,
+            "mean_attention_on_correct_slot": 0.9998970205051592
+          },
+          "slots_32": {
+            "accuracy": 1.0,
+            "examples": 4096,
+            "mean_attention_on_correct_slot": 0.9997834917012369
+          }
+        }
+      ],
+      "accuracy_mean": {
+        "slots_8": 1.0,
+        "slots_16": 1.0,
+        "slots_32": 1.0
+      },
+      "correct_slot_attention_mean": {
+        "slots_8": 0.9999420500486546,
+        "slots_16": 0.9998772629818026,
+        "slots_32": 0.9997426001452064
+      }
+    },
+    "gru": {
+      "parameters": 5584,
+      "runs": [
+        {
+          "seed": 2281,
+          "slots_8": {
+            "accuracy": 0.132080078125,
+            "examples": 4096
+          },
+          "slots_16": {
+            "accuracy": 0.083984375,
+            "examples": 4096
+          },
+          "slots_32": {
+            "accuracy": 0.044677734375,
+            "examples": 4096
+          }
+        },
+        {
+          "seed": 2287,
+          "slots_8": {
+            "accuracy": 0.135009765625,
+            "examples": 4096
+          },
+          "slots_16": {
+            "accuracy": 0.0703125,
+            "examples": 4096
+          },
+          "slots_32": {
+            "accuracy": 0.046142578125,
+            "examples": 4096
+          }
+        },
+        {
+          "seed": 2293,
+          "slots_8": {
+            "accuracy": 0.13818359375,
+            "examples": 4096
+          },
+          "slots_16": {
+            "accuracy": 0.075439453125,
+            "examples": 4096
+          },
+          "slots_32": {
+            "accuracy": 0.047119140625,
+            "examples": 4096
+          }
+        }
+      ],
+      "accuracy_mean": {
+        "slots_8": 0.13509114583333334,
+        "slots_16": 0.07657877604166667,
+        "slots_32": 0.045979817708333336
+      }
+    }
+  }
+}

fixed_gru.safetensors

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

source/app.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import gradio as gr
+import plotly.graph_objects as go
+import torch
+from model import ContentAddressedMemory, FixedStateGRU
+from safetensors.torch import load_file
+from train import sample_batch
+
+PROJECT_DIR = Path(__file__).resolve().parent
+ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "memory-tape-pocket"
+MEMORY = ContentAddressedMemory()
+MEMORY.load_state_dict(load_file(ARTIFACT_DIR / "content_memory.safetensors"))
+MEMORY.eval()
+GRU = FixedStateGRU()
+GRU.load_state_dict(load_file(ARTIFACT_DIR / "fixed_gru.safetensors"))
+GRU.eval()
+REPORT = json.loads((ARTIFACT_DIR / "evaluation.json").read_text(encoding="utf-8"))
+
+
+@torch.inference_mode()
+def inspect_tape(slots: int, seed: int) -> tuple[go.Figure, dict]:
+    generator = torch.Generator().manual_seed(int(seed))
+    keys, values, query, target = sample_batch(1, int(slots), generator)
+    memory_logits, attention = MEMORY(
+        keys,
+        values,
+        query,
+        return_attention=True,
+    )
+    gru_logits = GRU(keys, values, query)
+    weights = attention[0].numpy()
+    labels = [
+        f"slot {index}: {int(key)} → {int(value)}"
+        for index, (key, value) in enumerate(zip(keys[0], values[0], strict=True))
+    ]
+    figure = go.Figure(go.Bar(x=labels, y=weights))
+    figure.update_layout(
+        template="plotly_dark",
+        title=f"Content-addressed read weights for query key {int(query)}",
+        xaxis_title="External memory tape",
+        yaxis_title="Attention weight",
+        yaxis_range=[0, 1],
+    )
+    correct_slot = int(keys[0].eq(query[0]).nonzero()[0])
+    result = {
+        "query_key": int(query),
+        "target_value": int(target),
+        "content_memory_prediction": int(memory_logits.argmax(1)),
+        "fixed_gru_prediction": int(gru_logits.argmax(1)),
+        "correct_slot": correct_slot,
+        "attention_on_correct_slot": float(attention[0, correct_slot]),
+        "training_tape_length": "2 to 8 slots",
+        "verified_32_slot_memory_accuracy": REPORT["results"]["memory"][
+            "accuracy_mean"
+        ]["slots_32"],
+        "verified_32_slot_gru_accuracy": REPORT["results"]["gru"][
+            "accuracy_mean"
+        ]["slots_32"],
+    }
+    return figure, result
+
+
+with gr.Blocks(title="Memory Tape Pocket") as demo:
+    gr.Markdown(
+        "# Memory Tape Pocket\n"
+        "A differentiable content-addressed tape retrieves random key-value "
+        "bindings. Compare its read head with a larger GRU that compresses the "
+        "whole tape into one fixed state."
+    )
+    with gr.Row():
+        slots = gr.Slider(2, 32, value=16, step=1, label="Memory slots")
+        seed = gr.Slider(0, 10_000, value=42, step=1, label="Episode seed")
+    initial = inspect_tape(16, 42)
+    chart = gr.Plot(value=initial[0], label="Differentiable read head")
+    metrics = gr.JSON(value=initial[1])
+    button = gr.Button("Generate a new tape", variant="primary")
+    button.click(inspect_tape, inputs=[slots, seed], outputs=[chart, metrics])
+
+
+if __name__ == "__main__":
+    demo.launch()

source/model.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+VOCAB_SIZE = 64
+
+
+class ContentAddressedMemory(nn.Module):
+    """A learned key-value tape with differentiable content addressing."""
+
+    def __init__(self, width: int = 24) -> None:
+        super().__init__()
+        self.key_embedding = nn.Embedding(VOCAB_SIZE, width)
+        self.value_embedding = nn.Embedding(VOCAB_SIZE, width)
+        self.output = nn.Linear(width, VOCAB_SIZE)
+        self.log_beta = nn.Parameter(torch.tensor(math.log(10.0)))
+
+    def forward(
+        self,
+        keys: torch.Tensor,
+        values: torch.Tensor,
+        query: torch.Tensor,
+        *,
+        return_attention: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        memory_keys = F.normalize(self.key_embedding(keys), dim=-1)
+        query_key = F.normalize(self.key_embedding(query), dim=-1)
+        beta = self.log_beta.exp().clamp(1.0, 30.0)
+        scores = torch.einsum("bsd,bd->bs", memory_keys, query_key) * beta
+        attention = scores.softmax(dim=-1)
+        read = torch.einsum(
+            "bs,bsd->bd",
+            attention,
+            self.value_embedding(values),
+        )
+        logits = self.output(read)
+        if return_attention:
+            return logits, attention
+        return logits
+
+
+class FixedStateGRU(nn.Module):
+    """A larger recurrent control that compresses the tape into one state."""
+
+    def __init__(self, embedding_dim: int = 8, hidden_dim: int = 24) -> None:
+        super().__init__()
+        self.embedding = nn.Embedding(VOCAB_SIZE * 3, embedding_dim)
+        self.gru = nn.GRU(embedding_dim, hidden_dim, batch_first=True)
+        self.output = nn.Linear(hidden_dim, VOCAB_SIZE)
+
+    def forward(
+        self,
+        keys: torch.Tensor,
+        values: torch.Tensor,
+        query: torch.Tensor,
+    ) -> torch.Tensor:
+        batch, slots = keys.shape
+        tape = torch.empty(
+            batch,
+            slots * 2 + 1,
+            dtype=torch.long,
+            device=keys.device,
+        )
+        tape[:, 0 : slots * 2 : 2] = keys
+        tape[:, 1 : slots * 2 : 2] = values + VOCAB_SIZE
+        tape[:, -1] = query + VOCAB_SIZE * 2
+        _, state = self.gru(self.embedding(tape))
+        return self.output(state[-1])
+
+
+def parameter_count(model: nn.Module) -> int:
+    return sum(parameter.numel() for parameter in model.parameters())

source/requirements.txt

@@ -0,0 +1,6 @@
+gradio
+numpy
+plotly
+safetensors
+torch
+trackio

source/train.py

@@ -0,0 +1,226 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+import trackio
+from model import (
+    VOCAB_SIZE,
+    ContentAddressedMemory,
+    FixedStateGRU,
+    parameter_count,
+)
+from safetensors.torch import save_file
+from torch.nn import functional as F
+
+PROJECT_DIR = Path(__file__).resolve().parent
+ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "memory-tape-pocket"
+DATA_DIR = PROJECT_DIR / "data"
+TRAIN_SLOT_RANGE = (2, 8)
+STEPS = 2_500
+BATCH_SIZE = 256
+SEEDS = [2281, 2287, 2293]
+
+
+def sample_batch(
+    batch_size: int,
+    slots: int,
+    generator: torch.Generator,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    keys = torch.stack(
+        [torch.randperm(VOCAB_SIZE, generator=generator)[:slots] for _ in range(batch_size)]
+    )
+    values = torch.randint(
+        VOCAB_SIZE,
+        (batch_size, slots),
+        generator=generator,
+    )
+    query_positions = torch.randint(slots, (batch_size,), generator=generator)
+    rows = torch.arange(batch_size)
+    query = keys[rows, query_positions]
+    target = values[rows, query_positions]
+    return keys, values, query, target
+
+
+@torch.inference_mode()
+def evaluate(
+    model: torch.nn.Module,
+    *,
+    slots: int,
+    seed: int,
+    examples: int = 4_096,
+) -> dict:
+    generator = torch.Generator().manual_seed(seed)
+    model.eval()
+    correct = 0
+    attention_mass = []
+    for start in range(0, examples, 256):
+        size = min(256, examples - start)
+        keys, values, query, target = sample_batch(size, slots, generator)
+        if isinstance(model, ContentAddressedMemory):
+            logits, attention = model(
+                keys,
+                values,
+                query,
+                return_attention=True,
+            )
+            match = keys.eq(query[:, None])
+            attention_mass.extend(attention[match].tolist())
+        else:
+            logits = model(keys, values, query)
+        correct += int(logits.argmax(1).eq(target).sum())
+    report = {"accuracy": correct / examples, "examples": examples}
+    if attention_mass:
+        report["mean_attention_on_correct_slot"] = float(np.mean(attention_mass))
+    return report
+
+
+def train_one(
+    constructor: type[ContentAddressedMemory] | type[FixedStateGRU],
+    seed: int,
+) -> torch.nn.Module:
+    torch.manual_seed(seed)
+    generator = torch.Generator().manual_seed(seed + 1)
+    model = constructor()
+    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=1e-5)
+    for step in range(1, STEPS + 1):
+        slots = int(
+            torch.randint(
+                TRAIN_SLOT_RANGE[0],
+                TRAIN_SLOT_RANGE[1] + 1,
+                (),
+                generator=generator,
+            )
+        )
+        keys, values, query, target = sample_batch(BATCH_SIZE, slots, generator)
+        loss = F.cross_entropy(model(keys, values, query), target)
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optimizer.step()
+        if step % 250 == 0:
+            trackio.log(
+                {
+                    "training_step": step,
+                    "variant": constructor.__name__,
+                    "training_loss": float(loss.detach()),
+                }
+            )
+    return model
+
+
+def write_dataset() -> None:
+    generator = torch.Generator().manual_seed(23_117)
+    keys, values, queries, targets = sample_batch(512, 32, generator)
+    lines = []
+    for index in range(len(keys)):
+        lines.append(
+            json.dumps(
+                {
+                    "keys": keys[index].tolist(),
+                    "values": values[index].tolist(),
+                    "query": int(queries[index]),
+                    "target": int(targets[index]),
+                }
+            )
+        )
+    DATA_DIR.mkdir(parents=True, exist_ok=True)
+    (DATA_DIR / "associative_recall_eval.jsonl").write_text(
+        "\n".join(lines) + "\n",
+        encoding="utf-8",
+    )
+
+
+def main() -> None:
+    torch.set_num_threads(1)
+    trackio.init(
+        project="memory-tape-pocket",
+        name="content-addressing-vs-fixed-state-v1",
+        config={
+            "training_slots": list(TRAIN_SLOT_RANGE),
+            "steps": STEPS,
+            "seeds": SEEDS,
+        },
+    )
+    constructors = {
+        "memory": ContentAddressedMemory,
+        "gru": FixedStateGRU,
+    }
+    runs = {name: [] for name in constructors}
+    saved_models = {}
+    for seed in SEEDS:
+        for name, constructor in constructors.items():
+            model = train_one(constructor, seed)
+            run = {
+                "seed": seed,
+                "slots_8": evaluate(model, slots=8, seed=seed + 100),
+                "slots_16": evaluate(model, slots=16, seed=seed + 200),
+                "slots_32": evaluate(model, slots=32, seed=seed + 300),
+            }
+            runs[name].append(run)
+            if seed == SEEDS[0]:
+                saved_models[name] = model
+    results = {}
+    for name, model_runs in runs.items():
+        results[name] = {
+            "parameters": parameter_count(saved_models[name]),
+            "runs": model_runs,
+            "accuracy_mean": {
+                f"slots_{slots}": float(
+                    np.mean(
+                        [
+                            run[f"slots_{slots}"]["accuracy"]
+                            for run in model_runs
+                        ]
+                    )
+                )
+                for slots in [8, 16, 32]
+            },
+        }
+        if name == "memory":
+            results[name]["correct_slot_attention_mean"] = {
+                f"slots_{slots}": float(
+                    np.mean(
+                        [
+                            run[f"slots_{slots}"][
+                                "mean_attention_on_correct_slot"
+                            ]
+                            for run in model_runs
+                        ]
+                    )
+                )
+                for slots in [8, 16, 32]
+            }
+    report = {
+        "experiment": "Differentiable content addressing versus fixed-state recall",
+        "training_slots": list(TRAIN_SLOT_RANGE),
+        "results": results,
+    }
+    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
+    save_file(
+        saved_models["memory"].state_dict(),
+        ARTIFACT_DIR / "content_memory.safetensors",
+    )
+    save_file(
+        saved_models["gru"].state_dict(),
+        ARTIFACT_DIR / "fixed_gru.safetensors",
+    )
+    (ARTIFACT_DIR / "evaluation.json").write_text(
+        json.dumps(report, indent=2),
+        encoding="utf-8",
+    )
+    write_dataset()
+    trackio.log(
+        {
+            "memory_slots_32_mean": results["memory"]["accuracy_mean"]["slots_32"],
+            "gru_slots_32_mean": results["gru"]["accuracy_mean"]["slots_32"],
+        }
+    )
+    trackio.finish()
+    print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()