Transformer Decoder Architecture

Introduction

The transformer architecture, introduced in Attention Is All You Need (Vaswani et al., 2017), has become the foundation of modern large language models. While the original paper described an encoder-decoder model for machine translation, today’s LLMs use only the decoder half, stacking it dozens or even hundreds of layers deep and scaling parameters into the hundreds of billions.

Understanding the internals of this architecture is essential for anyone working on training, serving, or optimizing LLMs. Every decision in distributed training (how to shard model weights, where to place pipeline stage boundaries, whether to recompute activations) traces back to the per-layer cost profile.

This post walks through each component of the decoder-only transformer, providing:

• The math equations that define each layer’s computation
• Weight parameter counts — how many learnable parameters each layer contributes
• FLOP calculations — the computational cost of each stage in the forward pass
• Activation memory — which tensors must be stored during forward pass for backpropagation

By the end, we derive the well-known results: a transformer layer has $\approx 12d^2$ parameters, costs $24bsd^2 + 4bs^2d$ forward FLOPs, and training requires roughly $6P$ FLOPs per parameter per token.

Notation

FLOP convention: A matrix multiply of $(m, k) \times (k, n)$ costs $2mkn$ FLOPs (each of the $m \times n$ output elements requires $k$ multiply-accumulate operations).

For each layer, we analyze:

• Weight parameters — number of learnable parameters
• Activation memory — tensors stored during forward pass for backpropagation
• FLOPs — floating-point operations in the forward pass

Architecture Overview

A decoder-only transformer (GPT-style) consists of:

1. Token Embedding — maps token IDs to dense vectors
2. Positional Encoding — injects sequence position information
3. $L$ Decoder Layers, each containing:
   • Layer Norm (Pre-Norm)
   • Multi-Head Self-Attention (with causal mask)
   • Residual Connection
   • Layer Norm
   • Feed-Forward Network (MLP)
   • Residual Connection
4. Final Layer Norm
5. Output Projection (LM Head) — maps hidden states to vocabulary logits

Embedding

The token embedding layer maps each input token ID to a dense vector of dimension $d$.

Weight parameters: The embedding matrix $W_E \in \mathbb{R}^{v \times d}$ contains $vd$ parameters.

Forward pass: Given a sequence of token IDs, the embedding is a table lookup (indexing rows of $W_E$), producing output of shape $(b, s, d)$.

FLOPs: Essentially zero — a gather operation, not a matrix multiply.

Activation memory: Output tensor $(b, s, d)$.

Weight tying: Some models share the embedding matrix $W_E$ with the output projection (LM head), so the $vd$ parameters are counted only once. GPT-2 and early GPT-3 used weight tying; most modern models do not tie weights.

Positional Encoding

Multi-head attention is permutation-equivariant — it cannot distinguish whether one input token comes before another. Since word order is critical for language understanding, we must inject position information.

Sinusoidal Positional Encoding (Original Transformer)

Rather than learning an embedding for every possible position (which would not generalize to unseen lengths), the original transformer uses fixed sinusoidal functions that encode position as a pattern the network can identify:

$PE_{(pos,i)}$ is the positional encoding at sequence position $pos$ and hidden dimension index $i$. These values are added element-wise to the token embeddings.

The key property: $PE_{(pos+k,:)}$ can be expressed as a linear function of $PE_{(pos,:)}$, which allows the model to easily attend to relative positions. The wavelengths across dimensions range from $2\pi$ to $10000 \cdot 2\pi$.

Weight parameters: Zero (fixed, not learned).

Rotary Positional Embedding (RoPE)

Most modern LLMs use Rotary Positional Embedding (RoPE), which encodes position by rotating query and key vectors in 2D subspaces. RoPE applies a rotation matrix to each consecutive pair of dimensions $(2i, 2i+1)$:

where $\theta_i = 10000^{-2i/d_k}$. RoPE is applied to Q and K (not V) after the linear projections.

The key advantage: the dot product $q_m^T k_n$ depends only on the relative position $(m - n)$, enabling natural length generalization without learned position embeddings.

Weight parameters: Zero — RoPE is deterministic.

FLOPs: Negligible (element-wise rotations).

Self-Attention Block

Each decoder layer contains a multi-head self-attention block consisting of three stages: QKV projections, scaled dot-product attention with causal masking, and an output projection.

QKV Projections

The input $X \in \mathbb{R}^{b \times s \times d}$ is projected into queries, keys, and values:

For multi-head attention (MHA), each weight matrix is $\in \mathbb{R}^{d \times d}$. Conceptually, $W_Q = [W_Q^1, \ldots, W_Q^a]$ where each $W_Q^i \in \mathbb{R}^{d \times d_k}$ and $d_k = d/a$. After projection, Q, K, V are reshaped to $(b, a, s, d_k)$ for per-head attention.

For Grouped-Query Attention (GQA), K and V use fewer heads ($h_{kv} < a$), with each KV head shared across $a / h_{kv}$ query heads:

When $h_{kv} = 1$, GQA reduces to Multi-Query Attention (MQA) — all query heads share a single K and V head.

Scaled Dot-Product Attention

The core attention computation operates per-head on tensors of shape $(b, a, s, d_k)$:

The scaling factor is $\sqrt{d_k} = \sqrt{d/a}$ (the per-head dimension, not the full hidden dimension $d$). This prevents the dot products from growing large in magnitude as $d_k$ increases, which would push softmax into regions with vanishingly small gradients.

For decoder (autoregressive) models, a causal mask sets future positions to $-\infty$ before softmax, ensuring each token can only attend to itself and preceding tokens.

Step 1 — Attention scores: $QK^T$

FLOPs: $2 \cdot b \cdot a \cdot s \cdot d_k \cdot s = 2bs^2d$

Step 2 — Weighted values: $\text{softmax}(\cdot) \cdot V$

FLOPs: $2 \cdot b \cdot a \cdot s \cdot s \cdot d_k = 2bs^2d$

The attention FLOPs scale quadratically with sequence length $s$, which is why long-context models are compute-intensive. For large $s$, this term dominates the per-layer cost.

Output Projection

The multi-head attention output is concatenated from $(b, a, s, d_k)$ back to $(b, s, d)$ and linearly projected:

where $W_O \in \mathbb{R}^{d \times d}$.

Self-Attention Block Summary

For standard MHA:

Feed-Forward Network (MLP)

Each decoder layer contains a position-wise feed-forward network applied independently to each token position.

Standard FFN

Two linear layers with a GeLU non-linearity:

where $W_1 \in \mathbb{R}^{d \times d_{ff}}$, $W_2 \in \mathbb{R}^{d_{ff} \times d}$, and typically $d_{ff} = 4d$.

SwiGLU FFN (Modern LLMs)

Most modern LLMs (Llama, Mistral, Gemma, Qwen, etc.) replace the standard FFN with a SwiGLU variant that uses three weight matrices and no bias:

where $\odot$ denotes element-wise multiplication. The intermediate dimension is set to $d_{ff} \approx \frac{8d}{3}$ (rounded to a multiple of 256) so the total parameter count stays comparable to the standard FFN with $d_{ff} = 4d$.

With $d_{ff} = \frac{8d}{3}$: total weight $= 8d^2$ and total FLOPs $= 16bsd^2$ — matching the standard FFN.

Layer Normalization

Each decoder layer has two normalization sub-layers (one before attention, one before MLP in Pre-Norm architecture).

LayerNorm

For input vector $x \in \mathbb{R}^d$:

where:

Weight parameters: $2d$ — learnable scale $\gamma \in \mathbb{R}^d$ and shift $\beta \in \mathbb{R}^d$.

RMSNorm

RMSNorm (used by Llama, Mistral, and most modern LLMs) simplifies LayerNorm by removing mean centering:

Weight parameters: $d$ — learnable scale $\gamma$ only, no shift.

Analysis

Per decoder layer, there are 2 norm layers, contributing $\sim 4d$ (LayerNorm) or $\sim 2d$ (RMSNorm) parameters — negligible compared to $12d^2$ from attention + MLP.

Pre-Norm vs. Post-Norm

The original transformer applies normalization after the residual addition (Post-Norm). Modern LLMs apply it before the sub-layer (Pre-Norm), as proposed by Xiong et al., 2020:

Post-Norm:  x → Attention → Add(residual) → Norm → MLP → Add(residual) → Norm
Pre-Norm:   x → Norm → Attention → Add(residual) → Norm → MLP → Add(residual)

Pre-Norm provides better gradient flow through the residual path and removes the need for learning-rate warmup. The residual stream carries an un-normalized signal directly, making gradients more stable in deep networks. Post-Norm applies normalization on the residual path itself, which can cause gradient magnitude to shrink or explode across many layers. Pre-Norm avoids this by normalizing only the sub-layer branch, leaving the “gradient highway” through residual connections clean.

Residual Connections

Each sub-layer (attention and MLP) is wrapped in a residual connection:

Weight parameters: Zero.

FLOPs: Negligible — element-wise addition of $(b, s, d)$ tensors.

Residual connections are critical for training deep networks: they provide gradient shortcuts that bypass the sub-layers, preventing vanishing gradients and enabling networks with hundreds of layers.

Output Projection (LM Head)

After the final decoder layer, a layer norm is applied, followed by a linear projection to vocabulary logits:

where $W_{head} \in \mathbb{R}^{d \times v}$.

When weight tying is used, $W_{head} = W_E^T$ and contributes no additional parameters.

Per-Layer Summary

For one decoder layer with standard MHA and $d_{ff} = 4d$ (no bias):

Full Model Summary

For large models, the $12Ld^2$ term dominates. For example, Llama-2 70B has $d = 8192$, $L = 80$:

Forward FLOPs (per micro-batch)

The first term ($24bsd^2$) comes from the six large matrix multiplies per layer (QKV + output projection + two MLP layers). The second term ($4bs^2d$) comes from attention and scales quadratically with sequence length.

Training FLOPs

For each matrix multiply $Y = XW$ in the forward pass with cost $C$:

• Forward: $C$ FLOPs (compute $Y$)
• Backward w.r.t. input: $C$ FLOPs (compute $\frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y} W^T$)
• Backward w.r.t. weight: $C$ FLOPs (compute $\frac{\partial L}{\partial W} = X^T \frac{\partial L}{\partial Y}$)

Therefore, training cost $\approx 3 \times$ forward FLOPs:

Dividing by the number of tokens per micro-batch ($bs$) gives the per-token training FLOPs:

The 6P Rule and Its Corrections

The commonly cited $6P$ rule states that training costs $\approx 6$ FLOPs per parameter per token. Since $P \approx 12Ld^2$, the first term above is exactly:

This captures the cost of all dense linear layers (QKV, output projection, MLP) and is the dominant term in most training regimes. However, it ignores two additional contributions:

Attention correction — $+\;12Lsd$

This term comes from the quadratic attention computation ($QK^T$ and $\text{Attn} \times V$). It scales with sequence length and becomes significant when $s$ approaches $d$. For long-context models where $s \gg d$, this term can dominate the entire compute budget.

Embedding / LM head correction — $+\;6dv$

This term comes from the output projection to vocabulary logits. It is independent of model depth and matters most for small models or models with very large vocabularies.

Lower-order terms (softmax, norms, activations) contribute $\sim \mathcal{O}(Ld + Ls)$ FLOPs per token — negligible compared to the $d^2$ and $sd$ terms.

Putting it together:

For typical LLM training regimes ($d \gg s$), the $6P$ term dominates and the rule is a good approximation. But for long-context training (e.g. $s = 128\text{K}$ with $d = 8192$), the attention correction $12Lsd$ becomes comparable to or exceeds $6P$, and the simple rule breaks down.

Activation Memory (per layer)

During training, the forward pass must store intermediate tensors for backpropagation. The major stored activations per decoder layer (in number of elements, multiply by bytes-per-element for memory):

Approximate total per layer (with $d_{ff} = 4d$):

In mixed-precision training (fp16/bf16 activations at 2 bytes each):

The $bas^2$ term grows quadratically with sequence length and can dominate for long contexts. This is a key motivation for techniques like FlashAttention (which avoids materializing the full $s \times s$ attention matrix), activation checkpointing / recomputation, and sequence parallelism.