Intuition: Attention as Soft Dictionary Lookup

The cleanest single sentence about attention: it is a soft, differentiable dictionary lookup. A regular Python dictionary maps an exact key to an exact value: hit or miss. Attention generalizes this. Each token presents a query, every token in the context presents a key and a value, and the output is a weighted blend of all values, weighted by query-key similarity.

This single move buys three things at once. It's differentiable (so gradients flow through it during training), it's parallel (the n × n similarity matrix can be computed in one matmul), and it's relational (the same operation that lets a verb find its subject also lets a coreferent pronoun find its antecedent).

Three components, three roles

Query (Q): "What am I looking for right now?" — generated from the token doing the asking.
Key (K): "What kind of information am I?" — generated from each token being looked at.
Value (V): "What information do I provide if attended to?" — also generated from each token being looked at, but allowed to differ from K.

Why K and V can differ: the criterion for matching ("is this token the one I want?") may be different from the information you want to extract once you've decided to look. K is the lookup signature; V is the payload. Decoupling them gives the model more expressive power.

Code

Attention as a soft dictionary·python

import torch, torch.nn.functional as F

def soft_lookup(query, keys, values, temperature=1.0):
    # query: (d,)
    # keys, values: (n, d)
    scores = (keys @ query) / temperature       # (n,)
    weights = F.softmax(scores, dim=-1)          # (n,) sum to 1
    return weights @ values                      # (d,)

# Hard dict: 'paris' returns France only.
# Soft dict: 'paris' returns 0.9 * France + 0.05 * Berlin + 0.05 * Rome
# (the 'paris' query embedding is closest to France's key embedding).

Exercise

Implement soft_lookup as in the code snippet, then test it on a tiny synthetic example: 4 'documents' (random vectors) and a query that's 90% similar to document 2 and 10% to document 0. Verify that the output is approximately 0.9 * V[2] + 0.1 * V[0] (after softmax). Then try temperature=0.1 and temperature=10. What changes?

Intuition: Attention as Soft Dictionary Lookup

Three components, three roles

Code

External links

Exercise

Progress

Comments 0