Masked Training of Neural Networks with Partial Gradients

Mohtashami, Amirkeivan and Jaggi, Martin and Stich, Sebastian U.
(2022) Masked Training of Neural Networks with Partial Gradients.
In: AISTATS 2022, 28 Mar - 30 Mar 2022, online.
Conference: AISTATS International Conference on Artificial Intelligence and Statistics
(In Press)

2106.08895.pdf - Accepted Version

Download (4MB) | Preview
Official URL:


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.


Actions (login required)

View Item View Item