LLM Architecture Deep Dive

Introduction: Seven Years of Transformer Evolution

When Vaswani et al. introduced the transformer architecture in 2017 with their seminal paper "Attention is All You Need," they fundamentally changed the landscape of deep learning. Seven years later, as we analyze the architectures powering the most advanced language models in 2024-2025, we find something remarkable: the core transformer architecture remains largely unchanged, yet the refinements and innovations built upon this foundation have created models of unprecedented capability.

Figure 1: A comprehensive overview of the architectures analyzed in this guide, spanning from established models like GPT to cutting-edge systems like DeepSeek V3.

The Fundamental Building Blocks

To understand modern LLM architectures, we must first establish a solid foundation of the transformer's core components. The transformer architecture consists of several key building blocks that work in concert:

The Evolution Timeline

Let's trace the key milestones in transformer evolution:

Part 1: The Evolution of Attention Mechanisms

Understanding the Computational Challenge

The self-attention mechanism's power comes with a significant computational cost. To fully understand the innovations in attention, let's examine the mathematics and computational complexity:

Multi-Head Attention (MHA): The Original Design

Multi-Head Attention divides the model's representation into multiple "heads," each learning different types of relationships. This isn't just about parallelization—it's about specialization. Research has shown that different heads learn distinctly different patterns:

The mathematical formulation of MHA for h heads:

Figure 2: Visual comparison of Multi-Head Attention (left) with separate KV pairs per head, versus Grouped-Query Attention (right) where KV pairs are shared across head groups.

Grouped-Query Attention: Elegant Memory Optimization

Grouped-Query Attention (GQA) emerged from a critical observation: while queries benefit from many heads for expressiveness, keys and values show significant redundancy across heads. By sharing KV pairs across groups of query heads, GQA achieves substantial memory savings with minimal quality loss.

Multi-Head Latent Attention: Compression as a Feature

Multi-Head Latent Attention (MLA) takes a fundamentally different approach from GQA. Instead of sharing KV pairs, it compresses them into a lower-dimensional latent space. This isn't just about memory savings—the compression acts as a form of regularization that can actually improve model quality.

Figure 3: Multi-Head Latent Attention architecture showing compression of KV pairs into a compact latent representation c_t before projection to individual heads.

Figure 4: Ablation studies showing MLA achieves better perplexity than both standard MHA and GQA while using dramatically less memory.

Sliding Window Attention: Locality as an Inductive Bias

Sliding Window Attention exploits the observation that most linguistic dependencies are local. Instead of allowing every token to attend to all others, each token can only attend to a fixed window of surrounding tokens.

Figure 12: Comparison between global attention (left) where every token attends to all others, and sliding window attention (right) with local context windows.

Attention Pattern Analysis

Understanding how different attention mechanisms affect learned patterns is crucial for architecture selection:

Part 2: The Mixture of Experts Revolution

The Fundamental Insight Behind Sparse Models

The Mixture of Experts (MoE) architecture addresses a fundamental tension in language modeling: we want models with massive capacity to store knowledge, but we can't afford the computational cost of activating all parameters for every token. MoE elegantly resolves this by creating a sparse model where different parameters specialize in different types of inputs.

The Mathematics of Expert Routing

The routing mechanism is the heart of MoE architectures. It must decide which experts to activate for each token, balancing specialization with load distribution:

Figure 5: The Mixture of Experts architecture showing how tokens are routed to different experts based on the router's decisions.

Expert Specialization Patterns

Research has revealed fascinating specialization patterns that emerge naturally during MoE training:

The Shared Expert Innovation

DeepSeek introduced the concept of "shared experts" - experts that are always active regardless of routing decisions. This addresses a critical weakness in pure MoE designs:

Figure 6: Impact of shared experts on model performance, showing how they capture common patterns while routed experts handle specialization.

Scaling Laws for MoE Architectures

The optimal number of experts follows interesting scaling patterns:

Implementation Challenges and Solutions

Building efficient MoE models requires solving several engineering challenges:

Comparing MoE Implementations

Different organizations have taken varying approaches to MoE design:

Part 3: DeepSeek V3/R1 - Setting New Standards

Architectural Deep Dive

DeepSeek V3, released in December 2024, represents the current pinnacle of open-weight language model engineering. Its successor, DeepSeek-R1, adds reasoning capabilities through reinforcement learning while maintaining the same base architecture.

Training Innovations

┌─────────┬──────────────┬─────────────┐
│  Sign   │   Exponent   │   Mantissa  │
│  1 bit  │    4 bits    │    3 bits   │
├─────────┼──────────────┼─────────────┤
│    S    │    E E E E   │   M M M     │
└─────────┴──────────────┴─────────────┘

Range:     2^-6 to 2^8
Precision: ~1.5 decimal digits
Total:     8 bits per value

Performance Benchmarks

Deployment and Serving Optimizations

Part 4: Model-by-Model Deep Analysis

OLMo 2: The Transparent Blueprint

OLMo 2, developed by the Allen Institute for AI, provides unprecedented transparency in LLM development. Every decision, experiment, and failure is documented, making it an invaluable learning resource.

Figure 7: OLMo 2 achieves Pareto-optimal performance, matching larger models while using less compute through architectural refinements.

Normalization Innovation: The Best of Both Worlds

OLMo 2's key contribution is its hybrid normalization approach, combining benefits of Pre-Norm and Post-Norm:

Figure 8: Three normalization strategies - Post-Norm (training unstable), Pre-Norm (representation collapse), and OLMo's hybrid (stable + expressive).

Gemma 3: Google's Efficiency Champion

Gemma 3 focuses on deployment efficiency rather than pushing parameter counts, optimizing for real-world usage constraints.

Figure 11: Memory savings from Gemma's sliding window attention, enabling 1M+ token context on consumer GPUs.

Llama 4: Meta's MoE Evolution

Llama 4 represents Meta's embrace of sparse architectures after the dense-only Llama 1-3 series:

Figure 17: Architectural comparison between DeepSeek V3's aggressive sparsity and Llama 4's conservative approach.

SmolLM3: Extreme Efficiency at Small Scale

SmolLM3 demonstrates that architectural innovations benefit small models too:

Qwen3: Specialized Architecture Variants

The Qwen family demonstrates how architectural choices can be tailored for specific use cases:

Part 5: Implementation Guide

Building Key Components from Scratch

Let's implement the core architectural innovations to understand them deeply:

1. Grouped-Query Attention Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class GroupedQueryAttention(nn.Module):
    """
    Grouped-Query Attention as used in Llama 2/3, reducing KV cache by sharing
    keys and values across groups of query heads.
    """
    def __init__(self, d_model, n_heads, n_kv_heads, dropout=0.1):
        super().__init__()
        assert n_heads % n_kv_heads == 0, "n_heads must be divisible by n_kv_heads"

        self.d_model = d_model
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.n_rep = n_heads // n_kv_heads  # Repetition factor
        self.head_dim = d_model // n_heads

        # Projections
        self.w_q = nn.Linear(d_model, n_heads * self.head_dim, bias=False)
        self.w_k = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.w_v = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.w_o = nn.Linear(n_heads * self.head_dim, d_model, bias=False)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None, cache_k=None, cache_v=None):
        batch_size, seq_len, _ = x.shape

        # Compute Q, K, V
        q = self.w_q(x).view(batch_size, seq_len, self.n_heads, self.head_dim)
        k = self.w_k(x).view(batch_size, seq_len, self.n_kv_heads, self.head_dim)
        v = self.w_v(x).view(batch_size, seq_len, self.n_kv_heads, self.head_dim)

        # Transpose for attention: [batch, heads, seq_len, head_dim]
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # Handle KV cache for inference
        if cache_k is not None:
            k = torch.cat([cache_k, k], dim=2)
            v = torch.cat([cache_v, v], dim=2)

        # Repeat K, V to match number of Q heads
        if self.n_rep > 1:
            k = k.repeat_interleave(self.n_rep, dim=1)
            v = v.repeat_interleave(self.n_rep, dim=1)

        # Scaled dot-product attention
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # Apply attention to values
        attn_output = torch.matmul(attn_weights, v)

        # Reshape and project output
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, -1)
        output = self.w_o(attn_output)

        return output, (k, v) if cache_k is not None else None

# Example usage
model = GroupedQueryAttention(d_model=4096, n_heads=32, n_kv_heads=8)
x = torch.randn(2, 1024, 4096)  # [batch, seq_len, d_model]
output, _ = model(x)
print(f"Output shape: {output.shape}")  # [2, 1024, 4096]
print(f"Memory saved: {(32/8):.1f}x reduction in KV cache")

2. Mixture of Experts Layer

class MoELayer(nn.Module):
    """
    Mixture of Experts layer with top-k routing and load balancing.
    """
    def __init__(self, d_model, d_ff, n_experts, n_experts_per_tok, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.n_experts = n_experts
        self.n_experts_per_tok = n_experts_per_tok

        # Router (gate)
        self.router = nn.Linear(d_model, n_experts, bias=False)

        # Experts (simple FFN for this example)
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff, bias=False),
                nn.SiLU(),  # SwiGLU activation
                nn.Linear(d_ff, d_model, bias=False),
                nn.Dropout(dropout)
            ) for _ in range(n_experts)
        ])

    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        x_flat = x.view(-1, d_model)  # [batch * seq_len, d_model]

        # Compute router scores
        router_logits = self.router(x_flat)  # [batch * seq_len, n_experts]
        router_probs = F.softmax(router_logits, dim=-1)

        # Select top-k experts per token
        topk_probs, topk_indices = torch.topk(
            router_probs, self.n_experts_per_tok, dim=-1
        )

        # Normalize top-k probabilities
        topk_probs = topk_probs / topk_probs.sum(dim=-1, keepdim=True)

        # Initialize output
        output = torch.zeros_like(x_flat)

        # Route tokens to experts
        for i in range(self.n_experts_per_tok):
            # Get expert index for each token
            expert_idx = topk_indices[:, i]
            expert_weight = topk_probs[:, i].unsqueeze(-1)

            # Process each expert
            for expert_id in range(self.n_experts):
                # Find tokens routed to this expert
                mask = (expert_idx == expert_id)
                if mask.any():
                    token_indices = mask.nonzero(as_tuple=True)[0]
                    expert_input = x_flat[token_indices]
                    expert_output = self.experts[expert_id](expert_input)

                    # Weighted sum
                    output[token_indices] += expert_weight[token_indices] * expert_output

        # Reshape back
        output = output.view(batch_size, seq_len, d_model)

        # Compute load balancing loss (auxiliary)
        # This encourages uniform expert usage
        expert_usage = router_probs.mean(dim=0)  # Average prob per expert
        load_balance_loss = self.n_experts * (expert_usage * expert_usage).sum()

        return output, load_balance_loss

# Example usage
moe = MoELayer(d_model=768, d_ff=3072, n_experts=8, n_experts_per_tok=2)
x = torch.randn(2, 512, 768)
output, aux_loss = moe(x)
print(f"Output shape: {output.shape}")  # [2, 512, 768]
print(f"Load balance loss: {aux_loss:.4f}")

3. Multi-Head Latent Attention (Simplified)

class MultiHeadLatentAttention(nn.Module):
    """
    Multi-Head Latent Attention as used in DeepSeek V3.
    Compresses KV into latent space before caching.
    """
    def __init__(self, d_model, n_heads, d_latent, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_latent = d_latent
        self.head_dim = d_model // n_heads

        # Query projection (standard)
        self.w_q = nn.Linear(d_model, n_heads * self.head_dim, bias=False)

        # Latent compression for KV
        self.w_latent = nn.Linear(d_model, d_latent, bias=False)

        # Latent to KV projections (per head)
        self.latent_to_k = nn.Linear(d_latent, n_heads * self.head_dim, bias=False)
        self.latent_to_v = nn.Linear(d_latent, n_heads * self.head_dim, bias=False)

        # Output projection
        self.w_o = nn.Linear(n_heads * self.head_dim, d_model, bias=False)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape

        # Compute queries
        q = self.w_q(x).view(batch_size, seq_len, self.n_heads, self.head_dim)
        q = q.transpose(1, 2)  # [batch, n_heads, seq_len, head_dim]

        # Compress to latent space
        latent = self.w_latent(x)  # [batch, seq_len, d_latent]

        # Project from latent to K, V
        k = self.latent_to_k(latent).view(batch_size, seq_len, self.n_heads, self.head_dim)
        v = self.latent_to_v(latent).view(batch_size, seq_len, self.n_heads, self.head_dim)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # Standard attention computation
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        attn_output = torch.matmul(attn_weights, v)
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, -1)

        output = self.w_o(attn_output)

        # For caching, we only need to store the latent representation!
        cache_size = latent.shape[-1]  # d_latent instead of n_heads * head_dim * 2
        compression_ratio = (2 * self.n_heads * self.head_dim) / self.d_latent

        return output, latent  # Return latent for caching

# Example usage
mla = MultiHeadLatentAttention(d_model=4096, n_heads=32, d_latent=512)
x = torch.randn(2, 1024, 4096)
output, latent_cache = mla(x)
print(f"Output shape: {output.shape}")  # [2, 1024, 4096]
print(f"Latent cache shape: {latent_cache.shape}")  # [2, 1024, 512]
print(f"Compression ratio: {(32*128*2)/512:.1f}x")  # ~16x compression

4. RoPE (Rotary Position Embeddings)

class RotaryPositionEmbedding(nn.Module):
    """
    Rotary Position Embedding (RoPE) as used in most modern LLMs.
    """
    def __init__(self, dim, max_seq_len=8192, base=10000):
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.base = base

        # Precompute rotation frequencies
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)

        # Precompute cos and sin for all positions
        self._precompute_cache()

    def _precompute_cache(self):
        seq_idx = torch.arange(self.max_seq_len, dtype=self.inv_freq.dtype)
        freqs = torch.outer(seq_idx, self.inv_freq)

        # Create rotation matrix elements
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos()[None, None, :, :])
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :])

    def forward(self, q, k):
        # q, k: [batch, n_heads, seq_len, head_dim]
        batch_size, n_heads, seq_len, head_dim = q.shape

        # Apply rotary embeddings
        cos = self.cos_cached[:, :, :seq_len, :]
        sin = self.sin_cached[:, :, :seq_len, :]

        # Rotate half pattern (more efficient than complex number rotation)
        q_rot = self._rotate_half(q)
        k_rot = self._rotate_half(k)

        q_embed = q * cos + q_rot * sin
        k_embed = k * cos + k_rot * sin

        return q_embed, k_embed

    def _rotate_half(self, x):
        """Rotates half the hidden dims of the input."""
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)

# Example usage
rope = RotaryPositionEmbedding(dim=128)
q = torch.randn(2, 32, 1024, 128)  # [batch, heads, seq_len, head_dim]
k = torch.randn(2, 32, 1024, 128)
q_rotated, k_rotated = rope(q, k)
print(f"Q shape after RoPE: {q_rotated.shape}")

Part 6: Performance Analysis and Benchmarks

Comprehensive Performance Comparison

Memory Efficiency Deep Dive

Training Efficiency Analysis

Scaling Laws and Efficiency

Part 7: Architectural Failures and Lessons Learned

Lessons from Failed Optimization Attempts

Critical Insights from Failures

Part 8: Global Innovation Patterns

The Geography of LLM Innovation

The development of large language models has become a global endeavor, with different regions contributing unique innovations shaped by their constraints and priorities:

How Hardware Constraints Drive Innovation

Open Source vs. Closed Source Dynamics

Part 9: Future Directions

Emerging Architectural Trends

The Path to 100 Trillion Parameters

Fundamental Questions Remaining

Conclusion: The Architecture Convergence

As we survey the landscape of LLM architectures in 2024-2025, a remarkable pattern emerges: despite starting from different philosophies and constraints, the field is converging on a common set of architectural patterns:

The evolution from GPT-2 to DeepSeek V3 represents not a revolution but a careful refinement—thousands of small improvements that compound into dramatic gains. The transformer architecture, now seven years old, has proven remarkably robust and scalable.

Perhaps most importantly, the global nature of LLM development—with crucial innovations coming from China, the US, Europe, and beyond—demonstrates that advancing AI is truly a human endeavor. Different constraints breed different innovations, and the diversity of approaches strengthens the field.

As we look toward the future, the path seems clear: continued scaling, increased sparsity, and careful engineering optimization. The age of architectural exploration may be ending, replaced by an age of architectural refinement. The fundamental building blocks are in place; now we build higher.