# The answer is not a magic trick
We train the model with a very simple cycle, executed at a massive scale:
training_loop.txt
This cycle repeats thousands or millions of times. At each step, the model makes slightly fewer mistakes and its weights shift slightly. Training neural networks in PyTorch is built on exactly these core pieces: data, model, error computation, gradient propagation, and parameter updates.
# What does the model know at its very first moment?
A common misconception is that "training a model" means storing information inside it. But a model doesn't work like a library of sentences. In a neural network, what changes during training is primarily the model's numeric parameters.
initial_weights.py — python3
import torch
from torch import nn
# A freshly created layer — before any training:
layer = nn.Linear(3, 3)
print(layer.weight.data)
print(layer.bias.data)
▶ tensor([[ 0.12, -0.71, 0.33],
[ 0.52, 0.08, -0.41],
[ -0.29, 0.44, 0.17]])
▶ tensor([0.15, -0.04, 0.22])
These numbers are generated at the start of training using proper initialization schemes — not blindly random, but with distributions that make learning more stable (like Kaiming/He init for ReLU networks). The model doesn't "know" anything yet.
Why does initialization matter? If all weights are zero or identical, every neuron receives the same gradient and effectively becomes a duplicate — the network never learns. If they're too large, signals explode across layers. That's why the initial weight distribution is a critical design detail.
# Step one: data
To train, we give the model examples. Suppose we want to build a model that classifies clothing images. In machine learning, we deal with two things: the Input (what the model sees) and the Label (the correct answer).
dataset.txt
# Input → Label
image_001.png → T-shirt
image_002.png → Shoe
image_003.png → Bag
image_004.png → Dress
PyTorch manages such data with
Dataset and DataLoader. The Dataset holds samples and their labels; the DataLoader serves them to the model in batches.
Why don't we feed all the data at once?
Suppose we have 60,000 images. We don't need to push all of them through the network every time. Data is split into batches:
batches.txt
# Dataset: 60,000 samples, batch_size = 64
Batch 1 → 64 samples
Batch 2 → 64 samples
Batch 3 → 64 samples
...
Batch 938 → 32 samples # last batch (60000 % 64 = 32)
▶ 60000 / 64 ≈ 938 batches per epoch
This is more memory-practical, and it's also a core part of how modern models are trained: updating weights based on a gradient estimate from each batch (mini-batch gradient descent) is faster than sample-by-sample, and its noise actually helps avoid getting stuck in bad spots.
# Now the model must guess — and we measure how wrong it is
Suppose we feed the network an image of a shoe. In the earliest stages of training, the model probably produces something like this:
predict.py — early training
# Input: image of a shoe (label = Shoe)
# Model output after softmax (for intuition):
T-shirt 0.42
Shoe 0.15
Bag 0.33
Dress 0.10
✗ wrong prediction — but how wrong exactly?
The model was wrong. But a more important question remains: how do we know how wrong it was? That's where Loss comes in — a number that says how far the model's output is from the target. For classification, one of the most common functions is
CrossEntropyLoss:
loss.txt — over training
# Cross-Entropy Loss (lower = better)
▶ start of training: Loss = 2.41
▶ after a while: Loss = 0.83
▶ later: Loss = 0.21
▶ near the end: Loss = 0.07
Technical note — raw logits: In PyTorch,
nn.CrossEntropyLoss is actually a combination of LogSoftmax and NLLLoss. You feed it the raw output of the last layer (logits) directly — you must not apply Softmax yourself first. Doing so applies it twice and silently breaks training.
# But how does the model know which weight to change?
Here enters one of the most important parts of deep learning: Backpropagation. Suppose the model consists of thousands or millions of weights and ended up with
Loss = 1.82. Now we need to know which weights contributed to this error and by how much each should change. For that, we compute the gradient. Simply put, the gradient tells us: "if this parameter changes a little, how will the Loss change?"
gradients.py — python3
# After loss.backward():
# each parameter carries its own gradient
∂Loss/∂W1 = 0.72 # increasing W1 raises loss → decrease it
∂Loss/∂W2 = -0.14 # increasing W2 lowers loss → increase it
∂Loss/∂W3 = 0.03 # almost no effect → tiny change
These numbers specify the direction and magnitude of the required change. PyTorch performs this differentiation through its Automatic Differentiation (autograd) system: every operation on tensors is recorded in a computational graph, and
loss.backward() traverses it in reverse to compute every parameter's gradient via the chain rule.
# Gradient Descent: now change the weights
One of the simplest forms of weight update is:
update.py — python3
# W_new = W_old - learning_rate × gradient
W_old = 0.80
gradient = 0.20
learning_rate = 0.10
W_new = 0.80 - (0.10 × 0.20)
▶ W_new = 0.78
In a real network, this computation happens for a huge number of parameters. But one critical question remains: by how much? This is exactly where the Learning Rate becomes decisive:
learning_rate.txt
If
lr = 0.000001, the changes are tiny and training may crawl. If lr = 1.0, the model may overshoot the good spot, and the Loss may oscillate or never converge. So training isn't just "feeding data" — choosing hyperparameters is a crucial part of the process.
Beyond plain SGD: In practice, smarter optimizers like
Adam are usually used, which compute an adaptive step size per parameter (by keeping moving averages of gradients and their squares). Still, the core idea is the same for all of them: move against the gradient direction.
# Let's actually train a tiny AI
This is no longer just theory. With Python and PyTorch, we'll build and train a real neural network, using FashionMNIST — a dataset of 60,000 training and 10,000 test images (28×28) of clothing items. First, install:
terminal — bash
pip install torch torchvision
Step 1: Device and data
train.py — part 1
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
# 1) Device: GPU if available, else CPU
if torch.cuda.is_available():
device = torch.device("cuda")
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
device = torch.device("mps")
else:
device = torch.device("cpu")
# 2) Dataset + DataLoader
train_data = datasets.FashionMNIST(
root="data", train=True,
download=True, transform=ToTensor()
)
train_loader = DataLoader(
train_data, batch_size=64, shuffle=True
)
▶ Using: cuda (or 'cpu' on machines without GPU)
shuffle=True randomizes the training sample order every epoch — this prevents the model from learning the data order and makes each batch's gradient more representative of the whole dataset.
Step 2: Building the network
train.py — part 2
class NeuralNetwork(nn.Module):
def __init__(self):
super().__init__()
self.flatten = nn.Flatten()
self.network = nn.Sequential(
nn.Linear(28 * 28, 128),
nn.ReLU(),
nn.Linear(128, 64),
nn.ReLU(),
nn.Linear(64, 10)
)
def forward(self, x):
x = self.flatten(x)
return self.network(x)
model = NeuralNetwork().to(device)
architecture.txt
nn.Linear(28*28, 128) means the input has 784 features and this layer outputs 128. The network finally maps these down to 10 outputs — one per class. These outputs are logits, not probabilities — and as noted, CrossEntropyLoss takes exactly these raw logits.
Step 3: Loss, optimizer, and the training loop
train.py — part 3
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
epochs = 10
for epoch in range(epochs):
model.train()
for X, y in train_loader:
X, y = X.to(device), y.to(device)
prediction = model(X) # Forward
loss = loss_fn(prediction, y) # Loss
optimizer.zero_grad() # Clear old grads
loss.backward() # Backprop
optimizer.step() # Update weights
Three lines arguably matter more than everything else:
the_core.txt
▼ optimizer.zero_grad() — clear previous gradients
▼ loss.backward() — compute new gradients
▼ optimizer.step() — update the weights
✓ // this exact trio powers everything from this tiny MLP to billion-parameter LLMs
Why is zero_grad() needed? In PyTorch, gradients are accumulated by default — each
backward() adds the new gradient on top of the previous one. If you don't zero them before each batch, gradients from multiple batches mix and weight updates become wrong. This is one of the most common PyTorch training bugs.
Step 4: Evaluating the model
We must not test the model only on the data it was trained on — it may have memorized it. That's why a portion of the data (here: 10,000 test images) is set aside:
evaluate.py — python3
model.eval() # disable training-only behavior
with torch.no_grad(): # no gradient graph → faster, less memory
for X, y in test_loader:
prediction = model(X)
predicted_class = prediction.argmax(dim=1)
correct += (predicted_class == y).sum().item()
torch.no_grad() skips building the gradient graph during evaluation, saving memory and compute. model.eval() turns off training-only behaviors (like Dropout).
# What happens when you run it?
train.py — output
$ python train.py
▶ Epoch 1/10 | Loss: 0.52 | Accuracy: 0.812
▶ Epoch 2/10 | Loss: 0.39 | Accuracy: 0.848
▶ Epoch 5/10 | Loss: 0.29 | Accuracy: 0.876
▶ Epoch 10/10 | Loss: 0.23 | Accuracy: 0.889
Exact numbers depend on your environment, library versions, hardware, and settings — what matters is the trend. This chart shows the general shape of training this network — Loss descends smoothly and Accuracy climbs toward its ceiling:
training_curves.txt
Chart values are illustrative: The curves show a realistic shape for training this MLP on FashionMNIST (~88–89% test accuracy is typical for this simple architecture), but your exact numbers will differ based on hardware and random seed.
What is an epoch?
When the model has seen every training sample once, one Epoch is complete. With 60,000 samples and batches of 64, each epoch has about 938 batch steps. With
epochs = 10, the model sees the full training set about 10 times — and its parameters get updated roughly 9,380 times along the way.
# Overfitting: when the model studies too hard
Suppose the model performs superbly on training data but poorly on data it has never seen:
overfitting_check.txt
# A classic overfitting signature:
Training Accuracy 99.8%
Test Accuracy 72.0%
✗ gap = 27.8% → the model memorized, it didn't generalize
The model has probably become overly dependent on the quirks of the training set instead of learning generalizable patterns. This is called Overfitting. Depending on the problem, you can fight it with more data, data augmentation, regularization, dropout, early stopping, or architecture changes.
That's why serious projects split data into parts: Training for learning parameters, Validation for making decisions about settings and model selection, and Test for final evaluation on independent data. If you only watch the Training Loss, you might think the model keeps improving — while past some point the Validation Loss starts climbing, meaning generalization is degrading.
# The more exciting part: training a language model
The network we built could say "this image is probably a shoe." But language models (LLMs) tackle a far bigger problem. The objective of an autoregressive language model can be written as:
objective.py
# The objective of an autoregressive language model:
P(xₜ | x₁, x₂, ..., xₜ₋₁)
# "Given all previous tokens, how likely is each possible next token?"
"Given the previous tokens, how likely is the next one?" — the same idea we explored in the article on AI hallucinations: instead of having a "truth-checking module" at the core of generation, the model predicts the next token based on learned patterns.
First, text must be tokenized
tokenizer.py — python3
# "The cat is sleeping."
tokens = tokenizer.encode("The cat is sleeping.")
print(tokens)
▶ [464, 2415, 318, 7421, 13]
// a token is not necessarily a whole word —
// "unbelievable" may split into smaller pieces
The model no longer works with letters directly; it works with tensors and numeric representations. During training, we also have the true answer (the next token), so we can compare the model's output against it, compute the loss, backpropagate, and update weights — exactly the same cycle we saw for shoe detection.
Why Transformers?
For modern language models, the Transformer architecture plays a central role. The famous 2017 paper Attention Is All You Need showed that sequences can be processed with the Attention mechanism, without relying on traditional recurrent structures. The idea is simple but powerful: "while processing one token, which other parts of the input should I pay more attention to?"
attention.txt
"Ali put the book on the table because it was heavy."
# What does "it" refer to?
book 0.91
table 0.06
Ali 0.02
// attention lets the model link "it" to "book" across the sequence
But training an LLM is vastly bigger than the FashionMNIST example. The parameter count and data volume grow so large that training on a single ordinary computer becomes impractical — multiple GPUs, multiple machines, and distributed training infrastructure are required.
Four concepts you must not confuse
pipeline.txt
In the well-known InstructGPT work, researchers used human-written examples for supervised fine-tuning, then compared model outputs against human preferences to better align behavior with what people wanted. This isn't the only approach — methods like Preference Optimization also exist.
# The most important part we rarely talk about: data
Suppose you have the world's best architecture, but the data is duplicated, biased, incorrect, or noisy. What does the model learn? Most likely, a portion of exactly those problems. This is the famous Garbage In, Garbage Out: in machine learning systems, the quality of the data pipeline can matter as much as the model architecture.
data_pipeline.txt
Does the model just memorize?
Not quite that simple. A quick experiment: if the training data were only these three samples:
generalization_test.txt
# Training data:
2 + 2 = 4
2 + 3 = 5
2 + 4 = 6
# A model that only memorized:
17 + 23 = ??? # never seen → fails
# A model that learned the structure:
17 + 23 = 40 # generalizes to new inputs
Models may sometimes memorize parts of their training data, but machine learning cannot be reduced to pure memorization. During training, the model finds parameters that capture patterns in the data — and this difference between Memorization and Generalization is one of the field's most important topics.
# So what exactly was learned?
The short answer: the model's weights changed. We started from initial random weights (W₀), and at every training step the weights moved slightly against the gradient — W₁, W₂, W₃ — until reaching W_final: a set of parameters that makes the Loss better.
A small network might have a few hundred thousand parameters; a large model can have billions. But the core of the story remains the same:
universal_loop.txt
On one side, just a few lines of code:
prediction = model(X), loss = loss_fn(prediction, y), and those three golden lines zero_grad / backward / step. On the other side, one of the most complex computational systems ever built by humans. This contrast is perhaps the most fascinating part: the complexity of the result doesn't necessarily come from the complexity of the initial idea — sometimes it comes from the scale at which you run that idea.
takeaway.txt
AI starts with numbers, not intelligence.
It guesses, it errs, its weights get corrected —
and this cycle repeats billions of times.
That's what "learning" is:
slightly better than last time, every time.
It guesses, it errs, its weights get corrected —
and this cycle repeats billions of times.
That's what "learning" is:
slightly better than last time, every time.
