Softmax Output Approximation for Activation Memory-Efficient Training of Attention-based Networks
Authors: Changhyeon Lee
Seulki Lee
Venue: NeurIPS 2023 poster (2023)
Keywords: Memory efficient
Activation saving memory
NLP
Transformer
TL;DR: Memory efficient softmax output approximation method.

Abstract: In this paper, we propose to approximate the softmax output, which is the key product of the attention mechanism, to reduce its activation memory usage when training attention-based networks (aka Transformers). During the forward pass of the network, the proposed softmax output approximation method stores only a small fraction of the entire softmax output required for back-propagation and evicts the rest of the softmax output from memory. Then, during the backward pass, the evicted softmax activation output is approximated to compose the gradient to perform back-propagation for model training. Considering most attention-based models heavily rely on the softmax-based attention module that usually takes one of the biggest portions of the network, approximating the softmax activation output can be a simple yet effective way to decrease the training memory requirement of many attention-based networks. The experiment with various attention-based models and relevant tasks, i.e., machine translation, text classification, and sentiment analysis, shows that it curtails the activation memory usage of the softmax-based attention module by up to 84% (6.2× less memory) in model training while achieving comparable or better performance, e.g., up to 5.4% higher classification accuracy.

Softmax Output Approximation for Activation Memory-Effcient Training of A Attention-based Networks
Changhyeon Lee and Seulki Lee
UNIST (Ulsan National Institute of Science and Technology) [changhyeon,seull seulki. lee) @unist.cc. ac..tt........ ac. kr
Abstract
this paper, we propose to approximate the softmax output, which is the key product of the attention mechanism, to reduce its activation memory usage when training attention- based networks (aka Transformers). During the forward pass of the network, the proposed softmax output approximation method stores only a small fraction of the entire softmax output required for back pooo-propagation and evicts the rest of the softmax output from memory. Then, during the backward pass, the evicted softmax activation output is approximated to compose the gradient to perform back- propagation for model training. Considering most attention-baa models heavily rely on the softmax-based attention module that usually takes one of the biggest portions of the network, approximating the softmax activation output can be a simple yet effective way to decrease the training memory requirement of many attention- based networks. The experiment with various attention-bart based mod- els and relevant tasks, i.c., machine translation, text classification, and sentiment analysis, shows that curtails the activation memory usage of the softmax-based tention module by up to 84% (6.2x less memory) in model training while achieving comparable or better performance, e.g., up to 5.4% higher classification accuracy.
Introduction
Recently, many attention-basee networks, aka Transformers 50 widely used in many natural language processing tasks such as machine translation [50
13,, sentiment analysis [52 29 34 text
classification [52] [54, and question/answering [47 53 20 Due to their superior performance, they also have become the backbone of many large language and general task-agnostic models, e.g., GPT series 5), and are actively expanding other domains ranging from computer vision (Vision Transformers
4 338) to protein structure predictions (AlphaFold AM (2]). However, the stellar performance of intention-baseed networks comes with significant resource requirements in training, especially massive memory (RAM) usage. For example, the small-size GPT-2 model [41 is known to require 27.5 GB of peak memory for training with a mini-batch of size eight Hence, the memory bottleneck problem makes training of attention- based models challenging when massive memory capacity is not available, e.g., BERT takes 16 TPUs (1], each equipped with TB of built-in memory, for training, and the regular GPT-2 [1 is reported to be trained with 32 TPUs.
During deep model training, it is required parameters, input data, optimizer state, etc., along with activation output, i.e., the intermediate output of each layer. As the activation output of each layer must be stored in memory for gradient-based back-propagnn propagation [42). the memory capacity required to train model increases proportionally to the data sequence length (width), depth of the network, and mini-batch size of the model. Consequently, as the model gets huge and the training data gets bigger, the activation easily becomes to take up more memory than other components. Especially, is more evident in NLP (natural language processing) tasks, as provided in Figure| ((a) 22], which shows the memory composition required for training four different attention- based model configurations, ranging from 22 billion to trillion parameters.
various components in memory, i.e., model
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

For GPT-3 [S] that has 175B parameters, the activation output requires approximately 1.45x more memory than the parameter and optimizer state, and for MT-NLG [45] that has 530B parameters, is about 3.6x more. Therefore, to popularize such large models in practice, it is necessary to reduce the massive amount of activation memory required in model training.
Optimizer
Memory usage during training process
<Refore>>
<After>
MI-NLG
160
Other layers
Optimizer state
Optimizer state
Activation memory
Activation memory
Memory (GB)
120
GPT-3
Attention besed model
Attension mechanism
Softmax
MatMul
Softmax activation
Oth er activation
Softmax activation
Other activation
132 YG2
Mask (opt.)
Dutput [5 tiy (1
Output [o Mugit (
80
"11002
Scale
MatMul k
Input Pey Xe 2
Input (altihee
40
Madel parameters
Model parametres
16.3G
Lincar
Input data
Iaput data
31.388
B Batch size JF: Number of multi head atlention
226
1758
5308
1T
Input data
Saving for backward process
pir seralelle andd uptinion siale manuay activation memory
B
Suftmax output approximation
(a)
(b)
Figure 1: (a) The amount of memory required by 1) parameters and optimizer state and 2) activation
memory 25 in training attention-based models (GPT-3 [5] with 22B and 175B parameters and
MT-NLG 45 with 530B and 1T parameters). The dashed red line represents 80GB of memory capacity available in NVIDIA A100 GPU [35]. (b) The detailed memory composition when training an attention--asee network (blue boxes). The softmax output approximation reduces the activation memory usage by storing only a small subset of them (e.g. 81, 8n, Ot, (3) in memory (red boxes).
In this paper, we propose an approximation method of the softmax output [ produced in the attention module (50) to reduce its activation memory requirement in model training. as shown in Figure [(b) Instead of storing the entire elements of the softmax output vector that represents attention scores, it stores only small subset of the elements and evicts the rest from memory during the forward pass of the network. Later, during the backward pass, the discarded softmax output elements are approximated to construct the gradients for back- preck--ropggaiinn. As only a small fraction of the softmax activation output is stored during the forward pass, a substantial amount of memory can be saved. The proposed method is not to efficienate the computation complexity of the attention module, such as 99 48 3 [I], but to reduce the amount of activation memory used in training by storing only selected parts of the attention score, i.e., softmax output, in memory, not all of them. Thus, it is applicable to softmax output computed from any form of the attention module 9 48 31
[T] in
model-agnostic way. To the best of our knowledge, this is the first and unique attempt to curtail the activation memory usage of the attention module in training by approximating the softmax output.
The efficacy of approximating the softmax output is substantial because is the key component of the attention block used to compute the normalized attention score [48). Since many attention-based models repeatedly use multi-head attentions in their encoders and/or decoders [50). the softmax output can be accounted as the dominant activation in the network. For instance, the softmax operation, i.c., Softmax (QKT/ /dd), takes up 80% of the attention module itself, and 64.7% and 62.2% of the entire layer activation output during the forward pass of the classic Transformer [50] and BERT [I]] model, respectively. Therefore, streamlining the activation memory requirement of the repeatedly used saftmax-based attention module is a simple yet effective way of improving the training memory efficiency of attention- based networks, e-g., the reduced peak-memory usage and memory footprint.
The experiment results demonstrate that the softmax output approximation drastically decreases the training memory required by the attention module by up to 84% while achieving competitive or im- proved performance of various attention-based models, i.c., basic Transformer [50), XLNet (52). and
ALBERT 30 e.g., up to 5.4% higher accuracy, on machine translation 13 text classification [54] and sentiment analysis [34 46 tasks. It validates that the softmax output approximation 1) can be readily applied to any existing attention- based models without modifying the network architecture. 2) does not significantly increase computation time for activation memory saving, and 3) enables models to learn longer sequences with a larger batch size when given the same memory budget. We implement the proposed softmax output approximation with Py and make it public"
https: wttp:///ee.../.ottttttttuuttrarp

---

2 Training of Softmax Operation
This section discusses the training procedure of the softmax operation and its memory usage.
Softmax. Let 5 = [21, #2,.... 2n be a vector in R1x=, where n denotes the sequence length. The vector is obtained from the rows of Z E R"x", which is the result of the query/key multiplication (Z=QKT/ /vd)). Then, the element s; in the softmax vector [$1, $2, ...,...,, Sn is defined as 4
for all i = 1,2, n S.L. 0 S 81 1 and Σ
=1
Given the definition of softmax in Equation [ let us assume there is a I-layer DNN (deep neural network) having softmax operation at the k-th layer (k < l). To train the network via the gradient descent [2], the gradient of the loss T.(z) is computed with the input w.r.t. the weight matrix W... m.
i.e., 0T (T) /OW my at the m-th layer of the network (m < k). By using the chain rule, we get:
OL(E) OL(T) 03 0E θz
05
O
(2)
aw m
aw m
where म is the output of the 1-th layer, z = 5 and - in the middle is the output and input of the softmax operation at the k-th layer, respectively. Thus, to get @L(0//ow.. the gradient of 3 w.r.t. z i.e. 05/05_11, should be computed.
Gradient Matrix. To obtain 05/0% in Equation 2 the partial derivative of the i-th softmax element 8; w.r.t. its j-th input element 2ą 1.j is computed as follows:
Ds: s (Li=) - sy)
Ūzk-
where liny is the indicator function. Since it only involves and sj. not Ek-1, the gradient matrix of the softmax vector w.r.t. its input vector žk-1 1 given by:
11 - 81)
-182
81 8n
0s
8281
82
82)
828n
05k-1
81
82
s,,(1 - -s,)_
which allows the back-propagation of @L(F//OW in Equation 2 if 0 for all / is available.
Memory Requirement. Thus, to compute 05/0žz-1 in the back-propagation step. the softmax vector output needs to be stored for each input H as activation in memory during the forward pass of the network. Considering that 1) many DNNs, e.g.. Transformers [50] and convolutional networks (26 18), have deep and wide layers, and 2) propppropagttion is usually performed with multiple examples (mini- batch, the total amount of activation memory required in training increases linearly to the number and length of softmax layers and the mini-batch size. For example, network with softmax layers of length n trained with mini-batch consisting of b examples requires a total of b of activation memory for single step of back-propagatiion.
Memory Saving, Motivated by this observation, we reduce the amount of activation memory required to back-propagate through softmax layer by storing only small subset of softmax output elements = [$1, 52, 5n] in memory, while evicting the rest of the elements from memory during the forward pass of the network. During back-propagation the evicted elements approximated such that the gradient matrix (Equation 4) obtained from the approximated softmax output minimizes the expected loss of training performance. The amount of activation memory saved by the softmax output approximation is multiplied by the number of softmax layers / and the mini-batch size b, i.e., L. (n- m) b, where n is the total number of softmax output elements (the length of the attention score), and m is the number of elements to be stored in memory. Also, the input of the softmax, E, i.e., query/key multiplication QKT / /vdz, which is the output of the previous layer, be obtained from the approximated softmax output from the inverse of Equation as:
2 In(s;) + Ind Σ
(5)
Thus, Z of length can also be removed from the memory, which saves additional memory of L L.
3

---

3 Softmax Output Approximation
3.1 Approximation Strategy
Given softmax output vector 5 E [0, 1J", its approximation, 8", is computed with the following steps. First, among all softmax output elements of s, the highest and m lowest elements are stored in memory during the forward pass of the network. Then, the remaining n 2m elements are discarded from memory and approximated later in the backward pass. Lastly, using the stored and approximated elements together, the gradient matrix (Equation 4 is constructed for back propaaaoonn (Equation 2
Approximated Softmax. Given the softmax output vector 5, its approximation s' is defined as:
g = hi h22 ... Mmm 81 62...--221
lm lm-11 ... &
m highests n- 2m approximations
XYL lowests
where hi, f22, .... hm and lm fm-1s... are stored intact in memory, and 11 8/2 ... 8 non-stored elements, $1, $2, .... 8n 2m- Figure depicts an example of softmax output approximation.
.... l are the m highest and m lowest elements, respectively, that
are the rest n 2m elements that approximate the
2m-
Among total of n=21 soft- max output elements sorted in descend- ing order, the m=3 highest elements (h1, h2, hy) and the m=3 lowest ele- ments (lz, f2, E1) are stored in mem- ory. The remaining n-2m elements (8),..,, 22 are not stored but ap- proximated by the softmax output ap- proximation function described below.
Figure 2:
Softmax output
Approximated softmax output
Softmax value
Stored m high elements
Stored elements
R4 Rg. Rig
Approximated 2m elements
0.1
Sn 2m
2m
20
3.2 Element Selection for Memory Storing
Keeping the m highest and lowest softmax output elements in memory minimizes the possible loss of training performance caused by the softmax output approximation. The following shows doing so minimizes the difference between the gradient matrix constructed from the original softmax š and the gradient matrix constructed from the approximated softmax 3", so that the weight update computed from the approximated softmax becomes close the original ones.
Gradient Error. Given a n x n softmax gradient matrix 05/0 in Equation 4 and its approximate gradient matrix 0s' / obtained from the approximated softmax vector s the gradient error between them is defined as (5,
8 BE We first decompose the gradient error c(5, s
0s
into each i-th row-wise eπror (5 5²) consisting of 1) i-th column element's 's error and 2) the sum of the remaining j # i column element's errors in the i-th row. Then, the row-wise error c, s') for >i= <x(5, 87) as:
all rows are summed up to obtain the total gradient error, i.e., (8,55)
72
c($
(3,55)
(1) s( $
| ((,5)1
i-th element's error
11
o i elements' error sum
Error Bound. The second term Σ}=1, |8;8; 5 in Equation
is calculated in two cases. The
set S comprises indices where 8,8j 5/58/ for j # positive, while S2 consists of the opposite.
|s,(1 (1) <1(1 - +1)1 + + +2 jes, 558;-665 if s (1 8;) < s( [s,(111 si) <((1)()) s1) - 2 Ejes, -18--866 d/ol
(8) If s,sy - S/S5 for all # i belongs to either S₁ or Sz, the second term of Equation is bounded by its maximum value. In the case of S,, it becomes:
Σ (185)-818]1
if s1 si) 2 (1-8)) x))
2 > 81S, - 2(s:
2 = 2[ss(1 (s) 54(1 (1- x() 1 0.5
(9)
JES
JES,
JES
We denote the second term of Equation becomes to be bounded by:
8
as 8. From Equations
and 9 the gradient error (5,,5)
(55,6)) = Σ s:(1-8.) <((1-()) +8 max(s:(1-a;) s:), 0.25 - 0.5
(10)

---

where the inequality comes from the property of the softmax, i.e., 0 S s;(1 - s;) S 0.25.
Error Minimization. Hence, to minimize the gradient error c(5, 8*) in Equation 10 select set of s; as close as possible to some constants, i.c. {0, } for the first term, s; - 8;), and 0.5 for the second term, 0.25 s; 8.). First, by keeping the highest m and the lowest m elements in 6, can select elements closest to the first two constants, {0, 1). Second, similarly, when selecting the highest element h from 6, it becomes the one closest to 0.5 among all softmax elements since the remaining n-l elements are at most 1 1-h, from the properties of softmax (Equation D. Thus, their distances from 0.5 become least 1h, 0.5], i.e., |0.5 h4| = |0.5 (1 hn) > /// 0.5) for all i A 1. It shows that selecting the m highest and m lowest softmax elements minimizes the gradient approximation error, c(5, 8"), leading to the minimization of possible performance loss in training.
3.3 Approximation Method
During the back- propaaaapat the elements in the middle of Equation 6 i.e., $1. S 2 approximate the corresponding original elements, $1, $2, ... 5m-2m that are evicted from memory.
Exponential Distribution. We use the exponential distribution 443 i.e., λe- for the softmax approximation because 1) it naturally models a softmax in descending order well, 2) it can properly represent the non-linear changes of softmax its slope changes rapidly; its derivative is -X2e-A and 3) it can be efficiently computed with a small number of parameters. Unlike the original exponential distribution using single parameter X, we introduce two separate parameters a and / to better represent the characteristics of the softmax operation in our approximation.
Approximation Function. Given the non-stored original softmax output elements sorted in de- scending order as $1 2 $2 2 2 8n-2m that are computed but discarded from memory during the forward pass of the network, its k-th element 8y is approximated as s in the backward pass as:
= ae 8- for k = [0, 1, 1,.,, N --11
(11)
where N = n 2m + 2 is the number of softmax output elements to be approximated, including hm and fm the first and last reference values of the approximation, i.e., we let s hm and N'N = Em The two parameters a and 33 are given by:
Am
B
1
1
In fmm
+ 2mmmmmmmmmm -mm+t)))
(12)
N
1
1-hm
Em
s for k = O, 0,1, - can be computed to construct the gradient matrix in Equation 4 during pack-propagation
By using Equations 11 and 12
while having only 2m values in activation memory, i.e. the m 1 highest and m - 1 lowest softmax elements plus the two parameters and 3.
As mentioned above, the two parameters a and 3 are defined such that the first and last approximation value becomes to be the same as hm and lm. respectively, i.e., s = hm and 8 Ln- The parameter a can be easily obtained as a = hm since s/ = Qx 5-0 hm- The parameter 3 consists of three components, i.e., 1) normalization factor 1 2) range value o, and 3) rate value 3 as: hm and Ammt
3 = N 1m a. where
In 2
Ii -mmm1
(13)
Range value (o)). The range value a in Equation 13 is determined to map the domain of the approximation function, i.e., k [0, N -1], to new domain k' = v'000,, [0. Since the continuous domain of exponential distribution ranges from 0 to xc, we define a new finite range of domain from which the corresponding approximations (codomain) set the slope of the softmax approximation function in Equation IT at A = - 1 as symmetrical domain given an arbitrary length of softmax.
computed. To determine the range σ,
to the slope at k = 0 to provide the length-invarian Given the differentiation of & = ar B-k w.r.t. as 0sk/0k = - -0,3 8k the slope at 4: = 0 becomes fok = -a 8.0 = a,3. By making the slope at ie = N 1 as symmetrical to -a,, i... 00/y-1/0k= = -ase (N--1) = a3. replace N with new domain a,, the range a becomes:
1 1
σ = 3
22
(14)
5

---

Algorithm 1: Model Training with Softmax Output Approximation
Input: M: 1-layer model, B: training mini-batch, m: number of softmax elements to be stored in memory for k <- 1 to I do [forward pass of the model M] execute the operation the k-th layer with B:
k-th layer == seftmax then
foreach training sample in B do
① store m-1 highest and m-1 lowest softmax output elements and 3 [Eq 12 in memory;
discard the remaining n-2m+2 softmax output elements from memory; compute and store In( ( e²x) [Eq 5 in memory;
end
clse
store the activation output of k-th layer obtained from B in memory; end
end
for k + to do [backward pass of the model M]
if 1--th layer == softmax then
foreach training sample B
2 approximate n 2-2mm++2 softmax output elements using a and |3 [Eq make the softmax output vector 5' from D. 2 [
make the softmax imput vector ž from the softmax output vector & [Eg
make the gradient matrix 05' - from the softmax output vector Eq
end end
perform hack-propagation for the k-th layer with its activation output;
end
Normalization Factor. As shown in Equation 13 the range is normalized by 1 1/N, as a total of N softmax elements are approximated from it so that the discrete domain of 1: = [0, 1, 1..., -1] can be used approximate the k-th element. Since the approximation function in Equation [T converges to O when k + co, a is additionally multiplied with 1/(1 - ax) where & = /m (Fquation (2) to increase the range inverse proportional to hm based the property that the of a softmax equals one.
Rate Value (δ). The rate value 3 in Equation 13 determines the shape of the approximation function, playing a similar role to in the exponential distribution. To determine 3, the approximation a is set as the same as fm given Equation and and 14 i.e., we let K = fm and obtain 3 as:
B.
3
from b we
o, fm
(15)
Iim km
where a = hm as in Equation 12 Since the approximation is based on the exponential distribution, the convergence rate is proportional to 3. To utilize this characteristic, the rate value / is multiplied by the distance between the hn, and hm-1 in the exponential function, which given by:
A hm41
(16)
hmmem
By replacing 33 with 3 in Equations 14 and to distinguish them from 3 the left-hand side of the first equation in Equation [3 they become equivalent to the right two equations in Equation
1 Multiplying N 1-hh,. a. and 3 altogether in Equation 13 gives the final /3 parameter in Equation 12
Approximation Overhead. The proposed softmax output approximation 1) sorts the softmax elements, 2) performs the approximation, and 3) reorders the approximated softmax elements back to the original sequence, which takes O(n logg n) of time complexity and n log, n bits of space (memory) complexity, respectively, where is the length of softmax vector. However, these overheads are negligible in general when compared to the massive amount of the training workload.
3.4 Model Training with Softmax Output Approximation
Algorithm I illluttrasss the model training procedure that uses the proposed softmax output approxi- mation. Although its primary target is attention-basee models (Transformers), it is applicable to any models having the softmax layers in them without modifying the network architecture. It reduces the amount of the softmax activation memory without affecting other parts of the network.

---

4 Experiment
Implementation. We implement custom softmax operation that applies the proposed softmax approximation (Algorithm() I with PyTorch [39] Since our implementation is fully compatible with the existing normal softmax operation, it can be easily applied any model that uses the original softmax by just replacing it with our custom softmax without changing any other model components.
Tasks and Models. As the attention mechanism is being widely used and developed primarily in the field of NLP, we evaluate the proposed softmax output approximation with various NLP tasks, i.c., machine translation (Multi30K [T3), text classification (AG News and DBpedia [54)), and sentiment
analysis (SST-2 (29] and IMDb (4 ) on different attention- based models, i.e., Transformer [50],
XLNet [52]. and ALBERT 30. The details of each experiment are given below.
4.1 Activation Memory Saving and Model Performance
The first experiment is to evaluate the amount of activation memory reduction in the training process achieved by the proposed softmax output approximation and the consequent end-to-end model performance. To better understand its behavior, the models are 1) fine-tuned for down down-stream tasks from pre-trained model in Table [ and 2) fully trained from scratch in Table 2 on an RTX A6000 GPU, We recompose the datasets of some tasks by limiting the maximum length of the input sequence due to the large capacity of attention based models used in the experiment. The memory usage in the below three tables (Table aa), (b), and Table 2 represents the relative memory usage ratio to that of the normal softmax 1) without and 2) with the memory overhead incurred to save the softmax sequence, i.e., n log2 n, described above. We evaluate it with different sets of m given a fixed n in Equation 6 to restrict the activation memory usage of the softmax module during training and then measure their performance. The overhead remains constant with respect to n, regardless of m.
Classificution
Sentiment Analysis
Model
XLN
XLNa
Model Tisk
ALBERT
XLNNI
Task
News (OU)
D8pedia
(00=
SST-2
IMD6
f1-micro
Mem usage (%)
[1-icco
Mem usage (%)
f1-micro
Mcm usage (%)
fl-micro
Mem usage (%)
Normal
0.323
100% / TOIPS
039889
100% / TOORS
Normal
0,8322
100% 100%
0.8769
100% 7 TOP%
m=11
0.175
1%711.9%
0.743
1% 11.9%
m=1
0,8268
1% 11.9%
0.8833
1%7 11.9%
0.205
3% 7 15.9%
0.9699
3%6 7 15.9%
m=5
08314
5% 7 15.9%
0.8843
5%7 5%/11599 15.9%
10
0.335
10% 20.9%
0.977
10% 3198
10=10
08211
10% 20.9%
0.825
10% 20.9%
m==21
0.9353
20% 7 30.9%
0.9881
207 7 30.%
m=
0.8257
20% / 30.9%
0.8822
20% 7 30.9%
0.9771
30% /40.9%
80% 440.%
0,8003
30% /41
08840
30% /41,6%
m 40
09226
40% 50.9%
0.9835
40% 7 50.9%
0
40
0.8360
40% 7 55099
0.8800
40% 7 750.9%
m==50
0.9869
50% 7 760.9%
0.9862
50% 7 66.9%
w==50
018337
50% 7 160.9%
0,8838
50% 7 760.9%
(a) Text Classification result (fine-tuiig)
(b) Sentiment Analysis result (fine-tuiig))
Table 1: Among all n=100 100 softmax elements, only the m highest and m lowest elements are stored
in memory. 'Normal' indicates the normal softmax that stores all output elements in memory.
Text Classification. Table T shows the relative activation memory usage ratio to the normal softmax operation and the fl-micro score [15] of two text classification tasks, i.e., AG News [54] and DBpedia [54]. They are fine-tuned from the pre-trained XLNet model (52). In the AG News task, the 20.9% memory usage shows slight difference (0.0012) in the fl-score from the normal softmax, i.e., 0.9335 vs. 0.9323. In addition, there is an improvement of 0.0546 when using 60.9% activation memory. The accuracy improvement can be regarded as one of the positive effects of the proposed approximation method as the approximation errors are broadcasted to the gradient matrix and act as white noise that can improve the learning performance. Similarly, in the DBpedia task, the 15.9% memory usage shows a slight difference (0.002) in the fl-score from the normal softmax, i.e., 0.9869 vs. 0.9889. shows that sufficient training achieved from a large number of approximated softmax elements and a small number of memory-stored elements (less than 5%) in text classification tasks. When the length of input sequence, n, is increased to 1000, the softmax activation output, which takes 34.33 GB of memory, is reduced to 8.80 GB (AG News task) and 7.08 GB (DBpedia task).
Sentiment Analysis. Table [bb) shows the relative activation memory usage ratio to the normal softmax operation and the fll -miiooo score [15] of two sentiment analysis tasks, i.c., SST-2 [46) and IMDb (34. They are fine-tuned from pre-trained ALBERT 30 and XLNet 52 model, respectively, In the SST-2 task, the 11.9% memory usage shows slight difference (0.0058) in the fl-score from the normal softmax, i.e., 0.8268 vs. 0.8326. When n becomes 1000, the softmax activation output requiring 34.33 GB of memory space is reduced to 5.70 GB. In the IMDb task, the 11.9% memory usage shows a slight difference (0.0064) in the fl-score from the normal softmax, i.e., 0.8833 vs. 0.8769. The results of applying the softmax approximation show similar or higher scores to the normal softmax. They result from a combination of the characteristics of IMDb data, which contains many unstructured sentences and the noise caused by the softmax approximation.
7

---

Machine Translation
Training Perplexity
Test BLEU Score
Model
Transformer-biu base
Normal
m=8
Multt
(yt- 40)
MalitO Fn.
40)
63
m==1
m2-1
BLEU
Mem usage (%)
BLEU
Mem usage (98)
mb=2
m6
Normal
36.52
100S/ / 100%
34.05
100% / 100%
40
1
35.33
2.5% / 9%
32.76
2.5% 71 71.5%%
m==
36.29
5% / 14,4%
33.97
3% / 14.4%
m=1
19.4%
19.4%
mn
36.46
20% 29.49
34.28
20% 29.4%
50%
33.91
70% 39.45
91=
16
36.41
40% 49.4%
35.96
40%
499,4%
15
W0=20
36.97
50% 59.4%
33,95
50% 399.4%
Epoch
Epoch
Table 2: Machine translation result (full training). Among all n=40 softmax elements, only the
m highest and m lowest elements are stored in memory. 'Normal' indicates the normal softmax operation that stores all softmax output elements in memory for training.
The train perplexity (left) and test BLEU score (right) in machine translation (Table 2 over training epochs. The curves of the normal and approximated softmax almost overlap each other.
Figure 3:
Machine Translation. Table 2 shows the relative activation memory usage ratio and the BLEU score 37 of two machine translation tasks, i.c. Multi30K on German- Englishh (De-En) and English- German (En-De) translations [3] Transformer basic models [50] are fully trained from scratch for them. As observed in Table 2 the 14.4% memory usage case shows only a small difference (0.23) in the BLEU score from the normal softmax, i.c., 36.29 vs. 36.52. When more than 19.4% of memory is allowed, it starts to outperform the normal softmax. The result demonstrates that the approximated softmax outputs enable efficient yet effective learning for full training of the Transformer model while requiring much less activation memory, which can be readily extended to other types of attention-basee based models in various architectures and sizes. Figure their train perplexity and test BLEU score over training epochs, which shows almost the same learning trajectory is achieved when compared to the normal softmax for different m.
4.2 Softmax Output Approximation
The second experiment is to evaluate the approximation performance of the proposed method by comparing it against the original non-approximated softmax. We analyze the output of the softmax function over the training epochs. Figure 4(a) shows the comparison between the approximated softmax (Equation () and the normal softmax (Equation[] [T) at the 1, 8. 15, and 20 epochs. As shown in the figure, the approximation function effectively emulates the original softmax with low errors, i.e. only 0.0500 in MAE (mean absolute error) on average, and tends to get more correctly fitted as the training progresses.
4.3 Gradient Matrix Reconstruction
The third experiment is to evaluate the approximated gradient matrix (Equation 4 obtained from the approximated softmax output (Equation 6 by comparing it against the gradient matrix constructed from the original non-approximated softmax (Equation (T]. ( the heatmaps repre- senting the gradient error (Equation 7 measured in MAE between them over the training epochs, showing the approximated gradient matrices are almost the same as the original hon-approximated ones with consistently low errors. At the beginning of training (epoch 1), the average error is 0.00031, which is about O. 0.1244 12% of 0.25, the maximum value that each gradient element can have. As the training progresses, the average error reduces to 0.088% (epoch 15), and most gradient elements in the remaining epochs retain low errors around 0. 14%. Considering that the value of the gradient matrix is crucial it is used in the back-propagation chain to compute the weight parameter update, it implies that the approximated softmax outputs enable the accurate update of the model weight parameters in the training process as effectively reconstructs the gradient matrix with low error. Additionally, we convert the absolute gradient matrix error to the relative percentage error, i.c., 7.70% in epoch 1, 13.33% in epoch 8, 8.28% in epoch 15, and 12.74 12.44% in epoch 20. The elements of the gradient matrix range from -0.25 to 0.25, and since it has elements close to 0, the error appears to be unstable compared to MAE.
5 Related Work
The proposed softmax output approximation takes unique approach orthogonal to many existing memory-efficient training techniques summarized below, enabling additional improvement training memory efficiency when used on top of them, they do not curtail the softmax activation memory.
8

---

0.3
Epoch:
Epoch:
16
20
0.012
Normal
0,010
In Performer B an efficient matrix-wise calculation method, the row vectors of the query and key are mapped from d dimensions to r dimensions using random feature map á. Then, attention computed by taking the inner product of the mapped @(g:) and o. The memory efficiency of this Appraximatod softmax
MAC:0.0573
MAE-0. 0.0484
0.002
8.0
Epoch:
Fpoch:
Epoch:1, MAE:0.00031
Epoch:8, MAE:0.00029
00000
0.012
0.2
0,010
Suftmas
0.3
MAE: 0.039.
MAE 0.0547
000
200
12
0.002
Sorted autput
Sorted outpul
Epoch:15, MAE:0.00022 Epoch:20, MAE:0.00032
-00000
(a) Comparison of original and approximation
(b) The heatmaps of the gradient error (MAE)
Figure 4: (a) An example of the normal softmax (blue lines) and the approximated softmax output (red lines) sorted in descending order over the training epochs when n=21 and m=3. Among all 21 softmax elements, only 6 elements are stored in memory, while the remaining 15 elements are approximated with low error. (b) The heatmaps of the gradient error (RE) between the gradient matrix reconstructed from the approximated softmax outputs and the gradient matrix generated from the original softmax outputs. The attached plots on the left of each heatmap present the unsorted softmax output and the corresponding approximated softmax output, showing significant overlap.
Diverse Approaches to Efficient Attention Mechanisms. Efficiency improvements in attention- based models involve various approaches to streamline the attention mechanism, categorized as follows: 1) Efficient changes in the matrix-wise calculation of the attention mechanism, e.g., Flash Attention 9 Performer [ Linear Attention [23] Random Feature Attention [40], etc. 2) By compressing the sequence which the attention mechanism is applied, the size of the attention matrix is reduced, e.g.. Longformer (2) Sparse Transformers 77]. Reformer [24]. etc. 3) Reducing
the number of tokens the attention mechanism, e.g., Linformer (51). Luna (33 Perceiver [19, etc. 4) Utilizing quantization for memory-efficient training, e.g., Mesa (36), GACT 322. etc.
Flash attention 9 efficient matrix-wist calculation method, enhances the efficiency of the attention mechanism by employing varying memory access rates for each GPU memory layer. It speeds up memory- bouudd operations like matmul and softmax, where memory access exceeds computational requirements. Tiling enhances the efficiency of matrix matmul, whereas softmax is optimized through recomputation. To save memory and avoid the O(n?) memory requirement for softmax recalculation, the input is reconstructed after storing softmax normalization statistics. In this method, the proposed softmax approximation has limited applicability due to the absence of storage or processing of softmax output.
approach differs from the proposed softmax approximation method since it eliminates the softmax step in the process of streamlining matrix-wise operations.
Longformer 22, a sequence compression method, utilizes techniques like dilated sliding and global attention to overcome the limitations of existing methods, which face challenges in learning lengthy sentences due their (nn²) complexity. This approach proposes compressing the attention matrix from a larger size to a smaller one, effectively resolving the issue. In this approach, the operation of QKT is streamlined, and the attention score is computed using softmax the same way as existing attention mechanisms. Based on this, the proposed softmax approximation method can be utilized in this process. It efficiently approximates the whole from a part of the softmax output, making suitable for restoring already compressed sequence matrix, regardless of the input type.
Linformer [51] a method for reducing the number of tokens, simplifies complexity by decomposing the attention matrix into low-rank matrices and performing operations these low-rank matrices. The results obtained from these low-rank matrix operations are then processed through softmax operation to determine the attention scores matrix. This method overlaps with the proposed softmax approximation method due to the presence of activation memory resulting from the softmax operation.

---

Mesa [36], a memory--avinn method that employs quantization, utilizes quantization to store low- precision activation outputs, thereby reducing the memory used for activation storage during forward propagation. Additionally, during back-propagation, these low-precision activation outputs are dequantized to compute the gradients. In the context of mevory-saving through quantization, it's evident that efficiency limitations depend on the physical memory unit. However, differences arise in the proposed softmax approximation method because it can adapt flexibly based on the sentence.
The various methods described above have been employed to enhance efficiency through hardware- dependent measures, matrix-win calculation efficiency, and the compression or removal of segments within the attention matrix. The softmax approximation proposed in this paper is orthogonal to such approaches, as improves activation memory efficiency by approximating the entire softmax output from a subset. Notably, it can be used without requiring structural modifications to the existing model or attention mechanism, These distinctions make it easily applicable to many widely used attention-based models, and it can also be used concurrently with the above-mentionel efficient methods.
Re-materialization. Instead of saving activation outputs of each layer in memory, re-materialization methods [6] [10), also referred to as gradient checkpointing [16 14 28 generate them on the fly by re-executing (re-computing) the network. Although a substantial amount of memory can be saved by re-execution of the network, it comes at the cost of increased computation time (i.e. time-memory trade-off) as the network is repeatedly re-executed proportional to the number of model layers. Since 1) the complexity of the attention module is high, i... O(n)) in general [50], and 2) it appears multiple times in the network architecture, re- re-materializing attention-base in practice. Also, to handle the time- memooy trade-off, an appropriate schedule should be chosen for a
networks is not feasible
graph [27 17). Unlike them, the proposed softmax output approximation 1) saves the activation memory in training without increasing computation the network is executed only once, not multiple times, and 2) does not require scheduling a nack-propagation graph.
6 Discussions and Limitations
Approximation Overhead. To implement the proposed softmax approximation using the exponential distribution, which is a decreasing function, the softmax output elements are initially sorted in descending order. After approximating the sorted softmax output, they are reordered back to their original sequence to perform back- proppaaapaa........ In terms of space (memory) complexity, it requires storing the original sequence of softmax elements in memory, taking n log, n bits, where n is the number of softmax elements. When it comes to time (computation) complexity, it involves I) sorting the softmax elements, 2) performing the approximation, and 3) reordering the approximated softmax elements back to the original sequence.
The optimal implementation of the softmax approximation requires a lower-level programming language, like the CUDA kernel, rather than Python. Thus, the above processes would better be carried out at the "CUDA kernel' level to properly implement the softmax output approximation and assess the actual overhead. In this paper, we implement the softmax approximation in non-optimal way using Python, unlike the original softmax function highly optimized with the CUDA kernel. Our ongoing work involves further refinements, and one of the objectives is to optimize the softmax approximation algorithm with CUDA kernels. That said, the aforementioned overheads are negligible in general when compared to the massive training workload. Hence, with relatively little overheads in both memory and computation, the proposed method can efficiently approximate the softmax output.
Performance Improvement. In some cases, we have observed that the model's performance improves when the proposed softmax approximation is applied. When certain approximation errors in the softmax output are confined to a small range, they behave like white noise when delivered and broadcast to the gradient matrix. This white noise effect can have a positive impact on performance.
7 Conclusion
We propose to approximate the softmax output in the attention module to reduce its activation memory required to train attention- based deep models. By storing only a subset of the softmax output elements in memory during the forward pass and approximating the evicted elements during the backward pass, a significant amount of activation memory can be saved in an architecture-agnostic way when training attention-based deep models. The experiment shows that the proposed softmax output approximation, which significantly reduces softmax activation memory by up to 84% while maintaining comparable performance, is likely to be effective for models dealing with large datasets and substantial memory.
10

---

Acknowledgments and Disclosure of Funding
This work was supported by the National Research Foundation of Korea(NRF) grant funded by the Korea MoVernmen((MSIT) Technology Planning& & (Val)ation((ITTP) grant funded by the Korea (oVerament(MSIT) (No.2020-0- 01336, Artificial Intelligence graduate school suppor(UUNIST).
and Institute of Information & communications
References
[1] Olivier Beaumont, Julien Herrmann, Guillaume Pallez, and Alena Shilova. Optimal memory-aware back- propagation of deep join networks. Philosophical Transactions the Royal Society A, 37112166),20100449
2020.
[2] Tz Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05/ 2020.
[3] Trenton Bricken. Transformer Memory Requirements, 2022. https:/wwww.trentonnriikkmm comm/ Trans-
[4] John Bridle. Training stochastic model recognition algorithms as networks can lead to maximum mutual information estimation of parameters. Advances in neural information processing systems, 2. 1989. [5] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Ncelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 13:1877-1901, 2020.
[6] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv: arXiv:1604.00 1604.060447700 74, 2016.
[7] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 1904. 10509, 2019.
[8] Choromankii..... Valerii Likhosherstov. David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins. Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers, 2022.
[9] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness, 2022.
[10] Benjamin Dauvergne and Laurent Hascoët. The data-flow equations of checkpointing in reverse automatic diffcrentiation. In International Conference on Computational Science, pages 566-573. Springer, 2006. [11] Jacob Devlin. Ming-We Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre- Fre-training of deep bidirec- tional transformers for language understanding, arXiv preprini 100.:01810.04805 2018.
[12] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenbom, Xiaohua Zhai, Thomas Unterthiner. Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. /9999, 2020. [13] Desmond Elliott, Stella Frank, Khalil Sima an. and Lucia Specia. Multi30k: Multilingual english-german image descriptions. In Proceedings of the 5th Workshop on Vision and Language, pages 70-74. Association for Computational Linguistics. 2016.
[14] Jianwei Feng and Dong Huang. Cutting down training memory by re-fowarding. 2018.
[15] Margherita Grandini, Enrico Bagli, and Giorgio Visani. Metries for multi-class elassification: an overview.
arXiv preprint arXiv:2008. 05756, 2020.
[16] Andreas Griewank and Andrea Walther. Algorithm 799: revolve: an implementation of checkpointing for the reverse or adjoint mode of computational differentiation. ACM Transactions on Mathematical Software
(TMS)) 26( 1):19 19-4499,9944 2000.
[17] Audrunas Gruslys, Rémi Munos, Ivo Danihelka, Mare Lanctot, and Alex Graves. Memory-efficient
backpropagation through time. Advances in Neural Information Processing Systems, 29, 2016.
[18] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 770-778, 2016.
[19] Andrew Jaegle, Felix Gimeno, Andrew Brock, Andrew Zisserman, Oriol Vinyals, and Joao Carreira. Perceiver: General perception with iterative attention, 2021.
[20] Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoee. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint arXiv:1705.0355: 2017.
[21] Norman P Jouppi. Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, et al. In-datacenter performance analysis of tensor processing unit In Proceedings of the 44th annual international symposium on computer architecture, pages 1-12, 2017.
[22] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figumov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-588) 2021.
[23] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156-5165. PMLR, 2020.
[24] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04455 2020.

---

[25] Vijay Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models, 2022. Alex Krizhevsky, Ilya Sutskever, and Geotfrey E Hinton. Imagenet classification with deep convolutional neural networks. Communications of the ACM, 6016):84-90, 2017.
[26]
[27] Ravi Kumar, Manish Purohit, Zoya Svitkina, Erik Vee, and Joshua Wang. Efficient rematerialization for deep networks. Advances in Neural Information Processing Systems, 32, 2019.
[28] Mitsuru Kusumoto, Takuya Inouc, Gentaro Watanabe, Takuya Akiba, and Masanori Koyama, A graph theoretic framework of recomputation algorithms for memory-efficicn backpropagation. Advances in Neural Information Processing Systems, 32. 2019.
[29] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. ALBERT: A lite BERT for supervied learning of language representations. CoRR, abs/1909. 1999,,994, 2019.
[30] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel. Piyush Sharma, and Radu Soricut. Al- bert: A lite bert for self-supervised learning of language representations. arXiv preprint arXiv:1909. 1909.1999111111 1194, 2019.
[31] Peter Liu, Mohammad Salch, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. Generating wikipedia by summarizing long sequences. arXiv preprint arXiv: 1801,10/9,, 2018.
[32] 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. 2022.
[33] Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer.
Luna: Linear unified nested attention. 2021.
[34] Andrew Maas, Raymond E Daly. Peter T Pham, Dan Huang, Andrew Y Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies, pages 142-150. 2011.
[35] NVIDIA. Nvidia a100 tensor core gpu. ftttps://wwy nviiia. com/ en- us/ dardddaaaattteae/l000 2020.
[36] Zizheng Pan, Peng Chen, Haoyu Hc, Jing Liu, Jianfci Cai, and Bohan Zhuang. Mesa: memory-saving
training framework for transformers, 2022.
[37] Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th Annual Meeting of the Association for Computational Linguistics, pages 311 1118,,,, Philadelphia, Pennsylvania, USA, July 2002. Association for Computational Linguistics.
[38] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International conference on machine learning, pages 4055-4064 PMLR, 2018
[39]
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary De Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. NIPS-W,
2017.
[40] Hao Peng. Nikolaos Pappas, Dani Yogatama. Roy Schwartz, Noah A. Smith, and Lingpeng Kong. Random feature attention, 2021.
[41] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 188)9 2019.
[42] Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pages 400-407.11 1951
[43] Sheldon M Ross. Introduction to probobility and statistics for engineers and scientists. Academic press, 2020.
[44] Andrew W Senior, Richard Evans, John Jumper, James Kirkpatrick, Laurent Sifre, Tim Green, Chongli Qin, Augustin Zidek. Alexander WR Nelson, Alex Bridgland, et al. Improved protein structure prediction using potentials from deep learning. Nature, (777799))776-7111 2020.
[45] Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, Elton Zhang, Rewon Child, Reza Yaz- dani Aminabadi, Julie Bernauer, Xia Song, Mohammad Shoeybi, Yuxiong He, Michael Houston, Saurabh Tiwary, and Bryan Catanzaro. Using deepspeed and megatron to train megatron turing nig 530b, a large-scale generative language model, 2022.
[46] Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 conference empirical methods in natural language processing, pages 1631-1642, 2013.
[47]
SQuAD Squad. The stanford question answering dataset (2021).
[48]
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM
Computing Surveys (CSUR), 2020.
[49] Hugo Touvron, Matthicu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolics, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning, pages 10347-10357. PMLR, 2021.
[50] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems.
12

---

30, 2017.
[51] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.0477 2020.
[52]
Zhilin Yang. Zihang Dai, Yiming Yang. Jaime Carbonell, Russ R Salakhutdinov, and Quoe V.e Xlnet: Generalized autoregressive pretraining for language understanding. Advances neural information processing systems, 32, 2019.
[53] Zhilin Yang. Peng Qi, Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhutdinov, and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. arXiv preprint Xiv: 809, 09600, 2018.
[54] Xiang Zhang, Junbo Zhao, and Yann LeCun. Character- level convolutional networks for text classification. Advances in neural information processing systems, 28, 2015.
13


Reviews:
[{"review_id": "yTQzjS1ZNY", "reviewer": "Reviewer_fWPe", "summary": "This paper targets improving the memory efficiency of attention networks by reducing the activation storage of the softmax output. The authors propose to store only m highest and m lowest softmax output values, together with some auxiliary variables, and infer the missing part during backpropagation through interpolation. Results on several NLP tasks show encouraging results.\n", "questions": "* There are many approaches also aiming to improve the efficiency of attention, such as FlushAttention and Checkpointing. As a reader, I expect to see how your method compares with them and whether they could be complementary to deliver further improvements.\n* Apart from memory consumption, please provide training and inference time change. \n* The authors only experiment with Transformer base models, but the objective of memory reduction is to enable larger models. The authors should explore how their approach performs as the model size increases substantially.\n* For translation, a more convincing benchmark is WMT tasks. Examples in Multi30K follow simple patterns which artificially makes the translation task trivial. Please update your experiments with WMT benchmarks.", "limitations": "I didn't see particular limitations of their approach: the proposed method seems general to all softmax-based attention models.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea is easy to follow and implement\n2. The model obtains good performance and memory reduction on a set of tasks.", "weaknesses": "1. Important baselines are missing.\n2. The author only reports memory consumption but doesn't show the effect on training time and the impact of larger model sizes.\n3. Translation benchmarks should be improved.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for providing valuable and positive comments regarding our work, and the questions that helped us to refine our work.\n\n**Q1** Comparison of our method with  FlashAttention and Checkpointing\n> Flash attention tries to speed up the attention mechanism by using different memory access rates for each memory layer on the GPU. In order to improve efficiency using different access speeds, it accelerates matmul and sofmtax, that is, memory-bound operations, in which memory access is greater than the amount of computation. For this purpose, tilling and recomputation are used. In the case of matrix matmul, it is efficiently changed using tilling, and in the case of softmax, it is made efficient through recomputation. In this case, in order to avoid using $O(n^2)$ memory required to recalculate the softmax, the input is reconstructed after saving the softmax normalization statistics. In this case, our method is less applicable as there is no storage or operation of the softmax output. If possible, we think that we can increase the storage efficiency by not storing the input and normalization statistics and applying our method. This possibility will definitely be confirmed and evaluated through further research. In contrast, we would like to mention the merits of our method. In our method, if there is a large language model that prioritizes performance over any efficiency, and if it is an attention-based model, our method can be used depending on the presence or absence of softmax for attention scoring, regardless of the structure. Also, in the case of flash attention, it is a hardware-dependent method, but our method is software-dependent and efficient through an algorithm, so there may be a difference. Of course, flash attention is already sufficiently efficient. The reason for this view is that Flash attention v2, which is written at the level of 'CUDA kernel', exists. To compare at the same level, our method will also have to be compared with efficiency at the level of the 'CUDA kernel'. To this end, we will continue to study this matter and revise our manuscript. \n\n> Checkpointing is a method of inference between the corresponding layer and the checkpointed state whenever a gradient is needed without storing the output and input values of some layers required in the backward process. This increases time complexity due to inference and reduces activation memory. A similarity to our method is that there is an additional process to find the unsaved value when the gradient is needed. The difference is that checkpointing is selectively stored according to the model and situation, and since our method can be used with any softmax function, there will be a difference in that there are no restrictions on the model. Furthermore, if we can restore some values to the state to be checkpointed, similar to our method, the efficiency of checkpointing will increase.\n\n\n**Q2** About time complexity\n> We have conducted a new experiment on the training time of our method. The result can be found in the Global rebuttal response. Since our current code implementation is written in slow and naïve \"Pytorch\" and \"Python\", we would like to point that it is not obvious and not fair to compare our method with \"torch.nn.functional.softmax\" optimized with the CUDA kernel. Given this situation, we observe that the training time increases when compared to the baseline. The detailed indicators are included in the Global rebuttal response. We will optimize the implementation of our algorithm through continuous code optimization, especially with the CUDA kernel, which we expect to bring a huge speed up.\n\n**Q3** Experiments with large models\n> We appreciate the reviewer’s constructive suggestion on the need for experiments on large language models. Although we experimented with XLNet (0.12 billion parameters), similar in size to the GPT-2 small (0.124 billion parameters), which showed competitive learning performance, we agree that additional experiments are needed to see if our method works effectively for much larger language models. We will conduct in-depth experiments on larger models as much as possible we can by utilizing all the sources we have and will include the experiment result in our revised manuscript. Before presenting the experiment results on larger models, here we can provide some hints on how our method will work when applied to larger models. An important aspect of our method is how many and which softmax output values should be stored in memory to effectively perform softmax approximation that does not affect the training of a model. In general, in the case of a larger model that needs to learn longer and presumably more difficult data (e.g., longer sentences), the number of stored softmax output values in memory will also get longer if the ratio of the stored softmax output values is maintained. Thus, although the data gets longer and possibly gets harder to learn, we can expect that the softmax output values would be effectively approximated with marginal errors, as more softmax output values are stored in memory proportionally to the data length. However, again, we will conduct experiments to see if this surmise is correct or not. It would be appreciated if the reviewer would bear with us on the experiment results.\n\n**Q4** WMT task experiment\n> Thanks for suggesting to experiment with WMT. As the reviewer requested, we have conducted a new experiment on \"WMT2014 German-English\" to validate the reliability of our method. The result shows that our method achieves the same model performance (perplexity) on the WMT task up to using 7.75 times less activation memory. The details of the experiment results are provided in the PDF file in the Global rebuttal response. We kindly ask the reviewer to take a look at it."}, {"review_id": "0VsMPpUrUX", "reviewer": "Reviewer_jkF6", "summary": "This paper is about approximating softmax function, by storing only a fraction of the entire softmax output in a memory.\nThe authors argue that by applying approximation on the softmax function, they were able to save memory usage of softmax activation up to 84%.\nThe motivation starts from the observation that the softmax activation module takes up a huge amount of memory consumption in widely-used Transformer-based models.\nInstead of storing all softmax outputs, the work suggests storing only top-m highest and top-m lowest values of the outputs.\nFrom error analysis, the optimality of storing only highest and lowest values is justified.\nUsing those stored fraction of values, the rest of values are approximated in the assumption that the values will follow a modified exponential distribution.\nFrom experiments, the authors empirically prove that the proposed method has advantages of reducing memory usage while keeping performance degradation negligible.", "questions": "- In equation 5, the inverse of softmax to compute $z_i$ (i.e. input) from $s_i$ (i.e. output) depends on $\\sum_{j=1}^n e^{z_j}$. Can you elaborate more on how $z_i$ are recovered from $s_i$? Does that imply that we don't need to care about $\\sum_{j=1}^n e^{z_j}$, since it's partition function and the softmax function doesn't change output by transition?\n- In equation 8, the first and the second term seem to be equivalent only when $s_i s_j \\ge s_i' s_j'$ or $s_i s_j \\le s_i' s_j'$ for all $i \\neq j$. If the assumption does not hold, how does the error bound change?\n- In Table 1, as far as I understood, $m$ indicates the number of values to be kept among highest and lowest values; in such case $2m$ values will be stored. But for $n=100$, it is written that memory usage is 50%; doesn't $m=50$ mean we save all output values?\n- One of the old beliefs about softmax or attention mechanism is that having lower rank may suffer the performance. Is the proposed approximation method free from making softmax's rank lower?", "limitations": "I believe the limitations are addressed well in the checklist.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed method is easy to apply.\n- Experiments are done across multiple domains and tasks.\n- Reducing the overhead of softmax has high impact on further research", "weaknesses": "- It's not clear how large or small the MAEs of ~0.05 (output error) or ~0.0003 (gradient error) are. Rather than mean absolute error, relative error would be more appropriate method to present the significance of error introduced by the approximation.\n- Performance and memory usage comparison with other efficient softmax methods (e.g. LongFormer, Performer) would be useful.\n\n(typo)\n\nline 25: GTP-2 → GPT-2", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the time and effort you have dedicated to providing insightful feedback on our paper, including the typo checking.\n\n**Q1**  Inverse of softmax \n>We understand that Equation 5 may be confusing. In Equation 5, we can restore $z_i'$ (i.e., approximated softmax input) if we have $s_i'$ (i.e., approximated softmax output) and $\\ln(\\sum^n_{j=1} e^{z_j})$. In the forward execution, when the input goes through the softmax function, it stores the intact values of selected elements of the softmax output along with $\\ln(\\sum^n_{j=1} e^{z_j})$ in memory. Then, during the backward pass, by using $s_i'$ (i.e., approximated softmax output) and the stored $\\ln(\\sum^n_{j=1} e^{z_j})$, the input $z_i'$ is restored from Equation 5.\n\n**Q2** Error bound\n>As the reviewer correctly pointed out, the assumption may not always be held. However, we expect that the error bound does not change even if the assumption does not hold. That is because 1) the error bound is the sum of absolute error and 2) both the original softmax output ($s_i$) and the approximated output ($s'_i$) are positive numbers with a value between 0 and 1. Thus, the total amount of error given on the left-hand side in Equation 8 should not change regardless of the signs of each term. The reason we explicitly add the assumption is that we hope it would helps readers understand the process of the error-bound analysis. Although we believe the assumption does not affect the error bound, we will keep looking into it to see if there is a mistake or something we are missing. We appreciate again your careful and crucial comment on this.\n\n**Q3** $m=50$\n> Yes, as the reviewer mentioned, all the elements are stored in memory when $n=100$ and $m=50$ since it stores $2m$ elements in memory.\n\n**Q4** Lower rank\n> We appreciate the reviewer letting us consider an important aspect of softmax operation in the proposed method. We know about the low-rank performance degradation of the attention mechanism and softmax. Fortunately, the proposed method does not lower the rank of the softmax. It may seem to do so because the proposed approximation mechanism saves only $m$ selected softmax outputs in memory during the forward pass. However, the forward execution fully utilizes the full rank of the original softmax as a normal feedforward running, and then only the $m$ selected softmax outputs are stored in memory to be used later for the backward pass. During the backward pass, the non-stored softmax output elements are resorted to fully construct the exactly-same rank of the original softmax for gradient computation. Hence, the performance degradation caused by the lower rank of softmax is not the case that happens in the proposed method.\n\n**Gradient error**\n> We appreciate the reviewer’s advice regarding the readability and understandability of our manuscript. We agree with the reviewer that the current MAE metric may not help readers get the feel of the relative amount of error. So, as the reviewer suggested, we have converted the absolute gradient matrix error to the relative percentage error, i.e., [Epoch: 1 / RE: 7.70%], [Epoch: 8 / RE: 13.33%], [Epoch: 15 / RE: 8.28%], [Epoch: 20 / RE: 12.74%]. We will include the relative percentage error in our revised manuscript. In addition, through these improved indicators, we will try to further reduce the approximation error of the gradient matrix in future studies. Thanks again for your careful suggestion.\n\nWe truly appreciate the time and effort you invested in reviewing our work and providing constructive feedback. Your contributions have had a positive impact on the quality of our research, and We are sincerely thankful for your support."}, {"review_id": "OZa0hbrKUp", "reviewer": "Reviewer_UEEi", "summary": "This paper proposes to approximate softmax to improve memory efficiency for training attention models. During forward pass, first the softmax output is computed. Only the $m$ highest and $m$ lowest elements are stored along with sorted order from the the $n$ elements. Rest of the $n-2m$ entries are discarded to reduce the memory footprint. During backward pass, the discarded entries are approximated using exponential distribution whose parameters can be approximated using the stored elements. The $m$ highest and $m$ lowest elements are selected so as to minimize error in approximated gradient during the backward pass. Experiments on machine translation,  text classification, and sentiment analysis showcase the efficacy of proposed softmax approximation.", "questions": "Please see the discussion on weakness. In particular on missing justifications/claims, comparison to baselines/LRA tasks, and time-memory trade-off.", "limitations": "The discussion on negative societal impact is not required.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- While the idea to approximate softmax using the top-$m$ entries to reduce attention memory footprint has been explored before, the use of $m$ lowest entries and its motivation from gradient error minimization is novel. \n- The proposed method is simple and easy for the community to reproduce. The method also offers an easy way to plugin the approximation to pre-trained models.", "weaknesses": "  - Incorrect claims and/or missing justifications\n    - Lines 286-290 Efficient attention models do not focus on training memory reduction. This statement is false as there are many efficient attention mechanisms that explicitly reduce the memory for attention by (a) changing the order of computation with kernelized softmax: Linear [1], Performer [2], RFA [3]; (b) reducing the number of tokens to attend: sparse [4], Reformer [5], local [6], longformer [7], or (c) condensing the sequence: Linformer [8], compressive [9], perceiver [10]\n    - Line 60-61:  the softmax operation takes up 80% of the attention module itself, and 64.7% and 62.2% of the entire layer activation output during the forward pass. It is not clear that at what sequence lengths is this claimed.\n    - Error bound (Line 133): The assumption $s_i s_j \\geq s'_i s'_j$ or $s_i s_j \\leq s'_i s'_j$ is missing justification or any discussion. I am not sure how well this is going to hold in practice.\n  - Notations:\n    - Sometimes the paper uses confusing notations. For instance see line 74, $z$ is defined to be a vector of length $n$ with query/key multiplication $z= QK^T$ as an example. Typically $Q$ and $K$ are matrices of size $n \\times d$ which would lead $z$ to be a matrix.\n- Weak experiments:\n    - Baselines: The paper only compares against the scaled dot product attention as the baseline. The following two baselines are particularly close and also target memory reduction.\n      - FLASH Attention [11] exploits lazy softmax computation (normalization can be delayed until the end of the attention) to avoid storing softmax values and directly compute output values leading to memory savings.\n      - Memory-efficient Transformers via Top-k Attention [12]. This is another closely related work that not only targets memory improvements for attention but also improves the memory for feed-forward layers.\n  - Tasks:\n      - While the paper conducts experiments on MT, text classification, and sentiment analysis, it is unclear how the proposed technique would impact more general tasks in NLP or other domains. More real world tasks such as pre-training would help understand this better. Either Long range Arena [13] or more recent Comprehensive Attention Benchmark [14] would also provide more coverage on tasks and help understand better the tradeoffs of softmax approximation.\n      - Time-memory trade off as a function of sequence length. While the paper presents the relative memory improvement for sequence length of $100$, it would be good to see a general trend for memory consumption and time computation overhead as sequence length varies. The overhead on computation time is not discussed anywhere in the paper.\n- Missing references to other works that approximate softmax to improve the computational and memory complexity. Some of these work could also be good candidates as baselines\n    - SMYRF: Efficient Attention using Asymmetric Clustering (NeurIPS 2020)\n    - Fast transformers with clustered attention (NeurIPS 2020)\n    - Nystromformer (AAAI-21)\n    - Sparse attention with learning-to-hash (ICLR 2022)\n\n*References*:\n- [1] Fast Autoregressive Transformers with Linear Attention (ICML 2020)\n- [2] Rethinking Attention with Performers (ICLR 2021)\n- [3] Random Feature Attention (ICLR 2021)\n- [4] Generating Long Sequences with Sparse Transformers (Arxiv 2019)\n- [5] Reformer: The Efficient Transformer (ICLR 2020)\n- [6] Generating wikipedia by summarizing long sequences (Arxiv 2018)\n- [7] Longformer: The Long-Document Transformer (Arxiv 2020)\n- [8] Linformer: Self-Attention with Linear Complexity (Arxiv 2020)\n- [9] Luna: Linear Unified Nested Attention (NeurIPS 2021)\n- [10] Perceiver: General Perception with Iterative Attention (ICML 2021)\n- [11] Fast and Memory-Efficient Exact Attention with IO-Awareness (NeurIPS 2022)\n- [12] Memory-efficient Transformers via Top-k Attention (Proceedings of the 2nd Workshop on Simple and Efficient Natural Language Processing 2021)\n- [13] Long Range Arena : A Benchmark for Efficient Transformers (ICLR 2021)\n- [14] CAB: Comprehensive Attention Benchmarking on Long Sequence Modeling (ICML 2023)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Incorrect claims and/or missing justifications\n**W1** Line 286-290\n> Thanks for letting us know about the related works we missed. What we meant was that existing works on efficient Transformers do not explicitly try to save activation memory, unlike our method. Although they can save memory not in terms of activation memory, we would like to mention that our method is unique and orthogonal to them. By the way, we will go over them and revise our paper. In Performer[2] in (a), the row vector of query and key is mapped from $d$ dimension to $r$ dimension through random feature map $ϕ$. Then, attention is calculated by inner producing the mapped $ϕ(q_i)$ and $ϕ(k_j)$. In this case, memory efficiency is different from our method because the matrix-wise operation of the attention mechanism is removed from the efficient softmax. In Longformer[7] in (b), the existing attention-based model fails to learn long sentences due to the complexity of $O(n^2)$. This is a method using dilated sliding, global attention, etc. In this case, the operation of $QK^T$ is streamlined, and the attention score using softmax is obtained in the same way as the existing attention. We may be able to apply our method to this process. This will be tested later and used for comparison with transformer variants. In Linformer[8] in (c), since it simply decomposes the attention matrix into a low-rank matrix and uses it to obtain an attention score, it seems that it can be used overlapping with our method applicable depending on the presence or absence of softmax for attention score. Thank you for letting us know about the various perspectives of memory efficient attention as above. They may be different from the efficiency of activation memory that occurs during the training process like our method, but they will be greatly helpful for our future research.\n\n**W2** Line 60-61\n> It is the ratio of how much softmax is passed through among the layers that each input passes through inference in BERT and Transformer-base. It has nothing to do with the length of the sentence.\n\n**W3** Line 133 Error bound\n> As the reviewer correctly pointed out, the assumption may not always be held. However, we expect that the error bound does not change even if the assumption does not hold. That is because 1) the error bound is the sum of absolute error and 2) both the original softmax output ($s_i$) and the approximated output ($s'_i$) are positive numbers with a value between 0 and 1. Thus, the total amount of error given on the left-hand side in Equation 8 should not change regardless of the signs of each term. The reason we explicitly add the assumption is that we hope it would helps readers understand the process of the error-bound analysis. Although we believe the assumption does not affect the error bound, we will keep looking into it to see if there is a mistake or something we are missing. We appreciate again your careful and crucial comment on this.\n\n**W4** Line 74 Notations\n> Thank you for your accurate point. In the case of that notation, it's our mistake. Since $Z$ is a matrix and $\\vec{z}$ is a vector, as you said  $Z=QK^T$ of $Z \\in \\mathbb{R}^{n \\times n}$. Here, $n$ is the sequence length, and since softmax is applied based on $\\vec{z_i} \\in \\mathbb{R}^{1\\times n}$ to obtain the attention score, $\\vec{z}$ is expressed as $\\vec{z} \\in \\mathbb{R}^n$. Looks like it needs a fix.\n\n### Weak experiments\n**W1** Baselines\n> Flash attention speeds up the attention mechanism by using different memory access rates for each memory layer on the GPU. To improve efficiency using different access speeds, it accelerates matmul and sofmtax, that is, memory-bound operations, in which memory access is greater than the amount of computation. For this purpose, tilling and recomputation were used. In the case of matrix matmul, it is efficiently changed using tilling, and in the case of softmax, it was made efficient through recomputation. In this case, in order to avoid using $O(n^2)$ memory required to recalculate the softmax, the input is reconstructed after saving the softmax normalization statistics. In this case, our method is less applicable as there is no storage or operation of the softmax output. If possible, we think that we can increase the storage efficiency by not storing the input and normalization statistics and applying our method.\n\n**W2** Tasks\n> Our experiments were evaluated against representative NLP tasks. In the case of Text Classification and Sentiment Analysis, the pre-training model was fine-tuned, and in the case of MT, full-training was performed, which implicitly includes the pre-training stage. As the reviewer pointed, we were unable to experiment with a wide range of tasks due to time and resource limitations. Nevertheless, our method can restore the original softmax outputs using only a part of the softmax output, so it is expected to be capable of providing a similar performance unless data modality or type changes significantly. We will continue to conduct verification experiments with various tasks, including many interesting works the reviewer introduced, to see if there is some deviation from our expectations and revise our manuscript accordingly.\n> For Memory-efficient Transformers via Top-k Attention[12], a sparse attention method similar to Longformer[7] is borrowed from \"Generating Long Sequences with Sparse Transformers\" for softmax efficiency. We will evaluate the applicability and comparison of our method. We will also check the benchmarks and methods presented in the paper as much as possible we can.\n\n**W3** Time-memory trade off\n> We have attached a PDF file in Global rebuttal response, which provides new experiment results on memory consumption and computation time. We kindly ask the reviewer to see the Global response."}, {"review_id": "TPbZMDEE63", "reviewer": "Reviewer_s2ms", "summary": "The paper proposes an effective scheme for compressing the backpropagation of Softmax, the main idea of which is to only retain the maximum and minimum m values of the output, with the middle part using linear interpolation. Impressively, this approximation scheme not only saves memory but also achieves better results than precise backpropagation in some tasks.", "questions": "The method in the paper does indeed save the memory required for the backpropagation of Softmax, but I have a question, that is, is the proportion of memory occupied by the Softmax part in the entire backpropagation crucial?\n\nIntuitively, the memory occupied by the Softmax part is not too much, but perhaps my intuition is wrong, so I think this part needs further quantitative estimation, which determines the importance of this work (especially in the current large language model scenario).", "limitations": "I couldn't find any discussion about limitations in the main text and the appendix of the paper. Perhaps the authors should include it in the paper.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The memory usage of Softmax backpropagation is effectively reduced through approximation, which can be used in Attention models;\n\n2. There is a strict theoretical analysis of approximation errors;\n\n3. The performance is even better than the accurate Softmax backpropagation in some tasks.", "weaknesses": "1. There is no quantitative analysis of the amount of memory that can be saved theoretically;\n\n2. The memory efficiency has not been validated on larger models;\n\n3. The reason why the approximate softmax is better (in some tasks) has not been further analyzed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for providing valuable and positive comments regarding our work, and the questions that helped us to refine our work. Here, we have provided a detailed response.\n\n**Q1** The proportion of memory occupied by the softmax\n> Thanks for asking a crucial question about our work. It may seem the proportion of the softmax output is not significant from the perspective of the entire network. However, our analysis and experiment show that it actually takes a large portion of activation memory. For instance, the softmax operation takes up 80% of the attention module itself, and 64.7% and 62.2% of the entire layer activation output in the forward pass of the classic Transformer and BERT model. However, to validate the effectiveness of our method, we will conduct in-depth experiments on large models, as the reviewer suggested.\n\n> In principle, the size of the softmax activation memory of an attention-based model is equivalent to [Batch size $\\times$ Multi-head $\\times$ Sequence length $\\times$ Sequence length]. Therefore, for large language models that take longer sequences and/or a larger number of multi-heads, the percentage, as well as the absolute size, of softmax activation memory can be easily increased.\n\n> Although we will conduct experiments on larger models, we can take a simple example for an intuitive understanding of how big is the softmax activation output. In the case of the XLNet model, the softmax operation is passed 12 times for each forward execution. By assuming that the batch size is 256, the multi-head is 12, and the sequence length is 100, the total amount of memory taken by the softmax output is [256 $\\times$ 12 $\\times$ 100 $\\times$ 100] $\\times$ 12 $\\times$ 4 bytes, which is about 1.37GB. We can expect that for a large language model with more encoders, decoders, and multi-heads, softmax activation will take up more portion of activation memory."}, {"review_id": "ED9WhRmmYV", "reviewer": "Reviewer_r1xy", "summary": "This paper tries to reduce the memory footprint while training Transformer models by only keeping a fraction of softmax output and estimate them back during back propagation.", "questions": "Could you report the training profiling comparisons w/ and w/o the approximation overhead?\n\nCould you elaborate more on how this tech works with other Transformer variants like Performer?", "limitations": "n/a", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This is a very interesting and practical topic. With this tech, we can reduce the computation cost significantly while training large models.", "weaknesses": "I don't see any obvious weakness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We wanted to take a moment to express our heartfelt gratitude for reviewing our paper. Received your positive assessment has given us renewed confidence in our work.\n\n**Q1** Training profiling\n> Thanks for a good question. We felt that the training profiling the reviewer mentioned means two aspects (model performance and training time). Please let us know if we have mistaken your question.\n\n> First, profiling in terms of model performance (accuracy) is expected to be acceptable through our experiments presented in our current manuscript. In addition, in the case of the IMDb task, which had a rather low score, we have re-experiment it, which shows comparable performance to the baselines. Our Global rebuttal response includes the details of the new experiment on the IMDb task. \n\n>Second, profiling in terms of time is deemed necessary. Basically, we expect that the time complexity increases when applying our method, compared to the baseline that does not apply our method, as our method needs to perform additional processes such as sorting. Regarding this, we have conducted a new experiment on time overhead and included the result in the Global rebuttal response. We kindly ask the reviewer to check out the Global rebuttal response.\n\n**Q2** Transformer variants\n> Please let us describe the gist of common Transformer variants and their applicability to our method. The biggest goal of the Transformer variants is to reduce the time complexity of $O(n^2)$ of the matrix-wise calculation ($QK^T$) of the attention mechanism in the inference phase. In addition, some of these methods of reducing time complexity also streamline all of the $softmax(\\frac{QK^T}{\\sqrt{d}})V$, which are the computational processes of the attention mechanism. However, since these attention mechanisms are new variants of attention mechanisms, they are often accompanied by particular models suitable for them. However, our method can be applied in the case of an efficiency method with no change to softmax among attention mechanisms. The reason is that the applicability of our method depends on the existence of a softmax function for the input. If the softmax function exists and takes up activation memory, we can try our method. Also, the advantage of our method is that it can be applied to all variants of attention-based models. We believe this has been proven by our experiments on XLNet (0.12 billion parameters) with a similar number of parameters to GPT-2 small (0.124 billion parameters)."}, {"review_id": "P7GqIfYtSu", "reviewer": "Reviewer_8HU2", "summary": "\nThis paper proposes the Softmax Output Approximation algorithm, which approximates the softmax output during the forward pass and reconstructs it during the backward pass. In attention-based models like Transformer, the softmax output activations consume a significant amount of memory. By approximating them, a considerable amount of training memory can be saved. The algorithm successfully achieves memory savings of up to 84% compared to the existing softmax in various tasks.", "questions": "I think the idea presented in this paper is highly novel. If the following questions are addressed satisfactorily, **I am willing to raise my score to 5-7**.\n\n**1. Low accuracy on IMDb tasks**\n\n:It is disappointing to see a significant performance decrease on the IMDb dataset, while the SST-2 dataset shows minimal performance degradation. In my opinion, with appropriate hyperparameter tuning, it should be possible to achieve performance comparable to the baseline. If it is not feasible, it would be helpful to provide a qualitative explanation that softmax output activation may work well for general tasks, and IMDb tasks are an exceptional case.\n\n**2. Approximation time overhead**\n\n:It would be beneficial to report the increase in training time compared to the baseline. Specifically, it would be valuable to demonstrate that the approximation time overhead remains low even for larger input sizes.\n\n**3. Lack of Comparison:**\n\nAdding various algorithms [1-5] that approximate activations to the Related Work section would be beneficial. Particularly, for GACT [2] and Mesa [3], which also target Transformer and approximate softmax outputs, it is important to provide specific comparisons with proposed Softmax Output Approximation algorithm in terms of accuracy, memory usage, and training time.", "limitations": "This paper lacks a discussion on the limitations of the proposed algorithm and possible future work. It would be beneficial to include these aspects in the revised version of the paper.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and easy to understand overall.\n\n2. The mathematical analysis of gradients for approximating softmax outputs is reasonable. In particular, **the highly novel aspect of this approach** is that it approximates the distribution by assuming it and then uses sorting, without relying on pruning or quantization techniques.\n\n3. The experimental results demonstrate that the proposed method effectively reduces memory usage during pretraining or fine-tuning on certain datasets while achieving comparable accuracy to the baseline.\n\n\n\n", "weaknesses": "**1. Low accuracy on IMDb tasks**\n\n: The paper acknowledges the low performance on certain datasets, such as IMDb, and attributes it to the specific characteristics of the task. However, if the proposed softmax approximation struggles with tasks like IMDb, it may be challenging to apply it to other difficult tasks.\n\n**2. Approximation time overhead**\n\n: The paper claims that the overhead of the approximation is negligible compared to the general training workload. However, considering the potential overhead caused by sorting when dealing with large softmax outputs, it raises doubts about whether the overhead can truly be disregarded.\n\n**3. Lack of Comparison**\n\n: There are existing algorithms that approximate activations in various ways. For example, ActNN [1] and GACT [2] compress all activations based on their sensitivity, while [3] Mesa compresses all activations in a Transformer-specific manner. AAL [4] and DropIT [5] approximate input activations of linear and convolutional layers using auxiliary activations and pruning, respectively. It would be beneficial to compare the proposed approach with these existing methods.\n\n[1] ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training\n\n[2] GACT: Activation Compressed Training for Generic Network Architectures\n\n[3] Mesa: A Memory-saving Training Framework for Transformers\n\n[4] Learning with Auxiliary Activation for Memory-Efficient Training\n\n[5] DropIT: Dropping Intermediate Tensors for Memory-Efficient DNN Training", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We deeply appreciate the constructive feedback and the clear guidelines provided by the reviewer to make our work better. \n\n**Q1** Low accuracy on IMDb tasks\n> We have conducted a new experiment with the IMDb task and obtained comparable performance to the baseline, i.e., baseline accuracy 0.876 vs. our method accuracy 0.883 ($m=1$), by tuning the hyperparameters as the reviewer suggested. Please take a look at the details of the new experimental result (PDF file) posted on the Global rebuttal. The reason we got the low-performance result on the IMDb task in our submission might be some subtleties and difficulties in approximating the exact softmax outputs that tend to diverge randomly from the (continuous) exponential function. The IMDb task was the one that revealed this phenomenon most in our experiment. Although our method now works well with the IMDb task, as shown in our new experiment, we will keep working on validating the efficacy of our method for other difficult and general tasks and include them in our revised paper.\n\n**Q2** Approximation time overhead \n> As the reviewer correctly pointed out, our method inevitably increases the training time to some extent because it includes additional sorting and softmax output approximation processes. As the reviewer suggested, we have measured the time overhead incurred by our method over different batch sizes, sequence (input data) lengths, and the number of softmax output values stored in memory, of which details are included in the PDF file in the Global rebuttal. From this measurement experiment, we can observe that 1) it takes almost the same amount of time no matter how many softmax output values are stored in memory given the same input data length, and 2) the increase in time complexity caused by the change in batch size is greater than the sequence length. However, it looks like that our method takes at least about 1.6$\\times$ and up to about 6$\\times$ more time when compared to the baseline that does not apply our methods. The reason for it, we speculate, is that our method is implemented with the relatively slow and naïve “Python (Pytorch)” rather than fast “CUDA kernel (C++)”, unlike the original softmax operation that is highly optimized and run with “CDUA kernel”. In more detail, our method simply selects the “to-be stored softmax output values” after sorting them in the forward pass. And then, during the backward pass, it approximates the non-stored softmax output values and places them into their original positions. When turning this algorithm into actual code, we needed to use many additional matrix operations such as slicing, the scatter function, etc., not only the sorting mechanism, which takes a much longer time than sorting several hundreds or thousands of softmax elements, without implementation optimization (e.g., CUDA kernel optimization). That might explain the increased time complexity of our method compared to the baseline. Since the sorting time complexity we used is just $n \\log_{2} n$, which does not seem to incur significant time overhead compared to the compute-intensive back-propagation (gradient computation) when $n$ is around several thousands, we believe that the time gap between our method and the baseline would be dramatically reduced through the implementation optimization of our method, making our method more practical. Thanks again for your careful comment on this matter, and we will keep working on how to optimize our method to decrease the time overhead.\n\n**Q3** Lack of Comparison\n> We appreciate the reviewer providing reference papers that can be of great help to our submission and future research. We will include those papers in our related work and thoroughly discuss them against our method. Also, we will conduct a comparison experiment in terms of accuracy, memory usage, and training time. Here, we briefly discuss them as follows. \n\n> In the case of GACT, to our understanding, unit bits are reduced through quantization with a maximum of 4 bits and a minimum of 2 bits to minimize overlapping and errors of each layer under specific conditions. Thus, when applied to the softmax layer, it could achieve 16x memory reduction at maximum, i.e., from 32 bits (floating point) to 2 bits. And in the case of Mesa, similar to GACT, the gradient is quantized. As a preliminary experiment, we have briefly checked the difference between the gradient matrix obtained from the softmax output approximated by our method and the gradient quantized from Mesa. Although additional experiments seem necessary as it is unclear how this will affect actual learning, their gradient values generated from the same input data do not differ significantly from the result of our preliminary experiment. In other words, the error rates of the gradient matrix created from the existing normal softmax and the gradient matrix from Mesa and our method seem to be similar, meaning that there may be a little difference in their model accuracy.\n\n> In sum, we speculate our method would have more room to reduce softmax activation memory thanks to its flexible ratio of approximation ($m$ can be changed accordingly given $n$) while achieving similar model accuracy to GACT and Mesa, as the gradient matrix obtained from our method and Mesa seem to have similar error from the gradient obtained from the original softmax operation.\nHowever, we will conduct in-depth experiments with GACT and Mesa and compare them against our method (including training time) to see if our conjecture is correct or not, which will be included in our revised paper."}]

Author rebuttal:
We want to extend our heartfelt thanks to all the reviewers for taking the time to review our research paper from diverse angles and offering constructive critiques. Your valuable insights have greatly enriched the quality of our research.

We have conducted new and re-experiments requested by the reviewers and included the results in the attached PDF file. We kindly ask the reviewers to check out the attached PDF file.

**Experiments on IMDb and WMT (Table 1 and Figure 1)**

>***Re-experiment with the IMDb task***
> >  The result on the left side of Table 1 is the re-evaluation of the IMDb task described in the paper. All experimental parameters are the same as in the paper. The sentence length is limited to 100 and the batch size is 128. The result using normal softmax (baseline) is 0.876. By applying our method, the value of the highest 1 and lowest 1 softmax elements out of 100 softmax elements are stored in memory ($m=1$), and it achieves an accuracy of 0.883 by using only 1.9% of softmax activation memory, which is slightly higher than the baseline. This variability seems to arise from subtle difficulties in approximating outputs from softmax inputs (from unstructured sequence data) that tend to deviate randomly from the exponential function, which we will further study.

> ***New experiment on the WMT-14 De-En task***
> > The right side of Table 1 is the new experiment result on Machine Translation learning using WMT-14 De-En data for the Transformer-base model. In the case of the WMT14 dataset, there are about 4 million train data and 3,000 validation and test data. We have experimented with WMT-14 De-En using around 20,000 train data samples due to the tight time limitation of the rebuttal. The result shows the perplexity of the test dataset every $m$. The result using normal softmax (baseline) is 1.004. With our method, the same perplexity (1.004) is obtained by storing the values of the highest 2 and lowest 2 softmax outputs out of 100 elements. Figure 1 shows the learning performance trajectory with various $m$ over training epochs, which shows that our method draws almost the same trajectory as the baseline (normal).


**Experiments on time and memory complexity (Figure 2 and Figure 3)**

> ***Time and memory evaluation according to sentence length and batch size at a single softmax layer***
> > Figure 2(a) shows the activation memory usage and time complexity of the original softmax function and the softmax function to which our method is applied, as described in the caption. It can be observed that the time complexity changes according to the sequence length regardless of the number of softmax elements stored in memory. We can also observe that the time increases rapidly from a minimum of 4$\times$ to about 30$\times$ due to inefficient code implementations in the process of sorting, cutting, and rearranging the matrix, which are required by the proposed softmax output approximation. This is because we used a non-optimal way to implement our algorithm (Python, not CUDA kernel), unlike the original softmax function highly optimized with CUDA kernel. This will continue to be modified in our work, and we plan to implement it with the highly optimized CUDA kernel as the final goal.


> ***Evaluation of the effect of batch size and sequence length on the turnaround time of our method***
> > Figure 2(b), we present the factors that increase the time complexity when our method is applied. Each x-axis is batch size, and each line represents sentence length. As we can see, the batch size changing the softmax output shape to $n$ has a more significant effect on the inference speed than the sequence length ($n^2$). Based on this experiment result, we will optimize our code and make it faster.

> ***Time complexity evaluation of the entire model***
>> Figure 3, we proceed to the time evaluation spent by all the layers of the Transformer-base model, unlike Figure 2, which measures the amount of time spent at a single softmax layer. In this experiment, three different batch sizes of 32, 64, and 128 are used, which are divided into two groups; blue bars (original softmax; baseline) and green bars (our method). For the x-axis of $n=100$, the increase in time complexity due to our method occurs on average by a factor of about  1.7$\times$. In the case of $n = 200$, the complexity increases by about 5 $\times$ on average, and the amount of time required is about 6.3 $\times$ at most, which is smaller than the amount of time spent at a single softmax layer. Also, the trend of linear change is observed with the batch size and sequence length. Please note that these results may vary depending on the situation of the GPU or CPU. However, as we mentioned, we will improve our implementation of the proposed softmax output approximation method and will evaluate the time overhead again with the optimized code that can execute our algorithm with more efficiency while the significant memory-saving capability is maintained.

Meta-review: Strong paper on softmax output approximation to reduce memory requirements. Most of the reviewers think this paper should be accepted. Reviewer UEEi has some concerns but I think the authors have sufficiently addressed them in the rebuttal. I recommend acceptance.
Decision: Accept (poster)