Dual-Key Memory Module

• Dual-Key Memory Modules are neural architectures that use two key banks to exponentially increase memory slots while keeping computational costs low.
• They employ sparse top-k product addressing and dual-memory retrieval mechanisms to efficiently manage both short-term and long-term memories.
• Applications of these modules include retrieval-augmented dialogue systems and long-context language models, enhancing performance in tasks like conversational agents and episodic memory.

A dual-key memory module is a neural memory architecture designed to enable efficient storage, addressing, and retrieval of information by leveraging a pair of key banks—each representing a distinct dimension in the memory space. This structure underlies both static and dynamic (fast-weight) variants, notably Product Key Memory (PKM) and its fast-weight extension (FwPKM), as well as hybrid architectures that combine short-term and long-term memory for retrieval-augmented LLMs. Such modules provide scalable, high-capacity context tracking suited to applications such as conversational agents, LLMs, and long-context episodic memory systems (Zhang et al., 2024, Zhao et al., 2 Jan 2026).

1. Structure of Dual-Key Memory Modules

Dual-key memory modules adopt a bifurcated key space to enable addressing exponentially more slots than single-key systems while maintaining computational tractability. In the canonical PKM, two key matrices are maintained:

• $K^{(1)}\in\mathbb{R}^{\sqrt{N}\times D_K}$
• $K^{(2)}\in\mathbb{R}^{\sqrt{N}\times D_K}$

Each memory slot is indexed uniquely by a pair $(i,j)$, where $i$ and $j$ correspond to the indices in the two key matrices, producing $N = \sqrt{N} \times \sqrt{N}$ total slots. Queries $q\in\mathbb{R}^{2D_K}$ are split into $q^1, q^2\in\mathbb{R}^{D_K}$ that independently address $K^{(1)}$ and $K^{(2)}$, forming a product key space (Zhao et al., 2 Jan 2026).

In dual-memory LLM applications, "dual-key" refers more broadly to system-level duality: short-term memory (STM) acting as a sliding window buffer, and long-term memory (LTM) as a persistent, entity-indexed store. Both are independently keyed and can be addressed conditionally at inference (Zhang et al., 2024).

2. Addressing and Retrieval Mechanisms

Sparse Top-$K^{(2)}\in\mathbb{R}^{\sqrt{N}\times D_K}$0 Product Addressing

To achieve sublinear complexity in retrieving from $K^{(2)}\in\mathbb{R}^{\sqrt{N}\times D_K}$1 slots, PKM and FwPKM implement sparse top-$K^{(2)}\in\mathbb{R}^{\sqrt{N}\times D_K}$2 retrieval as follows:

1. Independently score query sub-vectors:
   • $K^{(2)}\in\mathbb{R}^{\sqrt{N}\times D_K}$3 for $K^{(2)}\in\mathbb{R}^{\sqrt{N}\times D_K}$4
   • $K^{(2)}\in\mathbb{R}^{\sqrt{N}\times D_K}$5 for $K^{(2)}\in\mathbb{R}^{\sqrt{N}\times D_K}$6
2. Select top-$K^{(2)}\in\mathbb{R}^{\sqrt{N}\times D_K}$7 indices $K^{(2)}\in\mathbb{R}^{\sqrt{N}\times D_K}$8, $K^{(2)}\in\mathbb{R}^{\sqrt{N}\times D_K}$9 in each bank.
3. Form $(i,j)$0 Cartesian product candidates $(i,j)$1.
4. Compute combined score $(i,j)$2 or a variant.
5. From $(i,j)$3, select top-$(i,j)$4 pairs $(i,j)$5 based on $(i,j)$6.
6. Compute normalized weights $(i,j)$7.
7. Retrieve value by weighted sum: $(i,j)$8.

This approach reduces per-token cost to $(i,j)$9 (Zhao et al., 2 Jan 2026).

Attention-based Dual-memory Retrieval

In conversational agents, STM and LTM are addressed by computing projections:

• $i$0, $i$1, $i$2, $i$3 via learnable $i$4, $i$5.
• Queries $i$6 attend over both banks using cosine or dot-product similarity.

Attention weights are calculated and resulting context vectors $i$7, $i$8 are concatenated for downstream prompt assembly (Zhang et al., 2024).

3. Update Rules and Memory Management

Static vs Dynamic (Fast-weight) Updates

Static PKM maintains trained, frozen weights $i$9, $j$0 at inference, updated only during offline optimization.

Fast-weight PKM (FwPKM) enables episodic updates via local chunk-level gradient descent. For each chunk of $j$1 tokens:

• Project $j$2 to queries and targets: $j$3, $j$4
• Compute PKM readout $j$5
• Define loss $j$6
• Aggregate over the chunk: $j$7 (with optional gate $j$8)
• Update $j$9 with
  • Value update from $N = \sqrt{N} \times \sqrt{N}$0
  • Key update from entropy regularization to encourage slot diversity
• One-step “rewriting” is achieved with $N = \sqrt{N} \times \sqrt{N}$1

STM (e.g., FIFO window) and LTM (entity store) applications use windowed append and periodic summary update via LLM API calls. Optional write gates can interpolate updates (Zhang et al., 2024, Zhao et al., 2 Jan 2026).

4. Integration in Retrieval-Augmented and Conversational Systems

Dual-key and dual-memory modules are pivotal in retrieval-augmented generation (RAG) architectures, particularly for LLM-enabled dialogue agents. The pipeline typically includes:

1. Privacy module: Anonymizes user utterances $N = \sqrt{N} \times \sqrt{N}$2, mapping PII to placeholders.
2. STM: Maintains sliding window of recent dialogue turns; updated and retrieved per user/agent turn.
3. LTM: Entity store mapping anonymized entities $N = \sqrt{N} \times \sqrt{N}$3 to LLM-generated summaries $N = \sqrt{N} \times \sqrt{N}$4. Updated every $N = \sqrt{N} \times \sqrt{N}$5 turns based on recent context for mentioned entities.
4. Conditional retrieval assembles:
   • Recent STM messages
   • LTM entity summaries (if mentioned in $N = \sqrt{N} \times \sqrt{N}$6)
   • Relevant therapist knowledge base (KB) examples based on vector similarity ($N = \sqrt{N} \times \sqrt{N}$7)
5. Retrieved contexts are concatenated as the LLM input prompt. LLM response is de-anonymized using stored mappings and returned (Zhang et al., 2024).

The following table summarizes dual-key memory usage in two research lines:

Line of Work	Dual-key Encapsulation	Retrieval/Update Dynamics
Product Key Memory (PKM)	$N = \sqrt{N} \times \sqrt{N}$8	Static or fast-weight updates
Dual-memory LLM agent (STM/LTM)	STM: window, LTM: entity→summary	STM: per-turn, LTM: periodic/conditional

5. Computational Efficiency and Empirical Results

Complexity

• PKM read: $N = \sqrt{N} \times \sqrt{N}$9 per token; $q\in\mathbb{R}^{2D_K}$0 storage.
• FwPKM update: $q\in\mathbb{R}^{2D_K}$1 for value matrix per chunk; $q\in\mathbb{R}^{2D_K}$2 for key matrix; amortized cost per token further reduced by chunked updating (Zhao et al., 2 Jan 2026).

Empirical Performance

• On Fineweb-Edu: PKM yields $q\in\mathbb{R}^{2D_K}$311% perplexity reduction; combined PKM + FwPKM $q\in\mathbb{R}^{2D_K}$415% reduction versus baseline.
• On very long-context tasks (e.g., LC64, LAMBADA): FwPKM shows 27–30% perplexity reduction.
• Needle-in-Haystack (NIAH) tasks with up to 128K context: multi-iterate passes (up to $q\in\mathbb{R}^{2D_K}$5) achieve $q\in\mathbb{R}^{2D_K}$690–95% retrieval accuracy, indicating strong generalization and retrieval capacity for FwPKM even when training context is limited to 4K tokens (Zhao et al., 2 Jan 2026).

Application in Conversational LLMs

In SoulSpeak, the dual-memory module enables personalized, privacy-preserving, and contextually coherent response generation in psychotherapy domains, demonstrating practical efficacy in controlled evaluation settings (Zhang et al., 2024).

6. Comparative Significance and Implications

Dual-key memory modules address the longstanding tradeoff between memory scalability and computational cost in neural sequence models. The product key structure enables exponential memory slot growth with sublinear addressing cost, while fast-weight updates provide adaptivity and rapid memorization of new episodic information. Integration with STM/LTM paradigms supports structured, hierarchical memory for long-context reasoning and user-adaptive generation.

The results from FwPKM suggest that introducing dynamic, locally updatable memory into transformer backbones both reduces sequence model perplexity and enables retrieval-based pattern completion even at context lengths orders-of-magnitude longer than those observed during training. A plausible implication is that such modules constitute an effective neural substrate for lifelong learning, contextual adaptation, and knowledge-augmented generation.

• "Advancing Conversational Psychotherapy: Integrating Privacy, Dual-Memory, and Domain Expertise with LLMs" (Zhang et al., 2024).
• "Fast-weight Product Key Memory" (Zhao et al., 2 Jan 2026).