Introduction to the AdamW Optimizer
Adaptive optimizers such as AdamW store additional parameters for each model weight, and that data is separate from the weights themselves.
The Belief That Weights Are Enough
When people first save a trained model, they often treat it as a single object: the numbers that make up the weights. The reasoning goes that if those numbers are written to disk and read back later, training can pick up exactly where it stopped. This belief feels natural, because the weights are what the model learned. It is also incomplete, and the gap matters most when the optimizer is an adaptive one such as AdamW.
This article is built around that belief. You will first see what goes wrong when the optimizer is left out of a checkpoint, then follow the code that saves and restores both pieces, and finally use the same reasoning to find the point where a resumed run stops matching the original.
What AdamW Keeps Beyond the Weights
Adaptive optimizers such as AdamW do not treat every update the same way with a fixed rule. They use historical data to adjust the learning rate for each model parameter dynamically. To make that possible, AdamW stores additional parameters for each model weight. Those parameters belong to the optimizer, not to the model's weights, and that separation is exactly why saving the model alone is not enough.
The diagram shows the key idea. Each weight tensor has its own companion values held by the optimizer. The weights and those companion values are separate objects, so the model's state dictionary does not contain them. Read the diagram as two columns that are linked weight by weight: the left column is what the model learned, and the right column is what the optimizer has accumulated about how each weight has been changing.
What a Reset Optimizer Does to Training
The source is direct about the consequence. If the optimizer state is not saved, the optimizer resets on restart. The model may then learn suboptimally, or even fail to converge properly. The source describes that failure in strong terms, noting that the model may lose the ability to generate coherent text.
The misconception this article targets is the assumption that a reset optimizer is harmless because the weights were preserved. The weights tell the resumed run where the model stands. The optimizer's saved history tells the run how it has been moving. A fresh optimizer starts with no such history, so the answer to the prediction question is no: the new optimizer does not behave as though it has seen the earlier training. Its updates are built from nothing, even though the model is already partway through learning.
Saving Both States into One File
PyTorch lets you save the model and the optimizer together. The model contributes its state_dict, and the optimizer contributes its own state_dict. Both are placed in one dictionary under two keys, and that dictionary is written to a single file with torch.save. The keys are names you choose; what matters is that the restore step later uses the same names.
- The optimizer is built from
model.parameters(), so the model must exist first - Two state dictionaries are placed under two keys in one dictionary
- torch.save writes that single dictionary to one file named checkpoint.pt
Notice what this does and does not do. It does not save the training loop, the data loader, or the code that defines the model. It saves the two state dictionaries, and only those, into the file named in the call. Restoring a run later depends entirely on what this one file contains.
Restoring in the Correct Order
Restoring is a two-stage operation. First the saved file is read back into memory with torch.load. Then the pieces are applied to freshly created objects with load_state_dict. The file itself does not reconnect to any live model or optimizer. That connection is made by you, in code, and the order of the steps matters.
- torch.load reads the file into memory as a dictionary
- The model is created before the optimizer, because the optimizer needs
model.parameters() - Each
load_state_dictcall is applied only after the matching object exists - The optimizer restore is the step a resumed run most easily loses
Read the sequence as a chain of dependencies. The file can be read at any point, but the two load_state_dict calls cannot run until the objects they apply to exist. The optimizer is created after the model's parameters exist, because its constructor needs model.parameters(). Its saved state is applied only after the optimizer itself has been built. Each step depends on a matching object already being in place.
Finding Where a Resumed Run Diverges
To locate a divergence, walk the restore sequence one step at a time and ask what the resumed run now holds that the original run also held. The first step where the resumed state stops matching the original is the place to look. In the sequence above, the weights are restored at step 4, but if step 5 is missing, the run has correct weights and a fresh optimizer. That mismatch is the divergence, and it appears before any training step runs.
A useful tool for this is inspecting what the checkpoint actually holds. The keys below are the ones used in the save example, so the check confirms whether the optimizer entry is really in the file before you blame the restore code.
- Lists the top-level keys stored in the file
- Confirms the model entry exists
- Confirms the optimizer entry exists, which is the entry a resumed run most often ignores
Common mistake Saving only the model's state_dict
What it looks liketorch.save(model.state_dict(), 'checkpoint.pt')Why it failsThe file then holds the weights but none of the optimizer's per-weight history, so the resumed optimizer starts from nothing.
Do this insteadSave a dictionary that holds both the model state dictionary and the optimizer state dictionary under two keys.
Common mistake Restoring the model and assuming the run has resumed
What it looks likemodel.load_state_dict(checkpoint['model_state']) with no optimizer load afterwardsWhy it failsA clean load does not show that the optimizer was restored. The source says a reset optimizer can leave the model learning suboptimally or failing to converge.
Do this insteadApply the optimizer's saved state with optimizer.load_state_dict after the optimizer has been created.
Common mistake Trying to apply optimizer state before the optimizer exists
Why it failsEach
load_state_dictcall depends on a matching object already being in place, and the optimizer is created frommodel.parameters().Do this insteadCreate the model, then create the optimizer from its parameters, then apply both saved states.
Common mistake Changing key names between save and restore
Why it failsThe restore code reads entries by key. If the names do not match, the lookup fails or the wrong entry is used.
Do this insteadUse the same key names in the save and restore code, and check them with a keys inspection when a restore misbehaves.
Diagnosing a Resume Script
- The file is read first
- The model and optimizer are created in dependency order
- Only the model's saved state is applied in this script
A resumed run produces noticeably different updates from the original run, even though the model weights loaded without error. Using the order of operations from the restore example, identify the single line most likely responsible for this script's behavior. Explain your reasoning in terms of what the optimizer is missing.
Hints
- Compare this script with the restore sequence: which of the five numbered steps is not present here?
- Ask what the new AdamW optimizer holds at the moment training continues, and where that history was saved in the checkpoint.
- A successful model load confirms only the weights; it does not restore the optimizer entry.
Key Takeaways
- AdamW keeps additional per-weight parameters that belong to the optimizer and are separate from the model's weights, so saving the model alone does not capture the training history.
- If the optimizer state is not saved, the optimizer resets on restart, which the source says can leave the model learning suboptimally or failing to converge.
- torch.save can write the model's
state_dictand the optimizer'sstate_dictinto one dictionary under two keys, saved to a single file. - A restore reads the file with torch.load, creates the model and then the optimizer, and applies each saved state with
load_state_dictin that order. - To find a divergence, walk the restore sequence step by step; the first step where the resumed state stops matching the original is the point to fix.