What the optimiser updates

A tensor attached to a module is one of three things, and the difference decides what happens to it.

                      parameters()   state_dict()   .to()   gradients
nn.Parameter               yes            yes        yes       yes
register_buffer            no             yes        yes       no
a plain attribute          no             no         no        no

Parameter is learned. It goes to the optimiser and receives gradients.

Buffer is state that belongs to the model and is not learned: BatchNorm’s running mean and variance, a causal mask, RoPE frequencies, a quantisation scale. Saved, moved, and never updated by an optimiser.

Plain attribute is none of the above, and the devices track covered what that costs.

Getting it wrong in the interesting direction

Registering something as a Parameter that should be a buffer is the subtler mistake, because everything appears to work:

self.running_mean = nn.Parameter(torch.zeros(n))

It is saved, it moves, and now the optimiser is applying gradient descent to a running average. Weight decay shrinks it toward zero every step. The model trains and slowly gets worse in a way that looks like a hyperparameter problem.

requires_grad=False on a Parameter is not the fix either: it stays in parameters(), so it still reaches the optimiser, and some optimisers will still touch it.

The question to ask

Does gradient descent have any business changing this? If yes, Parameter. If no but the model needs it, register_buffer. If neither, it probably should not be on the module at all.

Your task

Complete Norm so that weight is learned and running_mean is model state that is saved and moved but never updated, then:

def norm_facts(n: int) -> dict

reports the names in parameters(), the keys in state_dict(), and whether running_mean requires gradients.