LLM Architecture and Training Design Choices

See the notes on Transformers first before this, as everything is based on the original transformer paper. This note catalogs the design choices that distinguish a modern decoder-only LLM from the 2017 baseline: how attention shares its keys and values, how position is encoded, where normalization goes, which activation the feedforward block uses, whether the model is dense or sparse, and how depth trades against width.

Attention variants: MHA, MQA, GQA

The bottleneck at inference is the KV cache: for every past token, keys and values must be stored per head. Sharing them shrinks the cache.

Variant	Query heads	Key/Value heads	KV cache
MHA (multi-head)	$H$	$H$	Largest
MQA (multi-query)	$H$	1	Smallest, some quality loss
GQA (grouped-query)	$H$	$G$ ($1<G<H$)	Tunable middle ground

GQA groups query heads to share a smaller number of KV heads and is the current default (Llama 2/3, Mistral): near-MHA quality with a fraction of the memory traffic.

Position embeddings

Position Embeddings

• Sinusoidal embeddings
• Absolute embeddings
• Relative embeddings
• RoPE

• RoPE (rotary): rotates query and key vectors by an angle proportional to position, so the attention dot product depends on relative offset. Extrapolates and extends to long context by adjusting the base frequency.
• ALiBi: adds a linear, head-specific bias to attention scores penalizing distant tokens; no learned position parameters and strong length extrapolation.
• See Positional Encoding for the full treatment.

Norm placement

Pre-vs-post norm

• Pre-norm has become the standard (it is more stable in training)
• ‘Double’ norm has become a trend (found in Grok, Gemma, etc)
  • adding in Layer Norm both before and after the FFN

LayerNorm vs RMSNorm

• RMSNorm is faster and just as good
  • fewer operations (no mean calculation)
  • fewer parameters (no bias term to store)
• Dropping bias terms:
  • original transformer: $FFN(x)=max(0,x{W}_1+b_1)W_2+b_2$
  • most implementations nowadays: $FFN(x)=\sigma (x{W}_1)W_2$

Pre-norm keeps a clean residual stream (identity path with no norm on it), which is why it trains more stably at depth. See LayerNorm for the normalization details; RMSNorm normalizes by root-mean-square only.

Activation: SwiGLU

Activation functions

• ReLU: $FF(x)=\max (0,x{W}_1)W_2$
• GeLU:
• Gating Linear Units
  • GeGLU
  • SwiGLU

SwiGLU is the modern default: $\text{SwiGLU}(x)=(\text{Swish}(x{W}_1)\odot x{W}_3)W_2$. The gating (elementwise product of two projections) gives better quality per parameter. Because it uses three matrices instead of two, the hidden dimension is scaled down to $\frac{2}{3}\times 4d_{\text{model}}$ to keep the parameter count matched.

Mixture-of-Experts (MoE)

Replace the dense FFN with many expert FFNs and a router that sends each token to only the top-$k$ experts. This decouples total parameters from per-token compute: capacity grows without proportional FLOPs. Key knobs are the number of experts, top-$k$, and a load-balancing loss to keep experts evenly used. See Mixture-of-Experts.

Tokenizer

The choice of tokenizer and vocabulary size is an architectural decision: it fixes the embedding and output projection sizes and affects sequence length. See Tokenization.

Depth vs width

Hyperparameters

• feedforward dim $d_{f}f$, model_dim $d_{m}odel$
  • $d_f=4d_{m}odel$, $d_{f}f=2.66d_{m}odel$ if you use GLU variants
• $\text{Head dim}\ast \text{Num heads}=D_{model}$
• Aspect ratio ($d_{m}odel/n_{l}ayer$) within the range of 100 to 200
• Still regularize LMs using Weight decay (limited dropout - 0 to 0.1)

For a fixed parameter budget, deeper models express more compositional functions but are harder to parallelize and less stable; wider models are easier to train and shard but plateau. The aspect ratio ($d_{\text{model}}/n_{\text{layer}}$) around 100 to 200 captures the sweet spot most large models converge to.