Checkpoint Averaging

Vaswani et. al. 2017 and Izmailov et. al. 2018 average the last N checkpoints during training, which leads to better generalisations and flatter optima. Wortsman et. al. 2022 propose model soups, whereby weights of multiple fine-tuned models are averaged to produce a single model that combines the strengths of the individual fine-tuned models. We have also found checkpoint averaging to improve accuracy (~1-2% relative reduction in WER), by averaging the best N checkpoints (from the same training run) based on validation WER.

To average the weights of a set of checkpoints, run the following command:

python caiman_asr_train/export/checkpoint_averaging.py --ckpts path/to/ckpt1.pt path/to/ckpt2.pt path/to/ckpt3.pt --output_path path/to/avg_ckpt.pt --model_config path/to/config.yaml

This script simply averages the weights of a list of checkpoints, and does not find the best N checkpoints from a given training run - this must be done manually at the moment, for example, by checking tensorboard logs.

The averaged checkpoint carries the blank bias that the input checkpoints were tuned to. Only the checkpoints saved at the last training step store a tuned value, so averaging checkpoints from one run stores the value tuned for its last checkpoint, not for the average. To retune, validate the averaged checkpoint at a few --blank_bias values and pass the best as --blank_bias here, which overrides the stored value. If the inputs were tuned to different values, --blank_bias is required.