Skip to content

← nn.Module Mechanics step 3 of 7

Medium Primitives

What the optimiser updates

Jump to the task

A module can hold learned weights, running statistics and ordinary tensor attributes. How you register a tensor determines which module tools find it.

Tensor on the module In parameters()? In state_dict()? Moved by Module.to()?
nn.Parameter Yes Yes Yes
Registered buffer No Yes, unless persistent=False Yes
Plain tensor attribute No No No

A parameter represents a model parameter, usually learned during training. A buffer represents other tensor state, such as BatchNorm's running mean. Plain attributes are useful for ordinary configuration and temporary data, but a plain tensor does not automatically follow the model's device or checkpoint.

Registration and gradients are separate

Parameters require gradients by default. Buffers are usually created without gradients, but registering a buffer does not change its requires_grad flag. A plain tensor can also require gradients.

import torch
import torch.nn as nn

m = nn.Module()
m.weight = nn.Parameter(torch.ones(2))
m.register_buffer("mean", torch.zeros(2))
m.scratch = torch.ones(2, requires_grad=True)
print(list(dict(m.named_parameters())))  # ['weight']
print(sorted(m.state_dict()))            # ['mean', 'weight']
print(m.scratch.requires_grad)           # True

Passing m.parameters() to an optimizer gives it the registered parameters, not the buffers or plain attributes. Registration alone does not cause an update: the tensor must be in that optimizer and meet its update conditions. For example, PyTorch's SGD and Adam skip a parameter whose grad is None.

Why a running mean should be a buffer

A running mean is updated by the layer's statistics rule, rather than learned through gradient descent. Registering it as a parameter incorrectly includes it in parameters() and enables gradients by default. Whether an optimizer actually changes it depends on how it is used and on its gradient.

A frozen Parameter with requires_grad=False is valid for a frozen model weight. It is still a parameter, so freezing does not reclassify a running statistic as a buffer. Also, disabling new gradient recording does not clear an existing grad tensor; account for that when freezing a trained weight.

Buffers are not immutable. BatchNorm updates its running statistics during training. Keeping them out of parameters() prevents an optimizer constructed from that iterator from receiving them.

Your task

Complete Norm so weight is a trainable Parameter initialized to ones, and running_mean is a persistent buffer initialized to zeros without gradients. Then implement:

def norm_facts(n: int) -> dict

Build Norm(n) and return:

  • params: sorted names from named_parameters().
  • state: sorted keys from state_dict().
  • mean_needs_grad: whether running_mean requires gradients.

For n=4, the result is {"params": ["weight"], "state": ["running_mean", "weight"], "mean_needs_grad": False}. This exercise inspects registration; it does not run a training update.