Transformer Layers

The multi-head self-attention mechanism allows the model to jointly attend to information from different representation subspaces at different positions. It computes a weighted sum of values, where the weights (attention scores) are based on the compatibility between a query and a set of keys.

• Mechanism: For each token, it calculates how much to "pay attention" to every other token in the sequence.
• Heads: Multiple attention heads run in parallel, each learning different types of dependencies (e.g., syntactic, semantic).
• Output: The outputs from all heads are concatenated and linearly projected to form the sub-layer's final output.

The position-wise feed-forward network (FFN) is applied independently and identically to each token's representation. It consists of two linear transformations with a non-linear activation function (e.g., ReLU, GELU) in between.

• Function: Provides additional non-linear processing capacity to transform the attention output.
• Position-wise: The same FFN weights are used for every position in the sequence, making it highly efficient.
• Architecture: Typically expands the dimensionality (e.g., 512 to 2048) and then projects it back down, allowing for complex feature interactions.

A residual connection (or skip connection) adds the input of a sub-layer directly to its output. This creates a pathway for gradients to flow directly backward through the network, mitigating the vanishing gradient problem in deep architectures.

• Purpose: Enables the training of very deep transformer stacks (dozens of layers) by ensuring that information from earlier layers is not lost.
• Formula: Output = LayerNorm(x + Sublayer(x)), where x is the sub-layer input.
• Impact: Critical for the stability and convergence of models like BERT and GPT.

Layer normalization stabilizes the activations within a layer by normalizing the inputs across the feature dimension for each data point independently. It is applied after the residual addition.

• Operation: Normalizes the mean and variance of the summed inputs for each token.
• Benefit: Reduces training time and improves generalization by maintaining consistent activation scales.
• Placement: In the original Transformer paper, it is applied post-addition (LayerNorm(x + Sublayer(x))), though some variants (e.g., GPT) apply it before the sub-layer.

Cross-attention (or encoder-decoder attention) is a key component in transformer decoder layers, enabling modalities to interact. It allows queries from the decoder to attend to keys and values from the encoder's output.

• Use Case: Essential for sequence-to-sequence tasks like machine translation, image captioning, and multimodal fusion.
• Mechanism: The decoder's self-attention output serves as the query, while the encoder's final hidden states provide the keys and values.
• In VLAs: This mechanism is fundamental for vision-language-action models, allowing language or action tokens to attend to and ground themselves in visual features from an image encoder.

While the original Transformer uses a specific order, common variants exist that affect training dynamics and performance.

• Pre-LN: Layer normalization is applied before the sub-layer (e.g., x + Sublayer(LayerNorm(x))). This is more common in modern architectures (e.g., GPT) as it often leads to more stable training.
• Post-LN: The original scheme: LayerNorm(x + Sublayer(x)). Can be harder to train for very deep networks.
• Sandwich Norm: Applying layer normalization both before and after the sub-layer for additional stability. The choice influences gradient flow and is a key hyperparameter in model design.

The foundational Transformer Encoder Layer introduced by Vaswani et al. (2017) and used in BERT consists of two sub-layers:

• Multi-Head Self-Attention: Computes attention between all tokens in the input sequence, allowing the model to understand contextual relationships.
• Position-wise Feed-Forward Network: A small two-layer MLP applied independently to each token's representation. Each sub-layer is wrapped with residual connections and layer normalization, formulated as LayerNorm(x + Sublayer(x)). This design enables stable training of very deep networks.

Decoder-only layers, as used in the GPT family, are specialized for autoregressive text generation. Key modifications include:

• Masked Multi-Head Attention: Uses a causal attention mask to prevent positions from attending to future tokens, preserving the autoregressive property.
• Cross-Attention Omission: Standard decoder layers omit the encoder-decoder cross-attention sub-layer, as they are not part of a sequence-to-sequence architecture.
• Layer Order: Typically uses Pre-Layer Normalization (applying normalization before the sub-layer), which improves training stability for very large models.

Models like T5 and BART use the full encoder-decoder transformer architecture for sequence-to-sequence tasks (e.g., translation, summarization). Each decoder layer contains three core sub-layers:

1. Masked Self-Attention over previous decoder outputs.
2. Cross-Attention where decoder queries attend to the final encoder representations (keys/values).
3. Feed-Forward Network. This allows the decoder to condition its text generation on the full, bi-directional context from the encoded source input.

ViT adapts the standard transformer encoder for images by treating non-overlapping image patches as a sequence of tokens. The layers are identical to a BERT encoder but process a fundamentally different input:

• Patch Embeddings: Linear projection of image patches replaces word embeddings.
• Position Embeddings: Added to retain spatial information since the transformer itself is permutation-invariant.
• Class Token: A special [CLS] token prepended to the sequence aggregates global image representation for classification. The success of ViT demonstrated the transformer's generality beyond sequential data.

Models like CLIP use a dual-encoder architecture with independent transformer layers for each modality:

• Image Encoder: Often a ViT or CNN-based transformer.
• Text Encoder: A standard transformer encoder like BERT.
• Contrastive Training: The layers are trained not to fuse modalities internally, but to project image and text into a joint embedding space using a contrastive loss (InfoNCE). Fusion occurs at the task level via similarity search in this shared space, enabling zero-shot transfer.

The number of transformer layers (depth) is a primary scaling factor. Notable configurations:

• BERT-Base: 12 encoder layers.
• BERT-Large: 24 encoder layers.
• GPT-3: 96 decoder layers.
• PaLM: Up to 118 decoder layers. Increasing depth improves model capacity and reasoning but introduces challenges like vanishing gradients and training instability, addressed by advanced normalization (e.g., RMSNorm) and initialization schemes. The optimal depth is often determined empirically given compute budgets.

Self-attention is the fundamental mechanism within a transformer layer that allows each token in a sequence to attend to all other tokens. It computes a weighted sum of value vectors, where the weights are determined by the compatibility (dot product) between the query vector of the current token and the key vectors of all tokens. This enables the model to capture long-range dependencies and contextual relationships within the input, regardless of distance.

• Scaled Dot-Product Attention: The standard formulation, which scales the dot products by the square root of the key dimension to prevent vanishing gradients.
• Multi-Head Attention: Runs multiple self-attention operations in parallel, allowing the model to jointly attend to information from different representation subspaces.

The position-wise feed-forward network (FFN) is the second core sub-layer in a standard transformer block. It is applied independently and identically to each token position. Typically, it consists of two linear transformations with a non-linear activation function (like GELU or ReLU) in between, serving to further process and transform the representations output by the attention mechanism.

• Architecture: FFN(x) = W₂ * Activation(W₁ * x + b₁) + b₂, where W₁ expands the dimensionality (e.g., 4x) and W₂ projects it back down.
• Function: Provides capacity for complex, non-linear transformations that are applied per token, complementing the token-mixing function of self-attention.

Residual connections and layer normalization are critical stabilization techniques applied within each transformer sub-layer. A residual connection adds a sub-layer's input directly to its output (x + Sublayer(x)), which mitigates the vanishing gradient problem and enables the training of very deep networks. Layer normalization is applied to the sum, stabilizing the activations' mean and variance.

• Pre-LN vs. Post-LN: In the original Post-LN transformer, normalization is applied after the residual addition. Modern architectures like GPT often use Pre-LN, where normalization is applied before the sub-layer, leading to more stable training.
• Effect: These components are essential for enabling the effective training of transformers with dozens or hundreds of layers.

Cross-attention (or encoder-decoder attention) is a variant of the attention mechanism where the queries come from one sequence (e.g., the decoder) and the keys and values come from another sequence (e.g., the encoder output). It is a foundational component for sequence-to-sequence tasks and is a key mechanism for multimodal fusion, allowing one modality (like language) to attend to and integrate information from another (like vision).

• Use in Transformers: Central to the original transformer architecture for machine translation.
• Use in VLMs: In vision-language models, cross-attention layers allow text tokens to attend to visual feature patches, enabling grounded reasoning and generation.

The standard transformer layer has inspired numerous architectural variants optimized for efficiency, scale, or specific data types.

• Decoder-Only Block: Used in autoregressive LLMs like GPT. It features a masked self-attention sub-layer (which prevents attending to future tokens) followed by the FFN. It lacks cross-attention.
• Sparse/Mixture-of-Experts (MoE): Replaces the dense FFN with a router that selects a few experts from a large pool for each token, dramatically increasing model capacity without a proportional increase in compute per token.
• Parallel Attention: An architectural tweak where the self-attention and FFN sub-layers are computed in parallel rather than sequentially, reducing critical path length and potentially training faster.

Positional encodings are injected into the input embeddings to provide the transformer model with information about the order of tokens in a sequence, which it inherently lacks due to the permutation-invariant nature of self-attention.

• Absolute (Sinusoidal): The original method uses fixed, pre-defined sine and cosine functions of different frequencies.
• Learned (Absolute): Uses a trainable lookup table for each position index, common in models like BERT.
• Relative Positional Encoding: Encodes the relative distance between token pairs (e.g., key_position - query_position), which can generalize better to longer sequences than absolute encodings. Used in models like T5 and Transformer-XL.