Transformers Architecture Cheat Sheet
A cheat sheet for the Transformer architecture covering self-attention, multi-head attention, positional encoding, and Hugging Face model usage.
3 PagesAdvancedMar 2, 2026
Scaled Dot-Product Attention
The core attention computation from first principles.
python
import torchimport torch.nn.functional as Fdef scaled_dot_product_attention(Q, K, V, mask=None): d_k = Q.size(-1) scores = Q @ K.transpose(-2, -1) / (d_k ** 0.5) # scale by sqrt(d_k) if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf')) weights = F.softmax(scores, dim=-1) return weights @ V, weights
Multi-Head Attention
Using PyTorch's built-in attention module.
python
import torch.nn as nnmha = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)attn_output, attn_weights = mha(query, key, value) # self-attention: query = key = value
Hugging Face Usage
Load a pretrained transformer and run inference.
python
from transformers import AutoTokenizer, AutoModeltokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')model = AutoModel.from_pretrained('bert-base-uncased')inputs = tokenizer('Hello world', return_tensors='pt')outputs = model(**inputs)last_hidden_state = outputs.last_hidden_state # shape: (batch, seq_len, hidden_size)
Key Concepts
Core theory behind the Transformer architecture.
- Self-attention- Each token computes a weighted sum over all other tokens, letting the model relate any two positions directly
- Query, Key, Value- Learned linear projections of the input used to compute attention scores and weighted outputs
- Multi-head attention- Runs several attention operations in parallel on different learned subspaces, then concatenates the results
- Positional encoding- Injects order information into token embeddings since attention itself is permutation-invariant
- Feed-forward network- Position-wise MLP applied after attention within each encoder/decoder block
- Residual connections & layer norm- Stabilize training and let gradients flow cleanly through very deep stacks
Sinusoidal Positional Encoding From Scratch
The original fixed (non-learned) encoding from Vaswani et al., useful when you can't afford learned position embeddings.
python
import torchimport mathdef sinusoidal_positional_encoding(seq_len, d_model): pe = torch.zeros(seq_len, d_model) position = torch.arange(0, seq_len).unsqueeze(1).float() div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) return pe # (seq_len, d_model), added elementwise to token embeddings
Causal (Autoregressive) Attention Mask
Prevent decoder tokens from attending to future positions during training.
python
import torchdef causal_mask(seq_len): # Upper-triangular ones above the diagonal indicate 'forbidden' positions mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool() return mask # True = mask out (future position)scores = scores.masked_fill(causal_mask(scores.size(-1)), float('-inf'))# In nn.MultiheadAttention / nn.TransformerDecoderLayer pass this as attn_mask,# or use is_causal=True on recent PyTorch versions to skip building it explicitly
Full Transformer Encoder Block
Attention + feed-forward with residual connections and pre-norm, the pattern most modern LLMs use.
python
import torch.nn as nnclass EncoderBlock(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout=0.1): super().__init__() self.attn = nn.MultiheadAttention(d_model, num_heads, dropout=dropout, batch_first=True) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.ff = nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model) ) self.dropout = nn.Dropout(dropout) def forward(self, x, attn_mask=None): # Pre-norm: normalize before the sublayer, then add the residual normed = self.norm1(x) attn_out, _ = self.attn(normed, normed, normed, attn_mask=attn_mask) x = x + self.dropout(attn_out) x = x + self.dropout(self.ff(self.norm2(x))) return x
KV Cache for Autoregressive Generation
Hugging Face's generate() reuses cached keys/values so each new token costs O(1) attention passes instead of O(n).
python
from transformers import AutoTokenizer, AutoModelForCausalLMimport torchtok = AutoTokenizer.from_pretrained('gpt2')model = AutoModelForCausalLM.from_pretrained('gpt2')input_ids = tok('The capital of France is', return_tensors='pt').input_idspast_key_values = Nonefor _ in range(10): with torch.no_grad(): out = model(input_ids[:, -1:] if past_key_values else input_ids, past_key_values=past_key_values, use_cache=True) next_token = out.logits[:, -1, :].argmax(-1, keepdim=True) past_key_values = out.past_key_values # cached K/V, reused instead of recomputed input_ids = torch.cat([input_ids, next_token], dim=-1)
Attention & Efficiency Variants
Modifications used in production LLMs beyond the vanilla multi-head attention in the encoder/decoder paper.
- FlashAttention- fuses the softmax(QK^T/sqrt(d))V computation into IO-aware kernels, avoiding materializing the full attention matrix in HBM; same math, much less memory bandwidth
- Multi-query / grouped-query attention- shares key/value projections across multiple query heads (MQA: one KV head, GQA: a few groups), shrinking the KV cache for faster inference at slight quality cost
- Rotary positional embeddings (RoPE)- rotates query/key vectors by an angle proportional to position instead of adding a positional vector, giving relative-position awareness and better extrapolation to longer contexts
- ALiBi- adds a fixed, non-learned linear bias to attention scores based on distance between tokens, enabling extrapolation to sequence lengths unseen during training without any positional embedding table
- Sparse / sliding-window attention- restricts each token's attention to a local window or fixed pattern to cut the O(n^2) cost, used in Longformer, Mistral's sliding-window layers, etc.
- Cross-attention- decoder queries attend over encoder keys/values (as opposed to self-attention where Q=K=V come from the same sequence); the mechanism behind encoder-decoder models like T5 and translation systems
Pro Tip
Attention has quadratic time and memory complexity in sequence length because every token attends to every other token — for long contexts, look at memory-efficient exact implementations like FlashAttention or sparse/linear attention approximations.
Was this cheat sheet helpful?
Explore Topics
#TransformersArchitecture#TransformersArchitectureCheatSheet#DataScience#Advanced#Scaled#Dot#Product#Attention#MachineLearning#CheatSheet#SkillVeris
