Neural Network Abstractions: Own Parameters, Preserve the Graph

Autograd answers “how does this loss depend on this scalar?” It does not know which scalars belong to a model or which ones should be updated. Module supplies ownership and composition. Linear supplies a trainable transformation. Sequential connects transformations without losing their parameters.

After this lesson, you should be able to count parameters before running code, trace a linear layer numerically, and distinguish parameter registration from graph construction. Implement nn-module first, then nn-linear; later exercises supply these building blocks so you can focus on language modeling.

The two structures in a model

The module structure persists across training steps. It contains parameter leaves and child modules. The computation graph is rebuilt on each forward pass and contains temporary activations, losses, and references to those leaves.

Calling model(x) delegates to model.forward(x). It does not create parameters. Construct weights in __init__, then reuse the same objects in forward. Creating fresh weights inside forward discards learning on every call; wrapping a weight's .data in a new Value disconnects the registered parameter from the loss.

Register a nested structure

Initialize _parameters, _modules, and training with object.__setattr__ before enabling automatic registration. The custom __setattr__ inspects assigned values:

Assignment	Registration
self.bias = Value(...)	One parameter
self.layer = Linear(...)	One child module
self.weight = [[Value(...), ...], ...]	Scalar parameters under names such as weight.0.1
self.layers = [layer_a, layer_b]	Child modules under layers.0, layers.1
self.width = 4	Ordinary metadata

Finally perform the actual assignment with object.__setattr__; otherwise Python's attribute and registry views disagree. The exercise assumes the structure is constructed once. In-place list mutation, removing attributes, and changing their kinds after construction are outside this small registry's contract.

An alias is still one parameter: if two attributes reference the same Value, parameters() must return it once by identity. Likewise, a child exposed through both a list and an attribute must not double the optimizer updates. Traverse with sets of visited parameter and module identities. named_parameters() uses the first path in traversal order (own parameters before children) as the canonical name for an alias. Equal numeric values in distinct objects remain distinct parameters.

Trace a linear layer

Linear(in_features=2, out_features=2) owns weight rows for output features:

W = [[2, -1], [0.5, 1]]
b = [1, -2]
x = [3, 2]

Compute each row's dot product with the entire input:

y[0] = 2*3 + (-1)*2 + 1 = 5
y[1] = 0.5*3 + 1*2 - 2 = 1.5

The result has length 2. More generally, weight has shape (out, in) and output has shape (out,). For loss L=y[0]+y[1], dL/dx=[2.5, 0]; each weight row receives gradient [3,2], and each bias receives 1. These values give you a stronger check than output shape alone.

The exercise initializes each weight uniformly in [-1/sqrt(in), 1/sqrt(in)] and biases at zero. This is a fan-in-scaled uniform choice. Xavier uniform, with gain 1, uses sqrt(6/(in+out)) as its bound, so these names are not interchangeable. See the initialization definitions.

Composition and parameter accounting

Consider Sequential(Linear(2,3), Tanh(), Linear(3,2)):

input (2,) -> affine (3,) -> tanh (3,) -> affine (2,)

The first layer owns 3*2+3=9 scalars. The second owns 2*3+2=8. The activation owns none. Total: 17 unique trainable scalars. Every call must return a list, including a linear layer with one output feature; changing a one-element vector to a scalar breaks composition.

Without an activation, two affine layers collapse algebraically to one affine layer. The activation introduces a nonlinear decision surface. Apply tanh, relu, or sigmoid to each Value; do not run math functions over .data in the model's forward pass.

Sequential stores execution order separately from the meaning of registration. Whether children are registered through a list or named attributes, apply each child exactly once in the requested order and expose every unique parameter once.

Training mode is a flag, not a gradient switch

train(mode) recursively sets child flags and returns self; eval() calls train(False). The linear and activation layers here behave identically in both modes because they have no dropout or running statistics. Evaluation mode does not disable autograd, clear gradients, or freeze parameters. Only the training loop performs updates.

Reconstruct and diagnose

Before coding, draw the module tree for the 17-parameter example and write five expected names, including layer_0.weight.0.0. Then inspect the implementation's parameter identities and names.

Checkpoint 1. A model returns 34 parameters for that architecture, but only 17 distinct identities. What failed?

Check your reasoning

a registry traversal visits aliased children twice. Gradients may be correct, yet an SGD loop applies each update twice.

Checkpoint 2. Forward outputs are correct but registered weights have zero gradients after backward. What should you inspect?

Check your reasoning

look for .data conversions or fresh Value(weight.data) objects in the forward calculation. Registration cannot repair a disconnected graph.

Once ownership, shapes, and gradients agree, the next lesson defines the token sequences those layers will learn from.