Optimizer state-to-parameter association (source code)

= Optimizer state-to-parameter association

<Optimizer state-to-parameter association> identifies which parameter each saved piece of <optimizer state> belongs to. An <optimization checkpoint> can contain every stored number and still resume incorrectly if those numbers are attached to different parameters. Equal shapes and data types do not establish mathematical identity.

The <PyTorch> optimizer dictionary separates per-parameter state from parameter-group settings. Its loader pairs saved parameter identifiers with the current parameters in group order, without independently verifying the intended association.<torch.optim in the PyTorch documentation>[PyTorch2.15, Optimizer.state_dict, state/param_groups and loading note]{ref}

For an original two-parameter example, use $b'=0.5b+g$ and $\theta'=\theta-0.1b'$. Let the named parameters be $a=1$ and $c=10$, their current gradients both be 1, and their saved <momentum buffers> be $b_a=2$ and $b_c=20$. Correct restoration gives
$$
b'=(2,11),\qquad(a',c')=(0.8,8.9).
$$
If a reversed parameter ordering attaches the same saved buffers to the wrong scalar parameters, the result is instead
$$
b'=(11,2),\qquad(a',c')=(-0.1,9.8).
$$
Both buffers had valid scalar shapes, so a shape check would not expose this error.

Consistently permuting the parameters, gradients and state is merely a change of ordering; permuting only the state changes the algorithm's continuation. Record a stable parameter schema and group order with the checkpoint, and check any intentional remapping against that schema. Parameter values and optimizer dictionaries are distinct saved objects, so restoring one does not automatically restore the other's meaning.<Saving and Loading Models in the PyTorch tutorials>[What is a state_dict?; Saving & Loading a General Checkpoint]{ref}