Neural Network Abstractions: Module, Linear, Sequential

Course role: v1 is a conceptual bridge with compact scalar exercises. For the cumulative implementation sequence, numerical boundary cases, and learned multi-head projections, continue with deep learning v2. The v1 transformer exercise uses a single attention head even though its API accepts num_heads.

Goal

Replace ad-hoc classes with a small framework that matches PyTorch's mental model:

• Module base class with automatic parameter tracking
• Linear layer and common activations
• Sequential container for composition

Prerequisites: Autograd module.

---

1) The Module contract

A Module should:

1. Be callable via __call__ -> forward
2. Track parameters (Values) automatically
3. Track nested modules recursively
4. Support train/eval mode flags

Minimal state:

• _parameters: dict[str, Value]
• _modules: dict[str, Module]
• training: bool

---

2) Auto-registration via __setattr__

When a user assigns attributes, intercept and register:

class Module:
    def __init__(self):
        object.__setattr__(self, "_parameters", {})
        object.__setattr__(self, "_modules", {})
        object.__setattr__(self, "training", True)

    def __setattr__(self, name, value):
        if isinstance(value, Value):
            self._parameters[name] = value
        elif isinstance(value, Module):
            self._modules[name] = value
        elif isinstance(value, (list, tuple)):
            self._register_list(name, value)
        object.__setattr__(self, name, value)

Nested lists/tuples should register recursively (e.g., list of layers).

---

3) Parameter traversal

A parameter may have several names but is still one trainable scalar. Traverse registered parameters and children while retaining sets of visited parameter and module identities. Yield each parameter only on its first visit; named_parameters uses that first path as its name. Equality of numeric data is not enough: two distinct Value(0) objects are different parameters.

Work a shared-layer example: a parent stores Linear(2,3) under both layer and a one-element layers list. The child has 6 weights and 3 biases, so both enumerations must return 9 unique scalars, not 18. Plain recursive concatenation of both attribute paths would cause an SGD loop to update each twice even if its gradient were correct.

Sketch the traversal after attempting it

Maintain seen_modules and seen_parameters. On visiting a new module, emit its unseen own parameters in insertion order, then recurse through children. Retain these sets for the whole traversal rather than starting new ones in each child. parameters() can discard the names from named_parameters().

---

4) Linear layer (fully connected)

Shapes:

• input x: length in_features
• weight: (out_features, in_features)
• bias: (out_features,) or None

Forward:

y_i = sum_j weight[i][j] * x[j] + bias[i]

The initialization formulas distinguish fan-in from Glorot/Xavier conventions. Initialization here is uniform in [-1/sqrt(in_features), +1/sqrt(in_features)], a fan-in convention. Glorot/Xavier uniform instead uses a bound based on both fan-in and fan-out (with gain 1, sqrt(6/(fan_in+fan_out))).

---

5) Sequential container

Sequential stores modules in order and feeds outputs through each:

class Sequential(Module):
    def __init__(self, *modules):
        super().__init__()
        self.module_list = list(modules)
        for i, m in enumerate(modules):
            setattr(self, f"layer_{i}", m)

    def forward(self, x):
        for m in self.module_list:
            x = m(x)
        return x

---

6) Activation modules

Each activation maps a list of Value to a list of Value:

• Tanh: xi.tanh()
• ReLU: xi.relu()
• Sigmoid: use 1/(1+exp(-x)) for nonnegative x and exp(x)/(1+exp(x)) for negative x, preserving Value operations. This avoids exponent overflow at finite logits such as -1000.

These have no parameters but still behave like Modules.

---

7) Train/Eval mode

def train(self, mode=True):
    self.training = mode
    for m in self._modules.values():
        m.train(mode)
    return self

def eval(self):
    return self.train(False)

Used by layers like dropout/batchnorm (not implemented here, but the flag is part of the API). Calling eval() changes this mode flag; it does not disable Value graph construction or backward.

---

Key takeaways

1. Module centralizes parameter tracking and nesting.
2. __setattr__ is the hook that makes it automatic.
3. Linear + activations + Sequential are enough for many models.
4. Train/eval mode is an API contract, even if unused now.

Next: embeddings and positional encodings for sequence models.

Reconstruct a layer before composing it

Set a two-output layer's weight rows to [[1,2],[-1,3]], biases to [0.5,-0.5], and input to [2,1]. Calculate each output and the gradient of their sum with respect to the input before revealing the check.

Check the linear trace

Outputs are [4.5,0.5]. The sum's input gradient is [1+(-1),2+3]=[0,5]. Each weight receives its input coordinate's contribution; do not create fresh Value objects from weight.data during forward, or registration and the graph will disagree.

The ML pack's batched X@W stores weight columns per output; this scalar layer stores rows per output. Translate the storage layout explicitly. Then implement a new three-layer composition and predict its unique parameter count before running it.