Transformers and Multi-Head Attention: How the Core Engine of Modern Artificial Intelligence Works
An in-depth exploration of the Transformer architecture and its mathematical process, which grants AI the ability to understand language by simultaneously measuring relationships between words.
Today, behind major language models discussed worldwide such as ChatGPT, Claude, or vision models, a single foundational technology is at work, which is called the Transformer architecture. This model was first introduced on June 12, 2017, in the historic research paper "Attention Is All You Need" published by Google researchers. At that time, the main objective of the research was to make machine translation processes more accurate and faster, but this architecture completely changed the landscape of the entire artificial intelligence world.
The greatest feature of the Transformer model is its unmatched capability for parallel processing. Before the advent of Transformers, neural networks used to read text sequentially, word by word, whereas a Transformer can scan an entire sentence all at once. The main driving force behind this revolutionary change is the 'Multi-head attention' mechanism, which connects each word with all other words based on their context and importance.
From Sequential Processing to Parallel Computation History
Before the arrival of Transformers, Recurrent Neural Networks (RNNs), such as the 1990 Elman network, were widely used for sequential data. However, they suffered from a major flaw known as the vanishing-gradient problem, due to which information about the initial words was lost by the time the end of a long sentence was reached. To overcome this flaw, Long Short-Term Memory (LSTM) was introduced in 1997, which could preserve information through multiplicative gating units. Even so, since LSTMs still worked on only one token at a time, they could not be trained quickly in parallel on supercomputers or Graphics Processing Units (GPUs).
In 2014, encoder-decoder based sequence-to-sequence (seq2seq) models arrived, in which a machine translation model with 380 million parameters used two LSTM networks. Another model with 130 million parameters used Gated Recurrent Units (GRUs). However, compressing an entire sentence into a fixed-size vector resulted in an information bottleneck. To solve this problem, the attention mechanism was added via 'RNN search'. In 2016, Google launched Google Neural Machine Translation (GNMT), which had 8 layers of bidirectional LSTMs. Nevertheless, the fundamental problem remained the sequential processing itself, until Jakob Uszkoreit and their colleagues conceived developing a fully attention-based network without any recurrence.
| Model | Parameters | Architecture |
|---|---|---|
| Machine Translation Model | 380 million | Two LSTM Networks |
| Other Model | 130 million | Gated Recurrent Unit (GRU) |
| GNMT (Google) | Not applicable | 8 layers of Bidirectional LSTM |
The Internal Structure and Positional Encoding of the Transformer
To understand the functioning of a transformer, it is necessary to understand its data flow. When any text, image, or sound enters the model, it is first divided into digital pieces called tokens (a token being a word or a small part of it). Each of these tokens is converted by the embedding layer into a long collection of numbers, known as a vector (a sequence indicating mathematical dimensions).
However, a major challenge arises here: if all tokens are processed simultaneously, how does the model know the order of the words? For example, "Ram killed Ravan" and "Ravan killed Ram" contain the same words, but the meaning reverses when the order changes. Self-attention is fundamentally permutation-invariant. Therefore, the transformer adds positional encoding to the vector of each word, which indicates the exact position of the word in the sentence through mathematical signals based on sine and cosine waves.
The Mathematics of Attention: Query, Key, and Value
Scaled dot-product attention is the heart of the transformer. To understand it, a great analogy of a library can be used. When you go to the librarian and say, "I want a book on the history of Gujarat," this demand of yours is the 'Query' (Q). The subject label written on the spine of every book in the library is the 'Key' (K). The librarian compares your query with the keys of all books and gives you the actual information contained inside the book that matches the most, which is the 'Value' (V).
In the language of mathematics, when a dot-product multiplication of the query matrix Q and the key matrix K is performed, a relative score of every word with every other word is obtained. If the key dimension d_k is large, the value of the dot-product can become very large, causing gradients to vanish. Therefore, it is scaled by dividing it by the square root of d_k. Subsequently, a softmax function is applied to convert the scores into probabilities between 0 and 1, which indicate how much weight to place on which word. Finally, this weight is multiplied by the value matrix V.
- Dot-Product MultiplicationRaw relevance score is obtained by multiplying Query matrix (Q) and Key matrix (transpose of K).
- Scaling Factor DivisionThe score is divided by the square root of the key dimension (√d_k) to keep excessively large values under control.
- Softmax ProbabilityPrecise attention weights between 0 and 1 are determined by applying the softmax function to each row.
- Weighted Sum with ValueThe final context vector is generated by multiplying the calculated attention weights with the Value matrix (V).
Why is Multi-Head Attention Necessary?
A single sentence can have multiple types of relations for a single word. For example, in the sentence "The bank opened a new branch on the river bank," the grammatical relationship of the word 'bank' is with the verb 'opened', but its contextual meaning as a financial institution is linked with 'branch'. If there were only a single attention mechanism, it might only be able to focus on one aspect at a time.
To overcome this limitation, Multi-Head Attention was introduced into the transformer. In the original model with 100 million parameters, independent attention heads were used. Each such 'head' projects the word vectors into different dimensions and computes them concurrently in parallel. One head might keep track of grammar, a second head finds what pronouns refer to, and a third head captures the deeper meaning of the content. Finally, the results of all these heads are concatenated to prepare the final output matrix.
Causal Masking and Training Stability
In the encoder of the transformer, self-attention can freely look at all preceding and succeeding words, but at the decoder level this freedom can create a problem. When the model is generating a new sentence, it is imperative to ensure that it does not 'cheat' by looking at words that appear in the future. Causal masking is used for this. In this process, minus infinity (-∞) is placed in the positions of future words, so that after the softmax calculation their weight becomes zero and the model predicts the next word based solely on past words.
Initial plain transformer models were difficult to keep stable during training. The original research paper recommended learning rate warm-up, in which the learning rate was increased from zero to maximum up to the initial 2% of steps. A 2020 study found that placing layer normalization before the attention layer (Pre-LN) makes training very stable. Furthermore, techniques like dropout and label smoothing are also used to prevent overfitting.
Global Impact and Future Direction
The advent of the transformer architecture was not limited merely to language translation. In 2018, Google introduced encoder-only BERT, which began processing Google search queries starting in 2019. In 2020, Google Translate also removed the old system and adopted a transformer-based system. On the other hand, the decoder-only GPT series initiated by OpenAI since 2018 sparked a new revolution in generative AI worldwide in late 2022 in the form of ChatGPT.
Today, this very principle is used in Vision Transformers (ViT) for computer vision, robotics, audio, and advanced image and video generators like DALL-E of 2021, and Stable Diffusion 3 and Sora of 2024. To train massive models, data is distributed across thousands of GPU networks through tensor and pipeline parallelism. In standard self-attention, as the sentence length increases, computational cost increases at a quadratic rate, which is a known limitation of it. Yet, this attention journey that began in 2017 has enabled today's artificial intelligence to understand the nuanced depths of human language.
