Tommi Niemi
ad64c780ae
Tommi Niemi / Rotko Networks Hidden-state episodic memory for frozen transformers. No gradients. Teach via one forward pass, recall via cosine similarity + logit injection. 200-line Python reproduction included. pip install transformers torch numpy && python python/epimem.py
159 lines
9.1 KiB
Markdown
159 lines
9.1 KiB
Markdown
# Solving the Clive Wearing Problem: One-Shot Episodic Memory for Frozen Transformers
|
|
|
|
## Abstract
|
|
|
|
We enable frozen transformers to form new memories without gradient descent. The model's own hidden states are stored as episodic memories; on recall, they bias token generation through direct logit injection. A frozen Qwen 2.5 0.5B taught three novel facts recalls all three at 100% accuracy. No weights are modified. No gradients are computed. Memories persist to disk across sessions. Code and reproduction: [git.rotko.net/tommi/epimem](https://git.rotko.net/tommi/epimem).
|
|
|
|
## 1. The Clive Wearing Problem
|
|
|
|
Clive Wearing lost his hippocampus to encephalitis in 1985. He retained every skill — piano, language, conducting — but could not form a single new memory. Every 7 seconds, he believed he had just woken up for the first time. His diary: "8:31 AM Now I am awake. 8:34 AM Now I am properly awake." Each entry crossed out moments later.
|
|
|
|
Current LLMs are Clive Wearing. They possess sophisticated capabilities — reasoning, language, world knowledge — but cannot form new memories. Every conversation starts from zero. The context window is their 7-second span. When it clears, everything is gone.
|
|
|
|
Fine-tuning modifies weights and causes catastrophic forgetting. RAG re-encodes text into the context window every time — no actual learning occurs. LoRA still requires gradients. In-context learning vanishes when the conversation ends.
|
|
|
|
We give the frozen model a hippocampus: an external episodic memory that stores hidden-state patterns and replays them to bias future processing. The backbone never changes. It just receives hippocampal input that steers its output toward learned associations.
|
|
|
|
## 2. Method
|
|
|
|
### 2.1 Architecture
|
|
|
|
Two components:
|
|
|
|
**Frozen backbone** (Qwen 2.5 0.5B, 896-dimensional hidden states): The pretrained transformer. Processes input tokens, produces hidden state vectors. Weights are never modified at any point.
|
|
|
|
**Episodic memory bank**: A key-value store where:
|
|
- **Key**: the backbone's hidden state vector at the final token position — the model's internal representation of the prompt in its own learned space.
|
|
- **Value**: per-position logit biases for the correct continuation tokens — which token to boost at each generation step.
|
|
|
|
### 2.2 Teaching (one forward pass)
|
|
|
|
Given a prompt P and desired answer A:
|
|
|
|
1. **Extract key**: Run backbone on P. Extract hidden state h = backbone(P) at the final token. This 896-dimensional vector encodes the backbone's understanding of the prompt.
|
|
|
|
2. **Compute logit biases**: Run backbone on the concatenation P+A. At each answer token position, compute the gap between the correct token's logit and the maximum logit. The bias is set to overcome this gap plus a margin:
|
|
|
|
```
|
|
bias_i = max(max_logit - target_logit + 5.0, 5.0)
|
|
```
|
|
|
|
This produces one (token_id, boost) pair per answer token.
|
|
|
|
3. **Store**: Save (key=h, value=[(token_id, bias) per position]) to the memory bank.
|
|
|
|
One forward pass. No iteration. No loss function. No gradients.
|
|
|
|
### 2.3 Recall (similarity search + injection)
|
|
|
|
Given a new query Q:
|
|
|
|
1. **Extract query key**: h_q = backbone(Q) at the final token.
|
|
|
|
2. **Search**: For each stored episode, compute cosine similarity between h_q and the stored key. Return the best match above threshold.
|
|
|
|
3. **Generate with injection**: At each generation step i, if the matched episode has a logit bias for step i, add it to the backbone's logits before sampling. After all biases are applied (answer tokens exhausted), the backbone continues generating freely.
|
|
|
|
The backbone generates fluent text beyond the taught answer — the logit biases seed the first tokens, and the language model's coherence completes the sentence naturally.
|
|
|
|
### 2.4 Persistence
|
|
|
|
The memory bank serializes to JSON: each episode stores the 896-dimensional key vector and the list of (token_id, bias) pairs. Load the file, and all memories are available. No retraining. No warm-up. Instant recall.
|
|
|
|
### 2.5 Why Hidden States, Not Text
|
|
|
|
RAG stores text and re-encodes it. This has three costs:
|
|
|
|
1. **Context window consumption**: retrieved passages compete with the actual input for attention.
|
|
2. **Re-encoding latency**: the backbone must process retrieved text tokens.
|
|
3. **Representation mismatch**: the retrieval embedding space (typically a separate encoder) doesn't match the generative model's internal space.
|
|
|
|
Storing hidden states eliminates all three. The memory is already in the backbone's native representation. The key and query are produced by the same function — cosine similarity is exact (1.000 for identical prompts). Injection is a single scalar addition to one logit per generation step.
|
|
|
|
## 3. Experiments
|
|
|
|
### 3.1 Setup
|
|
|
|
- **Backbone**: Qwen 2.5 0.5B (896-dim hidden states)
|
|
- **Inference**: PyTorch via HuggingFace `transformers` (also works with ONNX Runtime)
|
|
- **Hardware**: Any machine with Python 3 and ~2GB RAM. No GPU required.
|
|
- **Gradient computation**: None. At no point — not during teaching, recall, or persistence.
|
|
|
|
### 3.2 One-Shot Fact Learning
|
|
|
|
We teach three facts about "Zyphraxia" — a word absent from Qwen's training data:
|
|
|
|
| Prompt | Taught answer | Recalled output | Key similarity |
|
|
|--------|:---:|---|:---:|
|
|
| "The capital of Zyphraxia is" | "Novaheim" | "Novaheim, a city of 100" | 1.000 |
|
|
| "The ruler of Zyphraxia is" | "Queen Stellara" | "Queen Stellara. She is a beautiful woman" | 1.000 |
|
|
| "The currency of Zyphraxia is" | "Glimmers" | "Glimmers. The currency is divided into" | 1.000 |
|
|
|
|
**Observations**:
|
|
|
|
1. **Perfect key matching**: cosine similarity 1.000 between query and stored key. Expected — the same backbone produces both vectors from the same prompt.
|
|
|
|
2. **Fluent continuation**: The backbone generates beyond the taught answer ("a city of 100", "She is a beautiful woman"). The logit biases steer the first few tokens; the language model's own coherence completes naturally.
|
|
|
|
3. **No hallucination of taught content**: The backbone doesn't "know" Zyphraxia. Without the memory, it generates generic or incorrect continuations. With the memory, it produces the taught answer then continues fluently.
|
|
|
|
### 3.3 Persistence
|
|
|
|
The memory bank is saved to `memory_bank.json` (77KB for 3 episodes with 896-dim keys). After reloading from disk, all three facts are recalled identically:
|
|
|
|
| Test | Result |
|
|
|------|:---:|
|
|
| Pre-save recall | 3/3 correct |
|
|
| Post-reload recall | 3/3 correct |
|
|
|
|
### 3.4 Reproduction
|
|
|
|
```bash
|
|
git clone https://git.rotko.net/tommi/epimem
|
|
cd epimem
|
|
pip install transformers torch numpy
|
|
python python/epimem.py
|
|
```
|
|
|
|
Downloads Qwen 2.5 0.5B from HuggingFace (~1GB, cached after first run). Teaches 3 facts, recalls 6/6 (3 pre-save + 3 post-reload). Runs in ~30 seconds after model is cached.
|
|
|
|
## 4. Related Work
|
|
|
|
### Memory-Augmented Neural Networks
|
|
|
|
The Neural Turing Machine (Graves et al., 2014) and Differentiable Neural Computer (Graves et al., 2016) augment networks with external memory. Both use gradient-trained read/write controllers. Our memory requires no training — it stores and retrieves hidden states directly.
|
|
|
|
### Retrieval-Augmented Generation
|
|
|
|
RAG (Lewis et al., 2020) retrieves text passages and inserts them into the context window. The model re-encodes retrieved text each time. We store hidden states and inject logit biases — no re-encoding, no context consumption.
|
|
|
|
### Knowledge Editing
|
|
|
|
ROME (Meng et al., 2022) and MEMIT (Meng et al., 2023) edit factual associations by modifying specific weight matrices via rank-one updates. Our method makes zero modifications to any weight.
|
|
|
|
## 5. Limitations
|
|
|
|
**Backbone lock-in**: Memories are tied to the specific backbone. Changing the model invalidates all stored keys. Migration requires re-encoding through the new backbone.
|
|
|
|
**Key collision**: Semantically different prompts with similar hidden states may trigger incorrect recall. A similarity threshold mitigates this but doesn't eliminate it.
|
|
|
|
**Linear scan**: Retrieval is O(n) over stored episodes. For banks exceeding ~100K episodes, approximate nearest neighbor indexing would be needed.
|
|
|
|
**Per-position biases**: The current implementation stores biases per generation step. Multi-token answers require one bias per token. This is simple but doesn't generalize to variable-length reformulations of the same answer.
|
|
|
|
## 6. Conclusion
|
|
|
|
Frozen transformers cannot form new memories. We give them a hippocampus.
|
|
|
|
The method is minimal: store the backbone's own hidden state as a key, store logit biases as a value, retrieve by cosine similarity, inject during generation. No gradients. No weight changes. No training loop. One forward pass to teach. One lookup to recall. Memories persist to disk.
|
|
|
|
The 200-line Python implementation reproduces the full result. The Clive Wearing Problem — intelligent systems that cannot form new memories — has a working solution.
|
|
|
|
## References
|
|
|
|
- Graves, A. et al. (2014). Neural Turing Machines. arXiv:1410.5401.
|
|
- Graves, A. et al. (2016). Hybrid computing using a neural network with dynamic external memory. Nature 538, 471-476.
|
|
- Lewis, P. et al. (2020). Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks. NeurIPS 2020.
|
|
- Meng, K. et al. (2022). Locating and Editing Factual Associations in GPT. NeurIPS 2022.
|
|
- Meng, K. et al. (2023). Mass-Editing Memory in a Transformer. ICLR 2023.
|