1 Introduction↩︎

Despite the impressive capabilities of Large Language Models (LLMs), a long-standing challenge persists: how can a deployed LLM continually absorb new knowledge, while faithfully retaining what has already been stored? Previous solutions can be broadly grouped into three classes. (1) Token-level memory [1], [2] stores past information as explicit text tokens, while it is bottlenecked by the context window and is vulnerable to the lost-in-the-middle phenomenon [3]. (2) Parametric memory either encodes knowledge into the base model [4] or into auxiliary adapters such as LoRA [5]. This form incurs no extra inference cost but requires costly tuning and is susceptible to catastrophic forgetting [6]. (3) Latent memory carries information implicitly in the model’s internal representations such as KV-caches, activations, or per-layer hidden states. The long-term form (e.g., MemoryLLM [7], M+ [8], and NextMem [9]) maintains a persistent pool of per-layer hidden states. By avoiding the repeated re-encoding of the full context and compressing information more densely in latent representations than discrete text tokens, this latent paradigm offers a compelling foundation for continually updatable LLMs.

However, the current long-term latent memory paradigm encounters the following limitations. (1) Positional encoding misalignment: When a new knowledge fragment is appended, the concatenated memory fragments are indexed with positions that no longer match those they were formed at, distorting the attention computation at inference time [10]. (2) The absence of any tracing mechanism: Existing methods treat all memory fragments as a flat, undifferentiated prefix and cannot dispel the noise information or single out fragment(s) relevant to the current query.

To address these issues, we propose MemDefrag, a training-free and model-agnostic framework that reinvents the way long-term latent memory is used and evolved. MemDefrag is fully plug-and-play: it requires neither additional training nor auxiliary modules to the underlying LLM. Specifically, the contributions of this work are:

(1) We first uncover an inherent tracing signal in transformer layers. Inspired by the query-aware sparsity of tracing token criticality across layer-wise KV cache pages during self-attention [11], we first ask a diagnostic question: is there any intrinsic signal in a transformer, that can trace specific relevant and crucial fragments inside the latent memory? By probing the layer-wise attention density that each stored fragment receives with respect to a given prompt, we find a clear and robust pattern: a small set of middle transformer layers consistently concentrate the highest attention density on the target fragment.

(2) In light of this finding, we propose inference-time memory defragmentation. At each prompt, MemDefrag computes the attention density of every stored fragment at a tracer layer and conducts defragmentation (ranks, reorders, and filters the memory fragments by their density). This simultaneously mitigates the irrelevant semantic noise and the accumulated positional distortion.

(3) Considering a bounded memory and compute budget, we propose an informativeness-based proportional forgetting strategy for memory evolution. When memory capacity is reached, MemDefrag allocates forgetting quotas proportionally across fragments, and within each fragment, evicts the tokens with the lowest self-information.

(4) Experiments show that MemDefrag substantially outperforms MemoryLLM and M+ on both knowledge retention and long-context QA. For instance in knowledge retention, on NaturalQA at step 50, MemDefrag (based on Llama-3.1-8B-Instruct) retains 43.0% accuracy, compared with 17.4% and 17.6% for the two baselines. Moreover, experiments support the generality and compatibility of our MemDefrag across models, latent-memory variants, and prompt compression.

2 Investigation: Positional Distortion of Latent Memory↩︎

2.1 Preliminaries: Vanilla Latent Memory↩︎

To introduce the latent memory paradigm and investigate its positional distortion, we first introduce a vanilla framework of latent memory, aiming to exclude irrelevant or redundant designs.

A language model equipped with latent memory is denoted as \(\mathcal{M}_{\theta,\phi}\) comprising two sets of parameters: memory parameters \(\theta\) and model parameters \(\phi\). Specifically, \(\phi\) consists of \(L\) transformer layers, denoted as \(\phi = \{\phi^l\}^L_{l=1}\). The memory parameters \(\theta\) are prefix hidden vectors within each transformer layer, denoted as \(\theta = \{\theta^l\}^L_{l=1}\). Each \(\theta^l\) has dimension \(N \times d\), corresponding to \(N\) hidden states with hidden dimension \(d\) in \(\phi\).

Model parameters \(\phi\) remain static, while the memory parameters \(\theta\) evolve upon encountering new knowledge. Crucially, in a training-free manner, the evolution requires neither gradient descent nor backpropagation. Given a new textual knowledge fragment as \(x\), the evolution is carried out by an update function \(U_{\phi}\) (which is based on the model parameters \(\phi\)), that stores \(x\) into \(\theta\) to yield updated memory parameters \(\theta^\prime\): \[\theta^\prime = U_{\phi}(\theta,x).\]

Extending this to multi-step, long-term evolution, consider a sequence of textual knowledge fragments \((x_1,\dots,x_n)\), where \(x_n\) denotes the knowledge fragment received at time step \(n\). The model successively integrates all knowledge to obtain the memory parameters \(\theta_{n}\) with its length \(N_n\): \[\theta_{n} = U_\phi(\cdots(U_\phi(\theta_0,x_1)),x_n).\]

2.1.0.1 Memory Formation by Prefilling

Consider storing a new knowledge fragment \(x_n = [x_n^{(0)}, x_n^{(1)},\dots,x_n^{(|x_n|)}]\), a sequence of tokens at time step \(n\). The formed memory fragment \(m_n = \{m_n^l\}_{l=1}^{L} = f_{\text{form}}(x_n)\) for \(x_n\) consists of the per-layer hidden states produced by passing (prefilling) \(x_n\) through the model \(\mathcal{M}_{\theta = \emptyset,\phi}\) (i.e., without latent memory): \[\quad m_n^l = \phi_l(h_n^{l-1}) \in \mathbb{R}^{|x_n| \times d},\] where \(h_n^{l-1}\) denotes the input hidden states to layer \(l\), with \(h_n^{0}\) being the token embeddings of \(x_n\).

2.1.0.2 Memory Evolution by Concatenation

Memory parameters are initialized as empty, i.e., \(\theta_0 = \emptyset\) (with initial length \(N_0 = 0\)). Given the current memory parameters \(\theta_{n-1}\) and the formed memory fragment \(m_n\), the evolution proceeds by concatenation: \(\theta_n = \text{concat}(\theta_{n-1}, m_n)\).

2.1.0.3 Inference with Latent Memory

During inference, at each layer \(l\), latent memory \(\theta_n^l\) serves as the prefix of the hidden states \(h^l\).

For the whole processes of the paradigm of vanilla latent memory, Figure 1 shows the sketch of memory formation, evolution and model inference. MemoryLLM [7], M+ [8], and NextMem [9] represent the latent memory paradigm, and follow the vanilla memory evolution by concatenation and inference with the whole latent memory.

From Figure 1, it can be inferred that, in such systems, the position encodings of concatenated memory fragments are misaligned with those expected during inference. Note that this positional distortion refers to the position encodings of memory fragments (including the target fragments and other fragments) being misaligned with those expected. This distortion can be accumulated and significantly degrade generation quality [10], [12]. We are therefore motivated to examine the effects of positional distortion.

2.2 Effects of Positional Distortion↩︎

To investigate how positional distortion affects inference performance, we conduct experiments based on vanilla latent memory, using Llama-3.1-8B-Instruct³ with greedy decoding (max_new_tokens = 32). The dataset used is NaturalQuestions (NaturalQA) [13]⁴, where each item is a (context, question, answer) triple. We sample 500 groups for evaluation. Within each group, the protocol (largely following the setup of MemoryLLM and M+) is defined as follows:

(1) Each context serves as the knowledge \(x_i\) (\(|x_i| \leq 512\)) stored at time step \(i\).

(2) After every memory update that obtains \(\theta_i\) (\(1 \leq i \leq 20\)), the first knowledge \(x_1\) is evaluated.

(3) The model \(\mathcal{M}_{\theta_i,\phi}\) is queried with the question corresponding to \(x_1\). A response is deemed correct if the gold answer appears as a substring of the model output.

Figure 2 (a) shows that the accuracy drops sharply as \(n\) and \(N_n\) (memory length at step \(n\)) increase, with performance collapsing once \(N_n\) exceeds 2000. It indicates that positional distortion of latent memory is much more severe than the degradation observed in the phenomenon [3] of long contexts where positional encodings are correctly assigned throughout the input sequence. Experiments and analyses based on the setting without positional distortion are shown in Appendix 9.1.

2.3 Attention Density as an Indicator↩︎

Although positional distortion is a fundamental issue inherent to the latent-memory paradigm, it is not directly addressable within this framework. This motivates us to seek a quantifiable indicator correlated with performance degradation. To this end, we introduce layer-wise attention density \(\rho_i^l\) for each knowledge item \(x_i\), which captures query-aware cruciality at the granular level of fragments. A related but distinct line of work, exemplified by SnapKV [14], uses per-token attention cruciality for KV-cache compression rather than long-term memory retention.

2.3.1 Attention Density↩︎

We define a per-layer, per-fragment relevance score \(\rho_i^l\) called the attention density of the stored knowledge fragment \(x_i\) at layer \(l\), which measures how much attention the current prompt \(p\) pays to \(x_i\).

Let \(l\) denote a designated transformer layer. At inference time, the memory is stored as a per-layer prefix: the layer-\(l\) input is formed by prepending the bos token and the stored memory hidden states to the prompt’s layer-\(l\) states, \[\label{eq:concat-seq} H^{l} = [\texttt{bos}^{l};\, m_1^{l};\, \dots;\, m_n^{l};\, H_p^{l}] \in \mathbb{R}^{S \times d},\tag{1}\] with \(S = 1 + N_n + |p|\) and \(N_n = \sum_{i=1}^{n} L_i\) the total number of retained memory positions (where \(L_i\) is the current retained length of \(x_i\)). Since only the prompt tokens need to query the memory, we run multi-head self-attention on \(H^{l}\) under a causal mask, restrict the queries to the \(|p|\) prompt positions, and average across heads. This yields a head-averaged prompt-to-sequence attention matrix \(\bar{A}^{l} \in \mathbb{R}^{|p| \times S}\) (full derivation in Appendix 8.1). Taking its last row gives the per-position attention score from the prompt’s last token, \(\bar{a}_t^{l} = \bar{A}_{|p|,\, t}^{l}\).

We then collapse these per-position scores into a single relevance score for each fragment \(x_i\) by averaging over the positions \(P_i\) that \(x_i\) occupies in the concatenated memory (excluding bos): \[\label{eq:density} \rho_i^l \;=\; \frac{1}{L_i} \sum_{t \in P_i} \bar{a}_t^l, \qquad i = 1, \dots, n,\tag{2}\] Intuitively, \(\rho_i^l\) summarizes how strongly layer \(l\) routes the prompt’s attention toward the memory region occupied by \(x_i\).⁵

2.3.2 Analyses of Attention Density↩︎

Figure 2 (b) presents a heatmap of the attention density \(\rho_1^l\) for the target knowledge across each layer \(l\) and time step \(n\), following the setup described in Section 2.2. A clear trend emerges: in several layers (particularly layers 13–20) attention density decreases markedly as the time step increases. Figure 2 (c) uses two curves to show that QA accuracy is highly aligned with the attention density allocated to the target memory fragment at layer 13.

To quantify the layer-wise correlation between accuracy degradation and attention density reduction, we compute the Pearson and Spearman correlation coefficients at each layer \(l\), as shown in Figure 2 (d).⁶ Notably, both coefficients exceed 0.95 across the majority of layers 13-20, indicating a near-perfect monotonic relationship between attention density reduction and accuracy decline. This strong correlation suggests that attention density is a key indicator associated with accuracy.

Moreover, according to Eq. 2 and its previous derivation, a straightforward approach to increasing the attention density (correlated with the QA performance) of the target memory fragment is to remove non-target ones (temporarily) during inference. In this case, the first step is to trace and identify the target memory fragment.

3 Attention Density as a Tracing Signal↩︎

In this section, we provide an insight that attention density serves not only as an indicator correlated with performance, but also as a tracing signal: Certain layers exhibit an inherent tracing capability, consistently identifying the target fragment through attention density ranking. Our insight can be illustrated by the following experiments.

We evaluate 500 groups from the NaturalQA dataset. For each group, 20 knowledge fragments are stored, where the target fragment is placed at each of the 20 positions in turn (from first to last) to mitigate positional bias. Table 1 reports the average attention density rank of the target fragment across all positions for each layer. Layer 13 achieves the lowest mean rank of 1.66, with a Top-1 accuracy of 85.6% and a Top-3 accuracy of 95.7%, indicating that it consistently assigns the highest attention density to the target fragment regardless of its position.

Notably, this aligns with findings in activation steering studies [15]–[17], which identify specific middle layers as the most effective intervention point to control model behaviors. One possible explanation is that intermediate layers can encode richer representations than both earlier layers and final layers, holding implications for feature relevance and extraction [18], [19]. Further, analyses about the tracing capabilities based on the setting without positional distortion are shown in Appendix 9.2.

4 MemDefrag↩︎

We propose MemDefrag of which latent memory defragmentation, given a prompt, traces and consolidates the most relevant memory fragments for inference while discarding irrelevant ones to mitigate accumulated positional distortion and to improve the attention density allocated to the relevant fragments. Tracing the relevant fragment(s) is guided by attention density at one tracer layer \(l^*\) (such as layer 13 or 9 listed in Table 1 in Section 3).

As MemDefrag is generally applicable to latent memory designs, without loss of generality, we present it as an extension of the memory formation and concatenation-based memory evolution in vanilla latent memory (Section 2.1).

4.1 Inference-Time Memory Defragmentation↩︎

At inference time, given a prompt \(p\) and the accumulated memory \(\theta_n = \{m_1, \dots, m_n\}\), the function \(f_{\text{defrag}}(\theta_n, p)\) defragments them to a prompt-specific inference-time temporary memory \(\theta_n^p\).

To defragment the whole memory \(\theta_n\), all \(n\) memory fragments are first sorted by their attention density in ascending order: \[\pi = \mathrm{argsort}(\rho^{l^*}_1, \dots, \rho^{l^*}_n),\] such that \(\rho^{l^*}_{\pi(1)} \leq \rho^{l^*}_{\pi(2)} \leq \cdots \leq \rho^{l^*}_{\pi(n)}\). The entire memory \(\theta_{n}\) is then reordered by \(\pi\) at each layer: \[\tilde{\theta}_n^l = [m_{\pi(1)}^l;\, m_{\pi(2)}^l;\, \dots;\, m_{\pi(n)}^l], \;l = 1, \dots, L.\] This places the most relevant memory fragments at the end (i.e., closest to the prompt), thereby leveraging the recency bias of causal attention.

Finally, we filter only the Top-\(K\) most relevant memory fragments by truncating the reordered sequence from the front: \[\theta_n^p = \{m_{\pi(n-K+1)},\, \dots,\, m_{\pi(n)}\},\] i.e., the hidden states at each layer are: \[(\theta_n^p)^l = [m_{\pi(n-K+1)}^l;\, \dots;\, m_{\pi(n)}^l] \in \mathbb{R}^{N_K \times d},\] where \(N_K = \sum_{k=n-K+1}^{n} L_{\pi(k)}\), in which \(L_i\) denotes the current retained length of \(x_i\). The prompt-specific memory \(\theta_n^p\) is then used for inference: \(\text{output} = \mathcal{M}_{\theta_n^p, \phi}(p)\). Algorithm 4 shows the process of defragmenting \(\theta_n^p\) from \(\theta_n\).

4.2 Memory Evolution with Informativeness-Guided Forgetting↩︎

Beyond defragmentation, our MemDefrag further incorporates an informativeness-based forgetting mechanism for memory evolution (e.g., updating \(\theta_{n-1}\) to \(\theta_n\)). The presence of forgetting is essential for latent memory, as practical deployment constraints demand an effective eviction strategy.

Unlike MemoryLLM [7] and M+ [8] that randomly drop memory states independently at each layer, indiscriminately losing information, MemDefrag introduces a selective strategy that operates coherently across all layers of each token.

4.2.1 Sketch of Memory Evolution↩︎

Given the current length \(N_{n-1}\), the formed memory \(m_{n}\) of new knowledge \(x_{n}\), and a maximum length \(N_{\text{max}}\) (set empirically), \(N_{\text{forget}} \;(N_{\text{forget}} = N_{n-1} + |x_{n}| - N_{\text{max}})\) hidden states are needed to be dropped. The evolution proceeds in two cases: \(\theta_n = U_{\phi}(\theta_{n-1}, f_{\text{form}}(x_n)) = U_{\phi}(\theta_{n-1}, m_n)\) equals \(\text{concat}(\theta_{n-1}, m_n)\) when \(N_{\text{forget}} \leq 0\), and equals \(\text{concat}(f_{\text{forget}}(\theta_{n-1}, N_{\text{forget}}), m_n)\) when \(N_{\text{forget}} > 0\), where \(|x_n| \leq N_{\text{max}}\), and the function \(f_{\text{forget}}(\cdot)\) prunes \(N_{\text{forget}}\) states from \(\theta_{n-1}\) to store the incoming memory. Algorithm 8 in Appendix sketches the memory evolution with the forgetting mechanism.

4.2.2 Informativeness-Based Proportional Forgetting of Memory↩︎

Our forgetting function \(f_\text{forget}(\cdot)\) performs proportional forgetting: when \(N_\text{forget}\) hidden states need to be dropped, the quota is distributed proportionally across all stored memory fragments.

Inspired by the ideas proposed in Selective Context [20] and other techniques of prompt compression [21], [22], which identify and prune redundant tokens, \(f_\text{forget}(\cdot)\) prunes hidden states within each stored memory fragment based on the informativeness of the corresponding tokens. Detailed descriptions of this forgetting strategy is shown in Appendix 10.

5 Experiments↩︎

We adopt MemoryLLM [7] and M+ [8], latent-memory methods that explicitly account for forgetting and knowledge retention, as baselines for our MemDefrag. Following their evaluation protocols, we assess MemDefrag across knowledge retention (Section 5.1) and long-context QA benchmarks (Section 5.2). We also provide analyses on generality and compatibility on various aspects (Section 5.3).

Memoryllm-8b⁷ and M+-8b⁸ serve as baseline models. Although MemDefrag is model-agnostic and applicable to backbone architectures, we conduct all comparative experiments (in Section 5.1 and 5.2) on Llama-3.1-8B-Instruct for consistency. All three models are built upon Llama-3.1-8B [23], ensuring a fair comparison. The maximum length \(N_{\text{max}}\) of MemDefrag is set as 12,800, which is equal to the memory token numbers of Memoryllm-8b and M+-8b. The layer 13 serves as the tracer layer of MemDefrag because of its highest tracing capability in Llama-3.1-8B-Instruct discussed in Section 3 and Appendix 8.2. All experiments are conducted on \(4 \times\) GPUs that feature 141 GB of HBM3e memory.

5.1 Knowledge Retention Experiments↩︎

To evaluate the ability of MemDefrag to recall long-term knowledge, we follow the setup in MemoryLLM and M+ on datasets NaturalQA [13] and SQuAD [24], where each item is a (context, question, answer) triple. The implementation details follow Section 2.2. We calculate last-token attention for NaturalQA and all-token attention for SQuAD as discussed in Appendix 8.2. To eliminate positional bias of evaluation, i.e., to prevent the target fragment from consistently occupying the first position, the stored memory fragments are shuffled after each memory update. For each dataset, 500 groups are sampled for evaluation.⁹ Besides, we report the QA accuracy obtained without storing any memory fragments (NaturalQA: 14.6%; SQuAD: 7.4%), referred to as the borderline accuracy.

5.1.1 Ablation Study of Different \(K\)↩︎

We first study the effects of different \(K\), which is used to keep Top-\(K\) memory fragments for inference-time memory Defragmentation (Section 4.1). The maximum time step is 20. As the maximum memory length \(N_{20}\) is less than \(N_{\text{max}} = 12,800\), which does not trigger any forgetting strategy. Figure 5 shows the results of knowledge retention, which reveals several key findings:

(1) Vanilla latent memory suffers from severe degradation as the number of time steps increases: accuracy drops from 73.2% to below 9.4% on NaturalQA and from 77.6% to 19.0% on SQuAD. On NaturalQA, vanilla accuracy eventually falls below the borderline, indicating that the accumulated positional distortion can compromise even the standard inference capacity of the language model.

(2) All MemDefrag variants substantially slow down the degradation curve compared to the vanilla. Even Top-all, which only reorders fragments without filtering, maintains higher accuracy throughout the 20 steps (e.g., 22.6% vs.% at step 20 on NaturalQA), demonstrating that solely placing the target memory fragment at a position closer to the prompt contributes to performance improvement.

(3) Top-2 achieves the highest accuracy on NaturalQA (67.2%) at step 20, while Top-3 achieves the highest on SQuAD (72.2%) at step 20. Top-6 (as well as Top-all) achieves lower accuracy in both datasets, revealing a trade-off: a larger \(K\) improves the chance of tracing the target memory fragment but reduces its attention density. Attention density is highly correlated with QA accuracy (discussed in Section 2.3.2 and Figure 2 (d)).

5.1.2 Ablation Study of Forgetting Strategies↩︎

We study the differences of forgetting strategies: informativeness-based (Section 4.2.2) vs. random proportional forgetting. The detailed results to validate the informativeness-based forgetting are shown in Appendix 11.1.

5.1.3 Comparison with MemoryLLM and M+↩︎

In MemoryLLM and M+, each 512-token chunk is compressed into 256 memory vectors per layer in a fixed 2:1 ratio.¹⁰ Figure 6 presents the results comparing MemoryLLM, M+ and, our MemDefrag from Top-1 to Top-4 memory defragmentation with informativeness-based forgetting. Several findings stand out:

(1) All MemDefrag variants substantially and consistently outperform both MemoryLLM and M+ across the 50 steps on both datasets. Table 2 summarizes the accuracy at \(\{10,20,30,40,50\}\) time steps. By step 50, Top-1 retains an accuracy of 43.0% on NaturalQA vs. 17.4% for MemoryLLM and 17.6% for M+.

(2) The optimum \(K\) remains small and does not increase with the number of steps. This suggests that concentrating attention density on the traced target fragment (achieved by filtering fewer candidates) can be more beneficial than broadening the candidate pool to increase the likelihood of including the target.

(3) Even at the earliest time steps, MemDefrag substantially outperforms MemoryLLM and M+, despite all three methods being built upon the same Llama-3.1-8B backbone. This indicates that: The fixed compression of MemoryLLM and M+ incurs a non-trivial information loss during memory formation, leading to a significant accuracy reduction before any long-term degradation occurs.

Further, Appendix 11.3 analyzes the running-time efficiency of conducting an experimental group which includes 50 memory fragment formation along with per-step evaluation.

5.2 Long Context Experiments↩︎

In this section, we evaluate the long-context modeling capabilities of MemDefrag. Following the experiment settings of MemoryLLM and M+, six datasets of the LongBench [25] are adopted, where the maximum length of each item is 16,384 tokens. Four metrics are used for evaluation: (1) String match (the rule referring to Section 2.2, (2) F1 score (3) ROUGE-L [26], and (4) BERTScore [27]. Detailed settings are shown in Appendix 11.2.

Table 3 reports the results, of which Top-2 memory defragmentation for MemDefrag is adopted, and the attention computation is based on all prompt tokens (Appendix 8.2). MemDefrag consistently outperforms MemoryLLM and M+ in the majority of datasets and metrics, achieving the best results in 20 out of 24 dataset-metric combinations. It suggests that MemDefrag effectively preserves long-context information and generalizes well across diverse question-answering scenarios.

5.3 Generality and Compatibility↩︎

5.3.0.1 Generality across Models.

Beyond Llama-3.1-8B-Instruct, we study the effects of MemDefrag in another three language models: Qwen2.5-7B-Instruct [28], Mistral-7B-Instruct-v0.3 [29], Gemma-2-9b-it [30] (see Appendix 11.4).

5.3.0.2 Compatibility with Latent-Memory Variants.

We study the compatibility of the inference-time memory defragmentation (Section 4.1) with MemoryLLM and M+ (see Appendix 11.5).

5.3.0.3 Compatibility with Prompt Compression.

MemDefrag supports flexible-length memory fragments, making it compatible with prompt compression. We study the effects of compressing knowledge fragments (see Appendix 11.6).

6 Conclusion↩︎

We identified two long-standing obstacles that cause degradation under repeated updates in the long-term latent memory paradigm: positional-encoding misalignment and the absence of a tracing mechanism. Through a layer-wise probing of attention density, we showed that a small set of middle transformer layers consistently concentrates the highest density on the target fragment, exposing a reliable tracing signal. Building on this signal, we proposed MemDefrag, which (1) defragments the stored memory by ranking, reordering, and filtering fragments at inference time, and (2) operates an informativeness-guided proportional forgetting strategy. The experiments show the superiority of MemDefrag compared to MemoryLLM and M+. As the tracing signal and the defragmentation procedure are decoupled from any specific backbone, they are transferable to other models and latent-memory variants. More broadly, according to the analyses in Appendix 9.2, defragmentation based on attention density is also potentially applicable to plain long-context inference.

Limitations↩︎

Top-\(K\) filtering is chosen empirically and held static throughout inference, rather than being predicted per prompt and per step. A natural remedy would be to make \(K\) adaptive, for instance by selecting it from the shape of the per-fragment attention-density distribution at the tracer layer. In practice, however, we find that the benefit for such an adaptive scheme is narrow. Across the LLMs and the QA benchmarks we study, the best accuracy is almost always concentrated at Top-1 or Top-2, which leaves very little room for an adaptive \(K\) to improve on a well-chosen static K. The one exception is Gemma-2-9b, where Top-1 deteriorates beyond roughly 20 update steps while Top-2 remains stable (Figure 11 (c)). This suggests that an adaptive \(K\) is primarily useful for coping with long-term drift on particular (model, dataset) pairs, rather than as a universal replacement for the static choice. In this paper, therefore, we focus on a fixed small \(K\), treating adaptive \(K\) as an orthogonal design choice rather than a core component of MemDefrag.

7 Related Work↩︎

LLM memory serves as a cornerstone for transforming static LLMs into adaptive tools, particularly when their parameters cannot be promptly updated in response to interactive environment changes [31]. Moreover, LLM memory is widely recognized as one of the most important capabilities (alongside reasoning, planning, perception, and tool use) that supports LLM-based agents [32]–[34].

7.1 LLM Memory Forms↩︎

We categorize LLM memory into three types: (1) token-level memory, (2) parametric memory, and (3) latent memory.

Token-level memory is organized as explicit and discrete text tokens that can be persistently stored in a structured form. Substantial methods in this category typically focus on storing, summarizing, managing, and preventing the forgetting of dialog contents. For example, MemGPT [1] proposes an operating-system metaphor with hierarchical management, and Mem0 [2] establishes standardized operations for storing and retrieving information from ongoing conversations. Although token-level memory offers great modularity and interpretability, these approaches suffer from the lost-in-the-middle phenomenon [3] and context window bottleneck [8].

Parametric memory stores information directly within the model’s parameters. One major type is memory encoded in the model’s original parameters. For example, LMLM [4] embeds retrievable knowledge during the pre-training phase, while CharacterGLM [35] fine-tunes the LLM to exhibit distinct character traits. The other type resides in auxiliary parameter sets, such as adapters and LoRA modules [5], which are attached alongside the base model. The primary advantage of parametric memory is that it incurs no additional inference-time overhead. However, it suffers from costly retraining requirements and the risk of catastrophic forgetting [6].

Latent memory refers to information carried implicitly in the model’s internal representations, such as key-value caches, activations, and hidden states. Most existing methods in this category operate as short-term (working) memory, primarily reusing the key-value (KV) cache as a form of latent storage [36], [37]. The most closely related to MemDefrag proposed in this paper is long-term latent memory, as exemplified by MemoryLLM [7], M+ [8], and Memory^3 [38]. They avoid repeatedly processing the full context and introduce substantially less inference latency. However, transferring information into latent representations may incur information loss or semantic bias, and training a dedicated module to generate such representations introduces additional computational overhead.

7.2 Long-Term Latent Memory↩︎

Our proposed MemDefrag falls into the category of long-term latent memory, which generally supports persistent knowledge retention, efficient retrieval across extended interactions, and graceful integration of new information without catastrophic forgetting. Our method stores information in per-layer hidden states, especially aligning with the paradigm of MemoryLLM [7], M+ [8], and NextMem [9].

MemoryLLM incorporates a substantial portion of self-updatable prefix hidden states (referred to as a memory pool) enabling the model to integrate new knowledge both effectively and efficiently. M+ extends MemoryLLM by introducing a co-trained retriever that provides scalable long-term memory; specifically, the retriever extracts and reuses evicted tokens during inference.

Beyond their strategies for updating the memory pool, both MemoryLLM and M+ rely on a series of continual training strategies for knowledge incorporation and forgetting mitigation. The computational overhead of such continual training has limited their publicly released models to relatively small scales (based on Llama-2-7B or Llama-3.1-8B). Moreover, their paradigm of concatenating hidden states introduces substantial positional encoding misalignment, as the appended states no longer correspond to their original positions.

8 Supplementary Descriptions for Attention Density↩︎

8.1 Derivation of Attention Density↩︎

This appendix provides the full derivation of the head-averaged attention matrix \(\bar{A}^l\) used in Equation 2 .

8.1.0.1 Notation.

Let \(W_{Q,h}^{l} \in \mathbb{R}^{d \times d_k}\) and \(W_{K,h}^{l} \in \mathbb{R}^{d \times d_k}\) be the query and key projection matrices of layer \(l\) for the \(h\)-th attention head, where \(H\) is the total number of heads and \(d_k = d / H\).

8.1.0.2 Memory as a per-layer prefix.

The memory is stored as a per-layer prefix: at every transformer layer \(l\), the stored per-layer hidden states \[H_\theta^{l} = [\texttt{bos}^{l};\, m_1^{l};\, m_2^{l};\, \dots;\, m_n^{l}] \in \mathbb{R}^{(1+N_{n}) \times d}\] are prepended to whatever prompt-side states the layer is currently processing. The prompt hidden states are therefore produced inside this memory-conditioned forward pass: denoting the prompt-side output of layer \(l\) as \(H_p^{l} \in \mathbb{R}^{|p| \times d}\), the full sequence seen by layer \(l\) is \[H^{l} = [H_\theta^{l};\, H_p^{l}] \in \mathbb{R}^{S \times d}, \qquad S = 1 + N_n + |p|,\] which is exactly the input used in Equation 1 . Note that \(H_p^{l}\) depends on \(H_\theta^{l}\) through the layer’s self-attention, so the memory prefix shapes the prompt representation at every layer.

8.1.0.3 Per-head attention.

For each head \(h \in \{1, \dots, H\}\), we only need the attention from the prompt tokens to the full sequence. We therefore project queries from the prompt hidden states \(H_p^{l}\) and keys from the full sequence \(H^{l}\), and apply scaled dot-product attention under a causal mask \(M_{\mathrm{causal}} \in \mathbb{R}^{|p| \times S}\) that prevents each prompt position from attending to later prompt positions: \[\begin{align} Q_h^l &= H_p^{l}\, W_{Q,h}^l \in \mathbb{R}^{|p| \times d_k}, \\[4pt] K_h^l &= H^{l}\, W_{K,h}^l \in \mathbb{R}^{S \times d_k}, \\[4pt] A_h^l &= \mathrm{softmax}\!\left( \frac{Q_h^l\, K_h^{l\top}}{\sqrt{d_k}} + M_{\mathrm{causal}} \right) \in \mathbb{R}^{|p| \times S}. \end{align}\]

Equivalently, \(A_h^l\) is the last \(|p|\) rows of the full \(S \times S\) self-attention matrix, which is all that is needed for computing \(\rho_i^l\).

8.1.0.4 Head averaging and last-position extraction.

The per-head attention matrices are averaged across heads, \[\bar{A}^l = \frac{1}{H} \sum_{h=1}^{H} A_h^l \in \mathbb{R}^{|p| \times S},\] and the last-position attention score used in Equation 2 is the final row of \(\bar{A}^l\) (i.e., the attention emitted by the prompt’s last token): \[\label{eq:last-position-attention} \bar{a}_t^l = \bar{A}_{|p|,\,t}^l, \qquad t = 1, \dots, S.\tag{3}\]

8.2 Variants of Attention Computation↩︎

Equation 3 extracts the attention at the last position. This extends to the mean attention across all prompt tokens: \[\label{eq:all-position32attention} \bar{a}_t^{\,l} = \frac{1}{|p|}\sum_{\mathrm{pos}=1}^{|p|} \bar{A}_{\mathrm{pos},\,t}^{\,l}, \quad t = 1,\dots,S.\tag{4}\]

Correspondingly, we extend the investigation of attention density as a tracing signal (Section 3 and Table 1) (originally based on NaturalQA with last-token attention) to additional configurations: NaturalQA with all-token attention, and SQuAD with both last-token and all-token attention. The results are reported in Table 4. Layer 13 consistently exhibits the highest tracing capacity across all configurations. In this experiment set, last-token attention is better suited to NaturalQA, while all-token attention is better suited to SQuAD.

9 Results and Analyses in Settings without Positional Distortion↩︎

9.1 Analyses of Attention Density without Positional Distortion↩︎

This section presents a mirrored version of the experiments in Section 2.2 and Section 2.3, where the latent-memory paradigm is removed and all updates are performed over text tokens. As a result, positional encodings are correctly assigned throughout the input sequence. Figure 7 reports the corresponding results on Llama-3.1-8B-Instruct and NaturalQA. Comparing Figure 7 with the positional-distortion results in Figure 2, we make two observations.

(1) In settings without positional distortion, accuracy remains more stable as the number of update steps increases, even though layer 10, which exhibits the highest correlation with accuracy, shows a substantial drop.

(2) As shown in Figure 7 (d), the layer-wise Pearson and Spearman correlation coefficients between attention density and QA accuracy are generally lower in settings without positional distortion.

These observations suggest that, when positional distortion is absent, attention density is no longer as significant an indicator of performance as it is when positional distortion is present.

9.2 Analyses of Tracing Capacities without Positional Distortion↩︎

This section presents a mirrored version of the experiments in Section 3, where the latent-memory paradigm is removed and the concatenated input is entirely text tokens. Thus, positional encodings are correctly assigned throughout the input sequence. Table 8 reports the average attention density rank with Top-\(K\) accuracy of the target fragment across all positions for all layers. We make the following observations:

(1) Layer 13 exhibits the highest tracing capability, which is consistent with the results observed under positional distortion in Table 1. This result is also implied by the high attention density of layer 13 of the target fragment shown in Figure 7 (b).

(2) Overall, layer-wise tracing capabilities are substantially higher in settings without positional distortion than in positional-distortion settings. This suggests that positional distortion can compromise the ability of attention density to trace relevant information across layers, even though this tracing capability remains effective under the latent-memory paradigm, as discussed in Section 3.

(3) As shown by the layer distribution in Table 8, intermediate layers exhibit much higher tracing capabilities than earlier or final layers. This observation is consistent with the explanation proposed in Section 3: Intermediate layers can encode richer representations, which benefits feature relevance and extraction [18].

(4) An interesting finding stands out: Tracing capabilities can vary substantially between neighboring layers (especially the gap observed between layer 12 and layer 13). This sharp transition and the general varying tendency suggest that tracing capabilities do not evolve smoothly across depth.

10 Descriptions of Informativeness-Based Proportional Forgetting↩︎

Let \(L_i \;(i = 1,\dots,n-1)\) denote the current number of hidden states retained for knowledge \(x_i\) (initially \(L_i = |x_i|\); this quantity decreases monotonically as forgetting is applied). For each retained position \(t\) in the current memory, we measure the self-information of the corresponding token as: \[I(x_i^{(t)}) = -\log p_\phi(x_i^{(t)} \mid x_i^{(<t)}),\] where \(p_\phi(\cdot)\) denotes the token probability estimated by the language model. Tokens with lower self-information are deemed more redundant. Given a forgetting quota of \(N_{\text{forget},i}\) for knowledge \(x_i\), we select the \(N_{\text{forget},i}\) positions with the lowest self-information for removal: \[D_i = \operatorname{argtopk}_{t \in {1,\dots,L_i}} \big(-I(x_i^{(t)}), N_{\text{forget},i}\big),\] where \(D_i\) is the index set of positions to be dropped, and \(\operatorname{argtopk}(\cdot,k)\) returns the indices of the \(k\) largest values. The pruned memory at layer \(l\) is then obtained by removing these positions: \[m_i^{l} \leftarrow m_i^l\big[{1, \dots, L_i} \setminus {D}_i; :\big] \in \mathbb{R}^{(L_i - N_{\text{forget},i}) \times d},\] and the retained length is updated accordingly: \[L_i \leftarrow L_i - N_{\text{forget},i}.\] This ensures that the most informative hidden states are preferentially retained under a fixed capacity budget, while the retained length of each knowledge fragment shrinks over successive evolution steps. Algorithm 9 details the process of proportional forgetting.

11 Supplementary Analyses and Descriptions for Experiments↩︎

11.1 Supplementary Results of Ablation Study of Forgetting Strategies↩︎

The maximum time step is 50, which allows the memory length \(N_i\) (when \(i \geq 26\)) to reach its maximum length \(N_{\text{max}}\) and triggers the forgetting strategy after time step 26. Both Top-2 and Top-3 defragmentation settings are evaluated on NaturalQA and SQuAD. Figure 10 presents the results, from which we find that: The advantage of informativeness-based forgetting is validated and their gap increases over time. Specifically, on NaturalQA with Top-2, informativeness-based forgetting sustains an accuracy of 37.8% at step 50, compared to 20.8% for random forgetting, indicating an improvement of 81.73%.

11.2 Setups for Long Context Experiments↩︎

As MemoryLLM and M+ adopt a fixed (512 to 256) compression strategy, for each long context, it is split into 512-token chunks for memory injection. For MemDefrag which can take advantage of flexible chunking size, for the six datasets, there are two chunking strategies:

(1) Passage-based chunking: For multi-doc QA datasets where the context is explicitly structured into labeled passages (2WikiMultihopQA, HotpotQA, and MuSiQue), the context is split at each delimiter of “Passage N”, treating each passage as a natural chunk. This preserves the inherent document boundaries provided by the dataset.

(2) Token-based chunking: For single-doc QA datasets where the context is a continuous document without explicit passage markers (MultiFieldQA-en, NarrativeQA, and Qasper), the context is split into chunks of at least 256 tokens each, breaking at sentence boundaries (periods or newlines). If the last chunk is shorter than 256 tokens, it is merged into the preceding chunk to avoid overly short trailing segments.

11.3 Efficiency Analyses↩︎

For the two datasets, NaturalQA and SQuAD, respectively, we list the efficiency of various methods in Table 9, including vanilla latent memory, MemoryLLM, M+, and several Top-\(K\) variants of our MemDefrag. From this table, we observe that:

(1) The running time of MemDefrag generally increases as \(K\) becomes larger. This is expected, since a larger \(K\) requires the model to keep more candidate memory fragments at each step, thereby introducing additional computational overhead for inference.

(2) When \(K \in \{1,2,3,4\}\), MemDefrag not only outperforms MemoryLLM and M+ in terms of accuracy, as shown in Figure 6, but also achieves higher efficiency than M+. Thus, this range of \(K\) provides a favorable trade-off between effectiveness and efficiency.

11.4 Generality across Models↩︎

In addition to Llama-3.1-8B-Instruct, we study the effects of MemDefrag in Qwen2.5-7B-Instruct [28], Mistral-7B-Instruct-v0.3 [29], Gemma-2-9b-it [30]. Table 10 reports the Top-ranked tracer layer for each model on NaturalQA and SQuAD under last-token or all-token attention computation strategy. Two consistent patterns emerge:

(1) The Top-ranked tracer layer is consistently the middle transformer layers, concentrating within layers 13–15 for the four models.

(2) At these top layers, Top-2 tracing accuracy generally exceeds 90% across both datasets (reaching up to 99.3% on Gemma-2-9b-it and SQuAD). This shows that the tracing signal is not an artifact of a single model, but a general phenomenon across various instruction-tuned LLMs.

Detailed results of the tracing capabilities of the four models are shown in Tables 4, 5, 6, and 7.

Further, Figure 11 presents the results of 50-step knowledge retention based on the models: Qwen2.5-7B-Instruct, Mistral-7B-Instruct-v0.3, and Gemma-2-9b-it. The attention computation strategy for each (model, dataset) pair follows the configurations reported shown in Table 10. We make the following observations.

(1) The performance achieved by MemDefrag with Top-1 to Top-4 filtering substantially outperforms that of vanilla latent memory. This suggests that MemDefrag is broadly effective and can be generalized to various language models.

(2) Similar to experiment results (Figure 6 and Table 2) of knowledge retention based on Llama-3.1-8B-Instruct, Top-1 or Top-2 filtering generally achieves the best performance across all time steps on Qwen2.5-7B-Instruct and Mistral-7B-Instruct-v0.3.

(3) However, for Gemma-2-9b-it, the performance of Top-1 filtering drops significantly once the time step exceeds 20 and Top-1 cannot serve as the optimum choice any longer. These results indicate that, in some cases, as the time step increases, adaptive Top-\(K\) may have the potential to serve as an effective strategy for maintaining optimal performance across different time steps.

11.4.0.1 Plug-and-play justification of the tracer layer.

A natural concern is whether selecting the tracer layer requires running an investigation for every new backbone. We argue that this is not the case. As summarized in Tables 4, 5, 6, and 7, across all four instruction-tuned LLMs we examined—Llama-3.1-8B-Instruct (layer 13/32), Qwen2.5-7B-Instruct (layer 14/28), Mistral-7B-Instruct-v0.3 (layer 15/32), and Gemma-2-9b-it (layer 15/42)—the top-ranked tracer layer consistently falls within the band \([L/3,\, L/2]\), where \(L\) denotes the total number of transformer layers. This band is narrow in absolute terms (5–7 candidate layers per model in our setup). In practice, this means that for a new backbone a user does not need to repeat the full investigation study of exploring the best tracing layer: A lightweight sweep over the \(\lfloor L/3 \rfloor\)–\(\lceil L/2 \rceil\) band is sufficient to identify a high-quality tracer layer. This keeps MemDefrag genuinely plug-and-play in deployment.

11.5 Compatibility with Latent-Memory Variants↩︎

We study the compatibility of our inference-time memory defragmentation (Section 4.1) with the other two latent-memory variants: MemoryLLM (Memoryllm-8b) and M+ (M+-8b). Considering that these two models have been with a fixed-length prefix latent memory, the defragmentation procedure in this section adopts only the ranking and reordering strategies, without Top-\(K\) filtering. The experiment setup follows that of the knowledge retention experiments in Section 5.1.

Figure 12 and Table 11 show the 50-step accuracy results under the presence or absence of defragmentation for MemoryLLM and M+. Last-token attention computation strategy based on layer 13 (the tracing layer) is used in defragmentation for both MemoryLLM and M+ and both datasets. Two observations stand out.

(1) Defragmentation consistently improves both backbones. Across nearly all checkpoints on both NaturalQA and SQuAD, applying our defragmentation procedure yields a clear and stable accuracy gain over the corresponding vanilla baseline. This indicates that explicitly reorganizing memory fragments makes the stored knowledge substantially more retrievable. Further, as MemoryLLM and M+ compress knowledge fragments by their own models, they support the compatibility of our MemDefrag with model-level compression.

(2) Random forgetting still drives a steady performance decay. Despite these gains, both MemoryLLM and M+ rely on a random forgetting strategy. As a result, even after defragmentation, accuracy continues to drop monotonically as \(n\) grows on both datasets.

11.6 Compatibility with Prompt Compression↩︎

Unlike MemoryLLM or M+ that construct memory fragments whose lengths are 256 in a fixed manner, our MemDefrag can take advantage of approaches of prompt compression to save the memory storage. In this section, we study the effects of LLMLingua-2 [22] on MemDefrag. The experiment setup follows that of knowledge retention experiments (Section 5.1). Specifically, the filtering strategy is Top-3 and the adopted model is Llama-3.1-8B-Instruct. Each knowledge fragment is compressed by LLMLingua-2 with a compression ratio before being formed into its memory fragment.

Figure 13 and Table 12 show the 50-step accuracy results under different compression ratios (from 30% to 100%, i.e., no compression). The results show that MemDefrag effectively leverages prompt compression to improve long-term knowledge retention. Compared with the no-compression setting, moderate compression ratios maintain competitive accuracies in early steps while achieving better performance as the step number increases. This improvement can be attributed to the shorter length of the compressed knowledge and the ability of the adopted compression technique to preserve informative tokens. The memory grows more slowly over time, which delays the triggering of the forgetting mechanism. This explains why compressed settings become particularly advantageous at later time steps, especially when \(n \geq 40\), where uncompressed knowledge (and memory) fragments have suffered from earlier forgetting and performance degradation.

References↩︎

---

1. This work was done during internship at LIGHTSPEED.↩︎

2. Corresponding author.↩︎

3. https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct.↩︎

4. We use a filtered subset of the dataset to ensure low data leakage and mutual independence among contexts: https://huggingface.co/datasets/YuWangX/KnowledgeRetentionProcessed.↩︎

5. We also consider an all-token variant that averages \(\bar{a}_t^l\) over all prompt positions rather than only the last one; details are in Appendix 8.2.↩︎

6. The two coefficients are complementary: Pearson captures linear co-variation, while Spearman, being rank-based, is robust to the heavy-tailed distribution of attention density and reflects monotonic (but possibly non-linear) dependence.↩︎

7. https://huggingface.co/YuWangX/memoryllm-8b↩︎

8. https://huggingface.co/YuWangX/mplus-8b↩︎

9. Note that the sampled groups used here are disjoint from those in Section 2, ensuring that hyperparameter selection is not biased by prior observations on the same data.↩︎

10. To ensure a fair comparison that accounts for this compression, each knowledge fragment \(x_i\) injected into MemoryLLM and M+ is 512 tokens.↩︎