AuthorsZijun Wu, Yongchang Hao, Lili Mou
TL;DR
TokMem encodes reusable task procedures into single trainable memory tokens, letting a frozen LLM match or beat LoRA and RAG (e.g., 67.0 ROUGE-L vs 66.5 Replay Memory on 1,000 SNI tasks).
Read our summary here, or open the publisher PDF on the next tab.
THE PROBLEM
Prompting and RAG rely on long textual instructions that must be reprocessed for every query, increasing quadratic self attention cost and truncation risk.
On Super Natural Instructions with 1,000 tasks, Sentence BERT retrieval accuracy drops to 79.7%, causing noisy procedure selection and degraded task performance for RAG style systems.
HOW IT WORKS
TokMem introduces a Memory Bank of trainable memory tokens and uses memory routing plus conditional generation to encode each reusable procedure into a single token.
You can think of TokMem like a CPU instruction cache: each memory token is a compact opcode that triggers a full multi step skill without re reading a long prompt.
This renormalization stabilized memory bank lets TokMem chain procedures at inference, add new skills continually, and avoid context window bloat that plain prompting or RAG cannot escape.
DIAGRAM
This diagram shows how TokMem routes queries to memory tokens and chains procedures during inference, as described in Section 2.3.
DIAGRAM
This diagram shows how TokMem is trained on procedure response pairs and evaluated on SNI and APIGen benchmarks.
PROCESS
Textualized context engineering
TokMem starts from a standard Transformer input sequence and avoids expanding prompts, replacing long textual instructions with compact memory tokens from the Memory Bank.
TokMem procedural memory as a token
TokMem concatenates the query with a memory token and response tokens, training the Memory Bank embeddings via next token prediction while keeping the backbone frozen.
Inference with memory tokens
TokMem uses memory routing to select the most probable memory token for the query, appends it, and runs conditional generation, optionally chaining multiple procedures.
Stabilizing new memories
TokMem applies renormalization to newly added memory embeddings, aligning their norms with existing tokens to prevent routing domination and catastrophic forgetting.
KEY CONTRIBUTIONS
One token procedural memory
TokMem encodes each reusable task procedure into a single trainable memory token in a Memory Bank, enabling constant size overhead even across 1,000 SNI tasks.
Frozen backbone with modular skills
TokMem keeps the LLM backbone frozen and stores procedural knowledge entirely in memory tokens, supporting continual learning without degrading existing procedures.
Superior recall and compositional tool use
TokMem achieves up to 67.0 ROUGE L on 1,000 SNI tasks and 99.1 tool selection F1 on APIGen, surpassing RAG and LoRA fine tuning with far fewer parameters.
RESULTS
ROUGE L Avg
67.0
+0.5 over Replay Memory on Llama 3.1 8B (66.5)
Routing Accuracy 1000 tasks
97.5%
+17.8 points vs Sentence BERT retriever at 79.7%
Tool Selection F1 Avg
99.1
+1.0 over Fine Tuning on Llama 3.1 8B (98.1)
Trainable Params TokMem 8B
0.20M
vs 3.41M for LoRA fine tuning on Llama 3.1 8B
On Super Natural Instructions (atomic recall) and APIGen (compositional tool use), TokMem demonstrates high quality procedural recall and routing. These results show TokMem can rival or exceed LoRA and RAG while using an order of magnitude fewer trainable parameters.
BENCHMARK
On Super Natural Instructions (atomic recall) and APIGen (compositional tool use), TokMem demonstrates high quality procedural recall and routing. These results show TokMem can rival or exceed LoRA and RAG while using an order of magnitude fewer trainable parameters.
BENCHMARK
ROUGE L score on Super Natural Instructions after learning 1,000 tasks.
BENCHMARK
Average tool selection F1 across 2–4 calls on APIGen.
KEY INSIGHT
TokMem with only 0.20M trainable parameters on Llama 3.1 8B reaches 99.1 tool selection F1, beating LoRA fine tuning with 3.41M parameters.
This is surprising because parameter efficient fine tuning is expected to be strictly stronger than a few learned embeddings, yet TokMem’s routed tokens capture richer procedural control.
WHY IT MATTERS
TokMem makes it practical to attach thousands of reusable procedures as single tokens to a frozen LLM, with high routing accuracy and minimal forgetting.
Builders can now grow modular skill libraries, chain procedures like tools, and personalize behavior via user specific memory banks without retraining or blowing up context windows.
Nikhil Verma
· 2026
Focus Agent adds start_focus, complete_focus, a persistent Knowledge block, and an optimized Persistent Bash plus String-Replace Editor scaffold to actively compress context during long software-engineering tasks. On five hard SWE-bench Lite instances against a Baseline ReAct agent, Focus Agent achieves 22.7% token reduction (14.9M → 11.5M) while matching 3/5 = 60% task success.
Xingyu Lyu, Jianfeng He et al.
· 2026
ADAM combines Anchor extraction, Distribution estimation, Anchor selection, and Query generation to adaptively probe agent memory via an auxiliary generator and entropy based selection. On the EHRAgent benchmark with Llama2-7b-chat, ADAM reaches EQ=77 and ASR=1.00, compared to MEXTRA’s EQ=44 and ASR=0.89.
Shannan Yan, Jingchen Ni et al.
· 2026
AdaMem organizes dialogue history into Working Memory, Episodic Memory, Persona Memory, and Graph Memory coordinated by a Memory Agent, Research Agent, and Working Agent. On LoCoMo with GPT-4.1-mini, AdaMem achieves 44.65 F1 overall, beating the best baseline LangMem at 41.76 F1 by +2.89.
Answers use this explainer on Memory Papers.
Checking…