Towards architecture-aware optimisation
TL;DR We’ve derived an optimiser called automatic gradient descent (AGD) that can train ImageNet without hyperparameters. This removes the need for expensive and time-consuming learning rate tuning, selection of learning rate decay schedulers, etc. Our paper can be found , and for a clean Pytorch implementation, or my
We use this normalisation because it has nice properties that PyTorch default parameterisation does not, including stability with width, resistance to blow-ups in the activations, and promotion of feature learning. This is similar to
Each layer is updated using eta multiplied by the layer’s weight norm, multiplied by normalised gradients, and divided by depth. The division by depth is responsible for scaling with depth. It is interesting that gradient normalisation drops out of the analysis, as other optimisers like Adam incorporate similar ideas heuristically.
Experiments
The goal for these experiments was to test AGD’s ability to (1) converge across a wide range of architectures and datasets, and (2) achieve comparable test accuracy to tuned Adam and SGD.
Figure 4 shows the learning curves of four architectures, from a fully-connected network (FCN) to ResNet-50, on datasets from CIFAR-10 to ImageNet. We compare AGD, shown with solid lines, to a standard optimiser, shown with dotted lines (SGD for ImageNet and tuned Adam for the other three). The top row shows the train objective (loss) and automatic learning rate η. The bottom row shows the train and test accuracies. Figure 5 compares AGD vs tuned Adam vs tuned SGD on an 8-layer FCN. We see very similar performance from all three algorithms, reaching near identical test accuracy.
Figure 6 shows that AGD trains FCNs over a wide range of depths (2 to 32) and widths (64 to 2048). Figure 7 shows the dependence of AGD on batch size (from 32 to 4096), on a 4-layer FCN. It seems to converge to a good optimum no matter the batch size!
Conclusion
To summarise, here is an “architecture-aware” optimiser: automatic gradient descent (AGD), capable of training small systems like an FCN on CIFAR-10 to large-scale systems like ResNet-50 on ImageNet, at a range of batch sizes, without the need for manual hyperparameter tuning.
While using AGD has not removed all hyperparameters from machine learning, those that remain — batch size and architecture — typically fall under into the “make them as large as possible to fill up time/compute budget”.
However, there is lots still to be done. We do not explicitly take into account stochasticity introduced into the gradient due to batch size. We also haven’t looked into regularisation like weight decay. While we’ve done a little bit of work in adding support for affine parameters (in batch norm layers) and bias terms, we haven’t tested it extensively, nor is it as well justified by theory as the rest of the results here.
Perhaps most importantly, we still need to do the analysis required for transformers, and test AGD on NLP tasks. Progress is being made on this front as well!
Finally, check out Jeremy’s for a developmental version with support for bias terms and affine parameters, if you want to try AGD! We hope you will find it useful.
Appendix A
We will go through a sketch of the important steps of the proof here. This is designed for anyone who wants to see how the main ideas come together, without going through the full proof, found in our paper was originally published in Towards Data Science on Medium, where people are continuing the conversation by highlighting and responding to this story.
SOCIAL SHARE CARD GENERATOR