Masked Training of Neural Networks with Partial Gradients

TL;DR

This paper introduces a unified theoretical framework for various stochastic gradient descent variants, including methods for communication efficiency and model compression, and demonstrates its effectiveness through joint training of networks like Slimmable Networks.

Contribution

It provides a comprehensive theoretical analysis of diverse SGD variants and applies this framework to improve joint training of neural networks such as Slimmable Networks.

Findings

Unified framework encompasses multiple SGD variants

Joint training of low-rank and standard Transformers improves performance

Framework guides efficiency improvements and generalization

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.