RandHan/ResMem-Mistral-7B-v0.3

ResMem for Mistral-7B-v0.3

Parametric memory from From Lookups to Lessons: Turning Retrieved Knowledge into Parametric Memory.

Code: github.com/SSRand/From-Lookups-to-Lessons

ResMem is a feed-forward memory (eight blocks, hidden size 4096, intermediate size 14336, 1.54B parameters) that reads the input of the final feed-forward sublayer of a frozen mistralai/Mistral-7B-v0.3 and emits a logit residual s. It is trained by residual self-distillation from a BM25-RAG teacher: the frozen base reading the top-3 BM25 passages for each 64-token span. At inference no retrieval is needed; the next-token distribution is softmax(z_B + γ s).

This is the knowledge-QA checkpoint reported in Table 1 of the paper, taken after one epoch over the December 2021 English Wikipedia. Weights are stored in BF16; the evaluation code loads them in FP32.

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from resmem.memory import capture_memory_input, load_memory  # from the GitHub repository

tok = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.3")
base = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.3", torch_dtype=torch.bfloat16).cuda().eval()
memory = load_memory("RandHan/ResMem-Mistral-7B-v0.3", base, "cuda")

store = {}
capture_memory_input(base, store)
ids = tok("Q: Who wrote Hamlet?\nA:", return_tensors="pt").input_ids.cuda()
with torch.no_grad():
    z_base = base(ids).logits[:, -1].float()
    s = memory(store["hidden_states"][:, -1].float())
    z = z_base + 0.4 * s  # gamma = 0.4
print(tok.decode(z.argmax(-1)))

In the paper, γ is selected per dataset on validation EM; Appendix C uses γ = 0.4 for every dataset.

Citation

@article{han2026lookups,
  title   = {From Lookups to Lessons: Turning Retrieved Knowledge into Parametric Memory},
  author  = {Han, Changjiang and Yuan, Ye and He, Bowei and Zhang, Weixu and Wu, Haolun and Liu, Xue and Kang, Jikun},
  journal = {arXiv preprint},
  year    = {2026}
}