The loss (MSE)

Up to here your network can compute, but it has no opinion about whether the result was any good. That opinion is the loss: a single number measuring how far off the prediction was. Small is good, zero is perfect.

That it has to be exactly one number is not a convention but a necessity. The backward pass needs a starting value, and that is only unambiguous if there is a scalar at the top. With a vector of ten outputs there would be ten starting points and no answer to the question which one applies.

The goal

Learning means making one single number smaller. That is the entire aim of a training run, and it is remarkably reductionist: whether a model recognises cats or translates sentences is decided by which number you define and in which direction you push it.

The mean squared error is the simplest of those numbers:

1. Take the difference between prediction and target.
2. Square it - that way deviations up and down count the same, and large errors weigh disproportionately heavily.
3. Average over all elements.

Three steps, and all three consist of operations your tensor already knows. Only one is new: collapsing many numbers into one, that is sum. It needs a backward step too, and it is the simplest in the whole course - every element of the input contributed exactly one to the sum.

Zero magic in this step

NumPy is allowed, torch, jax, tinygrad and autograd are forbidden. Even np.mean on the raw array would be a mistake here - not because NumPy were forbidden, but because a computation outside the tensor class falls out of the graph. What is not recorded cannot run backwards, and the loss would be mute.

The exercise

First a new method in meintorch/tensor.py:

    def sum(self):
        out = Tensor(self.data.sum(), (self,), "sum")

        def _backward():
            # TODO(du): every element contributed the same to the sum.
            raise NotImplementedError

        out._backward = _backward
        return out

Then meintorch/loss.py:

from .tensor import as_tensor


def mse_loss(pred, target):
    """Mean squared error as a scalar tensor."""
    target = as_tensor(target)
    # TODO(du): difference, square, sum, divide by the count.
    # All of it via tensor operations, so the graph stays intact.
    raise NotImplementedError

Hint 1 - the direction

For sum(): if the sum changes by one, every single summand has the same share in that. The incoming gradient is a scalar and has to be spread across the shape of the input.

For mse_loss: squaring does not need a ** operator. You already have something that multiplies two tensors elementwise - and two identical tensors are also two tensors. Dividing by the count is a multiplication by the reciprocal, and your __mul__ already knows that as a scalar.

Hint 2 - more concrete

• sum() backward step: self.grad += np.ones_like(self.data) * out.grad.
• mse_loss: diff = pred - target (the __sub__ from chapter 02), n = diff.data.size, then (diff * diff).sum() * (1.0 / n).
• You may read the count n straight off the NumPy array. It is a constant, not an intermediate result, so it does not have to enter the graph.

    # in Tensor.sum
        def _backward():
            self.grad += np.ones_like(self.data) * out.grad


# in loss.py
def mse_loss(pred, target):
    target = as_tensor(target)
    diff = pred - target
    n = diff.data.size
    return (diff * diff).sum() * (1.0 / n)

In goal mode

The autonomous run does not just check that something comes out, but that the right number comes out - the gradient of the MSE can be recomputed by hand and is therefore a hard test:

Extend `meintorch/tensor.py` with a method `sum()` that collapses all elements
into a scalar and brings its own backward step. Create `meintorch/loss.py` with
`mse_loss(pred, target)`: mean squared error, exclusively via tensor operations,
so that the graph stays intact.
Mandatory: NumPy allowed, torch/jax/tinygrad/autograd forbidden. No arithmetic
on raw arrays outside the tensor class.
Write `tests/test_04_loss.py`: the loss has to be a scalar (shape ()), and
`backward()` has to run from it through a network of two linear layers with a
ReLU. Additionally check the gradient at the prediction tensor against the
analytic formula 2*(pred-target)/n.
Run `python -m pytest` until everything is green.

The milestone test

tests/test_04_loss.py. Two things have to hold: the loss is a scalar, and backward() runs from it all the way into every weight.

import numpy as np

from meintorch.loss import mse_loss
from meintorch.nn import Linear, ReLU, Sequential
from meintorch.tensor import Tensor


def test_loss_is_a_scalar():
    loss = mse_loss(Tensor([[1.0], [2.0]]), Tensor([[0.0], [0.0]]))
    assert loss.shape == ()
    assert np.isclose(loss.data, 2.5)


def test_a_perfect_prediction_has_zero_loss():
    t = Tensor([[3.0, -1.0]])
    assert np.isclose(mse_loss(t, Tensor([[3.0, -1.0]])).data, 0.0)


def test_gradient_at_the_prediction_tensor():
    pred = Tensor([[1.0], [2.0]])
    loss = mse_loss(pred, Tensor([[0.0], [0.0]]))
    loss.backward()
    # 2 * (pred - target) / n, so 2 * [1, 2] / 2 here
    assert np.allclose(pred.grad, [[1.0], [2.0]])


def test_backward_runs_through_the_whole_network():
    net = Sequential(Linear(3, 4, seed=0), ReLU(), Linear(4, 1, seed=1))
    x = Tensor(np.random.default_rng(7).normal(size=(5, 3)))
    y = Tensor(np.random.default_rng(8).normal(size=(5, 1)))

    loss = mse_loss(net(x), y)
    loss.backward()

    assert loss.shape == ()
    for p in net.parameters():
        assert p.grad.shape == p.data.shape