Toolkit: attention and transformers

Why this is in a driving course

Every learned component in the rest of the course is built from the same block: sensor encoders that fuse cameras, predictors that relate agents, the language model that writes Alpamayo-R1's reasoning and the expert that reads its cache. This page gives you the working knowledge in one sitting. The full treatment, with exercises, is the Transformers module.

Attention is a soft lookup

Think of a dictionary lookup. You have a query; each stored item has a key and a value; you return the value whose key matches. Attention does the same thing softly: it scores the query against every key, turns the scores into weights that sum to one, and returns the weighted average of the values.

In a transformer, every token produces all three from its own vector, through three learned matrices. Each token then asks "which other tokens are relevant to me?" and mixes their information into its own representation. The $\sqrt{d_k}$ keeps the scores from growing with the vector size, so the softmax does not saturate.

The whole computation is four lines:

def scaled_dot_product_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    causal: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Computes softmax(q k^T / sqrt(d_k)) v.

    Args:
        q: Queries, shape (batch, heads, n, d_k).
        k: Keys, shape (batch, heads, m, d_k).
        v: Values, shape (batch, heads, m, d_v).
        causal: If True, position i may only attend to positions <= i.

    Returns:
        The attended values (batch, heads, n, d_v) and the weights (batch, heads, n, m).
    """
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
    if causal:
        future = torch.triu(torch.ones_like(scores, dtype=torch.bool), diagonal=1)
        scores = scores.masked_fill(future, float("-inf"))
    weights = scores.softmax(dim=-1)
    return weights @ v, weights

Two properties matter for driving models.

• Cost grows with the square of the sequence. The score matrix is $n \times n$. Doubling the number of camera tokens quadruples attention's cost, which is why image tokens get compressed (the Sensor encoders and Deploying a driving VLA modules).
• Attention has no notion of order. Shuffle the tokens and each output shuffles the same way. Order has to be added as a positional encoding. That property is also what makes a set of agents or map polylines easy to encode: their order should not matter.

Heads and blocks

One attention head learns one kind of relevance. Multi-head attention runs several heads in parallel, each in a smaller space, and concatenates their outputs, so one head can track the car in front while another tracks the lane. A transformer block is attention followed by a small MLP applied to each token, each wrapped with a residual connection and a normalisation. Models stack tens of blocks.

There are three ways to arrange them, and the course uses all three:

Generating one token at a time

A language model writes text by repeating one step: run the transformer on everything so far, read a probability for every possible next token, pick one, append it, repeat. A causal mask makes sure each position can only see earlier positions, so training on whole sentences teaches exactly this one-step-ahead prediction.

The key-value cache

When the model generates token 101, the keys and values of tokens 1 to 100 are the same as they were one step ago. The key-value cache stores them, so each step only computes the new token's query, key and value and reads the rest.

Hold on to this idea. In Alpamayo-R1 the cache is more than a speed-up: it is the interface between the reasoning model and the action expert, which attends to the keys and values the reasoning produced (Reading the Alpamayo-R1 code). It is also why decoding speed is limited by memory: every step reads the whole cache and all the weights.

Where this shows up in the course

• Sensor encoders: BEV queries attending to camera features, tokens per camera.
• Prediction: agents attending to each other and to the map.
• VLAs and reasoning: a decoder writing text, the cost of each reasoning token.
• Deployment: the quadratic cost, the cache and memory bandwidth.

Check yourself

1. In your own words, what do the query, key and value of a token each represent?
2. Why does attention's cost grow with the square of the number of tokens?
3. A set of map polylines is fed to attention in a random order. Does the output depend on the order?
4. What does a causal mask prevent, and why is it needed to train a model that writes text?
5. What does the key-value cache store, and what does it save at each generation step?

Quiz

Answer each question; feedback appears as soon as you choose.

Go deeper

The Transformers module builds attention, heads, positions and a tiny GPT step by step, with an exercise and the original paper's details.