PyTorch nn.MultiheadAttention

nn.MultiheadAttention is the library module that runs scaled dot-product attention in several heads at once. A head is one attention run over a slice of the vectors, and the module joins the slices back together.

Call the module

batch_first=True lets x arrive as batch, sequence, features. The call returns the output and the weights, averaged over the 2 heads.

import torch

torch.manual_seed(0)
module = torch.nn.MultiheadAttention(embed_dim=8, num_heads=2, batch_first=True)
x = torch.randn(2, 4, 8)
out, weights = module(x, x, x)
print("Output shape:", tuple(out.shape))
print("Weights shape:", tuple(weights.shape))

The default order differs

Without batch_first, the first dimension is the sequence and not the batch. The same kind of call then reads y as 4 steps from 2 sequences.

y = torch.randn(4, 2, 8)
plain = torch.nn.MultiheadAttention(embed_dim=8, num_heads=2)
out2, _ = plain(y, y, y)
print("Output shape:", tuple(out2.shape))

Three projections in one tensor

in_proj_weight stacks the query, key, and value tables in that order. Slicing the packed tensor into thirds recovers the three tables.

W = module.in_proj_weight
print("Packed shape:", tuple(W.shape))
print("Query slice:", tuple(W[:8].shape))
print("Key slice:", tuple(W[8:16].shape))
print("Value slice:", tuple(W[16:].shape))

Example

import math

import torch

torch.manual_seed(0)
module = torch.nn.MultiheadAttention(embed_dim=8, num_heads=2, batch_first=True)
x = torch.randn(1, 3, 8)

with torch.no_grad():
    module_out, module_weights = module(x, x, x)

W = module.in_proj_weight.detach()
b = module.in_proj_bias.detach()
W_q, W_k, W_v = W[:8], W[8:16], W[16:]
b_q, b_k, b_v = b[:8], b[8:16], b[16:]

Q = x @ W_q.T + b_q
K = x @ W_k.T + b_k
V = x @ W_v.T + b_v

def split_heads(t):
    return t.view(1, 3, 2, 4).transpose(1, 2)

scores = split_heads(Q) @ split_heads(K).transpose(-2, -1) / math.sqrt(4)
weights = torch.softmax(scores, dim=-1)
heads = weights @ split_heads(V)
joined = heads.transpose(1, 2).reshape(1, 3, 8)
manual_out = joined @ module.out_proj.weight.detach().T + module.out_proj.bias.detach()

print(f"Module output shape: {tuple(module_out.shape)}")
print(f"Averaged weights shape: {tuple(module_weights.shape)}")
print(f"Max output difference: {(manual_out - module_out).abs().max().item()}")
print(f"Outputs equal: {torch.allclose(manual_out, module_out, atol=1e-6)}")

Key points

Projection: a learned Linear map that turns token vectors into queries, keys, and values
The three projection tables are packed into in_proj_weight, stacked into one tensor
embed_dim is split evenly across the heads, so 8 with 2 heads gives slices of 4
batch_first=True makes the input order batch, then sequence, then features
The module returns the mixed output and the attention weights averaged over the heads
In short, the module is the scratch computation with learned projections and parallel heads

In this example, the module and a by-hand copy run on the same input. The seed fixes the module's random projection tables before x is drawn. The by-hand copy slices in_proj_weight into its query, key, and value thirds. It projects x with those tables and adds the packed bias in the same way. Each head then runs the scaled dot-product maths from the earlier lesson on its slice of 4. The head outputs are joined and passed through out_proj, the final Linear layer of the module. The module output shape prints as (1, 3, 8), and the averaged weights shape prints as (1, 3, 3). The max output difference prints as 0.0, and the equality check prints True.

Note

With batch_first left off, the module expects the sequence first and the batch second

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import math

import torch

torch.manual_seed(0)
module = torch.nn.MultiheadAttention(embed_dim=8, num_heads=2, batch_first=True)
x = torch.randn(1, 3, 8)

with torch.no_grad():
    module_out, module_weights = module(x, x, x)

W = module.in_proj_weight.detach()
b = module.in_proj_bias.detach()
W_q, W_k, W_v = W[:8], W[8:16], W[16:]
b_q, b_k, b_v = b[:8], b[8:16], b[16:]

Q = x @ W_q.T + b_q
K = x @ W_k.T + b_k
V = x @ W_v.T + b_v

def split_heads(t):
    return t.view(1, 3, 2, 4).transpose(1, 2)

scores = split_heads(Q) @ split_heads(K).transpose(-2, -1) / math.sqrt(4)
weights = torch.softmax(scores, dim=-1)
heads = weights @ split_heads(V)
joined = heads.transpose(1, 2).reshape(1, 3, 8)
manual_out = joined @ module.out_proj.weight.detach().T + module.out_proj.bias.detach()

print(f"Module output shape: {tuple(module_out.shape)}")
print(f"Averaged weights shape: {tuple(module_weights.shape)}")
print(f"Max output difference: {(manual_out - module_out).abs().max().item()}")
print(f"Outputs equal: {torch.allclose(manual_out, module_out, atol=1e-6)}")

Exercise

With embed_dim=8 and num_heads=2, how wide is the slice each head works on?

4 8 2 16