Queries, Keys, and Values

Dive into Deep Learning · §10.1

Queries, keys, and values
attention as soft lookup · softmax weights · Nadaraya–Watson pooling · why learn the kernel

The fixed-size bottleneck

Every network so far assumes inputs of fixed, known size — 224 \times 224 images, or a sequence squeezed through a fixed-dimensional RNN state.

Databases don’t have this problem. A database is a set of (\text{key}, \text{value}) pairs; a query retrieves the matching value.

  • The query stays simple no matter how large the database is.
  • The same query gets different answers from different databases.
  • No need to compress the database to make lookup work.

We want a differentiable layer with these properties.

Attention as soft database lookup

Over a database \mathcal{D} = \{(\mathbf{k}_1, \mathbf{v}_1), \ldots, (\mathbf{k}_m, \mathbf{v}_m)\}:

\textrm{Attention}(\mathbf{q}, \mathcal{D}) = \sum_{i=1}^m \alpha(\mathbf{q}, \mathbf{k}_i)\, \mathbf{v}_i.

One-hot \alpha → exact lookup; uniform \alpha → average pooling. Everything in between is attention.

Attention pooling: a linear combination of values, with weights from query–key compatibility.

Softmax makes any score a weight

Pick any scoring function a(\mathbf{q}, \mathbf{k}), exponentiate, normalize:

\alpha(\mathbf{q}, \mathbf{k}_i) = \frac{\exp(a(\mathbf{q}, \mathbf{k}_i))}{\sum_j \exp(a(\mathbf{q}, \mathbf{k}_j))}.

  • Nonnegative, sums to one — a convex combination of the values.
  • Differentiable; available in every framework.
  • The rest of the chapter is about the choice of a and where \mathbf{q}, \mathbf{k}, \mathbf{v} come from.

Visualizing attention weights

The (queries × keys) heatmap is the standard diagnostic — here the identity, i.e. exact lookup:

attention_weights = torch.eye(10).reshape((1, 1, 10, 10))
d2l.show_heatmaps(attention_weights, xlabel='Keys', ylabel='Queries')

A 1964 attention mechanism

Nadaraya–Watson regression is attention pooling with a hand-picked similarity kernel:

f(\mathbf{q}) = \sum_i \mathbf{v}_i \frac{\alpha(\mathbf{q}, \mathbf{k}_i)}{\sum_j \alpha(\mathbf{q}, \mathbf{k}_j)}.

Keys = training inputs, values = labels, query = where to predict. No training at all — and consistent, if the kernel narrows with more data.

Gaussian, boxcar, constant, and triangular kernels.

Nadaraya–Watson in action

y = 2\sin(x) + x + \epsilon, 40 noisy points. Four lines of code: distances → kernel → normalize over keys → weighted sum of labels.

def nadaraya_watson(x_train, y_train, x_val, sigma):
    dists = x_train[:, None] - x_val[None, :]
    k = torch.exp(-dists**2 / (2 * sigma**2))
    attention_w = k / k.sum(0)  # Normalize over keys for each query
    return y_train @ attention_w, attention_w

sigmas = (0.1, 0.5, 2.0)
estimates = [nadaraya_watson(x_train, y_train, x_val, s)[0] for s in sigmas]
d2l.plot(x_val, estimates + [y_val], 'x', 'y',
         legend=[f'sigma = {s:g}' for s in sigmas] + ['truth'])
d2l.plt.plot(x_train, y_train, 'o', alpha=0.4);

  • The kernel’s shape barely matters; the bandwidth \sigma decides everything.

The weights explain the fits

  • Narrow kernel: sharp, local, noisy — weight on a handful of keys.
  • Wide kernel: smooth, global — weight spread across the dataset.
  • Either way the kernel is chosen, not learned, and every query gets the same notion of similarity.

Recap

  • Attention = differentiable soft database lookup: \sum_i \alpha(\mathbf{q}, \mathbf{k}_i)\, \mathbf{v}_i.
  • Softmax of any scoring function gives valid weights; exact lookup and average pooling are the extreme weight patterns.
  • Works on databases of any size with a small, fixed amount of machinery.
  • Nadaraya–Watson: attention with fixed kernels already works — learning the queries and keys is what the rest of the chapter adds.