2. Building Models
At the held-out input (0.25,-0.75), the network in our
running example initially predicts -0.088261. The target is 0.2.
After 200 Adam updates, its prediction is 0.228325. The network still has the same two linear
layers, the same eight hidden units, and the same ReLU between them. What changed was the set of
numbers inside its four parameter tensors. Understanding how those numbers determine a function,
and how a training program changes them, gives us a way to read more than the final loss in a log.
Write the parameters as \theta=(W_1,b_1,W_2,b_2). For an input x, the network computes
\begin{aligned}
h &= \operatorname{ReLU}(W_1x+b_1),\\
f_\theta(x) &= W_2h+b_2.
\end{aligned}
Each of the eight hidden units takes a weighted sum of the two input coordinates, adds a bias, and
replaces a negative result with zero. Over the input plane, a nonzero weight row and its bias define
a line where that unit switches on. Changing the weight row can rotate the line, while changing the
bias shifts it. The unit's output weight controls its contribution on the active side. This is how
a small ReLU network can represent a surface with several slopes. Our regression target has
precisely this structure: its slope changes along x_1+x_2=0 and x_2-x_1=0. The training data
gives us sampled values of that surface, and the optimizer uses the errors at those samples to
adjust the network.
The parameter layout follows directly from this calculation. W_1 has shape [8, 2], b_1
has shape [8], W_2 has shape [1, 8], and b_2 has shape [1]: 33 scalar parameters in
all. In Lean, these dimensions appear in the tensor types. The first layer produces eight values,
ReLU preserves that shape, and the second layer consumes eight values. Connecting it to a layer
that expects seven is a type error at the model definition. Once those dimensions agree, we can
change the parameter values throughout training while preserving the same interfaces between
layers.
The data has a corresponding structure. The quickstart stores 25 input pairs in a tensor of shape
[25, 2] and their targets in a tensor of shape [25, 1]. Each training item is one input of shape
[2] paired with one target of shape [1]. The leading 25 counts the available examples; it does
not mean that every update uses all of them. The recorded run takes one example per update. A
batched version carries a batch dimension through the model as well. Keeping that distinction
visible matters when we compare a per-step loss with a mean over the whole dataset, or compare
runs that process different numbers of examples.
For one input and target, mean squared error reduces to the squared difference between the network's single output and the target. That scalar connects a prediction to every parameter that contributed to it. Reverse mode propagates the loss sensitivity through the output layer, the ReLU, and the input layer, producing four gradient tensors with the same shapes as the parameters. Adam uses those gradients together with its running first and second moments to compute an update. The next example is evaluated with the new parameter values, so even a repeated input can produce a different prediction.
TorchLean's interfaces let us inspect each of these objects separately. The model definition
specifies the layer composition and parameter layout. Initialization supplies a starting state;
the dataset supplies examples; the objective and optimizer specify how learning proceeds.
Trainer.new brings these choices together, and trainer.train returns a result whose learned
parameters we can inspect, save, and use for prediction. This separation becomes useful as soon as
we change an experiment: trying another optimizer can preserve the model and initialization,
while loading a checkpoint supplies a particular parameter state for evaluation. The tensor and
model definitions below make those relationships explicit in the code.