Scaled Dot-Product Attention › Matrix Form for Batched Computation

Michael BrenndoerferMay 25, 20253 min read
Language AI HandbookAll chapters
Part 13 of 67Self-Attention
Chapter 92Scaled Dot-Product Attention › Matrix Form for Batched Computation
Reading section 8 of 15 · Matrix Form for Batched Computation · 3 minSection 8 of 15 · 3 min

Book contents

Language AI Handbook67 parts · 513 chaptersView the whole book →
Active chapter 92 of 513

Section 8 of 15 · 3 min read

Scaled Dot-Product Attention › Matrix Form for Batched ComputationOne section · Sections in this chapter8 / 15 ↓
Sections in this chapter15
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

In practice, transformers process batches of sequences simultaneously. Training on a single sequence at a time would leave GPUs severely underutilized: modern GPUs achieve peak throughput when they can process hundreds or thousands of examples in parallel. The attention formula naturally extends to batched computation by adding a batch dimension. For a batch of BB sequences each of length nn:

  • Q∈RB×n×dkQ \in \mathbb{R}^{B \times n \times d_k}: query matrices for all sequences in the batch
  • K∈RB×n×dkK \in \mathbb{R}^{B \times n \times d_k}: key matrices for all sequences in the batch
  • V∈RB×n×dvV \in \mathbb{R}^{B \times n \times d_v}: value matrices for all sequences in the batch
  • Output ∈RB×n×dv\in \mathbb{R}^{B \times n \times d_v}: enriched representations for all sequences

The matrix multiplication QKTQK^T uses batched matrix multiply: for each example in the batch, we compute the n×nn \times n score matrix independently. NumPy and PyTorch handle this efficiently through the @ operator or torch.bmm. The key point is that sequences in a batch are completely independent: the attention weights for sequence 1 do not influence the attention weights for sequence 2. The batch dimension is purely an efficiency mechanism, not a semantic one.

In[26]:
Code
def batched_attention(Q_batch, K_batch, V_batch, mask=None):
    """
    Batched scaled dot-product attention.

    Args:
        Q_batch: (B, n, d_k)
        K_batch: (B, n, d_k)
        V_batch: (B, n, d_v)
        mask: optional (n, n) mask, applied to all examples

    Returns:
        output: (B, n, d_v)
        weights: (B, n, n)
    """
    d_k = Q_batch.shape[-1]
    # Batched matrix multiply: (B, n, d_k) @ (B, d_k, n) -> (B, n, n)
    scores = Q_batch @ K_batch.transpose(0, 2, 1) / np.sqrt(d_k)

    if mask is not None:
        scores = np.where(mask[np.newaxis, :, :], -1e9, scores)

    # Softmax over last dimension (keys)
    weights = softmax(scores, axis=-1)
    # (B, n, n) @ (B, n, d_v) -> (B, n, d_v)
    output = weights @ V_batch
    return output, weights


# Create a batch of 3 sequences
B = 3
Q_batch = np.random.randn(B, n, d_k)
K_batch = np.random.randn(B, n, d_k)
V_batch = np.random.randn(B, n, d_v)
Out[27]:
Console
Batch size: 3, sequence length: 5, d_k: 4, d_v: 4
Batched output shape: (3, 5, 4)
Batched weights shape: (3, 5, 5)

The batched formulation is identical mathematically to processing each example independently, but the vectorized computation is substantially faster on GPUs due to their parallel execution model. On a modern GPU, processing a batch of 32 sequences takes roughly the same wall-clock time as processing a single sequence, because all 32 matrix multiplications can execute in parallel across thousands of GPU cores.