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.^[ref]

For an original two-parameter example, use $b^{\prime }=0.5b+g$ and ${\theta }^{\prime }=\theta -0.1b^{\prime }$. 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 givesIf a reversed parameter ordering attaches the same saved buffers to the wrong scalar parameters, the result is insteadBoth 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.^[ref]