Blogs
Transformer Architecture

Transformer Architecture

The Transformer architecture, introduced by Vaswani et al. in the paper “Attention is All You Need”, revolutionized the field of Natural Language Processing (NLP) and deep learning models for sequence transduction tasks. Unlike earlier recurrent architectures like LSTMs and GRUs, the Transformer leverages a mechanism called self-attention and completely eschews recurrence, leading to better parallelization and handling of long-range dependencies.

Key Components of the Transformer

The Transformer model is built from an encoder-decoder architecture, where both the encoder and decoder are composed of multiple identical layers stacked on top of each other.

1. Encoder-Decoder Overview 

graph TD;
    Input-->Encoder;
    Encoder-->Intermediate_Representation;
    Intermediate_Representation-->Decoder;
    Decoder-->Output;

The encoder processes the input sequence into a continuous representation, and the decoder uses this representation along with its own input to generate the final sequence.

  • Encoder: Converts the input sequence into a fixed-size continuous representation.
  • Decoder: Generates output sequences based on the representation provided by the encoder.

Each encoder and decoder is made up of multiple layers of self-attention mechanisms and feedforward neural networks.

2. Self-Attention Mechanism

The core of the Transformer is the self-attention mechanism, which allows each word to attend to other words in a sequence, effectively capturing dependencies regardless of their distance. It computes a weighted sum of values (V) by calculating the compatibility of a query (Q) with corresponding keys (K).

The self-attention can be formulated as:

$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V $$

Where:

  • $ Q $ = Query
  • $ K $ = Key
  • $ V $ = Value
  • $ d_k $ = Dimension of the key vectors

This operation allows the model to attend to different parts of the input, capturing context more effectively.

graph LR;
    subgraph Self-Attention
    Input[Input Tokens]-->QKV[Generate Q, K, V];
    QKV-->Attention[Attention Mechanism];
    Attention-->Weighted_Sum[Weighted Sum];
    Weighted_Sum-->Output[Attention Output];
    end

3. Multi-Head Attention

Instead of applying a single attention mechanism, the Transformer uses multiple attention heads to jointly attend to information from different subspaces at different positions. This is called Multi-Head Attention.

For each attention head:

$$ \text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) $$

Where $W_i^Q$, $W_i^K$, and $W_i^V$ are projection matrices for the i-th head.

The outputs from each head are concatenated and projected again, enabling the model to capture different types of relationships in the input data.

graph TD;
    subgraph Multi-Head Attention
    QKV-->Head1[Attention Head 1];
    QKV-->Head2[Attention Head 2];
    QKV-->Head3[Attention Head 3];
    QKV-->HeadN[Attention Head N];
    Head1-->Concat[Concatenate Heads];
    Head2-->Concat;
    Head3-->Concat;
    HeadN-->Concat;
    Concat-->FinalOutput[Projected Output];
    end

4. Position-wise Feedforward Networks

Each encoder and decoder layer also includes a fully connected feedforward network applied to each position separately. This consists of two linear transformations with a ReLU activation in between:

$$ \text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2 $$

This component introduces non-linearity and increases the representational capacity of the model.

5. Positional Encoding

Since the Transformer model doesn’t have a built-in mechanism to understand the order of tokens (unlike RNNs), positional encodings are added to the input embeddings to inject information about the relative positions of tokens in the sequence.

The positional encoding for the i-th position and j-th dimension is defined as:

$$ PE(i, 2j) = \sin\left(\frac{i}{10000^{2j/d_{model}}}\right) $$ $$ PE(i, 2j+1) = \cos\left(\frac{i}{10000^{2j/d_{model}}}\right) $$

These encodings are added to the input embeddings to provide the model with information about the sequence order.

6. Layer Normalization and Residual Connections

To aid with optimization and training stability, each sub-layer (e.g., attention and feedforward) in the Transformer is surrounded by residual connections and followed by layer normalization.

The output of each sub-layer is:

$$ \text{LayerNorm}(x + \text{sublayer}(x)) $$

Residual connections help in training deeper networks by ensuring that gradients flow efficiently during backpropagation.

7. Decoder Architecture

The decoder mirrors the encoder but includes additional masking to ensure that the model cannot “peek” at future tokens during training. This is implemented through a masked multi-head attention mechanism.

graph LR;
    Input[Input Tokens]-->Masked_Multihead[Masked Multi-Head Attention]-->Multihead[Multi-Head Attention];
    Multihead-->Feedforward[Feedforward Layer];
    Feedforward-->Output[Final Output];

The decoder also incorporates encoder-decoder attention, where it attends to the output of the encoder while generating each token.

Summary of Key Features

  • Self-Attention: Allows each token to attend to all others in the sequence, capturing long-range dependencies.
  • Multi-Head Attention: Parallel attention heads allow for different aspects of input relationships to be captured.
  • Positional Encoding: Provides the model with positional information, crucial for understanding sequence structure.
  • Layer Normalization and Residual Connections: Facilitates training and ensures stable gradient flow.
  • Parallelization: The absence of recurrence allows for significant parallelization, speeding up training time compared to RNNs.

Conclusion

The Transformer architecture has been instrumental in shaping modern NLP and deep learning techniques. Its attention-based mechanism enables it to capture complex relationships between input elements efficiently. Furthermore, the model’s scalability and parallelization capabilities make it highly effective for large-scale training on contemporary hardware.

Transformers have since been the foundation for many state-of-the-art models such as BERT, GPT, and T5, which continue to push the boundaries of NLP and machine learning research.

Last updated on
Probabilistic Graphical Models