Einstein Summation in Geometric Deep Learning
Patrick Nicolas
Director Data Engineering @ aidéo technologies |software & data engineering, operations, and machine learning.
The einsum function in NumPy and PyTorch, which implements Einstein summation notation, provides a powerful and efficient way to perform complex tensor and matrix operations. This function plays a crucial role in two major libraries for geometric deep learning: Geomstats and PyTorch Geometric.
This article explores the various applications and use cases of this essential function.
Introduction
Overview
Numpy einsum
PyTorch einsum
Applications
Matrix Operations
Neural Network Linear Layer
References
What you will learn: How to leverage NumPy’s einsum method to efficiently perform complex linear algebra operations, as demonstrated in PyTorch Geometric and the Geomstats Python library.
✅ The complete article, featuring design principles, detailed implementation, in-depth analysis, and exercises, is available Dive into Einstein Notation in Geometric Deep Learning
Notes:
Introduction
Overview
The Einstein summation convention is widely used in differential geometry, particularly in fields such as physics, quantum mechanics, general relativity, and geometric deep learning. Two prominent libraries in geometric learning, PyTorch Geometric and Geomstats, heavily rely on Einstein summation for efficient tensor operations.
But what is the Einstein summation?
The Einstein convention implies a summation over a set of indexed terms in a formula ref 1]. Here are some example for Linear Algebra:
📌 The Einstein summation convention is also extensively used in differential calculus. The following formulas represent the divergence, gradient, and Laplacian of a function or vector field [ref 2].
Numpy einsum
The numpy.einsum function provides a flexible and efficient way to perform tensor operations using Einstein summation notation [ref 3]. It simplifies operations such as dot products, matrix multiplications, transpositions, outer products or tensor contractions.
The method einsum implements the Einstein summation indices convention with 3 different operational modes.
PyTorch einsum
The torch.einsum method in PyTorch provides a flexible and efficient way to perform tensor operations based on Einstein summation notation [ref 4]. It allows for concise, readable, and optimized implementations of dot products, matrix multiplications, outer products, transpositions, and more complex tensor contractions.
As expected, the PyTorch definition follows the same protocol or convention that its counter part in Numpy.
Applications
Matrix Operations
The NumPy einsum method offers an elegant alternative to the commonly used @operator or dedicated NumPy methods.
Matrix multiplication
einsum can be used as an alternative to the Numpy operator @ and matmul method.
The sequence of operations are
Output
[[0.7 2.25]
[2. 4.75]]
Let’s look at PyTorch implementation of einsum for matrix multiplication:
Recommended by LinkedIn
tensor([[ 1.1400, -0.4900, -0.2500],
[ 0.1300, -0.3000, 0.3700],
[-0.3600, 0.5700, -0.5100]])
Dot product
Let’s look at the computation of the dot product of two vectors using einsum methods in Numpy and PyTorch.
3.5
Matrix outer product
The NumPy outer method is specifically designed for vector operations. When applied to matrices, the outer product is flattened into a single matrix, requiring reshaping to properly compare results with the einsum function.
This demonstrates another advantage of using einsum, as it avoids the need for reshaping and provides a more intuitive representation of tensor operations.
Outer Product Vectors
[[0.1 2. 1.2 ]
[0.05 1. 0.6 ]
[0.2 4. 2.4 ]
[0.07 1.4 0.84]
[0.09 1.8 1.08]]
Outer Product Matrices
[[[[ 0.1 2. ]
[ 1.2 0.5 ]]
[[ 0.05 1. ]
[ 0.6 0.25]]
[[ 0.04 0.8 ]
[ 0.48 0.2 ]]]
[[[ 0.05 1. ]
[ 0.6 0.25]]
[[-0.03 -0.6 ]
[-0.36 -0.15]]
[[ 0.02 0.4 ]
[ 0.24 0.1 ]]]
[[[-0.1 -2. ]
[-1.2 -0.5 ]]
[[ 0.04 0.8 ]
[ 0.48 0.2 ]]
[[-0.03 -0.6 ]
[-0.36 -0.15]]]]
Matrix Transpose
The simplest way to transpose a matrix in both NumPy and PyTorch is by using the .T operator or calling the matrix.transpose and torch.transpose methods
This operation can also be effortlessly performed using einsum, simply by reversing the order of indices.
[[ 1. 0.5 -1. ]
[ 0.5 -0.3 0.4]
[ 0.4 0.2 -0.3]]
Matrix Trace
The trace of a square matrix is the sum of its main diagonal elements and can also be interpreted as the sum of some of its eigenvalues. It serves as a linear mapping for any square matrix and remains invariant under transposition, meaning the trace of a matrix equals the trace of its transpose. In NumPy, the trace is easily computed using the trace() method:
The einsum implementation is self-explanatory.
3.99999
Neural Network Linear Layer
I thought it would be interesting to look at few applications that leverages einsum function.
The linear transformation of a layer of multi-layer perceptron [ref 5] is defined as y = W.x+b
with
[2.3 2.46 1.29]
✅ Thank you for reading this article! I hope you found this overview insightful. For a detailed exploration of the topic, check out Dive into Einstein Notation in Geometric Deep Learning
References
Patrick Nicolas has over 25 years of experience in software and data engineering, architecture design and end-to-end deployment and support with extensive knowledge in machine learning.He has been director of data engineering at Aideo Technologies since 2017 and he is the author of "Scala for Machine Learning", Packt Publishing ISBN 978-1-78712-238-3 and Hands-on Geometric Deep Learning newsletter.
#EinsteinSummation #einsum #Numpy #PyTorch #PyG