Peripheral Memory for LLMs: Integration of Sequential Memory Banks with Adaptive Querying

1. Question about baselines

Thank you for raising this important concern. We provide detailed clarifications regarding the baseline implementation:

(1) Baseline Implementation

All baselines for Knowledge-based Model Editing (KME) were implemented using the widely adopted toolkit EasyEdit (https://github.com/zjunlp/EasyEdit), ensuring reproducibility, reliability and alignment with established practices. For instance, in the case of WISE, we strictly followed the official implementation guidelines, with our editing code closely mirroring the provided examples:

from easyeditor import WISEHyperParams
hparams = WISEHyperParams.from_hparams('./hparams/WISE/llama-3-8b.yaml')
editor = BaseEditor.from_hparams(hparams)
metrics, edited_model, _ = editor.edit(
        prompts=prompts,
        rephrase_prompts=rephrase_prompts,
        target_new=target_new,
        subject=subject,
        sequential_edit = True,
        locality_inputs=locality_inputs,
)

The hyperparameters are suggested by the EasyEdit.

(2) Low Baseline Performance

We observed that several baselines (e.g., PMET, MEMIT) exhibit zero performance in consecutive editing scenario. Initially surprised by these results, we rigorously repeated the experiments 10+ times and consistently observed the same outcome.

Key Insight: When the number of sequential edits exceeds 1K, these localization-based editing methods suffer from catastrophic performance degradation, mainly arising from the repeated modifications progressively destabilize the LLM’s parameters. This phenomenon aligns with observations in recent works [1,2], which highlights the fragility of localized edits under high edit counts.

(3) Reproducibility of our method

Full implementation details and hyperparameters are provided in the supplementary material. For instance, refer to memory.py (Lines 1073–1075) for key parameters such as memory depth and grid size. Additionally, they have also been discussed in Section 4.3 for clear description. We welcome further discussion or code review to address any remaining concerns.

2.Questions about Llama3.1-8B as backbone

Thank you for your insightful feedback. We address your concerns as follows:

(1) Token Limit of Llama3-8B

The token limitation of Llama3-8B (8k tokens) does not impact our experiments, as the input sequences in KME datasets (ZsRE and CounterFact) are well within this constraint: ZsRE: Maximum input length = 36 tokens and CounterFact: Maximum input length = 56 tokens. Thus, even for 3,000 sequential edits, the total token count remains far below the 8k limit, impossibly impacting model performance. Thus, the observed performance limitations of baseline methods (e.g., MEMIT, PMET) are not attributable to token constraints but to inherent challenges in parameter-localized editing [2].

(2) Performance under longer token limit

We acknowledge that token limits could theoretically hinder prompt-based methods. For comprehensive comparison, we conducted additional experiments on Llama3.1-8B (128k token limit) and compared our method with IKE, a state-of-the-art prompt-based editing method that uses in-context learning without parameter updates.

Our method achieves superior performance even with extended context lengths, confirming that token limitations are not the bottleneck in our setup. This also confirms that our framework’s advantages are architecture-agnostic and not contingent on context length.

3. Explanation of Figure 1 and Figure 2 and The function $g$ and $h$

Figure 1 provides an overview of our framework, while Figure 2 illustrates the peripheral memory.

The functions $g_i^k$ and $h_{i1}^k$ are learnable univariate functions, parameterized as B-spline curves. For a detailed understanding, please refer to Chapters 5.1 and 5.2 [3] for more details.

Due to current character constraints, we are unable to provide a more detailed explanation here but will offer a comprehensive discussion in later response.

References

[1] Wang et al. 2024. WISE: Rethinking the Knowledge Memory for Lifelong Model Editing of Large Language Models. In NeurIPS.

[2] Li et al. 2024. Consecutive Batch Model Editing with HooK Layers. In EMNLP, 13817–13833. Association for Computational Linguistics.

[3] Prautzsch et al. 2002. Bézier and B-Spline Techniques. Mathematics and Visualization. Springer Science & Business Media. doi:10.1007/978-3-662-04919-8. ISBN 978-3-540-43761-1.

1. Handling conflict

Thank you for raising this critical point. Below, we clarify our method’s current behavior and outline potential enhancements:

(1) Last-Write-Wins Policy

In sequential editing scenarios, our memory defaults to a temporal priority strategy: the most recent edit overwrites previous entries for the same fact. This is reflected in the confidence mechanism, where newer entries are assigned higher trust values due to their proximity in the query distribution. While this ensures consistency in retrieval (only the latest edit is returned), it does not explicitly resolve semantic contradictions. However, the confidence bank can be easily configured into detect contradictions using off-the-shelf entailment classifiers. $\Downarrow$

(2) Conflict-Aware Confidence Integrate semantic similarity checks between new edits and existing memory entries. For example, if a new edit contradicts a stored fact (measured via entailment models, the confidence bank could trigger a conflict resolution protocol (e.g., human-in-the-loop verification or probabilistic truth maintenance). Notably, this enhancement can be easily implemented as a configurable extension to our peripheral memory without altering the core LLM.

Actually, we are now actively studying a version-controlled peripheral memory inspired by database systems, where conflicting edits are stored as alternate branches. This could enable users to query the memory with temporal constraints (e.g., “What was believed about xxx in 2023?”).

2. Scale to even more edits

Our framework currently employs a direct memory querying strategy, which achieves reliable storage accuracy for up to 10K edits (see Figure 7). However, as noted in Section 5.4, while the memory retains high storage fidelity at scale, its generalization degrades for semantically equivalent queries due to geometric misalignment in hidden feature spaces (see Appendix C). To balance robustness and scalability, we adopt a memory archival protocol: when the active memory reaches 1K entries (optimal for generalization), it can be archived, and a fresh memory will be initialized. This approach enables theoretically unbounded storage capacity while maintaining generalization performance thanks to the unique properties of our peripheral memory component. Notably, even with a 1K entry limit, our memory achieves high effective storage density (4.95, see Table 4).

While this archival strategy significantly improves scalability, it trades off per-bank memory utilization. To address this, we are developing a memory management component (MMC), inspired by the Memory Management Unit (MMU) in operating systems. The module maps semantically equivalent inputs to unified memory space using contrastive learning, decoupling semantic alignment from storage operations. This allows the peripheral memory to specialize in efficient storage/retrieval, while the MMC handles query normalization and address translation.

3. Importance of KAN

Response #3 of Reviewer gZ6v provides detailed experimental results. These indicate that the KAN-based memory bank surpasses the MLP-based counterpart within our peripheral memory.

(1) Accuracy vs. Generalization Trade-off

KAN demonstrates exceptional performance in both accuracy and generalization, while MLPs (even at 4$\times$ params) plateau at 89.7% accuracy and 57.8% generalization. MLPs’ noisy memory outputs degrade Locality, as seen in CounterFact: KANs retain 100% Locality vs. MLPs’ 2.3%-3.3%.

(2) Scaling MLPs Fails to Close the Gap

Increasing MLP hidden dimensions marginally improves storage accuracy (e.g., 89.7% at 4× params vs. 86.7% at 1× on ZsRE), but harms generalization due to overfitting on memorizing data. This reflects the inability of MLPs to learn smooth mappings.

(3) Catastrophic Collapse at Extreme Scaling

At 8$\times$ parameters, MLP performance collapses to 0 across all metrics. We attribute this to:

• Optimization Instability: Overparameterized MLPs suffer from vanishing/exploding gradients, exacerbated by the memory’s sequential architecture (sequence length = 512).
• Loss Landscape Degradation: High-dimensional MLP weights create chaotic loss surfaces, preventing convergence. In contrast, KANs’ spline parameterization inherently regularizes the optimization landscape.

(3). Why MLPs Underperform

• Approximation Theoretic Limitations: MLPs struggle to model the compositional structure of memory mappings (query→key→value), which KANs explicitly encode via Kolmogorov-Arnold superposition[1].
• Noise Amplification: MLPs’ fixed activations amplify high-frequency noise in memory queries, degrading generalization. KANs’ adaptive splines act as low-pass filters, suppressing noise[2].

References

[1] Liu et al. (2024). Kan: Kolmogorov-arnold networks.In ICLR.

[2] Prautzsch et al. (2002). Bézier and B-Spline Techniques. Mathematics and Visualization. Springer Science & Business Media. ISBN 978-3-540-43761-1.

1.Question about the semantic drift

Thank you for this insightful question. We acknowledge that semantic drift and retrieval degradation at extreme scales (>10K updates) for semantically equivalent queries remain challenges. Below, we summarize our empirical observations and outline explicit mitigation strategies under development:

As discussed in Appendix-C, our memory is directly queried using the hidden state of the final input token. While efficient, this design inherits a challenge in representation learning: geometric misalignment in high-dimensional spaces. That is, small differences in inputs can lead to disproportionate shifts in representations. Thus, semantically equivalent queries may map to distinct regions of the memory space due to minor differences in token-level representations. This reduces the robustness of rephrased inputs, as shown in Table 2. Additionally, as memory utilization scales, the module becomes overly specialized to the original query distribution (stability), sacrificing generalization to representationally divergent queries (plasticity).

To address the limitations of the direct memory querying strategy, we are now actively developing a memory management module inspired by Memory Management Unit (MMU) in operating systems. The module acts as an abstraction layer between the LLM and peripheral memory, decoupling semantic alignment from storage operations. This allows the peripheral memory to specialize in efficient storage/retrieval, while the MMU handles query normalization, introducing three key innovations:

(1) Semantic-Aware Querying

The MMU parses and refines raw query signals (e.g., token-level hidden states) into semantically enriched descriptors, mitigating hypersensitivity to input variations. For example, paraphrased queries like "What is the capital of France?" and "Name France’s capital city" would be mapped to unified descriptors, enabling robust retrieval regardless of surface-form differences.

(2) Optimized Memory Operations

The MMU supports bulk memory read/write operations, reducing overhead for large-scale edits. Inspired by virtual memory paging, it dynamically groups related memory entries (e.g., knowledge about a specific entity) into contiguous blocks, improving cache utilization.

(3) Adaptive Scheduling Policies

Leveraging reinforcement learning, the MMU learns optimal policies for memory allocation and eviction. This balances hot (frequently accessed) and cold (rarely used) memory regions, addressing the stability-plasticity trade-off while minimizing fragmentation.

This MMU introduces OS-inspired abstractions (e.g., memory pages) to LLMs, enabling systematic memory control. We are currently refining this architecture and will introduce it in future work.

2.Question about sensibility of query The performance of our framework exhibits moderate sensitivity to the choice of query representation. Below, we analyze two key dimensions of this sensitivity:

• Token Position Sensitivity: Using the last token hidden state as the query signal is an intuitive choice since it aggregates semantic information from the entire input. In addition, we also do an experiment using the average of all hidden token features as query, and find similar results. This shows that the selection of token features is relatively stable. |Type||ZsRE||||||CF||| |:-----|:-----:|:-----:|:-----:|:-----:|:-----:|:-----:|-----:|:-----:|-----:|-----:| ||Efficacy|Generality|Locality|Score|||Efficacy|Generality|Locality|Score| |Last Token|0.9774|0.6432|1.0000|0.8735|||0.9915|0.3108|1.0000|0.7674| |Average|0.9698|0.6399|1.0000|0.8699|||0.9800|0.2520|1.0000|0.7440|

• Layer-Wise Variability: Deeper layers (e.g., layer 8-31 in Llama3-8B) yield more stable query representations due to their focus on high-level semantics, while shallower layers (e.g., layer <8) exhibit higher variance (See Figure 12 in Appendix D.3). This means that query performance is relatively stable as long as it is within reasonable limits.

Additionally, as noted in prior responses, the improved query mechanism (MMU) will decouple semantic alignment from query representation by aggregating hidden states across multiple layers at the same time. This could balance low-level syntactic and high-level semantic features, reducing positional bias and improving robustness.

3.Comparison with Recent RAG Thank you for emphasizing the importance of comparison with recent RAG. We have provided experimental analysis in Appendix D.1. Briefly, we evaluate our method against recent LongRAG:

LongRAG: Results with * use Document-level retrieval; others use Passage-level (1 retrieval unit).

4.​Limitation discussion:see Response #1.

5.​Additional results on diverse LLM architecture will be provided later due to current character constraints.

1.Question about $\mathbf{W}_0$ and $\mathbf{W}_1$

Thank you for raising this important concern. We clarify the training protocol and evaluation fairness as follows:

(1) Role and Training of Convertors

The mapping matrices $\mathbf{W}_0$ and $\mathbf{W}_1$ (see Section 3.2) serve solely as feature space adapters, analogous to "plug-and-play" connectors between the LLM’s hidden states and the memory module. Their sole purpose is to align dimensionality between the LLM’s hidden space and the memory’s representation space, not to encode task-specific knowledge. Specifically, the convertors were trained solely to memorize contexts into memory (See Line-986, 587, 273 of code), with no exposure to query data or ground-truth answers during training. This ensures the convertors do not learn task-specific knowledge.

(2) Ensuring No Data Memorization

To rigorously verify that convertors do not memorize evaluation data, we also implemented an empty_fn (similar to the memory_fn in Line-508 of supplementary code):

def empty_fn(self, layer_idx, hidden_states, causal_masks, **kwargs):
    if layer_idx == self.merge_layer_idx:
        empty_feats = self.convertor(self.query_signal.to(self.convertor.device))
        hidden_states[:,self.replace_idx,:] = empty_feats.to(hidden_states.device)
        causal_masks = self._replace_attn_fn(causal_masks, 1., self.replace_idx)
    if layer_idx > self.merge_layer_idx:
        causal_masks = self._replace_attn_fn(causal_masks, 1., self.replace_idx)
    return hidden_states, causal_masks

This function skips memory retrieval and aligns the LLM’s output with its original predictions. However, extensive experiments show no statistically significant difference in model performance ($\Delta$ < 0.3%) between models with and without empty_fn, confirming that convertors introduce no hidden memorization. In consideration of efficiency and simplicity, we omitted this part directly and did not provide it in the supplementary materials.

2.Question about the variable l in Section 2.1

Thank you for highlighting this important point. The variable $l$ in Section 2.1 denotes the index of the hidden layer from which query features are extracted for memory retrieval. While $l$ is indeed a critical hyper-parameter, we relegated its analysis to Appendix D.3 due to space constraints. Findings can also be found in Response#2 of Reviewer UVbc.

3. Question about using KAN

We appreciate your thoughtful critique and address the motivation for using KANs (Kolmogorov-Arnold Networks) as follows:

(1) Motivation for KANs

As formalized in Eq.2, the memory query process can be abstracted as a smooth mapping from query signals to memory data. Physical RAM approximates this mapping with an indicator function, which is ill-suited for neural memory systems requiring gradual state transitions. KAN was chosen over MLP due to its superior symbolic regression capabilities:

• Superior Function Approximation: KANs achieve higher precision with fewer parameters by leveraging spline-based nonlinearities, avoiding MLPs’ reliance on rigid activation functions (e.g., ReLU).
• Smoothness: KANs’ piecewise polynomial basis functions enable $C^2$-continuous mappings, critical for stable gradient propagation during sequential memory interactions. We will revise the title of Section 2.2 in future editions and provide further clarification in this section.

(2) Empirical Comparison with MLPs

We compared KANs against MLPs with varying parameter counts:

Note: Where the MLP is used to model a specific memory bank. Each MLP memory bank mirrors KAN’s structure:

• First&Confidence bank: torch.nn.Linear(1,9) $\to$ torch.nn.Linear(9,1)
• Other banks: torch.nn.Linear(2,9) $\to$ torch.nn.Linear(9,1)
• Hidden dimensions scaled via torch.nn.Linear(1,9$\times X$) $\to$ torch.nn.Linear(9$\times X$,1), #Epoch=10.

Results indicate that the KAN-based memory bank surpasses the MLP-based counterpart within our peripheral memory.

For a more detailed analysis, please refer to Response #3 of Reviewer TWi2 due to chararacter limitation of current response.