Attention as Routing

Attention is routing: at each position, information from other positions is weighted, summed, and delivered according to a learned compatibility score.

The word "routing" is doing precise work. Attention does not add information; it moves information between positions. The query decides what to look for, the key decides what is findable, and the value carries the payload. When a head attends sharply to a previous token, that token's value is routed to the current position.

This note builds self-attention from a ten-line lookup, derives why the naïve version fails at scale, and ends with the causal mask and the head counts that production models use.

The naïve version: a soft dictionary lookup#

The smallest useful formulation treats each token as a record and each output position as a query against every record. Start with plain Python:

import math
 
def softmax(scores):
    m = max(scores)
    exps = [math.exp(s - m) for s in scores]
    total = sum(exps)
    return [e / total for e in exps]
 
def attend_single(query, keys, values):
    d = len(query)
    scores = [sum(q * k for q, k in zip(query, key)) / math.sqrt(d)
              for key in keys]
    weights = softmax(scores)
    out = [sum(w * v[i] for w, v in zip(weights, values))
           for i in range(len(values[0]))]
    return out

The structure is three steps. First, compatibility: the dot product $q \cdot k$ scores how well a query matches a key. Second, normalisation: softmax turns the scores into a distribution over positions. Third, delivery: the output is a weighted sum of values. Nothing here creates content. The math is a convex combination of existing vectors.

The scaling by $\sqrt{d}$ matters more than it looks. Consider the variance. For a dot product of two random vectors whose components have variance $\sigma^2$:

As the dimension $d$ grows, the magnitude of raw dot products grows with $\sqrt{d}$. Unscaled logits push softmax into its saturated corners, where gradients vanish. Dividing by $\sqrt{d}$ keeps the pre-softmax scores at roughly unit variance regardless of dimension. That single division is the difference between a trainable mechanism and one whose gradients collapse after a few hundred dimensions.

Transpose the computation and the same function serves every position at once:

def attend_all(queries, keys, values):
    """queries: (n, d), keys: (n, d), values: (n, dv) -> (n, dv)"""
    d = len(queries[0])
    outputs = []
    for query in queries:
        outputs.append(attend_single(query, keys, values))
    return outputs

The loop is the bottleneck that the matrix form removes:

where $Q, K \in \mathbb{R}^{n \times d}$ are stacked queries and keys and $V \in \mathbb{R}^{n \times d_v}$ holds values. The attention matrix $A = \operatorname{softmax}(QK^\top / \sqrt{d_k})$ has shape $(n, n)$: entry $A_{ij}$ is the route weight from position $j$ to position $i$. Training the router is training the key and query spaces, because $A$ is a function of their geometry alone.

First principles: why positions need a router#

A feed-forward block processes each position independently. Token $i$'s output is a function of token $i$'s input. That architecture cannot, in one layer, let the representation of "it" depend on the noun that "it" refers to three sentences back. Recurrence solves the problem by threading state forward through time, at the cost of serial computation and a fixed-size bottleneck.

Attention threads information by route rather than by time. No position is special in the routing math: every query can reach every key in one step. The price is that the router must encode positional information explicitly, because the attention operation itself is permutation-invariant.

The permutation invariance is easy to see in the code above. Shuffle the rows of $K$ and $V$ together and the outputs shuffle identically; shuffle only one of them and the mechanism breaks silently. Attention alone has no notion of "before" and "after". Position has to be added as data.

The standard fix is sinusoidal or learned positional embeddings added to the input embeddings before any layer runs:

import math
 
def positional_embedding(pos, d):
    pe = [0.0] * d
    for i in range(0, d, 2):
        divisor = math.exp(i * math.log(10000.0) / d)
        pe[i] = math.sin(pos / divisor)
        if i + 1 < d:
            pe[i + 1] = math.cos(pos / divisor)
    return pe
 
def add_positions(embeddings):
    return [[e + p for e, p in zip(emb, positional_embedding(i, len(emb)))]
            for i, emb in enumerate(embeddings)]

Each position receives a distinct vector that a linear layer can read as distance. The sinusoidal version has a useful inductive bias: the representation of position $p + k$ is a linear function of the representation of $p$, shifted by a fixed rotation that depends only on $k$. Relative offsets are learnable with one step.

Bottleneck one: the quadratic bill#

The naïve implementation costs $O(n^2 d)$ time and $O(n^2)$ memory for the attention matrix. At $n = 4096$ tokens the matrix holds about $16.7$ million entries per head; at $n = 32768$ it holds over a billion. Sequence length is the dominant cost of the architecture.

Three structural fixes follow from the same matrix, each trading exactness for speed.

Causal masking is the cheapest and non-negotiable for generation. The query at position $i$ must not see keys from positions $j > i$, or training leaks the future into the prediction:

def causal_attention(queries, keys, values):
    d = len(queries[0])
    n = len(queries)
    outputs = []
    for i, query in enumerate(queries):
        scores = [sum(q * k for q, k in zip(query, keys[j])) / math.sqrt(d)
                  for j in range(i + 1)]
        weights = softmax(scores)
        outputs.append([sum(w * values[j][t] for j, w in enumerate(weights))
                        for t in range(len(values[0]))])
    return outputs

In matrix form the fix is a mask added before softmax, with $-\infty$ in the upper triangle. Entries set to $-\infty$ receive zero weight after exponentiation. The mask costs nothing in memory traffic beyond the comparison; the savings appear at decode time, where the key-value cache grows linearly with generated length and each step attends over the whole prefix.

KV caching. Training computes $K$ and $V$ once per sequence. Autoregressive generation recomputes them for the entire prefix at every new token in the naïve version:

class KVCache:
    def __init__(self, d):
        self.d = d
        self.keys = []
        self.values = []
 
    def step(self, query, key, value):
        self.keys.append(key)
        self.values.append(value)
        d = self.d
        scores = [sum(q * k for q, k in zip(query, k)) / math.sqrt(d)
                  for k in self.keys]
        weights = softmax(scores)
        return [sum(w * v[t] for w, v in zip(weights, self.values))
                for t in range(len(value))]

A step that attended over the full cache costs $O(T d)$ for $T$ cached tokens. Generating $T$ tokens without a cache costs $O(T^2 d)$ total, because every step re-encodes the whole history. The cache converts a quadratic rebuild into linear extension. The memory cost is $2 T d$ per layer per head, which becomes the binding constraint at long context.

Multi-head split. A single $d$-dimensional router has one compatibility function. Splitting the projections into $h$ heads gives each head a $d/h$-dimensional space and its own attention matrix:

def multihead_attention(queries, keys, values, W_q, W_k, W_v, W_o):
    """W_*: list of projection matrices, one per head."""
    heads = []
    for w_q, w_k, w_v in zip(W_q, W_k, W_v):
        q = [dot(w_q, row) for row in queries]
        k = [dot(w_k, row) for row in keys]
        v = [dot(w_v, row) for row in values]
        heads.append(attend_all(q, k, v))
    concat = list(zip(*[h for head in heads for h in head]))
    return [dot(W_o, row) for row in concat]

The compute cost is unchanged — the total projection dimension is the same — but the per-head dot products live in lower-dimensional spaces, and each head can specialize. An optional output projection $W_o$ mixes the heads back into the model dimension.

Bottleneck two: quadratic again, at long context#

Causal masking and caching tame the constant factor, not the $O(n^2)$ growth in training cost. The ProbSparse mechanism in Informer (arXiv:2012.07436) attacks the growth by noting that most attention rows are close to uniform: only a few keys receive most of the weight from any given query. Scoring each query by a sparsity measurement — the maximum of its scaled scores minus their mean, a proxy for its divergence from uniform — and keeping the top $k \sim c \log n$ queries reduces time and memory toward $O(n \log n)$ while preserving the accuracy of full attention on long time-series.

def logsumexp(scores):
    if not scores:
        return float('-inf')
    m = max(scores)
    return m + math.log(sum(math.exp(s - m) for s in scores))
 
def query_measure(query, keys, d):
    scores = [sum(q * k for q, k in zip(query, key)) / math.sqrt(d)
              for key in keys]
    mean = sum(scores) / len(scores)
    return logsumexp(scores) - mean
 
def top_queries(queries, keys, d, u):
    ranked = sorted(range(len(queries)),
                    key=lambda i: -query_measure(queries[i], keys, d))
    return ranked[:u]
 
def probsparse_attention(queries, keys, values, d, u):
    n = len(queries)
    outputs = [[0.0] * len(values[0]) for _ in range(n)]
    active = set(top_queries(queries, keys, d, u))
    for i in range(n):
        if i in active:
            scores = [sum(q * k for q, k in zip(queries[i], keys[j])) / math.sqrt(d)
                      for j in range(n)]
        else:
            scores = [0.0] * n
        weights = softmax(scores)
        outputs[i] = [sum(w * values[j][t] for j, w in enumerate(weights))
                      for t in range(len(values[0]))]
    return outputs

Queries outside the top $u$ receive the lazy approximation — a uniform distribution over keys, equivalent to the mean of $V$. The selection step is the only new cost, and its measurement can be computed on a subsample of keys rather than the full matrix. The mechanism preserves full softmax for the queries where attention peaks, and replaces the flat region with its expectation.

Bottleneck three: one head cannot be everywhere#

A single head learns one compatibility function. Local recency, long-range copying, subject tracking, and delimiter handling impose incompatible constraints on one geometry. Multiple heads provide parallel routers at identical layer cost.

The specialization is observable, not merely hypothesized. Ablating a single head in a trained model often removes a specific behaviour while leaving overall accuracy nearly intact, which indicates that the function was carried by that route. The canonical circuit is the induction head, a two-layer motif that implements pattern completion.

The circuit has two steps. In the first layer, a previous-token head moves each token's identity forward by one position:

def previous_token_head(tokens):
    """Each position receives the identity of the token before it."""
    return [None] + tokens[:-1]

In the second layer, the induction head matches the current token against the carried identity at every position and routes the following token's value forward:

def induction_head(carried_tokens, tokens):
    """At position i, find position j where carried[j] == tokens[i],
    then route tokens[j + 1] to position i."""
    result = []
    for i, tok in enumerate(tokens):
        match = None
        for j in range(len(carried_tokens)):
            if carried_tokens[j] == tok and j + 1 < len(tokens):
                match = tokens[j + 1]
        result.append(match)
    return result

The pattern reproduced is "…a b… a → b": find a prior occurrence of the current token and copy what followed it. Training on sequences with repeated structure rewards this circuit, and it appears reliably in small models trained from random initialisation.

What routing cannot do#

Attention moves information between positions. It does not add information at a position. The additions, transformations, and feature engineering happen in the feed-forward blocks between attention layers, which operate position-wise:

def feedforward(x, W1, W2, bias1, bias2):
    hidden = [max(0.0, sum(w * v for w, v in zip(row, x)) + b)
              for row, b in zip(W1, bias1)]
    return [sum(w * v for w, v in zip(row, hidden)) + b
            for row, b in zip(W2, bias2)]

Placing a position-wise block between two routing layers lets the network compute new features after each round of information movement. Attention alone, stacked without the interlayer transformation, composes convex combinations of convex combinations and quickly saturates.

Nor is routing selective in the sense of a hard switch. Attention is a softmax — a distribution, not a selection. Soft weights mean a low-attention position still contributes a little of its value, and the mechanism stays differentiable, which is what allows gradient descent to learn the routing.

What to measure#

Attention can be read directly, and should be before it is tuned. The following checks separate a healthy router from a degenerate one:

def attention_entropy(weights):
    return -sum(w * math.log(w + 1e-12) for w in weights)
 
def check_attention(matrix):
    for i, row in enumerate(matrix):
        entropy = attention_entropy(row)
        top = max(range(len(row)), key=lambda j: row[j])
        print(f"row {i}: entropy {entropy:.3f}, argmax position {top}, "
              f"max weight {row[top]:.3f}")

The checklist for reading a trained model:

• Histogram the per-row entropy. Rows stuck at high entropy indicate a head observing but not routing; rows collapsed to a single position indicate a head that has become an identity map.
• Check which heads survive ablation. A model whose accuracy is unchanged when half its heads are removed has allocated its routing to the remaining half, and the redundant heads can be pruned.
• Verify the causal mask in code and in the matrix. A shifted mask leaks the target token and produces training loss that looks excellent and generalises to nothing.
• Confirm the $\sqrt{d_k}$ scaling is present. Removing it produces saturated softmax weights and vanished gradients at high dimension, with symptoms indistinguishable from a learning-rate problem.
• Measure the KV cache footprint during generation. At long context the cache, not the parameters, dominates memory, and a deployable model is the one whose cache plus weights fit the device.

A checklist#

• Use the exact form $\operatorname{softmax}(QK^\top / \sqrt{d_k}) V$ and verify the scaling factor is inside the softmax, not after it.
• Apply a causal mask for any autoregressive task, and confirm that future positions receive exactly zero weight.
• Cache keys and values during generation; recomputing them per token is quadratic and wasteful.
• Split the model into multiple heads, then ablate to find which heads carry which function.
• Add positional information explicitly; attention has none by default.
• Keep the feed-forward interleaving; routing without position-wise transformation saturates after a few layers.
• For long sequences, consider sparse or ProbSparse attention so that the $O(n^2)$ term does not dominate the training budget.

Attention is a learned traffic map from position to position. The queries state demands, the keys state offers, and the values carry the freight. Heads specialize because a single routing structure cannot serve every demand at once, and circuits like induction heads emerge because copying from a previous occurrence is one of the most valuable routes a language model can learn.