TokMem: One-Token Procedural Memory for Large Language Models
Authors: Zijun Wu
Yongchang Hao
Lili Mou
Venue: ICLR 2026 Poster (2026)
Keywords: Procedural Memory
Memory tokens
Continual adaptation
Large language models
Primary area: foundation or frontier models, including LLMs

Abstract: Large language models are typically controlled via prompts, which must be repeatedly re-processed for every new query and are difficult to reuse modularly. We introduce TokMem, a procedural memory framework that compiles each reusable task procedure into a single trainable memory token. Each token serves as both a procedure index and a generation control signal that steers generation, enabling targeted behaviors with constant-size overhead. TokMem keeps the backbone LLM frozen and stores procedural knowledge entirely in these dedicated units, so new procedures can be added continually without interfering with existing ones. We evaluate TokMem on two settings: atomic recall over 1,000 Super-Natural Instructions tasks and compositional recall on multi-step function-calling. Our results show that TokMem consistently outperforms retrieval-augmented prompting while avoiding repeated context overhead. Moreover, it matches or exceeds parameter-efficient fine-tuning with substantially fewer trainable parameters.

Published as a conference paper at ICLR 2026
TOKMEM: ONE-TOKEN PROCEDURAL MEMORY FOR LARGE LANGUAGE MODELS
Zijun Wu', Yongchang Hao', Lili Mou
Dept. Computing Science & Alberta Machine Intelligence Institute (Amii), University of Alberta CCnada CIFAR Al Chair {zi jun4, yonghaaoyyooocccege/aeearr..
ABSTRACT
Large language models are typically controlled via prompts, which must be re- peatedly re-processed for every new query and are difficult to reuse modularly. We introduce TokMem, a procedural memory framework that compiles each reusable task procedure into a single trainable memory token. Each token serves as both a procedure index and generation control signal that steers generation, enabling targeted behaviors with constant-size overhead. TokMem keeps the backbone LLM frozen and stores procedural knowledge entirely in these dedicated units, new procedures can be added continually without interfering with existing ones. We evaluate TokMem on two settings: atomic recall over 1,000 Super-Natural Instructions tasks and compositional recall on multi-step function-calling. Our results show that TokMem consistently outperforms gurieval-auumenter prompt- ing while avoiding repeated context overhead. Moreover, it matches or exceeds parameter-efficient fine-tuning with substantially fewer trainable parameters
INTRODUCTION
Large language models (LLMs) have become the foundation modern natural language processing. powering wide range of applications in text understanding, generation, and coding (Brown et al. 2020 Llama Team 2024 Chen et al., 2021). Prompting is a widely adopted way to steer LLM be- havior, where in-context learning enables adaptation to new tasks without parameter updates (Brown) ct al. 2020). Consequently, prompt and context engineering has emerged as a dominant interface for specifying tasks, obtaining relevant information, and guiding multi-step reasoning tool invo- cation (Wei et al. 20226 Yao et al., 2023; Sahoo et al. 2025).
Despite their effectiveness, long prompts are inherently inefficient. They are labor-intensive con- struct and maintain, and they do not scale well when conditioning diverse queries across different tasks (Liu al., 2023). At inference time, longer contexts increase compute because self-attention scales quadratically (Vaswani et al.) (017). They also shrink the usable context window for inputs and outputs, often causing truncation and information loss (Liu al., 2024a). Moreover, tasks accumulate, these costs make it increasingly difficult to execute task-specific behaviors efficiently.
To mitigate these issues, recent work externalizes prompts into retrieval-based memory. Retrieval- augmented generation (RAG; Lewis et al. 2020) is used to retrieves and reinserts documents or conversational state at inference time, a strategy exemplified by MemGPT (Packer et al.) 2024). Although retrieving in-context demonstrations (Wei et al. (022b) can supply procedural cues, the retrieved knowledge remains explicit text that must be repeatedly interpreted, resembling declarative memory in cognitive science. This leads to two drawbacks: (1) retrieved content still occupies the context window, reintroducing quadratic compute and truncation pressure, and (2) frequently used procedures repeatedly re-read as text rather than distilled into compact, reusable representations.
We propose One-Token Procedural Memory (TokMem), a modular framework that encodes task procedures into compact, trainable tokens while maintaining frozen LLM backbone. Here, procedure refers to a reusable context-response mapping that captures a specific task behavior; the terminology inspired by physiological research (Anderson & Lebiere 1998), which posits
Our code is publicly available at httpp:///ggitttb.ccoo/////

---

Published as a conference paper at ICLR 2026
(a) Encoding Procedures with TokMen
Memory Bank
Response
<508>
Backbone
Query
Resporie
Final Respponee
(b) Inference Memory Recall
Query:
Step
Step 2
dinner shapping
diet Hewlthy)
[carse diet] pr procecure
Recal [search food)
Generated shopping Ist:
Item
Quantity
(fooch (t0f--200,
Tofu
230g
Step
braccell: 150,
brown rice: 180).
Junction coll
Brocooli
150g
Recall
metric: gram)
Brown rice
180g
output]
Figure I: Overview of our TokMem approach. (a) New (colored) memory tokens are interleaved with text sequences and trained via next-token prediction while the LLM backbone remains frozen. (b) At inference time, a query recalls and chains memory tokens (e.g., parse, search, format), en- abling multi-step procedural behavior without long prompts.
procedural knowledge compiled form of skill that is more efficient than the slow interpretation of facts. Unlike factual knowledge that can be expressed in plain text, procedural knowledge (e.g., riding bicycle) is often more nuanced and sophisticated.
In TokMem, a memory bank stores learned procedures (Figure Th), where each memory token serves both index procedure and control signal that steers generation, allowing the model to execute targeted tasks with constant-size overhead regardless of the procedure's 's complexity. Rather than front-loading procedures as long prompts, TokMem integrates memory tokens directly into the generation process by recalling them on demand at each stage. As illustrated in Figure Db. TokMem supports multi-step workflows by sequentially recalls procedure tokens such as parsing, searching, and formatting; after the model completes specific segment of the response, it retrieves the next relevant token to condition the following stage, allowing for modular task composition.
A key advantage of TokMem is that its memory tokens are parameter-isolated from the backbone LLM. Because each procedure distilled into its own dedicated token, new skills can be added with- out interfering with or degrading existing ones. This isolation naturally supports continual learning, enabling the model accumulate procedural skills over time while remaining stable. This mirrors human procedural memory, where skills are gradually acquired through practice and later invoked by contextual cues (Anderson & Lebiere 1998). Overall, TokMem enables efficient learning and continual expansion of procedural knowledge.
We evaluate TokMem in two complementary settings. First, in atomic recall, we treat each of the 1,000 tasks from Super-Natural Instructions (Wang al. 2022a) distinct procedure. In this setting, TokMem successfully stores and retrieves these procedures efficiently without catastrophic forgetting, outperforming retrieval-based baselines eliminating the noise and computational over- head of text-based prompts. Second, in compositional recall using function-calling benchmark (Liu al. 2024b), each tool invocation is treated as an atomic procedure, and answering query requires chaining multiple such procedures together. Here, TokMem consistently matches or surpasses the performance of parametric fine-tunie tine-tuning while using significantly fewer trainable parameters.
METHOD
In this section, we first review the Transformer processing pipeline and the limitations of current textualized prompting ($ 2.1 We then introduce TokMem, detailing how procedures encoded into a single memory token (S (2.2). how these tokens are recalled and chained at inference time 2.3 and the renormalization strategy used to stabilize the memory bank ($ 2.4
2

---

Published as a conference paper at ICLR 2026
2.1 TEXTUALIZED CONTEXT ENGINEERING
A Transformer (aaswani et al. 2017 processes a sequence of tokens (aI,...,a,, ap) E N where each a, is an integer representing the index of token (usually sub-word). The model retrieves the corresponding embedding vector from the embedding layer and produces an input sequence (z1....,,n,) Raxd, which then used predict the next token sequence.
Recent advances in prompting can be viewed as textualized context engineering, where the goal is choose input tokens that steer the model toward desired behavior. For example, chain-of-thought prompting (Wei et al., 2022a) augments the input with intermediate reasoning steps to improve logical inference. Similarly, retrieval-based methods like RAG (Lewis et al. (020) fetch external documents to ground generation, while recent memory systems such as MemGPT (Packer et al.) 2024) retrieve text history into the active context to maintain conversational state. Critically, these approaches all rely on expanding the textual prompt, consuming the limited context window and increasing inference latency due to the quadratic cost of self-attention.
2.2 TOKMEM: PROCEDURAL MEMORY AS TOKEN
Our key idea is that frequently reused task procedures can be effectively "compressed" and stored by encoding them into a dedicated memory token, significantly reducing the verbosity of task in- structions. Consider I memory tokens. They added to the vocabulary special tokens, each represented by an embedding. Thus, they form a memory bank of special embeddings:
m
M
E R
m. E TRd
m
Each m, a trainable vector with no direct textual form and represents a unique procedure. For simplicity, we assign each memory token m, special token index Om, N.
To connect these tokens with training, we first describe single training instance. We define procedure-response pair, where procedure is invoked by memory token m; and the response is a sequence of textual tokens (rr E N that realizes the procedure in text. For example, the response may specify tool-call arguments follow a particular output format implied by the procedure.
We formulate training supervised task where each procedure-responee pair is provided, rep- resented by concatenating a procedural memory token Om, with the embeddings of (ra...., Tin). Each training instance may contain multiple secedure-responee turns, modeling tasks that require multi-step reasoning or composition for a query q. Formally, the sequentialized training sequence has the layout
(01 (q1..... Qk. Omny Orgs Grasn Omgs Orgs argannn .....
prossdureresponse pair
poocedure-responce puir
We adopt the standard next-token prediction loss:
C(a; M) = Σ i>k
log Pr(a; aass; M).
The memory embeddings (m1,...; m.) are shared between the input embedding layer and the LM head. During optimization, these embeddings are trainable, whereas the pre-trained text token embeddings and the backbone LLM remain frozen. Across the training corpus, each memory em- bedding is exposed to diverse queries and responses, allowing it to learn the underlying procedure representation. We visualize the training process in Figure T..
2.3 INFERENCE WITH MEMORY TOKENS
At inference time, TokMem recalls procedures via memory routing and conditional generation, where routing selects the appropriate memory token for query. Given = (qı,..., Qk). the model predicts distribution over memory tokens from the final hidden state hy:
P(am. g) exp (logit(m;) fr))),
(4)
3

---

Published as a conference paper at ICLR 2026
We select the most probable token Om- and append to form [q; am-)) after which the model gen- erates the response ruto-regressivelyy For multi-procedure queries, the model can predict additional memory tokens after generating each response segment, enabling token chaining as in Figure П a query does not match any learned procedure, all memory-token logits may remain low and the model defaults generating regular text.
In summary, the model decides whether (and how many) memory tokens to recall based on the training layout in (Eq. 2
2.4 STABILIZING NEW MEMORIES
While TokMem supports the modular addition of procedures, integrating new tokens into frozen backbone presents stability challenges. In continual learning settings, newly optimized memory embeddings prone norm inflation (Hou et al., 2019). These new embeddings dominate the routing logits, suppressing the retrieval of old memories regardless semantic relevance.
To address this, introduce renormalization, lightweight post-update calibration of the memory bank M E TRixd_ Let A and 7 denote the indices of the new and existing procedural memories, respectively. We estimate the typical scale from the existing
ūr meanje? m; l1g
(5)
and rescale each new embedding as
f1
m:
i A.
(6)
mm +
This operation preserves the directions of newly added embeddings while aligning their magnitudes to the established scale the memory bank, ensuring smooth integration without overwhelming routing dynamics. The computational overhead is negligible, scaling as O(IAld).
EXPERIMENTS
We evaluate TokMem in two complementary scenarios. In atomic memory recall, each task from Super-Natural Instructions (Wang et al.] (002aa) is treated as a standalone procedure, where query maps directly to the desired response. In compositional memory recall, evaluate multi-procedum tool use on a function-calling dataset (Liu et al.) 2024b); each tool invocation is modeled as an atomic procedure, and solving query requires composing multiple calls. Experiments are con- ducted on the Qwen (Qwen Team 2025) and Llama (Llama Team 2024) families, ranging from 0.5B 8B.
3.1 EXPERIMENTAL SETUP
Baselines. Across both settings, compare TokMem against textualized context engineering, eerienel-augmeeneed memory, and parameter-efficient fine-tuning baselines.
Base: In the atomic setting, the model answers queries without demonstrations, providing an empirical lower bound that highlights the need to recall task knowledge.
ICL: In the compositional setting, we augment the input with all tool descriptions and prepend two compositional proceduue-response
response demonstrations as a context-engineeriig baseline.
RAG: We Sentence-BERT (Reimers & Gurevych 2019) retrieve relevant demonstrations
or tool-use examples and prepend them to the query, following
genera-
tion (Packer al. 2024 Chhikara al. 2025 Xu et al. 2025
We fine-tune low-rank adapters (LoRA; Hu al. 2022). which are inserted into the query and value projections of the transformer, with a parameter count comparable to TokMem. This provides parametric form of procedural memory.
Fine-tuning:
Replay Memory: To mitigate catastrophic forgetting during fine-tuning, adopt experience replay (Mmih et al.] (2015) by maintaining buffer previously seen tasks or tools and mixing them into the current training data.
4

---

Published as a conference paper at ICLR 2026
Table 1:
Atomic recall performance SNI (ROUGE-L.). TokMem consistently outperforms fine-
tuning and RAG across models and scales, maintaining strong performance even at 1,000 tasks.
Number of Tasks
Model
Method
10
50
200
500
1000
Avg.
Base
33.9
39.0
38.8
39.1
38.5
37.9
RAG
50.4
43.2
38.8
36.2
34.7
40.7
Qwen 2.5 0.5B
Fine-Tuing
52.4
48.0
40.6
41.7
43.2
45.2
Replay Memory
52.4
49.5
47.2
47.7
46.7
48.7
TokMem
52.8
51.3
49.3
50.2
50.0
50.7
TokMem+DC
53.8
50.5
50.2
50.9
50.0
51.1
Base
16.6
19.9
20.0
18.7
18.2
18.7
RAG
60.0
48.7
45.8
42.3
39.9
47.3
Llama 3.2 3B
Fine-Tuing
67.1
59.1
59.5
58.4
57.9
60.4
Replay Memory
67.1
61.1
60.6
61.4
60.0
62.0
TokMem
68.0
62.3
61.2
61.5
61.5
62.9
TokMem+DC
68.8
62.5
58.7
61.7
61.1
62.6
Base
27.2
27.8
29.6
29.5
28.9
RAG
63.8
53.9
49.1
45.3
42.6
50.9
Llama 3.1 8B
Fine- Fine-Tuing
75.8
64,3
63.2
58.7
61.6
64.7
Replay Memory
75.8
65.2
64.5
63.4
63.6
66.5
TokMem
75.4
65.5
65.1
64.4
64.8
67.0
To
75.6
65.8
64.2
64.4
66.7
Training Details. All methods are implemented in HuggingFace Transformers and trained on single NVIDIA A6000 GPU (48GB) using mixed-precision (bfloat16) training. The backbone mod- els remain frozen: for fine-tuning, only the adapter weights (rank 8) are updated, whereas for TokMem, only the embeddings of newly added procedure IDs are trainable. We expand the tokenizer vocabulary with these procedure IDs and initialize their embeddings by averaging the pretrained em- beddings (Hewitr) 2021). For Replay Memory, we mix 20% replayed samples into each batch, using buffer of 500 examples refreshed every 10 tasks in the atomic setting and 1,000 examples updated each round in the compositional setting.
We optimize with AdamW using learning rate of 5 10-5 for fine-tuning and 10-3 for TokMem; weight decay is 10-2 for fine-tuning and zero for TokMem. Training runs for one epoch with batch size and maximum sequence length 1024, using teacher forcing and applying the loss only to memory-token and response positions.
Evaluation Metrics. We evaluate TokMem from two perspectives: (1) memory-toker routing ac- curacy, which measures whether the correct memory tokens are selected, and (2) task performance, which measures task-level generation quality (e.g., ROUGE-L and Fl).
3.2 ATOMIC MEMORY RECALL
Dataset Details. We evaluate Super-Natural Instructions (SNI: Wang al. 2022a). which pro- vides diverse QA-style natural language tasks. We use ROUGE-L (Lin 2004 to measure generation quality, Each task is treated an individual procedure: query directly invokes the Iearned pro- cedure to produce the desired response. We sample 1,000 English tasks; each task contains 500 training samples and test samples. To reflect how memories are typically acquired over time, we introduce tasks sequentially during training rather than all at once, while shuffling examples within each task. We scale the number of tasks from to 1,000 and record checkpoints after training on {10, 50, 200, 500, 1,000) tasks. This resembles incremental domain adaptation (Asghar et al. 2020). where at each checkpoint we evaluate performance across all previously seen tasks. Additional task details provided in Appendix A.I
Results and Findings. As shown in Table I baseline methods such Base stable but fail to achieve competitive performance. RAG performs reasonably well when the memory load small, but degrades quickly the number of stored task memories increases, indicating its sensitivity to
5

---

Published as a conference paper at ICLR 2026
Table 2: Task routing accuracy. TokMem achieves near-perfect routing accuracy 1,000 tasks, far exceeding the RAG retriever, whose accuracy falls below 80%.
Number of Tasks
Model
Method
10
200
500
1000
Sentence-BERR
RAG
99.6
92.6
88.7
83.2
79.7
TokMem
99.4
98.6
97.4
96.9
94.7
Qwen 2.5 0.58
To
99.4
99.2
98.4
97.2
96.1
TokMem
100.0
96.1
Llama 3,2 3B
Tc
99.8
95.4
TokMem TokMem+DC
99.8
99.6
98.9
97.7
97.5
Llama 3.1 8B
99.7
99.4
97.8
97.2
97.2
Qwen
3.2
3.1
72.4
74.5
73.8
75.0
75.
730
60.0
390
66
67.5
68.0
000
63.
86
50
512
222
R72
40.
1,5
36
342
27
20.2
Baselines
10
100
500
Baselines
10
100
500
Baselines
10
100
500
Training Samples Task
Training Samples
Task
Training Samples per Task
Base
MAG
Fine Tuning
TokMem
TokMem+DDC
Sample efficiency on 10-task mixture from SNI. Tok Mem consistently outperforms fine- tuning in the low-data regime and can surpass RAG with only training samples, demonstrating strong few-shot capability.
Figure 2:
the retriever's quality. Parametric methods such fine-tuning achieve stronger initial accuracy but suffer from forgetting as tasks accumulate; replay memory alleviates this issue but still falls short of TokMem. By contrast, TokMem maintains high accuracy with minimal performance drop as acquires new task memories, achieving the best average results across all settings.
We also include a variant where each memory token is decoupled into an address token and a steering token. This decoupling separates the roles of memory tokens and increases the capacity of TokMem. We refer this variant as TokMem with decoupled embeddings (TokMem+DC); further details are provided in Appendix A.2
The decoupled variant (TokMem+DCC) yields modest gains for the smaller Qwen-0.5B model but no improvement for larger Llama models, and in some cases it underperforms TokMem when scaling to many tasks. Overall, although TokMem+DC tempting variant, does not provide additional benefits. We therefore focus on the simple yet effective TokMem (without DC).
Table further highlights TokMem's robustness in memory routing. Its accuracy remains above 94%
1,000 tasks with the smallest 0.5B model, significantly outperforming the Sentence- BERT retriever used in RAG. By contrast, the retriever's accuracy drops below 80% when routing over 1,000 tasks. This high-fidelity routing enables TokMem to sustain strong performance with- out external retrieval or fine-tuning, demonstrating its advantage for continual and large-scale task acquisition.
Analysis of Training Sample Efficiency. We compare the training sample efficiency of LoRA fine-tuning and TokMem on 10 SNI tasks. We vary the number of training samples per task from 10 (few-shot) to 500, and report performance across these sample budgets in Figure 2 We set the adapter rank to r = 1, which helps reduce overfitting in the low-data regime and matches the parameter scale of TokMem. Results show that TokMem consistently outperforms fine-tuning all budgets, with the largest gains at small sample sizes. The decoupled variant (TokMem+DC) yields modest gains this controlled setting, mainly at higher sample budgets. Overall, these results

---

Published as a conference paper at ICLR 2026
Table 3: Compositional tool-use performance APIGen. TokMMm achieves strong tool selection and argument FI across multiple tool calls, outperforming ICL and RAG while requiring less input augmentation, and surpassing fine-tuning with far fewer trainable parameters.
Tool Selection
Argument
Model
Method
#Params
calls
3 calls
4 calls
Avg.
calls
calls
calls
Avg.
ICL
27.6
11.1
10.5
16.4
0.6
0.7
0.0
0.4
RAG
29.5
10.8
10.5
16.9
7.2
0.0
2.7
Llama 3.2
Fine-Tuing
0.85M
10.4
9.5
7.0
77.3
72.6
55.8
68.6
TokMem (w/o adapt)
0.10M
86.8
80.9
90.8
86.2
68.9
61.1
73.0
67.7
TokMem (w/ adapt)
0.10M
98.4
98.0
98.9
98.4
84.3
84.3
87.8
85.5
ICI.
66.8
59.2
61.9
42.2
42.3
38.8
44.1
RAG
78.1
71.2
69.3
72.8
54.8
53.1
62.7
56.9
Llama 3.2 3B
Fine-Tuing
2.29X1
99.77
38.[
96.8
97.9
88.9
36.6
82.9
85.8
TokMem (w/o adopt)
0.15M
82.6
79.3
67.2
76.4
65.4
57.2
$0.2
57.6
TokMem (w/ adapt)
0.15M
99.2
98.2
100.0
99.2
85.9
86.7
88.3
86,3
ICL
79.7
72.9
75.4
76.0
51.5
52.6
57.3
53.8
RAG
79.6
75.3
93.0
82.6
53.3
57.1
69.2
59.9
Llama 3.1 8B
Fine-Tuing
3.41M
98.8
97.2
98.2
98.
87.7
86.8
68.2"
87.6
TokMem (w/o adopt)
84.9
82.0
56.7
65.9
62.8
TokMem (w/ adapt)
0.20M
99.4
97.9
100.0
88.1
86.5
93.4
89.3
highlight TokMem's 's ability to acquire new procedures with limited data, making it both parameter- and data-efficient.
3.3 COMPOSITIONAL MEMORY RECALL
Dataset Details. We construct benchmark from the APIGen dataset (Liu et al. 2024b) by sam- pling 50 frequently used tools. Each tool invocation is treated as an atomic procedure, and solving a query requires composing multiple such procedures. We synthesize 5,000 training queries and 500 test queries, each capped at four tool calls. Additional dataset details are provided in Appendix A.3
We report performance using Fl metries: (1) Tool Prediction FI, which measures whether the model invokes the correct tools; and (2) Argument Generation F1, which evaluates the correctness of tool-call arguments. To account for semantic equivalence, both gold and predicted outputs are normalized into Abstract Syntax Trees before scoring (Patil et al. 2025).
Adaptation for Compositionality. We find that TokMem benefits from a brief adaptation phase that exposes the backbone to the compositional structure of memory tokens. Concretely, we fine- tune the backbone on a held-out auxiliary tool set using the same LoRA setup as the baseline. We then merge the adapted weights, after which the backbone remains frozen for memory acquisition and evaluation (see Appendix [A. for details). Note that the auxiliary tool set is different from the evaluation set, so this procedure does not violate our Frozen-backbone setup.
This adaptation teaches the model how to interleave responses with memory tokens. We therefore treat it as part TokMem for compositional tasks, and use TokMem with adaptation as the default configuration.
Results and Findings. Table shows that TokMem achieves the strongest overall performance. Even without adaptation for compositionality, it outperforms RAG while avoiding the added com- plexity of an externa retrieval mechanism. ICL and RAG perform poorly on both tool prediction and argument generation, especially with the smaller Llama IB model, likely due to its weaker instruction-followe ability.
TokMem consistently matches or exceeds LoRA fine-tuning's performance while requiring an order of magnitude fewer trainable parameters. For example, when applied to Llama 8B, LoRA requires 3.41M trainable parameters, whereas TokMem needs only 0.2M and achieves higher performance.
Notably, TokMem exhibits stronger correlation between tool selection and argument generation: improvements in the former translate directly into gains in the latter. By contrast, LoRA fine-tuning shows weaker alignment. For example, the IB model, it often generates plausible arguments

---

Published as a conference paper at ICLR 2026
Liama 3.2 1B
Liama 3.2 30
Liama 3.1 BB
Argument Scooe
20
20
Training Round
Training Round
Training Round
TolMemm
Replay Memory
Tools 1-20
Tools 11:20
Tools 21:30
Tools 31:40
Tools 41:50
Figure 3: Forgetting analysis under continual adaptation. As new tools are introduced. fine-tuning with replay memory suffers sharp performance drops on earlier tasks, whereas TokMem remains stable. Larger models exhibit stronger retention, likely due to their greater capacity.
Liama
Liama 3.2 38
Argumant F1 Score
Training Round
Training Round
Training Round
Renorm
Rerorm
Tools 1-30
Tools 11:20
Tools
Tools 31:40
41:50
Figure 4: Effect of renormalization on TokMem. Without renormalization, newly learned tokens dominate and older tokens are forgotten, particularly in smaller models with limited embedding capacity. Renormalization mitigates this by balancing token embedding norms.
even when tool selection is incorrect, suggesting that argument generation is not properly grounded in the chosen tool.
TokMem also exhibits stronger compositional generalization than fine-tuning, successfully execut- ing queries with more function calls than it has seen during training. Detailed results are provided in Appendix B.I
Analysis of Forgetting. We compare TokMem with replay memory in a continual adaptation set- ting. where tools are introduced sequentially round, 11 11-20 in the second, and so forth). As shown in Figure replay memory struggles to prevent catastrophic forgetting new tools introduced. By contrast, TokMem maintains higher perfor- mance across tool groups, with only mild declines that mainly reflect the increasing number of tools. Larger models exhibit better retention for both approaches, likely due their greater capacity, which reduces interference with previously learned tools.
five training rounds (e.g., tools 1-10 in the first
We further investigate the effect of the renormalization step introduced in Section 2.4 Without renormalization, the norms of newly added memory tokens can dominate older tokens in the softmax (see Appendix B.2). As shown in Figure 4 TokMem without renormalization exhibits noticeable forgetting, especially for smaller models. Larger models are more robust even without renormaliza- tion, likely due to their greater embedding capacity. Overall, renormalization improves TokMem's resistance to forgetting by balancing routing between new and old memory tokens. We provide fur- ther analysis in Appendix [B.3 on how keeping the backbone frozen helps prevent forgetting during continual memory acquisition.
3.4 ANALYSIS ON MEMORY PLACEMENT
An important design choice in TokMem is where memory tokens are placed within the input se- quence, this directly affects how the backbone model attends to and integrates procedural knowl-
8

---

Published as a conference paper at ICLR 2026
Table 4: Comparison of TokMem and prefix tuning for memorizing text from the Fanfics dataset, evaluated by perplexity. The suffix "-k" refers to the number of embeddings, and Steps @ 9est represents the training steps needed reach 90% of the lowest achieved perplexity. TokMem con- verges faster and achieves lower perplexity than prefix tuning, especially with a small number of memory tokens.
1024 tokens
2048 tokens
4096 tokens
Method
Steps @ 90%Best
PPL
Steps Neeps@90%Best
PPL
Steps@90%Bess
PPL
Prefix tuming-1
1700
3.81
2300
8.77
2200
14.32
TokMem-1
1200
3.28
1400
7.07
1700
12.27
Prefix tuning-2
500
1.13
1700
1800
8.38
TokMem-2
600
1.09
1300
2.75
1700
7.21
Prefix unnnng-5
300
1.07
500
1.17
1400
2.39
TokMem-5
200
1.03
500
1.15
1400
1.91
edge. Although TokMem introduces a memory-routing mechanism for generating tokens, its ef- fectiveness also depends on the placement strategy. Without routing, TokMem reduces prompt- tuning method (Li Liang 2021 Lester et al. (021) with learnable embeddings; however, it adopts an infi placement: query MEM ① response. remains unclear whether this placement is consistently better than the more common prefix formulation, MEM ① query ① response, used in prior prompt-tuning work, where prefix tokens influence generation before observing the query. To study the impact of placement, we compare prefix and infix placements under matched token budgets in the single-task setting.
Setups. We compare TokMem with infix memory placement against prefix tuning by stress-testing the capacity of memory tokens (Sastre & Rosá, 2025) on the recent Fanfics dataset, which was collected after the pretraining of the LLMs (Kuratov al. 2025). We fix the sequence length 128 tokens and vary the batch size from to 32, compressing batches of 1024 to 4096 response tokens into to 5 memory tokens. For each sequence, we prepend randomly generated query that serves only as marker to distinguish the two placements; the target to be learned remains the response.
We measure learning speed using (rpp@99%Bess. defined as the number of training steps (eval- uated every 100 steps) required to reach 90% of the best perplexity. Results are averaged over five runs. Additional experiments on generalization to a math reasoning dataset are provided in Appendix B.4
Results. Table 4 shows that TokMem consistently achieves lower perplexity and often converges faster than prefix tuning. With single token, TokMem reaches 90% of its best perplexity roughly 30% sooner than prefix tuming, suggesting that conditioning on memory after the query helps the model learn more efficiently. Interestingly, when more tokens are available (e.g., five tokens), the performance gap narrows. This suggests that prior work (Li & Liang 2021], which typically uses dozens or even hundreds tokens, may have underestimated the importance of memory placement in low-token regimes, where each token must compress more procedural information.
RELATED WORK
Equipping LLMs with memory has been explored along multiple directions. Most existing ap- proaches emphasize declarative memory, where the objective is to store and retrieve explicit infor- mation such as facts and conversation history (Packer al. 2024 Chhikara al. [2025] Wang al.) 2024). By contrast, parameter-based approaches internalize task-specific behaviors within model parameters, resembling procedural memory. TokMem builds on this latter view while emphasizing modularity and compositionality.
Text-Based External Memory. A common approach is to externalize memory as textual con- generation (RAG; Lewis al. 2020) and its
tent retrieved at inference time. variants (Guu et al. 2020; Karpukhin et al.) 2020 Borgeaud et al., 2022] Khandelwal et all 2020)
attach relevant textual chunks during inference, while RET-LLM (Modarressi et al. 2024 encodes

---

Published as a conference paper at ICLR 2026
knowledge symbolic triplets. Building these ideas, recent systems extend these ideas to conversational settings through hierarchical or summarization- based memory states; examples in- clude MemGPT (Packer et al. 2024), Memo (Chhikara et al. 2025), and A-Mem (Xu et al. (025 While effective for factual recall, these approaches are not optimized for procedural control and often incur significant inference-time overhead due repeatedly re-reading textual memory.
Parameter-Basee Memory. Another line work encodes memory directly into model parame- ters. Fine-tuning and multitask instruction tuning (Wei et al. 2022a| [anh al. 2022) allow models to acquire new procedures, but task knowledge often becomes entangled within the shared weights. Parameter-efficient fine-tuning methods such as LoRA (Hu et al. 2022) and Flora (Hao al. (024) disentangle new knowledge from the pre-trained base model, but still suffer from interference when leaming multiple tasks. Hixturreof-LLoke dynamically routing imputs to task-specifi experts: but the mixtures are typically invoked indepen- dently and are not designed for memory composition. Memoryl.LMM (Wang et al. 2024) introduces latent memory pools but remains entangled. Prompt-based methods such as prompt tuning (Lester) ctt al. 2021: Wu et al. 2024 2026) store knowledge implicitly as global embeddings without se- lective routing, and L2P (Wang et al. 2022b) introduces modular prompt pools but still relies on an external controller to determine which prompts are retrieved. Prompt compression methods (Mu et al., 2023] Chevalier et al. 2023) compress prompts into context-agnostic representations, which may distort prompt information. ToolGen (Wang et al., 2025) compresses tools into virtual tokens but focuses on post-training the backbone through multi-stage fine-tuning. By contrast, our Tok- Mem keeps the backbone frozen and introduces dedicated memory embeddings that can be added or composed without retraining, supporting continual adaptation.
(Feng al.) (024) address this entanglement issue by
Compositional Memory. A complementary direction explores how models compose skills from simpler building blocks. Chain-of- prompting (Wei al. 2022b) and tool-angmented rea- soning frameworks such as Toolformer (Schick et al. 2023) enable multi-step reasoning, but they rely on textual instructions that must be re-interpreted at each step. Modular parameter meth- ods (Rosenbaum et al.) 2018:) Pfeiffer et al.] (202T) create specialized adapters that can be recom- bined, but composition typically requires parameter merging or heuristic routing. TokMem differs by representing procedures as discrete tokens that can be chained directly in generation, enabling lightweight, parameter-isolated isolated composition.
CONCLUSION
We introduced TokMem, a parameter-efficient framework that encodes procedural memory as com- pact tokens. TokMem enables selective recall and compositional use of procedures without mod- ifying backbone parameters, achieving strong performance across multitask and nol--augmented reasoning benchmarks.
Future directions are discussed in Appendix C including the automation of sneeryproccedral de- composition for synthesizing training trajectories, reinforcement learning for stronger compositional generalization and personalization through user-specific memory banks. These extensions pave the way for scalable, compact, and user-adaptive memory systems large language models.
ACKNOWLEDGMENTS
We thank the reviewers and chairs for their efforts. The research is supported in part by the Natural Sciences and Engineering Research Council of Canada (NSERC), the Amii Fellow Program, the Canada CIFAR AI Chair Program, donation from DeepMind, and the Digital Research Alliance of Canada (alliancecan.ca)
REFERENCES
John R. Anderson and Christian Lebiere. The Atomic Components of Thought. Lawrence Erlbaum Associates, 1998. hRI.https: cooi..org/000 4324/498113111885555588
10

---

Published as a conference paper at ICLR 2026
Nabiha Asghar, Lili Mou, Kira A. Selby, Kevin D. Pantasdo, Pascal Poupart, and Xin Jiang. Progres- sive memory banks for incremental domain adaptation. In International Conference on Learning Representations, 2020. URL htlthttp:::
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Mil- lican, George Bm Van Den Driessche, Lean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, Diego De Las Casas, Aurelia Guy, Jacob Menick, Roman Ring, Tom Hennigan, Saffron Huang, Loren Maggiore, Chris Jones, Albin Cassirer, Andy Brock, Michela Pagannni, Geoffrey Irv- ing, Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack Rae, Erich Elsen, and Laurent Sifre. Improving language models by retrieving from trillions of tokens. In Proceedings of the 39th International Conference Machine Learning, pp. //proceedings. , mlr. press http:///////oeeeeeaaaaa.a.cc
2022. URL https:
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems, pp. 1877-1901, 2020.
URL htttps::/papppr:::/ b8ac 42444464444 -ilf6rAAAbbttttc..hhhh
CC
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374, 2021. BRRh:ttp::///
xiv. xiv.org/abs /aabb////002 03374
Alexis Chevalier, Alexander Wettig, Anirudh Ajith, and Danqi Chen. Adapting language models to compress contexts. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 3829-3846, 2023. URL https:://ccaatnooloryy wwtp://ooooyyyo//23....2/
emn ma in. 232/
Prateek Chhikara, Dev Khant, Saket Aryan, Taranjeet Singh, and Deshraj Yadav. Mem0: Building production-ready ai agents with scalable long-term memory. arXiv preprint arXiv:2504. 19413,
2025. URL pattps://arxiln wwtps//rgaaaa....aaaa//
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichrro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint 1rXiv:2110.1416 14168, 2021.
abs wtt.:///////...........
Wenfeng Feng, Chuzhan Hao, Yuewei Zhang, Yu Han, and Hao Wang. Mixtureoof-LoRAs: An efficient multitask tuning method for large language models. In Proceedings of the 2024 Joint International Conference on Computational Linguistics, Language Resources and Evaluation, pp. 11371-11380, 2024. URL
http://ccaaathoooco wrtp:/rrg////// 1re.2maireemmmmmm 994/
Kelvin Guu, Kenton Lee, Zora Tung. Panupong Pasupat, and Ming- Wei Chang. Retrieval augmented language model pre-training. In International Conference on Machine Learning, pp. 3929-3938, 2020. UBLL:tttp///poceeiiiiit
mlr. wwwppprsssse..............e vt///vlllll//////// 1111199guu2a..hmm
Yongchang Hao, Yanshuai Cao, and Lili Mou. Flora: Low-rank adapters are secretly gradient
compressors. In Forty-first International Conference on Machine Learning, 2024. URL httpst
7/openreview htt//formmpp um?
wwt.//ooom..duuuaaio. d-ubbB2K9999
John Hewitt. Initializing new word embeddings for pretrained language models. 2021.
https//nlp.stanffrd.cmm johnhew//vocab-exxpssii//
Saihui Hou, Xinyu Pan, Chen Change Loy, Zilei Wang, and Dahua Lin. Learning a unified classifier incrementally via rebalancing. In Proceedings of the IEEE/CVE Conference on Computer Vision
and Pattern Recognition, 2019. URL https: tps http::///ppccccsss.oo thecff.eee...cc. com/ Hou Learning Unified 2019 hap.r..tmlm.........
la Rebalanccng L wttp://www..ooooootettteeeeeeoee content
CVPR_
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Alhe Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Con-
ference on Learning Representations, 2022. URL
.ntt/forumtt
d=n7eeKeeFyFF

---

Published as a conference paper at ICLR 2026
Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. In Pro- ceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, pp- 6769-6781, 2020.
emn pomnlp-mainn .500/
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. In International Conference on Learn ing Representations, 2020. URL htttttt://pppreee/a//aa
net forum w Hkk
Yuri Kuratov, Mikhail Arkhipov, Aydar Bulatov, and Mikhail Burtsev. Cramming 1568 tokens into single vector and back again: Exploring the limits of embedding space capacity. In Proceedings of the 63rd Amnual Meeting of the Association for Computational Linguistics, pp. 19323-19339, 2025. URL /ttttts::/aaaattootolooo wwwpgy..rgg2222222 o
acl acl- long. 948/
Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficieet prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Pro- cessing, pp. 3045-3055 2021. URL htttps://cllanttolosee orgp://2orrr/220 emnlp-mai n. n4.44/
Patrick Lewis, Ethan Perez, Aleksandar Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval-augmentie genera- tion for knowledge-intensive NLP tasks. In Advances Neural Information Processing Systems, pp. 9459- 2020. URL http:::/proceeednnnggg neurips wtcpilp..cccppppe sh/ htw.9/88000007100022222666666660
ha
html
Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference Natural Language Processing, pp. 4582-4597, 2021. URL buttps://aclantteet http://aaaaaatttaoaaee org/ 2021 long. Long.3537
Chin- Lin. ROUGE: A package for automatic evaluation of summaries. In Proceedings of the Workshop on Text Summarization Branches Out, pp. 81- 2004. Uttp::tss///acanttt
logy. org/ W04 1013/
Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, pp. 157-173, 2024a. URL btttss:/acclatth
2ocy.org/2024.ta00 w2w.0002//2222...o
Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre- train, prompt, and predict: A systematic survey of prompting methods in natural language pro- cessing. ACM Computing Surveys, 2023. URL http:::tpp///..... //tp:s:///oiii/// org 1145/ /4455/5555155
Zuxin Liu, Thai Quoe Hoang. Jianguo Zhang, Ming Zhu, Tian Lan, Shirley Kokane, Juntao Tan, Weiran Yao. Zhiwei Liu, Yihao Feng, Rithesh R N. Liangwei Yang, Silvio Savarese, Juan Carlos Niebles, Huan Wang, Shelby Heinecke, and Caiming Xiong. APIGen: Automated Pipeline for generating verifiable and diverse function-calling datasets. In Advances in Neural Information Processing Systems Datasets and Benchmarks Track, 2024b. URL htttps://openreeee
et//orrm?idd-jfg3wee
Llama Team. The Llama 3 herd of models. arXiv preprint arXiv: 2024. URL https: //axxiv.org/abs////
21783
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human- level control through deep reinforcement learning. In Nature, pp. 529-533, 2015. URL
Ttttp:://ww..aa... www.:/atturee..eeeeeee
com/ wwtoo//atticees/oooomomooo wtw.riicciee/aaauee13333535
Ali Modarressi, Ayyoob Imani, Mohsen Fayyaz, and Hinrich Schüizee RET-LLM: Towards general read-write memory for large language models. arXiv preprint arXiv:2305.14333 2024. URL
https://erxivv. http:///wrxvv..o//././o///
14322
12

---

Published as a conference paper at ICLR 2026
Jesse Mu, Xiang Lisa Li, and Noah Goodman. Learning to compress prompts with gist tokens. In Advances in Neural Information Processing Systems, 2023. URL httt/t:tt///sp//e / et St/Foruum11-=2DxxPPC0
Charles Packer, Sarah Wooders, Kevin Lin, Vivian Fang, Shishir G. Patil, Ion Stoica, and Joseph E. Gonzalez. MemGPT: Towards LLMs as operating systems. arXiv preprini arXiv:2310.0855
2024. URL httthtt/p//axtttt /atpps//aaa...oorrr///a////////p 08560
Shishir G Patil, Huanzhi Mao, Fanjia Yan, Charlie Cheng-Jie Ji, Vishnu Suresh, lon Stoica, and Joseph E. Gonzalez. The berkeley function calling leaderboard (BFCL): From tool use to agen- tic evaluation of large language models. In Forty-second International Conference on Machine Learning. 2025. httpps::/ppperevieeeee wet.eee.eeeeeeee.eee Eorum? id- 2GmDthBdDD
Jonas Pfeiffer, Aishwarya Kamath, Andreas Rücklé, Kyunghyun Cho, and Iryna Gurevych. Adapter-
Fusion: Non-destructive task composition for transfer learning. In P Proceedings of the 16th Con- ference of the European Chapter of the Association for Computational Linguistics, pp. 487-503, 2021. URL.htts::/acllathoon wow.oooyy..ooooooooooooo 2221 eacl -main. 39/
Qwen Team. Qwen2.5 technical report. arXiv preprint arXiv:2412.151155 2025. URL https:
//arx.v.oro/abs/22222111
Nils Reimers and Iryna Gurevych. Sentence-BERT. Sentence embeddings using siamese BERT- networks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing, pp. -992. 2019. BuRlattps://arxii www.iii../g..o...oooooooo abs//s08//11 10084
Clemens Rosenbaum, Tim Klinger, and Matthew Riemer. Routing networks: Adaptive selection of non-linear functions for multi-task learning. In International Conference on Learning Represen- tations, 2018. http:://t///
net/ fote://fommmmmidryfdoeoe
Pranab Sahoo, Ayush Kumar Singh, Sriparna Saha, Vinija Jain, Samrat Mondal, and Aman Chadha. A systematic survey prompt engineering large language models: Techniques and applica- tions. arXiv preprint arXiv:2402.009927, 2025. URL http:::/aaxxxvv.r/./oo 7927
Victor Sanh, Albert Webson, Colin Raffel, Stephen Bach, Lintang Sutawika. Zaid Alyafeai, Antoine Chaffin, Arnaud Stiegler, Arun Raja, Manan Dey, M Saiful Bari, Canwen Xu, Urmish Thakker, Shanya Sharma Sharma, Eliza Szczechla, Taewoon Kim, Gunjan Chhablani, Nihal Nayak, De- bajyoti Datta, Jonathan Chang, Mike Tian-Jian Jiang, Han Wang. Matteo Manica, Sheng Shen, Zheng Xin Yong, Harshit Pandey, Rachel Bawden, Thomas Wang. Trishala Neeraj, Jos Rozen, Abheesht Sharma, Andrea Santilli, Thibault Fevry, Jason Alan Fries, Ryan Techan, Teven Le Scao, Stella Biderman, Leo Gao, Thomas Wolf, and Alexander M Rush. Multitask prompted training enables zero-shot task generalization. In International Conference on Learning Repre- sentations, 2022. URL http/ttpss//ppeeeriie.... http:iiwwwweeeeeoooooeeeeeeeee www.:/eeeeerrooiii////eooo0
Ignacio Sastre and Aiala Rosá. Memory tokens: Large language models can generate reversible sentence embeddings. In Proceedings of the First Workshop on Large Language Model Memo- rization, pp. 183-189, 2025. WtRLhttps:// htt///aaaaaatltooooeoe /
-.11/
Timo Schick, Jane Dwivedi-Yu, Yu, Roberto Dessi, Roberta Raileanu, Maria Lomeli, Eric Hambro, Luke Zettlemoyer, Nicola Cancedda, and Thomas Scialom. Toolformer: Language models can teach themselves to use tools. In Advances in Neural Information Processing Systems, 2023, URL hittps://openreet
net / forum iddidd=yaoo d=t.===mxx222777
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances Neural Infor- mation Processing Systems, pp. 5998-6008, 2017. URL http:itttttsssptss.srrt nips www.srssiiips.ccppppppp.a cc/pap er htt.a//pppee///aaaaaaaaaaaaaaaa0 act html
Renxi Wang, Xudong Han, Lei Ji. Shu Wang. Timothy Baldwin, and Haonan Li. ToolGen: Unified tool retrieval and calling via generation. In International Conference on Learning Representa- fions, 2025. httthhttts:///eeett www.eeeeeeeeeetttnieiiiiee..eeee
13

---

Published as a conference paper at ICLR 2026
Yizhong Wang, Swaroop Mishra, Pegah Alipoormolabashi, Yeganeh Kordi, Amirreza Mirzaei, Atharva Naik, Arjun Ashok, Arut Selvan Dhanasckaran, Anjana Arunkumar, David Stap, et al. Generalization via declarative instructions 1600+ NLP tasks. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp- org 2022 emn 1-mmnlpmmmnn-.. 3307
aaaperrNattulllnstuuttoonn
5085-5109, 2022a. URL
Yu Wang, Yifan Gao, Xiusi Chen, Haoming Jiang, Shiyang Li, Jingfeng Yang. Qingyu Yin, Zheng Li, Xian Li, Bing Yin, Jingbo Shang. and Julian McAuley. MemoryLLM: Towards self self-updatable large language models. Forty-first International Conference on Machine Learning, 2024. URI, https:://ppenreet ew net. forum? d=p0 kwzdikzzkkkk
Zifeng Wang, Zizhao Zhang, Chen- Yu Lee, Han Zhang, Ruoxi Sun, Xiaoqi Ren, Guolong Su, Vin- cent Perot, Jennifer Dy, and Tomas Pfister. Learning to prompt for continual learning. In Pro- ceedings the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 139-149, 2022b. URL pttttpss://penacccssss www.....c...ccccm. wttp:/t..ccoooooooeeoeee
Wang Learning To Pro_Prompt___ for_ Continnnnnnnnn....o.ooo..... Leattinna__earrinng..o CVPR 2022 p df
Jason Wei, Maarten Bosma, Vincent Y Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoc Le. Finetuned language models are zero-shot learners. Interna- tional Conference on Learning Representations, 2022a. URL https:
net
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, brian ichter, Fei Xia, Ed H. Chi, Quoc V Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems, 2022b. URL httpp:
review. www.neetfforuue 7id= ViQLMeSB_
Zijun Wu, Yongkang Wu, and Lili Mou. Zero-shot continuous prompt transfer: Generalizing task semantics across language models. In The Twelfth International Conference on Learning Repre- sentations, 2024. URL http:/ttppss/eenneveeeeeeeettet http:/eeeeeooooonoddd222e22hppo
Zijun Wu, Yongchang Hao, and Lili Mou. Ulra-llo-dimennsional prompt tuning via random projec- tion. arXiv preprint arXiv:2502.04501, 2026.
http://aaxdd..o//a//////0000 045
Wujiang Xu, Zujie Liang, Kai Mei, Hang Gao, Juntao Tan, and Yongfeng Zhang. A-Mem: Agentic memory for LLM agents. In Advances in Neural Information Processing Systems, 2025. URL
https://openrev//
net/ httt//oormmmddaeeeee FIMCMBgcct
Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik R Narasimhan, and Yuan Cao. ReAct: Synergizing reasoning and acting language models. In International Conference on Learning Representations, 2023. URL https: /owpnreviewwweeeeoooeeeeee net/
vlu
YUL-X X
14

---

Published as a conference paper at ICLR 2026
Indexing
Response
Response
<EOS>
Steoring
LLM Backbone
Query
Response
Final Response
Figure 5: Overview of the decoupled variant of TokMem embeddings, which learn separate memory matrices for memory indexing and generation steering.
A EXPERIMENTAL DETAILS
A.1 DETAILS OF SUPER- NATURAL INSTRUCTIONS
We sample 1,000 English tasks from the SNI dataset, where each task labeled with a task ID and short descriptive name. The full list of sampled tasks is provided in Figure Tasks are introduced to the model sequentially in ascending order of their IDs, for instance, the model first sees task 1, then task 2, and so on.
After training on the first tasks, we save a checkpoint and evaluate performance on the test set of each of the & tasks encountered far. This simulates continual learning setup where the model is expected to acquire new procedures while retaining previously learned ones. Once the model has been trained all 1,000 tasks, is expected be able to perform all of them without forgetting earlier tasks.
A.2 DETAILS OF DECOUPLED EMBEDDING FOR TOKMEM
In the standard TokMem formulation, each memory token embedding m; E Id is shared across two roles: (1) indexing for memory routing and (2) condition steer for generation. We consider decoupled (DC) variant that separates these functions into two embedding matrices:
M Mindex -uii.... ur) E Rx
M (8ii... 81) E R/xd
As seen in Figure 5 Mindex provides index embeddings at the output layer. The model produces
a distribution over indices according to Mindex when memory token is predicted. The chosen index then used to retrieve the corresponding steering embedding s; from M which is then appended into the sequence and conditions subsequent generation.
Training follows the standard next-token prediction objective, analogous to Equation B
L(a; M Mindex Mstter)
Σ log Pr(a: acs Mindex M
(8)
where & denotes the query length. During optimization, only Mindes and Stter updated, while the backbone remains frozen. addition, the renormalization treatment introduced Section 2.2 applied only the index embeddings Mindex
This decoupled formulation provides a clean separation of functionality: routing is handled via Mindex while steering is controlled by M However, experiments do not show consistent improvements over the coupled formulation, particularly on larger models where their larger em- bedding capacity is sufficient to jointly support both roles. We therefore advocate for the simple yet effective TokMem (without DC).
A.3 DETAILS OF FUNCTION- CALLING DATASET
To evaluate compositional memory recall, we sample 50 tools from the APIGen dataset (Liu et al. 2024b). The list of tools and their corresponding descriptions is provided in Table
15

---

Published as a conference paper at ICLR 2026
Algorithm Adaptation Phase for Compositional Memory Recall
Require:
Pre-trained backbone fau adaptation traces D Dadags from held-out procedures 1: Initialize the backbone B <- 8g and temporary memory embeddings. M
2: Set learning rates T/e and 17M
3: for each minibatch in Dndape do
4: Insert M into sequence; run forward pass with fo
5: Compute the loss L on memory and response tokens
6: - VoL
7: M TM TMMMC
8: end for
9: Discard the temporary memory M and freeze the backbone
10: return the adapted backbone fe
For each tool, we collect 50 query-call pairs, some of which may involve multiple calls to the same tool. We categorize the use the same tool as non-compositional tool use, yielding a total of 50 50 = 2,500 samples.
Additionally, we augment the dataset by synthesizing complex queries by combining calls across different tools. These multi-step queries require the model to invoke multiple tools in sequence. To avoid data leakage, we split these samples into training and test sets, capping the synthesized samples at 5,000 for training and 500 for testing.
A.4 DETAILS OF ADAPTATION PHASE
In compositional scenarios, the model should not only recall individual procedures but also compose them to solve multi-step queries. To prepare TokMem for such use, construct a held-out auxiliary training set comprising an additional detailed in Appendix A.3 The backbone is then fine-tuned for epoch on this set using LoRA, jointly with the temporary memory embeddings.
tools (5,000 samples) that disjoint from the tool set
The intuition behind this adaptation phase is to let the base LLM learn to the embeddings with compositional memory recall, rather than memorizing on the tool descriptions themselves. After adaptation, the temporary embeddings are discarded, while the adapted backbone is retained for training on new tools. This procedure provides a general inductive bias for modular composition, enabling the model generalize new tools and procedures without further retraining.
Algorithm summarizes this lightweight procedure. Temporary memory embeddings are inserted into the input sequence, the loss is optimized jointly over memory and response tokens, and the temporary memory bank is discarded once the backbone has adapted.
B ADDITIONAL RESULTS
B. RESULTS ON COMPOSITIONAL GENERALIZATION
We observe that TokMem provides clear advantages in compositional generalization over fine- tuning. reports Argument Fl when the Llama 3B model evaluated on queries that require more function calls than those observed during training.
Notably, when trained solely on single-call data, TokMem achieves much stronger performance than fine-tuning when evaluated test data with 2-4 calls. This demonstrates that memory tokens trained for atomic procedures can be effectively composed at test time, enabling strong zero-shot generalization to multi-step behavior.
As the training regime is expanded to include more calls (e.g., up to 3 or 4), the performance gap narrows, but TokMem remains competitive slightly ahead across configurations. These results suggest that TokMem naturally supports compositionality. enabling flexible chaining of learned pro- cedures without requiring task-specific fine-tuning.
16

---

Published as a conference paper at ICLR 2026
Table 5: TokMem generalizes longer tool-call chains at time, significantly outperforming
fine-tuning in zero-shot multi-step settings (trained with I call and tested with 2-4 calls).
Test Time
Train Maximum Calls
Method
calls
3 calls
calls
Avg.
Fine-tuning TokMem
34.9
21.3
14.1
23.4
1-call
60.3
54.3
48.9
54.5 (+31.1)
Fine-tuning
86.2
78.8
64.8
76.6
2-call
TokMem
82.0
81.8
82.3
82.0 (+5.4)
Fine-tuning TokMem
86.9
85,5
80.3
84.2
3-call
86.8
84.0
84.7
85.2 (+1.0)
Fine-tuning 87.9
86.6
82.9
85.8
4-call
TokMem
85.9
86.7
88.3
86.3 (+0.5)
6.0
5.
Norm
3.5
3.0
Training
Figure 6: &a norm of newly added memory tokens. Each round introduces 10 new tool memories. Error bars indicate the standard deviation across the 10 tokens added in each round.
B.2 ANALYSIS NORM INFLATION FOR NEWER MEMORIES
We analyze the E norm of the learned memory embeddings when tools are introduced sequentially with Llama 3.2 3B. We follow the setup in Figure by training TokMem for rounds and adding 10 tools per round without our renormalization treatment detailed in Section 2.4] Figure 6 shows that newly added memory tokens gradually increase their & norms, which leads competition with existing frozen tokens in the softmax used for memory routing.
B.3 ANALYSIS OF UNFREEZING LLM BACKBONE FOR TOKMEM
We further examine the importance freezing the backbone when adding new tool memories, reflecting real-world usage where procedural knowledge grows incrementally time. This setting contrasts with recent approaches (Wang et al., (025) that aim to compress tool use description by post-training the backbone. Such methods rely on modifying backbone parameters, which hinders continual adaptation and risks overwriting prior knowledge.
As shown Figure 7 We confirm that unfreezing the backbone during TokMem adaptation leads to severe forgetting of previously learned tools. By contrast, our choice of freezing the backbone meaningfully preserves prior capabilities while allowing new tool memories to incorporated with- out loss. highlighting TokMem's 's advantage for incremental adaptation. This analysis suggests that TokMem strikes effective balance between performance and continual adaptation.
B.4 ADDITIONAL ANALYSIS ON MEMORY PLACEMENT
In Section 3.4 have stress-tested the memorization with different memory token placement (pre- fix vs. infix) by varying numbers of memory tokens. We now turn to the GSM8K math reasoning dataset (Cobbe al. 202T) to evaluate generalization and training sample efficiency.
17

---

Published as a conference paper at ICLR 2026
Liama
Liama
Liama
Argument F1 Score
20
80
80
60
60
60
20
20
20
Training Round
Training Round
Training Round
Frozen LLM
Treinable LLM
Tools 1-10
Tools 11-20
Tools 21-30
Tools 31-40
Tools 41-50
Figure 7: Comparison between freezing and unfreezing the backbone. Allowing the backbone to update when adding new tool memories severe forgetting, whereas freezing preserves prior tools while enabling new ones.
Table 6: Comparison of prefix tuning and TokMem conditioning embeddings on GSM8K with two different Llama model sizes. TokMem achieves higher compliance with required output formats and stronger exact-match accuracy than prefix tuning, especially in low-data regimes.
Llama 3.2 1B
Llama 3.2 3B
Data %
Method
Compliancet
EM
Compliancet
EM
Prefix tuning
0.0
0.0
45.9
33.1
20%
TokMem
98.0
37.7
94.6
65.6
Prefix tuning
82.8
30.0
97.2
64.1
100%
TokMem
97.4
39.1
98.2
66.9
Our experiments use Llama 3.2 1B and 3B backbone models and compare prefix tuning against TokMem under two training setups: using only 20% of the training set (a low-data regime) or the full dataset. We report two evaluation metrics.
Compliance measures whether the model follows the required answer format, i.e., producing the final answer after the delimiter "####" This metric isolates procedural-memory recall from the reasoning abilities already present in the backbone.
Exact Match (EM) measures the correctness of the final answer after standard text normalization (e.g., removing commas or extraneous symbols).
As shown in Table 6, TokMem significantly outperforms prefix tuning, particularly in the low-data setting. With only 20% of the data, prefix tuning fails to provide meaningful results, yielding zero compliance and EM on the 1B model and underperforming on the 3B model. By contrast, Tok- Mem achieves pear-perfect compliance and substantially higher EM scores across both backbones. When trained on the full dataset, prefix tuning improves considerably, yet TokMem continues to deliver stronger compliance and higher EM, showing its superior data efficiency and more reliable procedural control.
C FUTURE WORK
Our experiments are conducted on the research-oriented SNT and APIGen datasets, which provide controlled environment for analyzing atomic and compositional recall. While these settings demon- strate the feasibility and effectiveness of one-token procedural memory without backbone training, they do not fully capture the diversity and ambiguity of real-world procedures.
In particular, richer forms of composition, such as interleaving function calls from APIGen with NLP tasks from SNI, are bottlenecked by TokMem's 's reliance manually curated, procedurally decomposed datasets. To bridge this gap. a future direction is the automation of query-procedie decomposition. By leveraging highly capable teacher LLMss synthesize interleaved training tra-
18

---

Published as a conference paper at ICLR 2026
jectories, automatically map complex queries to discrete procedural tokens, providing
scalable alternative to manual annotation.
Beyond automated data synthesis, additional future directions include incorporating reinforcement learning to improve generalization over complex compositional structures, and enabling personal- ization by allowing users attach custom, user-specific memory banks while keeping the backbone frozen. These extensions pave the way for scalable, compact, and user-adaptive memory systems in large language models.
D USE OF LARGE LANGUAGE MODELS
We used ChatGPT general-purpose assistant to improve the writing of our paper, including grammar, readability, and clarity. Additionally, we used it to search for potentially missing related work, which we then manually read and discussed in the paper. All ideas, analyses, and conclu- sions presented in this paper are our own, and we take full responsibility for the content.
19

---

Published as a conference paper at ICLR 2026
Table 7: Details (names and descriptions) of the sampled tools from the APIGen dataset.
ID
Tool
Description
cato_compeete
Fetch nuto-complete suggestions for given query using the Wayfair API. Adds two binary numbers and returns the result as a binary string.
2
binary addiiion
hinary search cagr
Performs binary search sorted list find index of target
4
Calculates the Compound Annual Growth Rate (CAGR) of investment.
calculate factorial grlcul
Calculates the factorial of a non-negative intoger.
Calculates the weighted average grade based scores and their correponddnng weights.
calculate uediian can att.adttcaa.. all meetings
Calculates the median of a list numbers. Determines if a person can attend all meetings given a list of meeting time intervals.
8
9
cosine similarity
Calculates the cosine similarity betwoen two voctors, Counts the number of set bits (1's) in the binary representation number. Create histogram based on provided data.
10
count_bis
11
create histogram
12
Fetches the route information hetween two geographical locations including distance, duration, and steps,
13
Abbonacci
Calculates the nth Fibonacci number. Calculates the final velucity an object given its initial velocity, acceleration, and time.
final, velocity
15
lind. codd..uulibriumminie
Finds the equilibrium index of a list, where the sum of elements on the left is equal to the of elements the right.
16
find mmn_firs_noneeeatting Jepeaatnnnnngg....mmm find longest word
Finds the first non-repeating character in string.
17
Finds longest word list of words.
18
find mas suburay sum find minimum rolated yoreed aarra Hasten
Finds the maximum sum of contiguous subarray within a list of integers. Finds minimum
rotated sorted array.
20
Hattens nested list into single-level list.
format date
Converts a date string from one formal to another.
22
generate passward
Generates random password of specified length and character types. Generates random string of spocified length and character types. Retrieves the city name a given ZIP code using the Zipustic API. Retrieves information about a Pokémonn's 's move using the PokéAPI Fetches product details from an using given product
23
generate_randon random_string
get city zipcodo
25
get pol.pokmoon. move info
get get_pruduct
27
get we..prddatt.inn.coue .orcllgrrr.... greatest commnon. divisor integrate
Fetches products specilied calegory from the demo project's catalog. Computes the greatest common divisor (GCD) of two non negative integers. specified function between two values. Calculates the profit from an investment based on the initial amount, annual return rate, time.
28
Calculate the area under
30
investment.pret proffit
31
is M.Lamarram,pnn phrase
Checks two phrases are anagrams of each other, ignoring whitespace and ponctuation.
32
Jeup
Checks if year leup year.
33
is palindrome
Checks if string is palindrome.
34
is. popower
Checks number power given base.
35
rstattooon
Chocks one string is a rotation of another string.
is.valdd ip.iddiss.ddresss
Checks string valid address (IPv4).
37
is yalid, palindreee
Checks a string is valid palindrome, considering only alphanumeric charac. ters and ignoring case.
38
valid sudoku montlyy mortgage.payy payment
Checks a 9x9 Sudoku board is valid. Calculates the montlly morigage payment based on the loan amount, annual interest rate, and loan
39
40
note_duration
Calculates the duration between two musical notes based on their frequencies and tempa.
41
place safeway_annaaa
Order specified items from a Safeway location.
42
polygon_area aoro_aarr_sshllee
Calculates the area of a polygon using the shoclace formula. Calculates the electrostatic potential energy given the charge and vollage.
43
potential energy
44
popaaatpoon
Projects the population size after a spocified number of years. string.
reverse.string
Reverses characters
46
sulve quadratic
Computes the roots of quadratic equation given its coeflicients.
tripezoidal, integration whois 48
Calculates the definite integral function using trapezoidal nule.
Fetch the WholS lookup data for given domain using the specified Toolhench
RapidAPI key.
49
whole foods order rerssresit
Places order Whole Foods.
50
Calculates the resistance of wire based on its length, cross-sectional area, and
material resistivity.
20

---

Published as a conference paper at ICLR 2026
Figure 8: Overview of the 1,000 English tasks from the SNI dataset used in the atomic recall setting.
21


Reviews:
[{"review_id": "4Kvpyf3RYv", "reviewer": "Reviewer_9FXM", "summary": "This paper introduces TokMem, a trainable memory module that can augment existing language models and steer their behavior. The memory module is represented as a bank of memory tokens that can be retrieved and appended to the context on the fly, and each memory token is trained to be associated with a response. Only the memory embeddings are updated during training while the base model remains frozen.\nThe experiments include popular approaches, such as ICL, RAG,  FT, and reply memory, and test them on Super Natural Instructions and APIGen, representing atomic and compositional memory.", "questions": "How are the memory tokens retrieved during inference?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper presents a novel approach to managing memory for language models by fine-tuning just a memory module. The results are solid with comprehensive comparisons with different baselines. The design choices are solid and backed by ablation studies. It’s also great that the paper consider different types of memories like atomic and compositional memory.", "weaknesses": "In terms of the experimental settings, the original Super Natural Instruction paper uses unseen tasks for testing, but this paper uses the same tasks for training and testing. As a result, it’s unclear how the TokMem can help existing language models generalize to new tasks. Do you have evaluations on applying TokMem to the seen unseen tasks?\n\nThe presentation could benefit from how the memory tokens are used during inference. (see below for some questions on the procedure)\n\nTypo: line 312 “fine-tine” → “fine-tune”", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kyJinyIAQ5", "reviewer": "Reviewer_3EcY", "summary": "The paper introduces a method augmenting a language model with a bank of trainable “procedure” embeddings. The model is trained on sequences of query tokens with interleaved \"procedure\" embeddings and response tokens.\n\nThe method is evaluated on \"Atomic Memory Recall\" and \"Compositional Memory Recall\" setups.", "questions": "All the questions are listed in the weaknesses above.", "limitations": "", "rating": null, "confidence": 2, "soundness": 1, "presentation": 1, "contribution": 2, "strengths": "At present I’m unable to identify clear strengths because the core method and its usage at inference are unclear. As a result I can’t assess the contribution or empirical value with confidence.\n\nConceptually, representing procedures as tokens could be interesting.", "weaknesses": "The paper is highly unclear and have multiple problems, for example:\n1. The inference process is not described at all. During training, the model learns to predict a sequence of [query tokens]+interleaved [memory token][procedure tokens] -- what happens during inference? Are we appending procedure tokens to the input? In such case, it contradicts the motivation where the method is presented in opposition to e.g. RAG methods that inflate the context window. If the method is not injecting textual tokens of the \"procedures\", then it's another PEFT method, and should be compared to different PEFT baselines, including prompt-tuning, as well as other dynamic adapter approaches like Mixture-of-LoRAs.\n2. As the inference is unclear, so is the evaluation. For the \"Atomic Memory Recall\" setup, the tasks from Super-Natural Instructions dataset are regarded as the \"procedures\" in the memory bank. What is exactly evaluated here? Given test query, model predicts the memory token, then we append the text tokens of corresponding \"procedure\" and evaluate that with Rouge-L?\n3. Model is trained on data in the sequential order, which is contrary to the standard of shuffling the data for _stochastic_ gradient descent.\n4. Fine-tuning baseline includes only LoRA applied to query & key projections of the attention layers. It should involve also full fine-tuning, or at least LoRA applied to all linear layers, and not an arbitrary subset of them.\n5. What was the procedure of selecting hyperparameters for training? For example, learning rate for training LoRAs seems to be too low. Moreover, optimal learning rate usually differs across model sizes -- here the authors provided only a single value.\n6. \"Routing\" is mentioned multiple times in the paper but never introduced or defined.\n7. The method is described (e.g. in the abstract) as keeping the backbone frozen. However, in the second evaluation setup the model is initially finetuned before applying the TokMem.\n\nAs the authors claim in the Appendix, the paper was written with the help of ChatGPT, which might partially explain the level of presentation of the paper.\n\nAfter clarifications, I’m open to revising my score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lGGoqme4W2", "reviewer": "Reviewer_unGL", "summary": "This paper proposes TokMem, a tokenized memory that stores recurring procedures as compact, trainable embeddings. In particular, frequently reused procedures are “compressed” and stored by encoding them into an internalized memory token. To support continual adaptation, TokMem keeps the backbone model frozen, allowing new procedures to be added without interfering with existing ones.", "questions": "1. Is TokMem+DC using two tokens for each procedure/task? If yes, is it a fair comparison to TokMem? And can you further extend to more tokens?\n\n2. Why is renormalization needed, what is the intuition behind? I see it's helpful empirically but not sure if I understand the reason at Line 150-151. In particular, why new embeddings would develop inflated norms?\n\n3. Clarity: I found the following points confusing to me, requiring more time to catch the ideas.\n\n    - At Line 214-215, it is mentioned that both RAG and TokMem are with 'explicit' memory routing. IMO TokMem is doing 'implicit' routing since the memory token is chosen via next-token prediction, unlike RAG or MoE who has dedicated routing mechanism/architectural component. Can you clarify?\n    - Are the training examples formulated by inserting memory tokens between query and response of the original datasets?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is generall well-written and easy to follow.\n\n- The proposed method is compared to multiple valid baselines, with complementary experiments to understand the design choices.\n\n- Memory tokenization naturally reduces context size and supports continual learning without interfere\n\n- The backbone is frozen during the training, offering effciency and avoiding catastrophic forgetting.", "weaknesses": "The method requires an additional adapation stage for compositional tasks, which seems to break the 'frozen backbone' claim? Also, is it fair to compare 'TokMem + adapt' to fine-tuning? While it is true that the stored procedures are modular via independent memory tokens, the comparison of TokMem without adaptation to fine-tuning in Table 3 seems to suggest limited capability of composing the modular procedures.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xmBYjbgHt0", "reviewer": "Reviewer_BfBf", "summary": "The paper studies the problem of efficiently storing and applying frequently-used procedures (such as calling certain tools) in large language models as opposed to using lengthy prompts or retrieval mechanisms. Specifically, the paper proposes TokMem, which encodes procedural knowledge as trainable continuous memory tokens. The paper evaluates on two settings: using TokMem for single tasks in Super-Natural Instructions and compositionally invoking on function-calling tasks. The results suggest that TokMem outperforms retrieval-augmented generation and parameter-efficient fine-tuning while using fewer parameters.", "questions": "See weakness", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "The experimental setting covers multiple models across different scales. The experiment also covers both single task setting and compositional setting.\n\nThe results demonstrate good empirical results of TokMem, showing its improvements over LoRA finetuning.\n \nThe proposed approach is carefully designed. For instance, TokMem uses the renormalization technique to prevent new tokens from dominating routing.", "weaknesses": "I feel the paper needs substantial improvements in terms of clarity.\n* The paper heavily uses the term \"procedure\", but does not give a precise definition of \"procedure.\"\n* In particular, it's unclear whether the response in line 129 represents a fixed set for each procedure or can vary.\n* The implementation details of memory tokens are ambiguous. Are they essentially implemented as newly introduced special tokens added to the vocabulary?\n* I am also not so sure about the routing mechanism. Does the model predict a newly added special token that indexes a procedure during inference?\n* DC (decoupled embeddings) seems to also play an important role in some results but it is mainly covered in Appendix.\n\nThe choice of baselines mainly uses LoRA finetuning. It would be beneficial to include full fine-tuning as a reference to discuss the trade-offs between the number of parameters updated and performance.\n\nThe paper misses discussion of important related work on prompt compression. Methods like Gisting (Mu et al., 2023) and AutoCompressor (Chevalier et al., 2023) that compress prompts into tokens share similar motivations and should be\n\nI'm curious whether new special tokens are necessary. What if we used short natural language descriptions with special token wrappers like <call>tool name</call>, but similarly only tuned representations of these tokens? This approach wouldn't require adding new embeddings for every new procedure.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)