LUMIA-Group/MemoryDecoder-at-Scale

TL;DR: Memory Decoder at Scale pretrains a standalone "memory" transformer that distills a kNN-LM's retrieval distribution into parameters, then runs alongside a frozen backbone at inference by interpolating output distributions. The new result is scale: general memories up to 6.9B params trained on 300B tokens, where frozen Pythia-410M + 6.9B memory reportedly beats Pythia-12B across 17 benchmarks, and swappable 1.7B domain memories add >9 points on bio/law/finance to Qwen3 backbones from 0.6B to 14B — no backbone modification, no retrieval infrastructure at inference time.

The idea

This builds directly on the original Memory Decoder work (the repo credits and adapts that codebase). The recipe: run a base model over a corpus, store contextual representations in a Faiss datastore, and for each training token compute the kNN-LM next-token distribution (a nonparametric memory known to fix long-tail factual recall). Then train a separate decoder-only model to reproduce those kNN distributions — effectively compressing the datastore into parameters. At inference, the memory decoder and the frozen backbone both do a forward pass and their output distributions are interpolated with a weight λ. You get kNN-LM-style behavior with no index, no retrieval latency, and a memory you can hot-swap (general vs. domain-specialized) under a fixed backbone.

The engineering contribution here is making that pretraining feasible at LM scale: a distributed Faiss indexing/search pipeline and sparse, batch-wise loading of kNN target distributions (the targets are stored as flattened sparse shards — token ids, probs, offsets — rather than dense vocab-sized distributions, which would be intractable at 300B tokens).

What's actually there

• Code (Apache-2.0): full pipeline — corpus preprocessing, distributed datastore construction (train_base.py), Faiss kNN target generation, sparse-shard flattening, and memory pretraining (train_memdec_trainer.py) with FSDP configs for 8×GPU Qwen3 training.
• Weights and data: a HuggingFace collection with general memories (e.g. MemoryDecoder-Pythia-2.8B-general), domain memories, and data artifacts.
• Eval infrastructure: bundled lm-evaluation-harness and OpenCompass trees with an hf-memdec model adapter, covering general tasks, multi-hop QA (2Wiki, HotpotQA, Bamboogle), HaluEval, FinEval, BioInst, LawBench.

Headline numbers from the README: Pythia-410M + 6.9B general memory averages 37.34 over 17 benchmarks, surpassing Pythia-12B at 39% fewer total parameters; 1.7B domain memories give >9 points average improvement on bio/law/finance across every Qwen3 Base scale from 0.6B to 14B. Per-task tables are in the paper, not the README.

Caveats worth checking

• "39% fewer parameters" ≠ cheaper inference. You run two forward passes per token. 410M + 6.9B is ~7.3B params of active compute; the fair latency/FLOPs comparison is against a ~7B dense model, not against 12B parameter count. The README doesn't report inference throughput.
• Per-task λ tuning. The eval adapter "automatically selects the task-specific interpolation weight" for released memories — i.e., λ was tuned per benchmark. That's standard for kNN-LM-family work but softens the plug-and-play claim.
• Tokenizer lock-in. Distribution interpolation requires shared vocabulary; memories are family-specific (Pythia memories for Pythia, etc.). Swappable across scales within a family, not across families.
• Eval mix skews toward likelihood-scored tasks (ARC, PIQA, MMLU) and QA. Whether gains hold for free-form generation quality is the open question; HaluEval inclusion is a partial signal.
• Training cost: building a 300B-token Faiss datastore and kNN targets is a serious one-time compute investment — that's the point (amortized across backbones), but reproducing it yourself is nontrivial.

Why it matters

This is a concrete data point in the "memory vs. parameters" debate: factual knowledge can be factored out of the backbone into a separately pretrained, composable module — a middle path between fine-tuning (destructive, per-model) and RAG (inference-time retrieval cost, context pollution). If per-task λ can be replaced by something learned, this pattern could genuinely spread.

Try it

Evaluate a released memory against a frozen Pythia backbone:

conda env create -f eval/lm-evaluation-harness/environment.yml
conda activate lmeval
cd eval/lm-evaluation-harness

lm-eval \
  --model hf-memdec \
  --model_args pretrained=EleutherAI/pythia-1.4b-deduped,memdec_path=Rubin-Wei/MemoryDecoder-Pythia-2.8B-general \
  --tasks arc_easy,piqa,mmlu \
  --batch_size 1

Checkpoints: the MemoryDecoder-at-Scale collection on HuggingFace.