Basic Transformer

• Symbols
  • B – batch size
  • L – sequence length
  • d_model – embedding dims
• Step 1 – input pre-processing ( embeddings and position )
  • get token embedding
  • get pos embedding of size d_model for each token as func(pos, index)
    • there are d_model(say 128) index for each each token
    • even index with sign and old with cosine
• Step 2 – Multi head projection
  • Query, Key and Values
    • Different metrix for each attention head
    • If k heads, d_k = d_model / k
  • Intuition
    • allows model to look at same sentence from multiple affront angles in parallel
    • Group of experts looking at painting
      • one expert looks at colors
      • another at shops
      • 3rd historical context of artist
• Step 3 – Attention equation
• $$\mathbf{A}_h = \text{Softmax}\left( \frac{\mathbf{Q}_h \mathbf{K}_h^T}{\sqrt{d_k}} + \mathbf{M} \right) \mathbf{V}_h, \quad \text{where } \mathbf{M}_{ij} = \begin{cases} 0 & \text{keep} \\ -\infty & \text{mask} \end{cases}$$
  • Intuitions
    • 1) query matches all keys softly ( weight 0 to 1 ), values are multiplied by the weight and summed
      • In python dictionary, query matches just one key, returning its value
    • 2) query – what am I looking for ?
      • key – what do I contain ?
      • value – what information do I provide ?
• Step 4 – head aggregation
  • concat and project
  • intuition
    • head never talked to each other before
    • weighed matrix act as committee chair
• Step 5 – residual connection and layer norm
• Step 6 – Position wise FFN
  • Token wise, there is no mixing
    • Same weight applied to each token
  • Two linear transformation
    • Expansion-contraction loop
    • First one has activation, second one does not
  • Intuition
    • 1) Information funnel
      • expansion creates features
      • activation selects those
      • contraction brings them back to match residual connection
      • ReLU after contraction will clamp negative values to zero resulting in information loss
    • 2) Residual connection
      • Model will loose ability to subtract value, dampen signals or reverse feature directions
    • 3) Separation of concerns
      • Attention layer as token mixer
        • Look across and let tokens talk to each other
      • FFN layer as channel mixer
        • Look inside individual token’s vector to process its features
• Step 7 – second residual connection
• Intuition of residual connections
  • 1) Arch argument
    • Creates residual stream, nothing overwrites, everything accumulates
    • Allows to learn identify function
      • without it, it becomes hard to reconstruct input using non linear weights
    • Preserving low level features ( position, syntax )
      • deeper layer extract higher level semantic meaning like sentiment
  • 2) Optimization argument
    • Avoid vanishing in deep networks
      • even if gradient through complex layer shrinks to 0, +1 guarantees gradient can pass till first layer
• Math behind attention scaling
  • when d_k grows large, dot product grows large in magnitude, pushing softmax to region with extremely small gradients ( vanishing gradient problem )
  • Mathematical proof shows variance of d_k and std deviation of sqrt(d_k)
  • We kept the variance of 1, we don’t want to loose expressive power either

Norms

• Moved layer norms in the residual branch, keeping the highway clear.
  • Its also called pre-norm

• Recent models are moving from layerNorm to RMS because
  • 1) Few params means better bandwidth utilization ( memory bandwidth being GPU perf bottleneck )
  • 2) Single pass kernel fusion
    • don’t need to calculate and store mean
  • 3) Fewer intermediary states in memory for backprop
• Inspired by this people are dropping bias terms too
• On batch norms
  • Not used in transformers
  • Batch statistics are unstable and unreliable for variable length data ( padding zero skews mean, variance )
  • Autoregressive decoding makes it impossible

Position Embeddings

• RoPE
  • Rotary position embeddings
  • Goal – relative position <f(x,i), f(y,j)> = g(x, y, i-j)
  • The distance between token matters, wherever they occur in the sentence
  • Core idea
    • We know inner product are invariant to arbitrary rotation
    • Position i vector gets rotated by i position

Activation Units

• Gated activation units
  • ReLU shuts down some channel in single shot
    • 𝐹𝐹 𝑥 = max(0, 𝑥𝑊 1 )𝑊 2
  • GLU allows to do that as a function, looking at all channels
    • FFReGLU 𝑥 = (max (0, 𝑥𝑊 1) ⊗ 𝑥𝑉) 𝑊 2
• GLU idea can be applied to GeLU and Swish as well
  • Which results in GeGLU and SwiGLU
  • On their own GeLU and Swish are self gated – x is multiplied by gate computed from x itself
    • Swish = x * sigma(beta * x)
    • GeLU = x * phi(x)
    • GLU unities two giving gate, its own weight and more freedom

Group query attention

Attention Type	Query Matrices	Key Matrices	Value Matrices
Multi-Head (MHA)	H distinct matrices	H distinct matrices	H distinct matrices
Grouped-Query (GQA)	H distinct matrices	G distinct matrices (G < H)	G distinct matrices (G < H)
Multi-Query (MQA)	H distinct matrices	1 shared matrix	1 shared matrix

• Intution
  • MHA – asking 32 detective to solve mystery, asking them to read same book from personal topic
  • GQA
    • 8 groups of 4
    • Every 4 heads share single key/value ( shared book )
    • Unique query matrix means detective still form their individual conclusions
• Advantages
  • Very successful, performance close to MHA
  • Reduces KV cache size

Mixture Of Experts

• Different FFN layer for different tokens
  • Not at the sentence level
• Router is important architecture component
  • Decides which token should go to which expert
    • Top k based on inner product
      • Most common approach
      • Have some learnable weight vector for each expert
      • Use softmax to select it, given its differentiable gradient can propagate back
      • But “topK” is not
        • problem of “expert collapse” – some experts are abandoned
        • Noisy gating, add some gaussian noise to dot product logits before “topK”
    • RL to learn router
      • Lot of overhead, not used anymore, 2013 paper
    • Hashing
      • Simple yet works
• Intuition
  • Receptionist at hospital route to expert doctor
    • You don’t ask everyone yourself
  • Sparse activation
    • Only tiny fraction of model is activated for a given token
• Key Ideas
  • Specializaton
    • Individual sub network naturally learn to specialize – code, grammar, language
  • Router
    • Assigns probabilities
  • Separating size from cost
    • Scale model params without cost
• Main advantages
  • Scale without FLOPs
  • Parallizable across devices
    • GPU 0 hosts expert 1 and 2
    • Earlier we were slicing the matrix
• Memory bottleneck
  • Results are showing improvement by keeping the params constant with MoE

On Regularization

• SGD with single pass is haed to memorize
• In the compute constraint deep learning its not needed
  • Dropout is not used anymore
  • But weight decay ( L2 regularization) is still used with different purpose though
    • Studies have shown that it helps with optimization

Sparse/Sliding Window Attention

• Instead of entire attending entire context, can we focus on some nearby ( sliding windows )
• Standard practice –
  • Interleave full and local attentions
  • Ex – one full attention after every 3 layers