A value is a transformed attention vector whose information is combined according to attention weights. Query and key vectors determine which relationships matter; value vectors provide the information that attention returns. section 3.3.3
Attention uses value vectors through a sequential computation:
- Project each input sequence X through a learned linear layer to obtain V, alongside separate query Q and key K projections. section 3.3.3
- Compare queries with keys using scaled dot products, then apply softmax to produce attention weights whose rows sum to 100%. section 3.3.3
- Multiply the attention weights by V so each output combines information from the value vectors according to the learned relationships. section 3.3.3
- Repeat the process across heads, where each head uses smaller value vectors to capture different aspects before combining the head outputs. section 3.3.3
From tokens to contextual information
For the German phrase “Wie geht es dir,” token embeddings and positional encodings form the input sequence. Each attention head projects that sequence into query, key, and value vectors; the resulting weights combine the value vectors to produce contextual representations used for translation. section 2.1.1 section 3.3.3
Listing 3.4 Calculating self-attention in each attention head
class AttentionHead(nn.Module):
def __init__(self, hidden_size, attention_head_size, bias=True):
super().__init__()
self.hidden_size = hidden_size
self.attention_head_size = attention_head_size
self.query = nn.Linear(hidden_size, attention_head_size, bias=bias)
self.key = nn.Linear(hidden_size, attention_head_size, bias=bias)
self.value = nn.Linear(hidden_size, attention_head_size, bias=bias)
def forward(self, x):
query = self.query(x)
key = self.key(x)
value = self.value(x) #1
attention_scores = torch.matmul(query, key.transpose(-1, -2))
attention_scores = attention_scores / math.sqrt(\
self.attention_head_size)
attention_probs = nn.functional.softmax(attention_scores,
dim=-1) #2
attention_output = torch.matmul(attention_probs, #3
value) #3
return (attention_output, attention_probs)
The implementation makes the value pathway explicit: a learned linear layer creates value, softmax creates attention probabilities, and matrix multiplication combines them. section 3.3.3
Listing 7.3 The attention mechanism in the U-Net
class Attention(nn.Module):
def __init__(self, dim, heads=4, dim_head=32):
super().__init__()
self.scale = dim_head**-0.5
self.heads = heads
hidden_dim = dim_head * heads
self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
self.to_out = nn.Conv2d(hidden_dim, dim, 1)
def forward(self, x):
b, c, h, w = x.shape
qkv = self.to_qkv(x).chunk(3, m=1) #1
q, k, v = map(
lambda t: rearrange(t, 'b (h c) x y -> b h c (x y)', h=self.heads),
qkv) #2
q = q * self.scale
sim = einsum('b h d i, b h d j -> b h i j', q, k)
attn = sim.sofax(dim=-1)
out = einsum('b h i j, b h d j -> b h i',
attn, v) #3
out = rearrange(out, 'b h (x y) d -> b (h d) x y', x=h, y=w)
return self.to_out(out) #4
The U-Net implementation shows the same value mechanism operating on image feature maps, reshaped into heads and recombined into spatial output. section 7.3.1
