Reframe as Conditioned Reflex Injection (CRI) with multi-model test results
- 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>
This commit is contained in:
220
python/epimem.py
220
python/epimem.py
@@ -1,14 +1,20 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Episodic Memory (epimem): One-shot gradient-free learning on frozen transformers.
|
||||
Episodic Memory (epimem): Conditioned Reflex Injection on Frozen Transformers.
|
||||
|
||||
This is the minimal Python reproduction of the paper.
|
||||
Teaches a frozen Qwen 2.5 backbone new facts via hidden-state episodic memory,
|
||||
then recalls them with logit bias injection. No gradients at any point.
|
||||
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
|
||||
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
|
||||
@@ -24,18 +30,22 @@ from pathlib import Path
|
||||
# ─── Memory Bank ─────────────────────────────────────────
|
||||
|
||||
class EpisodicMemory:
|
||||
"""Hidden-state episodic memory bank.
|
||||
"""Hidden-state conditioned reflex bank.
|
||||
|
||||
Stores (key, value) pairs where:
|
||||
key = backbone hidden state (the model's internal representation of the prompt)
|
||||
value = logit biases (which tokens to boost for the correct answer)
|
||||
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):
|
||||
"""Store a new episodic memory. One-shot, no gradients."""
|
||||
"""Condition a new reflex. One-shot, no gradients."""
|
||||
self.episodes.append({
|
||||
"key": key / (np.linalg.norm(key) + 1e-8), # normalize
|
||||
"logit_biases": logit_biases,
|
||||
@@ -45,7 +55,7 @@ class EpisodicMemory:
|
||||
})
|
||||
|
||||
def recall(self, query_key: np.ndarray, threshold: float = 0.5):
|
||||
"""Retrieve best matching episode via cosine similarity."""
|
||||
"""Fire matching reflex via cosine similarity on activation pattern."""
|
||||
query_norm = query_key / (np.linalg.norm(query_key) + 1e-8)
|
||||
|
||||
best_sim = -1.0
|
||||
@@ -62,7 +72,7 @@ class EpisodicMemory:
|
||||
return None, best_sim
|
||||
|
||||
def save(self, path: str):
|
||||
"""Save memory bank to JSON."""
|
||||
"""Save reflex bank to JSON."""
|
||||
data = []
|
||||
for ep in self.episodes:
|
||||
data.append({
|
||||
@@ -77,7 +87,7 @@ class EpisodicMemory:
|
||||
print(f"Saved {len(data)} episodes to {path}")
|
||||
|
||||
def load(self, path: str):
|
||||
"""Load memory bank from JSON."""
|
||||
"""Load reflex bank from JSON."""
|
||||
with open(path) as f:
|
||||
data = json.load(f)
|
||||
self.episodes = []
|
||||
@@ -94,24 +104,74 @@ class EpisodicMemory:
|
||||
|
||||
# ─── 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:
|
||||
"""Qwen 2.5 backbone via HuggingFace transformers (PyTorch)."""
|
||||
"""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.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
|
||||
self.model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name, torch_dtype=torch.float32, trust_remote_code=True)
|
||||
self.model.eval()
|
||||
self.torch = torch
|
||||
self.model_name = model_name
|
||||
|
||||
self.hidden_dim = self.model.config.hidden_size
|
||||
self.vocab_size = self.model.config.vocab_size
|
||||
self.target_layer = self.model.config.num_hidden_layers - 1
|
||||
print(f" hidden_dim={self.hidden_dim}, vocab={self.vocab_size}")
|
||||
# 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."""
|
||||
@@ -121,14 +181,23 @@ class TransformersBackbone:
|
||||
"""Decode token IDs to text."""
|
||||
return self.tokenizer.decode(token_ids)
|
||||
|
||||
def get_hidden(self, token_ids: list) -> np.ndarray:
|
||||
"""Extract hidden state at the last token position."""
|
||||
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():
|
||||
outputs = self.model(ids, output_hidden_states=True)
|
||||
# Hidden state from target layer (pre-final)
|
||||
hidden = outputs.hidden_states[self.target_layer][0, -1]
|
||||
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:
|
||||
@@ -136,9 +205,13 @@ class TransformersBackbone:
|
||||
import torch
|
||||
ids = torch.tensor([token_ids])
|
||||
with torch.no_grad():
|
||||
outputs = self.model(ids)
|
||||
logits = outputs.logits[0] # [seq_len, vocab]
|
||||
return logits.numpy()
|
||||
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:
|
||||
@@ -150,7 +223,12 @@ class TransformersBackbone:
|
||||
for step in range(max_new):
|
||||
ids = torch.tensor([generated])
|
||||
with torch.no_grad():
|
||||
logits = self.model(ids).logits[0, -1] # [vocab]
|
||||
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):
|
||||
@@ -161,14 +239,21 @@ class TransformersBackbone:
|
||||
next_token = int(logits.argmax())
|
||||
generated.append(next_token)
|
||||
|
||||
if next_token == self.tokenizer.eos_token_id:
|
||||
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:
|
||||
"""Qwen 2.5 backbone via ONNX Runtime (faster, no PyTorch needed)."""
|
||||
"""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
|
||||
@@ -177,8 +262,13 @@ class OnnxBackbone:
|
||||
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(
|
||||
f"{model_dir}/tokenizer", trust_remote_code=True)
|
||||
tokenizer_path, trust_remote_code=True)
|
||||
|
||||
# Probe dimensions
|
||||
test_ids = np.array([[1, 2, 3]], dtype=np.int64)
|
||||
@@ -222,7 +312,11 @@ class OnnxBackbone:
|
||||
next_token = int(np.argmax(logits))
|
||||
generated.append(next_token)
|
||||
|
||||
if next_token == self.tokenizer.eos_token_id:
|
||||
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):]
|
||||
@@ -231,13 +325,13 @@ class OnnxBackbone:
|
||||
# ─── Teaching Protocol ───────────────────────────────────
|
||||
|
||||
def teach_fact(backbone, memory: EpisodicMemory, prompt: str, answer: str):
|
||||
"""Teach one fact. One forward pass, no gradients.
|
||||
"""Condition one reflex. One forward pass, no gradients.
|
||||
|
||||
1. Extract hidden state for prompt (= memory key)
|
||||
2. Get logits for prompt+answer vs prompt alone (= logit biases)
|
||||
3. Store in memory bank
|
||||
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
|
||||
"""
|
||||
# Key: hidden state of prompt
|
||||
# Trigger: hidden state of prompt
|
||||
prompt_ids = backbone.encode(prompt)
|
||||
key = backbone.get_hidden(prompt_ids)
|
||||
|
||||
@@ -268,7 +362,7 @@ def teach_fact(backbone, memory: EpisodicMemory, prompt: str, answer: str):
|
||||
|
||||
def recall_fact(backbone, memory: EpisodicMemory, query: str,
|
||||
max_tokens: int = 10) -> tuple:
|
||||
"""Recall a fact. Hidden-state lookup + logit injection.
|
||||
"""Fire conditioned reflex. Hidden-state trigger match + logit injection.
|
||||
|
||||
Returns (generated_text, similarity, episode).
|
||||
"""
|
||||
@@ -278,11 +372,11 @@ def recall_fact(backbone, memory: EpisodicMemory, query: str,
|
||||
episode, sim = memory.recall(query_key, threshold=0.3)
|
||||
|
||||
if episode is None:
|
||||
# No match — generate without memory
|
||||
# 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 logit bias injection
|
||||
# Generate with conditioned logit bias injection
|
||||
new_ids = backbone.generate(
|
||||
query_ids, max_new=max_tokens,
|
||||
logit_biases=episode["logit_biases"])
|
||||
@@ -292,9 +386,13 @@ def recall_fact(backbone, memory: EpisodicMemory, query: str,
|
||||
|
||||
# ─── Main ────────────────────────────────────────────────
|
||||
|
||||
DEFAULT_MODEL = "Qwen/Qwen2.5-0.5B"
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Episodic Memory: gradient-free learning on frozen transformers")
|
||||
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",
|
||||
@@ -305,12 +403,12 @@ def main():
|
||||
if args.onnx:
|
||||
backbone = OnnxBackbone(args.onnx)
|
||||
else:
|
||||
backbone = TransformersBackbone("Qwen/Qwen2.5-0.5B")
|
||||
backbone = TransformersBackbone(args.model)
|
||||
|
||||
memory = EpisodicMemory()
|
||||
|
||||
# ─── Teaching ─────────────────────────────────────────
|
||||
print("\n=== Teaching 3 facts ===")
|
||||
print("\n=== Conditioning 3 reflexes ===")
|
||||
facts = [
|
||||
("The capital of Zyphraxia is", "Novaheim"),
|
||||
("The ruler of Zyphraxia is", "Queen Stellara"),
|
||||
@@ -321,7 +419,7 @@ def main():
|
||||
teach_fact(backbone, memory, prompt, answer)
|
||||
|
||||
# ─── Recall ──────────────────────────────────────────
|
||||
print("\n=== Recall test ===")
|
||||
print("\n=== Reflex trigger test ===")
|
||||
all_ok = True
|
||||
for prompt, expected in facts:
|
||||
text, sim, ep = recall_fact(backbone, memory, prompt)
|
||||
@@ -331,6 +429,28 @@ def main():
|
||||
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)
|
||||
|
||||
@@ -350,10 +470,12 @@ def main():
|
||||
# ─── Summary ─────────────────────────────────────────
|
||||
print(f"\n{'='*50}")
|
||||
if all_ok:
|
||||
print("ALL TESTS PASSED: gradient-free episodic memory works.")
|
||||
print("ALL TESTS PASSED: conditioned reflex injection works.")
|
||||
else:
|
||||
print("SOME TESTS FAILED: check output above.")
|
||||
print(f"Memory bank saved to: {args.save}")
|
||||
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.")
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user