What Neural Networks Learn
A neural network is a parameterized function that maps inputs to outputs. Training adjusts weights and biases so the model minimizes a loss function on data. For language models, the input is a sequence of token embeddings and the output is a probability distribution over the next token.
Depth matters because stacking layers with non-linear activations lets the network approximate complex functions. A single linear layer can only learn linear boundaries; multiple layers with ReLU, GELU, or SiLU activations can represent hierarchical patterns in text, images, and other modalities.
Refer: Understanding Gradient Flow in Deep Neural Networks · Deep Learning (Goodfellow et al.)
Forward and backward passes
The forward pass computes predictions layer by layer. The backward pass uses the chain rule to propagate gradients from the loss back to every parameter. Backpropagation is efficient because it reuses intermediate values instead of recomputing gradients independently for each weight.
import torch
import torch.nn.functional as F
# Next-token cross-entropy: compare logits to true token ids
logits = model(input_ids) # shape: [batch, seq_len, vocab]
loss = F.cross_entropy(
logits.view(-1, logits.size(-1)),
labels.view(-1),
ignore_index=-100,
)Loss Functions and Optimization
Language models are trained with cross-entropy loss: the model predicts a distribution over the vocabulary, and loss measures how far that distribution is from the true next token. Minimizing cross-entropy is equivalent to maximum likelihood estimation on the training corpus.
Optimizers
- SGD uses a single learning rate for all parameters. It works well with careful tuning but can be slow and unstable on large transformer stacks.
- Adam maintains per-parameter adaptive learning rates using running estimates of gradient mean and variance. It converges faster on noisy, high-dimensional problems.
- AdamW decouples weight decay from the gradient update. Adam applies L2 regularization inside the adaptive step, which interacts poorly with per-parameter scaling. AdamW applies decay directly to weights and is the default for LLM pretraining and fine-tuning.
optimizer = torch.optim.AdamW(
model.parameters(),
lr=2e-4,
betas=(0.9, 0.95),
weight_decay=0.1,
)Normalization and Stability
Deep networks suffer from vanishing or exploding gradients when repeated multiplication of small or large values shrinks or blows up signal during backprop. Modern architectures address this with residual connections, careful initialization, and normalization layers.
Batch, layer, and RMS normalization
- BatchNorm normalizes across the batch dimension per feature. It depends on batch statistics and behaves poorly with variable-length sequences or small batches.
- LayerNorm normalizes across the feature dimension for each sample independently. It is batch-size agnostic and was standard in early transformers.
- RMSNorm rescales activations by their root-mean-square without mean-centering. It is cheaper than LayerNorm and is used in LLaMA, Mistral, and many recent LLMs with similar stability.
Residual (skip) connections add the input of a sublayer to its output, giving gradients a direct path backward and enabling very deep stacks without degradation.
Activation Functions
Non-linear activations prevent stacked linear layers from collapsing into one linear map. Common choices in LLMs include:
- ReLU - simple, fast, but can produce dead neurons when inputs are always negative.
- GELU - smooth, probabilistic gating used in BERT and GPT-2 era models.
- SwiGLU / SiLU - gated variants common in modern decoder-only LLMs for better expressiveness in feed-forward blocks.
Generalization and the Bias-Variance Tradeoff
Underfitting (high bias) means the model is too simple to capture patterns in data. Overfitting (high variance) means it memorizes training examples and fails on new inputs. In LLM training, underfitting shows up as poor benchmark scores when the model or data scale is too small. Overfitting during fine-tuning appears when the model memorizes a small instruction set and loses broader capabilities.
Mitigations include more data, larger models (within compute budget), regularization (dropout, weight decay), early stopping, and parameter-efficient fine-tuning methods that limit how much of the base model can change.
From Foundations to Transformers
Everything above applies directly to transformer LLMs: cross-entropy next-token prediction, AdamW with warmup and decay, RMSNorm inside blocks, residual connections around attention and MLP sublayers, and gated activations in feed-forward networks. Understanding these primitives makes it easier to reason about training instability, fine-tuning behavior, and why architectural choices in modern LLMs look the way they do.