Attention and transformers diagrams

The original (encoder-decoder) Transformer architecture features separate input text and output text which are processed by the decoder and encoder respectively. The decoder uses (unmasked) self-attention and produce final token embeddings. The encoder uses both self-attention (often masked/causal) and cross attention (to the output of the final encoder block).

This diagram depicts one the two pre-training tasks of BERT (Masked Language Model): some tokens are masked and the model must predict them.

In contrast to subsequent models, it uses Post-Norm (the LayerNorm is after the residual branch instead of being at the beginning of the residual branch).

LLama 3 is an example of a typical “modern” LLM. Compared to GPT-2, it uses RMSNorm instead of LayerNorm, GQA instead of MHA, RoPE instead of learned positional encoding, a SiLU FFN, untied weights for token unembedding and does not use biases.

The d_t query vectors (Q_i) independently attend to the d_s key-value pairs (K_i, V_i). Each query vector (Q_i) is transformed into a weighted sum (and more precisely a convex combination) of the d_s value vectors (V_i). The attention weights W_{i,j} are related to the (dot-product) similarity between the Q_i queries and the K_j keys.

Dot-product attention with QK norm is actually cosine attention with a learnable temperature parameter.

An typical attention module, includes projections for the query (Q), keys (K) and values (V) as well as an output projection.

With MHA, the attention is split between different attention heads. Each attention head can attend to different aspects of the tokens. A conventional design uses d_K = d_V = d_m / n_h.

The previous diagram represents per head QKV projections (Q_h = X_Q \, W_{Q,h} + b_{Q,h}). In practice, the QKV projections are usually implemented as a packed projection (Q = X_Q \, W_Q + b_Q). This second representation is especially useful in the context of LoRA: the low-rank factorization is done on the packed projection matrix W'_Q = W_Q + A_Q B_Q and not per-head in the original paper. See TensLoRA: Tensor Alternatives for low-rank adaptation for discussion on this topic.

MHA can (equivalently) be understood as having per head output projections. When used in a residual structure, each attention head output projection is added to the residual stream (y = x + \sum_h h(x) + b).

MQA reduces the size of the KV cache by using the same K and V for all attention heads.

GQA is an intermediate between MHA (n_h=1) and MQA (n_g=1). The same K and V are shared between several attention heads.

In MLA, a shared compressed latent representation of both keys and values (C_{KV}) is stored in cache.

Through some linear algebra manipulations, we can do the operations in the compressed latent space.

Query (Q) is decomposed into a content query (Q_C) and a positional one (Q_R). Same for the key. The cache stores the positional key K_R and the compressed (content) KV C_{KV}.

DSA uses a indexer heads to select K positions (Top K) to attend to for a given position (sparse attention). Here DSA is instanciated on top of MHA for simplicity.

The DSA paper instantiates DSA on top of MLA with RoPE. This diagram show how RoPE is integrated into the DSA indexer.

In self-attention, the source sequence and the target sequence are the same (and therefore d_t = d_s). For auto-regressive models (eg. LLMs), causal attention prevents a position from attending to future positions.

Deformable DETR uses deformable attention. For each query, each attention heads attends to a restricted number (K) of sampling positions of a feature map. Both the sampling positions and the attention weights are learned projections of the query.

Multi-scale Deformable attention module extends Deformable attention with support for multi-scale feature maps. For each query, each attention head attends to K sampling position per scale level.

In RPB, learned (or fixed) positional encoding is adding to the attention scores (before the softmax). These biases depend on the relative position of the tokens.

In this scheme, (learned or fixed) positional encoding is added before Q and K projection. The same positional encoding is used for all attention heads and for all attention blocks. For self-attention, the same positional embeddings are used for Q and K projections.

When using RoPE, the positional information is typically included in the Q and K of each attention block. The positional information is encoded by applying 2D rotations to pairs of features at different frequencies.

A typical design uses d_h = 4 × d_x.

Mixtral replaces the FFN with sparsely-gated Mixture of Experts (MoE) where only two experts are active for each token (top 2).

DeepSeekMoE replaces the FFN with sparsely-gated Mixture of Experts (MoE) where only K experts out of N are active for each token (top K). In addition, they include a small number of “shard experts“ which are always active (and are not gated).

Gemma 4 uses per-layer token embedding (PLE). It is included through an additional gated residual sub-block at the end of the transformer block. This (comparatively-)small-dimension per-layer token embedding is combined with a down projection of the global token embedding.

The original MTP design forks several independent prediction heads before the last transformer block, predicting respectively tokens at position n+1, n+2, etc.

DeepSeek-V2 uses a chained MTP. The MTP prediction for token at position n+k uses the sampled tokens at positions n+1 to n+1-k.

EAGLE-3 is an auto-regressive draft model. It takes as input the hidden features of the main model at three different levels and the token embeddings.

The Perceiver pattern can be used to handle large input sequences. It mitigates the quadratic complexity (wrt. sequence length N) of the self-attention mechanism (Q \, K^T). It introduces a sequence of learned latent vectors of size M \ll N which are used to extract information from the input through cross-attention. The resulting features are then used in a typical transformer. This reduces the complexity from \mathcal{O}(L \, N^2) to \mathcal{O}(M \, N + L \, M^2) . This process can be repeated several times either using the same input data (in order to extract more information from it) or additional data.

Perceiver IO adds another cross-attention block (decoder) where output queries attend to the latent features. For simple tasks, the output queries can be learned. For other tasks, the output queries can contains features describing the “task”. It is usually better to omit the skip connection around the output cross-attention (as depicted here) in this case.

The Perceiver Resampler uses a Perceiver architecture to downscale a high number of visual token (eg. coming from a video) into a fixed number (R) of visual tokens. R latent queries attend to the visual tokens and to themselves in a mixed self/cross-attention operation.

The input image is partitioned into (non-overlapping) patches of P×P pixels. These patches are seen as a sequence of tokens: each patch is first projected in the model embedding space and the resulting patch tokens are processed using a classic transformer decoder model. A special learned class token ([CLS]) is inserted at the beginning of the sequence: the final features of this token are fed into a classification head.

DeiT follows the architecture of ViT modified for the efficient distillation from a teacher model (eg. a ConvNet). An additional learned distillation token ([DIST]) is introduced. The student model is trained to match the teacher's predicted class through the distillation token (hard-label distillation), while simultaneously tracking the ground truth label through the class model.

At inference, the reference source code takes the mean of the logits of both heads.

Swin encourages some form of locality by restricting the attention within windows of 4×4 patches and by using a hierarchical structure (patch merging). It produces a pyramid of multi-scale feature maps.

MAE (Masked autoencoders) uses self-supervised (pre-)training using an autoencoder network. The input image is severely degraded by masking a large proportion of patches. The masked patches are omitted from the sequence in the encoder and replaced by a learned masked token in the decoder. The network is trained must reconstruct the masked patches.

DINO (Distillation with no labels) uses self-supervised learning method through self-distillation (for pre-training). A student network is trained to follow the output of a teacher model (distillation). The student and the teacher receive a different “part” (crop/transformation) of the same image and produce a probability distribution in a high-dimensional “prototype” space. The student is trained to produce prototype distribution similar to the the distribution produced by the teacher, even if has received a different view of the same image. This forces the network to extract semantic information out of the image. The trick is that the teacher network is derived from previous version of the student network (self-distillation).

Notice the usage of weight normalization (torch.nn.utils.weight_norm, torch.nn.utils.parametrizations.weight_norm) in the output layer.

DINOv2 use the same style of self-distillation training as DINO but adds an iBOT-style objective. An additional head (the iBOT head) with the same architecture is added to the network. The second student head processes images where some patch tokens have been masked: it must learn to reproduce the prototypes produced by the teacher network (with unmasked data) for these masked patches.

The teacher network is not displayed in this diagram.

DETR is a transformer-based detection head which can plugged to an existing visual backbone (such as a ViT or a CNN). Compared to other detection heads, it does not requires Non Max Suppression (NMS) on the resulting bounding boxes because it is trained to avoid creating duplicate detections.

Deformable DETR but replaces the MHA in the encoder self-attention and the decoder cross-attention with multi-head Deformable attention modules. A variant, not depicted here use mutli-scale deformable attention (attending to multi-scale feature maps).

Maskformer uses a multi-scale backbone (/4, /8, /16, /32). The multi-scale feature maps are fed into a pixel decoder (FPN-like CNN) which produce final pixel embeddings as /4 scale. A transformer decoder uses fixed segment queries which attends to the low-resolution feature map (/32). The output of the transformer decoder is fed into two heads in order to produce two outputs. First, a classification head produces class prediction (classification) for each candidate mask/segment. Second, a MLP is used to produce per-mask embeddings for each segment query: the spatial masks (logits) are obtained by matching (dot product) the per-mask embeddings produced by the MLP with the pixel embeddings produced by the pixel decoder.

This diagrams details the FPN-like pixel decoder.

In CLIP (Contrastive Language–Image Pre-training), is trained using cross-modal contrastive learning. A visual model and a text models are jointly pre-trained using pairs (image, label) pairs. For each batch of B (image, label) pairs, the output of the visual encoder (a ViT or CNN) and the text encoder (an auto-regressive transformer decode model) are projected into a shared multi-modal embedding space. The model is trained to maximize the similarity of the representations of matching images and labels (and minimize the similarity of the representations of unmatched images and labels).

In the first stage of pre-training, the Qformer (Querying Transformer) learns to extract useful information from the output of frozen image decoder. The Qformer is initialized from an pre-trained LLM transformer (here BERT).

In the second stage of pre-training, the Qformer is plugged into an existing frozen LLM and must learn to adapt the information from the visual encoder to visual tokens (soft visual prompt) understood by the LLM.

LLaVA uses a simple linear layer for projecting the visual patches (coming from a frozen image encoder) into the LLM token feature space. In the first stage (pre-training), only this small layer is trained (both vision encoder and LLM are frozen). In the second stage (fine tuning), both the linear adapter and the LLM are fine-tuned while the visual encoder stay frozen.

LLaVA 1.5 uses a (still very simple) FFN/MLP instead of the plain linear adapter.

CLM is basically CLIP applied to state/action pairs with a frozen LLM instead of the image encoder and a deep projection network. After some pre-training, it is used a zero-shot action classifier.

This diagram represent the first pre-training phase (contrastive pre-training). After, that the model is also trained to distinguish the correct action among a selection of possible choices (synthetic hard negatives).