Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Transformers and Neural State-Space Architectures › Worked Example: Attention as Content-Based Lookup

Michael BrenndoerferJune 26, 20262 min read

World Models HandbookAll chapters

Part 5 of 12World-Model Architectures

Chapter 26Transformers and Neural State-Space Architectures › Worked Example: Attention as Content-Based Lookup

Reading section 6 of 10 · Worked Example: Attention as Content-Based Lookup · 2 minSection 6 of 10 · 2 min

Section 6 of 10 · 2 min read

Transformers and Neural State-Space Architectures › Worked Example: Attention as Content-Based LookupOne section · Sections in this chapter6 / 10 ↓

Sections in this chapter10

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Before the full experiment, let us build direct intuition for attention with a tiny hand-computed lookup. Suppose a trajectory contains three tokens, each describing an object with a one-hot identity feature and an associated position value. A query token asks, in effect, "where is the red object?" In a real transformer the query, keys, and values are learned projections, but the hand-designed example uses the same weighted-lookup operation. Training optimizes the projections for an objective; useful semantic alignment must be learned rather than assumed.

In[3]:

Code

import numpy as np

## Query: "looking for red" (third feature)
q = np.array([0.0, 0.0, 1.0])

## Keys: one-hot object identities for blue block, green block, red key
K = np.array(
    [
        [1.0, 0.0, 0.0],  # token 1: blue block
        [0.0, 1.0, 0.0],  # token 2: green block
        [0.0, 0.0, 1.0],  # token 3: red key
    ]
)

## Values: observed positions of each object
V = np.array([[2.0], [7.0], [9.0]])

scores = K @ q / np.sqrt(3)
w = np.exp(scores - scores.max())
w = w / w.sum()
out = w @ V

print("attention weights:", np.round(w, 3))
print("attended output (retrieved position):", np.round(out, 3)[0])

Out[3]:

Console

attention weights: [0.264 0.264 0.471]
attended output (retrieved position): 6.62

With the one-hot scores divided by $\sqrt{3}$ , the matching token receives weight about 0.47 and each other token about 0.26. The output is 6.62, not the red key's stored position 9.0. Increasing the matching score's margin would concentrate the weights toward that value. This example illustrates standard scaled dot-product attention and maximum-subtracted softmax; these details are not universal to every attention variant. A retained token can be accessed at a long lag if useful keys, queries, and values are learned. When it is removed, direct access is lost, although another retained representation may preserve its content. In a fixed stable SSM, an input's contribution propagates through powers of the transition matrix. Its amplitude can decay while usable recall either deteriorates gradually or fails abruptly through finite precision, interference, or decoding thresholds. Selective retention can change that behavior; it is not implemented by the fixed-coefficient SSM below.