No history yet

Tensor-Based Self-Attention

Implementing Scaled Dot-Product Attention

The core computational unit of a transformer is the scaled dot-product attention mechanism. Our goal is to implement this not as a conceptual black box, but as a series of optimized tensor operations in PyTorch. The entire mechanism is defined by a single formula.

Attention(Q,K,V)=softmax(QKTdk)V\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V

The primary challenge in implementing this is performing the QKTQK^T matrix multiplication efficiently across both a batch of sequences and multiple attention heads. A naive approach using loops would be prohibitively slow. The standard solution is to use tensor contractions, for which torch.einsum is exceptionally well-suited. It provides a concise and powerful way to specify complex tensor operations.

import torch

# Assume Q, K have dimensions [batch_size, num_heads, seq_len, head_dim]
batch_size = 32
num_heads = 8
seq_len = 512
head_dim = 64

Q = torch.randn(batch_size, num_heads, seq_len, head_dim)
K = torch.randn(batch_size, num_heads, seq_len, head_dim)

# Use einsum for batched matrix multiplication across heads
# 'b h i d, b h j d -> b h i j' means:
# For each batch (b) and head (h),
# multiply matrix i x d with the transpose of j x d,
# resulting in matrix i x j.
scores = torch.einsum('bhid,bhjd->bhij', Q, K)

print(scores.shape) # torch.Size([32, 8, 512, 512])

The resulting scores tensor has dimensions [batch_size, num_heads, seq_len, seq_len], representing the attention score from every token to every other token in the sequence for each head and batch item.

Causal Masking for Autoregression

For decoder-style models like GPT, we must prevent tokens from attending to future positions in the sequence. This autoregressive property is enforced with a causal mask. The mask is a lower triangular matrix that sets all upper-triangular elements of the attention scores matrix to a large negative number, typically negative infinity. After the softmax operation, these positions will have a weight of zero, effectively blocking information flow from the future.

# seq_len is the length of the sequence
seq_len = 512

# Create a lower triangular matrix of ones
mask = torch.tril(torch.ones(seq_len, seq_len))

# Convert the mask to a boolean tensor where we want to mask
# The upper triangle (where mask == 0) should be masked
causal_mask = mask == 0

# In PyTorch, the attention mask is often additive.
# We set masked positions to negative infinity.
scores = torch.randn(32, 8, seq_len, seq_len) # Example scores
scores.masked_fill_(causal_mask, float('-inf'))

# Applying softmax
attention_weights = torch.softmax(scores, dim=-1)

# Verify that the upper triangle is all zeros
print(attention_weights[0, 0, 0, 5])  # Should be tensor(0.)
print(attention_weights[0, 0, 5, 0])  # Should be a non-zero value

Using masked_fill_ with -inf is a numerically stable way to ensure the softmax output for masked positions is zero. The operation is performed in-place for efficiency.

Full Attention Head Implementation

Now we can combine these components into a single, cohesive attention head function. This involves scaling, masking, applying softmax, and finally multiplying by the Value tensor V.

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    """Calculates scaled dot-product attention."""
    head_dim = Q.size(-1)
    
    # 1. MatMul Q and K^T and scale
    # Resulting shape: [batch, num_heads, seq_len, seq_len]
    scores = torch.einsum('bhid,bhjd->bhij', Q, K) / (head_dim ** 0.5)
    
    # 2. Apply optional mask
    if mask is not None:
        scores.masked_fill_(mask, float('-inf'))
    
    # 3. Apply softmax to get attention weights
    # Softmax is applied on the last dimension (keys)
    weights = F.softmax(scores, dim=-1)
    
    # 4. MatMul weights and V
    # Resulting shape: [batch, num_heads, seq_len, head_dim]
    output = torch.einsum('bhij,bhjd->bhid', weights, V)
    
    return output, weights

# --- Example Usage ---
batch_size = 32
num_heads = 8
seq_len = 512
head_dim = 64

Q = torch.randn(batch_size, num_heads, seq_len, head_dim)
K = torch.randn(batch_size, num_heads, seq_len, head_dim)
V = torch.randn(batch_size, num_heads, seq_len, head_dim)

# Create causal mask
mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()

output, attn_weights = scaled_dot_product_attention(Q, K, V, mask=mask)

print("Output shape:", output.shape)
print("Attention weights shape:", attn_weights.shape)

Note the second einsum call for multiplying the attention weights with the V matrix. The notation 'bhij,bhjd->bhid' ensures that for each batch and head, the [seq_len_q, seq_len_k] weights matrix multiplies the [seq_len_k, head_dim] value matrix, resulting in the desired [seq_len_q, head_dim] output tensor. This function now serves as the fundamental building block for a complete multi-head attention layer.