Memory Augmentation › Memory Retrieval Mechanisms: From Soft to Efficient

The basic memory read operation we just implemented has an elegant mathematical formulation, but it harbors a serious scalability problem. Let's understand this issue deeply before exploring solutions.

Our memory read computes attention over all memory slots:

$$\alpha_i = \text{softmax}\left(\frac{\mathbf{q} \cdot \mathbf{k}_i}{\sqrt{d}}\right) \quad \text{for } i = 1, \ldots, n$$

This soft attention approach has an important advantage: it's fully differentiable. Gradients flow smoothly from the output through the softmax and into the query and key representations. The model can learn what to store and how to query through standard backpropagation.

But soft attention also inherits the same dilution problem we identified earlier for long-context transformers. When $n$ is large, each slot's weight $\alpha_i$ becomes vanishingly small. The information we need gets drowned in a weighted average dominated by irrelevant entries.

One way to combat dilution is to sharpen the attention distribution. We can introduce a temperature parameter $\tau$ that scales the scores before softmax:

$$\alpha_i = \text{softmax}\left(\frac{\mathbf{q} \cdot \mathbf{k}_i}{\sqrt{d} \cdot \tau}\right)$$

The temperature $\tau$ controls the "peakiness" of the distribution:

• Low temperature ($\tau < 1$): Amplifies score differences. Small advantages become large, concentrating attention on the best-matching entries. In the limit $\tau \to 0$, this approaches hard argmax selection.

• High temperature ($\tau > 1$): Dampens score differences. The distribution becomes more uniform, spreading attention across many entries.

We can measure this effect through entropy, which quantifies how spread out a distribution is:

$$H(\alpha) = -\sum_{i=1}^{n} \alpha_i \log \alpha_i$$

Low entropy means attention is concentrated on few slots; high entropy means it's spread across many.

The results demonstrate the entropy-focus tradeoff concretely. At temperature 0.1, the entropy is low and the maximum weight approaches 1.0, meaning the model concentrates almost all attention on a single memory slot. The softmax has become nearly a hard argmax. At temperature 2.0, entropy increases substantially as attention spreads more evenly across slots.

This tradeoff has practical implications:

• For fact retrieval (e.g., "What is the capital of France?"), use low temperatures. The answer exists in one specific memory location.
• For synthesis tasks (e.g., "Summarize the key themes"), use higher temperatures. Relevant information is distributed across many entries.
• For training, moderate temperatures (0.5-1.0) often work best. They allow gradient flow to multiple entries without completely diluting the signal.

However, temperature alone cannot solve the scalability problem. Even with sharp attention, we still compute similarities with all $n$ memory entries. When $n$ reaches millions, this becomes prohibitively expensive.

A more radical solution abandons the soft attention over all entries. Instead, we first find the $k$ most relevant memory slots using a discrete selection step, then apply soft attention only over those $k$ entries:

$$\text{Top-}k(\mathbf{s}) = \{i_1, i_2, \ldots, i_k\} \quad \text{where } s_{i_1} \geq s_{i_2} \geq \cdots \geq s_{i_k}$$

The top-$k$ operation returns the indices of the $k$ highest-scoring entries. We then compute attention only over this subset:

$$\alpha_j = \text{softmax}(s_{i_j}) \quad \text{for } j = 1, \ldots, k$$ $$\mathbf{r} = \sum_{j=1}^{k} \alpha_j \cdot \mathbf{v}_{i_j}$$

The key insight is that $k$ is fixed regardless of memory size $n$. Whether the memory contains 1,000 or 1,000,000 entries, we only attend over $k$ slots in the final step. This bounds the attention computation at $O(k)$ regardless of $n$.

There's a catch: the top-$k$ selection involves an argmax-like operation (finding the highest scores), which is not differentiable. During backpropagation, we cannot compute gradients through the discrete selection step.

This means entries that weren't selected receive no gradient signal, even if they were "almost" selected. The model cannot learn to make a near-miss entry more relevant. In practice, this limitation is often acceptable because:

1. The entries that were selected still receive full gradients
2. The key representations are typically trained with other objectives (like the main language modeling loss)
3. The selection step can be approximated during training using straight-through estimators or Gumbel-softmax tricks

For many applications, the efficiency gains vastly outweigh this differentiability cost.

The efficiency gains are dramatic. With $k=5$, we use only 0.05% of the memory slots for the attention computation, a 2000x reduction compared to attending to all entries. Even with $k=100$ to retrieve more entries, we still use only 1% of the memory.

Let's trace through the computational costs:

1. Scoring all keys: We still compute $\mathbf{q} \cdot \mathbf{k}_i$ for all $n$ entries: $O(n \cdot d)$
2. Finding top-k: Selecting the $k$ largest from $n$ scores: $O(n)$ (or $O(n \log k)$ for a heap-based algorithm)
3. Attention over k: Softmax and weighted sum over $k$ entries: $O(k \cdot d_v)$

The total is $O(n \cdot d)$ for exact top-k retrieval. We've bounded the attention step but not the scoring step.

For truly large memories (millions of entries), even linear scoring becomes prohibitive. This is where approximate nearest neighbor (ANN) algorithms enter. Methods like FAISS, ScaNN, or Annoy build index structures that support sublinear retrieval:

• Exact search: $O(n)$ to find top-k
• Approximate search: $O(\log n)$ or even $O(1)$ depending on the algorithm and acceptable accuracy loss

With ANN indexing, a memory of 1 million entries can be queried nearly as fast as a memory of 1,000 entries, at the cost of occasionally missing the truly best match.

Top-k retrieval solves the attention efficiency problem but treats all memory entries uniformly. When memories span multiple time scales, a more sophisticated organization helps: hierarchical memory separates recent detailed information from older compressed summaries.

Short-term memory holds recent context at full resolution, while long-term memory stores condensed representations from earlier. The model queries both levels, using short-term memory for detailed recent information and long-term memory for general patterns.