Several Heads in Python

One attention head gives each position one blend of the positions up to it. Its weights add up to 1, so more on one letter means less on the others. Multi-head attention runs several heads side by side on the same input. Each head has its own queries, keys and values, so after training each one can look for something different: the letter just before, say, or the start of the name.

Two new tools. nn.ModuleList([...]) is a list of layers that PyTorch looks inside, so their weights show up in parameters() and get trained (a plain Python list hides them). torch.cat(list, dim=-1) joins tensors along their last dimension and gives back one tensor that keeps every number:

a = torch.tensor([[1, 2], [3, 4]])
b = torch.tensor([[5], [6]])
print(torch.cat([a, b], dim=-1).tolist())   # [[1, 2, 5], [3, 4, 6]]

Two heads of 16 each give back [32, 8, 16]. Joined side by side, that makes [32, 8, 32], the size of the input. A last layer, the projection nn.Linear(32, 32), makes each output number from all 32, so it mixes what the two heads found. Run on a batch, your code and the checks after it print:

h: [32, 8, 32] out: [32, 8, 32]
each head gives: [32, 8, 16] and [32, 8, 16]
out is proj(head 1 and head 2 side by side): True
window 0: 'kelolo\nr'
head 1 at position 7 takes [0.157, 0.202, 0.079, 0.094, 0.182, 0.141, 0.041, 0.103]
head 2 at position 7 takes [0.058, 0.122, 0.203, 0.196, 0.095, 0.135, 0.139, 0.053]
change position 7: positions 0-6 stay the same: True | position 7 changes: True
weights to learn: 4128

The heads are untrained, so their weights are random for now, but they already differ. Changing the last letter changes only the last position's output: the mask inside each head still holds.

Your task: the Head class from the last lesson is ready, and so is the code that runs MultiHead on a batch. Finish MultiHead: in __init__, store two Head(16) in an nn.ModuleList called self.heads and an nn.Linear(n_embd, n_embd) in self.proj. Then write forward(self, x), which runs every head on x, joins their outputs with torch.cat and gives back self.proj of the result.

Assignment
Follow the Coding Tutorial to build the two heads, the projection and the forward method of MultiHead.

Hint
Collect each head's output in a list with a loop over self.heads. Join the list with dim=-1, the last dimension, before passing it to self.proj.