Module 25 — Modern LLM Architecture

Modern Large Language Models (LLMs) improve the original Transformer architecture by introducing more efficient normalization, activation functions, attention mechanisms, and inference optimizations.

These techniques significantly improve training stability, memory efficiency, long-context handling, and inference speed.

They are widely used in models such as LLaMA, Gemma, Qwen, DeepSeek, Mistral, Phi, Falcon, Mixtral, Claude, and Gemini.

Topics

RMSNorm
SwiGLU
KV Cache
Grouped Query Attention (GQA)
Rotary Position Embedding (RoPE)

1. RMSNorm

RMSNorm (Root Mean Square Normalization) normalizes activations using only the root mean square value instead of subtracting the mean like LayerNorm.

Formula

RMS(x) = \sqrt{ \frac{1}{d} \sum_{i=1}^{d} x_i^2 + \epsilon }

RMSNorm

RMSNorm(x) = \gamma \frac{x}{RMS(x)}

Expanded,

RMSNorm(x) = \gamma \frac{x} { \sqrt{ \frac1d \sum_{i=1}^{d} x_i^2 + \epsilon } }

where

$\gamma$ = Learnable Scale Parameter
$\epsilon$ = Numerical Stability Constant

Advantages

Faster than LayerNorm
Lower computation
Better scalability for LLMs

Used in

LLaMA
Gemma
Mistral
Qwen
DeepSeek

2. SwiGLU

SwiGLU is an improved Feed Forward activation that combines the Swish activation with a gating mechanism.

Swish

Swish(x) = x \sigma(x)

where

\sigma(x) = \frac1{1+e^{-x}}

SwiGLU

SwiGLU(x) = Swish(xW) \odot (xV)

where

$\odot$ = Element-wise Multiplication

Expanded,

SwiGLU(x) = (xW) \sigma(xW) \odot (xV)

Advantages

Higher model quality
Better gradient flow
Replaces ReLU/GELU FFNs

Used in

PaLM
LLaMA
Gemma
Mistral
Qwen

3. KV Cache

During autoregressive inference, previously computed Key and Value matrices are stored to avoid recomputation.

Key Cache

K_{cache} = [K_1,K_2,\ldots,K_t]

Value Cache

V_{cache} = [V_1,V_2,\ldots,V_t]

Attention Using KV Cache

Attention ( Q_t, K_{cache}, V_{cache} ) = Softmax \left( \frac {Q_tK_{cache}^T} {\sqrt{d_k}} \right) V_{cache}

Advantages

Faster decoding
Reduced computation
Lower inference latency

4. Grouped Query Attention (GQA)

Grouped Query Attention allows multiple Query heads to share one Key-Value group.

Query

Q_h = XW_Q^{(h)}

Shared Key

K_g = XW_K^{(g)}

Shared Value

V_g = XW_V^{(g)}

Attention

Attention ( Q_h, K_g, V_g )

Advantages

Smaller KV Cache
Lower memory usage
Faster inference

Used in

LLaMA 2
LLaMA 3
Gemma
Mistral
Qwen

5. Rotary Position Embedding (RoPE)

RoPE encodes positional information by rotating Query and Key vectors.

Rotation Matrix

R(\theta) = \begin{bmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{bmatrix}

Rotated Query

Q' = R(\theta) Q

Rotated Key

K' = R(\theta) K

RoPE Attention

Attention = Softmax \left( \frac {Q'K'^T} {\sqrt{d_k}} \right) V

Advantages

Long-context capability
Better positional generalization
No learned position embeddings

Used in

GPT-NeoX
LLaMA
Gemma
DeepSeek
Qwen
Phi

6. Modern LLM Decoder Block

A modern decoder layer can be represented as

x \rightarrow RMSNorm \rightarrow RoPE \rightarrow GQA \rightarrow Residual \rightarrow RMSNorm \rightarrow SwiGLU \rightarrow Residual

7. Complete Modern LLM Pipeline

The complete architecture is

Input \rightarrow TokenEmbedding \rightarrow RoPE \rightarrow GroupedQueryAttention \rightarrow KVCache \rightarrow RMSNorm \rightarrow SwiGLU \rightarrow OutputProjection \rightarrow Softmax

8. Applications

These architectural improvements are used in

LLaMA 2
LLaMA 3
Gemma
Gemma 2
Mistral
Mixtral
Qwen
DeepSeek
Phi
Falcon
Claude
Gemini
GPT-NeoX
Modern Large Language Models (LLMs)

Summary

Concept	Formula
RMS	$RMS(x)=\sqrt{\frac1d\sum_{i=1}^{d}x_i^2+\epsilon}$