sparse-memory-qa

A learned sparse content-addressable memory for long-document question answering — a research prototype exploring new ways to run AI models with far less compute and memory than standard attention.

What it is

Instead of the usual O(n²) attention over the whole document (whose KV cache grows with every token), this model compresses a document into a fixed set of 2,048 memory slots (~4 MB, constant no matter how long the document is). At query time it retrieves the top-k relevant slots and fuses them with the question to pick an answer.

Five small neural components are learned on top of a frozen bert-base-uncased backbone (only ~2M trainable parameters):

• key / value / query encoders
• an importance scorer (decides what deserves a memory slot)
• a fusion gate (combines retrieved memory with the question)

Measured results (1,000 held-out synthetic long-doc QA examples)

metric	value
recall@1	0.508
recall@10	0.996
answer accuracy (4-way)	0.336
latency	~92 ms/query (Colab T4)
memory	4.2 MB (constant in document length)

For comparison, the untrained version reached 55% top-10 recall; chance is 3%. Training was 3 epochs (10 min) on a free Colab T4 GPU.

Files

• pytorch_model.bin — the trained weights (2M params, fp32)
• config.json — architecture config (backbone, slot count)
• model.py — the model definition (SparseMemoryQA) and slot memory
• train.py / eval.py / data_gen.py — training, evaluation and synthetic data generation code

Use

from model import SparseMemoryQA
import torch

model = SparseMemoryQA(backbone="bert-base-uncased", n_slots=2048)
model.load_trainable(torch.load("pytorch_model.bin"))
model.eval()

See model.py for the full forward_example API (hard top-k retrieval is the inference-time path).

Limitations

Research prototype. Answer accuracy (34% vs 25% chance) shows the retrieval works far better than the answer head — more training and real datasets are the natural next steps. Not production-ready.

License

MIT