Demystifying Transformers: A Step‑by‑Step Guide to Self‑Attention and Architecture

Introduction

Google's BERT model achieved state‑of‑the‑art results on many NLP tasks, largely thanks to the Transformer architecture. Originally designed for machine translation, the Transformer replaces slow RNNs with fast, parallelizable self‑attention layers and can be stacked deeply to improve accuracy.

Model Overview

The Transformer, introduced in the paper Attention Is All You Need , is now a reference model on Google Cloud TPUs. Implementations are available in TensorFlow (Tensor2Tensor) and PyTorch.

Encoder and Decoder

The model consists of an encoder stack and a decoder stack connected by attention layers. Each encoder contains identical encoder blocks (typically six), and each decoder contains the same number of decoder blocks.

Encoder blocks process input vectors through a self‑attention layer followed by a feed‑forward network. Decoder blocks have an additional encoder‑decoder attention layer.

Self‑Attention Mechanism

Each word is first embedded into a 512‑dimensional vector. The self‑attention layer creates three new vectors for each word: a query, a key, and a value (typically 64‑dimensional). Scores are computed by dot‑product of queries with all keys, scaled by √64, and passed through softmax to obtain attention weights.

These weights weight the value vectors, which are summed to produce the self‑attention output. The process can be expressed as a matrix operation:

Attention(Q,K,V)=softmax(QKᵀ/√d_k)·V

Multi‑Head Attention

Instead of a single attention head, the Transformer uses eight parallel heads. Each head has its own query, key, and value weight matrices, allowing the model to attend to information from different representation subspaces.

The outputs of all heads are concatenated and projected with a final weight matrix to produce a single combined representation.

Positional Encoding

Since the model contains no recurrence, positional encodings are added to the word embeddings to give the network a sense of order. The encodings use sinusoidal functions that can extrapolate to sequence lengths longer than those seen during training.

Residual Connections and Layer Normalization

Each sub‑layer (self‑attention, feed‑forward, and encoder‑decoder attention) is wrapped with a residual connection followed by layer‑normalization, which stabilizes training.

Training and Loss Function

During training, the model predicts a probability distribution over the target vocabulary for each position. The cross‑entropy loss (or KL‑divergence) compares the predicted distribution with the one‑hot ground‑truth vector.

Typical vocabularies contain thousands of tokens; the final linear layer projects the decoder output to logits of that size, followed by a softmax to obtain probabilities.

Inference (Beam Search)

At inference time, the model generates one token at a time. Greedy decoding selects the highest‑probability token, while beam search keeps the top‑k hypotheses (e.g., k=2) and expands them, improving translation quality.

Further Reading

For deeper study, read the original Attention Is All You Need paper, the Google Transformer blog post, the Tensor2Tensor announcement, and follow‑up works such as Depthwise Separable Convolutions, One Model To Learn Them All, Discrete Autoencoders for Sequence Models, Image Transformer, and training tips for Transformers.