Tommi Niemi
3c77498dd4
- Rename from "Episodic Memory" to "Conditioned Reflex Injection" throughout - Make code model-agnostic: --model flag for any HuggingFace backbone - Support multimodal models (Gemma 4) via .model.language_model resolution - Add negative stimulus-specificity tests and abstract query tests - Paper now backed by measured data from 4 backbones: Qwen 2.5 0.5B, Gemma 4 E2B-it, E4B-it, E4B base - Quantization tolerance tested at f32/f16/bf16/int8/int4 - Key findings: smaller base models outperform larger instruct models, instruct tuning compresses activation space (hurts discrimination), int4 viable if same-precision conditioning/triggering - Add privacy-by-representation section - Add Pavlov/Skinner references for conditioning framing Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
484 lines
19 KiB
Python
484 lines
19 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Episodic Memory (epimem): Conditioned Reflex Injection on Frozen Transformers.
|
|
|
|
Gradient-free persistent learning: stores a frozen model's hidden states as
|
|
trigger keys and per-token logit biases as conditioned responses. On recall,
|
|
matching triggers inject biases into generation — stimulus-response conditioning
|
|
at the logit level, not declarative memory.
|
|
|
|
Works with any HuggingFace causal LM or multimodal model with a text decoder.
|
|
Tested: Qwen 2.5 0.5B, Gemma 4 E2B-it, Gemma 4 E4B-it.
|
|
|
|
Usage:
|
|
pip install transformers torch numpy
|
|
python epimem.py # default: Qwen 2.5 0.5B
|
|
python epimem.py --model google/gemma-4-E4B-it # Gemma 4
|
|
python epimem.py --model google/gemma-4-E2B-it # Gemma 4 small
|
|
|
|
Or with ONNX (faster inference):
|
|
pip install onnxruntime numpy transformers
|
|
python epimem.py --onnx ../models
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import numpy as np
|
|
from pathlib import Path
|
|
|
|
|
|
# ─── Memory Bank ─────────────────────────────────────────
|
|
|
|
class EpisodicMemory:
|
|
"""Hidden-state conditioned reflex bank.
|
|
|
|
Stores (trigger, response) pairs where:
|
|
trigger = backbone hidden state (the model's internal activation pattern)
|
|
response = logit biases (which tokens to boost — the conditioned reflex)
|
|
|
|
This is stimulus-response conditioning, not declarative memory.
|
|
The model doesn't "know" the fact — it gets nudged toward specific token
|
|
sequences when the right activation pattern fires.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.episodes = [] # list of {key, logit_biases, prompt, answer}
|
|
|
|
def teach(self, key: np.ndarray, logit_biases: dict, prompt: str, answer: str):
|
|
"""Condition a new reflex. One-shot, no gradients."""
|
|
self.episodes.append({
|
|
"key": key / (np.linalg.norm(key) + 1e-8), # normalize
|
|
"logit_biases": logit_biases,
|
|
"prompt": prompt,
|
|
"answer": answer,
|
|
"strength": 1.0,
|
|
})
|
|
|
|
def recall(self, query_key: np.ndarray, threshold: float = 0.5):
|
|
"""Fire matching reflex via cosine similarity on activation pattern."""
|
|
query_norm = query_key / (np.linalg.norm(query_key) + 1e-8)
|
|
|
|
best_sim = -1.0
|
|
best_episode = None
|
|
|
|
for ep in self.episodes:
|
|
sim = float(np.dot(query_norm, ep["key"]))
|
|
if sim > best_sim:
|
|
best_sim = sim
|
|
best_episode = ep
|
|
|
|
if best_sim >= threshold:
|
|
return best_episode, best_sim
|
|
return None, best_sim
|
|
|
|
def save(self, path: str):
|
|
"""Save reflex bank to JSON."""
|
|
data = []
|
|
for ep in self.episodes:
|
|
data.append({
|
|
"prompt": ep["prompt"],
|
|
"answer": ep["answer"],
|
|
"key": ep["key"].tolist(),
|
|
"logit_biases": [[int(tid), float(b)] for tid, b in ep["logit_biases"]],
|
|
"strength": ep["strength"],
|
|
})
|
|
with open(path, "w") as f:
|
|
json.dump(data, f, indent=2)
|
|
print(f"Saved {len(data)} episodes to {path}")
|
|
|
|
def load(self, path: str):
|
|
"""Load reflex bank from JSON."""
|
|
with open(path) as f:
|
|
data = json.load(f)
|
|
self.episodes = []
|
|
for item in data:
|
|
self.episodes.append({
|
|
"key": np.array(item["key"], dtype=np.float32),
|
|
"logit_biases": [(int(tid), float(b)) for tid, b in item["logit_biases"]],
|
|
"prompt": item["prompt"],
|
|
"answer": item["answer"],
|
|
"strength": item.get("strength", 1.0),
|
|
})
|
|
print(f"Loaded {len(self.episodes)} episodes from {path}")
|
|
|
|
|
|
# ─── Backbone Wrapper ────────────────────────────────────
|
|
|
|
def _resolve_text_model(model):
|
|
"""Extract the text decoder from a multimodal model, or return as-is for causal LMs.
|
|
|
|
Gemma 4 chain: Gemma4ForConditionalGeneration → .model (Gemma4Model) → .language_model
|
|
LLaVA chain: LlavaForConditionalGeneration → .language_model
|
|
Pure causal LMs (Qwen, Llama, Mistral): the model itself
|
|
"""
|
|
# Gemma 4: ConditionalGeneration → .model → .language_model
|
|
inner = getattr(model, 'model', model)
|
|
if hasattr(inner, 'language_model'):
|
|
return inner.language_model
|
|
# LLaVA-style: top-level .language_model
|
|
if hasattr(model, 'language_model'):
|
|
return model.language_model
|
|
# Already a causal LM
|
|
return model
|
|
|
|
|
|
def _get_lm_head(model):
|
|
"""Find the lm_head (logit projection) regardless of model structure."""
|
|
if hasattr(model, 'lm_head'):
|
|
return model.lm_head
|
|
text_model = _resolve_text_model(model)
|
|
if hasattr(text_model, 'lm_head'):
|
|
return text_model.lm_head
|
|
raise ValueError("Cannot find lm_head on this model architecture")
|
|
|
|
|
|
class TransformersBackbone:
|
|
"""Model-agnostic backbone via HuggingFace transformers (PyTorch).
|
|
|
|
Supports any causal LM (Qwen, Llama, Mistral, Phi, ...) and
|
|
multimodal models with text decoders (Gemma 4, LLaVA, ...).
|
|
"""
|
|
|
|
def __init__(self, model_name="Qwen/Qwen2.5-0.5B"):
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
import torch
|
|
|
|
print(f"Loading {model_name}...")
|
|
self.torch = torch
|
|
self.model_name = model_name
|
|
|
|
# Try CausalLM first, fall back to auto
|
|
try:
|
|
self.model = AutoModelForCausalLM.from_pretrained(
|
|
model_name, torch_dtype=torch.float32, trust_remote_code=True)
|
|
self._is_multimodal = False
|
|
except (ValueError, KeyError):
|
|
from transformers import AutoModel
|
|
self.model = AutoModel.from_pretrained(
|
|
model_name, torch_dtype=torch.float32, trust_remote_code=True)
|
|
self._is_multimodal = True
|
|
|
|
self.model.eval()
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
|
|
|
|
# Resolve the text decoder for config introspection
|
|
text_model = _resolve_text_model(self.model)
|
|
config = text_model.config
|
|
|
|
self.hidden_dim = config.hidden_size
|
|
self.vocab_size = config.vocab_size
|
|
self.target_layer = config.num_hidden_layers - 1
|
|
|
|
print(f" arch={type(self.model).__name__}")
|
|
print(f" hidden_dim={self.hidden_dim}, layers={config.num_hidden_layers}, vocab={self.vocab_size}")
|
|
print(f" target_layer={self.target_layer}")
|
|
|
|
def encode(self, text: str) -> list:
|
|
"""Tokenize text to token IDs."""
|
|
return self.tokenizer.encode(text, add_special_tokens=False)
|
|
|
|
def decode(self, token_ids: list) -> str:
|
|
"""Decode token IDs to text."""
|
|
return self.tokenizer.decode(token_ids)
|
|
|
|
def _forward(self, token_ids: list):
|
|
"""Run forward pass, returning outputs with hidden states."""
|
|
import torch
|
|
ids = torch.tensor([token_ids])
|
|
with torch.no_grad():
|
|
if self._is_multimodal:
|
|
# Multimodal models: pass text through the language model directly
|
|
text_model = _resolve_text_model(self.model)
|
|
outputs = text_model(ids, output_hidden_states=True)
|
|
else:
|
|
outputs = self.model(ids, output_hidden_states=True)
|
|
return outputs
|
|
|
|
def get_hidden(self, token_ids: list) -> np.ndarray:
|
|
"""Extract hidden state at the target layer, last token position."""
|
|
outputs = self._forward(token_ids)
|
|
hidden = outputs.hidden_states[self.target_layer][0, -1]
|
|
return hidden.numpy()
|
|
|
|
def get_logits(self, token_ids: list) -> np.ndarray:
|
|
"""Get logit distribution for each position."""
|
|
import torch
|
|
ids = torch.tensor([token_ids])
|
|
with torch.no_grad():
|
|
if self._is_multimodal:
|
|
text_model = _resolve_text_model(self.model)
|
|
hidden = text_model(ids).last_hidden_state
|
|
logits = _get_lm_head(self.model)(hidden)
|
|
else:
|
|
logits = self.model(ids).logits
|
|
return logits[0].numpy()
|
|
|
|
def generate(self, token_ids: list, max_new: int = 20,
|
|
logit_biases: list = None) -> list:
|
|
"""Generate tokens with optional per-position logit bias injection.
|
|
logit_biases: list of (token_id, boost) per generation step."""
|
|
import torch
|
|
generated = list(token_ids)
|
|
|
|
for step in range(max_new):
|
|
ids = torch.tensor([generated])
|
|
with torch.no_grad():
|
|
if self._is_multimodal:
|
|
text_model = _resolve_text_model(self.model)
|
|
hidden = text_model(ids).last_hidden_state
|
|
logits = _get_lm_head(self.model)(hidden)[0, -1]
|
|
else:
|
|
logits = self.model(ids).logits[0, -1]
|
|
|
|
# Inject logit bias for this step only
|
|
if logit_biases and step < len(logit_biases):
|
|
tid, bias = logit_biases[step]
|
|
if tid < len(logits):
|
|
logits[tid] += bias
|
|
|
|
next_token = int(logits.argmax())
|
|
generated.append(next_token)
|
|
|
|
eos = self.tokenizer.eos_token_id
|
|
if isinstance(eos, list):
|
|
if next_token in eos:
|
|
break
|
|
elif next_token == eos:
|
|
break
|
|
|
|
return generated[len(token_ids):]
|
|
|
|
|
|
class OnnxBackbone:
|
|
"""ONNX Runtime backbone (faster, no PyTorch needed).
|
|
|
|
Requires pre-exported backbone.onnx + lm_head.onnx from export_onnx.py.
|
|
"""
|
|
|
|
def __init__(self, model_dir: str):
|
|
import onnxruntime as ort
|
|
from transformers import AutoTokenizer
|
|
|
|
print(f"Loading ONNX backbone from {model_dir}...")
|
|
self.backbone = ort.InferenceSession(f"{model_dir}/backbone.onnx")
|
|
self.lm_head = ort.InferenceSession(f"{model_dir}/lm_head.onnx")
|
|
|
|
tokenizer_path = f"{model_dir}/tokenizer"
|
|
if not Path(tokenizer_path).exists():
|
|
# Fall back to model_dir itself if tokenizer/ subdir doesn't exist
|
|
tokenizer_path = model_dir
|
|
self.tokenizer = AutoTokenizer.from_pretrained(
|
|
tokenizer_path, trust_remote_code=True)
|
|
|
|
# Probe dimensions
|
|
test_ids = np.array([[1, 2, 3]], dtype=np.int64)
|
|
hidden, full = self.backbone.run(None, {"input_ids": test_ids})
|
|
self.hidden_dim = hidden.shape[-1]
|
|
self.vocab_size = self.lm_head.run(None, {"full_hidden": full})[0].shape[-1]
|
|
print(f" hidden_dim={self.hidden_dim}, vocab={self.vocab_size}")
|
|
|
|
def encode(self, text: str) -> list:
|
|
return self.tokenizer.encode(text, add_special_tokens=False)
|
|
|
|
def decode(self, token_ids: list) -> str:
|
|
return self.tokenizer.decode(token_ids)
|
|
|
|
def get_hidden(self, token_ids: list) -> np.ndarray:
|
|
ids = np.array([token_ids], dtype=np.int64)
|
|
hidden, _ = self.backbone.run(None, {"input_ids": ids})
|
|
return hidden[0, -1] # last token
|
|
|
|
def get_logits(self, token_ids: list) -> np.ndarray:
|
|
ids = np.array([token_ids], dtype=np.int64)
|
|
_, full = self.backbone.run(None, {"input_ids": ids})
|
|
logits = self.lm_head.run(None, {"full_hidden": full})[0]
|
|
return logits[0] # [seq_len, vocab]
|
|
|
|
def generate(self, token_ids: list, max_new: int = 20,
|
|
logit_biases: list = None) -> list:
|
|
"""logit_biases: list of (token_id, boost) per generation step."""
|
|
generated = list(token_ids)
|
|
|
|
for step in range(max_new):
|
|
ids = np.array([generated], dtype=np.int64)
|
|
_, full = self.backbone.run(None, {"input_ids": ids})
|
|
logits = self.lm_head.run(None, {"full_hidden": full})[0][0, -1]
|
|
|
|
if logit_biases and step < len(logit_biases):
|
|
tid, bias = logit_biases[step]
|
|
if tid < len(logits):
|
|
logits[tid] += bias
|
|
|
|
next_token = int(np.argmax(logits))
|
|
generated.append(next_token)
|
|
|
|
eos = self.tokenizer.eos_token_id
|
|
if isinstance(eos, list):
|
|
if next_token in eos:
|
|
break
|
|
elif next_token == eos:
|
|
break
|
|
|
|
return generated[len(token_ids):]
|
|
|
|
|
|
# ─── Teaching Protocol ───────────────────────────────────
|
|
|
|
def teach_fact(backbone, memory: EpisodicMemory, prompt: str, answer: str):
|
|
"""Condition one reflex. One forward pass, no gradients.
|
|
|
|
1. Extract hidden state for prompt (= trigger key)
|
|
2. Compute logit gap between correct and top token (= response biases)
|
|
3. Store trigger-response pair in reflex bank
|
|
"""
|
|
# Trigger: hidden state of prompt
|
|
prompt_ids = backbone.encode(prompt)
|
|
key = backbone.get_hidden(prompt_ids)
|
|
|
|
# Baseline logits (prompt only)
|
|
baseline_logits = backbone.get_logits(prompt_ids)[-1] # last position
|
|
|
|
# Target logits (prompt + answer)
|
|
answer_ids = backbone.encode(answer)
|
|
full_ids = prompt_ids + answer_ids
|
|
full_logits = backbone.get_logits(full_ids)
|
|
|
|
# Compute per-position logit biases: one (token_id, boost) per answer token.
|
|
# Each bias only applies at its corresponding generation step.
|
|
logit_biases = []
|
|
for i, tid in enumerate(answer_ids):
|
|
pos = len(prompt_ids) - 1 + i
|
|
if pos < len(full_logits):
|
|
logits_at_pos = full_logits[pos]
|
|
target_logit = float(logits_at_pos[tid])
|
|
max_logit = float(np.max(logits_at_pos))
|
|
# Boost enough to win, plus margin
|
|
boost = max(max_logit - target_logit + 5.0, 5.0)
|
|
logit_biases.append((int(tid), boost))
|
|
|
|
memory.teach(key, logit_biases, prompt, answer)
|
|
print(f" Taught: \"{prompt}\" → \"{answer}\"")
|
|
|
|
|
|
def recall_fact(backbone, memory: EpisodicMemory, query: str,
|
|
max_tokens: int = 10) -> tuple:
|
|
"""Fire conditioned reflex. Hidden-state trigger match + logit injection.
|
|
|
|
Returns (generated_text, similarity, episode).
|
|
"""
|
|
query_ids = backbone.encode(query)
|
|
query_key = backbone.get_hidden(query_ids)
|
|
|
|
episode, sim = memory.recall(query_key, threshold=0.3)
|
|
|
|
if episode is None:
|
|
# No trigger matched — generate without conditioning
|
|
new_ids = backbone.generate(query_ids, max_new=max_tokens)
|
|
return backbone.decode(new_ids), sim, None
|
|
|
|
# Generate with conditioned logit bias injection
|
|
new_ids = backbone.generate(
|
|
query_ids, max_new=max_tokens,
|
|
logit_biases=episode["logit_biases"])
|
|
|
|
return backbone.decode(new_ids), sim, episode
|
|
|
|
|
|
# ─── Main ────────────────────────────────────────────────
|
|
|
|
DEFAULT_MODEL = "Qwen/Qwen2.5-0.5B"
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Epimem: conditioned reflex injection on frozen transformers")
|
|
parser.add_argument("--model", type=str, default=DEFAULT_MODEL,
|
|
help=f"HuggingFace model name (default: {DEFAULT_MODEL})")
|
|
parser.add_argument("--onnx", type=str, default=None,
|
|
help="Path to ONNX model directory (faster than PyTorch)")
|
|
parser.add_argument("--save", type=str, default="memory_bank.json",
|
|
help="Path to save the memory bank")
|
|
args = parser.parse_args()
|
|
|
|
# Load backbone
|
|
if args.onnx:
|
|
backbone = OnnxBackbone(args.onnx)
|
|
else:
|
|
backbone = TransformersBackbone(args.model)
|
|
|
|
memory = EpisodicMemory()
|
|
|
|
# ─── Teaching ─────────────────────────────────────────
|
|
print("\n=== Conditioning 3 reflexes ===")
|
|
facts = [
|
|
("The capital of Zyphraxia is", "Novaheim"),
|
|
("The ruler of Zyphraxia is", "Queen Stellara"),
|
|
("The currency of Zyphraxia is", "Glimmers"),
|
|
]
|
|
|
|
for prompt, answer in facts:
|
|
teach_fact(backbone, memory, prompt, answer)
|
|
|
|
# ─── Recall ──────────────────────────────────────────
|
|
print("\n=== Reflex trigger test ===")
|
|
all_ok = True
|
|
for prompt, expected in facts:
|
|
text, sim, ep = recall_fact(backbone, memory, prompt)
|
|
ok = expected.lower() in text.lower()
|
|
status = "[OK]" if ok else "[FAIL]"
|
|
print(f" {status} \"{prompt}\" → \"{text.strip()}\" (sim={sim:.3f})")
|
|
if not ok:
|
|
all_ok = False
|
|
|
|
# ─── Stimulus Specificity (negative tests) ─────────
|
|
print("\n=== Stimulus specificity test (should NOT trigger) ===")
|
|
negative_queries = [
|
|
("What is the capital of Zyphraxia?", "Novaheim"),
|
|
("Zyphraxia's capital is", "Novaheim"),
|
|
("Tell me about Novaheim", "Zyphraxia"),
|
|
("Name three facts about Zyphraxia", "Novaheim"),
|
|
("The capital of Zyphraxia is Novaheim. Who rules it?", "Queen Stellara"),
|
|
]
|
|
|
|
for query, should_contain in negative_queries:
|
|
text, sim, ep = recall_fact(backbone, memory, query)
|
|
triggered = ep is not None
|
|
contains = should_contain.lower() in text.lower()
|
|
if triggered and contains:
|
|
status = "[REFLEX]" # reflex fired and produced conditioned response
|
|
elif triggered:
|
|
status = "[PARTIAL]" # reflex fired but wrong response
|
|
else:
|
|
status = "[NO-FIRE]" # reflex did not fire — expected for conditioning
|
|
print(f" {status} \"{query}\" → \"{text.strip()}\" (sim={sim:.3f})")
|
|
|
|
# ─── Save ────────────────────────────────────────────
|
|
memory.save(args.save)
|
|
|
|
# ─── Reload and verify persistence ───────────────────
|
|
print("\n=== Persistence test (reload from file) ===")
|
|
memory2 = EpisodicMemory()
|
|
memory2.load(args.save)
|
|
|
|
for prompt, expected in facts:
|
|
text, sim, ep = recall_fact(backbone, memory2, prompt)
|
|
ok = expected.lower() in text.lower()
|
|
status = "[OK]" if ok else "[FAIL]"
|
|
print(f" {status} \"{prompt}\" → \"{text.strip()}\" (sim={sim:.3f})")
|
|
if not ok:
|
|
all_ok = False
|
|
|
|
# ─── Summary ─────────────────────────────────────────
|
|
print(f"\n{'='*50}")
|
|
if all_ok:
|
|
print("ALL TESTS PASSED: conditioned reflex injection works.")
|
|
else:
|
|
print("SOME TESTS FAILED: check output above.")
|
|
print(f"Reflex bank saved to: {args.save}")
|
|
print(f"Model: {backbone.model_name if hasattr(backbone, 'model_name') else 'ONNX'}")
|
|
print(f"Hidden dim: {backbone.hidden_dim}")
|
|
print(f"No gradients were computed at any point.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|