Johnny's Knowledge

Batched Matrix Multiplication

Properties1
tagsdeep-learning, pytorch

Batched matrix multiplication performs multiple matrix multiplications in parallel, where each batch operates independently. This is common in deep learning frameworks for efficient computation.

Concept

Instead of multiplying two matrices, we multiply two batches of matrices:

  • Input: Two 3D tensors of shape (batch_size, m, n) and (batch_size, n, p)
  • Output: 3D tensor of shape (batch_size, m, p)

Each pair of matrices in the batch is multiplied independently.

Example

import torch
 
# Create batched matrices
A = torch.randn(32, 10, 20)  # batch_size=32, m=10, n=20
B = torch.randn(32, 20, 15)  # batch_size=32, n=20, p=15
 
# Batched matrix multiplication
C = torch.bmm(A, B)  # shape: (32, 10, 15)

Using Einstein Summation (Einsum)

More flexible notation for complex tensor operations:

# Equivalent to bmm
C = torch.einsum('bmn,bnp->bmp', A, B)
 
# More complex example: batched attention
# Q: (batch, seq_len, dim)
# K: (batch, seq_len, dim)
attention = torch.einsum('bqd,bkd->bqk', Q, K)

Common Patterns

Standard Batched Multiplication

'bij,bjk->bik' - Matrix multiplication for each batch

Batched Matrix-Vector

'bij,bj->bi' - Matrix-vector product for each batch

Batched Outer Product

'bi,bj->bij' - Outer product for each batch

Advantages

  • Parallelization: GPU can process all batches simultaneously
  • Memory efficiency: Contiguous memory access
  • Code simplicity: Single operation instead of loops

Applications in Deep Learning

  • Attention mechanisms in transformers
  • Graph neural networks
  • Recurrent neural networks
  • Multi-head attention

Backlinks