AI Frameworks

An array with any number of dimensions, held on CPU or GPU. A number is 0-D, a vector 1-D, an image 3-D, a batch of images 4-D.

image batch: [N, C, H, W] = [32, 3, 224, 224]

x = torch.randn(32, 3, 224, 224, device="cuda")
x.shape, x.dtype, x.device

The numbers the model learns. "Weights" usually means the multiplicative ones, "parameters" means all of them including biases. Model size is quoted in parameter count — 8B means eight billion.

y = Wx + b   →   W and b are the parameters

sum(p.numel() for p in model.parameters())

Running input through the model to get an output. Nothing is learned yet; this is just the prediction.

logits = model(x)

A single number saying how wrong the prediction is. Training is the process of making it smaller.

L = f(prediction, target)   →   lower is better

loss = F.cross_entropy(logits, y)

The model's raw output scores, before they are turned into probabilities. They can be any real number, positive or negative.

logits ∈ ℝ   vs   probabilities ∈ [0, 1] summing to 1

logits = model(x)                 # e.g. [2.1, -0.4, 0.9]
probs  = torch.softmax(logits, -1)  # e.g. [0.68, 0.06, 0.21]

Turns a vector of logits into probabilities that sum to 1. Exponentiating first means the largest logit dominates, which is why it is called soft max.

softmax(z)ᵢ = e^(zᵢ) / Σⱼ e^(zⱼ)

probs = torch.softmax(logits, dim=-1)

Squashes one number into the range 0 to 1. Used for binary problems where each output is an independent yes/no rather than a choice between classes.

σ(z) = 1 / (1 + e^(−z))

p = torch.sigmoid(logit)          # single probability

The standard classification loss. It measures how far the predicted probability distribution is from the true one. Because the true distribution puts all its mass on the correct class, it reduces to the negative log of the probability the model gave that class.

H(y, ŷ) = −Σᵢ yᵢ · log ŷᵢ   →   for true class c: −log ŷ_c

Predict 0.9 for the right answer and the loss is 0.105. Predict 0.1 and it is 2.303. Predict near zero and it explodes — which is exactly the intent: confident and wrong should hurt.

loss = F.cross_entropy(logits, targets)   # takes LOGITS, applies softmax itself
# equivalent, the long way:
loss = F.nll_loss(torch.log_softmax(logits, -1), targets)

The two-class version, applied independently to each output. Used for yes/no problems and multi-label problems where several labels can be true at once.

L = −[ y·log ŷ + (1−y)·log(1−ŷ) ]

loss = F.binary_cross_entropy_with_logits(logit, target)  # fuses sigmoid, numerically safer

The standard regression loss: average squared difference between prediction and target. Squaring punishes large errors far more than small ones, which also makes it sensitive to outliers.

MSE = (1/n) Σ (yᵢ − ŷᵢ)²

loss = F.mse_loss(pred, target)
loss = F.l1_loss(pred, target)     # mean absolute error - less outlier-sensitive

The derivative of the loss with respect to each parameter — how much the loss would change if that parameter moved slightly, and in which direction. Training moves every parameter a little way against its gradient.

∇L = [ ∂L/∂w₁, ∂L/∂w₂, … ]

loss.backward()
model.fc.weight.grad          # the gradient for that layer's weights

Computing those gradients by applying the chain rule backwards through the network, layer by layer, reusing each layer's result for the one before it.

∂L/∂w = (∂L/∂out) · (∂out/∂w)  — the chain rule

loss.backward()               # fills .grad on every parameter

The machinery that records every operation performed on a tensor and can therefore differentiate the whole chain automatically. You never write a derivative by hand.

x = torch.tensor(3.0, requires_grad=True)
y = x ** 2
y.backward()
x.grad        # tensor(6.) - because dy/dx = 2x

The record of which operations produced which tensors. PyTorch builds it as the code runs and throws it away after each backward pass; graph-mode frameworks build it once and reuse it.

loss.grad_fn                  # the node that produced this tensor

The rule that turns gradients into parameter updates. Plain SGD steps directly against the gradient; Adam and AdamW keep running averages of the gradient and its square so each parameter gets its own effective step size.

SGD: w ← w − α·∇L     (α = learning rate)

opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
opt.step()          # apply the update
opt.zero_grad()     # clear gradients, or they accumulate

How big a step to take against the gradient. Too high and training diverges; too low and it crawls or gets stuck. The single most important hyperparameter.

typical: 3e-4 for Adam, 1e-5 to 5e-5 for fine-tuning a pretrained model

sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)

A batch is the group of examples processed together before one update. An epoch is one full pass over the training set. A step is one parameter update.

steps per epoch = dataset size ÷ batch size

for epoch in range(10):
    for batch in dataloader:      # one step per batch
        ...

Summing gradients over several small batches before stepping, to simulate a large batch that would not fit in memory.

effective batch = micro-batch × accumulation steps × number of GPUs

(loss / accum).backward()
if (i + 1) % accum == 0:
    opt.step(); opt.zero_grad()

Overfitting is learning the training data itself rather than the pattern — training loss keeps falling while validation loss rises. Underfitting is the model being too weak to capture the pattern at all.

overfit: train ↓   val ↑   |   underfit: both high

Anything that constrains the model so it generalises instead of memorising: weight decay, dropout, data augmentation, early stopping.

L_total = L + λ·‖w‖²  (L2 / weight decay)

opt = torch.optim.AdamW(model.parameters(), weight_decay=0.01)
nn.Dropout(p=0.1)

Stopping training when validation performance stops improving, and keeping the best checkpoint. Must use a validation set separate from the test set, or the reported score is optimistic.

clf.fit(X_tr, y_tr, eval_set=[(X_val, y_val)])   # xgboost: early_stopping_rounds=50

Train fits the parameters. Validation tunes the choices you make by hand — learning rate, when to stop. Test is touched once, at the end. Reusing test for decisions turns it into a validation set and inflates the number you report.

How many bits each number uses. Fewer bits means less memory and faster arithmetic, at the cost of range or accuracy. bf16 keeps fp32's exponent range with fewer mantissa bits, which is why it trains stably where fp16 overflows.

fp32: 4 bytes · fp16: 2 · bf16: 2 · fp8: 1 · int8: 1 · int4: 0.5

model.to(torch.bfloat16)
# an 8B model: 32 GB at fp32, 16 GB at bf16, 4 GB at int4

Doing most arithmetic in 16-bit while keeping a 32-bit copy of the weights and the loss accumulation, so you get the speed without the numerical instability.

with torch.autocast("cuda", dtype=torch.bfloat16):
    loss = loss_fn(model(x), y)
loss.backward()

Storing weights and activations as low-precision integers, with a scale factor to map them back. Post-training quantisation converts an already-trained model; quantisation-aware training simulates the rounding during training so the model adapts to it.

q = round(x / s) + z   →   x ≈ s · (q − z)

model = AutoModelForCausalLM.from_pretrained(mid, load_in_4bit=True)

Not storing every intermediate activation for the backward pass, and recomputing them instead. Trades roughly 30% more compute for a large memory saving — often the difference between a model fitting and not.

from torch.utils.checkpoint import checkpoint
out = checkpoint(block, x)        # recompute this block's activations later

During generation, each new token attends to every previous one. Rather than recomputing their keys and values every step, they are cached. It makes generation fast and it is the main consumer of memory when serving.

size ≈ 2 × layers × heads × head_dim × seq_len × batch × bytes

out = model.generate(ids, use_cache=True)   # on by default

GPU memory holds four things: weights, gradients, optimiser state, and activations. Adam keeps two extra values per parameter, so optimiser state alone is typically twice the model size.

fp32 training ≈ 16 bytes/param: 4 weights + 4 grads + 8 Adam state

Operations run the moment you write them, like normal Python. Easy to debug, slower than a compiled graph because each operation is dispatched separately.

y = x + 1      # runs right now, y is a real tensor you can print

Running the function once to record the operations, then executing that recorded graph from then on. Faster, but any Python that depends on actual values is frozen at the value it had while tracing.

@tf.function              # TensorFlow
traced = torch.jit.trace(model, example)   # PyTorch

Just-in-time: compiling at runtime, once the shapes and types are actually known, rather than ahead of time. The first call is slow; every later call is fast.

model = torch.compile(model)     # first call compiles, then it pays off

A single function that runs on the GPU. One line of Python may launch several. Launching a kernel has fixed overhead, so many tiny kernels waste time.

Combining several operations into one kernel so intermediate results stay in fast registers instead of being written to GPU memory and read back. Memory traffic, not arithmetic, is usually the bottleneck.

unfused: read → add → write → read → relu → write fused: read → add → relu → write

A compiler that takes a whole graph of array operations and emits optimised machine code, fusing aggressively. It is what JAX and TensorFlow use, and how TPUs are targeted.

A point where the compiler cannot trace further — usually Python control flow that depends on a tensor's value — so it falls back to eager execution for that part. Correctness is preserved; speed is quietly lost.

print(torch._dynamo.explain(model)(x).break_reasons)

The order dimensions are stored in. PyTorch defaults to channels-first, TensorFlow to channels-last. Same numbers, different memory order — which is why conversion between frameworks means transposing weights.

NCHW: [batch, channels, height, width]  ·  NHWC: [batch, height, width, channels]

x.permute(0, 2, 3, 1)                       # NCHW → NHWC
x.to(memory_format=torch.channels_last)     # same shape, different layout

An open file format describing a model as a graph of standard operators, so a model trained in one framework can be run by another runtime.

torch.onnx.export(model, dummy, "model.onnx", opset_version=17)

A saved snapshot of the parameters. PyTorch's state_dict is a plain dictionary of tensors and carries no code, which is why the model class must be defined to load it.

torch.save(model.state_dict(), "ckpt.pt")
model.load_state_dict(torch.load("ckpt.pt"))

Every GPU holds a full copy of the model and processes a different slice of the batch. Gradients are averaged across GPUs after each backward pass so all copies stay identical.

memory per GPU = full model  →  largest model is capped by one device

model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank])

The collective operation that sums each GPU's gradients and gives every GPU the result. It is the communication cost that decides how well data parallelism scales.

each rank ends with: g = (g₀ + g₁ + … + g_{n−1}) / n

Splitting something across devices so each holds only a slice, rather than replicating it. What lifts the single-GPU memory ceiling.

Sharding applied to parameters, gradients, and optimiser state. Each layer's parameters are gathered just before use and released straight after. ZeRO's three stages shard progressively more: optimiser state, then gradients, then parameters.

stage 1: optimiser · stage 2: + gradients · stage 3: + parameters

Splitting a single matrix multiplication across GPUs — each computes part of the output, then the parts are combined. Needs fast interconnect, so it is normally kept within one node.

Y = XW  →  W split by columns: Y = [XW₁ | XW₂]

Putting different layers on different GPUs and passing activations along. The problem is the "bubble" — GPUs idle while waiting — which is reduced by splitting the batch into micro-batches that flow through continuously.

Using data, tensor, and pipeline parallelism together. Standard for frontier-scale pretraining.

total GPUs = data × tensor × pipeline

Moving parameters or optimiser state to CPU RAM or NVMe when not in use, and fetching them back on demand. Makes otherwise impossible runs possible, at a real throughput cost.

World size is the total number of processes; rank is one process's index; local rank is its index on its own machine. A node is one machine, usually with 8 GPUs.

$ torchrun --nproc_per_node 8 --nnodes 4 train.py

A matrix multiply plus a bias. The basic building block; most of a transformer's parameters and compute live in these.

y = xWᵀ + b

nn.Linear(in_features=768, out_features=3072)

The non-linear step between layers. Without one, stacking linear layers collapses into a single linear layer and depth buys nothing.

ReLU(x) = max(0, x)  ·  GELU(x) ≈ x · Φ(x)

nn.ReLU(), nn.GELU(), nn.SiLU()

Sliding a small learned filter across an image and taking a weighted sum at each position. The same filter everywhere means far fewer parameters than a linear layer, and a feature detected in one place is detected anywhere.

out[i,j] = Σ_m Σ_n in[i+m, j+n] · k[m,n]

nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)

Rescaling activations to a stable mean and variance so training does not destabilise. BatchNorm normalises across the batch; LayerNorm and RMSNorm normalise across features, which is what transformers use.

LN(x) = γ · (x − μ) / √(σ² + ε) + β

nn.LayerNorm(768)

A lookup table mapping a discrete id — a token, a category — to a learned vector. Similar meanings end up as nearby vectors.

vocab 32000 × dim 4096 → 131M parameters in that one table

emb = nn.Embedding(32000, 4096)
emb(torch.tensor([15, 892]))      # [2, 4096]

Models read integers, not text. The tokeniser splits text into sub-word pieces and maps each to an id. Roughly 4 characters per token for English.

tok("Hello world")["input_ids"]    # [9906, 1917]
tok.decode([9906, 1917])           # "Hello world"

Each position produces a query, a key, and a value. It compares its query against every key to get weights, then takes a weighted sum of the values. That is how a token pulls in context from the rest of the sequence.

Attention(Q,K,V) = softmax(QKᵀ / √d_k) · V

The √d_k divisor keeps the dot products from growing with dimension and pushing softmax into a region with vanishing gradients. Cost is quadratic in sequence length, which is why long context is expensive.

out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

Running attention several times in parallel on different learned projections, then concatenating. Each head can specialise — one on syntax, another on long-range reference.

head_dim = model_dim / n_heads  (e.g. 4096 / 32 = 128)

nn.MultiheadAttention(embed_dim=4096, num_heads=32, batch_first=True)

Attention followed by a small feed-forward network, each wrapped in a residual connection and a normalisation. Stack a few dozen of these and you have a modern model.

x ← x + Attn(LN(x));   x ← x + FFN(LN(x))

Adding a layer's input to its output. It gives the gradient a direct path back through the network, which is what makes very deep models trainable.

y = x + F(x)

An encoder reads the whole input at once and produces representations — good for classification and retrieval. A decoder generates one token at a time and may only look backwards — good for generation.

encoder: bidirectional (BERT) · decoder: causal (GPT, Llama)

The backbone is the pretrained feature extractor; the head is the small task-specific layer bolted on top. Transfer learning is keeping the backbone and replacing the head.

model = timm.create_model("resnet50", pretrained=True, num_classes=10)  # new head

Reducing many vectors to one — by mean, by max, or by taking a designated token. How a sequence of token vectors becomes a single sentence embedding.

sentence = token_embeddings.mean(dim=1)     # mean pooling

Pretraining learns general structure from a very large unlabelled corpus and costs millions. Fine-tuning adapts that model to your task with a small labelled set and costs hours.

Marking parameters as not trainable so gradients are not computed or stored for them. The cheapest way to cut memory when only part of a model needs to change.

for p in model.parameters():
    p.requires_grad = False

Instead of updating a weight matrix, learn a low-rank correction to it. Two thin matrices multiply out to the same shape as the original, so the frozen weight stays untouched and the update is tiny.

W' = W + BA   where B: d×r, A: r×k, and r ≪ min(d,k)

With d = k = 4096 and r = 16, that is 33.5M parameters replaced by 131K — about 0.4%.

LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj", "v_proj"])

LoRA on top of a base model quantised to 4-bit. The frozen base takes a quarter of the memory, the adapters train in 16-bit, and a 70B fine-tune fits on a single large GPU.

Any small trainable module added to a frozen model. Adapters are megabytes, can be swapped at inference, and several can be composed over one base model.

Training on example prompt-and-answer pairs so the model learns the format and behaviour you want. The first stage of post-training.

Reinforcement learning from human feedback. Humans rank model outputs, a reward model is trained to predict those rankings, and the model is then optimised against that reward — classically with PPO.

SFT → reward model → PPO against the reward, with a KL penalty to the SFT model

Direct preference optimisation. It shows that the RLHF objective can be rewritten as a simple classification loss over preferred and rejected pairs — no separate reward model, no RL loop.

L = −log σ( β · [ log(π/π_ref)(chosen) − log(π/π_ref)(rejected) ] )

Fine-tuning hard on a narrow task degrades the general ability the model had. Lower learning rates, fewer epochs, and adapter methods that leave the base weights frozen all limit it.

Latency is how long one request takes. Throughput is how many complete per second. Batching improves throughput and worsens latency, so you cannot maximise both.

for LLMs: TTFT (time to first token) and TPOT (time per output token)

Percentiles of the latency distribution. p95 = 200 ms means 95% of requests finished within 200 ms. Report percentiles, not the mean — the mean hides the tail that users actually feel.

torch.tensor(times).quantile(0.95)

Static batching fixes the group before it runs. Dynamic batching waits a few milliseconds to gather arriving requests. Continuous batching, used for LLMs, lets a finished sequence leave the batch and a new one join mid-flight rather than waiting for the slowest.

Storing the KV cache in fixed-size blocks like operating-system memory pages, instead of one contiguous allocation per sequence. It removes the wasted space from over-allocating for a maximum length, so far more requests fit at once.

Requests that begin with the same tokens — a shared system prompt, the same tool definitions — reuse the cached KV entries for that prefix instead of recomputing it every time.

Prefill processes the whole prompt in parallel and is compute-bound. Decode generates one token at a time and is memory-bandwidth-bound. They have opposite bottlenecks, which is why serving engines schedule them separately.

The first calls to a model are slow — memory allocation, compilation, caches filling. Always discard them before measuring, or the numbers are meaningless.

for _ in range(20): model(x)      # warm up
torch.cuda.synchronize()          # GPU work is async - sync before timing

A model compiled for one specific GPU architecture, precision, and shape range, as TensorRT produces. Fastest available, and not portable — a new GPU generation means a new build.

A fixed-length list of numbers representing a piece of text, an image, or an item, positioned so that similar things are close together.

emb = model.encode(["a sentence"], normalize_embeddings=True)   # [1, 768]

The angle between two vectors, ignoring their length. 1 means identical direction, 0 unrelated, −1 opposite. On normalised vectors it is just the dot product.

cos(a,b) = (a · b) / (‖a‖ ‖b‖)

sim = F.cosine_similarity(a, b, dim=-1)
sim = a_norm @ b_norm.T          # same thing if both are unit length

A bi-encoder embeds query and document separately, so documents can be indexed in advance and search is fast. A cross-encoder feeds the pair through together, which is far more accurate and far too slow to run over a whole corpus.

bi-encoder: encode once, compare many · cross-encoder: one forward pass per pair

Retrieve a shortlist cheaply, then reorder it with the expensive accurate model. Most practical retrieval quality gains come from this two-stage pattern rather than a better first-stage model.

Approximate nearest neighbour. Exact search compares against every vector; approximate methods trade a little recall for orders-of-magnitude speed, which is what makes million-vector search practical.

Retrieval-augmented generation: fetch relevant documents and put them in the prompt so the model answers from supplied evidence rather than memory.

Intersection over union — the overlap between two boxes divided by their combined area. 0 means no overlap, 1 means identical.

IoU = area(A ∩ B) / area(A ∪ B)

from torchvision.ops import box_iou
box_iou(pred_boxes, gt_boxes)

Non-maximum suppression. A detector fires several times on the same object; NMS keeps the highest-confidence box and discards any box overlapping it beyond an IoU threshold.

keep = torchvision.ops.nms(boxes, scores, iou_threshold=0.5)

Mean average precision — area under the precision-recall curve, averaged over classes. COCO's mAP@[.5:.95] averages over ten IoU thresholds, so it rewards tight boxes and not merely finding the object.

An encoder-decoder with skip connections carrying high-resolution detail across from encoder to decoder. Standard for segmentation, and the standard denoiser inside diffusion models.

An autoencoder that encodes to a distribution rather than a point, giving a continuous latent space you can sample from. In latent diffusion it is the compressor that lets generation run in a small space instead of on pixels.

z = μ + σ · ε,   ε ~ N(0, 1)  — the reparameterisation trick

Noise is added to an image over many steps, and a network learns to remove it one step at a time. The scheduler decides how much noise each step carries and how sampling walks back; fewer steps means faster and rougher.

x_t = √(ᾱ_t)·x₀ + √(1−ᾱ_t)·ε  — jump to any noise level directly

Classifier-free guidance. The model predicts with and without the prompt, and sampling extrapolates away from the unconditional prediction. Higher values follow the prompt more literally and reduce variety.

ε = ε_uncond + w · (ε_cond − ε_uncond)  (w ≈ 6–8 typical)

A series of yes/no splits on feature values, ending in a prediction. Individually weak and prone to overfitting, which is why they are used in ensembles.

Building trees one after another, each fitted to the errors the ensemble has made so far. Each tree corrects its predecessors rather than voting independently.

F_m(x) = F_{m−1}(x) + η · h_m(x)   where h_m fits the residual

Bagging trains many models independently on resamples and averages them, reducing variance — that is a random forest. Boosting trains them in sequence, reducing bias.

Level-wise grows every node at a depth before going deeper; leaf-wise always splits whichever leaf reduces loss most. Leaf-wise is more accurate per tree and overfits more easily on small data.

{"num_leaves": 63, "min_data_in_leaf": 100}   # the two that control it

Bucketing continuous feature values into a few hundred bins before searching for splits, so the search is over bins rather than every distinct value. The main reason LightGBM trains fast.

Turning categories into numbers. One-hot creates a column per value and explodes on high-cardinality features; target encoding replaces a category with the mean outcome for it and leaks unless computed carefully within folds.

When one class is much rarer than the other, a model that never predicts it can still be 99% accurate. Fix by weighting the loss, not by trusting accuracy.

scale_pos_weight = n_negative / n_positive

Splitting the data into k folds, training on k−1 and validating on the held-out one, k times. Gives a more reliable estimate than a single split, and a spread that tells you how stable it is.

cross_val_score(pipe, X, y, cv=5, scoring="roc_auc")

Information in the training features that would not be available at prediction time — including preprocessing fitted on the full dataset before splitting. It inflates offline scores and disappears in production.

Pipeline([("scale", StandardScaler()), ("clf", model)])   # scaler refits per fold

ROC AUC is the probability the model ranks a random positive above a random negative — 0.5 is chance. PR AUC uses precision and recall instead, and is the more informative of the two when positives are rare.

precision = TP/(TP+FP)  ·  recall = TP/(TP+FN)

A setting you choose rather than one the model learns — learning rate, tree depth, LoRA rank, batch size.

Grid tries every combination and wastes effort on parameters that do not matter. Random sampling beats it in practice. Bayesian search builds a model of which regions look promising and samples there.

Stopping a hyperparameter trial early once it is clearly behind, so the budget goes to candidates that might win.

trial.report(auc, step)
if trial.should_prune(): raise optuna.TrialPruned()

Recording the config, code version, metrics, and outputs of every run, so a result months later can be traced back to exactly what produced it.

Any file a run produces and you may need again — a checkpoint, a dataset snapshot, a plot — stored with the run rather than beside it on someone's laptop.

A versioned catalogue of trained models with stages such as Staging and Production, and a link back to the run that produced each. The audit trail a governed release process needs.

Being able to rerun and get the same result. Needs the seed, the library versions, the data version, and the config — and on GPU, still not exactly bit-identical unless you force deterministic kernels.

torch.manual_seed(0)
torch.use_deterministic_algorithms(True)   # slower, but repeatable