Attention and transformers

“Pick up the cup.” A robot that hears this must find the cup among everything on the table, and ignore the rest.

How does a network decide which of its inputs matter for the answer it is computing?

Overview

The running example: a 1 m × 1 m table with four objects of different kinds at random places. A request names one kind, and the network must say where it is.

TokensOne list of numbers per object.AttentionAsk, match, average.LayersTwo questions in a row.ACTCameras in, a chunk of actions out.

The mental model: attention is a soft lookup. A query is compared with every key, and the best matches decide what is read.

Tokens

Each object becomes one list of 10 numbers:

Kind: 8 numbers, a 1 in the place of its kind (cube, ball, cup, pen, box, can, key, toy) and 0 elsewhere.
Position: x and y in metres.

The request is a list of 8 numbers: a 1 in the place of the kind asked for.

Token: one item of a model’s input as its own list of numbers: here an object, in ACT a patch of a camera image or the arm’s joint angles.

A scene is a set of tokens, not one fixed list. It can hold four objects or six, in any order.

Attention

Attention works in three steps. Each token is first turned into new lists of numbers by learned linear layers:

The query says what is being looked for. Here it comes from the request.
Each key says what a token offers, so it can be matched against the query.
Each value is what a token hands over if it is chosen.

Then:

Score every key against the query: multiply them number by number and add up, then divide by \(\sqrt{d}\), where \(d\) is the length of the lists.
Softmax the scores into weights that add up to 1, exactly as for the T-shirt sizes.
Average the values with those weights.

\[\text{attention}(q, K, V) = \text{softmax}\!\left(\frac{q\,K^{\top}}{\sqrt d}\right) V\]

pytorch_attention.py · attention by handOpen source ↗

def attention(query, keys, values):
    """Scaled dot-product attention, by hand: score every key against the query, softmax, average the values."""
    scores = query @ keys.T / math.sqrt(keys.shape[1])
    weights = scores.softmax(dim=-1)
    return weights @ values, weights, scores

Full file · pytorch_attention.py · 148 lines

"""How a network decides which inputs matter: attention, on a table with a few objects.

Run: python3 site/ml/foundations/examples/pytorch_attention.py

Each scene is a 1 m × 1 m table with 4 objects of different kinds at random places. A request names one kind.
Task "find": where is the requested object? Task "nearest": where is the object nearest to it?
Every network trains with Adam.
"""

import math

import torch
from torch import nn

ITEMS = ["cube", "ball", "cup", "pen", "box", "can", "key", "toy"]
KINDS = len(ITEMS)


def scenes(count, objects, generator):
    """Object tokens (count, objects, 10): kind as one-hot (8 numbers), then x and y in metres.
    Requests (count, 8): the requested kind as one-hot. Answers for both tasks: (count, 2) positions."""
    kind = torch.stack([torch.randperm(KINDS, generator=generator)[:objects] for _ in range(count)])
    position = torch.rand(count, objects, 2, generator=generator)
    asked = torch.randint(0, objects, (count,), generator=generator)
    rows = torch.arange(count)
    tokens = torch.cat([nn.functional.one_hot(kind, KINDS).float(), position], dim=2)
    request = nn.functional.one_hot(kind[rows, asked], KINDS).float()
    distance = (position - position[rows, asked][:, None]).norm(dim=2)
    distance[rows, asked] = float("inf")                           # not itself
    return tokens, request, {"find": position[rows, asked], "nearest": position[rows, distance.argmin(dim=1)]}


def attention(query, keys, values):
    """Scaled dot-product attention, by hand: score every key against the query, softmax, average the values."""
    scores = query @ keys.T / math.sqrt(keys.shape[1])
    weights = scores.softmax(dim=-1)
    return weights @ values, weights, scores


class Lookup(nn.Module):
    """One attention layer: the request asks, the objects answer."""
    def __init__(self, width=16):
        super().__init__()
        self.embed_object = nn.Linear(KINDS + 2, width)
        self.embed_request = nn.Linear(KINDS, width)
        self.attend = nn.MultiheadAttention(width, num_heads=1, batch_first=True)
        self.out = nn.Linear(width, 2)

    def forward(self, tokens, request):
        objects = self.embed_object(tokens)                        # (N, objects, 16)
        query = self.embed_request(request)[:, None]               # (N, 1, 16)
        answer, _ = self.attend(query, objects, objects)           # query, keys, values
        return self.out(answer[:, 0])


class Transformer(nn.Module):
    """The request token and the object tokens attend to each other, `layers` times; the request's token answers."""
    def __init__(self, layers=2, width=32, heads=2):
        super().__init__()
        self.embed_object = nn.Linear(KINDS + 2, width)
        self.embed_request = nn.Linear(KINDS, width)
        layer = nn.TransformerEncoderLayer(width, heads, dim_feedforward=64, dropout=0.0, batch_first=True)
        self.encoder = nn.TransformerEncoder(layer, layers)
        self.out = nn.Linear(width, 2)

    def forward(self, tokens, request):
        sequence = torch.cat([self.embed_request(request)[:, None], self.embed_object(tokens)], dim=1)
        return self.out(self.encoder(sequence)[:, 0])


class Plain(nn.Module):
    """A plain network on all numbers at once: 8 object slots of 11 numbers (10 + 'filled'), plus the request."""
    SLOTS = 8

    def __init__(self, hidden=256):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(self.SLOTS * (KINDS + 3) + KINDS, hidden), nn.ReLU(),
                                 nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, 2))

    def forward(self, tokens, request):
        count, objects, _ = tokens.shape
        slots = torch.zeros(count, self.SLOTS, KINDS + 3)
        slots[:, :objects, :KINDS + 2] = tokens
        slots[:, :objects, KINDS + 2] = 1
        return self.net(torch.cat([slots.flatten(1), request], dim=1))


MODELS = {"lookup": Lookup, "transformer-1": lambda: Transformer(1), "transformer-2": lambda: Transformer(2),
          "plain-16": lambda: Plain(16), "plain-256": lambda: Plain(256)}


LEARNING_RATE = {"find/lookup": 0.01, "find/plain-16": 0.003, "find/plain-256": 0.003,     # Adam rates chosen on
                 "nearest/lookup": 0.001, "nearest/transformer-1": 0.01, "nearest/transformer-2": 0.001,  # validation scenes
                 "nearest/plain-256": 0.0003}


def train(model, task, seed=0, lr=None, epochs=None):
    lr = lr or LEARNING_RATE[f"{task}/{model}"]
    epochs = epochs or {"find": 30, "nearest": 40}[task]
    torch.set_num_threads(1)
    torch.manual_seed(seed)
    generator = torch.Generator().manual_seed(seed)
    tokens, request, answers = scenes(20000, 4, generator)
    network = MODELS[model]()
    optimizer = torch.optim.Adam(network.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
    for epoch in range(epochs):
        for batch in torch.randperm(len(tokens), generator=generator).split(128):
            loss = loss_fn(network(tokens[batch], request[batch]), answers[task][batch])
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    return network


def miss_cm(network, task, objects, seed=99):
    tokens, request, answers = scenes(2000, objects, torch.Generator().manual_seed(seed))
    with torch.no_grad():
        return 100 * (network(tokens, request) - answers[task]).norm(dim=1).mean().item()


if __name__ == "__main__":
    network = train("lookup", "find")
    tokens, request, answers = scenes(1, 4, torch.Generator().manual_seed(5))
    with torch.no_grad():
        objects = network.embed_object(tokens[0])                  # the 4 object tokens, 16 numbers each
        query = network.embed_request(request)                     # the request token
        wq, wk, wv = network.attend.in_proj_weight.split(16)       # the layer's three projections
        bq, bk, bv = network.attend.in_proj_bias.split(16)
        mixed, weights, scores = attention(query @ wq.T + bq, objects @ wk.T + bk, objects @ wv.T + bv)
        by_hand = network.out(network.attend.out_proj(mixed))[0]
        by_torch = network(tokens, request)[0]
    kinds = [ITEMS[i] for i in tokens[0, :, :KINDS].argmax(dim=1).tolist()]
    print("objects:", kinds, "request:", ITEMS[request[0].argmax().item()])
    print("scores:", [round(v, 3) for v in scores[0].tolist()], " weights:", [round(v, 4) for v in weights[0].tolist()])
    print("answer by hand:", [round(v, 4) for v in by_hand.tolist()], " by the module:", [round(v, 4) for v in by_torch.tolist()],
          " true:", [round(v, 4) for v in answers["find"][0].tolist()])
    shuffled = tokens[:, torch.tensor([2, 0, 3, 1])]
    with torch.no_grad():
        print("same scene, objects in another order:", [round(v, 4) for v in network(shuffled, request)[0].tolist()])

    for model in ("lookup", "plain-16", "plain-256"):
        net = network if model == "lookup" else train(model, "find")
        print(f"find    {model:13s} {sum(p.numel() for p in net.parameters()):6d} parameters: miss {miss_cm(net, 'find', 4):.2f} cm "
              f"with 4 objects, {miss_cm(net, 'find', 6):.2f} cm with 6")
    for model in ("lookup", "transformer-1", "transformer-2"):
        net = train(model, "nearest")
        print(f"nearest {model:13s} {sum(p.numel() for p in net.parameters()):6d} parameters: miss {miss_cm(net, 'nearest', 4):.2f} cm")

Our first network is just this: the request’s query, the objects’ keys and values, then a linear layer from the averaged value to x and y. It has 1,442 weights and biases, and its lists are 16 numbers long.

pytorch_attention.py · the Lookup networkOpen source ↗

def forward(self, tokens, request):
    objects = self.embed_object(tokens)                        # (N, objects, 16)
    query = self.embed_request(request)[:, None]               # (N, 1, 16)
    answer, _ = self.attend(query, objects, objects)           # query, keys, values
    return self.out(answer[:, 0])

Full file: pytorch_attention.py, in its first card above ↑

self.attend is nn.MultiheadAttention(16, num_heads=1). It holds the three linear layers that make queries, keys and values.
(query, objects, objects): the query comes from the request; keys and values both come from the objects.

Choose a scene and a request. The lines show how much weight each object gets:

How to read the lab:

The square is the table seen from above, 1 m each side.
Each line from the request box to an object is as thick as that object’s weight.
The blue dot is the network’s answer; the ring is the true position of the requested object.

The trained network gives the requested object almost all the weight, so the averaged value is almost exactly its value, and the answer lands on it.

The numbers are the example’s own: the hand-written attention function, fed with the module’s weights, gives the same answer as nn.MultiheadAttention to every printed digit.

Query · key · value: what is looked for · what each token offers for matching · what each token hands over.
Attention weights: the softmax of the query–key scores: how much each token’s value counts.
Cross-attention: attention where the query comes from one source (the request) and the keys and values from another (the objects).

Test

Does the task really need attention? A plain network could take all the numbers at once: eight object slots of 11 numbers (10, plus 1 for “slot in use”), then the request. Training always fills the first four slots.

All networks train with Adam on 20,000 scenes of four objects; each rate is chosen on validation scenes by the mean miss over seeds 0–2. The tests use 2,000 new scenes, with four objects and with six:

Every seed and every rate tried

Order

Attention has no idea of order. Every object is scored and averaged the same way wherever it sits in the list.

The example checks it: the same scene with its four objects in another order gives exactly the same answer. The largest difference is m.

Here that is what we want, because each token already carries its position on the table. When order does matter, as for words in a sentence, each token gets a position embedding added: a list of numbers that says where in the sequence it is.

Position embedding: numbers added to each token to tell attention where the token came from.

ACT needs them. Its image tokens are cells of the ResNet feature map, and without a position each cell would be an anonymous feature. LeRobot adds a two-dimensional sine-wave pattern to each cell, which encodes its row and column.

Layers

Now a harder request: where is the object nearest to the requested one? That is two questions in a row. First find the requested object, then find what is closest to it.

A single attention layer can only ask one. So we let every token attend to every other, and stack two such layers:

pytorch_attention.py · the Transformer networkOpen source ↗

def forward(self, tokens, request):
    sequence = torch.cat([self.embed_request(request)[:, None], self.embed_object(tokens)], dim=1)
    return self.out(self.encoder(sequence)[:, 0])

Full file: pytorch_attention.py, in its first card above ↑

torch.cat puts the request token first, then the object tokens: five tokens in one sequence.
self.encoder is nn.TransformerEncoder: layers of self-attention, each followed by a small feed-forward network (32 → 64 → 32) and layer normalisation.
[:, 0] reads the answer from the request token’s place.

Self-attention: attention where every token makes a query, and keys and values come from the same tokens, so each can gather from all the others.
Transformer: a stack of layers, each self-attention (or cross-attention) followed by a small feed-forward network.
Heads: several attentions side by side in one layer, each with its own queries, keys and values on part of the numbers; here 2, in ACT 8.

Layer 1: what the request token attends to

Layer 2: what the request token attends to

The two layers split the job between them: layer 1 looks at the requested object, layer 2 at its nearest neighbour. Nobody told them to; it is what training found. The weights shown are averaged over the layer’s two heads.

ACT’s transformer

ACT is built from the same parts, larger. In LeRobot 0.6.1:

Encoder tokens: one for a learned “style” variable (the next lesson), one for the arm’s joint angles, and 300 per camera from ResNet-18’s 15 × 20 feature map. Each is a list of 512 numbers.
Encoder: 4 layers of self-attention with 8 heads, each followed by a feed-forward network 512 → 3,200 → 512.
Decoder: 100 query tokens, one per future time step. Each starts as zeros plus a learned position embedding, and cross-attends to the encoder’s tokens: our Lookup network’s pattern, a hundred times over.
Output: a linear layer turns each of the 100 tokens into one action, so one pass predicts a chunk of 100 actions.

The decoder has a single layer. LeRobot’s config explains why: the original ACT code asked for 7, but a bug meant only the first was ever used, so 1 matches the published results.

Recap

scenes(count, objects, generator) makes tokens: one list of 10 numbers per object, plus the request.
attention(query, keys, values) scores keys against the query, softmaxes, and averages the values.
Lookup is one cross-attention layer: the request asks, the objects answer. It finds the requested object to a few millimetres, with any number of objects.
Shuffling the tokens leaves the answer unchanged: position must come from the tokens themselves, or from position embeddings.
Transformer(2) stacks two self-attention layers: the first finds the requested object, the second its nearest neighbour.

Attention lets each output choose its inputs by content. Stacking layers lets it ask one question after another, and that is what ACT does with camera patches, joint angles and its chunk of 100 future actions.

Key concepts

Token: One item of input as its own list of numbers.
Query · key · value: What is looked for · what each token offers for matching · what it hands over.
Attention weights: Softmax of query–key scores, divided by \(\sqrt d\).
Cross-attention: Queries from one source, keys and values from another.
Self-attention: Queries, keys and values all from the same tokens.
Heads: Several attentions side by side, each on part of the numbers.
Position embedding: Numbers added to a token to say where it came from.
Transformer: Layers of attention, each followed by a small feed-forward network.
Encoder · decoder: The part that processes the input tokens · the part whose query tokens read from them to produce outputs.

Run it yourself

From the repository root:

python3 site/ml/foundations/examples/pytorch_attention.py
python3 site/ml/foundations/tools/build_attention_data.py

The first trains seed 0 of each network at its chosen rate, prints one scene’s scores and weights with the hand answer beside the module’s, the shuffled-order check, and the test misses (about a minute). The second chooses the rates, runs seeds 0–9 in parallel and saves this page’s data.

Further reading: Vaswani et al., “Attention Is All You Need” (2017) · Zhao et al., “Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware” (ACT, 2023) · nn.MultiheadAttention.