Decoding Einsum

Einstein summation, or einsum, is a powerful and flexible operation used for tensor manipulation. It allows you to perform summation over multiple dimensions concisely, making it a favourite among data scientists and machine learning practitioners.

Why Choose Einsum?

• Flexibility: Works with a variety of tensor shapes and operations.
• Conciseness: Eliminates the need for multiple intermediate steps like transpositions or reshaping.
• Simplicity: Reduces complex operations into a single line of code. Einsum is available across major libraries like NumPy and PyTorch. Let’s explore its versatility using PyTorch, starting with some common use cases like matrix multiplication, summation, permutation, and more.

Syntax of Einsum

Einsum uses a simple notation to define operations:

• Indices: Represent dimensions (e.g., i, j, k).

• Commas: Separate tensors being operated on.

• Arrow (**>**): Indicates the desired output shape. For example, 'ik, kj -> ij' means:

• The first tensor has dimensions ik.

• The second tensor has dimensions kj.

• The output should have dimensions ij, summing over the common dimension k.

Examples of Einsum Operations

Matrix Multiplication

a = torch.randn(3, 5)
b = torch.randn(5, 2)
c = torch.einsum('ik, kj -> ij', a, b)
print(c.shape)  # Output: torch.Size([3, 2])

Here, k is the common dimension, and the output is shaped as [3, 2].

Einsum simplifies operations that typically require transpositions:

a = torch.randn(3, 5)
b = torch.randn(2, 5)
c = torch.einsum('ik, jk -> ij', a, b)
print(c.shape)  # Output: torch.Size([3, 2])

Batch Matrix Multiplication

a = torch.randn(3, 5, 4)
b = torch.randn(3, 4, 6)
c = torch.einsum('ijk, ikl -> ijl', a, b)
print(c.shape)  # Output: torch.Size([3, 5, 6])

Summation

Einsum can be used for summation over dimensions:

a = torch.tensor((1, 2, 3))
b = torch.einsum('i ->', a)
print(b)  # Output: tensor(6)

If no output dimension is specified, einsum sums over all dimensions.

You can also sum along a specific axis:

a = torch.tensor(((1, 2, 3), (3, 4, 5)))
b = torch.einsum('i ->', a[1])
print(b)  # Output: tensor(12)

Permutation

Reordering tensor dimensions becomes straightforward:

a = torch.randn(3, 5, 4)
c = torch.einsum('ijk -> kji', a)
print(c.shape)  # Output: torch.Size([4, 5, 3])

Outer Product

The outer product of two tensors can be written concisely:

a = torch.randn(5)
b = torch.randn(3)
c = torch.einsum('i, j -> ij', a, b)
print(c.shape)  # Output: torch.Size([5, 3])

Attention Using Einsum

Einsum is instrumental in implementing self-attention mechanisms in deep learning:

q = torch.randn(2, 3, 4, 5)  # (bs, nh, sl, dim)
k = torch.randn(2, 3, 4, 5)
v = torch.randn(2, 3, 4, 5)

a = torch.einsum('abcd, abed -> abce', q, k)
print("After matrix multiplication of q & k:", a.shape)  # torch.Size([2, 3, 4, 4])

a /= math.sqrt(5.0)  # Scale by dimension size
a = F.softmax(a, dim=-1)

a = torch.einsum('abcd, abef -> abcf', a, v)
print("After matrix multiplication of attention & v:", a.shape)  # torch.Size([2, 3, 4, 5])

Explanation:

1. q and k are multiplied along the last dimension (d) to calculate attention weights.
2. We normalize the weights and apply softmax.
3. Finally, these weights are multiplied with v to produce the output. This enables efficient computation while keeping the code clean and readable. Try it today!