# Masked Training of Neural Networks with Partial Gradients

arXiv: https://arxiv.org/abs/2106.08895 · HTML: https://mlsanity.com/paper/2106.08895/

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 52 references that are in the map.
- First posted on arXiv 16 Jun 2021.

## Builds on (most cited first, 30 of 35)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://mlsanity.com/paper/1803.03635/) (2018)
- [Learning both Weights and Connections for Efficient Neural Networks](https://mlsanity.com/paper/1506.02626/) (2015)
- [fairseq: A Fast, Extensible Toolkit for Sequence Modeling](https://mlsanity.com/paper/1904.01038/) (2019)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [On First-Order Meta-Learning Algorithms](https://mlsanity.com/paper/1803.02999/) (2018)
- [Once-for-All: Train One Network and Specialize it for Efficient Deployment](https://mlsanity.com/paper/1908.09791/) (2019)
- [Wide & Deep Learning for Recommender Systems](https://mlsanity.com/paper/1606.07792/) (2016)
- [Ensemble Distillation for Robust Model Fusion in Federated Learning](https://mlsanity.com/paper/2006.07242/) (2020)
- [DeepFM: A Factorization-Machine based Neural Network for CTR Prediction](https://mlsanity.com/paper/1703.04247/) (2017)
- [Linear Mode Connectivity and the Lottery Ticket Hypothesis](https://mlsanity.com/paper/1912.05671/) (2019)
- [To prune, or not to prune: exploring the efficacy of pruning for model compression](https://mlsanity.com/paper/1710.01878/) (2017)
- [Sparsified SGD with Memory](https://mlsanity.com/paper/1809.07599/) (2018)
- [Slimmable Neural Networks](https://mlsanity.com/paper/1812.08928/) (2018)
- [Sparse Communication for Distributed Gradient Descent](https://mlsanity.com/paper/1704.05021/) (2017)
- [The Convergence of Sparsified Gradient Methods](https://mlsanity.com/paper/1809.10505/) (2018)
- [Universally Slimmable Networks and Improved Training Techniques](https://mlsanity.com/paper/1903.05134/) (2019)
- [BigNAS: Scaling Up Neural Architecture Search with Big Single-Stage Models](https://mlsanity.com/paper/2003.11142/) (2020)
- [Characterising Bias in Compressed Models](https://mlsanity.com/paper/2010.03058/) (2020)
- [Dynamic Model Pruning with Feedback](https://mlsanity.com/paper/2006.07253/) (2020)
- [Evaluating the Search Phase of Neural Architecture Search](https://mlsanity.com/paper/1902.08142/) (2019)
- [On the Convergence of SGD with Biased Gradients](https://mlsanity.com/paper/2008.00051/) (2020)
- [Unified Optimal Analysis of the (Stochastic) Gradient Method](https://mlsanity.com/paper/1907.04232/) (2019)
- [Pufferfish: Communication-efficient Models At No Extra Cost](https://mlsanity.com/paper/2103.03936/) (2021)
- [A Tight Convergence Analysis for Stochastic Gradient Descent with Delayed Updates](https://mlsanity.com/paper/1806.10188/) (2018)
- [meProp: Sparsified Back Propagation for Accelerated Deep Learning with Reduced Overfitting](https://mlsanity.com/paper/1706.06197/) (2017)
- [Coordinating Filters for Faster Deep Neural Networks](https://mlsanity.com/paper/1703.09746/) (2017)
- [Dynamic Sparsity Neural Networks for Automatic Speech Recognition](https://mlsanity.com/paper/2005.10627/) (2020)

## Built on by (most cited first, 16 of 16)

- [Gradient Routing: Masking Gradients to Localize Computation in Neural Networks](https://mlsanity.com/paper/2410.04332/) (2024)
- [Sparse Random Networks for Communication-Efficient Federated Learning](https://mlsanity.com/paper/2209.15328/) (2022)
- [CoTFormer: A Chain-of-Thought Driven Architecture with Budget-Adaptive Computation Cost at Inference](https://mlsanity.com/paper/2310.10845/) (2023)
- [Optimization with Access to Auxiliary Information](https://mlsanity.com/paper/2206.00395/) (2022)
- [SpaFL: Communication-Efficient Federated Learning with Sparse Models and Low computational Overhead](https://mlsanity.com/paper/2406.00431/) (2024)
- [Adaptive Compression in Federated Learning via Side Information](https://mlsanity.com/paper/2306.12625/) (2023)
- [Towards a Better Theoretical Understanding of Independent Subnetwork Training](https://mlsanity.com/paper/2306.16484/) (2023)
- [Aequa: Fair Model Rewards in Collaborative Learning via Slimmable Networks](https://mlsanity.com/paper/2502.04850/) (2025)
- [EncCluster: Scalable Functional Encryption in Federated Learning through Weight Clustering and Probabilistic Filters](https://mlsanity.com/paper/2406.09152/) (2024)
- [Straggler-resilient Federated Learning: Tackling Computation Heterogeneity with Layer-wise Partial Model Training in Mobile Edge Network](https://mlsanity.com/paper/2311.10002/) (2023)
- [Federated Fine-Tuning of Foundation Models via Probabilistic Masking](https://mlsanity.com/paper/2311.17299/) (2023)
- [Dropout Neural Network Training Viewed from a Percolation Perspective](https://mlsanity.com/paper/2512.13853/) (2025)
- [Federated Learning over Hierarchical Wireless Networks: Training Latency Minimization via Submodel Partitioning](https://mlsanity.com/paper/2310.17890/) (2023)
- [On the Convergence and Stability of Distributed Sub-model Training](https://mlsanity.com/paper/2511.06132/) (2025)
- [Delayed Random Partial Gradient Averaging for Federated Learning](https://mlsanity.com/paper/2412.19987/) (2024)
- [UAV-Assisted Cooperative Edge Inference for Low-Altitude Economy via MoE-based Hierarchical Deep Reinforcement Learning](https://mlsanity.com/paper/2605.19290/) (2026)

## Abstract (arXiv)

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.