Weight Prediction Boosts the Convergence of AdamW
This is an incremental improvement for deep learning practitioners seeking faster training convergence with AdamW.
The paper tackles the problem of slow convergence in AdamW optimizer for deep neural network training by introducing weight prediction, which uses predicted future weights for forward and backward passes, resulting in improved convergence and better accuracy on image classification and language modeling tasks.
In this paper, we introduce weight prediction into the AdamW optimizer to boost its convergence when training the deep neural network (DNN) models. In particular, ahead of each mini-batch training, we predict the future weights according to the update rule of AdamW and then apply the predicted future weights to do both forward pass and backward propagation. In this way, the AdamW optimizer always utilizes the gradients w.r.t. the future weights instead of current weights to update the DNN parameters, making the AdamW optimizer achieve better convergence. Our proposal is simple and straightforward to implement but effective in boosting the convergence of DNN training. We performed extensive experimental evaluations on image classification and language modeling tasks to verify the effectiveness of our proposal. The experimental results validate that our proposal can boost the convergence of AdamW and achieve better accuracy than AdamW when training the DNN models.