PythiloLesson 5 of 9 · progress kept in this browser

A self-attention layer

In a model, the queries, keys and values are not given: each is a learned linear map of the same input. That is self-attention: the sequence attends to itself.

import torch
from torch import nn

d = 8
q = nn.Linear(d, d, bias=False)
# batch, positions, d
x = torch.randn(2, 5, d)
print(q(x).shape)
# torch.Size([2, 5, 8])

The layer applies three such maps, scores every query against every key, masks the future, and mixes the values. Shapes stay (batch, positions, d), so layers can be stacked. With the causal mask, the output at a position depends only on that position and the ones before it, which is what lets a model learn to predict the next token from all of them at once.

Your task

Write the class SelfAttention(d), an nn.Module with three nn.Linear(d, d, bias=False) layers named q, k and v. Its forward(x) takes (batch, positions, d) and returns causal scaled dot-product attention of the three projections, of the same shape. You may compute it yourself or use F.scaled_dot_product_attention.

This lesson uses torch, which Python in this browser does not include. Read it here, and run and check it in the Pythilo app on iPhone or iPad.