DropIT: Dropping Intermediate Tensors for Memory-Efficient DNN Training
Authors: Joya Chen
Kai Xu
Yuhui Wang
Yifei Cheng
Angela Yao
Venue: ICLR 2023 poster (2023)
Keywords: dropping intermediate tensors
dropping activations
activation compressed training
top-k
vision transformer
cnn

Abstract: A standard hardware bottleneck when training deep neural networks is GPU memory. The bulk of memory is occupied by caching intermediate tensors for gradient computation in the backward pass. We propose a novel method to reduce this footprint - Dropping Intermediate Tensors (DropIT).  DropIT drops min-k elements of the intermediate tensors and approximates gradients from the sparsified tensors in the backward pass. Theoretically, DropIT reduces noise on estimated gradients and therefore has a higher rate of convergence than vanilla-SGD. Experiments show that we can drop up to 90\% of the intermediate tensor elements in fully-connected and convolutional layers while achieving higher testing accuracy for Visual Transformers and Convolutional Neural Networks on various tasks (e.g., classification, object detection, instance segmentation). Our code and models are available at https://github.com/chenjoya/dropit.

Published as a conference paper at ICLR 2023
DROPIT: DROPPING INTERMEDIATE TENSORS FOR DNN TRAINING
Joya Chen'r, Kai Xu Yuhui Wang', Yifei Cheng", Angela Yao
1 National University of Singapore
2 University of Science and Technology of China
Joyachengu. nuw.nneuuuuuuuuu..e.e..
[lexn, yuhuin, ayao) ( 2oomp. nus. mus.eddu.ee chengyi f8mmil www.u.ucuuc........ edu.cn
ABSTRACT
A standard hardware bottleneck when training deep neural networks GPU mem- ory. The bulk of memory is occupied by caching intermediate tensors for gradient computation in the backward pass. We propose novel method to reduce this foot- print Dropping Intermediate Tensors (DropIT). DropIT drops min-k elements of the intermediate tensors and approximates gradients from the sparsified tensors the backward pass. Theoretically, DropITT reduces noise on estimated gradients and therefore has a higher rate of convergence than vanilla-SGD. Experiments show that we can drop up to 90% of the intermediate tensor elements in fully- connected and convolutional layers while achieving higher testing accuracy for Visual Transformers and Convolutional Neural Networks on various tasks (e.g.. classification, object detection, instance segmentution). Our code and models are available https: / thub. com coynjoya///oooooooooooooooooaoar /tropoo/////rriii.tt.. dropit.
INTRODUCTION
The training of state-of-the-arr deep neural networks (DNNs) (Krizhevsky et al., 2017; Simonyan & Zisserman, 2015; He et al., 2016; Vaswani et al., 2017; Dosovitskiy et al., 2021) for computer vision often requires large GPU memory. For example, training simple visual transformer detection model ViTDet-B (Li et al., 2022). with its required input image size of 1024 1024 and batch size of 64, requires 700 GB GPU memory. Such high memory requirement makes the training of DNNs out of reach for the average academic or practitioner without access to high-end GPU resources.
When training DNNs, the GPU memory has six primary uses (Rajbhandari et al., 2020): network parameters, parameter gradients, optimizer states (Kingma & Ba, 2015), intermediate tensors (also called activations), temporary buffers, and memory fragmentation. Vision tasks often require train- ing with large batches of high-resolution images or videos, which can lead to a significant memory cost for intermediate tensors. In the instance of ViTDet-B, approximately 70% GPU memory cost (~-400 GB) is assigned to the intermediate tensor cache. Similarly, for NLP, approximately 50% of GPU memory consumed by caching intermediate tensors for training the language model GPT- 2 (Radford et al., 2019; Rajbhandari et al., 2020). As such, previous studies (Gruslys et al., 2016 6; Chen et al., 2016; Rajbhandari et al., 2020; Feng & Huang, 202 1) treat the intermediate tensor cache as the largest consumer of GPU memory.
For differentiable layers, standard implementations store the intermediate tensors for computing the gradients during back- pock-propagation. One option to reduce storage is to cache tensors from only some layers. Uncached tensors are recomputed on the fly during the backward pass this is the strategy of gradient checkpointing (Gruslys al., 2016; Chen et al., 2016; Bulo et al., 201 8; Feng & Huang, 2021). Another option is to quantize the tensors after the forward computation and use the quantized values for gradient computation during the backward pass (Jain et al., 2018: 8; Chakrabarti & Moseley, 2019; Fu et al., 2020; Evans & Aamodt, 2021; Liu et al., 2022) this is known as activation com- pression training (ACT). Quantization can reduce memory considerably, but also brings inevitable performance drops. Accuracy drops can be mitigated by bounding the error at each layer through adaptive quantization (Evans & Aamodt, 2021; Liu al., 2022), i.e. adaptive ACT. However, train- ing time consequently suffers extensive tensor profiling is necessary during training.
"Equal contribution.

---

Published as a conference paper at ICLR 2023
In this paper, propose to reduce the memory usage of intermediate tensors by simply dropping elements from the tensor. We call our method Dropping Intermediate Tensors (DropIT). In the most basic setting, dropping indices can be selected randomly, though dropping based on min-k ranking on the element magnitude more effective. Both strategies are much simpler than the sensitivity checking and other profiling strategies, making DropITT much faster than adaptive ACT.
During training, the intermediate tensor is transformed over to a sparse format after the forward computation complete. The sparse tensor then recovered to general tensor during backward gradient computation with dropped indices filled with zero. Curiously, with the right dropping strategy and ratio, DropIT has improved convergence properties compared SGD. We attribute this to the fact that DropIT can, theoretically, reduce noise on the gradients. In general, reducing noise will result in more precise and stable gradients. Experimentally, this strategy exhibits consistent performance improvements on various network architectures and different tasks.
To the best of our knowledge, we are the first to propose activation sparsification. The closest related line of existing work is ACT, but unlike ACT, DropIT leaves key elements untouched, which is crucial for ensuring accuracy. Nevertheless, DropIT orthogonal to activation quantization, and the two can be combined for additional memory reduction with higher final accuracy. The key contributions of our work are summarized as follows:
We propose DropIT, a novel strategy to reduce the activation memory by dropping the elements of the intermediate tensor.
We theoretically and experimentally show that DropIT can be seen as a noise reduction on stochastic gradients, which leads to better convergence.
DropIT can work for various settings: training from scratch, fine-tuning on classification, object detection, etc. Our experiments demonstrate that DropIT can drop up to 90% of the intermediate tensor elements in fully-connected and convolutional layers with a testing accuracy higher than the baseline for CNNs and ViTs. We also show that DropIT is much better regarding accuracy and speed compared to SOTA activation quantization methods, and it can be combined with them to pursue higher memory efficiency.
2 RELATED WORK
Memory-efficient training. Current DNNs usually incur considerable memory costs due to huge model parameters (e.g. GPT's (Radford et al., 2019; Brown et al., 2020)) or intermediate tensors (e.g.. high-resolution feature map (Sun et al., 2019; Gu et al., 2022)). The model parameters and corresponding optimizer states can be reduced with lightweight operations (Howard et al., 2017: Xie et al., 2017; Zhang et al., 2022), distributed optimization scheduling (Rajbhandari et al., 2020), and mixed precision training (Micikevicius et al., 2018). Nevertheless, intermediate tensors, which essential for gradient computation during the backward pass, consume the majority of GPU memory (Gruslys et al., 2016; Chen et al., 2016; Rajbhandari et al., 2020; Feng & Huang, 2021), and reducing their size can be challenging.
Gradient checkpointing. To reduce the tensor cache, gradient checkpointing (Chen et al., 2016; Gruslys al., 2016; Feng & Huang, 2021) stores tensors from only a few layers and recomputes any uncached tensors when performing the backward pass; in the worst-case scenario, this is equivalent duplicating the forward pass, So any memory savings come as an extra computational expense. InPlace- ABN (Bulo et al., 2018) halves the tensor cache by merging batch normalization and ac- tivation into a single in-place operation. The tensor cache is compressed in the forward pass and recovered in the backward pass. Our method is distinct in that it does not require additional recom- putation; instead, the cached tensors are sparsified heuristically.
Activation compression. (Jain et al., 2018; Chakrabarti & Moseley, 2019; Fu et al., 2020; Evans & Aamodt, 2021; Chen et al., 2021; Liu et al., 2022) explored lossy compression on the activation cache via low-precision quantization. (Wang et al., 2022) compressed high-frequency components while (Evans et al., 2020) adopted JPEG-style compression. In contrast to all of these methods, DropIT reduces activation storage via sparsification, which has been previously unexplored. In addition, DropIT is more lightweight than adaptive fow-precision quantization methods (Evans & Aamodt, 2021; 1; Liu et al., 2022).
2

---

Published as a conference paper at ICLR 2023
Layer in forward pass
Layer pass
Parameter 0
Layer
Parameeer
Forward
Foruard fonction
Input
Quiput
Taput Disply
Ouiput
Layeri +
Layer i+ +
Gradient
Parameter d
Layer ++1
Gradient V00
Layer ++1
Rockward function
Rockward function
Gradieut Vx
Gradient
Podded R(A)
Gradient
Gradient Va
Cochedl a
Layer
Tayer 1-
Layer In backward
Layer in backwarr paes
(a) Baseline
(b) DropIT
For regular baseline network (a), the intermediate tensor is fully cached in the forward pass be used for gradient computation during the backward pass. For DropIT, elements of the intermediate tensors are dropped during caching; only the retained elements with zero padding are used for gradient computation during the backward pass. DropIT can save GPU memory for two reasons. First, cached tensors are accumulated layer by layer during the forward pass, and DropIT sparsifies them, thereby reducing maximum memory allocation. Second, backward tensors are re- leased after making the memory cost associated with padding negligible. Best viewed in color.
Figure 1:
Gradient approximation. Approximating gradients has been explored in large-scale distributed training to limit communication bandwidth for gradient exchange. (Strom, 2015; Dryden et al., 2016; Aji Heafield, 2017; Lin et al., 2018) propose dropping gradients based on some fixed thresholds and sending only the most significant entries of the stochastic gradients with the guar- anteed convergence (Stich et al., 2018; Cheng et al., 2022: Chen et al., 2020). Instead dropping gradient components, DropIT directly drops elements within intermediate tensors objective to reduce the training memory.
3 METHODOLOGY
3.1 PRELIMINARIES
We denote the forward function and learnable parameters of the i-th layer as / and 0, respectively. In the forward pass, operates on the layer's input to compute the output z:
2 I(a, 0).
(1)
For example, if layer is convolution layer, would indicate convolution operation with @ repre- senting the kernel weights and bias parameter.
Given loss function F(O), where O represents the parameters of the entire network, the gradient, with respect to 0 at layer i, can be estimated according to the chain rule as
OF(O)
= V 0z = V2
al(a, 0)
VO
(2)
00
00
00
0F(@) is the gradient passed back from layer i + 1. Note that the computation of 82
where 2 (l(a.0) 00 requires a if the forward function | involves tensor multiplication between a and 9. This is the case for common learnable layers, such convolutions in CNNs and fully-connected layers in transformers. As such, a necessary for estimating the gradient and cached after it is computed
Note that the output from the previous layer i- 1, i.e. a' = ااي However, we assign different symbols to denote the input and output of given layer explicitly; this redundant notation conveniently allows for clarity purposes, drop the explicit reference of the layer index superseript.
3

---

Published as a conference paper at ICLR 2023
Layer Type
Parameter 0
Tensor a
Convolution Fully Connected
OlC.C.k²) O(C.C.)
O(BC,Lo) O(BLoCa)
Table I: Space complexity for parameters and intermediate tensors in single layer. B: batch size, L.: input sequence length (e.g.. width x height) C.. C₂: the number of input, output channels, K; convolutional kernel size. Typically, Cas Cz, K would be fixed once the model has been built, the complexity for intermediate tensors would be considerable with large B, La.
during the forward pass, illustrated in Figure ((a). A common way reduce storage for a is to store quantized version (Jain et al., 2018; Chakrabarti & Moseley, 2019, 19; Fu et al., 2020; Evans & Aamodt, 2021; Liu et al., 2022). Subsequent gradients in the backward pass are then computed using the quantized a. The gradient Va can be estimated similarly via chain rule as
aF(e) da
dz Da
a(a,
Va
= V2
V
(3)
da
Analogous Eq. 2. the partial (((a.)) model memory. However, the stored always shares memory with the model residing in the GPU, so it does not incur additional memory consumption. Furthermore, A typically occupies much less memory. In Table I, the intermediate tensor's space complexity becomes significant when B or La large, which is common in CV and NLP tasks.
da may depend on the parameter and 0 is similarly stored in the
3.2 DROPPING INTERMEDIATE TENSORS
Let X' denote the of all indices for intermediate tensor Suppose that X' is partitioned into two disjoint sets Na and X,, i.e. X, n Xa 0 and X, U Xa = X. In DropIT, we introduce dropping operation D(-) sparsify into where à consists of the elements ax, and the indices Nr, i.e. à = (aa) faxx. The sparse à can be used as a substitute for in Eq. 2. While sparsification can theoretically reduce both storage and computation time, we benefit only from storage savings in practice. We retain general matrix multiplication because the sparsity rate is insufficient for sparse matrix multiplication
provide meaningful computational gains. As such, the full intermediate tensors are recovered for gradient computation, i.e. VO s Vz - ((.R()))0) where 89
R(-) represents the process that inflates back to a general matrix with dropped indices filled with zero. The overall procedure is demonstrated in Figure 1(b).
Consider for convolutional layer with C2 kernels of size K K. For the jth kernel, where j [1, Cz]. the gradient at location (u, v) for the 4thh channel is given by convolving incoming gradient
z and input a:
(0,x(u,)==
Σ (2)( (N.x)))
(4)
where x'=x + u and yy=y y + v. The set X' in this case would denote the set of all sample indices [1, B] and all location indices (r,, y) E [1, W] [1, H) in the feature map. Without any loss in generality, we can partition X' into two disjoint sets X, and X' to split Eq. as
V0nu(u, = Σ (n.
Σ (n,x))).,
( (5)
T6,,fau,))
Assume now, that some element (x', 3/) is small or near-zero; in CNNs and Transformers, such an assumption is reasonable due to preceding batch/layer normalization and ReLU or GeLU activations (see Figure 3). Accordingly, this element's contribution to the gradient will also be correspondingly small. If assign the spatial indices (x. y) in sample of all small near-zero elements Xd, then we can approximate the gradient v0,u(au (u, v) with simply the second term of Eq. We denote the approximated gradient Gdropit T (u, v).
4

---

Published as a conference paper at ICLR 2023
X
Weight 0
Inpul a
Weight 0
Output 2
Gradient Va
Gradient V2
Dropping Function D-)
(egg, Random, 0.9)
Gradient T6 (approximated)
R(ay"
Gradient V2
Cache
(a) Forward
(h) Backward
Forward and backward of DropIT on the fully-connected layer (without bias). In the forward pass, we sparsify the cache tensor and drop 7 percentage storage. In the backward pass, only saved elements participate in the gradient computation.
Figure 2:
blocks. gkv
blocks. 11.tttn.pro
blocks. mip.fc1
blocks. 11.mlp.fc222
y 70%
y 70%
y 70%
y 70%
Count (105)
2
Value
Value
Value
Value
Figure 3: Distribution of element values in intermediate tensors' on DeiT-Ti. Dropped elements are shaded in grey. DropIT with min-k only discards elements that are close to zero. Here we only show final block while observing that the distributions of other blocks are similar.
For a fully connected layer, the approximated gradient can be defined similarly as
Idropát -0 = Σ V2 ax
(6)
nEX,
A visualization of the gradient approximation is shown in Figure 2. With the approximated gradient, we can use any standard deep learning optimization scheme to update the parameters.
3.3 DROPPING FUNCTION D(.)
|X: We define the overall dropping rate as 7 BC. for fully connected layer and y BC, HW for a convolutional layer. 7 can be varied and will be used later to define the dropping function D(-). As aim to drop elements with minimal contribution to the gradient, it is logical to perform min-k based selection on the elements' magnitudes before dropping the elements. As a baseline comparison, also select X₂ based uniform random sampling. We investigate the following options for D(-):
|X~1
Random Elements: fraction of elements dropped randomly within mini-batch.
Min-K Elements: Within a mini-batch, we drop the smallest y fraction elements according to their absolute magnitudes.
5

---

Published as a conference paper at ICLR 2023
3.4 THEORETICAL ANALYSIS
Below, analyze convergence for dropping min-k elements. The gradient of Stochastic Gradient Descent (SGD) is commonly viewed as Gradient Descent (GD) with noise:
Igd + r((0
Axyd
where n represents some zero-mean noise distribution with a variance ? introduced by variation in the input data batches.
With min-k dropping, the gradient becomes biased; we assume it can be modeled as:
9min-k 0g9gd (n.
(8)
That is, min-k dropping results in bias factor a while affecting noise by a factor of B. a and / vary each iteration, i.c., {a1 Q2 ar} and {31, 52. .... }. Additionally in Appendix A.2, we provide nonlinear approximation to Imin-k that achieves same convergence.
By scaling the learning rate with factor of a the gradient after min-k dropping as given in Eq. can also be expressed as:
Jmin-k Agd + (n,
We can formally show (see Appendix A.3) that Eja] 2 E 2 1 - and therefore E| < 1. This suggests that min-k dropping reduces the noise the gradient. With less noise, better theoretical convergence is expected.
Similar to convergence proofs in most optimizers, we will assume that the loss function is L-- smooth. Under the L-smooth assumption. for SGD with learning rate and min-k dropping with a learning rate N. we can reach the following convergence after T iterations:
VF(2))1 2(F(x,) - 2(F(())-F(x)) + +nLE²
SGD:
(10)
Tn
T
VF(x)) 2(F(x)) -F(")) + nLL
1 T
T
DropIT with min-k:
E
(11)
Tn
2
where indicates an optimal solution. Full proof can found in Appendix A.1. Note that the two inequalities differ only by the second term in the right-hand side. ar represents the bias caused by dropping at the t-th iteration and 3, measures the noise reduction effect after dropping. We further investigate and in the supplementary and show that under certain conditions, E[a] 2 E[β], thereby reducing the noise and improving the convergence of DropIT from standard SGD.
3.5 DROPIT FOR NETWORKS
For some layers, normalization and activations, Ha may also depend on a. In these cases, do not drop the cache of intermediate tensors as this will affect subsequent back-propagation. For DropIT,, dropping happens only when the gradient flows to the parameters, which prevents the aggregation of emrrrs from approximating the gradient.
Al(a.6)
Now, we have discussed dropping tensor elements from the cache of a single layer. DropIT is the- oretically applicable for all convolutional and fully-connected layers in network since does not affect the forward pass. For Visual Transformers (Dosovitskiy et al., 2021), we apply DropIT for most learnable layers, though we ignore the normalization and activations like LayerNorm (Ba et al., 2016) and GELU (Hendrycks & Gimpel, 2016)). The applicable layers include fully-connected lay- ers in multi-head attention and MLPs in each block, the beginning convolutional layer (for patches projection), and the final fully-connected classification layer. For CNNs the applicable layers in- clude all convolutional layers and the final fully tully-connected classification layer. We leave networks unchanged during inference.

---

Published as a conference paper at ICLR 2023
Strategy
Dropping Rate
y
(%( (ascline)
10%
20%
30%
40%
50%
60%
70%
80%
90%
Random
72.1*
72.4
72.4
72.0
71.7
70.8
69.6
68.1
65.8
60.8
Min-K
72.1*
72.1
72.1
72.2
72.4
72.5
72.4
72.1
70.8
66.4
From Touvron et al. (2021)'s official implementation, obtain 72.13 with public weights and our training.
Table 2: Ablation study on dropping strategy and dropping rate. Reported results are top-1 accuracy on the ImageNet-Ih k validation set, achieved by DeiT-Ti training from scratch on the ImageNet- Ik training set. We highlight that the accuracy is higher than baseline ( (2721)).
70
5.8
2020%
10%
4.25
y-70%
5.
50-
3.71
4.5
4.18
3.68
85
50
100
150
200
250
300
Eprch
Figure 4: Training loss curves of Min-K DropIT. Baseline (7 = 0%) is bolded. y 80%, 90% are hidden their losses are obviously higher than the baseline. y 10%. 30% are also hidden for easier viewing. 7 = 40%~70% achieve lower loss than baseline at the end. Best viewed in color.
EXPERIMENTS
In this section, we present a comprehensive evaluation of DropIT's effectiveness, leveraging experi- ments training from seratch on ImageNet- 1 (Russakovsky et al., 2015). Our results demonstrate that DropITT outperforms existing methods by achieving lower training loss, higher testing accuracy, and reduced GPU memory consumption. We showcase the versatility of DropIT in various fine- tuning scenarios, such as ImageNet- to CIFAR-100 (Krizhevsky et al., 2009), object detection, and instance segmentation on MS-COCO (Lin et al., 2014). Furthermore, we compare DropIT with recent state-of-the- ACT methods (Pan et al., 2022; Liu et al., 2022) and establish its superiority in terms of accuracy, speed, and memory cost.
4.1 EXPERIMENTAL DETAILS
For image classification, we employed DeiT' (Touvron et al., 2021) instead of vanilla ViT (Dosovitskiy al., 2021) since it doesn't require fine-tuning from ImageNet-2Ik. DeiT and ViT share the same architecture, differing only in their training hyper-parameters. Additionally, for transfer learning, we utilized Faster/Mask R-CNN models (Ren al., 2017; He al., 2017) to evaluate our approach in object detection and instance segmentation.
Models.
Implementation Details. We use the official implementations of DeiT (without distillation) and Faster/Mask R-CNN, and keep all pute gradients using DropIT. Our implementation consistent. The only difference is that we com- based on PyTorch 1.12 (Paszke et al., 2019), and we utilize torch. autograa rad package. During the forward pass, we use DropIT to convert the dense tensor to coordinate format, and recover it during the backward pass. The min-k strategy of DropIT is implemented by torch took, which retains elements with the largest absolute value, based on a proportion of 1 1- The corresponding indices of these elements are also maintained. For all experiments, we follow DeiT (Touvron et al., 2021) 1) and set fixed random seed of O. We mea- sure training speed and memory on NVIDIA RTX A5000 GPUs. Additional details can be found in Appendix A.8.
7

---

Published as a conference paper at ICLR 2023
(a) CIFAR-100 fine-tuning results. DeiT' net- works are initialized from their publicly avail- able pre-trained ImageNet-Ik Ik weights.
(b) Detection & instance segmentation on COCO.
R50-FPN denotes ResNet-50 with FPN (Lin et al., 2017), initialized from public ImageNet-IK Ik weights.
Network
DropIT
Top-1 Accuracy
Network
DropIT
Aplex
Apmesk
89.7
Faster R-CNN
37.0
n/a
DeiT-S
90%
90.1
(R50-FPN)
Y = 90%
37.2
n/a
90.8
Mask R- R-CNN (R50-FPN)
37.9
34.5
Deil-B
90%
91.3
7 = 80%
38.5
34.5
Table 3: Fine-tuning with DropIT on image classification, object detection & instance segmentation.
Dropping Rate y
Cached
= 0%
60%
70%
80%
Y = 90%
Tensor
11.26 G
4.50 G ( 60%) 3.38 G ( (0%%) 2.25 G (00%)
1.13 (-90%)
Table 4: Memory cost of DropIT cached tensors (without indices) for different y. DropIT can precisely reduce the memory by y. The measured model is DeïT-S with a batch size of 1024.
4.2 IMPACT ON ACCURACY
Training from scratch on ImageNet-1k, Table 2 shows that training DeiT-Ti from scratch without DropIT (baseline) has a top-1 accuracy of 72.1 on ImageNet-Ik: lk. Random dropping matches or improves the accuracy (72.4) when y 20%, but with higher (y 2 30%), accuracy progressively decreases from the baseline. The phenomenon can be explained by the following: (1) Small amounts of random dropping (7 S 20%) can be regarded as adding random noise the gradient. The noise has regularization effect on the network optimization to improve the accuracy, similar to what was observed in previous studies (Neclakantan et al., 2015; Evans Aamodt, 2021). (2) Too much random dropping (7 30%) results in deviations that can no longer be seen as small gradient noise, hence redueing performance.
With min-k dropping, DropIT can match or exceed the baseline accuracy over wide range of (< 70%). Intuitively, training from scratch should be difficult with DropIT, especially under large dropping rates, as the computed gradients are approximations, However, our experiments demon- strate that DropITT achieves 0.4% and 0.3% higher accuracy in 7 50% and 60%, respectively. In fact, DropIT can match the baseline accuracy even after discarding 70% of the clements.
Fig. 4 compares the loss curves when training from scratch on the baseline DeiT-Ti model without and with DropIT using min-k strategy. The loss curves of DropIT with various y values follow the same trend as the baseline; up to some value of y, the curves are also but are consistently lower than the baseline, with y 50%, 60% achieving the lowest losses and highest accuracies. As such, we conclude that DropIT accurately approximates the gradient while reducing noise, as per our theoretical analysis.
Fine-tuning on C1FAR-100. Table 3(a) shows that DeiT networks can be fine-tuned with DropIT achieve higher than baseline accuracies even while dropping up to 90% intermediate elements. Compared to training from scratch from Table 2, DropIT can work with a more extreme dropping rate (90% vs. 70%). We interpret that this because the network already has a good initialization before fine-tuning, thereby simplifying the optimization and allowing higher ? to be used.
Backbone fine-tuning, head network training from scratch, on COCO. We investigated DropIT in two settings: training from scratch and fine-tuning from pre-trained network. We also studied a backbone network initialized with ImageNet pre-training, while leaving others, such as RPN and R-CNN head, uninitialized, which is common practice in object detection. Table 3(b) shows that DropIT can steadily improve detection accuracy (Apbor), ). When 7 = 80%, we observed an impres- sive 0.6 AP gain in Mask R-CNN, although this gain was not observed in Apmask We believe that the segmentation performance may be highly related to the deconvolutional layers in the mask head, which are not currently supported by DropIT. We plan investigate this further in future work. These experiments demonstrate the effectiveness of DropIT on CNNs, and in Appendix A.4, we demonstrate the effectiveness of DropIT for ResNet training ImageNet-1k.
8

---

Published as a conference paper at ICLR 2023
Benchmark
FC Cache
Others Cache
Acc
MaxM
MaxM (-Index)
Speed (ms)
none
none
89.7
6,66 G
6.66 G
172
DropIT (7 90%)
none
90.1
5.59 G
5.29 G
212
DeiT-S on 100
MESA (8-bit)
MESA (8-bit)
89.7
3.52
3.52 G
416
DropIT(y 90%)
MESA (8-bit)
89.9
3.27 G
2.97 G
375
GACT (4-bit)
GACT (4-bit)
89.7
2.16
2.16 G
290+49"
DropIT(y 90%)
GACT (4-bit)
90.0
2.27 G
1,97 G
286+25'
Note: GACT has time-consuming sensitivity profiling computation every 1000 iterations. costs 19.8I and 25.43 (+DropIT) seconds in our benchmark. we add an average of this time over 1000 iterations).
Table 5: Compare and combine with state-of-the-art ACT methods. FC: fully-connected. MaxM: maximum memory. MaxM (- Index): maximum memory without index (moved to CPU). We follow MESA use batch size 128 to measure memory and speed single GPU.
4.3 IMPACT ON MEMORY SPEED, SOTA COMPARISON
Intermediate Tensor Cache Reduction. Table shows the intermediate tensor cache reduction achieved by DropIT. In DropIT applied layers (FC layers of DeiT-S), the total reserved activation (batch size = 1024) is 11.26 When we use DropIT to discard activations, the memory reduction precisely controlled by 7. i.e. 7 90% means the reduction is 11.26 x 0.9. DropITT does incur some memory cost for indexing. but as we show next, the maximum GPU memory can still be reduced.
Comparison and Combination with SOTA. In Table 5, we compare and combine DropIT with state-of-the-arr activation quantization methods. Mcasuring performance individually, with 7 = 90%, DropI'T improves accuracy by 0.4 and reduces maximum memory by 1.07 G (1.37G activations -0.3G indexing), and slightly increases the time (40 ms) per iteration. The max memory reduction is less than that shown Table 4 because activations from non-applicable layers still occupy considerable memory. Therefore, a natural idea to supplement DropIT is to perform activation quantization for layers without DropIT. We next present the combination results of DropIT with recent methods MESA (Pan et al., 2022) and GACT (Liu et al., 2022).
As shown in Table 5, MESA can reduce memory with 8-bit quantization and has no impact on the baseline accuracy (89.7). However, the time cost of the MESA algorithm is also considerable, and 416/172 R 2.4x slower than baseline and 416/212 improvement in CIFAR- 100 finetuning. MESA achieves 71.9 accuracy of DropIT-Ti on ImageNet- 1k, but DropIT can go up to 72.5 (Table 2, 7 = 50%). We can combine MESA with DropIT by applying DropIT in the conv/fc layers and applying MESA in the other layers. Together, the accuracy, memory, and speed are all improved over MESA alone, conclusively demonstrating the effectiveness of DropIT.
2x more than DropIT, with no accuracy
We compare similarly to GACT; Table shows that at 4 bits, it can reduce max-memory even further. Combining GACT with DropIT marginally increases the max-memory due to DropIT's indexing consumption; however, there are both accuracy and speed gains. Furthermore, GACT reports 0.2~0.4 Aphoz loss on COCO (Liu et al., 2022), though DropIT can produce 0.6 Apooz improvement on COCO (Table 3()).. To sum up, DropIT has its unique advantages in terms of accuracy and speed compared to existing activation quantization methods. Although it saves less memory than the latter, we can combine the two to achieve higher memory efficiency.
CONCLUSION
In this paper, we propose the Dropping Intermediate Tensors (DropI')) method to reduce the GPU memory cost during the training of DNNs. Specifically, DropIT drops elements in intermediate tensors to achieve memory--fficient tensor cache, and it recovers sparsified tensors from the maining elements in the backward pass to compute the gradient. Our experiments show that DropIT can improve the accuracies of DNNs and save GPU memory different backbones and datasets. DropIT provides a new perspective to reduce GPU memory costs during DNN training. For future work, DropIT can be explored in training large (vision-Nanguuage models.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
This research is supported by the National Research Foundation, Singapore under its NRF Fellow- ship for AI CEE-NRRFFAT-2019-000 tions expressed in this material are those of the author(s) and do not reflect the views of National Research Foundation, Singapore.
Any opinions, findings and conclusions or recommenda-
REFERENCES
Alham Fikri Aji and Kenneth Heafield. Sparse communication for distributed gradient descent. In EMNLP, pp. 440-445, 2017.
Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. arXiv: 1607.06450 2016.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In NeurlPS, volume 33, pp. 1877- 2020.
Samuel Rota Bulo, Lorenzo Porzi, and Peter Kontschieder. In-place activated batchnorm for memory-optimizee training of dnns. In CVPR, pp. 8639-5647, 2018.
Ayan Chakrabarti and Benjamin Moseley. Backprop with approximate activations for memory- efficient network training. NeurIPS, pp. 2426-2435, 2019.
Jianfei Chen, Lianmin Zheng, Zhewei Yao, Dequan Wang, Ion Stoica, Michael W. Mahoney, and Joseph Gonzalez. Actmn: Reducing training memory footprint via 2-bit activation compressed training. In ICML, pp. 1803-1813, 2021.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv, 1604.06177 74, 2016.
Zhao Chen, Jiquan Ngiam, Yanping Huang, Thang Luong, Henrik Kretzschmar, Yuning Chai, and Dragomir Anguelov. Just pick a sign: Optimizing deep multitask models with gradient sign dropout. NeurlPS, pp. 2039-222290 2399-2050, 2020.
Feng Cheng, Mingze Xu, Yuanjun Xiong, Hao Chen, Xinyu Li, Wei Li, and Wei Xia. Stochastic backpropagation: memory efficient strategy for training video models. In CVPR, pp. 8301- 8310, 2022.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2021.
Nikoli Dryden, Tim Moon, Sam Ade Jacobs, and Brian Van Essen. Communication quantization for data-parallel training of deep neural networks. In MLHPC workshop, pp. 1-8, 2016.
R. David Evans and Tor M. Aamodt. AC-GC: lossy activation compression with guaranteed conver- gence. In NeurlPS, pp. 27434-27448, 2021.
David Evans, Lufei Liu, and Tor M Aamodt. Jpeg-act: accelerating deep learning via transform- based lossy compression. In ISCA, pp. 860-873, 2020.
Jianwei Feng and Dong Huang. Optimal gradient checkpoint search for arbitrary computation graphs. In CVPR, pp. 1143--11442, 2021.
10

---

Published as a conference paper at ICLR 2023
Fangcheng Fu, Yuzheng Hu, Yihan He. Jiawei Jiang, Yingxia Shao, Ce Zhang, and Bin Cui. Don't waste your bits! squeeze activations and gradients for deep neural networks via tinyscript. In ICML, pp. 3304-3314, 2020.
Audrunas Gruslys, Rémi Munos, Ivo Danihelka, Mare Lanctot, and Alex Graves. Memory-efficient backpropagation through time. In NeurIPS, pp. 4125-4133, 2016.
Kerui Gu, Linlin Yang, and Angela Yao. Dive deeper into integral pose regression. In ICLR, 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. CVPR, pp. 770-778, 2016.
Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask R-CNN. In ICCV, pp. 2980-2988, 2017.
Dan Hendrycks and Kevin Gimpel. Gaussian enror linear units (gelus), arxiv:1606.08411 2016.
Andrew G. Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv: 7704 048611 2017.
Animesh Jain, Amar Phanishayee, Jason Mars, Lingjia Tang, and Gennady Pekhimenko. Gist: Efficient data encoding for deep neural network training. In ISCA, pp. 776-789, 2018.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convo- lutional neural networks. Communications of the ACM, 60(6):84990 90,44990,222 2017.
Yanghao Li, Hanzi Mao, Ross Girshick, and Kaiming He. Exploring plain vision transformer back- bones for object detection. In ECCV, pp. 280-296, 2022.
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. Microsoft COCO: common objects in context. In ECCV, pp. 740-755, 2014.
Tsung-Y Yi Lin, Piotr Dollár, Ross B. Girshick, Kaiming He, Bharath Hariharan, and Serge J. Be- longie. Feature pyramid networks for object detection. CVPR, pp. 936-944, 2017.
Yujun Lin, Song Han, Huizi Mao, Yu Wang, and Bill Dally. Deep gradient compression: Reducing the communication bandwidth for distributed training. In ICLR, 2018.
Xiaoxuan Liu, Lianmin Zheng, Dequan Wang, Yukuo Cen, Weize Chen, Xu Han, Jianfei Chen, Zhiyuan Liu, Jie Tang, Joey Gonzalez, Michael Mahoney, and Alvin Cheung. GACT: Activation compressed training for generic network architectures. In JCML, pp. 14139-14152, 2022.
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory F. Diamos, Erich Elsen, David Gareía, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed precision training. In ICLR, 2018.
Arvind Neelakantan, Luke Vilnis, Quoc V. Le, Ilya Sutskever, Lukasz Kaiser, Karol Kurach, and James Martens. Adding gradient noise improves learning for very deep networks. arXiv:1511.06800 2015.
Zizheng Pan, Peng Chen, Haoyu He, Jing Liu, Jianfei Cai, and Bohan Zhuang. Mesa: memory- saving training framework for transformers. arXiv:2111.11124 2022.
Adam Paszke, Sam Gross, and Francisco et al Massa. Pytorch: An imperative style, high- performance deep learning library. In NeurIPS, pp. 8026-8037, 2019.
Alec Radford, Jeff Wu. Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models unsupervised multitask learners. 2019.

---

Published as a conference paper at ICLR 2023
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: memory optimizations toward training trillion parameter models. In SC, pp. 20, 2020.
Shaoqing Ren, Kaiming He, Ross B. Girshick, and Jian Sun. Faster R-CNN: towards real-time object detection with region proposal networks. IEEE Trans. Pattern Anal, Mach. Intell., 39(6): 1137-1149, 2017.
Olga Russakovsky, Jia Deng. Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael S. Bernstein, Alexander C. Berg, and Fei-Fei Li. Imagenet large scale visual recognition challenge. Int. J. Comput. Vis., 115(3):211-255 2015.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In ICLR, 2015.
Sebastian U. Stich, Jean-Baae Baptiste Cordonnier, and Martin Jaggi. Sparsified SGD with memory. In NeurlPS, pp. 4452-4463, 2018.
Nikko Strom. Scalable distributed DNN training using commodity GPU cloud computing. In Inter- speech, pp. 1488-1492, 2015.
Ke Sun, Bin Xiao, Dong Liu, and Jingdong Wang. Deep high high-resolution representation learning for human pose estimation. CVPR, pp. 5693-5703,222 2019.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In ICML, pp. 10347-10357, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NeurIPS, pp. 6000-6010, 2017.
Guanchu Wang. Zirui Liu, Zhimeng Jiang, Ninghao Liu, Na Zou, and Xia Hu. A concise framework of memory efficient training via dual activation precision. arXiv:2208.04 2022.
Saining Xie, Ross B. Girshick, Piotr Dollár, Zhuowen Tu, and Kaiming He. Aggregated residual transformations for deep neural networks. In CVPR, pp. 5987-5995, 2017.
David Junhao Zhang, Kunchang Li, Yali Wang, Yunpeng Chen, Shashwat Chandra, Qiao, Luoqi Liu, and Mike Zheng Shou. Morphmlp: An efficient mip-like backbone for spatial-temppral representation learning. ECCV, pp. 230-248, 2022.
12

---

Published as a conference paper at ICLR 2023
A APPENDIX
A.1 COMPLETE CONVERGENCE ANALYSIS
Here we prove convergence of DropIT with min-k dropping strategy. By scaling the learning rate with factor of L, the gradient of min-k dropping is modeled as:
Jmin-k Igd + 2 (0,
(12)
where is zero-mean noise with variance of ç², and a, are varied each iteration.
We assume that the loss function F is I-smooth, i.e., differentiable and there exists a constant L 0 such that
I. F(v) S F(x) (VF()). x) ly x11? YR.eRR
(13)
Performing Taylor expansion we have:
B|F(I++) < F(x,) - (VF(x)), It+1 - vz.) + I²
E[IV (xx))1
2
15
< - (
2a
(14)
Rearranging the terms of the above inequality and dividing by 3. we obtain:
[VF(2)) 2(F(xx) B[F(x++1)))
nLE?
(15)
a
Summing up from to T' and divided by T, we get:
VF(x))) 2(F(x1) ((F(z))-F(aa) F(x*))
1
3
(16)
Tn
T
o
where x² indicates an optimal solution.
A.2 MODELING MIN-À DROPPING GRADIENT WITH NONLINEAR FUNCTION
We can replace Eq. (gradient model of min-k dropping) with nonlinear function and still achieve the same convergence as in Eq. 10.
The gradient is biased with min-k dropping, assume it can be modeled as:
Amin-k Agd + (0 <²) +
(17)
where is bias and ||b|] < (1 (1-aa 11)lgll) a and varies each iteration, i.c., a = {ar, a2,.... a..
and B {81, ,2 3.).
13

---

Published as a conference paper at ICLR 2023
Assuming the loss function F is L-smooth, we obtain:
z1) + T
& - ( -1))
E +
2
= -(()) 7 (VF(( VF + 1
2
S f(() KVF VF(()) b) ! (VF())) + 1
LE
[(x) 7 2(VF(()) (1.) (F()))) 1) (1 b11))
2
|VFF(x)) - [1b12) T
= f(x) !
2
2
[VF(,))] L5 2
S ((a)
1a 2
(18)
Rearranging the terms of the above inequality and dividing by 201, we obtain:
[ 2(F((z) - E(F(x++1)))
L
1a
(19)
Using a learning rate of 12 instead, we have:
[V((2)) 2(F(x)) - E[F(x(+1)]
n p/
of
1
(20)
Summing from = to 7 and divided by T, we get:
1 [F(x))1 2(F(z1) (((1)))((="))
1
33
T
Tn
n
(21)
where x* indicates an optimal solution. The convergence is exactly the same as in Appendix A.I.
test
Optimizer
Learning Rate
Convergence
SGD
77
AF
Ty
+ n2L
SGD
Ck Tn
AF
+ L
DropIT
AF Th
nLL T ΣT B
DropIT (modeling nonlinearity as in A.2)
"
AF T' + tnL T 2
Table 6: Theoretical convergences of SGD and DropIT under L-smooth condition
In Table 6, we compare convergence of SGD and DropIT under various learning rates. Under fixed learning rate, SGD and DropIT differ no both convergence speed (the 1st term convergence) and emor (the 2nd term in convergence formula). For a fair setting, we compare SGD with learning rate n and DropIT with learning rate n. With a fixed convergence speed, DropIT theoretically achieves lower enror.
A.3 THEORETICAL ANALYSIS ON AND 3
In this section we compare the gradients of SGD and DropIT with min-k dropping. Note we slightly change the notation of the gradients from Fogd and Jmink in the main paper to improve clarity for
14

---

Published as a conference paper at ICLR 2023
element- wise analysis. We denote the gradient of SGD G and DropIT G'. Both gradients are computed by an input tensor and intermediate tensor Z. In DropIT we drop y percent of the elements in A. Thus we have:
G = A x G' = (A o D) x
(22)
(23)
where @ is glement-wiss, multiplication and D is dropping mask where each element is either 1 or 0.
From element-wise viewpoint, rewrite the computation of gradients:
gij =
Σ Uij
(24)
9ij Σ Zuj
(25)
where dæ is mask, i.e., du is either 0 or I depending azk- For the simplicity of analysis, we assume and Z are independent. Let μ be the mean value of A and be the mean of dropped elements, after dropping, A ④ D has mean of μ Taking expectation over all possible inputs, we have:
El TE. (((gggg
a.k
k
9sj-
(26)
P
Therefore the bias caused by dropping expected to be μ=c Assuming the mean value of A μ and the mean of dropped value e, after dropping, ((A) has mean of μ c. Thus the bias caused by dropping is R[ax] = μ=e. Recall that we drop elements with small absolute values. In the extreme case where every element in has the same value as μ,, e will reach the upper bound y.p. Therefore, E[a] Y-
Now we analyze on noise and compute 3. Due to the variation on input samples, we have noise in A and Z. which results in noise in G and G'. To highlight the noise, rewrite a noisy element as ø + fz. where 7 is the mean value of x and n, is zero-mean noise. Applying it to Eq.24 and 25 arrive at:
Dij + ng =
Σ (āsk + (ak + nz)
(27)
g' + ng Σ der (ast + na) + nz).
(28)
Focusing the noise of gradients, obtain:
ng ((ann + EksNaa + na7)))
(29)
Rgg Edunāgn + da + d..nan)))
(30)
Recall that du is mask depending on and therefore depending 0æa, thus from Eq.26 we have
E(dggā=))
H aik Fo Tike > (1 7)
(31)
11
15

---

Published as a conference paper at ICLR 2023
Dataset
Method
Top-1
Top-5
Memory (GB)
ResNet-18 (32 32)
77.96
94.05
648
ResNet-18 (32 32) + DropIT (7 = 0.8)
78.17
94.19
598
CIFAR-100
ViT-B/I6 (224 224)
90.32
98.88
20290
ViT-B/16 (224 224) + DropIT (y = 0.9)
90.90
99.02
16052 x 4
ResNet-18 (224 224)
69.76
89.08
2826
ResNet-18 (224 x 224) + DropIT (y 0.8)
69.85
89.39
2600
ImageNet
ViT-B/16 (224 224)
83.40
96.96
20290x4
ViT-B/16 (224 224) + DropITT (7 0.9)
83.61
97.01
16056x4
More results of different network architecture achieved by DropIT. ResNet-18 results are training from scratch, and ViT-B/16 fine-tuning from public ImageNet-2Ik weights.
Table 7:
Because Y percent of D is 0 and the other 1 y percent of D is 1, we obtain E[dix] = Y Plugging them in Eq.30 we have:
E[ = B[
azkli (1 Eaykna + - m)m))s
H
1
CakMt + EkyMa + naa
B[angl.
(32)
where the inequality satisfied due to c S 7p-.
This result tells E[i]] s B[a)), which suggests DropIT with min-k dropping has noise reduction effect and should converge better than SGD.
A.4 RESNET-50 TRAINING FROM SCRATCH ON IMAGENET-IK
We follow torchvision training script to train ResNet with and without DropIT. No hyper- parameters are changed. When 70%, ResNet-50 with DropIT achieves 76.3 top-1 accuracy, slightly higher than baseline's 76. accuracy, demonstrating the effectiveness of DropIT.
A.5 MORE NETWORK RESULTS
We present more results of different network architectures shown in Table 7. ResNet-18 trained from scratch by 90 epochs on ImageNet, totally following archvision reference code. VIT-B/16 is fine-tuned in 3 epochs from its ImageNet-214 pretrained weights. Our proposed DropIT can improve the accuracy for these setting with lower GPÜ memory cost.
A.6 WHY DROPIT IS NOT USED FOR THE NETWORK FIRST FINAL LAYERS
We do not apply DropIT conv/ / fc layer if is the first/final layer of the network. The reason is that this does not save memory:
(1) The first layer: The logic of our DropIT on saving memory can be concluded as: creating smaller tensor Xdropped (i.e. by torch. topk) from input tensor x, then the input tensor x will be automatically released by python garbage collection. However, popular code style is like:
dataloačer
loss_
def model (x)
layerl
x layer2 (x) x can be released with DropIT
= Layeril (x) can be released with DropIT
layeri (x) x can be released with DropIT, but cannot. save maxi memory
16

---

Published as a conference paper at ICLR 2023
10
return
11
12
for x, in dataloader:
13
= model (x) # x will not be released with DropIT
loss func (x, .).bckwarrd))
15
As we can see, in the dat aataloader loop. the input x to model can only be recycled when model rumning is finished. So, using DropIT in layerl will not reduce maximum memory instead, it will increase the maximum memory as DropIT created new Xdrapped-
(2) The final layer: it is easy understand that DropIT using in the final layer has effect memory. See the code block, when running to layeri, the maximum memory should be 1 layerl layeril 1 cached tensors plus input layeri. If we DropIT layeri, then there would be an extra Xdrapped produced, making the maximum memory even higher.
A.7 HOw TO SELECT OF DROPIT
From our experiments, we recommend y 70% for training from scratch and y 80, 90% for fine- tuning. As DropIT incurs memory cost for indexing. 7 should be larger than 50% to be meaningful (assuming index data type is int with the same number bits of float for activation). Empirically. we observe that y is reflected consistently in both training loss and testing accuracy. A too-high y which will bias the gradient will have training losses higher than the baseline. As such, alternative way to select y is to observe the training loss after some iterations (e.g. 100); if it is lower than the baseline, then the testing accuracy likely to improve as well.
A.8 MORE EXPERIMENTAL DETAILS
We list the detailed key training hyper-parameters, though they are totally the same with the offical implementations:
DeïT-Ti, training from scratch, ImageNet- Ik, w/wo DropIT batch size 1024, A optimizer, learning rate 10-3 weight decay 0.05, cosine LR schedule, 300 epochs, with auto mixed precision (AMP) training;
DeiT-S, finetuning from official DeiT-S ImageNet-Ik weights, CIFAR-100, w/wo Drop/T batch size 768, SGD optimizer (momentum 0.9), learning rate 10-2, weight decay 10-4, cosine LR sched- ule, 1000 epochs, with AMP training;
DeiT-B, finetuning from official DeiT-B ImageNet-Ik weights, CIFAR-1O0, w/wo DropIT': batch size 768, SGD optimizer (momentum 0.9), learning rate 10 weight decay 10-4, cosine LR sched- ule, 1000 epochs, with AMP training;
Faster R-CNN. finetuning from torchvision ResNet-50 ImageNet- Ik weights (VI), COCO, w/wo DropIT": batch size 16. SGD optimizer (momentum 0.9), learning rate 0.02, weight decay 10-4, Multistep LR schedule (16,22 epochs), 26 epochs, without AMP training:
Mask R-CNN, finetuning from torchvision ResNet-50 ImageNet-Ik Ik weights (VI), COCO, w/wo DropIT": batch size 16, SGD optimizer (momentum 0.9), learning rate 0.02, weight decay 10 Multistep LR schedule (16,22 epochs), 26 epochs, without AMP training.
http:///p........oo.o..ccc/cccta mask-r-cnn
17


Reviews:
[{"review_id": "ynVMkIetau", "reviewer": "Reviewer_zGSK", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper introduces a novel idea of reducing memory footprint during training by dropping elements in the intermediate tensors whose absolute values are closest to zero.\n\nThis paper makes clear contribution in terms of the following aspects:\n* proposed a novel approach that is different than gradient checkpointing which requires recomputation, activation compression which uses quantization rather than sparsification or gradient approximation which has a pre-define fixed threshold.\n* induces a natural noise reduction that could benefit SGD convergence, which is theoretically proven and justified by experiments;\n* demonstrates the effectiveness of the approach with comprehensive experiments, which shows that in most of the cases, most (70% - 90%) of the elements in the intermediate tensor could be dropped with even better accuracy.\n", "summary_of_the_review": "This paper is clearly written with clear novelty and contribution. The reviewer believes that its strength outweighs its limitation, and would love to get some clarity in rebuttal, specifically on weakness 1, 2 and 3.\n", "strength_and_weaknesses": "Strength\n* The unsurprisingly positive side effect of helping noise reduction in SGD convergence makes this method outstanding, specifically compared with some of other gradient approximation or quantization works, which if not tweaked carefully to balance with memory saving, could lead to inferior performance;\n* The proposed approach can be simple enough as a pure python plugin to a PyTorch model (Appendix A.4), which is an advantage to practitioners who wanted to quickly test;\n* This approach does not require rematerialization, is orthogonal to quantization, and do not have performance overhead compared with methods like GACT which requires to time-consuming profiling.\n\nWeakness\n* Saving the indices, which is an 32bit integer or possibly 64bit on large tensors, could be extra non-trivial storage overhead. The possibility of swapping them to CPU, which as mentioned in the paper indeed helps to mitigate the issue, would incur another overhead of slow transfer between CPU and GPU. Consider an optimistic case where 90% of the elements are pruned (fp32), but needs one extra int32 index, the compression rate becomes 5x. Compared with 4bit quantization which is 8x, it does not have significant advantage; In the worse case, if only 70% of the elements are pruned to retain accuracy, then the compression rate drops to 1.67x, which would be less ideal in general.\n* The general tradeoff between training speed vs accuracy is unclear. While Table 5 presents comparison between the proposed method and activation quantization approaches, it is not clearly demonstrated how this approach is compared with other lines of work, for example, rematerialization.\n* It’s unclear why the proposed work is only used in FC cache in Table 5. Would be more positive evidence if it could show benefit of pruning other caches while retaining the same level of accuracy. Appendix A.5 provides some general explanation on the inapplicability to a network’s first and last layer, which is completely understandable, but it is not clear how it is connected to Table 5.\n* GPU sorting or min-K could be highly non-trivial to optimize beyond using `torch.topk`, and this could possibily incur non-negligible overhead in terms of training speed. The overhead of such operation is not carefully studied in the experiment section.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "eI95MljzbH-", "reviewer": "Reviewer_GQ3x", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose DropIT, where stashed intermediate tensors are pruned via a Top-K function and converted to a compressed sparse format. This sparsification reduces the memory consumption of stashing, enabling larger batch sizes and faster training. During the backward pass, the stashed tensors are first decompressed to dense representation since the sparsity is usually not high enough for sparse matmuls to be effective.\n\nThe authors show that by scaling the learning rate, DropIT can actually decrease the gradient noise.\n\nExperiments show that a variety of vision benchmarks such as DeiT and R-CNN can be trained on datasets like ImageNet and COCO with 70%-90% sparsity without reducing the final accuracy. This is a great result. However, wall-clock time improvements seem less impressive, without only about 10% run time improvement in Table 5.\n\nEDIT: raised score from 5 to 6 after rebuttal", "summary_of_the_review": "DropIT proposes Min-K pruning of stashed activations to save memory. This isn't a novel idea and has been investigated in at least the GIST paper[1]. However, this paper is still useful as it presents an isolated study on this technique and presents impressive results on memory savings for vision benchmarks.\n\nI would raise the score if I see additional discussion regarding the assumptions for the theoretical proof and results on more memory-intensive benchmarks.\n\n[1] https://ieeexplore.ieee.org/abstract/document/8416872?casa_token=0XtaIex9huQAAAAA:RwuzO_eOgBsEms6Y3XESj66oOFSaihBdAT0vVXlq_YETd7xlO_ifa9feFTDT0wrGlU1KKM7D", "strength_and_weaknesses": "Strenghts:\n - Straightforward and easy to implement method that seems fairly robust, at least on vision benchmarks\n\nWeaknesses:\n - The idea is not novel: it's already part of the GIST paper. GIST combines Min-K stashed activation compression with binarized stashed ReLU masks. The GIST paper is systems-focused and contains multiple techniques. This paper gives a more focused study of the Min-K sparsity technique and is still a useful result.\n - The memory savings is great but the actual run time savings seem unimpressive. Perhaps it would be better to focus on benchmarks which are much more memory-hungry such as Transformers. These models are typically trained with microbatching and increasing the microbatch size should grant a big performance boost.\n\nQuestion:\n - The proof that DropIT can decrease the gradient noise depends on Equation 8, where the noise introduced by the Min-K sparsification is modeled as $g_{min-k} = \\alpha g_{gd} + \\beta n(0, \\xi^2)$. It's not clear to me why this noise is zero-mean. In CNNs, we have non-symmetric activation functions like ReLU. Min-K sparsity would change the means of the stashed activations tensors, which in turn would change the means of the gradient update tensors. I am not an expert here so I would like to see the authors' response.\n\n[1]https://ieeexplore.ieee.org/abstract/document/8416872?casa_token=0XtaIex9huQAAAAA:RwuzO_eOgBsEms6Y3XESj66oOFSaihBdAT0vVXlq_YETd7xlO_ifa9feFTDT0wrGlU1KKM7D", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "EyFfxY-2Qkd", "reviewer": "Reviewer_xcyW", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a strategy, named DropIT, for activation compression during training. The proposed method can be applied with random dropping or min-k based dropping. The authors show that DropIT can also reduce the gradient noise during training, and provided empirical evidence that the proposed method can drop up to 90% intermediate tensor elements on FC and convolutional layers.", "summary_of_the_review": "Given the concerns on the theoretical analysis and novelty of this paper, I do not think it is ready to be accpeted to ICLR at its current format. ", "strength_and_weaknesses": "Strength\n\n- The paper is well-written and easy to understand. The graphical illustration really helps build understanding of the proposed technique.\n- The evaluation is on large scale vision datasets (eg. ImageNet and COCO), also vision transformers are included.\n\nWeakness\n\n- A large volume literatures on Dropout are ignored by the authors? I would think this line of work is directly related to this paper.\n- I am not comfortable with the theoretical analysis shown in Section 3.4.\n    - $g_{min-k} = \\alpha g_{gd} + \\beta n(0, \\sigma)$ (Equation 8) does not look like a good approximation to your min-k operation. Simply speaking, min-k is non-linear and I do not really understand how this linear function can approximate that well?\n    - The convergence proof is confusing. Surely your bound is tighter if you use a learning rate that is scaled by $\\frac{1}{\\alpha}$? Like this does make too much sense, if I say SGD has a smaller LR and its LR scale is greater than my noise scale, it is surely tighter under the L-smooth condition. What is the actual point of showing this? Do you consider this as a proof?\n- The theoretical gain is not achievable.\n    - The results in Table3 is confusing. In my opinion, the authors are saving ‘memory bandwidth’ but not really using smaller GPU memory. Are you sure your forward activation values are kept on-chip (in SRAM or Cache) in between the feedforward and backward computation? I thought the GPU model will have to push these values into DRAM after finishing one layer of forward computation and load them back from DRAM for backward computation? If this is the case, the proposed technique does not optimize for SRAM space, but in fact is optimizing for the DRAM bandwidth. I will say this is an equally important topic for efficiency since most training workloads are DRAM-bound, but I do not think this paper is making a clear discussion on this.\n    - ‘the memory reduction is precisely controlled by γ,\n    *i.e*. γ = 90% means the reduction is 11.26×0.9’ → I simply disagree with this statement and do not think this is achievable with out custom silicon.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "BJg3D9jdBC", "reviewer": "Reviewer_Ltqg", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents a method that reduces the memory usage of training/fine-tuning for convolutional neural networks and vision transformers. The idea is to prune away insignificant intermediate tensors (i.e., those that are low in value) required for the back-propagation. The method has been tested on DeiT and R-CNN models. It has been shown that up to 90% of intermediate tensors can be pruned away without any performance degradation. However, this reduction doesn't linearly translate into memory reduction. For example, it was shown that with sparsity rate of 90%, there is only about 1 GB memory reduction when fine-tuning DeiT on CIFAR-100.", "summary_of_the_review": "In my opinion, the contribution of this paper is novel and is addressing a real challenge which memory usage reduction during fine-tuning. However, the proposed method was not fully tested which makes me wonder about its generality. For the rebuttal, I would like to see how this method works when fine-tuning ViT-base model on CIFAR-10 and CIFAR-100 for 3 to 4 epochs, which is a common practice. It would be great to include some results from BERT on GLUE and SQuAD tasks too (again when fine-tuning it for a couple of epochs). I would also like to see some results for small-size CNNs such as MobileNet on ImageNet-1K or even ResNet-18 on CIFAR-10. I am also wondering why there is only 1GB memory reduction on DeiT when gamma is 90% in Table 5; how much of the memory usage is dedicated for the intermediate tensors out of 6.66GB in this case? ", "strength_and_weaknesses": "(Strengths:) In general, I believe the idea of reducing the memory usage of neural networks during training/fine-tuning is of paramount importance specially for transformer models. This work is an attempt to address this issue with a novel idea. The paper is also well-written and easy to understand.\n\n(Weaknesses:) However, I am not convinced about the effectiveness of the proposed method due to lack of sufficient experiments. More precisely, the choice of DeiT and fine-tuning it for 1000 epochs doesn't reflect the effectiveness of the proposed method. Instead, I believe ViT models should be used for fine-tuning which requires up to 4 epochs. Moreover, I expected to see more CNNs tested such as ResNet-18/MobileNet on CIFAR and ImageNet as an example. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster