🏛️ LLM Architecture Deep Dive: Understanding Modern Language Models

Input Sequence (Embeddings)
    ↓
[Multi-head Self-Attention]
├─ Parallel attention heads
├─ Different representation subspaces
└─ Combined outputs
    ↓
[Add & Norm] (Residual connection)
    ↓
[Feed-Forward Network]
├─ Dense layer 1 (expand)
│  Dimension: d_model → 4*d_model
├─ ReLU/GELU activation
└─ Dense layer 2 (contract)
   Dimension: 4*d_model → d_model
    ↓
[Add & Norm] (Residual connection)
    ↓
Output (same shape as input)

PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

Benefits:
- Learned automatically
- Consistent across sequences
- Extrapolates to longer sequences

┌─────────────────────────────────┐
│   Input Embeddings              │
│   + Positional Encoding         │
├─────────────────────────────────┤
│  Encoder Block 1 (Self-Att)     │
├─────────────────────────────────┤
│  Encoder Block 2 (Self-Att)     │
├─────────────────────────────────┤
│  ...                            │
├─────────────────────────────────┤
│  Encoder Block N (Self-Att)     │
├─────────────────────────────────┤
│  Decoder Block 1 (Att + X-Att)  │
├─────────────────────────────────┤
│  Decoder Block 2 (Att + X-Att)  │
├─────────────────────────────────┤
│  ...                            │
├─────────────────────────────────┤
│  Decoder Block N (Att + X-Att)  │
├─────────────────────────────────┤
│  Output Layer (Softmax)         │
└─────────────────────────────────┘

Attention(Q, K, V) = softmax(Q·K^T / √d_k)·V

Where:
- Q (Query): What we're looking for
- K (Key): What information is available
- V (Value): What information to retrieve
- d_k: Dimension of key vectors (for scaling)
- √d_k: Prevents attention weights from becoming too small

MultiHead(Q,K,V) = Concat(head_1, ..., head_h)·W^O

Where each head uses:
head_i = Attention(Q·W_i^Q, K·W_i^K, V·W_i^V)

Decoder attends to Encoder:
Attention(Q_decoder, K_encoder, V_encoder)

Enables:
- Seq2Seq translation
- Information flow
- Encoder-decoder coupling

Token 1 sees: [Token 1]
Token 2 sees: [Token 1, Token 2]  
Token 3 sees: [Token 1, Token 2, Token 3]

Mask prevents looking at future tokens:
Attention_matrix[i, j] = -∞ for j > i

Traditional: Add position to embedding

RoPE: Rotate query/key vectors
- Encodes relative positions
- Better extrapolation
- Used in GPT-3, LLaMA, QwQ

Multi-Head Attention:
- h heads each with unique K, V

Grouped Query:
- Multiple Q heads share one K,V
- Reduces memory and computation
- Maintains performance

Multi-Query (Extreme GQA):
- All Q heads share one K,V pair
- Maximum efficiency

Standard Attention: Reads entire matrices from memory
└─ O(N²) memory bandwidth

Flash Attention: 
- Block-wise computation
- Minimizes memory I/O
- 2-3x speedup
- No approximation loss

Predict next token given context
Loss = -log P(token_t | token_1...t-1)

Used by: GPT family

Randomly mask tokens, predict them
Loss = -log P(masked_tokens | unmasked)

Used by: BERT, RoBERTa

Maximize similarity between related pairs
Minimize similarity between unrelated pairs

Used by: Embedding models

Normalizes hidden states per sample
- Speeds up convergence
- Reduces internal covariate shift
- Applied before or after attention/FFN

Trade compute for memory
- Store activations at checkpoints
- Recompute others in backward pass
- Enables training larger models

Use float16 for computations, float32 for weights
- Reduces memory by ~2x
- Maintains accuracy
- Better hardware utilization

Loss ≈ a * N^(-α) + b * D^(-β)

Where:
- N: Model size (parameters)
- D: Dataset size (tokens)
- α, β: Scaling exponents (~0.07)

Key finding: Similar diminishing returns for model & data

Optimal Allocation:
- Equal compute ~50% to model, ~50% to data
- Larger dataset ≠ worse results
- Training longer helps

Compute Budget:
- For fixed budget, optimal model usually smaller than trained
- But inference cost matters
- Trade-off between size and training duration

For C compute operations:

Optimal N ≈ C / 6
Optimal D ≈ C

Example:
- 1 trillion token budget
- Optimal: 20B parameter model, trained on 1T tokens

Selective SSM (Mamba):
├─ Linear complexity in sequence length
├─ Competitive performance with transformers
├─ Great for long sequences
└─ Simpler hardware usage

Architecture:
Input → SSM layers (not attention) → Output

Input
  ↓
[Router: Assign to experts]
  ├─ Expert 1 (FFN)
  ├─ Expert 2 (FFN)
  ├─ Expert 3 (FFN)
  └─ Expert N (FFN)
  ↓
[Combine expert outputs]

Advantages:
- Activate only relevant experts
- More parameters, same compute
- Example: Google's Switch-100B

Query + Context → Model
  ├─ Retrieve relevant documents
  ├─ Include in context
  └─ Generate answer

Benefits:
- External knowledge
- Up-to-date information
- Reduced hallucination

Text ─┐
     ├─[Shared Transformers]─→Output
Image─┤
Audio─┘

Cross-modal attention enables transfer learning

Factors to Balance:
├─ Latency: Response time
├─ Throughput: Requests/second
├─ Cost: Infrastructure expense
├─ Quality: Accuracy/coherence
└─ Availability: Uptime requirement