Memory Networks & Attention-Based Architectures
Memory networks augment neural networks with external memory stores, enabling explicit read/write operations and long-term reasoning. This module covers memory-augmented neural networks, the Differentiable Neural Computer, and the connection between attention and memory.
1. Memory Network Basics
1.1 Architecture
A memory network consists of:
- Input feature map : Converts input to memory representation
- Generalization : Updates memory based on input
- Output feature map : Reads from memory given query
- Response : Produces final output
1.2 Memory Representation
Memory is a matrix of slots, each with dimensions:
1.3 Memory Addressing
Content-based addressing: Find memory slots similar to query :
Location-based addressing: Access memory by index.
Hybrid addressing: Combine content and location.
2. Memory-Augmented Neural Networks (MANN)
2.1 Neural Turing Machine (NTM)
Graves et al. (2014) introduced differentiable read/write operations.
Read operation:
where is the attention weight over memory slots.
Write operation:
where is the erase vector and is the add vector.
2.2 Addressing Mechanism
Content-based addressing:
where is the key and is the strength parameter.
Location-based addressing (shift weights):
Interpolation gate:
Convolutional shift:
Sharpening:
2.3 Controller
The controller is typically an LSTM or feedforward network:
The controller outputs: read key, shift, strength, erase, add, output.
3. Differentiable Neural Computer (DNC)
3.1 DNC Architecture
Santoro et al. (2018) extended NTM with:
- Usage vector: Track memory slot usage
- Link matrix: Record temporal ordering
- Priority: Weight by recency and frequency
3.2 Usage Vector
This tracks the probability each slot has been used.
3.3 Free Gate
Controls whether to write to existing slots or allocate new ones.
3.4 Allocation
Unused slots get higher allocation weight.
3.5 Write Weighting
3.6 Link Matrix
Tracks temporal ordering of writes.
3.7 Read Head
Forward link:
Backward link:
4. Attention as Memory
4.1 Self-Attention as Memory
In Transformers, the key-value pairs form a memory:
Read operation:
4.2 Linear Attention as Memory
Linear attention approximates softmax attention:
where is a feature map.
Recursive formulation:
This is a linear-time recurrent formulation with memory state .
4.3 Compressive Transformer
Compress past memories to maintain long-range context:
Store compressed memories alongside recent ones.
4.4 Memory Transformer (Memorizing Transformer)
Store all key-value pairs in external memory:
where includes all historical key-value pairs.
5. Implementation
import torch
import torch.nn as nn
import torch.nn.functional as F
class NTM(nn.Module):
def __init__(self, input_size, output_size, memory_slots=128, memory_dim=64):
super().__init__()
self.memory_slots = memory_slots
self.memory_dim = memory_dim
self.controller = nn.LSTM(input_size + memory_dim, 256, batch_first=True)
self.read_key = nn.Linear(256, memory_dim)
self.read_strength = nn.Linear(256, 1)
self.read_shift = nn.Linear(256, 3)
self.read_gamma = nn.Linear(256, 1)
self.write_key = nn.Linear(256, memory_dim)
self.write_strength = nn.Linear(256, 1)
self.write_shift = nn.Linear(256, 3)
self.write_gamma = nn.Linear(256, 1)
self.write_erase = nn.Linear(256, memory_dim)
self.write_add = nn.Linear(256, memory_dim)
self.output = nn.Linear(256 + memory_dim, output_size)
def content_addressing(self, key, memory, strength):
strength = F.softplus(strength) + 1e-8
similarity = F.cosine_similarity(
key.unsqueeze(1), memory.unsqueeze(0), dim=2
)
return F.softmax(strength * similarity, dim=-1)
def shift_addressing(self, w_prev, shift_weights):
batch_size = w_prev.shape[0]
shift = F.softmax(shift_weights, dim=-1)
w_shifted = torch.zeros_like(w_prev)
for i in range(3):
shift_amount = i - 1
w_shifted += shift[:, i:i+1] * torch.roll(w_prev, shift_amount, dims=1)
return w_shifted
def sharpen(self, w, gamma):
gamma = 1 + F.softplus(gamma)
w = w ** gamma
return w / (w.sum(dim=-1, keepdim=True) + 1e-8)
def forward(self, x):
batch_size, seq_len, _ = x.shape
memory = torch.zeros(batch_size, self.memory_slots, self.memory_dim).to(x.device)
w_prev_read = torch.zeros(batch_size, self.memory_slots).to(x.device)
w_prev_write = torch.zeros(batch_size, self.memory_slots).to(x.device)
h_prev = torch.zeros(1, batch_size, 256).to(x.device)
c_prev = torch.zeros(1, batch_size, 256).to(x.device)
outputs = []
for t in range(seq_len):
r_prev = (w_prev_read.unsqueeze(-1) * memory).sum(dim=1)
controller_input = torch.cat([x[:, t], r_prev], dim=-1)
h, (h_prev, c_prev) = self.controller(controller_input.unsqueeze(1), (h_prev, c_prev))
h = h.squeeze(1)
# Read
r_key = self.read_key(h)
r_strength = self.read_strength(h)
r_shift = self.read_shift(h)
r_gamma = self.read_gamma(h)
w_content = self.content_addressing(r_key, memory, r_strength)
w_shifted = self.shift_addressing(w_prev_read, r_shift)
w_read = self.sharpen(w_shifted, r_gamma)
r = (w_read.unsqueeze(-1) * memory).sum(dim=1)
# Write
w_key = self.write_key(h)
w_strength = self.write_strength(h)
w_shift = self.write_shift(h)
w_gamma = self.write_gamma(h)
erase = torch.sigmoid(self.write_erase(h))
add = torch.tanh(self.write_add(h))
w_content = self.content_addressing(w_key, memory, w_strength)
w_shifted = self.shift_addressing(w_prev_write, w_shift)
w_write = self.sharpen(w_shifted, w_gamma)
memory = memory * (1 - w_write.unsqueeze(-1) * erase.unsqueeze(1))
memory = memory + w_write.unsqueeze(-1) * add.unsqueeze(1)
w_prev_read = w_read
w_prev_write = w_write
output = self.output(torch.cat([h, r], dim=-1))
outputs.append(output)
return torch.stack(outputs, dim=1)
class LinearAttention(nn.Module):
def __init__(self, dim, num_heads=8, feature_map='elu'):
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.qkv = nn.Linear(dim, 3 * dim)
self.out = nn.Linear(dim, dim)
if feature_map == 'elu':
self.phi = nn.ELU()
elif feature_map == 'relu':
self.phi = nn.ReLU()
elif feature_map == 'softmax':
self.phi = lambda x: F.softmax(x, dim=-1)
def forward(self, x):
B, N, C = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
q, k, v = qkv.unbind(2)
q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v))
q = self.phi(q)
k = self.phi(k)
# Linear attention: (Q^T K) V instead of softmax(Q K^T) V
kv = torch.einsum('bhnd,bhne->bhde', k, v)
qkv = torch.einsum('bhnd,bhde->bhne', q, kv)
# Normalize
k_sum = k.sum(dim=2, keepdim=True)
denom = torch.einsum('bhnd,bhnd->bhn', q, k_sum).unsqueeze(-1) + 1e-6
out = qkv / denom
out = out.transpose(1, 2).reshape(B, N, C)
return self.out(out)
class DNC(nn.Module):
def __init__(self, input_size, output_size, memory_slots=64, memory_dim=32):
super().__init__()
self.N = memory_slots
self.M = memory_dim
self.controller = nn.LSTM(input_size + memory_dim, 128, batch_first=True)
self.read_head = nn.Linear(128, memory_dim + 3 + 1 + 1)
self.write_head = nn.Linear(128, memory_dim + memory_dim + 3 + 1 + 1 + 1)
self.output = nn.Linear(128 + memory_dim, output_size)
def forward(self, x):
B, T, _ = x.shape
memory = torch.zeros(B, self.N, self.M).to(x.device)
usage = torch.zeros(B, self.N).to(x.device)
link = torch.zeros(B, self.N, self.N).to(x.device)
w_prev = torch.zeros(B, self.N).to(x.device)
h = c = torch.zeros(1, B, 128).to(x.device)
outputs = []
for t in range(T):
r = (w_prev.unsqueeze(-1) * memory).sum(dim=1)
inp = torch.cat([x[:, t], r], dim=-1)
h, (h, c) = self.controller(inp.unsqueeze(1), (h, c))
h = h.squeeze(1)
# Read
read_params = self.read_head(h)
r_key = read_params[:, :self.M]
r_strength = read_params[:, self.M]
r_shift = read_params[:, self.M+1:self.M+4]
r_gamma = read_params[:, self.M+4]
# Content addressing
sim = F.cosine_similarity(r_key.unsqueeze(1), memory, dim=2)
w_content = F.softmax(F.softplus(r_strength).unsqueeze(1) * sim, dim=-1)
# Shift
shift = F.softmax(r_shift, dim=-1)
w_shifted = torch.zeros_like(w_prev)
for i in range(3):
w_shifted += shift[:, i:i+1] * torch.roll(w_prev, i-1, dims=1)
# Sharpen
gamma = 1 + F.softplus(r_gamma)
w_read = (w_shifted ** gamma) / ((w_shifted ** gamma).sum(dim=1, keepdim=True) + 1e-8)
# Write
write_params = self.write_head(h)
w_key = write_params[:, :self.M]
w_strength = write_params[:, self.M]
w_shift = write_params[:, self.M+1:self.M+4]
w_gamma = write_params[:, self.M+4]
erase = torch.sigmoid(write_params[:, self.M+5:self.M+5+self.M])
add = torch.tanh(write_params[:, self.M+5+self.M:])
# Write addressing
sim = F.cosine_similarity(w_key.unsqueeze(1), memory, dim=2)
w_content = F.softmax(F.softplus(w_strength).unsqueeze(1) * sim, dim=-1)
shift = F.softmax(w_shift, dim=-1)
w_shifted = torch.zeros_like(w_prev)
for i in range(3):
w_shifted += shift[:, i:i+1] * torch.roll(w_prev, i-1, dims=1)
gamma = 1 + F.softplus(w_gamma)
w_write = (w_shifted ** gamma) / ((w_shifted ** gamma).sum(dim=1, keepdim=True) + 1e-8)
# Update memory
memory = memory * (1 - w_write.unsqueeze(-1) * erase.unsqueeze(1))
memory = memory + w_write.unsqueeze(-1) * add.unsqueeze(1)
# Update usage
usage = usage + w_write - usage * w_write
# Update link
link = link * (1 - w_write.unsqueeze(2) * w_write.unsqueeze(1))
link = link + w_write.unsqueeze(2) * w_prev.unsqueeze(1)
w_prev = w_read
r = (w_read.unsqueeze(-1) * memory).sum(dim=1)
outputs.append(self.output(torch.cat([h, r], dim=-1)))
return torch.stack(outputs, dim=1)
6. SVG: Memory Network Architecture
7. SVG: DNC Read/Write Heads
8. Comparison of Methods
| Model | Memory | Addressing | Complexity | Use Case |
|---|---|---|---|---|
| NTM | Fixed | Content + Location | O(N) | Algorithmic tasks |
| DNC | Fixed | Content + Location + Links | O(N²) | Complex reasoning |
| Transformer | Self | Content only | O(N²) | General sequence modeling |
| Linear Attention | Recursive | Content (linear) | O(N) | Long sequences |
| Memorizing Transformer | External | Content | O(N²) | Long-context QA |
9. Open Problems
- Scalability: Scaling memory to millions of entries
- Hierarchical memory: Multiple levels of abstraction
- Memory compression: Efficient storage and retrieval
- Persistent memory: Learning what to remember across sessions
- Memory efficiency: Reducing attention's quadratic complexity
References
- Graves, A., Wayne, G., & Danihelka, I. (2014). Neural Turing Machines. arXiv.
- Sukhbaatar, S., Szlam, A., Weston, J., & Fergus, R. (2015). End-To-End Memory Networks. NeurIPS.
- Santoro, A., et al. (2018). Measuring Abstract Reasoning in Neural Networks. ICML.
- Katharopoulos, A., et al. (2020). Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. ICML.
- Wu, Y., et al. (2022). Memorizing Transformers. ICLR.