Memory Augmentation › Memory Network Fundamentals

Imagine you're reading a 500-page novel and someone asks you, "What color was the protagonist's childhood home?" You don't re-read the entire book. Instead, your mind somehow retrieves that specific detail from storage. Memory networks attempt to give neural networks this same capability: the ability to write information to an external store during processing and retrieve it on demand.

The fundamental innovation of memory networks is decoupling what the model remembers from how it processes information. In a standard neural network, "memory" exists implicitly in the hidden states, and information must survive through layer after layer of transformation. In a memory network, we provide an explicit storage mechanism that the model can query independently.

This separation immediately suggests two operations:

1. Writing: The model encounters important information and stores it
2. Reading: The model needs information and retrieves it from storage

But how should storage be organized? And how does the model know where to look when reading?

Consider how you might organize a library. You could assign each book a numbered shelf (positional addressing), or you could organize books by topic and keywords (content-based addressing). A positional system requires remembering exactly where you put each book. A content-based system lets you find books by describing what you're looking for.

Memory networks use content-based addressing. Each memory entry consists of two parts:

• A key $\mathbf{k}_i$: describes what the entry contains (like a book's title and keywords)
• A value $\mathbf{v}_i$: the actual stored content (like the book itself)

When the model wants to retrieve information, it provides a query $\mathbf{q}$ describing what it's looking for. The memory system compares this query against all stored keys and returns the values whose keys match best.

How do we measure whether a query matches a key? The most natural choice is the dot product. Given a query vector $\mathbf{q}$ and a key vector $\mathbf{k}_i$, their similarity is:

$$s_i = \mathbf{q} \cdot \mathbf{k}_i = \sum_{j=1}^{d} q_j \cdot k_{ij}$$

where:

• $\mathbf{q} \in \mathbb{R}^d$ is the query vector with dimension $d$
• $\mathbf{k}_i \in \mathbb{R}^d$ is the key for memory slot $i$
• $s_i$ is the scalar similarity score

This score tells us how well the query matches each key. Larger positive values indicate better matches. But we need to convert these raw scores into a probability distribution, a set of weights that sum to 1, telling us how much to trust each memory entry.

The softmax function accomplishes this. For $n$ memory slots with similarity scores $s_1, \ldots, s_n$:

$$\alpha_i = \frac{e^{s_i}}{\sum_{j=1}^{n} e^{s_j}}$$

The exponential $e^{s_i}$ ensures all weights are positive (since $e^x > 0$ for all $x$). The denominator normalizes so that $\sum_i \alpha_i = 1$. The weight $\alpha_i$ can be interpreted as "the probability that memory slot $i$ contains what we're looking for."

In practice, we add one refinement. When the dimension $d$ is large, dot products can become very large in magnitude. This pushes softmax outputs toward extreme values (nearly 0 or nearly 1), which causes gradient problems during training. The solution is scaled dot-product attention:

$$s_i = \frac{\mathbf{q} \cdot \mathbf{k}_i}{\sqrt{d}}$$

Why $\sqrt{d}$? If each component of $\mathbf{q}$ and $\mathbf{k}$ has variance 1, then the dot product has variance $d$ (since we're summing $d$ independent products). Dividing by $\sqrt{d}$ restores unit variance, keeping the scores in a well-behaved range regardless of dimension.

Once we have attention weights $\alpha_i$ over all memory slots, we retrieve information by taking a weighted sum of the stored values:

$$\mathbf{r} = \sum_{i=1}^{n} \alpha_i \cdot \mathbf{v}_i$$

where:

• $\mathbf{r} \in \mathbb{R}^{d_v}$ is the retrieved vector (output)
• $\mathbf{v}_i \in \mathbb{R}^{d_v}$ is the value stored in memory slot $i$
• $\alpha_i$ is the attention weight computed above

This weighted sum is fully differentiable, meaning gradients can flow back through the retrieval operation during training. The model learns which keys to associate with which values by receiving gradient signals about whether the retrieved information helped with the task.

The full memory read operation combines all these steps:

$$\mathbf{r} = \text{MemoryRead}(\mathbf{q}, \mathbf{K}, \mathbf{V}) = \sum_{i=1}^{n} \text{softmax}\left(\frac{\mathbf{q} \mathbf{K}^T}{\sqrt{d}}\right)_i \cdot \mathbf{v}_i$$

where:

• $\mathbf{K} \in \mathbb{R}^{n \times d}$ is the matrix of all $n$ keys stacked as rows
• $\mathbf{V} \in \mathbb{R}^{n \times d_v}$ is the matrix of all $n$ values stacked as rows
• $\mathbf{q} \mathbf{K}^T$ computes all $n$ similarity scores in parallel

This is identical to the attention mechanism used in transformers. The difference is conceptual: in standard self-attention, keys and values come from the input sequence; in a memory network, they come from an external storage that persists across inputs.

Let's translate these mathematical concepts into working code. The implementation directly mirrors the formulas above: