Masked Training of Neural Networks with Partial Gradients

Amirkeivan Mohtashami, Martin Jaggi, Sebastian Stich
Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, PMLR 151:5876-5890, 2022.

Abstract

State-of-the-art training algorithms for deep learning models are based on stochastic gradient descent (SGD). Recently, many variations have been explored: perturbing parameters for better accuracy (such as in Extragradient), limiting SGD updates to a subset of parameters for increased efficiency (such as meProp) or a combination of both (such as Dropout). However, the convergence of these methods is often not studied in theory. We propose a unified theoretical framework to study such SGD variants—encompassing the aforementioned algorithms and additionally a broad variety of methods used for communication efficient training or model compression. Our insights can be used as a guide to improve the efficiency of such methods and facilitate generalization to new applications. As an example, we tackle the task of jointly training networks, a version of which (limited to sub-networks) is used to create Slimmable Networks. By training a low-rank Transformer jointly with a standard one we obtain superior performance than when it is trained separately.

Cite this Paper

BibTeX
@InProceedings{pmlr-v151-mohtashami22a, title = { Masked Training of Neural Networks with Partial Gradients }, author = {Mohtashami, Amirkeivan and Jaggi, Martin and Stich, Sebastian}, booktitle = {Proceedings of The 25th International Conference on Artificial Intelligence and Statistics}, pages = {5876--5890}, year = {2022}, editor = {Camps-Valls, Gustau and Ruiz, Francisco J. R. and Valera, Isabel}, volume = {151}, series = {Proceedings of Machine Learning Research}, month = {28--30 Mar}, publisher = {PMLR}, pdf = {https://proceedings.mlr.press/v151/mohtashami22a/mohtashami22a.pdf}, url = {https://proceedings.mlr.press/v151/mohtashami22a.html}, abstract = { State-of-the-art training algorithms for deep learning models are based on stochastic gradient descent (SGD). Recently, many variations have been explored: perturbing parameters for better accuracy (such as in Extragradient), limiting SGD updates to a subset of parameters for increased efficiency (such as meProp) or a combination of both (such as Dropout). However, the convergence of these methods is often not studied in theory. We propose a unified theoretical framework to study such SGD variants—encompassing the aforementioned algorithms and additionally a broad variety of methods used for communication efficient training or model compression. Our insights can be used as a guide to improve the efficiency of such methods and facilitate generalization to new applications. As an example, we tackle the task of jointly training networks, a version of which (limited to sub-networks) is used to create Slimmable Networks. By training a low-rank Transformer jointly with a standard one we obtain superior performance than when it is trained separately. } }
Endnote
%0 Conference Paper %T Masked Training of Neural Networks with Partial Gradients %A Amirkeivan Mohtashami %A Martin Jaggi %A Sebastian Stich %B Proceedings of The 25th International Conference on Artificial Intelligence and Statistics %C Proceedings of Machine Learning Research %D 2022 %E Gustau Camps-Valls %E Francisco J. R. Ruiz %E Isabel Valera %F pmlr-v151-mohtashami22a %I PMLR %P 5876--5890 %U https://proceedings.mlr.press/v151/mohtashami22a.html %V 151 %X State-of-the-art training algorithms for deep learning models are based on stochastic gradient descent (SGD). Recently, many variations have been explored: perturbing parameters for better accuracy (such as in Extragradient), limiting SGD updates to a subset of parameters for increased efficiency (such as meProp) or a combination of both (such as Dropout). However, the convergence of these methods is often not studied in theory. We propose a unified theoretical framework to study such SGD variants—encompassing the aforementioned algorithms and additionally a broad variety of methods used for communication efficient training or model compression. Our insights can be used as a guide to improve the efficiency of such methods and facilitate generalization to new applications. As an example, we tackle the task of jointly training networks, a version of which (limited to sub-networks) is used to create Slimmable Networks. By training a low-rank Transformer jointly with a standard one we obtain superior performance than when it is trained separately.
APA
Mohtashami, A., Jaggi, M. & Stich, S.. (2022). Masked Training of Neural Networks with Partial Gradients . Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, in Proceedings of Machine Learning Research 151:5876-5890 Available from https://proceedings.mlr.press/v151/mohtashami22a.html.

Related Material