SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Algorithm 1: Model Training with Softmax Output Approximation
Input:  $\mathcal{M}$  :  $l$  -layer model,  $\mathcal{B}$  : training mini-batch,  $m$  : number of softmax elements to be stored in memory for  $k\gets 1$  to  $l$  do [forward pass of the model  $\mathcal{M}$ ] execute the operation at the  $k$  -th layer with  $\mathcal{B}$  . if  $k$  -th layer  $= =$  softmax then foreach training sample in  $\mathcal{B}$  do 1 store  $m - 1$  highest and  $m - 1$  lowest softmax output elements and  $\alpha ,\beta$  [Eq.12] in memory; discard the remaining  $n - 2m + 2$  softmax output elements from memory; compute and store  $\ln (\sum_{j = 1}^{n}e^{z_j})$  [Eq.5] in memory; end else store the activation output of  $k$  -th layer obtained from  $\mathcal{B}$  in memory; end   
for  $k\gets l$  to 1 do [backward pass of the model  $\mathcal{M}$ ] if  $k$  -th layer  $= =$  softmax then foreach training sample in  $\mathcal{B}$  do 2 approximate  $n - 2m + 2$  softmax output elements using  $\alpha$  and  $\beta$  [Eq.11]; make the softmax output vector  $\vec{s}^{\prime}$  from 1, 2 [Eq.6]; make the softmax input vector  $\vec{z}$  from the softmax output vector  $\vec{s}^{\prime}$  [Eq.5]; make the gradient matrix  $\partial \vec{s}^{\prime} / \partial \vec{z}_{k - 1}$  from the softmax output vector  $\vec{s}^{\prime}$  [Eq.4]; end end perform back-propagation for the  $k$  -th layer with its activation output;
Normalization Factor. As shown in Equation 13, the range  $\sigma$  is normalized by  $1 / N$ , as a total of  $N$  softmax elements are approximated from it so that the discrete domain of  $k = [0,1,\dots,N - 1]$  can be used to approximate the  $k$ -th element. Since the approximation function in Equation 11 converges to 0 when  $k\to \infty$ ,  $\sigma$  is additionally multiplied with  $1 / (1 - \alpha)$  where  $\alpha = h_{m}$  (Equation 12) to increase the range  $\sigma$  inverse proportional to  $h_m$  based on the property that the sum of a softmax equals to one.
Rate Value  $(\hat{\beta})$ . The rate value  $\hat{\beta}$  in Equation 13 determines the shape of the approximation function, playing a similar role to  $\lambda$  in the exponential distribution. To determine  $\hat{\beta}$ , the approximation at  $k = \sigma$  is set as the same as  $\ell_{m}$  given Equation 11 and 14, i.e., we let  $s_{\sigma}^{\prime} = \ell_{m}$  and obtain  $\beta$  as:
$$
\beta = \sqrt {\frac {1}{h _ {m} \ell_ {m}}} \text {f r o m} s _ {\sigma} ^ {\prime} = \alpha e ^ {- \beta \cdot \left(- \frac {1}{\beta} \ln \frac {1}{\alpha^ {2} \beta^ {2}}\right)} = \frac {1}{\alpha \beta^ {2}} = \ell_ {m} \tag {15}
$$
where  $\alpha = h_{m}$  as in Equation 12. Since the approximation is based on the exponential distribution, the convergence rate is proportional to  $\beta$ . To utilize this characteristic, the rate value  $\hat{\beta}$  is multiplied by the distance between the  $h_m$  and  $h_{m + 1}$  in the exponential function, which is given by:
$$
\hat {\beta} = \sqrt {\frac {1}{h _ {m} \ell_ {m}}} \cdot e ^ {h _ {m} - h _ {m + 1}} \tag {16}
$$
By replacing  $\beta$  with  $\hat{\beta}$  in Equations 14 and 16 to distinguish them from  $\beta$  on the left-hand side of the first equation in Equation 13, they become equivalent to the right two equations in Equation 13. Multiplying  $\frac{1}{N} \frac{1}{1 - h_m}$ ,  $\sigma$ , and  $\hat{\beta}$  altogether in Equation 13 gives the final  $\beta$  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  $\mathcal{O}(n\log_2n)$  of time complexity and  $n\log_2n$  bits of space (memory) complexity, respectively, where  $n$  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 1 illustrates the model training procedure that uses the proposed softmax output approximation. Although its primary target is attention-based 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.
Implementation. We implement a custom softmax operation that applies the proposed softmax approximation (Algorithm 1) with PyTorch [39]. Since our implementation is fully compatible with the existing normal softmax operation, it can be easily applied to 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.e., machine translation (Multi30K [13]), text classification (AG News and DBpedia [54]), and sentiment analysis (SST-2 [29] and IMDb [34]) 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-stream tasks from a pre-trained model in Table 1, 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 1(a), (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 \log_2 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$ .
Text Classification
Model	XLNet	XLNet
Task	AG News (n=100)	DBpedia (n=100)
f1-micro	Mem usage (%)	f1-micro	Mem usage (%)
Normal	0.9323	100% / 100%	0.9889	100% / 100%
m=1	0.9175	1% / 11.9%	0.9743	1% / 11.9%
m=5	0.9205	5% / 15.9%	0.9869	5% / 15.9%
m=10	0.9335	10% / 20.9%	0.987	10% / 20.9%
m=20	0.9353	20% / 30.9%	0.9881	20% / 30.9%
m=30	0.9371	30% / 40.9%	0.9877	30% / 40.9%
m=40	0.9226	40% / 50.9%	0.9855	40% / 50.9%
m=50	0.9869	50% / 60.9%	0.9862	50% / 60.9%
(a) Text Classification result (fine-tuning)
Sentiment Analysis
Model	ALBERT	XLNet
Task	SST-2 (n=100)	IMDb (n=100)
f1-micro	Mem usage (%)	f1-micro	Mem usage (%)
Normal	0.8326	100% / 100%	0.8769	100% / 100%
m=1	0.8268	1% / 11.9%	0.8833	1% / 11.9%
m=5	0.8314	5% / 15.9%	0.8843	5% / 15.9%
m=10	0.8211	10% / 20.9%	0.8826	10% / 20.9%
m=20	0.8257	20% / 30.9%	0.8822	20% / 30.9%
m=30	0.8303	30% / 41.9%	0.8840	30% / 41.9%
m=40	0.8360	40% / 50.9%	0.8860	40% / 50.9%
m=50	0.8337	50% / 60.9%	0.8838	50% / 60.9%
(b) Sentiment Analysis result (fine-tuning)
Table 1: Among all  $n = 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 1(a) shows the relative activation memory usage ratio to the normal softmax operation and the f1-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 a slight difference (0.0012) in the f1-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 f1-score from the normal softmax, i.e., 0.9869 vs. 0.9889. It shows that sufficient training is 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 1(b) shows the relative activation memory usage ratio to the normal softmax operation and the f1-micro score [15] of two sentiment analysis tasks, i.e., SST-2 [46] and IMDb [34]. They are fine-tuned from a pre-trained ALBERT [30] and XLNet [52] model, respectively. In the SST-2 task, the  $11.9\%$  memory usage shows a slight difference (0.0058) in the f1-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 f1-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.
Machine Translation
Model	Transformer-base
Task	Multi30k	De-En (n=40)	Multi30k	En-De (n=40)
BLEU	Mem usage (%)	BLEU	Mem usage (%)
Normal	36.52	100% / 100%	34.05	100% / 100%
m=1	35.53	2.5% / 11.9%	32.76	2.5% / 11.9%
m=2	36.29	5% / 14.4%	33.97	5% / 14.4%
m=4	36.72	10% / 19.4%	34.16	10% / 19.4%
m=8	36.46	20% / 29.4%	34.28	20% / 29.4%
m=12	36.74	30% / 39.4%	33.93	30% / 39.4%
m=16	36.41	40% / 49.4%	33.96	40% / 49.4%
m=20	36.97	50% / 59.4%	33.95	50% / 59.4%
Figure 3: 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.
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.
Machine Translation. Table 2 shows the relative activation memory usage ratio and the BLEU score [37] of two machine translation tasks, i.e., Multi30K on German-English (De-En) and English-German (En-De) translations [13]. 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.e., 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-based models in various architectures and sizes. Figure 3 plots 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 6) and the normal softmax (Equation 1) 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 1). Figure 4(b) depicts the heatmaps representing 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 non-approximated ones with consistently low errors. At the beginning of training (epoch 1), the average error is 0.00031, which is about  $0.124\%$  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.114\%$ . Considering that the value of the gradient matrix is crucial as 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 it effectively reconstructs the gradient matrix with low error. Additionally, we convert the absolute gradient matrix error to the relative percentage error, i.e.,  $7.70\%$  in epoch 1,  $13.33\%$  in epoch 8,  $8.28\%$  in epoch 15, and  $12.74\%$  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.
The proposed softmax output approximation takes a unique approach orthogonal to many existing memory-efficient training techniques summarized below, enabling additional improvement in training memory efficiency when used on top of them, as they do not curtail the softmax activation memory.
(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 [8], Linear Attention [23], Random Feature Attention [40], etc. 2) By compressing the sequence to which the attention mechanism is applied, the size of the attention matrix is reduced, e.g., Longformer [2], Sparse Transformers [7], Reformer [24], etc. 3) Reducing the number of tokens in the attention mechanism, e.g., Linformer [51], Luna [33], Perceiver [19], etc. 4) Utilizing quantization for memory-efficient training, e.g., Mesa [36], GACT [32], etc.
Flash attention [9], an efficient matrix-wise calculation method, enhances the efficiency of the attention mechanism by employing varying memory access rates for each GPU memory layer. It speeds up memory-bound 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^{2})$  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.
In Performer [8], an efficient matrix-wise calculation method, the row vectors of the query and key are mapped from  $d$  dimensions to  $r$  dimensions using a random feature map  $\phi$ . Then, attention is computed by taking the inner product of the mapped  $\phi(q_i)$  and  $\phi(q_j)$ . The memory efficiency of this approach differs from the proposed softmax approximation method since it eliminates the softmax step in the process of streamlining matrix-wise operations.
Longformer [2], 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 to their  $O(n^{2})$  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  $QK^T$  is streamlined, and the attention score is computed using softmax in 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 it suitable for restoring an 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 on these low-rank matrices. The results obtained from these low-rank matrix operations are then processed through a 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-saving 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 memory-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-wise 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 it 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-mentioned 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, 1, 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.e.,  $\mathcal{O}(n^2)$  in general [50], and 2) it appears multiple times in the network architecture, re-materializing attention-based networks is not feasible in practice. Also, to handle the time-memory trade-off, an appropriate schedule should be chosen for a back-propagation graph [27, 17]. Unlike them, the proposed softmax output approximation 1) saves the activation memory in training without increasing computation as the network is executed only once, not multiple times, and 2) does not require scheduling a back-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-propagation. In terms of space (memory) complexity, it requires storing the original sequence of softmax elements in memory, taking  $n \log_2 n$  bits, where  $n$  is the number of softmax elements. When it comes to time (computation) complexity, it involves 1) 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 a 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.
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.
This work was supported by the National Research Foundation of Korea(NRF) grant funded by the Korea government(MSIT) (RS-2023-00277383) and Institute of Information & communications Technology Planning & Evaluation(IITP) grant funded by the Korea government(MSIT) (No.2020-0-01336, Artificial Intelligence graduate school support(UNIST)).
30, 2017.