|$ curl https://forge-ai.dev/api/markdown?path=docs/ai/architecture

$cat docs/model-architecture.md

updated Recently·45 min read·published

Model Architecture

◆AI Engineering◆Architecture◆Intermediate to Advanced

Introduction

The Transformer architecture is the foundation of virtually every modern LLM. Introduced in the 2017 paper "Attention Is All You Need," it replaced recurrent and convolutional architectures with a purely attention-based design that scales dramatically better with data and compute.

This section provides a deep dive into the Transformer architecture: the attention mechanism, multi-head attention, feed-forward networks, positional encodings, layer normalization, and the differences between encoder-decoder and decoder-only designs. Understanding these internals helps you make informed decisions about model selection, fine-tuning, and debugging.

📝

note

You can build effective AI applications without understanding every detail of Transformer internals. But when things go wrong — unexpected outputs, poor fine-tuning results, or strange attention patterns — architectural knowledge becomes invaluable.

Transformer Overview

At the highest level, a Transformer is composed of stacked layers, each containing two main sub-layers: a multi-head self-attention mechanism and a position-wise feed-forward network. Residual connections and layer normalization surround each sub-layer.

preview

Each layer processes all tokens in parallel (unlike RNNs which process sequentially), which enables massive parallelism during training. The depth (number of layers), width (hidden dimension), and number of attention heads are the key architectural hyperparameters.

Hyperparameter	GPT-4o (est.)	Llama 3 70B	Claude 3.5 (est.)	Function
Layers	~96	80	~64	Depth of the network
Hidden Dim	~8,192	8,192	~6,144	Width of each layer
Heads	~128	64	~48	Parallel attention heads
Head Dim	64	128	128	Size per attention head
FFN Width	~32,768	28,672	~24,576	Hidden dim of feed-forward
Vocab Size	100K	128K	~80K	Number of unique tokens

ℹ

info

The architectures of proprietary models (GPT-4o, Claude) are not publicly known. The estimates in the table above are based on inference benchmarks, parameter counts from API behavior, and published research papers. Open-source models like Llama provide exact specifications.

The Attention Mechanism

Attention is the core innovation of the Transformer. It allows each token to directly interact with every other token in the sequence, computing context-aware representations. The operation is defined as:

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V

Where Q (queries), K (keys), and V (values) are linear projections of the input. The dot product QK^T measures compatibility between tokens, the softmax converts this to a probability distribution, and the weighted sum of values produces the output.

scaled-dot-product-attention.py

Python

1	# Scaled Dot-Product Attention (full implementation)
2	import torch
3	import torch.nn.functional as F
4
5	def scaled_dot_product_attention(
6	query: torch.Tensor, # (batch, heads, seq_len, d_k)
7	key: torch.Tensor, # (batch, heads, seq_len, d_k)
8	value: torch.Tensor, # (batch, heads, seq_len, d_v)
9	mask: torch.Tensor \| None = None,
10	dropout: float = 0.0
11	) -> torch.Tensor:
12	d_k = query.size(-1)
13
14	# Compute attention scores: Q @ K^T
15	scores = torch.matmul(query, key.transpose(-2, -1))
16
17	# Scale by sqrt(d_k) to prevent vanishing gradients
18	scores = scores / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))
19
20	# Apply causal mask (for decoder-only models)
21	if mask is not None:
22	scores = scores.masked_fill(mask == 0, float('-inf'))
23
24	# Softmax along the key dimension
25	attention_weights = F.softmax(scores, dim=-1)
26
27	# Apply dropout for regularization
28	attention_weights = F.dropout(attention_weights, p=dropout)
29
30	# Weighted sum of values
31	output = torch.matmul(attention_weights, value)
32
33	return output, attention_weights
34
35	# The attention weights tell us which tokens the model
36	# is focusing on — useful for interpretability

✓

best practice

The attention mechanism has O(n²) memory complexity with respect to sequence length. This is the bottleneck for long-context models. Techniques like Flash Attention (Dao et al., 2022) optimize memory usage by computing attention in tiles without materializing the full attention matrix. Flash Attention is now standard in all major LLM frameworks.

Multi-Head Attention

Instead of computing a single attention function, multi-head attention runs multiple attention operations in parallel (each with its own learned projections). The outputs are concatenated and linearly projected. This allows the model to attend to different types of relationships simultaneously — syntactic, semantic, positional, and more.

multi-head-attention.py

Python

1	# Multi-Head Attention implementation
2	import torch
3	import torch.nn as nn
4
5	class MultiHeadAttention(nn.Module):
6	def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
7	super().__init__()
8	assert d_model % n_heads == 0
9
10	self.d_model = d_model
11	self.n_heads = n_heads
12	self.d_k = d_model // n_heads # Dimension per head
13
14	# Linear projections for Q, K, V (all heads combined)
15	self.W_q = nn.Linear(d_model, d_model)
16	self.W_k = nn.Linear(d_model, d_model)
17	self.W_v = nn.Linear(d_model, d_model)
18	self.W_o = nn.Linear(d_model, d_model) # Output projection
19
20	self.dropout = nn.Dropout(dropout)
21
22	def forward(self, query, key, value, mask=None):
23	batch_size = query.size(0)
24
25	# Project and reshape for multi-head
26	Q = self.W_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
27	K = self.W_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
28	V = self.W_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
29
30	# Apply attention (see previous function)
31	attn_output, weights = scaled_dot_product_attention(Q, K, V, mask, self.dropout)
32
33	# Concatenate heads and project
34	attn_output = attn_output.transpose(1, 2).contiguous()
35	attn_output = attn_output.view(batch_size, -1, self.d_model)
36
37	return self.W_o(attn_output)
38
39	# Usage example:
40	# d_model=4096, n_heads=32 means each head works in 128 dimensions
41	# Each head learns different patterns:
42	# - Head 1-4: syntactic relationships (subject-verb, noun-adjective)
43	# - Head 5-8: positional relationships (next-token, previous-token)
44	# - Head 9-16: semantic relationships (coreference, synonyms)
45	# - Head 17-32: task-specific patterns learned during training

preview

🔥

pro tip

The number of heads is a key design choice. More heads = more parallel attention patterns but diminishing returns. Most models use between 16 and 128 heads. The head dimension (d_k) is typically 64 or 128 — this is small enough to be efficient but large enough to capture useful relationships.

Feed-Forward Networks

Each Transformer layer contains a position-wise feed-forward network (FFN) that processes each token independently. The FFN typically consists of two linear transformations with a non-linear activation in between. The hidden dimension of the FFN is usually 2-4x the model dimension.

feed-forward-variants.py

Python

1	# Feed-Forward Network variants used in modern LLMs
2
3	import torch
4	import torch.nn as nn
5	import torch.nn.functional as F
6
7	class FFN_GELU(nn.Module):
8	"""Standard FFN with GELU activation (GPT-4, BERT)."""
9	def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
10	super().__init__()
11	self.fc1 = nn.Linear(d_model, d_ff)
12	self.fc2 = nn.Linear(d_ff, d_model)
13	self.dropout = nn.Dropout(dropout)
14
15	def forward(self, x):
16	return self.fc2(self.dropout(F.gelu(self.fc1(x))))
17
18	class FFN_SwiGLU(nn.Module):
19	"""SwiGLU variant (Llama 2/3, Mistral). Uses gated activation.
20	Outperforms standard GELU in most modern LLMs."""
21	def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
22	super().__init__()
23	self.gate = nn.Linear(d_model, d_ff)
24	self.up = nn.Linear(d_model, d_ff)
25	self.down = nn.Linear(d_ff, d_model)
26	self.dropout = nn.Dropout(dropout)
27
28	def forward(self, x):
29	# SwiGLU: (x @ W_gate) * silu(x @ W_up) @ W_down
30	gate = F.silu(self.gate(x))
31	up = self.up(x)
32	return self.down(self.dropout(gate * up))
33
34	# SwiGLU requires 3 weight matrices instead of 2,
35	# so d_ff is typically 8/3 * d_model (instead of 4 * d_model)
36	# to keep total parameters similar.
37	# Example: d_model=4096, d_ff=10923 (SwiGLU) vs d_ff=16384 (GELU)

The FFN is where the model stores and applies learned knowledge. While attention decides which information to route between tokens, the FFN transforms the information at each token position. This is why FFNs constitute about two-thirds of the model's parameters.

✓

best practice

When fine-tuning a model, the FFN layers are often the most impactful to update. Techniques like LoRA (Low-Rank Adaptation) target the attention projections, but full fine-tuning or DoRA (Weight-Decomposed Low-Rank Adaptation) can achieve better results by also updating FFN weights.

Positional Encoding

Transformers process all tokens in parallel and have no inherent notion of token order. Positional encodings inject information about the position of each token in the sequence. Several approaches have been developed:

Sinusoidal Positional Encoding (Original)

The original Transformer used fixed sine and cosine functions of different frequencies. Each position gets a unique encoding vector that the model can learn to interpret. The sinusoidal pattern allows the model to extrapolate to longer sequences than seen during training.

sinusoidal-encoding.py

Python

1	# Sinusoidal positional encoding
2	import torch
3	import math
4
5	def sinusoidal_positional_encoding(
6	seq_len: int, d_model: int
7	) -> torch.Tensor:
8	"""Generate sinusoidal position encodings."""
9	pe = torch.zeros(seq_len, d_model)
10	position = torch.arange(0, seq_len, dtype=torch.float32).unsqueeze(1)
11
12	# Different frequencies for each dimension
13	div_term = torch.exp(
14	torch.arange(0, d_model, 2, dtype=torch.float32) *
15	-(math.log(10000.0) / d_model)
16	)
17
18	# Apply sine to even indices, cosine to odd indices
19	pe[:, 0::2] = torch.sin(position * div_term)
20	pe[:, 1::2] = torch.cos(position * div_term)
21
22	return pe
23
24	# Visualize: each position has a unique pattern
25	pe = sinusoidal_positional_encoding(100, 64)
26	# Position 0 might be [0, 1, 0, 1, ...]
27	# Position 1 might be [0.84, 0.54, 0.01, 0.99, ...]
28	# Nearby positions have similar encodings (smooth transitions)

Rotary Position Encoding (RoPE)

RoPE (Su et al., 2021) is the dominant positional encoding in modern LLMs (Llama, Mistral, GPT-4o). Instead of adding a position vector to the input, RoPE rotates the query and key vectors by an angle proportional to the position. This naturally encodes relative position information into the attention scores.

rope-encoding.py

Python

1	# Rotary Position Encoding (RoPE)
2	import torch
3
4	def apply_rotary_emb(
5	x: torch.Tensor, # (batch, heads, seq_len, d_k)
6	seq_len: int,
7	base: float = 10000.0
8	) -> torch.Tensor:
9	"""Apply rotary position encoding to queries and keys."""
10	d_k = x.shape[-1]
11
12	# Compute frequencies for each dimension pair
13	theta = 1.0 / (base ** (torch.arange(0, d_k, 2).float() / d_k))
14	positions = torch.arange(seq_len).float()
15
16	# Outer product: (seq_len, d_k/2)
17	freqs = torch.outer(positions, theta)
18
19	# Compute cos and sin
20	cos = freqs.cos().unsqueeze(0).unsqueeze(0) # (1, 1, seq_len, d_k/2)
21	sin = freqs.sin().unsqueeze(0).unsqueeze(0)
22
23	# Interleave: duplicate for paired dimensions
24	cos = cos.repeat_interleave(2, dim=-1)
25	sin = sin.repeat_interleave(2, dim=-1)
26
27	# Apply rotation: x * cos + rotate_half(x) * sin
28	x_rotated = torch.stack([-x[..., 1::2], x[..., ::2]], dim=-1).reshape_as(x)
29	return x * cos + x_rotated * sin
30
31	# RoPE benefits:
32	# 1. Relative position: attention depends on position difference
33	# 2. Decay: distant tokens have lower rotation alignment
34	# 3. Extrapolation: can generalize to longer sequences
35	# 4. Compatible with linear attention

🔥

pro tip

The RoPE base parameter (default 10000) controls how quickly the rotation frequencies decay. Increasing the base (e.g., to 500000 in Llama 3) allows the model to handle longer contexts by reducing the rotation angle between adjacent positions. This is a key hyperparameter for long-context models.

Layer Normalization

Layer normalization stabilizes training by normalizing activations across the feature dimension. It has become a critical component of the Transformer architecture, with modern variants using different placement and formulation.

Pre-Norm vs Post-Norm

The original Transformer placed layer normalization after each sub-layer (post-norm). Modern LLMs (GPT-4o, Llama, Claude) use pre-norm — normalization before each sub-layer. Pre-norm provides more stable gradients during training and enables training deeper models without gradient divergence.

layer-normalization.py

Python

1	# Pre-Norm vs Post-Norm Transformer block
2
3	import torch
4	import torch.nn as nn
5
6	class PostNormBlock(nn.Module):
7	"""Original Transformer: Norm after sub-layer (post-norm)."""
8	def __init__(self, d_model: int, n_heads: int):
9	super().__init__()
10	self.attention = MultiHeadAttention(d_model, n_heads)
11	self.norm1 = nn.LayerNorm(d_model)
12	self.ffn = FFN_GELU(d_model, d_model * 4)
13	self.norm2 = nn.LayerNorm(d_model)
14
15	def forward(self, x):
16	# Attention with residual -> Norm
17	x = self.norm1(x + self.attention(x, x, x))
18	# FFN with residual -> Norm
19	x = self.norm2(x + self.ffn(x))
20	return x
21
22	class PreNormBlock(nn.Module):
23	"""Modern LLM: Norm before sub-layer (pre-norm).
24	Used by GPT-4o, Llama 3, Mistral, Claude."""
25	def __init__(self, d_model: int, n_heads: int):
26	super().__init__()
27	self.attention = MultiHeadAttention(d_model, n_heads)
28	self.norm1 = nn.LayerNorm(d_model)
29	self.ffn = FFN_SwiGLU(d_model, int(d_model * 8 / 3))
30	self.norm2 = nn.LayerNorm(d_model)
31
32	def forward(self, x):
33	# Norm -> Attention -> Residual
34	x = x + self.attention(self.norm1(x), self.norm1(x), self.norm1(x))
35	# Norm -> FFN -> Residual
36	x = x + self.ffn(self.norm2(x))
37	return x
38
39	# RMS Norm (used by Llama 3 instead of LayerNorm)
40	class RMSNorm(nn.Module):
41	"""Root Mean Square Layer Normalization.
42	Simpler and faster than LayerNorm."""
43	def __init__(self, d_model: int, eps: float = 1e-6):
44	super().__init__()
45	self.weight = nn.Parameter(torch.ones(d_model))
46	self.eps = eps
47
48	def forward(self, x):
49	rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
50	return x / rms * self.weight

📝

note

The shift from post-norm to pre-norm was one of the key engineering advances that enabled training deeper Transformers. If you are implementing your own Transformer, always use pre-norm with RMSNorm — it is simpler, faster, and more stable.

Encoder-Decoder vs Decoder-Only

Two major architectural variants dominate modern LLMs: encoder-decoder models (like the original Transformer and T5) and decoder-only models (like GPT, Llama, Mistral, Claude). Understanding the differences is crucial for choosing the right architecture for your use case.

Feature	Encoder-Decoder	Decoder-Only
Architecture	Encoder + Decoder stacks	Single decoder stack
Attention	Bidirectional (encoder) + causal (decoder)	Causal (unidirectional)
Context Access	Full bidirectional context in encoder	Left-to-right only
Examples	T5, BART, UL2	GPT-4o, Llama, Mistral, Claude, Gemini
Best For	Translation, summarization, classification	Chat, code gen, creative, general purpose
Scaling	Less efficient (two stacks)	More efficient (one stack)
Parameter Count	Higher for same capacity	Lower for same capacity

Decoder-Only Architecture (GPT-style)

Decoder-only models stack identical decoder layers, each with masked self-attention (causal mask prevents attending to future tokens) and a feed-forward network. The causal mask ensures autoregressive generation — each token can only attend to itself and previous tokens. This is the dominant architecture for modern LLMs.

decoder-only-lm.py

Python

1	# Decoder-only transformer (GPT-style)
2	import torch
3	import torch.nn as nn
4
5	class DecoderOnlyLayer(nn.Module):
6	def __init__(self, d_model: int, n_heads: int):
7	super().__init__()
8	self.self_attn = MultiHeadAttention(d_model, n_heads)
9	self.norm1 = nn.LayerNorm(d_model)
10	self.ffn = FFN_SwiGLU(d_model, int(d_model * 8 / 3))
11	self.norm2 = nn.LayerNorm(d_model)
12
13	def forward(self, x, causal_mask):
14	# Self-attention with causal mask
15	x = x + self.self_attn(
16	self.norm1(x), self.norm1(x), self.norm1(x), mask=causal_mask
17	)
18	# Feed-forward
19	x = x + self.ffn(self.norm2(x))
20	return x
21
22	class DecoderOnlyModel(nn.Module):
23	def __init__(self, vocab_size: int, d_model: int, n_layers: int, n_heads: int):
24	super().__init__()
25	self.embedding = nn.Embedding(vag_size, d_model)
26	self.layers = nn.ModuleList([
27	DecoderOnlyLayer(d_model, n_heads) for _ in range(n_layers)
28	])
29	self.norm = nn.LayerNorm(d_model)
30	self.lm_head = nn.Linear(d_model, vocab_size)
31
32	def forward(self, input_ids):
33	seq_len = input_ids.shape[1]
34	# Causal mask: prevent attending to future tokens
35	mask = torch.tril(torch.ones(seq_len, seq_len)).view(1, 1, seq_len, seq_len)
36
37	x = self.embedding(input_ids)
38	for layer in self.layers:
39	x = layer(x, mask)
40	x = self.norm(x)
41	logits = self.lm_head(x)
42	return logits

Encoder-Decoder Architecture (T5-style)

The encoder processes the input with bidirectional attention (each token attends to all tokens), producing a representation. The decoder then generates output autoregressively, attending to both the encoder output and previously generated tokens via cross-attention.

encoder-decoder.py

Python

1	# Encoder-Decoder (T5-style)
2	class EncoderLayer(nn.Module):
3	def __init__(self, d_model: int, n_heads: int):
4	super().__init__()
5	self.self_attn = MultiHeadAttention(d_model, n_heads)
6	self.norm1 = nn.LayerNorm(d_model)
7	self.ffn = FFN_GELU(d_model, d_model * 4)
8	self.norm2 = nn.LayerNorm(d_model)
9
10	def forward(self, x):
11	# Bidirectional attention (no mask)
12	x = x + self.self_attn(self.norm1(x), self.norm1(x), self.norm1(x))
13	x = x + self.ffn(self.norm2(x))
14	return x
15
16	class DecoderLayer(nn.Module):
17	def __init__(self, d_model: int, n_heads: int):
18	super().__init__()
19	self.self_attn = MultiHeadAttention(d_model, n_heads)
20	self.cross_attn = MultiHeadAttention(d_model, n_heads)
21	self.norm1 = nn.LayerNorm(d_model)
22	self.norm2 = nn.LayerNorm(d_model)
23	self.norm3 = nn.LayerNorm(d_model)
24	self.ffn = FFN_GELU(d_model, d_model * 4)
25
26	def forward(self, x, encoder_output, causal_mask):
27	# Masked self-attention
28	x = x + self.self_attn(self.norm1(x), self.norm1(x), self.norm1(x), mask=causal_mask)
29	# Cross-attention: queries from decoder, KVs from encoder
30	x = x + self.cross_attn(self.norm2(x), encoder_output, encoder_output)
31	x = x + self.ffn(self.norm3(x))
32	return x

ℹ

info

For most modern AI engineering applications, decoder-only models are the better choice. They are simpler to deploy, more parameter-efficient, and perform better on chat and instruction-following tasks. Encoder-decoder models remain useful for tasks with strong asymmetry between input and output formats (like translation or summarization with very long inputs).

Modern Architectural Innovations

Beyond the classic Transformer, several architectural innovations have shaped modern LLMs. Understanding these helps you evaluate new models and make informed choices.

Mixture of Experts (MoE)

MoE replaces dense FFN layers with multiple "expert" FFNs and a routing mechanism that activates only a subset for each token. This increases total parameter count without proportionally increasing compute. GPT-4 and Mixtral 8x22B use MoE architectures. The router learns to assign different tokens to different experts, creating specialization.

moe-ffn.py

Python

1	# Simplified MoE FFN
2	class MoEFFN(nn.Module):
3	def __init__(self, d_model: int, d_ff: int, n_experts: int = 8, top_k: int = 2):
4	super().__init__()
5	self.experts = nn.ModuleList([
6	FFN_SwiGLU(d_model, d_ff) for _ in range(n_experts)
7	])
8	self.router = nn.Linear(d_model, n_experts)
9	self.top_k = top_k
10
11	def forward(self, x):
12	# Route tokens to top-k experts
13	routing_logits = self.router(x)
14	routing_weights = F.softmax(routing_logits, dim=-1)
15	top_k_weights, top_k_indices = torch.topk(routing_weights, self.top_k, dim=-1)
16
17	# Compute weighted expert outputs
18	output = torch.zeros_like(x)
19	for i, expert in enumerate(self.experts):
20	mask = (top_k_indices == i).any(dim=-1)
21	if mask.any():
22	output[mask] += top_k_weights[mask][:, top_k_indices[mask] == i].unsqueeze(-1) * expert(x[mask])
23
24	return output

Grouped Query Attention (GQA)

GQA reduces the number of key-value heads relative to query heads. Instead of each query head having its own KV head, groups of query heads share a single KV head. This dramatically reduces the memory footprint of the KV cache during inference, enabling longer context and larger batch sizes. Llama 3 70B uses 64 query heads and 8 KV heads (8:1 ratio).

KV Caching

During autoregressive generation, the key and value tensors from previous tokens can be cached and reused, avoiding redundant computation. Without KV caching, generating N tokens would cost O(N²) in attention compute. With caching, it drops to O(N). This is why frameworks like vLLM and TensorRT-LLM implement sophisticated KV cache management (PagedAttention, prefix caching).

Flash Attention

Flash Attention computes exact attention without materializing the full N×N attention matrix, using tiling and recomputation. It is 2-4x faster than standard attention for long sequences and uses O(N) memory instead of O(N²). Flash Attention 3 adds async processing and FP8 support for further speedups on H100 GPUs.

🔥

pro tip

When deploying models in production, KV cache memory often becomes the bottleneck, not model parameters. A 70B model with 128K context can require 40+ GB of KV cache per request. Use GQA models, compute-aware batching, and prompt compression to manage this. vLLM's PagedAttention is the standard solution for efficient KV cache management.

$Blueprint — Engineering Documentation·Section ID: AI-04·Revision: 1.0

X GitHub YouTube Instagram Facebook Buy me a coffee Sponsor