Lecture 8: Neural Network Library Implementation

What this lecture covers

This is a hands-on implementation lecture that builds the modular components described in the previous lecture (modules, loss functions, optimizers) directly in needle. It first addresses a subtle but important bug: naive weight updates inside an autograd framework silently accumulate the entire computational graph across iterations, which wastes memory. It then walks through implementing Parameter, Module, a loss function, and an SGD optimizer, and closes with numerically stable softmax and a brief look at tuple-valued operators.

After watching, you can explain why weight updates must use detached tensors, implement a minimal nn.Module with recursive parameter collection, wire together a module, loss function, and optimizer into a training loop, and explain why naive softmax computation overflows for large inputs.

Key ideas

• In-place graph accumulation is a memory bug: writing w = w - lr * grad inside a loop, without detaching, keeps extending w's computational graph every iteration, retaining references to all previous computations and causing memory to grow linearly with iteration count.
• .data / .detach(): returns a new tensor sharing the same underlying cached_data but with empty inputs and op fields, so it is disconnected from the computational graph; needle also automatically produces detached tensors in eager mode when none of an operation's inputs require gradient.
• In-place weight mutation: assigning through w.data = ... updates a parameter's content without breaking the reference other parts of the model (such as a module's parameter list) hold to that same tensor object.
• Numerical stability of softmax: because floating-point numbers have limited precision (a sign bit, mantissa, and exponent), naively computing exp(x) / sum(exp(x)) overflows for large inputs; subtracting the max value from x before exponentiating keeps the computation stable without changing the mathematical result.
• Parameter: a subclass of Tensor with no special behavior beyond marking a tensor as a trainable parameter, which lets Module.parameters() distinguish parameters from other tensors.
• Module.parameters(): recursively walks a module's __dict__, collecting any Parameter, recursing into dictionaries and lists, and recursing into any nested Module's own parameters() call, giving automatic parameter discovery for arbitrarily composed modules.
• Loss as a module: implementing a loss (such as a simple squared-difference L2 loss) as a module lets it plug into the same forward/backward flow as any other layer.
• Optimizer interface: holds a list of parameters, exposes reset_grad() (clears each parameter's .grad) and step() (applies an update rule using detached data), and can keep auxiliary per-parameter state such as momentum for more advanced update rules.
• Modularity payoff: swapping the hypothesis (model), the loss function, or the optimizer in the training loop each requires changing only one line, in contrast to hand-deriving gradients and writing update rules from scratch.
• Tuple-valued operators: needle's autograd mechanism extends to operators that return multiple tensors as a tuple (such as FusedAddScalars), with gradients still propagating correctly through indexing into the tuple.

Walkthrough

The graph-accumulation bug in weight updates (2:05)

Starting from a plain weight tensor w with requires_grad=True, the lecture shows that a loop like w = w + (-lr * grad) keeps extending w's op and inputs fields every iteration, since each new w still refers back to all previous computations. Inspecting w.op, w.inputs, and nested .inputs.inputs confirms the growing chain. This is identified as a common source of out-of-memory errors when writing training loops in imperative autograd frameworks like PyTorch.

Fixing it with detached tensors (7:15)

The lecture explains that computational graphs are only needed to support reverse-mode automatic differentiation, which weight update rules themselves don't need. Calling w.data (a shortcut for w.detach(), which calls Tensor.make_const(self.realize_cached_data())) returns a new tensor sharing the same underlying cached data but with op set to None and empty inputs, so updates written as new_w = w.data - lr * grad.data no longer accumulate history. In-place mutation via w.data = ... is shown as the preferred approach, since it updates content without breaking references other structures hold to w.

Numerically stable softmax (15:22)

A naive softmax implementation (exp(x) / sum(exp(x))) works for small inputs like [0, 0, 1] but returns NaN for larger inputs like [100, 100, 101], because floating-point numbers cannot represent the resulting huge exponentials. Subtracting the maximum value from x before exponentiating is shown to leave the mathematical result unchanged (since exp(x_i - c) normalizes the same way for any constant c) while keeping all values well within representable range; the same principle applies to logsoftmax and logsumexp.

Building Parameter and Module (22:28)

Parameter is defined as a plain subclass of Tensor used only as a marker. Module.parameters() is implemented via a recursive helper, _get_params, that checks whether a value is a Parameter (return it), a dictionary (recurse over its items), or a Module (call its own .parameters()), pulling from self.__dict__. A ScaleAdd module (x * self.s + self.b, with self.s and self.b as parameters) is built and shown to return both parameters when .parameters() is called, and to compute forward outputs correctly when applied to a tensor.

Composing modules recursively (34:08)

A MultiPathScaleAdd module combining two ScaleAdd sub-modules along parallel paths (self.path0(x) + self.path1(x)) demonstrates that .parameters() automatically collects all four parameters across both sub-modules without any extra code, which is the same mechanism used to build residual modules in the homework.

Loss function, backward pass, and manual verification (37:11)

An L2Loss module computes the squared difference between a prediction and a label. Running the multi-path model on an input, computing the loss, and calling .backward() produces a gradient that the lecture verifies by hand (chaining the derivative of the squared loss through the scale-add path), confirming the gradient value matches the automatic computation.

Implementing SGD as an optimizer (41:18)

An Optimizer base holds a list of parameters and exposes reset_grad() and step(). A minimal SGD subclass stores the parameter list and learning rate, and its step() applies w.data = w.data - lr * w.grad.data for each parameter, relying on detached data to avoid graph accumulation. The lecture notes that more advanced optimizers, such as SGD with momentum, need to track additional per-parameter state (for example, a momentum tensor initialized to zero for each parameter).

End-to-end training loop (45:29)

A full loop is assembled: create input and label tensors, instantiate the multi-path model, the L2 loss, and the SGD optimizer, then for each epoch call opt.reset_grad(), run the model forward, compute the loss, call loss.backward(), and call opt.step(). The loss is shown decreasing over epochs, and the lecture emphasizes that swapping the model, loss function, or optimizer each requires changing only a single line of this loop.

Initialization variance and tensor tuples (brief) (49:52)

The lecture briefly revisits deriving the initialization variance needed to keep output variance close to input variance through a ReLU layer, matching the Kaiming-style derivation from the optimization lecture, and notes the same approach extends to uniform-distribution initialization. It closes with a short look at TensorTuple, needle's mechanism for operators (like FusedAddScalars) that return multiple tensors as a tuple, showing that automatic differentiation still works correctly when indexing into tuple-valued outputs.

Before you watch

• Complete or be familiar with the automatic differentiation homework, since this lecture assumes a working Tensor/autograd implementation and reuses its detach and realize_cached_data mechanics.
• Watch the previous lecture on neural network abstractions, since this lecture implements the nn.Module, loss function, and optimizer concepts it introduced.

Check your understanding

1. Why does writing w = w - lr * grad in a loop without detaching cause memory usage to grow with each iteration?
2. What does .data (or .detach()) return, and why does using it in a weight update avoid the graph-accumulation problem?
3. Why does naive softmax computation produce NaN for large input values, and why does subtracting the maximum value fix this without changing the result?
4. How does Module.parameters() find every trainable parameter in a nested module like MultiPathScaleAdd without being told the module's structure in advance?
5. What are the two core responsibilities of an Optimizer, and why might a more advanced optimizer need to store additional state beyond the parameter list?