### Install Fast Weight Product Key Memory

Source: https://github.com/lucidrains/fast-weight-product-key-memory/blob/main/README.md

Install the package via pip.

```bash
$ pip install fast-weight-product-key-memory
```

--------------------------------

### Integrate fwPKM into Transformer Block

Source: https://context7.com/lucidrains/fast-weight-product-key-memory/llms.txt

Illustrates how to integrate the fwPKM layer into a standard transformer block, augmenting it with memory capabilities. This example shows the memory layer used alongside attention and feedforward components.

```python
import torch
from torch import nn
from fast_weight_product_key_memory import fwPKM

class MemoryAugmentedTransformerBlock(nn.Module):
    def __init__(self, dim=512, heads=8, num_memories=256*256):
        super().__init__()

        # Neural memory layer
        self.memory = fwPKM(
            dim=dim,
            heads=4,
            num_memories=num_memories,
            dim_queries_keys=dim,
            dim_values=dim,
            learning_rate=1.,
            topk=8,
            chunk_size=256,
            addressing_loss_weight=5.
        )

        # Standard transformer components
        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.norm2 = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim)
        )

    def forward(self, x, past_memories=None, return_memories=False):
        # Memory-augmented residual connection
        mem_out, next_memories = self.memory(
            x,
            past_memories=past_memories,
            return_next_memories=True
        )
        x = x + mem_out

        # Standard attention
        attn_out, _ = self.attn(self.norm1(x), self.norm1(x), self.norm1(x))
        x = x + attn_out

        # Feedforward
        x = x + self.ffn(self.norm2(x))

        if return_memories:
            return x, next_memories
        return x

# Usage
block = MemoryAugmentedTransformerBlock(dim=512)
tokens = torch.randn(2, 256, 512)

# First forward pass
out1, memories = block(tokens, return_memories=True)

# Subsequent pass with accumulated memories
out2, memories = block(tokens, past_memories=memories, return_memories=True)
print(f"Output shape: {out2.shape}")  # torch.Size([2, 256, 512])
```

--------------------------------

### Initialize and Use fwPKM Module

Source: https://context7.com/lucidrains/fast-weight-product-key-memory/llms.txt

Demonstrates initializing the fwPKM module with various parameters and performing a forward pass for memory retrieval. Shows how to optionally return next memories for chaining.

```python
import torch
from fast_weight_product_key_memory import fwPKM

# Initialize the memory module
mem = fwPKM(
    dim=512,                    # Input/output embedding dimension
    num_memories=256 * 256,     # Total memory slots (must have integer sqrt)
    dim_queries_keys=512,       # Dimension for query/key projections
    dim_values=512,             # Dimension for value projections
    heads=4,                    # Number of attention heads
    topk=8,                     # Number of memories to retrieve per query
    learning_rate=1.,           # Fast weight learning rate
    chunk_size=256,             # Chunk size for memory updates
    addressing_loss_weight=10., # Weight for addressing entropy loss
    mse_loss_weight_to_keys=1.  # Weight for MSE loss to keys
)

# Create sample input tokens (batch_size=2, seq_len=256, dim=512)
tokens = torch.randn(2, 256, 512)

# Forward pass with memory retrieval only (no memory chaining)
retrieved = mem(tokens)
print(f"Retrieved shape: {retrieved.shape}")  # torch.Size([2, 256, 512])

# Forward pass returning next memories for chaining
retrieved, memories = mem(tokens, return_next_memories=True)
print(f"Memory values shape: {memories.memory_values.shape}")  # torch.Size([65536, 4, 128])
print(f"Memory keys shape: {memories.keys.shape}")  # torch.Size([2, 4, 256, 128])
```

--------------------------------

### Initialize and Inspect fwPKM Memories

Source: https://context7.com/lucidrains/fast-weight-product-key-memory/llms.txt

Demonstrates how to initialize the fwPKM with specified dimensions and parameters, and then inspect the state of the Memories NamedTuple after a forward pass. Use this to understand the memory system's internal state.

```python
import torch
from fast_weight_product_key_memory import fwPKM, Memories

mem = fwPKM(
    dim=512,
    num_memories=256 * 256,
    dim_queries_keys=512,
    dim_values=512,
    topk=8,
    chunk_size=256
)

tokens = torch.randn(2, 256, 512)
retrieved, memories = mem(tokens, return_next_memories=True)

# Inspect the Memories state
print(f"memory_values: {memories.memory_values.shape}")  # Fast weight values
print(f"keys: {memories.keys.shape}")                     # Fast weight keys
print(f"last_token: {memories.last_token.shape}")         # Last token for causal chaining
print(f"cached_tokens: {memories.cached_tokens.shape}")   # Tokens cached for next chunk
print(f"token_count: {memories.token_count}")             # Total tokens processed
print(f"num_cached: {memories.num_cached}")               # Number of cached tokens

# Initialize fresh memories using init_memories property
fresh_memories = mem.init_memories
print(f"Fresh memories initialized: {fresh_memories.token_count}")  # 0
```

--------------------------------

### Initialize and Use fwPKM

Source: https://github.com/lucidrains/fast-weight-product-key-memory/blob/main/README.md

Initialize the memory module and perform forward passes with optional memory chaining.

```python
import torch
from fast_weight_product_key_memory import fwPKM

mem = fwPKM(
    dim = 512,
    num_memories = 256 * 256,
    dim_queries_keys = 512,
    dim_values = 512,
    topk = 8,
    learning_rate = 1.,
    chunk_size = 256
)

tokens = torch.randn(2, 256, 512)

# forward a chunk of tokens for retrieved and the fast weight episodic memories

retrieved, next_memories = mem(tokens, return_next_memories = True)

# chain memories

retrieved, next_memories = mem(tokens, return_next_memories = True, past_memories = next_memories)
retrieved, next_memories = mem(tokens, return_next_memories = True, past_memories = next_memories)
retrieved, next_memories = mem(tokens, return_next_memories = True, past_memories = next_memories)
```

--------------------------------

### Train Enwik8 Model

Source: https://github.com/lucidrains/fast-weight-product-key-memory/blob/main/README.md

Run the training script for the character-level language model.

```shell
$ uv run train_enwik8.py
```

--------------------------------

### Citation BibTeX

Source: https://github.com/lucidrains/fast-weight-product-key-memory/blob/main/README.md

BibTeX entries for citing the Fast Weight Product Key Memory and related research.

```bibtex
@misc{zhao2026fastweightproductkeymemory,
    title   = {Fast-weight Product Key Memory},
    author  = {Tianyu Zhao and Llion Jones},
    year    = {2026},
    eprint  = {2601.00671},
    archivePrefix = {arXiv},
    primaryClass = {cs.CL},
    url     = {https://arxiv.org/abs/2601.00671},
}
```

```bibtex
@article{Yaghoubietal2026,
    author  = {Yaghoubi, Mohammad and Nieto-Posadas, Andres and Mosser, Coralie-Anne and Gisiger, Thomas and Wilson, Émmanuel and Williams, Sylvain and Brandon, Mark P.},
    title   = {Predictive coding of reward in the hippocampus},
    journal = {Nature},
    year    = {2026},
    doi     = {10.1038/s41586-025-09958-0}
}
```

--------------------------------

### Truncated BPTT with detach_next_memories_every

Source: https://context7.com/lucidrains/fast-weight-product-key-memory/llms.txt

Shows how to enable truncated backpropagation through time by using `detach_next_memories_every`. This periodically detaches memory gradients to manage memory usage during training on very long sequences.

```python
import torch
from fast_weight_product_key_memory import fwPKM

mem = fwPKM(
    dim=512,
    num_memories=256 * 256,
    dim_queries_keys=512,
    dim_values=512,
    topk=8,
    chunk_size=256
)

# Long sequence that would otherwise cause memory issues
long_tokens = torch.randn(1, 1024, 512)

# Enable truncated BPTT - detach memories every 2 chunks
retrieved, memories = mem(
    long_tokens,
    return_next_memories=True,
    detach_next_memories_every=2  # Detach gradients every 2 chunks
)

# Gradients only flow back through the last 2 chunks
```

--------------------------------

### Token-by-Token Inference with fwPKM

Source: https://context7.com/lucidrains/fast-weight-product-key-memory/llms.txt

Shows how to perform token-by-token autoregressive inference using fwPKM by setting chunk_size=1. This is suitable for sequence generation tasks where tokens are processed one at a time.

```python
import torch
from fast_weight_product_key_memory import fwPKM

mem = fwPKM(
    dim=32,
    num_memories=16 * 16,
    dim_queries_keys=16,
    dim_values=16,
    chunk_size=1,  # Set chunk_size=1 for token-by-token processing
    topk=4
)

# Simulate autoregressive generation
batch_size, dim = 1, 32
sequence_length = 10

state = None
outputs = []

for step in range(sequence_length):
    # Generate or receive one token at a time
    token = torch.randn(batch_size, 1, dim)

    # Process single token with memory state
    out, state = mem(
        token,
        return_next_memories=True,
        past_memories=state
    )
    outputs.append(out)
    print(f"Step {step}: processed token, memory has {state.token_count} tokens")

# Concatenate all outputs
full_output = torch.cat(outputs, dim=1)
print(f"Full output shape: {full_output.shape}")  # torch.Size([1, 10, 32])
```

--------------------------------

### Train fwPKM for Sequence Memorization

Source: https://context7.com/lucidrains/fast-weight-product-key-memory/llms.txt

Trains a PyTorch model with fwPKM to memorize and recall repeating sequences. Loss is calculated only on the recall portion (second half of the sequence).

```python
import torch
import torch.nn.functional as F
from torch import nn
from torch.optim import Adam
from fast_weight_product_key_memory import fwPKM

class MemorizingModel(nn.Module):
    def __init__(self, num_tokens, dim, use_memory=True):
        super().__init__()
        self.embed = nn.Embedding(num_tokens, dim)

        self.memory = None
        if use_memory:
            self.memory = fwPKM(
                dim=dim,
                heads=1,
                num_memories=256,
                dim_queries_keys=64,
                dim_values=dim,
                learning_rate=1.,
                topk=4,
                chunk_size=4,
                addressing_loss_weight=1e-3
            )

        self.head = nn.Linear(dim, num_tokens)

    def forward(self, x):
        h = self.embed(x)
        if self.memory is not None:
            h = h + self.memory(h)
        return self.head(h)

# Training loop for sequence memorization task
num_tokens, dim, half_len = 32, 128, 8
model = MemorizingModel(num_tokens, dim, use_memory=True)
optim = Adam(model.parameters(), lr=1e-3)

for step in range(100):
    # Create a sequence that repeats: [A, B, C, D, A, B, C, D]
    half = torch.randint(0, num_tokens, (1, half_len))
    seq = torch.cat((half, half), dim=-1)
    x, labels = seq[:, :-1], seq[:, 1:]

    logits = model(x)

    # Loss only on second half (recall portion)
    loss = F.cross_entropy(
        logits[:, (half_len - 1):].reshape(-1, num_tokens),
        labels[:, (half_len - 1):].reshape(-1)
    )

    loss.backward()
    optim.step()
    optim.zero_grad()

    if step % 20 == 0:
        print(f"Step {step}, Loss: {loss.item():.4f}")
```

--------------------------------

### Memory Chaining for Sequential Processing

Source: https://context7.com/lucidrains/fast-weight-product-key-memory/llms.txt

Illustrates how to chain memories across multiple forward passes to process long sequences. The `past_memories` argument is used to maintain and accumulate the memory state.

```python
import torch
from fast_weight_product_key_memory import fwPKM

mem = fwPKM(
    dim=512,
    num_memories=256 * 256,
    dim_queries_keys=512,
    dim_values=512,
    topk=8,
    learning_rate=1.,
    chunk_size=256
)

# Process a long sequence in chunks with memory chaining
tokens_chunk1 = torch.randn(2, 256, 512)
tokens_chunk2 = torch.randn(2, 256, 512)
tokens_chunk3 = torch.randn(2, 256, 512)

# First chunk - initialize memories
retrieved1, memories = mem(tokens_chunk1, return_next_memories=True)

# Second chunk - use accumulated memories
retrieved2, memories = mem(tokens_chunk2, return_next_memories=True, past_memories=memories)

# Third chunk - memories continue accumulating
retrieved3, memories = mem(tokens_chunk3, return_next_memories=True, past_memories=memories)

# Memory state persists across chunks
print(f"Token count in memory state: {memories.token_count}")  # 768 (256 * 3)
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.