# Lowering PyTorch's Memory Consumption for Selective Differentiation

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 46 references that are in the map.
- First posted on arXiv 15 Apr 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks](https://mlsanity.com/paper/1506.01497/) (2015)
- [Visual Instruction Tuning](https://mlsanity.com/paper/2304.08485/) (2023)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://mlsanity.com/paper/1905.11946/) (2019)
- [Fully Convolutional Networks for Semantic Segmentation](https://mlsanity.com/paper/1411.4038/) (2014)
- [MobileNetV2: Inverted Residuals and Linear Bottlenecks](https://mlsanity.com/paper/1801.04381/) (2018)
- [Aggregated Residual Transformations for Deep Neural Networks](https://mlsanity.com/paper/1611.05431/) (2016)
- [Rethinking Atrous Convolution for Semantic Image Segmentation](https://mlsanity.com/paper/1706.05587/) (2017)
- [Searching for MobileNetV3](https://mlsanity.com/paper/1905.02244/) (2019)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [EfficientNetV2: Smaller Models and Faster Training](https://mlsanity.com/paper/2104.00298/) (2021)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [PaLI: A Jointly-Scaled Multilingual Language-Image Model](https://mlsanity.com/paper/2209.06794/) (2022)
- [Quantized Neural Networks: Training Neural Networks with Low Precision Weights and Activations](https://mlsanity.com/paper/1609.07061/) (2016)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Optimal Brain Compression: A Framework for Accurate Post-Training Quantization and Pruning](https://mlsanity.com/paper/2208.11580/) (2022)
- [Surgical Fine-Tuning Improves Adaptation to Distribution Shifts](https://mlsanity.com/paper/2210.11466/) (2022)
- [The Computational Limits of Deep Learning](https://mlsanity.com/paper/2007.05558/) (2020)
- [Training Quantized Nets: A Deeper Understanding](https://mlsanity.com/paper/1706.02379/) (2017)
- [Tuning LayerNorm in Attention: Towards Efficient Multi-Modal LLM Finetuning](https://mlsanity.com/paper/2312.11420/) (2023)
- [ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training](https://mlsanity.com/paper/2104.14129/) (2021)
- [Randomized Automatic Differentiation](https://mlsanity.com/paper/2007.10412/) (2020)

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

- [ReNO: Enhancing One-step Text-to-Image Models through Reward-based Noise Optimization](https://mlsanity.com/paper/2406.04312/) (2024)
- [Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Models](https://mlsanity.com/paper/2508.09968/) (2025)

## Abstract (arXiv)

Memory is a limiting resource for many deep learning tasks. Beside the neural network weights, one main memory consumer is the computation graph built up by automatic differentiation (AD) for backpropagation. We observe that PyTorch's current AD implementation neglects information about parameter differentiability when storing the computation graph. This information is useful though to reduce memory whenever gradients are requested for a parameter subset, as is the case in many modern fine-tuning tasks. Specifically, inputs to layers that act linearly in their parameters (dense, convolution, or normalization layers) can be discarded whenever the parameters are marked as non-differentiable. We provide a drop-in, differentiability-agnostic implementation of such layers and demonstrate its ability to reduce memory without affecting run time.