Few-bit Backward: Quantized Gradients of Activation Functions for Memory Footprint Reduction
Authors: Georgii Sergeevich Novikov
Daniel Bershatsky
Julia Gusak
Alex Shonenkov
Denis Valerievich Dimitrov
Ivan Oseledets
Venue: Submitted to ICLR 2023 (2023)

Abstract: Memory footprint is one of the main limiting factors for large neural network training. In backpropagation, one needs to store the input to each operation in the computational graph. Every modern neural network model has quite a few pointwise nonlinearities in its architecture, and such operation induces additional memory costs which --- as we show -- can be significantly reduced by quantization of the gradients.
We propose a systematic approach to compute optimal quantization of the retained gradients of the pointwise nonlinear functions with only a few bits per each element.
We show that such approximation can be achieved by computing optimal piecewise-constant approximation of the derivative of the activation function, which can be done by dynamic programming. The drop-in replacements are implemented for all popular nonlinearities and can be used in any existing pipeline. We confirm the memory reduction and the same convergence on several open benchmarks.

Under review as a conference paper at ICLR 2023
FEW-BIT BACKWARD: QUANTIZED GRADIENTS OF AC- TIVATION FUNCTIONS FOR MEMORY FOOTPRINT RE- DUCTION
Anonymous authors Paper under double-blind review
ABSTRACT
Memory footprint is one of the main limiting factors for large neural network training. In backpropagation, the computational graph. Every modern neural network model has quite a few pointwise nonlinearities in its architecture, and such operation induces additional memory costs which as we show can be significantly reduced by quantization of the gradients. We propose a systematic approach compute optimal quanti- zation of the retained gradients of the pointwise nonlinear functions with only a few bits per each element. We show that such approximation can be achieved by computing optimal
needs to store the imput each operation in
consstntt approximation of the derivative of the activation function, which can be done by dynamic programming. The drop-in replacements are implemented for all popular nonlinearities and can be used in any existing pipeline. We confirm the memory reduction and the same convergence on several open benchmarks.
INTRODUCTION
1.75
0.25
derivative
1.50
0.20
1.00
0.15
Signcid
0.75
0.10
b-bis Siansid
0.05
0.00
0.00
-10
-5
10
-10
10
-10
-5
10
Figure 1: Examples of 3-bit approximations for derivatives of popular nonlinearities: GELU, SELU, and Sigmoid.
Modern neural network models are getting larger and larger. One of the main bottlenecks in the training loop is the required device memory storage Ojika et al. (2020); Gao et al. (2020). In this paper, we propose universal approach that helps to reduce the model memory footprint during backpropagation. Note that this approach is complementary to other memory reducing techniques such as checkpointing Chen et al. (2016) or offloading Beaumont et al. (202I). Our method can be applied to any neural network without any additional preprocessing.
Memory consumed by the model during training (except intermediate tensors) can be split into two groups: 1) the model weights (including additional memory for the optimizer state), 2) activations saved for the backward pass, over which the computation is not carried out directly at the moment, but which will be required in the future to compute the gradients.
Every operation in the computational graph generates a memory footprint. It is typically overlooked, that the application of the pointwise non-linearity (such GELU sigmoid) results in storing the

---

Under review as a conference paper at ICLR 2023
input for the backward pass. We show that instead of keeping the full input tensor, it is possible to store a low-bit representation, which allows accurate gradients approximation.
In this work, propose to approximate the derivative of the activation function in piecewise- constant form. Such an approximation problem has to be solved once for each activation function, and we propose simple technique do that.
The proposed approximation divides all values into several bins and saves only their corresponding bin indices instead of storing all values. This is lossy compression, but the additional noise introduced by it is negligible as we will show on several benchmarks in Section
The main contributions of our paper are:
We propose approximate backward computation schemes that significantly reduce the memory consumption of neural network training.
We benchmark our approach on several tasks. We show that it provides up to 40% memory reduction on various tasks while maintaining accuracy on par with the model trained via the standard approach.
QUANTIZED GRADIENTS OF ACTIVATIONS
Gradients activations using automatic dif- ferentiation. Modern deep learning frame- works use the reverse mode automatic differ- entiation to calculate the gradients of the loss over the model parameters. Forward computa- tion can be associated with directed acyclic graph, depicted in Fig. 2 Each operation / com- putes the output X(+1 given the input X, and has to save some information S; that would be used on the backward pass in order calculate the derivative / from OL/8X1+1 and S,. Thus, in a typical training loop, the intermedi- ates S, of all operations in the graph are stored in the memory during the whole forward pass until they are no longer needed after the completion of the corresponding backward operation dur- ing backward pass. This generates additional memory, which can be quite significant and be larger than the total amount of parameters of the model.
Eoowant
Backkwent
Inpest
(0.05, 0,47, 1.12, 0.83)
3X
Tensors caved for bedkward
fnlermedinte lenases for Cradient computation
Save
Gradient f
Quantise
Cuartaaed saved for cochward
and
Output X..1
(0.5, 0.5, 0, 1,1) 01,00,
Figure 2: Computation graph of both forward and backward pass. Orange and purple parts of the graph correspond to standard and proposed ways of saving tensors for backward, respectively. Vector æbit stands for the tensor saved using 2-bit quanti- zation, while æ denotes its uncompressed version.
Pointwise activations. In this paper, we focus on a pointwise activation function, which is ubiqui- tous in modern neural network architectures. Given an input tensor X, we apply function each of the elements of this tensor:
f(X:) = (X (e..).))))))))))) "1) Ma......k
This operation very cheap compared to other operations the deep neural network model and does not attract much attention when analysing computational complexity. However, standard implementation in such a framework as PyTorch induces not a very small memory footprint and the whole input X, is saved for the backward pass.
The backward pass for such a function consists of element-wise multiplication of the propagated gradient tensor by the derivative of the nonlinearity function at the points of the input tensor: if X f(X;), then the gradient of the loss L with respect to X₁ is computed as
aL = 0x
aL
(((i),
(I)
0x.
2

---

Under review as a conference paper at ICLR 2023
where fX( is the tensor with elements, consisting of the derivative of f evaluated in each element of X₁. From Eq. (I). it follows that for the backward pass we have to store only f' and X is not needed.
ReLU activation function. To illustrate our idea, consider one of the most popular nonlinearities, f()) (c)U(x) max(0, 2). Its derivative f' takes only two values, 0 and 1 and it only requires 1 bit to store. single precision is used, then the compression is 32, which quite noticeable.
GELU activation function. In modern transformer architectures Vaswani et al. (0117) the GELU Hendrycks & Gimpel (2016) nonlinearity is typically used. The derivative no longer takes two values. Instead, we propose to approximate f' by a prrcewiss-constane function. For example, if we allow 8 different values, we will need only 3 bits per each element (Fig.)
Quantized gradients of activations. In stochas- tic optimization, if the gradient for given batch computed approximately. the optimization may still converge. The GELU derivative (see
81
Fig. D quite "similar" to piecewise-consten function: for large values of Jæe], it is almost ex-
actly equal to 1, and for small values of x, rather interesting transition from 0 to occurs.
Us
Instead of calculating the derivative exactly on the backward pass, approximate it using certain piecewise- biecewise-constane constant approximation:
y1
elels.v)
.33 86
y( Σ yelfa [8s; 1:41l
(2)
Figure 3: GELU derivative and its approximation
where s = (81, ,8++1) is sorted vector of intervals, on which approximation is constant, (....... ) is vector of the correspond- ing values of approximation and denotes an indicator function, which equals whenever its argument is true and 0 otherwise. That means, that y equals when E [8;; 8++1], see Fig. 3 for illustration. As noted above, if the approximation
s. with five piccewise-constant intervals
constant intervals, instead of storing the full input tensor X, it will be possible to save only log k bits of information (per element of the input tensor), which, accordingly, will reduce the memory consumption by 32/ log & times for single precision.
quantizatoin scheme Eq. (2 given, drop-in replacement for activation function / is very straightforward On the forward pass, instead of the full tensor X, have save only indices of intervals to which the elements of X belong, and on the backward pass, we need to multiply gradient output not with the actual derivative of f, but with values from y corresponding to stored indices. Pseudocode is presented in Alg.
Globally 1y stored pi piecewi se- constant. ixproximation parameters
y = I...I, ....]
def forward
X_pcs sortedsearch (s, X)
Lree_for_backwarddde return f (X)
def backward (dLdY)
10
p_p get_saved for feee__forbbaccwardo return dLdY y yoo_
11
Listing 1: Pseudo code for quantized backward layer. Arrays and are parameters of quantization Eq. (2). sortedsearch is binary search method.
Memory of Few-bit Appproximation. As it was mentioned above, by replacing all pointwise nonlinearity layers the neural network with Few-bit approximation consisting of k piecewise- constant intervals, the memory consumption of such layers during forward-backward pass will be
3

---

weights Under review as a conference paper at ICLR 2023
reduced by 32/k times for single-precision learning mode. However, how many times in total the neural network memory consumption is reduced depends on the particular architecture of the neural network and the optimizer used in the process. During training, the memory is spent (parameters) of the network, on optimizer statistics, and on all stored activations, some of which are activations of pointwise nonlinearity layers. For example, when training ResNet18 with the Adam optimizer on 256x256 images, the model weights take 44.6M 44.6Mb, 3 44.6 = 133. 133.8Mb is used by the optimizer store gradients and moments, 4S0+00M is needed store all activations during forward- pass, BS 11 .5Mb of which are pointwise nonlinearity layers and BS * 28. 5b is for all other layers, where BS the batch size. Therefore, the maximum possible batch size with standard nonlinearities is ([GPU_MEM size is BGGPU_MENT
4 440,) while the maximum batch size with Few-bit nonlinearities of
4 44.6)/(28.5 + 11.5 * log k/32)J, where GPU_MEM is the available GPU memory. In our example with R for standard nonlinearity layers, the maximum batch size for a video card with 32Gb memory is 813, while using 4-bit Few-bit approximation is 1086 (+33%). Memory consumption for different Few-bit mods and different neural network architectures is presented in Appendix
Speed of Few-hit Approximation The memory gain of Few-bit layer does not slow down the speed. The standard nonlinearity layer calculates the activation function in the forward pass and the activation function gradient in the reverse pass. The activation function gradient usually includes complex functions such as exponent, erf, and others. The Few-bit version of the layer also calculates the activation function on forward pass, but the gradient calculation during backward pass is replaced by one binary search and one lookup in the value table (see Alg. [ Our efficient implementation of this procedure using CUDA kernels runs several percent faster than the standard nonlinearity layer. However, this result may depend specific framework implementation and the used GPU, so in our experiments in Section we do not consider the time gain, assuming that both layers are roughly equally fast, but focus specifically memory savings.
4
3 OPTIMAL PIECEWISE-CONSTAE APPROXIMATION
Fig. D shows examples of an optimized 3- 3-bit sescewise-coosstan approximation for several nonlinearity function. Finding the optimal approximation parameters (boundaries of intervals and values on them) is challenging task. We propose to find them by minimizing the (weighted) L2 norm of the error.
Consider function f : IR > R and its derivative f'. We will measure the quality of a piecewise constant approximation Eq. with weighted L2 norm:
min L(s,y), y,a
k Σ/
(s, = / [ (((s))
(f'(z) (x))d, //
where w is some weight function reflecting our prior knowledge of the activation function argument distribution. Practical choices of w may be either [[x E [A; B]] (with some reasonable A and B, which should be large enough) which makes integral Eq. tractable, or maybe, e.g., standard normal distribution.
ys is differentiable w.r.. s and y. so optimal piecewise-constrn approximations can be found using standard gradient-based optimization techniques. But the minimization problem Eq. (3) has many local minima that are far from optimal. We suggest using dynamic programming to get some good initial approximation that can be further finetuned using gradient-baser methods (but also can be used as is because it is very accurate on its own).
Dynamic programming. We will assume that the weighting function w chosen such that w((z) 0 for & [A; BJ. Consider the following auxiliary value:
[ ((((()) q(0lly, ())
DP(1,K)
min BLiks
A.t.
tER
4

---

Under review as a conference paper at ICLR 2023
Essentially, (t,,, k)) is the optimal piecewise constant approximation of size for the given function f' on the interval [A; t]. The recurrent formula for this value is:
min min (P((t +) [f(())) - ult', ((td,,)))uu ((),)))))())
(5)
Se (())))
(((' = J w(x)dr
(6)
since a piecewise-constant approximation of size k consists of corresponding approximation of size k (first term) plus one constant interval (second term). Here 1' chooses the right bound of approximation of size k, and ytt', t) stands for the optimal value for the interval [t'; Eq. 8 Then the minimal value of ((s, y) of size is equal to (P(B,)) k).
To solve the minimization problem Eq. ( we suggest considering the discretization t: A 4g < <... t, B and reducing the calculation of DP k.) to its approximation only in the points of discretization:
DP k)) min (DP(j, k- - T(,,))
St f Sr w(r)dr
TÚi,) 11 (f(()) ((() - - v()) z(i.i)
Eq. (7) can be calculated in O(n time and O(nK) space, which is described in Appendix G in detail. Please note, that this routine should be evaluated only once, possibly by the framework developers, and the used indefinitely. Which means that number of discritization points n can be taken quite large, tens of thousends easily. That would make global solutoin of discrete problem, given in Eq. 7 very close to the global solution of the original problem Eq. (3). We give precalculated Few-bit approximations for many different pointwise nonlinearity functions in our implementation at https: /
github. ctwp http://aoomooosrepsssooooo
4.0
GELU
GELU
3.9
1-bit GELU
3.9
1-bit GELL
2.hit GELU
2-hit GELL
3.8
3-bit GELU
3.8
3-bit GELU
Train Loss
4-bit GELU
Valid Valid Loss
4-bit GELU
3.6
3.5
3.6
3.4
3.3
3.5
3.2
500
1000
1500
2000
2500
500
1000
1500
2000
2500
Step
Step
Figure 4: activations.
Dynamic of loss values in finetuning ruDALL-E Malevich with Few-bit GELU
EXPERIMENTS
The goal our experiments is not only to show that the Few-bit nonlinearity approach provides memory savings during neural network training without loss of the final model quality. In addition, we want to experimentally prove that this approach does not change the learning dynamic itself its application in practice is almost completely safe: there is memory gain
because, in this without loss of speed or quality, and without risks interference with other training factors under study (hence, additional search fitting of other hyperparameters is needed). To achieve this
5

---

Under review as a conference paper at ICLR 2023
goal, in addition to the main metrics of the trained model (which depend on specific tasks and benchmarks), we also compare the training loss and validation loss graphs during the whole training process. Further you will see that 1-bit and 2-bit f-bit approximations already almost the same as the original nonlinearity layers. And the 3- and 4-bit Few-bit approximations achieve the original quality of the model.
We have tested two of the most important and commonly used neural network architectures: con- volutional neural networks and transformer-based networks. We use standard popular open-source benchmarks with open hyperparameters for training in order to demonstrate the behavior of the Few-bit approach under drop-in replacement of standard nonlinearities without any hyperparam- eter optimization or specially selected training conditions. In SSection 4.1, we test the RoBERT-a transformer-baser neural network on the GLUU Wang et al. (2019) benchmark, which includes different NI.P tasks. In Sectioio 4.2 we test the training of the generative ruDALL-e model in the task of modeling the joint distribution of text and image tokens for the Russian Emoji dataset. We use the GELU nonlinearity for both transformer architectures, as it is the main nonlinearity function used in such models. In Section 4.3) we test the classical R architecture on the ImageNet dataset using the open benchmark frev Leclere et al. (2022). In the classical ResNet architecture, replace all ReLU nonlinearities with one of GELU, SELU, or Swish to demonstrate that the Few-bit approach works with a wide range of different popular activation functions.
The main analog of our Few-bit approach is the ActNN method. In Section 4.4 we make a detailed comparison with this method.
The code to reproduce all experiments is available at /https://gittiiii repository and all hyperparameters for training are presented in Appendix F
com/ www./ooo//oosss/////oooo
4.1 GLUE benchmark. In Table we report results for RoBERTa-base model Liu et al. (209)) on GLUE benchmark Wang et al. (2019) for standard GELU and 1-, 2-, 3- and 4-bits Few-bit GELU. 1- and 2-bits versions have minor performance degradation, while 3- and 4-bits GELU have no visible difference and closely match vanilla GELU performance, which can be seen more clearly on the dependence of the metric. averaged across all GLUE tasks, on the number of bits in Few approximation, represented in Fig. 6 The behaviour of loss during training is depicted in Fig. 5 3- and 4-bit versions are hardly distinguishable from the standard GELU.
Table I: RoBERTa- on GLUE benchmark with different quantization budgets. Metrie: mean accuracy/correlation (lask specific). Averaged across five runs.
1-bit GELU
2-bits GELU
3-bits GELU
4-bits GELU
Vanila GELU
stsb
0.906 (1: 0.002)
0.907 (= 0.002)
0.910 (1: 0.002)
0.909 (= 0.002)
0.909 (+: 0.001)
mnli-mm
0,870 (= 0.001)
0.870 - 0.002)
0.871 (# 0.002)
0.870 0.001)
0.871 (# 0.002)
mrpc
0.880 (+: 0.009)
0.884 (+ 0.008)
0.884 (1: 0.007)
0.885 (= 0.008)
0.882 (1: 0.005)
cola
0.595 (+: 0.016)
0.580 (+ 0.014)
0.596 (+:: 0.015)
0.607 (= 0.014)
0.604 (+: 0.013)
mnli
0,873 (= 0.001)
0.872 (= 0.002)
0.874 (# 0.001)
0.874 0.002)
0.874 0.001)
sst2
0.939 (+: 0.003)
0.938 (+ 0.003)
0.941 (1: 0.004)
0.941 (=: 0.003)
0.943 (1: 0.002)
rtc
0.752 (=: 0.021)
0.756 (= 0.023)
0.780 (# 0.014)
0,771 (= 0.025)
0.771 (= 0.017)
qqp
0.914 (= 0.001)
0.915 (= 0.000)
0.916 (# 0.001)
0.916 (= 0.001)
0.916 (# 0.001)
qnli
0.925 (+: 0.002)
0.925 (+ 0.002)
0.926 (+: 0.002)
0.927 (= 0.002)
0.927 (+: 0.002)
4.2 RuDALL-E. In Fig. we present training dynamic of ruDALL- E Malevich Ramesh et al. (2021) model on Russian Emoji dataset. The dataset Shonenkov et al. (2021) contains 2749 unique emoji icons and 1611 unique texts that were collected by web scrapping (the difference in quantities is due to the fact that there are sets, within which emojis differ only in color, moreover, some elements are homonyms in Russian). ruDALL-E Malevich is big multimodal pretrained transformer, which learns the conditional distribution of images given some string of text (more precisely it autoregressively models the text and image tokens as a single stream of data). ruDALL-E Malevich encoder part is 24 layer Transformer Vaswani et al. (201)) model with 16 attention heads, 2048 hidden dimensions and standard GELU nonlinearity, which in total has 1.3B parameters. works with 128 text tokens, which are prepared from the text input using YTTM tokenizef and 1024 image
Implemennatiion is taken from ht Implementation is taken from httpp///ggthhbb.cc/////////////p

---

Under review as a conference paper at ICLR 2023
GELU
0.21
GELU
0.375
GELU
6 x 10
1-bit GELU
1-bit GELU
1-bit GELU
Training Loss
2-bit GELU
Training Loss
0.22
2-bit GELU
Validation Loss
0.350
2-bit GELU
3-hit GELU GELUU
3-bit GELU
3- GBLU
x 4x100
0.325
4-hit GELJU
0.300
3 x 10
0.275
2 x 1071
0.250
0
2000
11000
3000
3500
4000
1500
2000
4000
Iteration. (a)
Iteration. (b)
Iteration. (c)
Figure RoBERTa- base on QQP task from GLUE benchmark, averaged across 10 runs. (a): Training loss. (b): Training loss zoomed into the last third of the training. (c): Validation loss.
1.000
Relative Top-1 Accuracy
0.995
Average metric
0.854
0.985
0.852
Few-bit
0.980
ActNN
Few-bit GELL
Stundard
0.975
Nonlinearity
0.850
Standard GELU
1-bit
2-bit
3-bit
4-bit
1-bit
2-bic
3-bir
4-bit
Figure 7: Relative top-1 accuracy for ResNet18 net- work on ImageNet dataset, averaged across three nonlincarities: GLUE, SELU and Swish. For each nonlinearity approximation top-1 accuracy (Few- bit approximation and ActNN approach) was mea- sured relatively to the top-1 accuracy of the model with corresponding unaltered nonlinearity.
Figure 6: Task-specific metric, averaged across all tasks in GLUE benchmark. Blue line de- pendence on the number of bits in Few-bit GELU and dashed red line is standard GELU, With 3 bits approximation we already match unaltered nonlin- earity quality.
tokens, which are obtained after encoding the input image using Sber- VQGAN Few-bit backward for ruDALL-E Malevich shows same behaviour as for RoBERTa- base architecture: 1- and 2- versions, although coping with training perfectly fine, demonstrates minor performance degradation, while 3- and 4-bit versions are indistigguishable from the original GELU.
4.3 ResNet Architecture. We trained ResNet18 model He al. (2016) on ImageNet Russakovsky et al. (2015) benchmark Leclerc al. (2022) dataset with ReLU replaced with GELU, Swish and SiLU nonlinearity functions. Graphs for Swish nonlinearity can be seen in Fig. an graphs for other nonlinearities can be seen in Fig. 13 in Appendix [P] 1- and 2- bits have minor performance drop, while and bits are on par with standard nonlinearity.
4.4 ActNN. As baseline, we use another quantization scheme ActNN Chen et al. (2021). It works in much wider spectrum of situations, as it can quantize not only pointwise nonlinearity layers but also all kinds of linear layers (convolutional and dense layers), normalization layers and pooling layers. Without going deep into details, ActNN divides the saved tensor H into chunks h, where each chunk of an equal size Then, given the quantization budget bits, each chunk h, (h/)))m(xxx)) (h,) min(h:)) and its randomly quantized version saved ú, = u, with prob. [u [u:] otherwise. Random rounding performed in order to guarantee that the quantization is unbiased. For each group, two additional values min (h.) and
normalized:
(h,
Impllmmetatton taken from hwwm:::pe.///thhhb.c////ebbrrbbr

---

Under review as a conference paper at ICLR 2023
10
Swish Swish 2-bil Swish
Training loss
Training loss
2.6
Top-1 accuracy
0.728
3-bit Swish
2.5
0.726
x 10"
4-hit Swish
Swish
0.724
bit Swish
2.4
2-bit Swish
0.722
bit Swish
Swish
2.3
0.720
50 76
80
1-bit
2-bit
3-bit
4-bit
Epoch. (a)
Epoch. (h)
(c)
Figure 8: ResNet18 with ReLU replaced with Swish nonlinearity trained on Imagenet. (a): Training
loss. (b): Training loss zoomed into the last third of the training. (c): Final validation top-1 accuracy. All graphs are averaged across three runs with different seeds. Error bars denote minimum and maximum values.
(hx are saved as well, but for the group size of G 256 it only 0.125 additional bits per element, which we ignore in our following tests.
ActNN by construction does not take into account the global behaviour of the nonlinearity derivative. We argue that for nonlinearity layers, it is very crucial, and thus our preoptimized quantization scheme more preferable. To confirm that, we consider ActNN behaviour on the QQP task from the GLUE benchmark with respect to different quantization budgets and compare it with our method (Fig. and Table 22). In general, our method with bit less budget works the same or better than ActNN, which is very important in the low-bit setting.
In Fig. 10 we compare ActNN and Few-bit for ResNet18 architecture on ImageNet dataset for SELU nonlinearity, while results for GELU and Swish nonlinearities can be found in Fig. 14 in Appendix [F Aggregated top-1 accuracy for all activation functions presented in Fig. Our method steadily outperform ActNN which is especially noticeable for 1-bit regime: ActNN experience strong downgrade of accuracy, while Few-bit Backward has much closer performance to standard nonlinearities. This means that one-bit Few-bit backward be used in cases where is very important to reduce memory consumption by neural network.
ActNN
Our
1-bit
0.8880
00.0008)
0.9080 (=000006)
2-bit
0.9072
( 00005
0.9097
(#00000))
3-bit
0.9106
(=0000033
0.9T14 (+F00000)
b-t
0.9113 (=0.0)
0.9114 (+0.0005)
Table 2: Accuracy on QQP task from GLUE benchmark for ActNN and Few-bit (Our). Averaged
across 5 runs. Few-bit approach better for each memory budget.
0.300
x 10
Training Loss
Oer OILE
AcD G4L GELL
Training Loss
0.275
GELU AuRON
Validation Loss
AdDii AA GELL"
DBLE)
Acdit 24a GELU
AcdiE 24a GELL
Acbik GELL
4xx 10
24b GILL
0.250
Our 2-8 GELD
0.35
Our 248 GELU
3 x 10
0.225
0.30
0.200
10
0.175
0.25
2000
4000
3000
3500
4000
4500
2000
4000
Iteration. (a)
Iteration. (b)
Iteration. (c)
Figure 9: Comparison of RoBERTa-banr base QQP task from GLUE benchmark with ActNN quanti- zation and Few-bit approximation. Averaged across ten runs. (a): Training loss. (b): Training loss zoomed into the last third of the training. (c): Validation loss.
8

---

Under review as a conference paper at ICLR 2023
SELU Fow-bit AcINN
Top-1 accuracy
Top-5 accuracy
0.900
Training loss
0.71
0.895
2.3
0.70
SELU
0.890
Fow-bit
Few-bit
2.2
0.69
ActNN
0.885
ActNN
1-bits
2-bits
3-bits
4-bits
1-bits
2-bits
3-bits
4-bits
1-bits
2-bits
4-bits
(a)
(b)
(c)
Figure 10: Comparison of AcINN SELU with Few-bit SELU (Our) for ResNet18 architecture on ImageNet dataset. (a) Training loss. (b) Top-1 accuracy. (c) Top-5 accuracy. Our method with 1-bit already matches unaltered nonlinearity performance and significantly outperform 1-bit ActNN.
RELATED WORK
The reduction of the memory footprint is an important topic. To save memory during training, in addition to working with stored activations, the memory used to store model parameters can be compressed. Quantization BBodarenko al. (202T): ): Bengio al. (2013): Banner al. (2099) Jacob eet al. (2018); Nagel et al. (2021): Krishnamoorth (0018) limits the admissible values of weights to some small finite set. Thus, less memory needed for storage. The low-rank representation of weights Hrinchuk et al. (2020); Phan et al. (2020); Gusak et al. 2019 2021); Cui et al. (2020); Novikov (2018)))); Lebedev et al. (2013) assumes some internal structure of model weights and saves memory by explicitly using this structure with low-rank methods from linear algebra. Low precision learning and low precision optimizers focus on using the lower precision floats to store weights, optimization parameters, and model gradients. All of these approaches are complementary to the proposed one and can be used together.
Checkpointing Beaumont et al. 20199] 2021); Chen et al. (2016) methods save memory by the cost of more calculations. It stores a fewer number of activations and repeats the calculation of the rest from the saved checkpoints. Offloading methods [eaummot al. (2020) send the saved activations to the computer's 's RAM and load them back to the video memory on the backwards passes, which also saves GPU memory at the cost of host-device communication time.
ActNN Chen et al. (202I) is a framework for quantizing stored activations adaptively on the fly. work, it allows quantizing not only layers of glement-by-element activations but also many others, including convolutional, normalization, and linear layers. However, this method depends the distribution of elements of quantizable tensors and, because of that, its performance may degrade. Our approach, on the other hand, selects data-agnostic optimal quantization, which in practice turns out to be sufficient and easier use.
In contrast
CONCLUSION
We have proposed method to reduce memory consumption during the training of deep neural network models by storing less information for backward pass in the element- wise activation functions. For effective training, there is no need to calculate the derivative of the activation functions precisely, but only its piecewise-constant approximation is sufficient. This makes it possible to save not the entire input tensor each application the activation function, but only the interval number in the connstate approximation. Experiments show that for wide class of models and problems, storing only bits of information per tensor element does not lead to degradation of the learning quality and saves about 20 percent of memory. We have proposed an efficient algorithm for constructing an optimal piecewise-constant approximation The proposed drop-in replacements for popular activation functions (ReLU, GELU, Swish, Sigmoid and others) do not depend the neural network model, the problem to be solved, or the peculiarities of data distribution. The replacement of the original activation functions by the proposed method can be performed any training stage (both to models trained from scratch and to tre-trained models for subsequent fine- the-tunnng) and does not require any changes in the training pipelines. An efficient CUDA implementation the proposed method, together with pre-computed piecewise-constant aptp:ontttaapppxximimmmom approximations for many popular activation functions, is available for PyTorch GitHub repository
Source code repository can be found http:///ghhhbb../////oyooo////pp

---

Under review as a conference paper at ICLR 2023
REFERENCES
Banner, R., Nahshan, Y., and Soudry, D. Post training 4-bit quantization of convolutional networks for rapid-deployment. In Wallach, H. M., Larochelle, H., Beygelzimer, A., d' Alché-Buc, F., Fox, E. B., and Garnett, R. (eds.), Advances in Newral Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NewrIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 7948-7956, 2019, pttp/tpps///pprocceegg.....eseet
ccc paper/ ccppaper/ .tmm
Beaumont, O., Eyraud- I.., Herrmann, J., Joly, A., and Shilova, A. Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory. CoRR, abs/1911.13214, 2019. URL httlhhtt::// www.vaiivvv/..m///vv///aa//aaaaa wttp://////////111.111111./
Beaumont, O., Eyraud-Dubois L.,, and Shilova, A. Optimal GPU-CPU offloading strategies for deep neural network training. In Malawski, M. and Rzadca, K. (eds.), Euro-Par 2020: Parallel Processing 26th International Conference Parallel and Distributed Computing, Warsaw, Poland, August 24-28, 2020, Proceedings, volume 12247 of Lecture Notes in Computer Science, URL BtLLhhtss:///oii
pp. 151-166. Springer, 2020. doi: 10.
10.10077976-3--30--77777
Beaumont, O., Eyraud-Dubois L., and Shilova, A. Efficient combination of rematerialization and offloading for training dnns. Advances in Neural Information Processing Systems, 34, 2021.
Bengio, Y., Léonard, N., and Courville, A. C. Estimating or propagating gradients through stochastic neurons for conditional computation. CoRR, abs/1308. 3432, 2013. URL httt::ttttxx....oot
/008/ 3432
Bondarenko, Y., Nagel, M., and Blankevoort, T. Understanding and overcoming the challenges of efficient transformer quantization. In Moens, M., Huang, X., Specia, L., and Yih, S. W. (eds.), Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November, 2021, pp. 7947- 7969. Association for Computational Linguistics, 2021. doi: 10. 00.18863/v1/202111 10.1863/w12001.emmbbpp m2in.627. URL httpt: doi wttp:/0.....//1//000000
10.
emnl p- ma ww..a.n........m...m........mw.. 627
Chen, J., Zheng, I.., Yao, 7.,, Wang, D., Stoica, I., Mahoney, M. W., and Gonzalez, J. Acton: Reducing training memory footprint via 2-bit activation compressed training. In Meila, M. and Zhang, T. (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 1803- 1813. PMLR, 2021. URL http:://ppocceedeegeeeeoo mlr wwwppppsssssss.eee wttp://///.1////e/e22l..e htmll
Chen, T., Xu, B., Zhang, C., and Guestrin, C. Training deep nets with sublinear memory cost. CoRR, abs/1 hass/604.061744, 2016. httt:: wwwpsiivv.rrraaeoo abs /104. 06174
Cui, C., Zhang, K., Daulbaev, T., Gusak, J., Oseledets, I. V., and Zhang. Z. Active subspace neural
networks: Structural analysis and universal attacks. SIAM J. Math. Data Sci., 2(4): 1122, 2020. doi: 10.1137/1MM1296070 URL_attpp:///aa..c///1111113 01070///M11999797
Dettmers, T., Lewis, M., Shleifer, S., and Zettlemoyer, 8-bit optimizeer via block-wise quantization. CoRR, abs/2110.02866 2021. CtL.https: www.aaivvv..ooo org/ hbtp://b///11/ 0121111.0811111111
Gao, Y., Liu, Y., Zhang, H., Li, Z., Zhu, Y., Lin, H., and Yang, M. Estimating gpu memory consumption of deep learning models. In Proceedings of the 28th ACM Joint Meeting on European Software Engineering Conference and Symposium on the Foundations of Software Engineering, pp. 1342-1352, 2020.
Gusak, J., Kholyavchenko, M., Ponomarev, E., Markeeva, L., Blagoveschensky, P., Cichocki, A., and Oseledets, I. V. Automated multi-stage compression of neural networks. In 2019 JEEF/CVF International Conference on Computer Vision Workshops, ICCV Workshops 2019, Seoul, Korea (South), October 27-28, 2019, pp. 2501-2508. IEEE, 2019. doi: 10.1009/CCCW.2209900000 URL httpt:
htRL::/ts:///oo /wwp///oiiia/ www.doi..orrggio wwt.:1111111.. 2019 2019.00306 00306
Gusak, J., Daulbaev, T., Ponomarev, E., Cichocki, A., and Oscledets, I. Reduced-order modeling of deep neural networks. Computational Mathematics and Mathematical Physics, 61(55:774-785, 2021.
10

---

Under review as a conference paper at ICLR 2023
He, K., Zhang. X., Ren, S., and Sun, J. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pp. 770-778 IEEE Computer Society, 2016. doi: 10.1109/CVPR.22166.00. URL https: / doi. www.oo........... 1109/ /tt.///////2222 201 6. 90
Hendrycks, D. and Gimpel, K. Bridging nonlinearities and stochastic regalaaattcrrgullrizieeeet with gaussian 1666 0141111165600016
linear units. CoRR, (tbs//606.00/ 0000008415, 2016.
Hrinchuk, O., Khrulkov, V., Mirvakhabova, L., Orlova, D., and Oseledets. I. V. Tensorized embedding layers. In Cohn, T., He, Y., and Liu, Y. (eds.), Findings of the Association for Compu- tational Linguistics: EMNLP 2020, Online Event, 16-20 November 2020, volume EMNLP 2020 of Findings of ACL, pp. 4847-4861 Association for Computational Linguistics, 2020. doi:
10.18653/v1/2020. 1118665/w122000..ninnn-eemmmmmmm
URL https: http:///.ooorog/..../ 1tt..ooo/1111113311/22000000
Findings emnlp. 436
Jacob, B., Kligys, S., Chen, B., Zhu, M., Tang, M., Howard, A. G., Adam, H., and Kalenichenko, D. Quantization and training of neural networks for efficient integer- arithmetic-only inference. In 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018, pp. 2704-2713. Computer Vision Foundation / IEER Computer Society, 2018. doi: 10.1109/CVPR. 2018.00286. URL httpte::///peeacceeeee /owpaacpeeccssseeeeeeeepeep.ssse theevf com/ content cvpr 2018/html Jacob Quant izationiaiion and Training CVPR 201e_paperrra_____ html
Krishnamoorthi, R. Quantizing deep convolutional networks for efficient inference: whitepaper. CoRR, abs//a/6.o0ee 806.08342, 2018. UtL/http::/aaxxi.v.o owg.abss.rrroo./oo..ooooo..bbsss 180.. 08342
Lebedev, V., Ganin, Y., Rakhuba, M., Oseledets, V., and Lempitsky, V. S. Speeding-up up convolutional neural networks using fine-tuned cp-decomposition. In Bengio, Y. and LeCun, Y. (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL htLhhtt::/// htttt:://axxv..oo.
abs 1412.412222 6553
Leclerc, G., Ilyas, A., Engstrom, L., Park, S. M., Salman, H., and Madry, A. ffev. https 2022. commit XXXXXXX.
//githaab.co//llfffe///
Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., and Stoyanov, V. Roberta: A robustly optimized BERT pretraining approach. CoRR, abs/1907.11692. 2019. URL TRLLtttp///axxiv.oom www.wwwwww..a.aaaaaaaaaaaooaaaea wtw.ab////99..1.
Nagel, M., Fournarakiss M., Amjad, R. A., Bondarenko, Y., van Baalen, M., and Blankevoort,
A white paper on neural network quantization. CoRR, abs/2106.08295, 2021. URL h
//arxiv.org. abs/
waws/////... /2002 08295
Novikov, A., Trofimov, M., and Oseledets, I. Exponential machines. Bulletin of the Polish Academy of Sciences: Technical Sciences, pp. 789-797, 2018.
Ojika, D., Patel, B., Reina, G. A., Boyer, T., Martin, C., and Shah, P. Addressing the memory bottleneck in Al model training. arXiv preprint arXiv:2003.08732 2020.
Phan, A. H., Sobolev, K., Sozykin, K., Ermilov, D., Gusak, J., Tichavsky, Glukhov, V., Oseledets,
I. V., and Cichocki, A. Stable low-rank tensor decomposition for compression of convolutional neural network. Vedaldi, A., Bischof, H., Brox, T., and Frahm, (eds.), Computer Vision ECCV 2020 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XXIX, volume 12374 of Lecture Notes in Computer Science, pp. 522 539. Springer, 2020. doi: 10. 1007/ 1007 3-
URL pttps://da..oog//111100000.
31
Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Radford, A., Chen, M., and Sutskever, I. Zero-shot text-to-image generation. In Meila, M. and Zhang, T. (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. http:://prose ngs mm www.srss.ee.o vl9 213399/ame222112e .hmmmll html
PMLR, 2021. URL
11

---

Under review as a conference paper at ICLR 2023
Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M. S., Berg, A. C., and Fei-Fei, L. Imagenet large scale visual recognition challenge. Int. J. Comput. Vis., 115(3)211-252, 2015. doi: 10.10077/111003--11100888661 URL httpss://de http:/ddo..ogg///// .rrg/1......... 10000001.0000000011100000000
- 0816-y
Shonenkov, A., Bakshandaeva, D., Dimitrov, D., and Nikolich, A. Emojich zero-shot emoji generation using russian language: a technical report. CoRR, 2021. URL https://arxis: org abs 2112.02443
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, I.., Gomez, A. N., Kaiser, I., and Polosukhin, Attention is all you need. In Guyon, I,, von Luxburg, U., Ben- gio, S., Wallach, H. M.. Fergus, R., Vishwanathan, S. V. N., and Garnett, R. (eds.), Ad- vances Neural Information Processing Systems 30: Annual Conference on Neural In- formation Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998 6008, 2017. URL htw.apppp//pppreeennnnneeeei////
Wang, A., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. R. GLUE: A multi-task bench- mark and analysis platform for natural language understanding, In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL [tttttttt/:/: ew weet:/formmi..-rreeee.
12

---

Under review as a conference paper at ICLR 2023
A
DETAILED EXAMPLES OF FEW-BIT APPROXIMATIONS FOR POPULAR NONLINEARITY LAYERS
1.0
1.0
1.0
1.0
GBL
GULU
0.5
0.5
0.5
0.5
0.0
0.0
10
1.0
1.0
1.0
1.0
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
0
10
10
10
-10
10
110
10
0.2
0.2
0.2
0.2
denulive Darrigpee
0.1
0.1
She Sigment
0.1
Joe Sigendd
0.1
ahs Sigeed
0.0
0.0
0.0
10
10
10
10
100
10
110
0
10
1.0
1.0
1.0
1.0
Tan
Tal
0.5
0.5
0.5
Jefeanc S4rThe
0.5
3 Tan
A. Terk
0.0
10
10
100
100
10
100
0
10
SILU
SELU
1.0
1.0
1.0
1.0
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
10
10
10
110
10
Figure 11: 1- to 4-bit approximations of popular nonlinearty layers.
13

---

Under review as a conference paper at ICLR 2023
B DETAILED MEMORY MEASUREMENTS FOR DIFFERENT MODELS
We provide memory measurements for different model architectures in Table Appendix B "Model size" the total memory used for storing model parameters (without model gradients and optimizator statistics). "All activations size" is the total memory used by tensors, saved for backward pass. "Non- linearity activations size" is the part of all activations used only by nonlinearity layers. "Percentage saving" is memory saved on all activation using our method compared to full precision non- non-linearities, and percentage value in the "Maximum Batch Size" row is the increase the batch size achievable by using our method compared to full precision non-linearities, taken in ideal circumstances. Maximum batch size is calculated with the assumption, that four model copies are stored on the device (model parameters, model gradients and optimizer statistics like two moments stored by Adam optimizer) for GPU with 32G memory.
Standard
All
Nonlin. Act.
1-hit batch
2 bit batch
3 3-bit batch
bit Max batch
Model Size (Mb)
Nonlin.
Max
Size
Size
batch
(Mb)
(Mb)
size
size
size
size
size
ResNet-18
44.6
40.0
11.5
813
1 1127 (438.0%)
1113 (436.9%)
1100 (435 3%%
1086 (433.0%)
ResNet-50
99.2
156.8
47.9
(+42.2%%)
(+40.3%)
285 (438.35))
281 (436.49)
ResNet-101
171.4
234.5
136
(+44.11%)
193 (+41
(4:39.97))
188 (438.2%)
ResNel-152
232.3
328.2
104.9
97
140 (+44,33%
138 (442.3%)
130 (46)
134 (438. 1%)
DenseNet-121
30.9
243.8
79.1
133
193 (+46.66%
192 (+44.45)
189 (+92.11%)
186 (439.8%6 8%
DenseNet-161
1124
458.8
147.0
70
(+45.7%)
(+42.99%)
99 (441448)
(-38.669
DenseNet-169
109
(-45.9%)
137 (+44.0%)
155 (4-4..29
(439.49)
DenseNet-201
77.4
382.2
123.9
84
123 (+46.4%)
122 (445.2%)
120 (4492,956)
118 (+-10.5%)
Efficient Net B0
20.4
112.4
32.4
290
403 (+39.0%)
398 (+37.2%%)
393 (+35.5%)
388 (4333.8%)
Efficient B3
47.5
218.6
59.5
149
202 (+35.66%)
200 (+34.2%)
197 (432.2%))
195 (+30.99%
Efficient
256.3
(433.099)
(+31 (99))
(+29.895)
VGG II
507.2
100.9
37.0
304
(+55.3%)
464 (452.6%)
456 (450.0%)
448 (447.4%)
VGG 16
528.2
163.8
68.5
187
314 (66).
307 (+64.29)
301 (+61.0%)
295 (+57.8%%
VGG 19
548.4
178.8
75.0
171
288 (+68.4%))
28T (+64.3%)
275 (460.8%)
270 (457.9%)
RoBERTa-base
480.7
185.6
36.0
166
(+22.9%%
(422.35)
201 (421.15)
200 (420.35%)
RoBERTa-large
1355.6
482.1
(0.425..098
(+23.2%)
69 (423.2%)
(+21.4%)
GPT2
491.0
297.1
146.2
103
198 (+92.2%%
192 (486.4%)
187 (481,6%)
182 (476.7%)
14

---

Under review as a conference paper at ICLR 2023
C
NUMERICAL RESULTS FOR DYNAMIC PROGRAMMING
1-bit
2-bits
3-bits
4-bits
ReLU
0.0
GELU
0.1410
0.0406
0.0119
0.0031
Swish
0.2150
0.0479
0.0170
0.0045
Sigmoid
0.0181
0.0038
0.0009
0.0002
Tanh 0.1584
0.0319
0.0073
0.0017
SELU 0.2554
0.1010
0.0184
0.0039
Softplus
0.2902
0.0541
0.0121
0.0029
Table 3: Numerical values of error Eq.
(3) with uniform weight on interval [-10; 10).
15

---

Under review as a conference paper at ICLR 2023
D EXPERIMENT SETUPS
D.1 GLUE
Benchmark implementation is based on opensource Huggingface implementation and is available at /ttttss://gttiii com ancnyymouss/rppssttoo..
The following parameters were used:
Batch
Learning
Number of epochs
Warmup
Task
Size
rate
length
Cola
32
0.00002
10
320
MNLI
32
0,00001
10
7432
MNLI-MM
32
0.00001
10
7432
MRPC
16
0.00001
10
137
QNLI
32
0.00001
10
1986
QQP
32
0.00001
10
28318
RTE
16
0.00002
10
122
SST2
32
0.00002
10
1256
STSB
16
0,00002
10
214
Common parameters are:
Parameter
Value
Optimizer
Adam
Adam B1
0.9
Adam 32
0.98
Adam
le-6
Weight Decay
0.1
Float Precision
fp16
D.2 RESNET
We use open source FFCV Leclere al. 2022 Imagenet benchmark with ResNet18 parameters for A100 Nvidia GPU https:// com http:///ff.f//eeeee.ome/////////
rn18
88 hep_aachhs__pp_hh_______ yaml
D.3 RUDALL-E
We used open source implementation that can be found https://github.com sberban-aa/zu-Ga11.
All experiments have following setup: training size 2474, valid size 275, loss image weight 1000, frozen MLP and attention layers, batch size 40, start Ir 4e-7, max Ir 1e-5, final Ir 2e-8, warmup 0.1. 8bit-Adam Dettmers Dettmmeemmtttessseeemm et al. (002)) weight decay 0.2, betas (0.9, 0.98), eps 1e-6, gradient checkpointing 24, trained for 6h using 1xA100.
"huggingface.co
16

---

Under review as a conference paper at ICLR 2023
E COMBINATION OF ACTNN AND FEWBIT
ActNN method is more general and can be applied to the broader class of layers, while our method only focus on one class of layers - pointwise nonlinearities. In the cases when it is not enough and more memory saving is required it is possible to join these two methods and to use Fewbit for pointwise nonlinearities and ActNN for everything else. Such a combination should work better than pure ActNN, since Fewbit works better than ActNN for pointwise nonlinearity layers. To check this hypothesis we train R on CIFARTO dataset. We replace standard ReLU pointwise nonlinearity with GELU, compress all layers except GELU with 4-bit ActNN (since 2-bit ActNN is too much of a compression and model diverges) and GELU layers are compressed with either 2-bit ActNN or 2-bit Fewbit. On Fig. you can see training loss and accuracy. ActNN Fewbit for pointwise nonlincarities works slightly better than pure ActNN, as expected.
ActMN + Faaflii2 GEL
3.0
ActhN ACDIN-3 GELII
0.9
2.5
0.8
Training Accuracy
Training Loss
0.7
2.0
0.6
1.5
0.5
0.4
0.5
AcMi-4 Tewtik-2 GEL
0.3
ACtiN-4 ACtiN-2 GEL
Standard Network
1000
2000
3000
4000
1000
2000
3000
4000
Step
Step
(a)
ActhN4 2extii
0.96
ActHW T
0.45
Standard Network
Standard Metwork
0.94
Trai ning Accuracy
0.92
Training Loss
0.35
0.30
0.90
0.25
0.88
0.20
0.86
0.15
0.84
2000
2500
3000
3500
4000
2000
2500
3000
3500
4000
Step
Step
(b)
Figure 12: ResNet18 CIFARIO dataset. All ReLUs replaced with GELU. All layers except pointwise nonlinearities compress their activations saved for backward with 4-bit ActNN, GELUs compress their activations saved for backward with either 2-bit ActNN (orange) 2-bit Fewbit (blue). ResNet18 without any compresssion is depicted with green. (a): Training loss and accuracy for the whole training course. (b): Training loss and accuracy zoomed to the last half of the training course. ActNN + Fewbit for pointwise nonlinearities works slightly better than pure ActNN.
17

---

Under review as a conference paper at ICLR 2023
F
MORE PLOTS FOR EXPERIMENTS
x 10
Seik
0.7300
Syth
116 Si.sk
2.6
1-8 Smik
Fru Seiik
242 Si-sh
242 Smisk
0.7275
x 10
2.5
0.7250
0.908
2.4
0.7225
0.906
2.3
0.7200
70
80
1-bir 2-bit
3-bit
4-bit
1-bit
2-bit
3-bit
4-bit
Epoch (a) Training loss
Epoch
(c) Top-1 accuracy on validatasfffop-5 accuracy on validation
(b) Training loss
x 10
101
SELL
2.7
Fow SELL
0.902
2.6
10
0.716
2.5
SRT 21
0.901
0.714
2-bit
3-bit
4-bit
I-bir
3-bit
4-bir
Epoch (d) Training loss
Epoch (e) Training loss
(f) Top-1 accuracy on validattodfop- accuracy validation
x 10
CBL
Toe CREU
2.6
0.728
0.912
CIRAU
tie CBS.U
0.726
x 10"
2.3
0.910
0.724
2.4
CRLA
0.908
20NCELU
0.722
2.3
70
80
3-bit
4-bit
Epoch
Epoch
(i) Top-1 accuracy validatioffop-S accuracy validation
(g) Training loss
(h) Training loss
ResNet18 with ReLU replaced with Swish, SELU and GELU nonlinearity trained on Imagenet. (a): Training loss. (b): Training loss zoomed into the last third of the training. (c): Final
Figure 13:
validation top-1 I accuracy. All graphs are averaged across three runs with different seeds. Error bars denote minimum and maximum values.
18

---

Under review as a conference paper at ICLR 2023
0.9125
GELU
Top-I accuracy
Top-5 accuracy
Training loss
Few-bit
0.725
0.9100
2.25
AcIN
0.9075
0.720
2:20
GELU
0.9050
GELU
Few-bir
Few-bir
0.715
AcINN
0.9025
AcINN
2.15
1-bits
2 bits
3-bits
4 bits
1-bits
bibis
3-biss
4ibis
1-bits
bits
bits
bits
(a)
(b)
(c)
SELU
Training loss
Few-bit
Top-I accuracy
Top-5 accuracy
0.900
ActNN
0.895
2.3
0.70
SELU
SELU
0.890
Few bit
Few hit
2.2
0.69
AcINN
AcINN
1-bits
2.bits
3-bits
4-hits
1-hits
2-bits
3-bits
4.bits
1-bits
2-hits
3-bits
4-hits
(a)
(b)
(c)
Swish
2.25
Top-1 I accuracy
Top-5 accuracy
Training loss
Few-bit
0.725
ActNN
2.20
0.720
Swish
Swish
Few-bit
Few-hit
0.904
2.13
ActNN
ActNN
0.715
1-bits
2-bits
3-bits
4-bits
1-bits
2-bits
bils
4-bits
1-bits
2-bits
3-bits
4-bits
(a)
(h)
(c)
Figure 14: Comparison of ActNN GELU, SELU and Swish with Few-bit GELU, SELU and Swish (Our) for ResNet18 architecture on ImageNet dataset. (a) Training loss. (b) Top-I accuracy. (c) Top-5 accuracy. Our method with 1-bit already matches unaltered nonlinearity performance and significantly outperform 1-bit ActNN.
19

---

Under review as a conference paper at ICLR 2023
G DYNAMIC PROGRAMMING
is easy to see that the optimal value of y for L(s in Eq. (3) with given s is:
fa (s(s) = Jo w(x)dx
(8)
Consider Eq. (7) both y(j. i) and T(j can be calculated in advance using analytical formulas (if possible) numerically for the corresponding 1-dimensional integrals. After that, the full array of (P( k) can be calculated On K) time and O(n²) space, where K is the required number of constant intervals in the approximation Eq. (2). Please note that this optimization has to be performed only once, so n can chosen quite large thus the result would be very close to the global minimum.
Note that the space complexity can be reduced to Onz) by adding three auxilliary arryys F2, W and FW and rewriting Eq. 7
FO(1) '' 2
((())((z)dz
w(1) SAA w(x)dr,
(9)
(W I'A (()x)w(x)d
v(j (FW() - FW(O)/(W()) W(O),
TÚ. - - ( v (W W
We can see that ultimately only (n) one-dimensional integrals have to be stored, and everything else can be easily evaluated in (1) time on the spot. The -ieedimnsioonal integrals can be calculated numerically in On)) time and space complexity well:
1(++1)== F²)) J. ((())((o)d,,
= 1 w(z)dz, ([
(10)
(il1 ((1) f((x)((z)d.
Numerical results. In Fig. T we provide some 3-bit examples for popular activation functions obtained with described method, and more fewbit approximations can be seen in Fig. II In Table we provide numerical values error Eq.
20


Reviews:
[{"review_id": "kpsmDWGxHG7", "reviewer": "Reviewer_Wx7m", "summary": "", "questions": "", "limitations": "", "rating": 3, "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 proposes a method to perform activation compressed training (ACT) for pointwise non-linear activation functions. Instead of storing a quantized version of the input x, the proposed approach stores a quantized version of the activation function's gradient f'(x), which is multiplied with the gradient in the backward phase. The paper proposes a nonlinear quantization method for f'(x) with a dynamic programming method to determine the optimal quantized approximation. Quantization can be reduced to table lookup then. The proposed method is evaluated on language model pretraining, text-to-image generation, and image classification tasks, while the proposed method is slightly better than existing ACT methods under the same bitwidth.", "summary_of_the_review": "The paper deals with an important problem in a reasonable manner. However, the paper might require more work to fully demonstrate its significance before publication.", "strength_and_weaknesses": "Strength:\n- Reducing training memory footprint is an important problem.\n- The proposed method technically sounds, and should be the correct way to deal with nonlinearities. \n\nWeaknesses:\n- The paper still needs polishing. For example, the introduction is too brief.\n- The improvement is incremental.\n- It is somewhat unclear whether the comparison with ActNN is fair. ActNN compresses both the linear and nonlinear layers. Does the proposed method also do so? (i.e., combine the proposed method for nonlinear layer and ActNN for linear layer. I think the combined result is what of practical interest.)\n- Table lookup can be very expensive. The time consumption is not reported. I'm not sure if the proposed nonlinear quantization is practical.", "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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "_StZphCD9n7", "reviewer": "Reviewer_1fcj", "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": "The paper proposes a method for reducing the memory needs of the backward pass of a neural network. It does so by quantizing the gradients of the activation functions meaning that they now only need a few bits to store as opposed to the standard 32-bit floating point representation.\n\nThe few bit representations are computed to minimize the l2 distance from the exact gradients. The paper proposes an algorithm based on dynamic programming to compute them. This algorithm only needs to be run once for each bit count/activation function pairing.\n\nThe experiments confirm that these few-bit representation do not significantly affect the training dynamics so the memory reduction has no downside.", "summary_of_the_review": "The paper is well written and good quality. My main issue is the thoroughness of the experiments. I would like to see the total memory reduction measured at runtime and it would also be good to see comparison against low-precision approaches (that use low-precision throughout the network not only the gradients of the non-linearities).", "strength_and_weaknesses": "The algorithm for finding the quantized representation is quite clever. Its makes good use of the dynamic programming principle. I wonder, however,\n* Is l2 loss the right metric to optimize? The goal is to minimize the bias that we are adding to the gradient updates and perhaps a different metric, maybe one that is optimized to minimize this bias would perform better.\n* Is dynamic programming needed here? Could we simply optimize the values and the beginnings/endings of the segments using a simple optimizer?\n* What is the resolution of t used in the experiments (and what interval [A, B])?\n\nThe paper makes a convincing case that the proposed method indeed reduces the memory footprint of the gradients and that this reduction does not significantly impact the training dynamics of the network. It would be nice to show concrete benefits of this memory reduction.\n* As stated in the paper, the memory reduction does not result in faster backward passes. Instead, it enables training with larger batch sizes. Does training with larger batch sizes then result in a speedup?\n* In the Abstract, it is stated that \"Memory footprint is one of the main limiting factors for large neural network training.\". Does it enable the training of larger networks that have better accuracy?\n* In the conclusion, it is stated that the method results in a 20% reduction in memory usage. How is this 20% computed?\n* In Appendix B, we see the potential increase in batch size. Is this table generated by measuring the size of the model running on a GPU or is it calculated according to the formula in Section 2? There may be significant differences in the theoretical memory footprints and the real memory footprints of these models.\n* The paper should also compare against low-precision neural networks, for example 16-bit floating point networks. It is claimed that the savings are `complementary and can be used together', but this is not fully true because the savings don't combine multiplicatively. The savings in a low-precision network are a smaller proportion of the overall memory cost than in a high-precision network.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "GQaTRw0V6F", "reviewer": "Reviewer_KDUc", "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 paper proposes to use point-wise activations to reduce the memory cost of back-propagation. An analysis is proposed to determine quantization levels minimizing a quantization MSE metric. Experimental results using GeLU and similar activations using several bitwidths are cometitve.", "summary_of_the_review": "The presented work is an interesting approach to code-book quantization of activations. The paper is well written and has a satisfactory amount of details. However, I have some concerns regarding the novelty of the proposed quantization scheme.", "strength_and_weaknesses": "The idea is interesting however, I have a couple of concerns with the analysis.\n\n1) The presented analysis in eq. (3) - (7) is not novel. It is simply re-deriving the famous Lloyd-Max algorithm [1].\n[1] Lloyd, S. P. \"Least square quantization in PCM. Bell Telephone Laboratories Paper. Published in journal much later: Lloyd, SP: Least squares quantization in PCM.\" IEEE Trans. Inform. Theor.(1957/1982) 18 (1957): 5.\n\n2) How does the method compare to low-precision integer quantization, which can also be optimized (see [2])? Unlike the proposed codebook quantization scheme, low-precision integer representation can use ultra-fast tensor cores on GPUs. Can the authors provide some comparisons taking that into account?\n[2] Sakr, Charbel, et al. \"Optimal Clipping and Magnitude-aware Differentiation for Improved Quantization-aware Training.\" International Conference on Machine Learning. PMLR, 2022.", "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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject