Deep Learning: Making It Learnable

Heman Shakeri

14 Kernel Regression: Attention Before It Was Learnable

Chapter 13 ended with a machine that could translate, and one uncomfortable picture: every source fact had to cross a single fixed-size bridge before the decoder wrote its first token. The encoder had computed a state at every source position. Then we threw all but the last one away.

Why throw them away? Keep the whole sequence of encoder states as a memory bank. At each decoding step, ask what the decoder needs, compare that query with every stored state, and retrieve a weighted mixture of the relevant content. Keeping a bank does not make its entries lossless, but it removes the rule that all source information must pass through the last entry alone.

One problem remains: how should a query decide the mixture? A nearest-neighbor rule changes its chosen slot abruptly. We first need a smooth similarity-weighted rule, and statistics had one in 1964, half a century before modern neural attention.

Start with three observed pairs:

observation	location \(x_i\)	response \(y_i\)
1	1	1.5
2	3	2.8
3	5	1.8

For a query \(q=3.5\), a Gaussian similarity with bandwidth \(h=0.6\) assigns weights \((0.0002, 0.9413, 0.0585)\). Their weighted mixture is \(2.7412\). The second observation dominates because its location is closest, the third contributes a little, and the distant first observation is nearly silent. No similarity function was learned; the chosen metric and bandwidth did the routing.

14.1 Chapter 1 already mixed old answers

The linear regression of Chapter 1 makes a prediction by mixing the observed targets. Recall from that chapter’s hat matrix that every fitted value is already a weighted combination of the training targets; the same algebra gives the weights at a new query.

Let the augmented design matrix \(\widetilde{\matr{X}}\in\mathbb{R}^{n\times(d+1)}\) include the intercept column, let \(\widetilde{\vect{q}}\in\mathbb{R}^{d+1}\) be a new query with its leading 1, and let \(\vect{y}\in\mathbb{R}^{n}\) contain the observed targets. Assuming full column rank,

\[ \hat{\vect{\beta}} = (\widetilde{\matr{X}}^\top\widetilde{\matr{X}})^{-1} \widetilde{\matr{X}}^\top\vect{y}. \]

The prediction at the query is therefore

\[ \hat f(\vect{q}) = \widetilde{\vect{q}}^\top\hat{\vect{\beta}} = {\underbrace{\left[ \widetilde{\matr{X}} (\widetilde{\matr{X}}^\top\widetilde{\matr{X}})^{-1} \widetilde{\vect{q}} \right]}_{\vect{\lambda}(\vect{q})}}^{\!\top}\vect{y}. \tag{14.1}\]

Trap: weighted combination does not mean weighted average

Chapter 1’s prediction \(\hat f(\vect{q})\) is exactly a weighted combination of the training targets. With an intercept, the weights sum to one because OLS reproduces a constant response. But weighted combination is not the same as weighted average. The entries of \(\vect{\lambda}\) may be negative.

For keys \((-2,-1,0,1,2)\) and query \(q=1.5\), the OLS weights are

\[ (-0.10, 0.05, 0.20, 0.35, 0.50). \]

They sum to one, yet the first is negative. If the targets are \((1,0,0,0,0)\), OLS predicts \(-0.10\), outside the observed range \([0,1]\). Linear regression is allowed to leave the observed targets’ convex hull even when the query lies inside the observed key range. For local smoothing, we impose three conditions: nearby evidence should receive more weight, every weight should be nonnegative, and the weights should sum to one.

14.2 Similarity weights and local prediction

A smoothing kernel is a similarity rule. It looks at a query and a stored key and returns a nonnegative affinity: large when they are close, small when they are far. To avoid confusing the kernel function with the matrix of keys later, we write it as \(\kappa_h\). The Gaussian choice is

\[ \kappa_h(\vect{q},\vect{k}) = \exp\!\left(-\frac{\norm{\vect{q}-\vect{k}}^2}{2h^2}\right), \qquad h>0. \tag{14.2}\]

The bandwidth \(h\) sets the kernel’s width, like the width of a magnifying glass. A small value sees only the closest keys; a large value blends a broad neighborhood.

Given observed pairs \((\vect{x}_i,y_i)\), normalize the affinities,

\[ \alpha_i(\vect{q}) = \frac{\kappa_h(\vect{q},\vect{x}_i)} {\sum_{j=1}^{n}\kappa_h(\vect{q},\vect{x}_j)}, \qquad \hat f_h(\vect{q})=\sum_{i=1}^{n}\alpha_i(\vect{q})y_i. \tag{14.3}\]

This is the Nadaraya–Watson estimator, proposed independently by Nadaraya and Watson in 1964. For the strictly positive Gaussian, each \(\alpha_i>0\) and \(\sum_i\alpha_i=1\). Scalar predictions remain between the smallest and largest observed values. Unlike OLS, this operator smooths locally while keeping every prediction inside the observed values’ convex hull.

As \(h\to0^+\), the weight concentrates on the unique nearest key, so the smoother approaches the nearest-neighbor rule of the chapter opening; exact nearest-key ties split the mass. At \(h=0\) the formula is undefined. As \(h\to\infty\), every finite distance looks alike and the weights approach \(1/n\), so every prediction approaches the unweighted mean of the observed responses.

We now call the stored input locations keys \(\vect{k}_i\), and their associated responses values \(\vect{v}_i\). The values may be vectors. For a query \(\vect{q}\), consider fitting one constant vector to these values, weighted by their keys’ similarity to the query.

Trap: the kernel supplies the weights; the constant fit supplies the average

It is tempting to say that attention solves an unrestricted regression problem with zero regularization. That is too loose. An unrestricted function could interpolate the stored values and remain undetermined elsewhere; nothing in that statement forces an average at a new query. The exact claim is local and smaller. For a fixed query \(\vect{q}\), fit one constant vector \(\vect{c}\):

\[ \vect{c}^{*}(\vect{q}) =\arg\min_{\vect{c}} \frac{1}{2}\sum_i \kappa(\vect{q},\vect{k}_i) \norm{\vect{c}-\vect{v}_i}^{2}. \tag{14.4}\]

Setting its gradient to zero gives \(\sum_i\kappa_i(\vect{c}^{*}-\vect{v}_i)=\vect{0}\), so \(\vect{c}^{*}=\sum_i\kappa_i\vect{v}_i/\sum_i\kappa_i\). With \(\kappa(\vect{q},\vect{k})=\exp(\vect{q}^{\top}\vect{k}/\sqrt d)\), that is exactly scaled dot-product softmax attention. The kernel supplies nonnegative weights; the local-constant model is what makes their solution an average.

Here is the whole operator in NumPy. The shapes come before the multiplication: queries are \((m,d)\), keys are \((n,d)\), values are \((n,p)\), and the returned mixture is \((m,p)\).

Fix the seed and set NumPy to print four decimals without scientific notation.
Define the reusable helpers: row_softmax and gaussian_attention.
Fixed Gaussian attention, from scores to mixtures.

from __future__ import annotations

import matplotlib.pyplot as plt
import numpy as np
import torch
torch.set_num_threads(4)  # pinned: results depend on the CPU thread count

# [1]
torch.manual_seed(6050)
np.set_printoptions(precision=4, suppress=True)

# [2]
def row_softmax(scores: np.ndarray) -> np.ndarray:
    """Normalize each row after a stability-preserving shift."""
    shifted = scores - scores.max(axis=1, keepdims=True)
    weights = np.exp(shifted)
    return weights / weights.sum(axis=1, keepdims=True)

def gaussian_attention(
    queries: np.ndarray,
    keys: np.ndarray,
    values: np.ndarray,
    bandwidth: float,
) -> tuple[np.ndarray, np.ndarray]:
    """Fixed Gaussian attention for Q:(m,d), K:(n,d), V:(n,p)."""
    if not np.isfinite(bandwidth) or bandwidth <= 0:
        raise ValueError("bandwidth must be positive and finite")
    if queries.ndim != 2 or keys.ndim != 2 or values.ndim != 2:
        raise ValueError("queries, keys, and values must be matrices")
    if queries.shape[1] != keys.shape[1] or keys.shape[0] != values.shape[0]:
        raise ValueError("Q/K feature dimensions and K/V row counts must agree")

    distances2 = ((queries[:, None, :] - keys[None, :, :]) ** 2).sum(axis=2)
    log_scores = -distances2 / (2.0 * bandwidth**2)  # (m, n)
    weights = row_softmax(log_scores)                # (m, n)
    return weights @ values, weights                 # (m, p), (m, n)

keys3 = np.array([[1.0], [3.0], [5.0]])
values3 = np.array([[1.5], [2.8], [1.8]])
query3 = np.array([[3.5]])
pred3, weights3 = gaussian_attention(query3, keys3, values3, bandwidth=0.6)
# [3]
print("weights:", weights3[0])
print(f"prediction: {pred3.item():.4f}")

weights: [0.0002 0.9413 0.0585]
prediction: 2.7412

A query star connects to three observed points with line widths proportional to their influence. In the adjacent bar chart, the nearby middle key receives most of the normalized weight. — Figure 14.1: A fixed Gaussian lookup. The bars give the normalized weights, and the connecting lines show which observations feed the query. The query lands close to the second location, so its response dominates the mixture; no similarity function was learned.

Watch the mechanism: distance becomes influence

As the query moves, which observation gains influence?

Three observations, one fixed bandwidth. Follow the calculation above, then move only the query.

● Observed value◆ Prediction Query line Distance bracket

Query q = 3.50Bandwidth h = 0.60 (fixed)

Key 1 · value 1.5

Distance: 2.50
Affinity: 0.00017
Weight: 0.0002
Weight × value: 0.0003

Key 3 · value 2.8

Distance: 0.50
Affinity: 0.70665
Weight: 0.9413
Weight × value: 2.6355

Key 5 · value 1.8

Distance: 1.50
Affinity: 0.04394
Weight: 0.0585
Weight × value: 0.1053

Affinity = exp(−distance² / (2h²)); weight = affinity / sum of affinities.

Weights sum to 1.0000.

Prediction = 2.7412, within the observed-value range [1.5, 2.8].

The nearby key at 3 receives most of the influence. The same normalized weights multiply the observed values and add to one prediction.

The calculation is readable without playback. Opening this panel loads the optional controls.

Only the query moves. These weights are computed, not learned; the observations and bandwidth stay fixed. Nonnegative weights summing to one keep the prediction within the observed-value range.

Check yourself. Move the query onto the first observation, q = 1. Which observation dominates, and what is the prediction?

The first one: its weight is about 0.996, so the prediction is about 1.505 — almost its own value, 1.5. With a bandwidth this narrow, a query sitting on a key is nearly a nearest-neighbour lookup.

Transcript and keyboard help

Keep keys 1, 3, 5 and their observed values 1.5, 2.8, 1.8 fixed. At q = 3.5, the distances are 2.5, 0.5, 1.5.
With bandwidth h = 0.6, the Gaussian affinities are approximately 0.00017, 0.70665, 0.04394. Divide each by their common sum, approximately 0.75076.
The normalized weights are approximately 0.0002, 0.9413, 0.0585. Multiply each by its observed value and add: the prediction is approximately 2.7412. Calculations use unrounded weights.
Now move the query across the same observations. A closer key generally gains influence relative to farther keys. Weights still sum to one; the prediction stays between 1.5 and 2.8. Return to q = 3.5, the printed witness.

Silent; paused initially; 1.5× playback. Focus the animation pane: Space or K plays/pauses, arrows seek, Home/End jump, Escape pauses. The scrubber and speed menu also work by keyboard. Reduced-motion mode uses discrete query positions.

The same code shows the difference between the signed OLS weights and the convex kernel weights:

Set five keys from -2 to 2 and the query 1.5, and add an intercept column to each.
Signed OLS weights versus convex kernel weights.
Print the OLS and kernel weights with their sums, and both predictions for targets (1, 0, 0, 0, 0).

# [1]
ols_keys = np.array([-2.0, -1.0, 0.0, 1.0, 2.0])
ols_query = 1.5
design = np.column_stack([np.ones_like(ols_keys), ols_keys])
query_design = np.array([1.0, ols_query])
# [2]
ols_weights = query_design @ np.linalg.solve(design.T @ design, design.T)

witness_values = np.array([[1.0], [0.0], [0.0], [0.0], [0.0]])
nw_pred, nw_weights = gaussian_attention(
    np.array([[ols_query]]), ols_keys[:, None], witness_values, bandwidth=1.0
)
# [3]
print("OLS weights:", ols_weights, "sum:", ols_weights.sum())
print("NW  weights:", nw_weights[0], "sum:", nw_weights.sum())
print(f"witness predictions: OLS {ols_weights @ witness_values[:, 0]:.4f}, "
      f"NW {nw_pred.item():.4f}; target range [0, 1]")

OLS weights: [-0.1   0.05  0.2   0.35  0.5 ] sum: 1.0
NW  weights: [0.001  0.0206 0.152  0.4132 0.4132] sum: 1.0
witness predictions: OLS -0.1000, NW 0.0010; target range [0, 1]

Which kernels make a convex average?

Only nonnegative smoothing kernels make this convex average. A Gaussian gives positive weights; a compact-support kernel may give zeros. Other objects also bear the name kernel. Equation 14.3 needs nonnegative numerators and a positive denominator.

14.3 Chapter 2’s scores-to-weights machine

With a strictly positive kernel, take the logarithm of each affinity and apply softmax. This gives the normalized weights:

\[ \alpha_i(\vect{q}) = \frac{\kappa_h(\vect{q},\vect{k}_i)} {\sum_j\kappa_h(\vect{q},\vect{k}_j)} = \operatorname{softmax}_i\!\left( \log\kappa_h(\vect{q},\vect{k}_i)\right). \tag{14.5}\]

The score is the log-kernel. With the Gaussian,

\[ a_i(\vect{q}) = -\frac{\norm{\vect{q}-\vect{k}_i}^2}{2h^2}. \]

This also makes the temperature connection exact. If the unscaled score is \(s_i=-\norm{\vect{q}-\vect{k}_i}^2\), then

\[ \vect{\alpha}(\vect{q}) = \operatorname{softmax}\!\left(\frac{\vect{s}}{\tau}\right), \qquad \tau=2h^2. \tag{14.6}\]

Small bandwidth \(\rightarrow\) low temperature \(\rightarrow\) sharp lookup. Large bandwidth \(\rightarrow\) high temperature \(\rightarrow\) diffuse averaging. The factor \(2\) belongs to this score convention; if we had defined \(s_i=-\norm{\vect{q}-\vect{k}_i}^2/2\), the temperature would be \(h^2\) instead.

Notice the implementation never computes \(\kappa_h\) and then takes its logarithm. It computes the log-score directly, subtracts the largest score in each row, and only then exponentiates. That is Chapter 2’s stable-softmax hygiene, applied to kernel scores.

14.4 Bandwidth: bias versus variance

The three-point example makes a narrow bandwidth look wonderful. That is not evidence that narrower is better. A tiny neighborhood follows every noisy bump; a huge neighborhood washes real structure away. Chapter 1 called this bias versus variance. Here the bandwidth controls the trade.

Because this is a synthetic laboratory, we can do something ordinary observed data does not permit: hold 60 key locations fixed, redraw the response noise 1,000 times, and compare every prediction with the known function

\[ f(x)=\sin(1.5x)+0.3\cos(4x), \qquad y_i=f(x_i)+\epsilon_i, \qquad \epsilon_i\sim\mathcal{N}(0,0.25^2). \]

We generate the key locations as seeded uniform draws on \([-3,3]\) and evaluate the predictions at 241 query locations on \([-2.9,2.9]\), none of which shares an exact coordinate with a key. Holding the design fixed lets us separate squared bias from variance over independent noise draws. The population identity

\[ \operatorname{MSE}=\operatorname{bias}^2+\operatorname{variance} \]

has no irreducible-noise term because we evaluate against the noiseless \(f\); predicting a fresh noisy response would add the irreducible variance \(0.25^2\) to every row. In our finite 1,000-world ensemble, the analogous empirical decomposition closes to floating-point roundoff.

Define the reusable truth helper.
Draw 60 sorted keys, one noisy sample, and 1,000 noise redraws, and fix 241 queries and eight bandwidths.
Fixed-design bias/variance sweep over 1,000 noisy worlds.
Print each bandwidth’s 2h^2, squared bias, variance, and MSE, then the one with the lowest MSE.

# [1]
def truth(x: np.ndarray) -> np.ndarray:
    return np.sin(1.5 * x) + 0.3 * np.cos(4.0 * x)

# [2]
data_rng = np.random.default_rng(605012)
keys = np.sort(data_rng.uniform(-3.0, 3.0, size=60))
observed_rng = np.random.default_rng(605013)
observed_values = truth(keys) + observed_rng.normal(0.0, 0.25, size=keys.size)
queries = np.linspace(-2.9, 2.9, 241)

mc_rng = np.random.default_rng(605014)
responses = truth(keys)[None, :] + mc_rng.normal(
    0.0, 0.25, size=(1_000, keys.size)
)
bandwidths = np.array([0.05, 0.08, 0.12, 0.18, 0.28, 0.42, 0.65, 1.00])

rows = []
# [3]
for bandwidth in bandwidths:
    _, weights = gaussian_attention(
        queries[:, None], keys[:, None], observed_values[:, None], bandwidth
    )
    predictions = responses @ weights.T                 # (world, query)
    mean_prediction = predictions.mean(axis=0)
    squared_bias = np.mean((mean_prediction - truth(queries)) ** 2)
    variance = np.mean(np.var(predictions, axis=0))
    mse = np.mean((predictions - truth(queries)[None, :]) ** 2)
    rows.append((bandwidth, 2 * bandwidth**2, squared_bias, variance, mse))

# [4]
print(" h     2h^2    bias^2  variance     MSE")
for row in rows:
    print(f"{row[0]:.2f}   {row[1]:.4f}   {row[2]:.4f}   "
          f"{row[3]:.4f}   {row[4]:.4f}")
best = rows[int(np.argmin([row[-1] for row in rows]))]
print(f"grid minimum: h={best[0]:.2f}, MSE={best[-1]:.4f}")

 h     2h^2    bias^2  variance     MSE
0.05   0.0050   0.0119   0.0343   0.0462
0.08   0.0128   0.0101   0.0247   0.0348
0.12   0.0288   0.0090   0.0169   0.0259
0.18   0.0648   0.0127   0.0109   0.0236
0.28   0.1568   0.0304   0.0069   0.0373
0.42   0.3528   0.0747   0.0047   0.0794
0.65   0.8450   0.1721   0.0031   0.1753
1.00   2.0000   0.3142   0.0022   0.3164
grid minimum: h=0.18, MSE=0.0236

The variance falls at every step, from \(0.0343\) at \(h=0.05\) to \(0.0022\) at \(h=1.00\). Bias squared is not perfectly monotone at the sharp end of this sparse, random design; it bottoms at \(0.0090\) for \(h=0.12\), then rises to \(0.3142\) at \(h=1.00\). Their sum is the U-shaped quantity we care about. On this predeclared grid, \(h=0.18\) has the smallest MSE, \(0.0236\).

One noisy world shows what those numbers mean:

Noisy observations surround a smooth black truth curve. The narrow-bandwidth fit wiggles around local noise, the middle fit follows the truth most closely, and the broad fit smooths away smaller oscillations. — Figure 14.2: One fixed noisy sample under three bandwidths. The narrow rule follows local noise, the middle rule tracks the signal, and the broad rule erases the function’s smaller-scale structure. The blue bandwidth is the grid minimum in the 1,000-world audit.

Selecting \(h\) when the truth is hidden

We know the black truth curve only because we generated the data ourselves. On real data, choose the bandwidth using a validation set, then report once on a sealed test set. Training error alone will usually reward an extremely narrow kernel for copying each noisy target back to itself. The generalization protocol from Chapter 6 still applies even when nothing is trained by gradient descent.

14.5 From one query to an attention matrix

Now write the operator for many queries at once, in the vocabulary that the rest of Part IV will keep. Let

\[ \matr{Q}\in\mathbb{R}^{m\times d},\qquad \matr{K}\in\mathbb{R}^{n\times d},\qquad \matr{V}\in\mathbb{R}^{n\times p} \]

hold \(m\) queries, \(n\) keys, and one \(p\)-dimensional value per key. The score and weight matrices are

\[ S_{ri}=-\frac{\norm{\vect{q}_r-\vect{k}_i}^2}{2h^2}, \qquad \matr{A}=\operatorname{rowsoftmax}(\matr{S}) \in\mathbb{R}^{m\times n}. \tag{14.7}\]

Every row of \(\matr{A}\) sums to one; columns generally do not. The pooled outputs are one matrix multiplication,

\[ \operatorname{Attention}_h(\matr{Q},\matr{K},\matr{V}) =\matr{A}\matr{V}\in\mathbb{R}^{m\times p}. \tag{14.8}\]

This is the book’s first attention-weight matrix. It is not a learned alignment. It is the allocation imposed by Euclidean distance and the chosen bandwidth.

Form the first fixed attention map and audit its shapes and row sums.

# [1]
map_queries = np.linspace(-2.9, 2.9, 31)
map_predictions, attention_map = gaussian_attention(
    map_queries[:, None], keys[:, None], observed_values[:, None], best[0]
)
print(f"Q {map_queries[:, None].shape}, K {keys[:, None].shape}, "
      f"V {observed_values[:, None].shape}, A {attention_map.shape}, "
      f"AV {map_predictions.shape}")
print("maximum row-sum error:",
      f"{np.max(np.abs(attention_map.sum(axis=1) - 1)):.2e}")

Q (31, 1), K (60, 1), V (60, 1), A (31, 60), AV (31, 1)
maximum row-sum error: 2.22e-16

The left panel shows variance falling as bandwidth grows, bias eventually rising, and their summed error reaching an interior minimum. The right heatmap has a narrow bright diagonal band where each query assigns weight to nearby sorted keys. — Figure 14.3: Left: the fixed-design bias–variance audit. Right: the same Gaussian rule written as a 31-query by 60-key attention matrix at \(h=0.18\). Bright cells receive more of a row’s unit mass; the white trace marks the \(k=q\) locus in sorted-key coordinates. The map is local and interpretable, but it was specified rather than learned.

Read the bright band as a moving magnifying glass. Each query row reallocates one unit of weight across the stored keys. Where keys are dense, that unit is divided among several neighbors. Where keys are sparse, one or two slots carry most of it. Changing \(h\) changes the width of the band, exactly as changing temperature changes the spread of a softmax distribution.

There is one numerical trap. A Gaussian is mathematically positive everywhere, but a distant query and small bandwidth can make every directly exponentiated kernel value underflow to zero. Normalizing those zeros produces 0/0. The log-score version still has a largest entry, so the shifted softmax remains finite. Appendix C collects this with the other stable algorithms.

Weight the 60 keys by normalizing the kernel and by row_softmax of its log, for queries 0.37 and 1,000.
Verify the log-kernel identity and survive underflow.

# [1]
identity_query, identity_h = 0.37, 0.28
log_kernel = -0.5 * ((identity_query - keys) / identity_h) ** 2
direct_kernel = np.exp(log_kernel)
direct_weights = direct_kernel / direct_kernel.sum()
softmax_weights = row_softmax(log_kernel[None, :])[0]

far_log_kernel = -0.5 * ((1_000.0 - keys) / 0.1) ** 2
far_kernel = np.exp(far_log_kernel)
with np.errstate(divide="ignore", invalid="ignore"):
    naive_far_weights = far_kernel / far_kernel.sum()
stable_far_weights = row_softmax(far_log_kernel[None, :])[0]

# [2]
print("max |normalize(K)-softmax(log K)|:",
      f"{np.max(np.abs(direct_weights - softmax_weights)):.2e}")
print(f"far query: naive denominator={far_kernel.sum():.1f}, "
      f"all finite={np.all(np.isfinite(naive_far_weights))}")
print(f"stable: sum={stable_far_weights.sum():.1f}, "
      f"nearest key={keys[np.argmax(stable_far_weights)]:.4f}")

max |normalize(K)-softmax(log K)|: 2.78e-17
far query: naive denominator=0.0, all finite=False
stable: sum=1.0, nearest key=2.9886

Trap: similarity weights are not uncertainty probabilities

The entries of \(\matr{A}\) are probability-like allocations: nonnegative and row-normalized. They are not calibrated probabilities that the corresponding key is “correct,” and a sharp row is not automatically confidence. It may simply mean the bandwidth is small.

14.6 Return the fixed operator to the memory bank

Map the regression names onto Chapter 13’s memory bank; each does a different job:

job	kernel regression	sequence memory preview
query	location \(\vect{q}\) where we need an answer	what the decoder needs now
keys	stored coordinates \(\vect{k}_i\)	descriptors used to match encoder memory slots
values	observed responses \(\vect{v}_i\)	content retrieved from those slots
weights	normalized fixed similarities \(\alpha_i(\vect{q})\)	how much each slot contributes
output	\(\sum_i\alpha_i\vect{v}_i\)	a query-specific context vector

The opening example fills every kernel-regression role. The query is the location \(q=3.5\). The keys are the stored locations 1, 3, and 5, and the values are their responses 1.5, 2.8, and 1.8. Gaussian similarity at \(h=0.6\), normalized, gives the weights \((0.0002, 0.9413, 0.0585)\), and the prediction is the mixture \(2.7412\). The sequence-memory roles have their own instance in Chapter 13’s date task. When the decoder is about to write the year, its current state is the query; the encoder states at the source characters serve as keys to match and as values to read; and useful weights concentrate on the last four source characters, which spell the year. Chapter 15 measures how much weight lands there.

Queries specify what is needed; keys provide what to compare against; values provide what to retrieve. In the simplest sequence model, the same encoder state can serve as both key and value, just as our one-dimensional example stores a location and its response together. Later, learned projections will let keys and values carry different representations.

Chapter 1 also used one more similarity primitive, the dot product:

\[ \vect{q}^\top\vect{k} =\norm{\vect{q}}\,\norm{\vect{k}}\cos\theta. \]

It is pure angular similarity only when the norms are controlled. If \(\vect{q}=(1,0)\), \(\vect{k}_1=(1,0)\), and \(\vect{k}_2=(2,0)\), both keys point in the same direction as the query, but their dot products are 1 and 2. Magnitude participates. That is not a defect: the score defines what “similar” means.

Our score today is a hand-written negative squared Euclidean distance, and \(h\) is selected rather than learned. The operator stores the data, computes a fixed map, and pools the values. It has attention’s form, a normalized, similarity-weighted average of the values, but it has no learned compatibility function.

At the end of Chapter 8 we had a fixed image kernel and asked what would happen if the filter became learnable. Chapter 9 answered with a CNN. Our similarity rule is fixed in the same way:

What if the similarity itself were learnable?

The memory bank is ready. Query, keys, values, scores, row-wise softmax, and weighted pooling are ready. Only one hand-designed piece remains: the rule that decides whether a query and key belong together. Chapter 15 makes that rule learnable, then returns to Chapter 13’s date task so the decoder can show us where it looked.

Check yourself

Close the book for one minute and reconstruct the weighted prediction.

How do scores become nonnegative weights that sum to one?
What does bandwidth change in a kernel smoother?
Which component is still hand designed at the end of this chapter?

14.7 Okay, so: attention is normalized memory mixing

Chapter 1’s linear prediction already mixed training targets. OLS gives \(\hat f(\vect{q})=\vect{\lambda}(\vect{q})^\top\vect{y}\), but its weights may be signed, so the prediction can leave the targets’ convex hull.
Nadaraya–Watson turns locality into a convex mixture. A nonnegative kernel scores query–key similarity; normalization makes the weights sum to one; their weighted values form the prediction.
Bandwidth determines the softmax temperature. For Gaussian distance scores, \(\tau=2h^2\): narrow is sharp and variable, wide is diffuse and biased. In the fixed-design sweep over the 1,000 noisy worlds, the lowest MSE on the grid was \(0.0236\), at \(h=0.18\).
Chapter 2’s softmax is the normalization exactly. Positive-kernel weights equal softmax(log kernel). Computing log-scores directly and shifting each row avoids the distant-query 0/0 trap.
Fixed attention is \(\matr{A}\matr{V}\). Queries score keys, row-softmax produces an \(m\times n\) allocation matrix, and that matrix pools the values. The algebra is complete; only the score is still hand designed.

Connection to test-time regression

Work from 2024–26 uses a forward pass as test-time regression: the prefix supplies key–value training pairs, the current token supplies a query, and a memory operator fits or updates a predictor before answering. Test-Time Training, the older Delta-style fast weights, Titans, and MesaNet differ in which part they change: the learned views of a token, the weighting of its history, the model or regularizer, or the online solver. Wang and colleagues’ Beyond Test-Time Memory changes the objective itself, from regression to finite-horizon optimal control.

That shared regression perspective is a design language, not a claim that the architectures are equivalent. Softmax attention performs a query-dependent local fit while retaining its key–value rows. Other methods compress the history into a fixed-size parametric state and no longer retain those rows. Chapter 17 develops this regression perspective by varying the kernel, predictor class, and update rule, then comparing the information retained and the computation required.

Sources and further reading

Nadaraya, On Estimating Regression: the 1964 two-page paper introducing one side of the estimator now bearing Nadaraya and Watson’s names.
Watson, Smooth Regression Analysis: the independent 1964 development of smooth regression by local weighting.
Bahdanau, Cho, & Bengio, Neural Machine Translation by Jointly Learning to Align and Translate: the neural-attention step that replaces one fixed context with learned soft alignment; Chapter 15 implements that step.
Wang, Shi, & Fox, Test-Time Regression: A Unifying Framework for Designing Sequence Models with Associative Memory (2025): develops test-time regression as a common design language for sequence memory.
Sun et al., Learning to (Learn at Test Time): RNNs with Expressive Hidden States (2024 preprint; 2025 proceedings): treats a sequence model’s hidden state as a learner updated by a self-supervised objective during the forward pass.
Schlag, Irie, & Schmidhuber, Linear Transformers Are Secretly Fast Weight Programmers (2021): connects linear attention to fast-weight updates and the delta rule.
Behrouz et al., Titans: Learning to Memorize at Test Time (2025): adds a learned long-term memory updated at test time.
von Oswald et al., MesaNet: Sequence Modeling by Locally Optimal Test-Time Training (2025): derives sequence layers from locally solved test-time objectives.
Wang et al., Beyond Test-Time Memory: State-Space Optimal Control for LLM Reasoning (2026): extends the regression framing toward test-time control; the epilogue treats that proposal as a frontier bet rather than a settled equivalence.

Exercises

(Pencil.) Prove that Equation 14.3 is a convex combination when \(\kappa_h\ge0\) and its denominator is positive. Then prove that a scalar prediction lies in \([\min_i v_i,\max_i v_i]\). Which step fails for the OLS weights in Equation 14.1?
(Pencil.) For keys \((-2,-1,0,1,2)\), derive the five OLS target weights at \(q=1.5\) without using code. Verify that they sum to one, identify the negative weight, and construct a second target vector for which OLS leaves the target range.
(Code.) Split a fresh noisy sample from this chapter’s function into training, validation, and sealed test sets. Select \(h\) on validation data, report the test MSE once, and compare your selected curve with the synthetic-truth grid minimum above. Explain why the two procedures need not select the same bandwidth.
(Pencil.) (a) Derive Equation 14.6. (Code.) (b) Use keys \([0,0,1]\), query 0, and successively smaller positive bandwidths. What weights does stable softmax assign to the exact nearest-key tie, and why is \(h=0\) still invalid?
(Code.) Extend gaussian_attention to two-dimensional keys whose coordinates have very different units. Compare the attention map before and after standardizing each feature using training-set statistics. What does this reveal about the metric hidden inside a fixed kernel?
(Pencil.) Differentiate Equation 14.4 with respect to the vector \(\vect{c}\). Solve the stationarity equation, then substitute \(\kappa(\vect{q},\vect{k}_i)=\exp(\vect{q}^{\top}\vect{k}_i/\sqrt d)\). Why would the stronger claim “an unrestricted \(M\) with no regularizer must average” be false?
(Code.) Return to Figure 14.1 and parameterize a positive inverse bandwidth as \(w=\operatorname{softplus}(\rho)\) in the score \(-w(q-k_i)^2\). Fit the scalar \(\rho\) by SGD on training queries, choose the stopping point on held-out queries, and evaluate once on a sealed test grid. Then replace \(w\) by one positive \(w_i\) per key and repeat under matched splits, updates, and seed. Plot the learned weights and train/validation errors. Named wrong answer: “the lower-training-loss bandwidth is the better estimator”. The sharp limit can interpolate the stored targets while degrading held-out prediction.