CausalAttention is a single-head attention module that processes input sequentially while preserving autoregressive token order. It supplies one attention head for a multi-head attention design. section 3.6
Multiple CausalAttention objects can be stacked, with each head independently computing attention before the head results are combined. An integrated MultiHeadAttention class later replaces this repeated-module structure by projecting once, reshaping into heads, and recombining the outputs. section 3.6.2 section 3.6
From one head to many
A multi-head wrapper can create several CausalAttention instances, such as two independent heads, concatenate their results, and thereby implement multi-head attention. The efficient alternative uses one larger query, key, and value projection and splits each result into heads. section 3.6.2 section 3.6
Listing 3.5 An efficient multi-head attention class
class MultiHeadAttention(nn.Module):
def __init__(self, d_in, d_out,
context_length, dropout, num_heads, qkv_bias=False):
super().__init__()
assert (d_out % num_heads == 0), \
"d_out must be divisible by num_heads"
self.d_out = d_out
self.num_heads = num_heads
self.head_dim = d_out // num_heads #1
self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
self.out_proj = nn.Linear(d_out, d_out) #2
self.dropout = nn.Dropout(dropout)
self.register_buffer(
"mask",
torch.triu(torch.ones(context_length, context_length),
diagonal=1)
)
def forward(self, x):
b, num_tokens, d_in = x.shape
keys = self.W_key(x) #3
queries = self.W_query(x) #3
values = self.W_value(x) #3
keys = keys.view(b, num_tokens, self.num_heads, self.head_dim) #4
values = values.view(b, num_tokens, self.num_heads, self.head_dim)
queries = queries.view(
b, num_tokens, self.num_heads, self.head_dim
)
keys = keys.transpose(1, 2) #5
queries = queries.transpose(1, 2) #5
values = values.transpose(1, 2) #5
attn_scores = queries @ keys.transpose(2, 3) #6
mask_bool = self.mask.bool()[:num_tokens, :num_tokens] #7
attn_scores.masked_fill_(mask_bool, -torch.inf) #8
attn_weights = torch.softmax(
attn_scores / keys.shape[-1]**0.5, dim=-1)
attn_weights = self.dropout(attn_weights)
context_vec = (attn_weights @ values).transpose(1, 2) #9
#10
context_vec = context_vec.contiguous().view(
b, num_tokens, self.d_out
)
context_vec = self.out_proj(context_vec) #11
return context_vec
Inspect the efficient integrated implementation that replaces separately instantiated CausalAttention heads with reshaping, transposing, and shared projections. section 3.6.2
