Skip to content

Chapter 2

Learn the Mathematics by Training One Small Neural Network

I introduce only the mathematics needed to inspect and train a small network: vectors, matrices, probability, derivatives, and optimization. A runnable PyTorch experiment makes learning visible while exposing overfitting and the cost of unnecessary complexity.

By Deepak Kumar · 13 min read · 2,986 words

How do vectors, matrices, and tensor shapes describe a neural network?

A neural network can run without learning anything useful, or fail before its first calculation because two arrays do not fit together. This chapter makes those failures inspectable by building a small question classifier for the study assistant, starting with the shapes of its numbers.

The classifier will distinguish definition requests from comparison requests. I would start with deliberately simple numerical features here, not a sophisticated text representation: you should be able to trace every input to its source. These features are teaching instruments, not a claim that keyword detection understands a question.

Give every axis a meaning

Represent each question using 3 numbers: whether it starts with “define,” whether it starts with “compare,” and its word count divided by 10. “Define a stack” becomes [1.0, 0.0, 0.3]. This ordered collection is a vector; each position has an agreed meaning. Change that order without changing the model, and you introduce a silent prediction error even though every shape still matches.

A single number is a scalar. Stack 4 question vectors as rows and you have a matrix with shape (4, 3): 4 examples, 3 features each. A tensor is the array container used for these numbers, including scalars, vectors, matrices, and higher-dimensional arrays. Its rank counts axes, not elements: this matrix has rank 2 and contains 12 values.

I write axis meanings beside shapes before writing layers. A shape is a useful contract only when you also know what its positions represent. Two unrelated arrays can both have shape (4, 3); the library cannot tell whether their columns describe questions or temperatures.

Follow one calculation through the network

A layer combines input features using adjustable numbers called weights, then adds adjustable offsets called biases. For one output unit, multiply each input by its corresponding weight and add the products plus its bias. With input [1.0, 0.0, 0.3], weights [0.5, -0.2, 1.0], and bias 0.1, the result is 0.9. Training will adjust those weights and biases; this section only follows the calculation.

The executable example below requires PyTorch in your project environment. Its first layer produces 4 intermediate values per question. nn.ReLU replaces negative values with zero, introducing a nonlinear operation. Without a nonlinear operation between these linear layers, their combined calculation could be expressed as one linear layer. The final layer produces 2 class scores.

import torch
from torch import nn

torch.manual_seed(7)

questions = [
    "Define a stack",
    "Compare stacks and queues",
    "Define a queue",
    "Compare arrays and linked lists",
]

def features(question):
    text = question.lower()
    return [
        float(text.startswith("define")),
        float(text.startswith("compare")),
        len(text.split()) / 10.0,
    ]

inputs = torch.tensor(
    [features(question) for question in questions],
    dtype=torch.float32,
)

model = nn.Sequential(
    nn.Linear(3, 4),
    nn.ReLU(),
    nn.Linear(4, 2),
)

hidden = model[0](inputs)
activated = model[1](hidden)
scores = model[2](activated)

print("inputs:", tuple(inputs.shape))
print("first weights:", tuple(model[0].weight.shape))
print("hidden:", tuple(hidden.shape))
print("scores:", tuple(scores.shape))
print("parameters:", sum(p.numel() for p in model.parameters()))

assert inputs.shape == (4, 3)
assert scores.shape == (4, 2)

The printed shapes are (4, 3), (4, 3), (4, 4), and (4, 2). PyTorch stores nn.Linear weights as output features by input features, so its calculation is inputs @ weight.T + bias. Matrix multiplication contracts the matching feature axes; it does not multiply corresponding elements. The bias is broadcast: its 4 values are added to every question’s intermediate row.

Catch errors before interpreting scores

This network has 26 adjustable parameters: 12 first-layer weights, 4 first-layer biases, 8 final-layer weights, and 2 final-layer biases. Adding parameters increases capacity, not guaranteed accuracy. I would keep this network small until held-out examples justify additional complexity.

If you transpose the inputs, the first layer raises RuntimeError with mat1 and mat2 shapes cannot be multiplied. If you construct integer inputs instead, the floating-point layer raises a dtype mismatch error. Keep dtype=torch.float32 explicit here. A more dangerous mistake produces no exception: swapping feature columns preserves dimensions while changing their meaning.

Interview check: Why does adding a fifth question change the output shape but not the parameter count? Because questions occupy the batch axis; each row uses the same weights. Try it and inspect scores.shape. The resulting scores are still untrained numbers—not probabilities. What would make one score represent a defensible preference for a class?

Why do probability, logarithms, and softmax turn scores into predictions?

A classifier can choose the correct category and still deserve a large penalty. Suppose your study assistant routes questions into 3 categories: concept explanation, calculation, and administrative information. For “What is a derivative?”, choosing concept explanation with probability 0.36 is a different prediction from choosing it with probability 0.95. Accuracy treats both as correct. Training needs a signal that distinguishes them.

The network’s final scores, called logits, are not probabilities. They may be negative, exceed 1, and need not sum to anything useful. I keep logits as the network’s output and convert them only where needed; this makes the training interface simpler and avoids a common numerical mistake.

Make the scores compete

For this classifier, each question has exactly one category. Softmax converts its 3 logits into positive probabilities that sum to 1: exponentiate each score, then divide by the sum of those exponentials. In symbols, p_i = exp(z_i) / sum_j exp(z_j), where z_i is one category’s logit.

Scores of 2, 1, and 0 become approximately 0.665, 0.245, and 0.090. The highest score still wins, but its probability depends on the competing scores. Adding 100 to every score leaves these probabilities unchanged because the common exponential factor cancels. Score differences matter; the absolute offset does not.

This construction fits mutually exclusive categories. If a question can carry several independent labels, forcing their probabilities to sum to 1 is the wrong modeling choice. Also, 0.95 is the model’s assigned probability, not a verified promise of 95% accuracy. Whether those numbers match observed frequencies requires checking held-out predictions.

Give confident mistakes a larger penalty

For one example with one correct category, cross-entropy loss is the negative natural logarithm of the probability assigned to that category: loss = -log(p_correct). A probability of 0.9 gives a loss near 0.105; 0.1 gives 2.303; 0.01 gives 4.605. Lower is better. The logarithm makes confident mistakes expensive without treating every incorrect prediction identically.

Why a logarithm? If you model labeled examples as independent, their joint probability is a product. Logarithms turn that product into a sum, so maximizing the joint probability becomes minimizing a sum of negative logarithms. PyTorch’s default mean reduction divides that sum by the number of examples. This supplies a useful training objective even when the winning category has not changed.

I would use torch.nn.functional.cross_entropy rather than manually combining torch.softmax and torch.log. Unless you need the probabilities for inspection, pass logits directly:

import torch
import torch.nn.functional as F

# Category order: concept, calculation, administrative.
logits = torch.tensor([[2.0, 1.0, 0.0]])
target = torch.tensor([0], dtype=torch.long)

probabilities = torch.softmax(logits, dim=1)
loss = F.cross_entropy(logits, target)

print("Probabilities:", probabilities.tolist())
print("Predicted category:", logits.argmax(dim=1).item())
print("Loss:", round(loss.item(), 6))

shifted_logits = logits + 1000.0
print("Shifted loss:", round(F.cross_entropy(shifted_logits, target).item(), 6))
assert torch.allclose(
    probabilities, torch.softmax(shifted_logits, dim=1)
)

The loss prints approximately 0.407606 for both inputs. Here dim=1 selects the category axis in the batch-by-category tensor. Normalizing over dim=0 instead makes examples compete with one another; with a batch of one, every category becomes 1. That is a silent failure, not a useful probability distribution.

Numerical stability is the other trap. Naively exponentiating 1000 in ordinary floating-point arithmetic overflows to inf; dividing infinities can produce nan. Stable softmax subtracts the largest logit before exponentiating. Stable log-softmax, used by cross-entropy, also avoids explicitly taking the logarithm of a probability that rounded to zero.

Do not feed softmax probabilities into F.cross_entropy: it treats them as logits and normalizes again, silently changing the objective. For these index targets, use integer category IDs; an ID of 3 for 3 categories raises IndexError because valid IDs are 0 through 2.

Try changing the correct category’s score from 2 to 4, then to −2. Explain why the loss changes even before checking the output. You now have a number that judges the prediction—but how does that number tell each weight which direction to move?

How do layers, activations, and loss functions express a learning problem?

A classifier can produce plausible-looking probabilities while learning the wrong task. If the study assistant must choose between a definition, comparison, and procedure question, its output needs 3 scores—not a number whose distance from an arbitrary category ID is minimized. I choose the output and loss before adding hidden layers. They specify what counts as a mistake.

Give the layers a concrete job

For this experiment, represent each question with 4 binary indicators: whether it contains “what,” “difference,” “how,” and “steps.” These deliberately crude features make the wiring inspectable; they are not a serious language-understanding system. The labels are definition = 0, comparison = 1, and procedure = 2. Those integers name categories; procedure is not twice comparison.

A nn.Linear layer computes weighted sums plus learned offsets called biases. Each output unit combines all 4 inputs, allowing “how” and “steps” to contribute to the same decision. A hidden layer is an intermediate representation rather than the final prediction. Here, 8 hidden units feed 3 output units, giving 67 trainable parameters, including biases.

Stacking these layers without a nonlinear activation wastes the extra machinery: two affine transformations still compose into one affine transformation. nn.ReLU keeps positive values and replaces negative values with zero. Its input-dependent gating prevents that collapse and lets the network express more complicated boundaries. It does not guarantee better predictions. I would keep a single linear classifier as the comparison unless held-out results justify these extra parameters.

Run this with PyTorch installed in your project environment. The three rows are illustrative questions, not a training dataset.

import torch
from torch import nn

torch.manual_seed(7)

# Features: contains "what", "difference", "how", "steps"
questions = [
    "What is a matrix?",
    "What is the difference between a vector and a matrix?",
    "How do I multiply matrices in steps?",
]
features = torch.tensor([
    [1.0, 0.0, 0.0, 0.0],
    [1.0, 1.0, 0.0, 0.0],
    [0.0, 0.0, 1.0, 1.0],
])
targets = torch.tensor([0, 1, 2], dtype=torch.long)

model = nn.Sequential(
    nn.Linear(4, 8),
    nn.ReLU(),
    nn.Linear(8, 3),
)

logits = model(features)
criterion = nn.CrossEntropyLoss()
loss = criterion(logits, targets)

assert logits.shape == (3, 3)
assert sum(p.numel() for p in model.parameters()) == 67

print("Logit shape:", tuple(logits.shape))
print("Mean loss:", loss.item())
for question, prediction in zip(questions, logits.argmax(dim=1).tolist()):
    print(question, "-> predicted category:", prediction)

Make the loss match the promise

The loss is a scalar penalty used to judge the current parameters. For one question with one correct category, cross-entropy penalizes the negative logarithm of the probability assigned to that category. Assigning the correct category probability 0.8 costs approximately 0.223; assigning it 0.1 costs approximately 2.303. Confident wrong predictions incur a larger penalty, even when their winning category remains unchanged.

nn.CrossEntropyLoss expects raw scores, or logits, and handles the log-probability calculation internally using numerically stable operations. Do not insert nn.Softmax before it. That usually produces no exception: the loss treats your probabilities as new logits, changing the objective and restricting their effective separation. This silent failure is more troublesome than a crash.

For this integer-target interface, targets must contain one valid category index per row and use torch.long. A target of 3 fails with Target 3 is out of bounds. because the available indices end at 2. Supplying only 2 targets for 3 rows raises a batch-size mismatch error. Neither failure means the network needs more layers; inspect the data contract first.

The default reduction averages the 3 example losses. The printed value therefore describes this batch, not accuracy, and no learning has happened yet. Random parameters can already produce a winning category through argmax; a prediction is not evidence of training.

Interview check: Why not use squared error on category IDs? A good answer explains that it imposes an invented ordering and distance: confusing category 0 with 2 becomes worse than confusing 0 with 1. Cross-entropy instead matches mutually exclusive category choices. If multiple categories may simultaneously be correct, that assumption fails and the output–loss pairing must change.

Exercise: Remove nn.ReLU and verify that the shapes remain unchanged. The program still runs, but its expressive capacity changes. What must happen to the weights for the loss to improve rather than merely print another number?

Can we follow gradients through a tiny PyTorch question classifier?

A training loop can print a falling loss while teaching you almost nothing about why the weights changed. I prefer to inspect one update before running 100 of them. If you cannot connect a gradient to a changed parameter, adding layers only gives you more places to hide a mistake.

For the study assistant, use a deliberately narrow task: distinguish definition questions, label 0, from comparison questions, label 1. Each question becomes 3 hand-built features: a definition cue, a comparison cue, and character count divided by 100. These are inspectable inputs, not a serious language-understanding system. Unless questions contain the chosen phrases, the cues miss them.

Run this in your existing environment with PyTorch installed. The network has 3 inputs, 4 hidden units, 2 output scores, and 26 trainable parameters, including biases. I use torch.optim.SGD here because its update is easy to audit: subtract the learning rate multiplied by the gradient.

import torch
from torch import nn

torch.manual_seed(7)

examples = [
    ("Define a matrix", 0),
    ("What is a vector?", 0),
    ("Define probability", 0),
    ("What is a derivative?", 0),
    ("Compare vectors and matrices", 1),
    ("What is the difference between mean and median?", 1),
    ("Compare training and validation", 1),
    ("Explain the difference between loss and accuracy", 1),
]

def features(question):
    text = question.lower()
    return [
        float("define" in text or "what is" in text),
        float("compare" in text or "difference" in text),
        len(text) / 100.0,
    ]

inputs = torch.tensor(
    [features(question) for question, _ in examples],
    dtype=torch.float32,
)
targets = torch.tensor(
    [label for _, label in examples], dtype=torch.long
)

model = nn.Sequential(
    nn.Linear(3, 4),
    nn.Tanh(),
    nn.Linear(4, 2),
)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

print("input:", tuple(inputs.shape))
print("parameters:", sum(p.numel() for p in model.parameters()))

for step in range(1, 101):
    optimizer.zero_grad(set_to_none=True)
    scores = model(inputs)
    loss = criterion(scores, targets)
    loss.backward()

    if step == 1:
        weight = model[0].weight
        before = weight.detach().clone()
        gradient = weight.grad.detach().clone()
        print("scores:", tuple(scores.shape))
        print("first-layer gradient:", tuple(gradient.shape))
        print("gradient norm:", gradient.norm().item())

    optimizer.step()

    if step == 1:
        expected = before - 0.1 * gradient
        print("update matches SGD:", torch.allclose(weight, expected))

    if step in (1, 20, 100):
        print(f"step={step}, loss_before_update={loss.item():.4f}")

with torch.no_grad():
    predictions = model(inputs).argmax(dim=1)
    accuracy = (predictions == targets).float().mean().item()
print(f"training accuracy={accuracy:.3f}")

Follow one update, not just the loss

The input shape is (8, 3); the scores have shape (8, 2). The first layer stores its weights as (4, 3), so model[0].weight.grad has that same shape. Each entry answers a local question: if this weight increased slightly, how would the current batch loss change, with every other parameter held fixed?

loss.backward() follows the recorded operations backward. For a first-layer weight, the derivative passes through the output layer, the hidden activation, and the input multiplication. The chain rule multiplies derivatives along each path; contributions from different paths and examples add together. With the default mean reduction in nn.CrossEntropyLoss, those example contributions are averaged across the 8 questions.

A positive gradient means a small increase in that weight would increase loss locally, so the update decreases it. A negative gradient reverses the direction. At learning rate 0.1, a gradient of 0.2 produces a change of −0.02. This is a local approximation, not a promise that every finite step lowers loss; an oversized step can overshoot.

The torch.allclose check should print True. It verifies the actual update against the subtraction, within numerical tolerance. The printed loss belongs to the forward pass before that step’s update. I label it explicitly because otherwise logs can appear to credit a weight change that has not been evaluated yet.

Failures worth causing deliberately

Remove optimizer.zero_grad and gradients accumulate across iterations instead of describing only the current batch. Accumulation is useful only when intentional; here it silently changes the experiment. Remove optimizer.step and no parameters change, even though gradients exist. Both failures can run without an exception.

Move the forward pass and loss calculation inside torch.no_grad, and backpropagation fails with element 0 of tensors does not require grad and does not have a grad_fn. Keep that context for the final measurement, not training. Also leave targets as torch.long: floating-point class indices commonly trigger expected scalar type Long but found Float.

Interview check: “Does backward train the model?” My answer is no: it computes and accumulates gradients; the optimizer changes parameters. As an exercise, replace the cues in one question with “Contrast vectors with matrices.” The network may classify it incorrectly even after achieving perfect training accuracy. Has it learned the distinction, or merely the wording used to teach it?

When does training accuracy hide failure, and what experiment would reveal it?

A question classifier can reach 100% training accuracy and still send every new question to the wrong destination. No exception appears. The tensors have the right shapes, the gradients are finite, and the optimizer is doing exactly what you asked. The failure is in what the examples reward.

For the study assistant, imagine routing administrative questions to course information and conceptual questions to study material. If all administrative training questions came from one file and all conceptual questions came from another, the classifier could learn the file identity instead of the question’s intent. I would check this possibility before adding another hidden layer. More capacity can make an accidental shortcut easier to learn.

A controlled shortcut experiment

This deliberately constructed experiment isolates that failure. Each input has 2 manually assigned features: an intent cue, encoded as −1 or +1, and a source-file cue, encoded as −4 or +4. Class 0 means conceptual; class 1 means administrative. The intent cue is perfectly informative here only to make the diagnosis unambiguous; real question features will not be this clean.

The 8 training examples have matching cues. In the 4 validation examples, the source cue reverses while the intent cue remains correct. These are synthetic feature rows, not measurements from real students. Run this in your existing PyTorch environment:

import torch
from torch import nn

torch.manual_seed(7)

train_x = torch.tensor([[-1.0, -4.0], [1.0, 4.0]] * 4)
train_y = torch.tensor([0, 1] * 4)
valid_x = torch.tensor([[-1.0, 4.0], [1.0, -4.0]] * 2)
valid_y = torch.tensor([0, 1] * 2)
loss_fn = nn.CrossEntropyLoss()

def run(use_source):
    model = nn.Linear(2, 2)
    nn.init.zeros_(model.weight)
    nn.init.zeros_(model.bias)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
    mask = torch.tensor([1.0, float(use_source)])
    x_train = train_x * mask
    x_valid = valid_x * mask

    for epoch in range(1, 41):
        model.train()
        optimizer.zero_grad()
        loss = loss_fn(model(x_train), train_y)
        loss.backward()
        optimizer.step()

        if epoch in (1, 10, 40):
            model.eval()
            with torch.no_grad():
                train_scores = model(x_train)
                valid_scores = model(x_valid)
                train_loss = loss_fn(train_scores, train_y).item()
                valid_loss = loss_fn(valid_scores, valid_y).item()
                train_acc = (
                    train_scores.argmax(dim=1) == train_y
                ).float().mean().item()
                valid_acc = (
                    valid_scores.argmax(dim=1) == valid_y
                ).float().mean().item()
            print(
                f"source={use_source} epoch={epoch:2d} "
                f"loss={train_loss:.3f}/{valid_loss:.3f} "
                f"accuracy={train_acc:.0%}/{valid_acc:.0%}"
            )

run(use_source=True)
run(use_source=False)

With use_source=True, training accuracy reaches 100% while validation accuracy is 0%. With use_source=False, both reach 100%. This second run is an ablation: remove one suspected cause while holding the remaining experiment fixed. I prefer this intervention to guessing from a loss curve because it directly tests whether the source feature causes the routing failure.

The mechanism is visible without a large network. In train_x, the second feature is always 4 times the first. Starting from zero weights, its gradient contribution is correspondingly larger. Training cannot distinguish “use intent” from “use source”: both explain every training label. On valid_x, those explanations disagree, and the stronger source contribution wins. More epochs reinforce the wrong rule rather than repair it.

Read the disagreement, not just the score

Watch train_loss fall while valid_loss rises. Accuracy counts wrong decisions; cross-entropy also penalizes increasing confidence in those decisions. If validation accuracy stays fixed while validation loss climbs, the classifier may be becoming more confidently wrong. That is a useful warning, not proof of one particular cause. Ordinary overfitting, inconsistent labels, or changed input distributions can produce similar symptoms.

model.eval() selects evaluation behavior for layers that need it; it does not disable gradients. torch.no_grad() prevents building the evaluation gradient graph. Neither repairs leakage or poor examples. Here nn.Linear behaves identically in both modes, but I keep the evaluation boundary explicit so later architectural changes do not silently change the experiment.

Interview question: “Training accuracy is high and validation accuracy is low. Should you increase model size?” My answer is no—not before checking representative mistakes, label consistency, and shortcuts. If training loss keeps falling after validation loss starts rising, an earlier checkpoint may help. Unless validation examples represent intended usage, however, selecting that checkpoint merely optimizes the wrong exam.

Exercise: Change the source magnitude from 4 to 0.5, keeping its reversed validation relationship. Predict the outcome before running it, then explain why feature scale changes the learned decision. Do not treat these 4 validation rows as evidence of deployment quality; they test one hypothesis.

The next chapter needs this habit of tracing behavior back to the training objective. Carry the distinction between fitting observed examples and learning a useful rule into the study of generative models.

All chapters

  1. 1
  2. 2
    Learn the Mathematics by Training One Small Neural Network
  3. 3
  4. 4
  5. 5
  6. 6
  7. 7
  8. 8
  9. 9
  10. 10
  11. 11
  12. 12
  13. 13
  14. 14
  15. 15
  16. 16
Message me