PaperlayerView exact revision on arXiv

Source-linked AI summary

Fast-weight Product Key Memory

Tianyu Zhao, Llion Jones

arXiv:2601.00671v2cs.CLcs.AI

TL;DR

Fixed-state sequence models limit retained information, while dense fast-weight networks are costly to update and query. FwPKM adds online sparse updates to Product Key Memory, and it generalizes from 4K-token training sequences to 128K-token contexts in Needle-in-a-Haystack evaluations.

  • Problem

    Sequence models trade memory capacity against compute: attention has quadratic sequence-length cost, fixed-state models limit retained information, and dense fast-weight MLPs are costly to update and query.

    Source: §1 · p. 2

  • Method

    FwPKM turns PKM into an episodic memory by updating sparse key-value matrices online through chunk-level local reconstruction.

    Source: §1 · p. 2, §3.3 · p. 5, §8 · p. 17

  • Results

    128K-token contexts are handled by models trained on 4K-token sequences, and iterative reading raises retrieval accuracy from < 10% to > 70% in difficult Needle-in-a-Haystack scenarios.

    Source: §8 · p. 17

  • Takeaways & Limitations

    FwPKM functions as episodic memory that selectively captures novel entities and complements semantic knowledge in standard slow weights.

    Source: §1 · p. 2, §8 · p. 17

    FwPKM components currently run slower than dense counterparts because their sparse operations rely on less specialized kernel implementations.

    Source: §6 · p. 15

Abstract

from arXiv · show

Sequence modeling layers in modern language models typically face a trade-off between storage capacity and computational efficiency. While softmax attention offers unbounded storage at prohibitive quadratic cost, linear variants are more efficient but suffer from limited, fixed-size storage. We introduce Fast-weight Product Key Memory (FwPKM), a sparse fast-weight memory layer that resolves this tension. FwPKM updates sparsely activated parameters at both training and inference time using chunk-level gradient descent on a local memory-rewrite objective. This performs Test-Time Training (TTT)-style gradient updates on activated slots in a sparse memory, enabling rapid memorization and retrieval of many new key-value associations while keeping per-token compute low and fixed. Experiments show that FwPKM functions as an effective episodic memory that complements the semantic memory of standard modules, yielding significant perplexity reductions on long-context datasets. Notably, in Needle-in-a-Haystack evaluations, FwPKM generalizes to 128K-token contexts despite being trained on only 4K-token sequences.

1. Introduction

Long-context models must retrieve local facts from far back in a sequence, but existing approaches trade memory capacity against computation: attention is costly, while fixed-state models have limited capacity and dense fast-weight networks are expensive to update. FwPKM uses sparse fast weights for context-specific episodic memory, with reported long-context retrieval and perplexity gains alongside retention challenges.

1. Introduction

  • Softmax attention stores each past token’s key-value pair and retrieves by comparing against all keys, yielding high-fidelity retrieval at quadratic sequence-length cost.Modern RNNs instead compress history into a fixed-size state for sub-quadratic retrieval, but that capacity limits retained information.Source: §1 · p. 2
  • Dense fast-weight MLPs adapt online through reconstruction loss, but their computation scales with parameter count, making large networks costly to update and query.The paper frames sparse associative memory as a way to retain many parameters while activating only a small subset per token.Source: §1 · p. 2
  • FwPKM turns sparse Product Key Memory into a dynamic episodic memory by updating key/value matrices online, while standard slow weights retain dataset-wide semantic knowledge.The paper reports lower perplexity on long-context datasets and generalization to 128K-token contexts after training on 4K-token sequences; it also reports retention challenges across continual-learning stages.Source: §1 · pp. 2–3

2. Preliminary

The preliminary section motivates sparse associative memory and describes Product Key Memory as a way to index large memories without scoring every slot. PKM factorizes keys into smaller codebooks to make sparse retrieval more scalable.

2. Preliminary

  • The preliminary review positions sparse key-value memories and PKM as the basis for the paper’s online-updated fast-weight module.The FwPKM layer is placed before a token mixer; slow weights update only during training, while fast key/value parameters update during training and inference.Source: §2 · p. 3, Figure 1 · p. 4

2.1. Sparse Top-𝑘key-value memory

  • Top-k retrieval still scores all N keys to select candidates, so its O(N) cost limits scaling to massive memories.It reduces the number of value rows accessed, but not the number of keys scored.Source: §2.1 · p. 3

2.2. Product Key Memory (PKM)

  • PKM splits each query into two sub-queries, scores sub-key sets, and searches their restricted Cartesian product before selecting top-k pairs and retrieving values.This factorization avoids scoring all N slots and can index millions of slots using thousands of scores.Source: §2.2 · p. 3

3. Fast-weight Product Key Memory

FwPKM retains PKM’s sparse retrieval while changing memory writing to online gradient updates on chunk-local objectives. Its forward path retrieves and gates memory values; its backward path reconstructs targets and balances slot use, with practical choices intended to improve performance and stability.

3. Fast-weight Product Key Memory

  • FwPKM keeps PKM retrieval but updates memory matrices online to store input-sequence key-value associations through chunk-level local reconstruction.Predictions within a chunk use fast weights from previous chunks, preserving causality.Source: §3 · p. 4, §3.1 · p. 4

3.1. Overview

  • For each chunk, slow-weight projections construct queries, values, and a gate; PKM retrieves a sparse value prediction, and fast weights are updated after the chunk.The value matrix is updated with a gated reconstruction loss, while key matrices use an auxiliary addressing loss to prevent slot collapse.Source: §3.1 · p. 4

3.2. Forward Pass: Inputs and Retrieval

  • FwPKM uses product-key sparse retrieval to produce a weighted value prediction, then interpolates it with a residual value path.The slow-weight network supplies query, value, and gate inputs.Source: §3.2 · p. 5

3.3. Backward Pass: Memorization by Learning to Reconstruct

  • Chunk-level reconstruction updates fast weights so retrieval reconstructs target values derived from the same input stream.The MSE residual acts as an explicit rewrite signal for the value rows; gradients for rows accessed multiple times are normalized by their read count as a consensus mechanism.Source: §3.3 · p. 5

3.4. Backward Pass: Prevent Slot Collapsing via Addressing Optimization

  • An auxiliary addressing objective encourages balanced slot usage across a chunk without requiring each query to use slots uniformly.The method minimizes negative entropy over marginal slot-usage distributions and updates keys using the resulting gradients.Source: §3.4 · p. 6

3.5. Practical Choices

  • Practical choices include lookahead value targets, inverse-distance key scoring, and z-score target normalization.Lookahead pairs each query with the next token’s value; inverse-distance weighting encourages key rows to become representation-space centroids.Source: §3.5 · p. 6

4. Experiments

The experiments assess memory roles through perplexity, synthetic retrieval, realistic long-context tasks, and online domain adaptation.

4.1. Training Setting

  • Experiments use approximately 112M-parameter, 12-layer models, place PKM and FwPKM at layers 2, 6, and 10, and train on 5B tokens each from LongContext64 and Fineweb-Edu at 4K sequence length.FwPKM retrieves Top-8 slots, while PKM retrieves Top-128 slots.Source: §4.1 · pp. 6–7

4.2. PPL Evaluation

  • FwPKM lowers PPL on long-context LC64 and LAMBADA, PKM brings the largest Fineweb-Edu gains, and combining them performs best across all metrics.The figure distinguishes models by memory-component type.Source: §4.2 · p. 7, Figure 2 · p. 8

4.3. Needle-in-a-Haystack Evaluation

  • For GDN and GDN+SWA, rereading raises NIAH accuracy from <10% to >50% in many cases and above 70% with more passes.The evaluation stacks results for 1–4 iterations across 4K–128K contexts.Source: §4.3 · p. 9, Figure 4 · p. 10
  • FwPKM maintains robust NIAH retrieval at 128K contexts despite training only on 4K-token sequences, while Full Attention baselines degrade rapidly.Figure 4 reports NIAH accuracy across 4K, 8K, 32K, and 128K test contexts.Source: §4.3 · p. 9, Figure 4 · p. 10

4.4. Longbench Evaluation

  • FwPKM reduces answer-token perplexity on five LongBench tasks for GDN, GDN+SWA, and GDN+FA models when it is actively used.Figure 5 compares perplexity across the five tasks and distinguishes models by memory components.Source: §4.4 · p. 10, Figure 5 · p. 11

4.5. Continual Learning Evaluation

  • After each adaptation stage, the corresponding domain's NLL decreases, but gains often fail to persist as fast-weight knowledge is replaced by newer domains.The evaluation tracks ΔNLL across six domain-adaptation stages.Source: §4.5 · pp. 10, 12, Figure 6 · p. 11

5. Interpretability Analyses

FwPKM’s explicit slot reads make retrieved information traceable to input tokens. In the NIAH case study, multiple retrieved slots support accurate passkey generation, while gating analyses show different behavior across layers and stronger activation on novel entities.

5. Interpretability Analyses

  • FwPKM’s explicit memory writes and reads let researchers trace retrieved information to specific input tokens.The NIAH case study follows a GDN+SWA|PKM@6+FwPKM@2,10 model generating a 6-digit answer after four memorization iterations.Source: §5 · p. 12, §5.1 · p. 12

5.1. Case Study: Probing Memory in NIAH

  • Sixteen slots across two FwPKM layers aggregate information despite individual errors, allowing the model to generate the correct 6-digit value.The paper characterizes this as distributed storage in which consensus across slots compensates for noisy retrievals.Source: §5.1 · p. 12

5.2. Case Study: Selective Gating

  • The lower-layer FwPKM tends to gate highly across tokens, whereas the higher-layer FwPKM activates more selectively.The visualization uses Wikipedia text about Sakana AI, a topic unseen during pre-training, with chunk size reduced from 512 to 32 tokens.Source: §5.2 · pp. 12, 14
  • Higher-layer gating spikes on rare named entities and proper nouns, indicating selective use of fast weights for novel, context-specific entities.The examples include “Sakana AI,” “David Ha,” “Llion Jones,” and “Ren Ito.”Source: §5.2 · p. 14

6. Cost Analyses

FwPKM increases total parameter count while keeping active parameters comparable to dense baselines, resulting in low theoretical compute. Its current sparse implementations nevertheless run slower than dense counterparts.

6. Cost Analyses

  • FwPKM models have about 5× more total parameters than baselines, yet active parameters remain comparable: GDN has 112M versus FwPKM’s 116M.The cost comparison covers model size and computation; the table caption specifies 128 active rows for PKM and 8 active rows per token for FwPKM.Source: §6 · p. 15, Table 1 · p. 15
  • Sparse PKM/FwPKM can require fewer FLOPs than standard dense MLP layers despite the models’ larger total parameter counts.The analysis distinguishes active parameters, theoretical FLOPs, and measured FLOPS.Source: §6 · pp. 14–15
  • FwPKM components currently run slower than dense counterparts because sparse operations use less specialized implementations than optimized dense kernels.The authors identify hardware-aware kernels for sparse fast-weight updates as a future direction.Source: §6 · p. 15

7. Related Work

Related work addresses sequence modeling through efficient attention, fast-weight updates, hybrid architectures, sparse memory, and continual learning. FwPKM builds on sparse PKM and test-time learning by making sparse memory a dynamic, online-updated fast-weight module.

Softmax attention and linear variants

  • Linear attention and recurrent or state-space variants reduce sequence-processing complexity, while softmax attention retains quadratic complexity and limits use on extremely long sequences.The surveyed approaches include Mamba, Mamba2, DeltaNet, Gated DeltaNet, RWKV7, and Memory Mosaics.Source: §7 · pp. 15–16

Fast weights and test-time training

  • Recent Test-Time Training and Titans approaches update parameters during inference to memorize the current context; FwPKM contributes a memory architecture and memorization rule for sparse structure.The cited TTT works are Sun et al. (2025) and Zhang et al. (2025b), and the cited Titans work is Behrouz et al. (2025c).Source: §7 · p. 16

Hybrid architectures

  • Hybrid architectures combine sequence-modeling components with different memory characteristics, and the paper reports that their interactions vary across tasks.The described combination includes GDN, softmax attention, slow-weight PKM, and fast-weight FwPKM.Source: §7 · p. 16

Memory models

  • PKM and PEER use sparse access to scale memory banks, while FwPKM builds on PKM by changing it from static slow memory to dynamic fast memory.The passage also describes Ultra Sparse Memory as extending PKM with more expressive keys and other improvements.Source: §7 · p. 16

Continual and episodic learning

  • Prior continual-learning work updates sparse slow weights as semantic memory; FwPKM instead updates episodic memory online, while long-term retention across domain shifts remains challenging.The paper describes low-frequency FwPKM updates as amortizing optimization cost and suggests hybrid memory systems as a future direction.Source: §7 · pp. 16–17

8. Conclusion

FwPKM turns sparse PKM into an online-updated fast-weight memory for memorizing and retrieving transient context. It generalizes from 4K-token training sequences to 128K-token contexts, but faces kernel-efficiency and long-term-retention challenges.

8. Conclusion

  • FwPKM performs sparse gradient updates to PKM key-value parameters, enabling efficient memorization of transient information and retrieval thousands of tokens later.The approach combines PKM’s large-scale sparse storage with Test-Time Training’s rapid adaptability.Source: §8 · p. 17
  • Models trained on 4K-token sequences generalize to 128K-token contexts without fine-tuning, a 32× length extrapolation.Iterative reading boosts retrieval accuracy from < 10% to > 70% in difficult Needle-in-a-Haystack scenarios.Source: §8 · p. 17
  • Sparse-update kernels currently limit speed, and FwPKM struggles with long-term retention across domain shifts.The authors point to hardware-aware sparse kernels and memory consolidation as future research directions.Source: §8 · p. 17

A. FwPKM Implementation Details

FwPKM’s implementation pairs each query with a lookahead value target for chunk-level updates, and uses IDW scoring and several update choices to stabilize and direct its sparse memory. The pseudo-code organizes processing into forward retrieval and separate backward update phases.

A.1. Lookahead Value Targets

  • Lookahead targets pair q_t with the next token’s value v_t+1, and their use improves performance enough to be the default.Fast-weight updates occur after the full chunk, so lookahead targets do not affect predictions within that chunk; inference accumulates pairs until the update chunk size is reached.Source: §A.1 · p. 22

A.2. Inverse-Distance Weighting (IDW) Scoring

  • IDW replaces only PKM’s query–key scoring within Top-k selection; its Euclidean-distance gradients make keys behave as query-cluster prototypes, and it outperforms dot-product scoring in this FwPKM setup.The scoring change leaves the other retrieval steps unchanged.Source: §A.2 · p. 22

A.3. Target Normalization and Gradient Clipping

  • Target values are z-score normalized, and fast-weight gradients are not clipped so FwPKM can match unbounded target-value scales.The passage gives the normalization and no-clipping choices as optimization measures.Source: §A.3 · p. 22

A.4. Loss Reduction and Effective Step Size

  • The memorization loss sums over tokens and features rather than averaging, avoiding a gradient scale factor of 1/(C d_v) that would change the effective fast-weight step size.The passages identify chunk size C and value dimension d_v as the dimensions that would shrink the update under mean reduction.Source: §A.4 · p. 22

A.5. Gating-weighted Memorization

  • The gate g_t weights each token’s memorization MSE according to how much the model relies on memory output at that position.The gate also interpolates the memory and other output paths.Source: §A.5 · p. 23

B. Pseudo Code

  • The pseudo-code processes chunks through a forward pass, then separate backward phases for memorization and addressing.The supplied algorithm fragments label the latter phases as memorization and anti-collapsing addressing.Source: B. Pseudo Code · p. 24

C. Detailed Experiment Settings

The experiments compare FwPKM and PKM variants under shared training and evaluation settings, then ablate design choices and test whether fast-weight updates matter. Removing lookahead targets causes the largest reported performance harm, while freezing fast weights raises perplexity on long-context datasets.

C. Detailed Experiment Settings

  • Training uses a global batch size of 128, micro-batch size 8, 100 warmup steps, and 20,000 total steps; experiments run on 4 H100 GPUs, while evaluation runs on 1 H100 GPU.Evaluation applies YaRN to RoPE embeddings when sample length exceeds 4096.Source: C. Detailed Experiment Settings · p. 25, Table 2 · p. 25
  • FwMLP is a dense fast-weight SwiGLU-MLP baseline trained on chunk-level lookahead-value MSE, whereas LaCT is a TTT baseline with sliding-window attention and fast- and slow-weight SwiGLU MLPs.FwMLP omits the sparse-memory loss aggregation, gradient shaping, and addressing optimization; LaCT’s fast weights use SGD with momentum or Muon.Source: C. Detailed Experiment Settings · pp. 25–26

D. Ablation Study

  • The ablation variants change Top-k or head count, value normalization, addressing loss, gating, gate-based loss weighting, or lookahead targets.The ablations use the GDN+SWA | PKM6+FwPKM@2,10 model and the same evaluation pipelines.Source: D. Ablation Study · p. 27

D.1. Evaluation

  • Removing lookahead values causes the greatest performance harm; addressing loss and gate-weighted MSE matter for NIAH accuracy, while 1 head × 8 slots balances task performance and compute efficiency.Increasing effective Top-k slightly improves some tasks; replacing IDW with dot-product scoring was omitted after often causing memory collapse and loss divergence.Source: §D.1 · p. 27, D. Ablation Study · p. 27

D.2. Freezing Fast Weights

  • Freezing FwPKM fast weights significantly increases perplexity on LC64 and LAMBADA, while Fineweb-Edu is less affected.The frozen parameters are the value matrix V and key matrices K^(1), K^(2).Source: §D.2 · p. 30, Figure 14 · p. 30

D.3. Addressing Metrics

The section measures slot-use balance through coverage, collision, and KLD metrics, averaging measurements over fixed token intervals and both FwPKM layers. Long-context datasets show more balanced use, but slot access remains nonuniform and better balance does not consistently track evaluation differences.

D.3. Addressing Metrics

  • Coverage rate measures the ratio of accessed slots, collision rate the ratio of competitive slot accesses, and KLD divergence from uniform slot use.The metrics are derived from row contribution, defined as the number of accesses to a value row.Source: §D.3 · p. 31
  • Metrics are computed every 32K tokens across 8M tokens, averaged over the two FwPKM layers, with 32K chosen because 8 slots per token would cover all 256K slots if selection were balanced.This yields 256 measurements per evaluation dataset and metric.Source: §D.3 · p. 31
  • Slot use is more balanced on LC64 and LAMBADA but remains far from uniform; addressing-loss ablation concentrates slot access.Increasing head count or Top-k can improve evenness but also raises collision rate.Source: §D.3 · p. 31, Table 4 · p. 31
  • Models activating 32 slots have higher coverage and lower KLD than Top-8 models, yet their evaluation differences do not correlate with those slot-use statistics.The comparison may reflect compound factors, including the Top-k mechanism itself; the authors still recommend mitigating severely unbalanced access, for example with addressing loss.Source: §D.3 · p. 31

E.1. Improving Perplexity over Sequence Boundary

FwPKM’s cross-sequence memory is associated with lower perplexity on tokens requiring cross-sequence dependencies. Its cumulative perplexity curves are jagged, which the authors hypothesize reflects use of information from previous sequences; a separate visualization shows token-level gating and update positions.

E.1. Improving Perplexity over Sequence Boundary

  • Carrying memory states across sequence boundaries helps FwPKM reduce perplexity on tokens requiring cross-sequence dependencies.The comparison plots cumulative perplexity reduction against vanilla GDN, GDN+SWA, and GDN+FA models on the first ten LC64 documents; FwPKM curves are jagged, which the authors hypothesize reflects use of prior-sequence information.Source: §E.1 · p. 32, Figure 15 · p. 32

E.2. Extra Selective Gating Example

  • The gating example visualizes layer-2 intensity in blue and layer-10 intensity in red, with black lines marking fast-weight updates every 32 tokens.Darker background colors indicate higher gating intensity.Source: §E.2 · p. 32, Figure 16 · p. 33
Loading 2601.00671v2…