Einstein Notation: A New Lens on Transformers
In this article, we'll embark on a playful journey through the world of transformers, unraveling the complexities of their architecture using the Einstein notation.
Introduction:
Transformer models have revolutionized the field of natural language processing (and beyond), achieving state-of-the-art results on a variety of tasks. They have impressive performance but the underlying mathematical operations can be complex and difficult to grasp – especially without breaking down the individual layers. In this article, I propose using the Einstein notation to express the mathematical operations within a transformer model.
Note that the Einstein notation is normally used in Physics and Mathematics such as in General Relativity, Electromagnetism, Quantum and Fluid Mechanics but also in Linear Algebra to represent matrix operations in a more compact form.
The goal is to write the mathematical operations of every layer in a concise and elegant way. By leveraging implicit summation over repeated indices, Einstein notation can simplify the representation of tensor operations, making it (potentially) easier to understand and therefore implement the individual layers of the transformer models. In more detail, I will show how the Einstein notation can be applied to various components of transformer models, including self-attention, feed-forward neural networks, and layer normalization.
For simplicity, I will focus only on the decoder part of the transformer which is currently a common best practice for generative large language models (LLMs).
Motivation
To this date, modern transformer models rely on computationally intensive operations, particularly within the self-attention mechanism. In other words, in research and development we experience that an increasing sequence / token length is a major bottleneck due to the exponential growth in computational cost. Thus, scaling models is hard relying on the current mathematical solutions for inference and training.
Back in the day, I learned the Einstein notation in Physics. Einstein notation is know for its elegance and efficiency. Recently, I sought to explore its potential in simplifying and optimizing the mathematical representations of the transformer architecture – thus, transforming the math of the transformer model to convey it to a lager audience of e.g. non-machine learning researchers -creating a new perspective which could lead to novel insights and optimizations.
The core concepts of self-attention are relatively straightforward. However, the explicit matrix operations and summation can obscure the underlying structure and make some mathematical steps difficult to follow.
Ultimately, the goal of this article is to contribute to the understanding of the transformer model by adopting a fresh point of view of its mathematical foundations.
Let's do a quick recap of the essential of the Einstein notation:
The inner product of two vectors is obtained by pairing up corresponding elements from each vector, multiplying the pairs, and then adding up all the resulting products.
The Levi-Civita symbol (ε) is used to concisely express the cross product. The repeated indices j and k are implicitly summed over. In other words, the Einstein summation convention implies that we sum over repeated indices (j and k in this case). While using the Levi-Civita symbol and Einstein notation, the cross product can be expressed in a compact and elegant way, without explicitly writing out the summation.
The cross product of two vectors a and b results in a new vector that is perpendicular to both a and b. The magnitude of the resulting vector is equal to the product of the magnitudes of a and b times the sine of the angle between them.
Again the in Matrix multiplication, the repeated index k is implicitly summed over. This means that to calculate the element at position (i, j) in the product matrix AB, we multiply corresponding elements of the i-th row of A with the j-th column of B and sum the products. Note – that the Matrix multiplication is a fundamental operation in Transformer models – and the inspiration of this article.
Based on the examples given above – one can clearly see that the Einstein notation has some advantages. It is more concise, as it reduces the number of symbols and operations required. It is more clear, as it ** highlights the underlying structure of ‘tensor' operations, also making them more intuitive. Efficient, as it could lead to easier implementations of algorithms, especially when dealing with high-dimensional matrices. And lastly, more genera**l, as it can be applied to a wide range of ‘tensor' operations, making it a versatile tool for expressing complex mathematical relationships. In a nutshell, utilizing the Einstein notation, researchers and practitioners could leverage those advantages when working with ‘tensor'-based models and other deep learning architectures in a mathematical sense.
Methodology
In this section, I will introduce the math behind the transformer model (decoder) in a standard fashion. In addition, I will demonstrate how Einstein notation can be used to represent the mathematical operations from a different perspective.
The Einstein notation, allows readers that are not familiar with state-of-the-art Machine Learning research and the corresponding math notation, to consume the mathematical foundation of the transformer in a more standard manner.
Token Embedding converts input tokens (words or sub-words aka ‘tokens') into dense numerical representations, enabling the model to process and understand the semantic and syntactic information within the text.
e: The embedding vector for the input token x at index i E: The embedding matrix, where i is the index of the token and j is the dimension x: The one-hot encoded representation of the input token x at index j i: Index of the token j: Index of the dimension of the embedding
Positional Encoding adds to the word (or token) embeddings and provides the model with information about the relative and absolute position of each word in the sequence. Here the Einstein notation does not change anything to the original formulation.
PE(pos, i): The positional encoding at position pos and dimension i pos: The position of the token in the sequence i: The dimension of the positional encoding d: The model dimension or embedding dimension
The Attention mechanism calculates the relevance of each input token to the current output token by computing a weighted sum of the input embeddings – where the weights are determined by the attention scores derived from the query (Q), key (K), and value (V) vectors of the input and output tokens in each head (i).
Q: Query matrix K: Key matrix V: Value matrix i, j, k, l: Indices used to access specific elements of the matrices d_k: Dimension of the key vectors
Feed-Forward Network (FFN): The important of the FFN is two-fold. First, it does introduce non-linearity via an activation function. In the original Attention is all you need paper – the ReLU activation function was used. Nowadays we see more advanced activation functions in the current decoder-only Large Langue Models. Just to recap – the non-linearity allows the network to learn complex mappings between input and output using backpropagation. Second, the FFN operates on the output of the attention layer, which captures long-range dependencies – hence, the FFN helps to extract meaningful features. Lastly, literature also states that the FFN adds to the capacity of the network by introducing additional layers and parameters.
xj: Input vector. W: Weight matrices for the first and second layers, respectively. b: Bias vector for the first and second layers, respectively. i, j, k: Indices used to access specific elements of the matrices and vectors.
Layer Normalization is a important component of attention mechanisms in LLMs, thus playing a significant role in their effectiveness and stability by (1) Stabilizing Training and (2) Enhancement . The main advantage of layer normalization is the idea of gradient clipping. In other words, normalization helps prevent the vanishing or exploding gradient problem during training – as it keeps gradients within a reasonable range, making training more stable. Second, layer normalization projects input vectors to a space where attention queries can attend to all keys equally. This offloads some burden of learning this behavior from the attention mechanism. Further, by scaling all vectors to the same norm, layer normalization ensures that no single key can dominate the attention process – thus, avoiding being biased towards certain inputs.
x: Input tensor of shape [batch size, sequence length, hidden size] μ: Mean of x over the hidden size dimension σ: Standard deviation of x over the hidden size dimension α, β: Learnable parameters (scale and shift)
Conclusion, Limitations and Future Research
In this article, we've explored the application of Einstein notation to the mathematical operations within a transformer model. To do this the implicit summation over repeated indices was leveraged. Therefore, we've presented a more concise and elegant representation of the complex tensor operations involved in attention, feed-forward neural networks, and layer normalization.
While Einstein notation offers a valuable perspective on transformer models, it's important to acknowledge its limitations and potential areas for future research. First, there is a learning curve the **** Einstein notation. Although using it can simplify complex expressions, it requires a certain level of mathematical maturity to fully grasp its nuances. Second, from a research and analytical perspective the Einstein notation makes sense to convey a fresh perspective, directly translating Einstein notation into efficient code can be challenging, especially for large-scale models.
Future Directions for this research could be exploring compiler optimizations and hardware acceleration techniques to leverage the potential performance benefits of Einstein notation. Also a hybrid approach could be useful combining Einstein notation with traditional matrix notation to strike a balance between conciseness and readability.
Most importantly the generation of theoretical insights could be very attractive future research **** direction, as it could lead to a deeper understanding of the underlying principles of transformer models and potentially inspire novel architectures and optimization techniques.
Clearly, by addressing these limitations and exploring future research directions, we can unlock the full potential of Einstein notation in advancing our understanding and development of transformer models.
References
- Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin, "Attention Is All You Need," in Advances in Neural Information Processing Systems 30 (2017)
