Transformer architecture

Transformer architecture is a neural network architecture introduced in 2017 by Google researchers in the paper "Attention Is All You Need". It revolutionized the field of natural language processing (NLP) and has become the foundation for most modern large language models (LLMs), such as BERT, GPT, and Gemini. The key innovation of the Transformer is the self-attention mechanism, which allows the model to weigh the importance of different parts of the input data and process sequences in parallel, abandoning the recurrence inherent in RNNs and LSTMs.

Historical Context and Prerequisites

Before 2017, the dominant architectures for processing sequential data, such as text, were recurrent neural networks (RNNs) and their advanced variant, long short-term memory (LSTM) networks.

Problems with RNN/LSTM Addressed by the Transformer

• Sequential Processing Limitations: RNNs and LSTMs process data token by token, which prevents intra-sequence parallelism and slows down training on large datasets.
• Vanishing and Exploding Gradients Problem: In long sequences, gradients propagated backward through many steps can either vanish or explode, making it difficult to train long-term dependencies.
• Long-Term Dependencies: Information from the beginning of a sequence can be lost by the time the end is reached.

The Transformer's approach is a complete departure from recurrence in favor of an attention mechanism. It provides a constant dependency path length between any two positions ($O(1)$), which facilitates the modeling of long-range dependencies, although the basic implementation of self-attention has a quadratic computational complexity with respect to the sequence length ($O(n^2)$). The vanishing/exploding gradients problem doesn't "disappear" but is mitigated by residual connections, LayerNorm, and the training regimen; modern implementations often use the Pre-LayerNorm (Pre-LN) variant for greater training stability.

Architecture and Key Components

The original Transformer architecture consists of two main parts: an encoder and a decoder. Both components are stacks of identical layers ($N=6$ in the original paper).

• Encoder Layer Structure: (1) multi-head self-attention (MHA), (2) a position-wise feed-forward network (FFN); each sub-layer is surrounded by a residual connection and LayerNorm.
• Decoder Layer Structure: (1) masked self-attention (a causal mask prevents access to future positions), (2) cross-attention to the encoder outputs, (3) an FFN—also with residuals and LayerNorm.

Attention Mechanism and Self-Attention

The attention mechanism computes a weighted sum of value vectors (Value), where the weights are determined by the compatibility of keys (Key) with queries (Query). The Transformer uses Scaled Dot-Product Attention:

$\mathrm{A}\mathrm{t}\mathrm{t}\mathrm{e}\mathrm{n}\mathrm{t}\mathrm{i}\mathrm{o}\mathrm{n}(Q,K,V)=\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}\left(\frac{Q{K}^{\mathrm{\top }}}{\sqrt{d_k}}\right)V$

Where $d_k$ is the dimension of the keys/queries; dividing by $\sqrt{d_k}$ prevents the softmax function from saturating. When $Q$, $K$, and $V$ are generated from the same sequence, the mechanism is called self-attention.

Multi-Head Attention

Instead of a single set of matrices $(W_Q,W_K,W_V)$, $h$ parallel "heads" are used, each of which projects $Q,K,V$ into lower-dimensional subspaces, computes attention independently, and then the results are concatenated and projected:

$\mathrm{M}\mathrm{u}\mathrm{l}\mathrm{t}\mathrm{i}\mathrm{H}\mathrm{e}\mathrm{a}\mathrm{d}(Q,K,V)=\mathrm{C}\mathrm{o}\mathrm{n}\mathrm{c}\mathrm{a}\mathrm{t}({\text{head}}_1,\dots ,{\text{head}}_h)W^O,\text{where }{\text{head}}_i=\mathrm{A}\mathrm{t}\mathrm{t}\mathrm{e}\mathrm{n}\mathrm{t}\mathrm{i}\mathrm{o}\mathrm{n}(Q{W}_i^Q,K{W}_i^K,V{W}_i^V).$

Variants for Inference Acceleration:

• MQA (Multi-Query Attention): All heads share a single key/value pair, which significantly reduces the size and memory traffic of the KV cache during decoding.
• GQA (Grouped-Query Attention): A compromise between MHA and MQA—several groups of heads share K/V; its quality is close to MHA with MQA-like speed.

Positional Encoding

Since self-attention is invariant to token order, positional encodings are added to the input embeddings.

• Original sinusoidal encodings (PE) from:

$\text{PE}(\text{pos},2i)=\sin (\text{pos}/10000^{2i/d_{\text{model}}}),\text{PE}(\text{pos},2i+1)=\cos (\text{pos}/10000^{2i/d_{\text{model}}}).$

• Modern relative/rotary variants:
  • RoPE (Rotary Position Embeddings) encodes relative shifts by rotating the $Q$/$K$ vectors; it is used in several modern LLMs.
  • ALiBi introduces a linear bias to attention scores, which improves extrapolation to lengths longer than those seen during training.

Position-wise FFN, Residuals, and Normalization

Each encoder and decoder layer, in addition to attention, contains a position-wise feed-forward network (FFN):

$\mathrm{F}\mathrm{F}\mathrm{N}(x)=\max (0,x{W}_1+b_1)W_2+b_2.$

Residual connections and Layer Normalization are used around each sub-layer: $\mathrm{L}\mathrm{a}\mathrm{y}\mathrm{e}\mathrm{r}\mathrm{N}\mathrm{o}\mathrm{r}\mathrm{m}(x+\mathrm{S}\mathrm{u}\mathrm{b}\mathrm{l}\mathrm{a}\mathrm{y}\mathrm{e}\mathrm{r}(x))$. The original paper used the Post-LN variant; modern LLMs often use Pre-LN for better training stability and less reliance on a long warm-up period.

Evolution and Modern Variants

Early methods based on recurrent neural networks (RNNs) and their advanced variants, such as LSTMs, processed text sequentially, one token at a time. Although this approach intuitively matched the structure of language, it created a significant limitation: it hindered parallel computation and made it difficult to capture dependencies between elements that were far apart in the text. In 2017, a group of researchers from Google introduced a paper titled "Attention Is All You Need." In it, they described a new architecture—the "Transformer." This model was the first to completely abandon the use of recurrent neural networks, replacing them with an "attention" mechanism. The main innovation was that the attention mechanism allowed the Transformer to weigh the importance of every word in the input sequence when generating the corresponding word in the output. Crucially, the model could process all words simultaneously. This capability for parallel processing made it possible to train much larger models on vast amounts of data. This resulted in the emergence of modern large language models (LLMs).

The Transformer architecture served as the basis for numerous models, which can be broadly divided into three classes.

1. Encoder-only Models

• Example: BERT (and RoBERTa, ALBERT).
• Principle: Pre-training on a masked language modeling (MLM) task with a bidirectional context.
• Application: Understanding tasks (classification, NER, etc.).

2. Decoder-only Models

• Example: The GPT series (GPT-1/2/3), LLaMA, Claude.
• Principle: Causal language modeling (CLM)—predicting the next token; a causal mask is applied to the attention mechanism.
• Application: Text generation, dialogue, and code.

3. Encoder-decoder Models

• Example: The original Transformer, T5, BART.
• Principle: The encoder builds a representation of the input, and the decoder generates the output using cross-attention to the encoder's features.
• Application: Seq2seq tasks (translation, summarization, etc.).

4. Multimodal and Alternative Architectures

• Vision Transformer (ViT)—an adaptation for images (by splitting them into patches); Swin Transformer—a hierarchical model using shifted windows.
• Alternatives for long sequences:
  • Mamba—selective state space models (SSMs) with linear complexity.
  • RWKV—an RNN-like architecture with parallelizable training and linear inference complexity.
  • Hybrids (e.g., Jamba): Alternate between Transformer and Mamba blocks; sometimes supplemented with MoE.

Training and Optimization Techniques

The effectiveness of the Transformer is closely tied to training techniques and infrastructure.

• Pre-training Strategies: CLM and MLM; also contrastive and denoising objectives (ELECTRA, T5).
• Fine-tuning Techniques:
  • Full fine-tuning of all parameters.
  • Parameter-Efficient Fine-Tuning (PEFT): LoRA introduces low-rank adapters while keeping the base weights frozen.
• Behavioral Alignment: RLHF—Reinforcement Learning from Human Feedback.
• System-level Inference Optimizations: PagedAttention/vLLM increases serving throughput via paged management of the KV cache; especially useful for long sequences and large batches.

Links

• The Illustrated Transformer — A visual explanation of the Transformer architecture

Literature

• Vaswani, A., Shazeer, N., Parmar, N., et al. (2017). Attention Is All You Need. NeurIPS. arXiv:1706.03762.
• Devlin, J., Chang, M.-W., Lee, K., Toutanova, K. (2019). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. arXiv:1810.04805.
• Radford, A., Narasimhan, K., Salimans, T., Sutskever, I. (2018). Improving Language Understanding by Generative Pre-Training. OpenAI Technical Report.
• Brown, T. B., Mann, B., Ryder, N., et al. (2020). Language Models Are Few-Shot Learners. NeurIPS. arXiv:2005.14165.
• Raffel, C., Shazeer, N., Roberts, A., et al. (2019). Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. arXiv:1910.10683.
• Dosovitskiy, A., Beyer, L., Kolesnikov, A., et al. (2020). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. arXiv:2010.11929.
• Liu, Z., Lin, Y., Cao, Y., et al. (2021). Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. arXiv:2103.14030.
• Tay, Y., Dehghani, M., Bahri, D., Metzler, D. (2020). Efficient Transformers: A Survey. arXiv:2009.06732.
• Xiong, R., Yang, Y., He, D., et al. (2020). On Layer Normalization in the Transformer Architecture. ICML. arXiv:2002.04745.
• Su, J., Lu, Y., Pan, S., et al. (2021). RoFormer: Rotary Position Embedding. arXiv:2104.09864.
• Press, O., Smith, N. A., Lewis, M. (2021). Train Short, Test Long: Attention with Linear Biases (ALiBi). arXiv:2108.12409.
• Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need (MQA). arXiv:1911.02150.
• Ainslie, J., Lee-Thorp, J., de Jong, M., et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP. arXiv:2305.13245.
• Kwon, W., Li, Z., Zhuang, S., et al. (2023). Efficient Memory Management for LLM Serving with PagedAttention (vLLM). arXiv:2309.06180.
• Touvron, H., Lavril, T., Izacard, G., et al. (2023). LLaMA: Open and Efficient Foundation Language Models. arXiv:2302.13971.
• Gu, A., Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752.
• Peng, B., et al. (2023). RWKV: Reinventing RNNs for the Transformer Era. arXiv:2305.13048.
• Lieber, O., Lenz, B., Bata, H., et al. (2024). Jamba: A Hybrid Transformer-Mamba Language Model. arXiv:2403.19887.
• Hu, E. J., Shen, Y., Wallis, P., et al. (2021). LoRA: Low-Rank Adaptation of Large Language Models. arXiv:2106.09685.
• Ouyang, L., Wu, J., Jiang, X., et al. (2022). Training language models to follow instructions with human feedback. OpenReview.