Inverted Activations: Reducing Memory Footprint in Neural Network Training

## Abstract

Abstract The scaling of neural networks with increasing data and model sizes necessitates the development of more efficient deep learning algorithms.
A significant challenge in neural network training is the memory footprint associated with activation tensors, particularly in pointwise nonlinearity layers that traditionally save the entire input tensor for the backward pass, leading to substantial memory consumption.
In this paper, we propose a modification to the handling of activation tensors in pointwise nonlinearity layers.
Our method involves saving the output tensor instead of the input tensor during the forward pass. Since the subsequent layer typically also saves its input tensor, this approach reduces the total memory required by storing only one tensor between layers instead of two. This optimization is especially beneficial for transformer-based architectures like GPT, BERT, Mistral, and Llama.
To enable this approach, we utilize the inverse function of the nonlinearity during the backward pass. As the inverse cannot be computed analytically for most nonlinearities, we construct accurate approximations using simpler functions.
Experimental results demonstrate that our method significantly reduces memory usage without affecting training accuracy or computational performance.
Our implementation is provided as a drop-in replacement for standard nonlinearity layers in the PyTorch framework, facilitating easy adoption without requiring architectural modifications. The code is available at https://github.com/removed/for/anonimity .

## 1 Introduction

The remarkable scaling of neural networks with increasing data and model sizes has created a continuous demand for training and deploying larger models more efficiently.
This necessity drives the ongoing development of deep learning algorithms that are both computation and memory-efficient.

A significant portion of memory consumption during neural network training arises from activation tensors, alongside model parameters and optimizer statistics.
Activation tensors are stored during the forward pass to facilitate gradient computations in the backward pass.
Even simple and computationally fast operations, such as pointwise nonlinearities, contribute substantially to this memory footprint by storing entire input tensors.
In most popular deep learning frameworks, pointwise nonlinearities like GELU Hendrycks & Gimpel () and SiLU Elfwing et al. () save the whole input tensor for the backward pass.
This practice leads to considerable memory usage, especially in transformer-based architectures where such activations are prevalent.

In this work, we propose a modification to the forward and backward passes of element-wise nonlinearity layers.
Specifically, we suggest saving the output tensor instead of the input tensor during the forward pass.
Since the subsequent layer in the computation graph typically also saves its input tensor, our approach ensures that only one tensor is stored between layers rather than two, effectively reducing the memory footprint by nearly 25% in practice.

For a novel approach in implementing nonlinearity functions to gain widespread adoption, it should ideally meet three key conditions:

- 1.
The method should not introduce any additional computational error, ensuring that model performance remains unaffected.
- 2.
It should not slow down the model, maintaining or improving the speed of forward and backward computations.
- 3.
The approach should be user-friendly, meaning it can serve as a straightforward, drop-in replacement for existing nonlinearity layers without requiring modifications to the architecture or specific tuning of adjacent layers.

For example, highly optimized building blocks like Triton  or CUDA-based GeGLU implementations for models such as Llama and Mistral, or efficiently fused MLP blocks for transformer architectures, can offer significant speed and/or memory advantages, but are applicable to a very narrow range of use cases.
In contrast, for general research purposes and R&D, truly versatile layers are highly desirable.
In this paper, we demonstrate that our proposed method meets all three of these criteria, making it a promising candidate for broader adoption in deep learning systems.

We implemented this approach as a drop-in replacement for standard nonlinearity layers within the PyTorch Ansel et al. () framework, ensuring seamless integration for researchers and practitioners without the need for architectural modifications.
In LABEL:fig:usage_example_and_memory, we illustrate the straightforward integration of our nonlinearity layers into existing pipelines, which requires only the substitution of standard nonlinearities with our InvAct (Inverted Activation) layers, without the need for additional modifications to the architecture or hyperparameters.

## 2 Method

Background Knowledge.
The computation graph during neural network training can be expressed as a series of operations $f_{i}$, each of which transforms an input tensor $\mathbf{X}^{(i)}$ into an output tensor $\mathbf{X}^{(i+1)}$.
This output tensor subsequently serves as the input for the next operation in the computation graph:

$$ $\ldots\rightarrow\mathbf{X}^{(i)}\xrightarrow{f_{i}}\mathbf{X}^{(i+1)} \rightarrow\ldots\rightarrow L.$ $$

Here, $L$ is the final loss that is optimized using gradient descent.
Each operation stores some additional information $\mathbf{S}^{(i)}$, referred to as activation tensors, for the backward pass. During the backward pass, these activation tensors are used to compute the gradient of the loss $L$ with respect to all intermediate tensors of the computation:

$$ $\ldots\leftarrow\frac{\partial L}{\partial{\mathbf{X}^{(i)}}}\xleftarrow{\text {gradient }f_{i},\mathbf{S}^{(i)}}\frac{\partial L}{\partial{\mathbf{X}^{(i+1) }}}\leftarrow\frac{\partial L}{\partial L}.$ $$

For pointwise nonlinearity layers, the forward computation $f_{i}$ is simply the element-wise application of a function $f:\mathbb{R}\rightarrow\mathbb{R}$ to the input tensor $\mathbf{X}^{(i)}$:

$$ $\mathbf{X}^{(i+1)}_{j_{1},\ldots,j_{k}}=f\left(\mathbf{X}^{(i)}_{j_{1},\ldots, j_{k}}\right),$ (1) $$

while the backward computation is the Hadamard product of the gradient with respect to the output tensor and the derivative of $f$ evaluated at the input tensor:

$$ $\frac{\partial L}{\partial{\mathbf{X}^{(i)}}}=\frac{\partial L}{\partial{ \mathbf{X}^{(i+1)}}}\odot f^{\prime}\left(\mathbf{X}^{(i)}\right).$ (2) $$

Here, $f^{\prime}\left(\mathbf{X}^{(i)}\right)$ is computed element-wise, similarly to Eq. 1.

Since both Eq. 1 and Eq. 2 are applied independently to each element of the corresponding tensors, we can simplify the notation by letting $x\in\mathbb{R}$ represent an input element of the nonlinearity and $y\in\mathbb{R}$ represent the corresponding output:

$$ $y=f(x),\quad\frac{\partial L}{\partial x}=\frac{\partial L}{\partial y}f^{ \prime}(x).$ (3) $$

It is evident that to perform the backward computation in Eq. 2 efficiently, one can store the entire input tensor $\mathbf{X}^{(i)}$ as an activation tensor $\mathbf{S}^{(i)}$.
This approach is adopted by most popular deep learning frameworks, such as PyTorch and Jax, for nonlinearity layers like GELU and SiLU, which are widely used in transformer-based neural networks.

Figure: Listing 2: Pseudocode of inverted activation approach for forward and backward passes of pointwise nonlinearity layer. Boolean array S should be saved in an optimal way (1 bit per element) to achieve memory footprint reduction.
Refer to caption: x1.png

Inverted Activation.
In this work, we propose modifying the forward and backward computations to save the output tensor $\mathbf{X}^{(i+1)}$ instead of the input tensor $\mathbf{X}^{(i)}$.
If the subsequent layer after the pointwise nonlinearity also saves its input tensor, the two layers together will only store one tensor instead of two, thereby significantly reducing the overall memory footprint.

The backward computation in Eq. 3 can be rewritten as:

$$ $\frac{\partial L}{\partial x}=\frac{\partial L}{\partial y}f^{\prime}(x)=\frac {\partial L}{\partial y}f^{\prime}\left(f^{-1}(y)\right).$ $$

Here, $f^{-1}$ is the inverse function of the nonlinearity $f$, which gives our method its name: Inverted Activations (or InvAct for short).

Space-efficient Boolean Indicator.
Both GELU and SiLU are non-invertible functions, but they consist of two monotonic parts, each of which is individually invertible.
Thus, we can save additional information during the forward pass – a Boolean indicator that specifies whether each element of the input tensor belongs to the first monotonic part or the second part, which helps disambiguate $f^{-1}$:

$$ $s=\begin{cases}1&\text{if }x<T\\ 0&\text{otherwise}\end{cases}$ (4) $$

Here, $T$ is the point of separation between the two monotonic halves (see Fig. 1 for a graphical explanation).
It is possible to store the tensor $\mathbf{S}$ in a space-efficient Boolean array, where each element requires only 1 bit of memory.
In our implementation, we achieve this by storing the Boolean array $\mathbf{S}$ in a tensor $\mathbf{S}_{\text{compressed}}$ with the data type ‘uint8‘, in which $\mathbf{S}[i]$ corresponds to the $(i\mod 8)$-th bit of the $(i\div 8)$-th element:

$$ $\mathbf{S}[i]=\mathbf{S}_{\text{compressed}}\left[i\div 8\right]\gg\left(i\mod 8 \right).$ $$

Here, $\gg$ denotes the bitwise right shift, and $\&$ denotes the bitwise AND operation. Pseudocode for the resulting pointwise nonlinearity layer can be found in LABEL:alg:forward-backward.

Approximation.
Both GELU and SiLU are complex functions, and to the best of our knowledge, neither their inverses $f^{-1}$ nor $f^{\prime}\circ f^{-1}$ can be analytically derived.
Therefore, we propose to use approximations based on simple base functions, similar to the fast GELU approximation Hendrycks & Gimpel ().

There are many decisions to be made when constructing the approximation for $f^{\prime}\circ f^{-1}$: whether to approximate only $f^{-1}$ and evaluate $f^{\prime}$ in the standard way, or to approximate $f^{\prime}\circ f^{-1}$ as a whole; which primitive functions to use; what form the approximation should take; and how to balance complexity and accuracy.
In this work, we chose to approximate $f^{\prime}\circ f^{-1}$ as a whole.
The resulting approximation is highly accurate – so much so that the training process using our approximation is indistinguishable from training with the original nonlinearity layer, as we will demonstrate in our experiments in Section 3.

Figure: Figure 2: Derivative of the nonlinearity and its approximation in two coordinate systems: plotted as a function of the output variable $y$ (left two columns) and the input variable $x$ (right column). The $y$-based representation highlights how equations Eqs. 5, 6, 7 and 8 approximate the respective regions of $f^{\prime}(f^{-1}(y))$. The $x$-based representation, in turn, helps assess the approximation quality of the layer baased on how it is used in a neural network context.
Refer to caption: x2.png

For GELU we use the following approximations:

$$ $\displaystyle x<T\Rightarrow f^{\prime}\left(f^{-1}(y)\right)$ $\displaystyle\approx q^{\text{left}}(y)=c_{0}\sqrt{y+c_{1}}*(2y+c_{2}\sqrt{-y} )(|c_{3}y^{2}+|c_{4}y+c_{5}|+c_{6}|+c_{7}),$ (5) $\displaystyle x\geq T\Rightarrow f^{\prime}\left(f^{-1}(y)\right)$ $\displaystyle\approx q^{\text{right}}(y)=1+\left(c_{0}+c_{1}\sqrt{\tilde{y}}+c _{2}\tilde{y}\right)e^{c_{3}\left(c_{4}-\tilde{y}\right)^{3}},\text{ where } \tilde{y}=y-f(T).$ (6) $$

For SiLU we use the following approximations:

$$ $\displaystyle x<T\Rightarrow f^{\prime}\left(f^{-1}(y)\right)$ $\displaystyle\approx q^{\text{left}}(y)=\left(c_{0}+c_{1}\sqrt{\tilde{y}}+c_{2 }\tilde{y}+c_{3}{\tilde{y}}^{2}\right)*(1-y)+y,$ (7) $\displaystyle x\geq T\Rightarrow f^{\prime}\left(f^{-1}(y)\right)$ $\displaystyle\approx q^{\text{right}}(y)=\left(1+\left(c_{0}+c_{1}\sqrt{\tilde {y}}+c_{2}\tilde{y}\right)e^{c_{3}\left(c_{4}-\tilde{y}\right)^{3}}\right)*(1- y)+y,$ (8) $\displaystyle\text{ where }\tilde{y}=y-f(T)$ (9) $$

The coefficients $c_{i}$ are precomputed and fine-tuned specifically for each function (GELU or SiLU) and each approximation region ($x<T$ or $x\geq T$).
All coefficient values used in this work are provided in Section A.2.
In Fig. 2, we illustrate the functions $f^{\prime}(f^{-1}(y))$ for $x<T$ and $x\geq T$, along with their respective approximations $q^{\text{left}}$ and $q^{\text{right}}$.
The plots also show the approximation errors, defined as $f^{\prime}(f^{-1}(y))-q^{\text{left}}(y)$ for $x<T$ and $f^{\prime}(f^{-1}(y))-q^{\text{right}}(y)$ for $x\geq T$.

It is important to note that the presented approximations are not necessarily the most optimal or minimal representations of $f^{\prime}(f^{-1}(y))$.
As we are not experts in constructing such approximations, it is possible that simpler, more accurate, and computationally efficient alternatives exist, which could further enhance the performance of InvAct.
Nonetheless, our current approximations are well-suited for the specific task at hand, as demonstrated by their effectiveness, accuracy and speed in the experiments detailed in Section 3.

### 2.1 Alternatives to Bit-Compressed Boolean Indicator

An essential component for inverting nonlinearities such as GELU or SiLU is an indicator function that identifies which monotonic part the original $x$ belongs to.
Currently, we store this information in a space-efficient, bit-compressed Boolean tensor, where each Boolean value occupies only 1 bit of memory.
In this section, we will discuss two alternative methods for storing this information without requiring even 1 bit per element.

Sign-bit Inverted Activation.
Both GELU and SiLU functions are bounded, which means there exists some constant $C$ (see Fig. 1) such that for all $x$, $f(x)\geq C$, where $f$ represents either the GELU or SiLU function.
This implies that the value $f(x)-C$ is always greater than zero.
Instead of storing $f(x)$ for the backward pass, we can store $f(x)-C$ with its sign bit modified according to the indicator function $s$ (see Eq. 4).

For most popular floating-point formats used in deep learning, such as float32, float16, and bfloat16, the sign bit is the last bit of a floating-point number.
This approach allows us to eliminate the need for additional memory to store the 1 bit per element for the indicator function, making the entire nonlinearity layer memory-free in terms of footprint. The modified value $y=f(x)$ does not pose an issue for subsequent computations, as $C$ is a predefined global constant for a given activation function.
For example, the computation for an activation function followed by a fully-connected linear layer can be expressed as follows:

Here, the forward pass takes the input tensor $X$ and the matrix $W_{\text{linear}}$ of the fully-connected linear layer and returns $f(X)\cdot W_{\text{linear}}$.

Unfortunately, to the best of our knowledge, it is not possible to implement such a layer in PyTorch that works independently of what comes after the nonlinearity function, meaning that you would have to write a custom implementation for each new architectural block.
As a result, the sign-bit approach does not satisfy condition 3 outlined in the Introduction (see Section 1), and thus we do not consider it a viable candidate for broad adoption.

Precision Bit Inverted Activation.
The second approach to store the indicator function $s$ (see Eq. 4) is to use the lowest precision bit of a floating-point number (for most popular floating-point data types, this is the first bit in a base-2 representation).
During the forward pass, we compute $y=f(x)$ and modify this bit of $y$ to be equal to $s$ (see Eq. 4).

Compared to the previous approach, this method does not violate condition 3, as it allows for a convenient drop-in replacement for classical activation functions on all modern deep learning frameworks.
However, it is more challenging to evaluate the accuracy of this approach.
Both the standard inverted activation and the sign-bit inverted activation perform the forward pass exactly and introduce only a minimal approximation error during the backward pass.
This method, on the other hand, sacrifices one precision bit for the pointwise nonlinearity computation during the forward pass.
Moreover, with the ongoing trend of reducing floating-point precision, loss of accuracy may become larger in the future.
For this reason, we do not consider it a viable candidate for broad adoption at this time, as it requires extensive experimentation to validate its accuracy.
However, we still regard it as an interesting approach that may find applications in the future.

This is a pseudocode for the forward and backward passes of this approach:

## 3 Experiments

The validation of our method’s efficiency is two-fold:

- 1.
We compare the computational efficiency of the forward and backward passes of our inverted activation nonlinearity layer in various settings.
- 2.
We demonstrate that the approximation of the nonlinearity gradient does not affect the quality of gradient descent during training.

With these results, we can confidently assert that inverted nonlinearity layers can be safely used in practice to reduce the memory required for training neural networks without any speed or accuracy degradation.
All tests and measurements were performed on a single NVIDIA A100 GPU.

### 3.1 Computational Efficiency

In this section, we compare the computational efficiency of PyTorch Ansel et al. () nonlinearity layers with our Triton-based  implementations of inverted activation nonlinearities.
We conduct three groups of tests (with detailed setups presented in Section A.3):

- 1.
Application of nonlinearity to a tensor.
- 2.
Application of several layers, one of which is a nonlinearity layer.
For this purpose, we selected several popular building blocks of modern neural networks:
•
Linear layer + nonlinearity
•
Transformer MLP block: Linear layer + nonlinearity + Linear layer
•
GeGLU Shazeer (): Hadamard product of a linear layer + nonlinearity with another linear layer
- 3.
Full neural network training iterations for the following popular network architectures:
•
BERT
•
Llama v3.1 8B

The corresponding measurements are presented in table 1.
The results indicate that the difference in computational efficiency between our inverted activation nonlinearity layer and the standard nonlinearity layer within various architectural blocks (such as MLP, Linear + activation, or GeGLU) is less than 1%.
Furthermore, for full-scale neural networks like BERT and Llama, this difference is even more negligible, being less than 0.1%.

**Table 1: Comparison of computation time (in milliseconds) for different activation function implementatoins and neural network architectures. We compare following architecture blocks: plain nonlinearity layer, fully-connected layer with nonlinearity, Transformer MLP block, GeGLU Shazeer (), BERT Devlin et al. () model, Llama3.1-8b model. For each architecture block we compare following activation functions: torch GELU, our inverted activation and precision-bit inverted activation. We show difference in percents w.r.t. torch nonlinearity. Our layer is lower then 1% slower compared to torch nonlinearity inside some computation block, and lower then 0.1% slower inside BERT and Llama.**
|  | Torch (ms) | InvAct (ms) | Precision Bit InvAct (ms) |
| --- | --- | --- | --- |
| Plain GELU | $0.82$ | $0.86$ $(+5.75\%)$ | $0.86$ $(+4.96\%)$ |
| Linear + GELU | $11.78$ | $11.80$ $(+0.13\%)$ | $11.79$ $(+0.06\%)$ |
| Tranformer MLP Block | $75.41$ | $75.46$ $(+0.06\%)$ | $75.41$ $(-0.00\%)$ |
| GeGLU | $94.08$ | $94.13$ $(+0.05\%)$ | $94.09$ $(+0.00\%)$ |
| BERT | $1112.12$ | $1112.62$ $(+0.05\%)$ | $1112.25$ $(+0.01\%)$ |
| Llama3.1-8b | $182.84$ | $182.94$ $(+0.06\%)$ | $182.80$ $(-0.02\%)$ |

### 3.2 Effect of the Approximation Error

Figure: Figure 3: Comparison of BERT model fine-tuning with original and inverted GELU activations. The first row (except the rightmost column) compares two runs with Torch GELU and inverted GELU using identical random seeds. The second row shows losses averaged across 16 runs for each activation function. The leftmost column displays the training loss during the first 100 iterations. The middle column shows the training loss during the last 100 iterations out of a total of 2000. The rightmost column compares the validation losses. The observed losses, both training and validation, are nearly identical, with any discrepancies attributable to minor approximation inaccuracies that are absolutely insignificant and negligible for practical purposes.
Refer to caption: x4.png

In this section, we investigate the impact of our approximation method on the training process and final model performance.
To evaluate this, we conducted experiments using two popular neural network architectures: BERT and Llama.
We trained these models on standard benchmark tasks using both the original PyTorch GELU activation function and our approximated inverted GELU function for BERT, as well as the PyTorch SiLU activation function and our approximated inverted SiLU function for Llama.

For BERT, we fine-tuned the model on the Yelp Review dataset Yelp () for sentiment classification.
We compared the training process of models using both the original GELU activation function and our inverted GELU activation.
Comparisons of the training and validation losses are presented in Fig. 3.
We observe that the approximation inaccuracy of the inverted activation does not affect BERT’s training quality, as both training and validation losses are nearly indistinguishable over all $2000$ steps of the training process.
The variance produced by different random seeds is significantly larger than any difference between the two activation functions, suggesting that any discrepancies are negligible and on par with machine precision.

We extended this experiment to a significantly larger model, Llama v3.1 8B Dubey et al. (), which uses a different activation function – SiLU.
We fine-tuned a 4-bit quantized version with LoRA Hu et al. () using ORPO on a combination of DPO datasets(^3^33[https://huggingface.co/datasets/mlabonne/orpo-dpo-mix-40k](https://huggingface.co/datasets/mlabonne/orpo-dpo-mix-40k)).
The results for Llama training, presented in Fig. 4, show an even closer match between the original SiLU activation and our inverted SiLU activation compared to the BERT experiment.
The training and validation losses are virtually indistinguishable throughout the entire training process.
This remarkable similarity further supports our conclusion that the approximation inaccuracy of the inverted activation has a negligible impact on training quality and final model performance.

Figure: Figure 4: Same as Fig. 3, but for the Llama v3.1 8B model with SiLU nonlinearity. The first row (except the rightmost column) compares two runs with Torch SiLU and inverted SiLU using identical random seeds. The second row shows losses averaged across 16 runs for each activation function. The leftmost column displays the training loss during the first 100 iterations. The middle column shows the training loss during the last 100 iterations out of a total of 2000. The rightmost column compares the validation losses. The observed losses, both training and validation, are nearly identical, with any discrepancies attributable to minor approximation inaccuracies that are absolutely insignificant and negligible for practical purposes.
Refer to caption: x5.png

Based on this comparison of model training using original and inverted activations, we can conclude that the approximation inaccuracy of the InvAct does not affect training quality in any way.
Combined with the speed measurements presented in Section 3.1, we can state that the inverted activation nonlinearity layer is a perfect drop-in replacement for GELU and SiLU activation functions.

Furthermore, the seamless integration of our method into existing architectures – without requiring any modifications to pipelines, code, or hyperparameters – highlights its potential for widespread adoption in various deep learning applications.
The ability to reduce memory requirements without compromising model accuracy or training dynamics makes our inverted activation approach a promising tool for researchers and practitioners working with memory-intensive models, especially in resource-constrained environments.

### 3.3 Comparison with FewBit

Another method that reduces the memory required by pointwise nonlinearities is FewBit Novikov et al. (), a gradient quantization method that saves a quantized version of the nonlinearity gradient.
FewBit with one bit per element suffers from a considerable accuracy drop, as demonstrated in Fig. 5.
The authors report that using 4-bit quantization, which requires four times more memory compared to our standard InvAct implementation, FewBit does not show any performance degradation.
However, even in that setting, our approach is superior: the approximation of $f^{\prime}(f^{-1}(x))$ is much more accurate than FewBit quantization up to 8 bits!
We show $L_{2}$ and $L_{\infty}$ approximation errors for FewBit and InvAct in Fig. 5, and our approach demonstrates significantly better approximation quality, even when compared to FewBit using 8 times more bits per element.

Figure: Figure 5: Comparison of InvAct (our) nonlinearity layer with the FewBit Novikov et al. () nonlinearity layer. The two left plots show the validation loss for fine-tuning BERT and Llama v3.1 8B, respectively. The BERT model uses the GELU nonlinearity, while the Llama model uses the SiLU nonlinearity. FewBit is used with a 1-bit per element memory budget to match the memory footprint of InvAct. The FewBit approach exhibits significantly worse validation loss, while our InvAct approach is indistinguishable from the standard Torch nonlinearity. The right column shows the $L_{2}$ approximation error (top plot) and $L_{\infty}$ approximation error (bottom plot) for FewBit GELU with respect to the number of bits used in gradient quantization (blue line), compared to the approximation quality of InvAct (red dashed line). Our approach always uses only 1 bit per element. InvAct demonstrates superior approximation quality, even when compared to 8-bit FewBit GELU (y-axis is in logarithmic scale).
Refer to caption: x6.png

## 4 Related Work

The efficient scaling of neural, has necessitated innovations in low-precision training, quantization techniques, and various compression methods.
These approaches address computational and memory bottlenecks, enabling the training of ever-larger models on hardware-limited resources.

Quantization methods aim to compress neural network weights and activations by representing them in lower-precision formats, such as 4-bit integers.
One can quantize model weights Jacob et al. (); Zafrir et al. (); Gholami et al. () or activations that are saved for backward Chen et al. (); Novikov et al. ().
Former are used more in during inference, but approaches to utilize it during training are also exist Liu et al. (); Shao et al. (); Xu et al. ().
These methods are especially popular for the fine-tuning stage of training, where reducing memory usage without extensive retraining is crucial.
While quantization effectively decreases the memory footprint, it can introduce computational inaccuracies. As a result, these errors lead to degradation in model performance, making the balance between memory savings and maintaining model quality a key consideration when applying quantization.

Low-precision training techniques have emerged as another powerful approach to reducing both the memory and computational requirements of large-scale models Micikevicius et al. ().
One widely adopted format is bfloat16 Kalamkar et al. (), which retains a wider range of values compared to traditional FP16 while providing faster computation and reduced memory consumption.
More recently work towards even lower precision training in float8 formats have also been explored Wang et al. () for their potential to further reduce memory and computational demands.
These formats are particularly attractive for their extreme compression capability, though their usage requires careful consideration to avoid significant accuracy loss.
Moreover, widespread adoption of these lower precision formats relies heavily on GPU manufacturers integrating support for them directly into hardware. Until this support is in place, these formats remain experimental and difficult to implement at scale.

Custom CUDA and Triton implementations Unsloth (); Hsu et al. () have become essential for optimizing critical building blocks in popular architectures, significantly improving memory efficiency and computational speed.
For instance, highly fused implementations of MLP blocks Müller et al. () or optimized GeGLU layers are widely used to reduce the overhead of separate operations by merging them into a single, streamlined kernel.
Another prominent example is FlashAttention Dao (), which significantly improves the efficiency of attention mechanisms by reducing memory bottlenecks while maintaining performance.
These custom kernels, often developed in CUDA or Triton, enable fine-grained optimizations tailored to specific hardware architectures, providing substantial performance improvements.
However, the development of such highly optimized components demands significant engineering effort and a deep understanding of both the underlying hardware and neural network operations.
Each optimization, whether in matrix multiplications, element-wise operations, or memory management, requires meticulous custom work, and the resulting implementations are typically neural netowrk architecture-specific.
Despite the complexity, continuous innovation in these low-level building blocks is critical, as it provides massive improvements in memory usage and training speed, making it possible to push the boundaries of model scaling and performance.

## 5 Limitations

It is important to note that inverted nonlinearities save memory only when the next layer after the nonlinearity also saves its input for the backward pass. Fortunately, this behavior is common in most popular neural network architectures, such as transformers.

Additionally, our method reduces memory usage only for nonlinearity layers. The percentage of memory savings depends on the architecture in use. Furthermore, the presence of complementary techniques may impact these gains. For instance, when training models with checkpointing Griewank & Walther (), the memory bottleneck could shift between activations saved during the forward pass and activations recomputed during the backward pass, potentially affecting the benefits of our approach.

Despite these limitations, our highly efficient implementation of inverted activations within the PyTorch framework using Triton kernels makes us confident that our method will see broad adoption. We believe that it will prove invaluable to deep learning practitioners who seek to reduce memory usage without compromising speed order and model accuracy.

## 6 Conclusion

In this work, we presented an innovative method to reduce the memory footprint during neural network training by modifying how activation tensors are saved for backward passes in pointwise nonlinearity layers.
By saving the output tensor instead of the input tensor, we significantly decreased the memory requirements for transformer-based architectures such as GPT, BERT, Mistral, and Llama.

Our proposed method includes the development of a space-efficient Boolean indicator to handle the non-invertible nature of these functions, enabling a practical implementation within the PyTorch framework.
Additionally, we introduced accurate approximations for the derivatives of the inverse functions of GELU and SiLU, ensuring that training performance remains unaffected while achieving substantial memory savings.

The experiments and approximations demonstrated the feasibility and effectiveness of our approach, with results showing that the training process using our modified nonlinearity layers is indistinguishable from using the original layers in terms of accuracy.

We set out to meet three key criteria for a successful replacement of standard nonlinearity layers: maintaining model speed, preserving model accuracy, and ensuring ease of use. Our proposed method satisfies all three conditions.
It does not slow down the neural network, it does not degrade model quality, and it is extremely easy to integrate, making memory savings essentially "free."

We have implemented this method for the PyTorch framework, and the code is available at [https://github.com/removed/for/anonimity](https://github.com/removed/for/anonimity).

## Appendix A Appendix

### A.1 Memory Table

We perform memory measurements for several popular neural network architectres from Hgggingface model repository: [https://huggingface.co/models](https://huggingface.co/models).
Memory saving percentage given in a table are calculated as how much less memory is used for stored activations, without taking model size and optimizer statistics in consideration.
The following models were used:

- •
BERT: [google-bert/bert-base-uncased](https://huggingface.co/google-bert/bert-base-uncased)
Measured for sequence length equals $1024$
- •
Audio Spectrogram Transformer: [MIT/ast-finetuned-audioset-10-10-0.4593](https://huggingface.co/MIT/ast-finetuned-audioset-10-10-0.4593)
Measured for $1024\times 128$ spectrogram shape
- •
ViT (Visual Transformer): [google/vit-base-patch16-224-in21k](https://huggingface.co/google/vit-base-patch16-224-in21k)
- •
CLIP: [openai/clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)
Measured for text length equals $77$ and number of image exaples equal to the number of text examples.

### A.2 Approximation Coefficients

Approximations Eqs. 6, 7 and 8 were built by hands.
Left part of GELU approximation Eq. 5 were buld with an assstance of PySR Cranmer () framework.
Parameters of approximation for GELU:

$$ $\displaystyle q^{\text{left}}:\;$ $c_{0}$ $+1.6311011311381$ $c_{1}$ $+0.16997246666667$ $c_{2}$ $-0.06261728$ $c_{3}$ $+1.2947087$ $c_{4}$ $+1.98055565$ $c_{5}$ $+0.22730362$ $c_{6}$ $-0.038978495$ $c_{7}$ $+1.3295193$ (18) $\displaystyle q^{\text{right}}:\;$ $c_{0}$ $-1.383717971214795$ $c_{1}$ $+1.558420184350027$ $c_{2}$ $+0.044045748018110$ $c_{3}$ $+0.032146736769376$ $c_{4}$ $-2.119885089843949$ (24) $$

Parameters of approximation for SiLU:

$$ $\displaystyle q^{\text{left}}:\;$ $c_{0}$ $-1.310856402130980$ $c_{1}$ $+0.848589647031652$ $c_{2}$ $-0.162990512595109$ $c_{3}$ $+0.002696163985044$ $c_{4}$ $-5.770613302664509$ (30) $\displaystyle q^{\text{right}}:\;$ $c_{0}$ $+0.217177007595768$ $c_{1}$ $-0.507684370508263$ $c_{2}$ $+0.079631397669175$ $c_{3}$ $+0.357494204859375$ (35) $$

### A.3 Experiment Setups

- •
Plain Nonlinearity
Batch size = $2^{25}$.
- •
Nonlinearity + Linear
Batch size = $2^{15}$, features dimension = $2^{10}$.
- •
MLP Block
Batch size = $2^{15}$, input/output dimension = $2^{10}$, hidden dimension = $4*2^{10}$.
- •
GeGLU Block
Batch size = $2^{15}$, input dimension = $2^{10}$, output dimension = $4*2^{10}$.
- •
BERT
Batch size = $64$, sequence length = $1024$
- •
Llama v3.1 8B
Batch size = $1$, sequence length = $512$

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. **Sample Size and Standard Deviation**: Are the results in Table 1 averaged over multiple trials? If so, what is the sample size and standard deviation?  If the results show metrics over only a single run, then reporting the results over multiple runs with reports on the standard deviations would better highlight robustness and reproducibility of the method.
2. **Precision Bit InvAct**: I am curious to know if you have performed any performance evaluation as seen in Sections 3.2 and 3.3 on Precision Bit InvAct. The results presented in Table 1 suggest this variant to have less overhead but it is not clear how much error it produces in downstream tasks.
3. **Clarification on Experimental Consistency**: For the FewBit comparison in Figure 5, was the same dataset used as in Section 3.2? If different datasets were used, specifying this in Section 3.3 would enhance clarity.
4. **Sequence Length in BERT Experiment**: In Appendix A.3, the experiment setup for BERT specifies a sequence length of 1024, which exceeds the model’s typical limitation of 512 tokens. Could you clarify how this was achieved? Was a specific technique used to expand the sequence limit to 1024? If so please share more details.
## Minor suggestions:
1. **Additional Details Needed in Listing 1**: In Listing 1, memory savings for the Llama model are not shown, despite Llama being used in Section 3 for performance analysis. Including this information would create a more comprehensive view of the memory savings achievable with this approach.
2. **Sentence Rewrite in Line 464**: In the sentence, "Former are used more in during inference," the word "in" is unnecessary. The corrected sentence should read:  
  > "Former are used more during inference."
soundness: 4
strengths: - **Originality**: The approach introduces a novel solution to the memory bottleneck in neural networks training. The proposed solution is straightforward yet impactful adjustment for deep learning workflows.
- **Clarity**: The method is clearly outlined, with pseudocode, detailed experimentation, and well written text.
- **Significance**: The reduction of memory footprint in training aligns well with current trends in model scaling. The method's compatibility with widely-used architectures and frameworks adds further value, especially for resource-constrained training environments.
summary: This paper proposes a method, called InvAct, to reduce the memory footprint associated with activation tensors during neural network training by modifying how activation tensors are saved for the backward pass in pointwise nonlinearity layers like GELU and SiLU.
Considering a block of two layers, instead of saving the input tensor for each layer, the authors suggest storing the output tensor of layer 1, which is the input tensor of layer 2. Then during the backward pass the input layer of layer 1 is reconstructed using the inverse of the nonlinearity. Since the nonlinearities considered in the paper do not have straightforward analytical inverses, the paper approximates the inverse functions using simpler functions. The paper shows great deal of memory savings for Transformer based architecture, around 25%, without noticeable impact on accuracy or computational efficiency, and can be used as a drop-in replacement for existing nonlinearity layers in PyTorch.
weaknesses: - **Experimental Scope**: The evaluation is limited to a few model architectures (BERT and LLama) and tasks (e.g., BERT fine-tuning on Yelp Reviews). Adding results from multiple tasks, such as [GLUE](https://huggingface.co/datasets/nyu-mll/glue), would strengthen claims of general applicability and robustness across varied scenarios.
- **Reporting of Experimental Details**: 
Including the fine-tuning details for the results reported in Sections 3.2 and 3.3, perhaps in the appendix could benefit the readers and help future research with reproducing the results reported in the paper.
- **Lack of Task-Specific Performance Metrics**: The results shown in Section 3.2 focus on training/validation loss without reporting task-specific metrics (e.g., accuracy, F1 score for Yelp Reviews). Including such metrics would better contextualize the impact of approximation errors on model quality.
- **Lack of Performance Evaluation for Other Models**: Although memory savings are reported for models like the Audio Spectral Transformer, ViT, and CLIP, no experimental analysis is provided to assess if or how these savings impact model performance. For instance, similar to the BERT and Llama evaluations in Sections 3.1 and 3.2, experiments could show if the method affects performance metrics and how much computational overhead is introduced for these models. Furthermore, in the abstract, the authors mention models like *GPT* and *Mistral* as potential beneficiaries of the method, yet there is no experimental data on memory savings, time overhead, or performance impact for these models. Including analysis or removing these model names would provide a more accurate representation of the paper's scope.

**P.S.**: I am more than willing to raise my score if the concerns mentioned in the **Weakness** and **Questions** section are addressed.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: - How were the coefficients $c_i$ obtained ?
- Why is there an increase in execution time for InvAct only for the Plain GELU? What about the execution time for SiLU?
- The Sign-bit Inverted activation should be explained more clearly.
- Figure 3 requires labels for the rows and columns to be more clear.

**Minor details**

- The abstract in the openreview summary has a formatting error.
- line 182 "based"
- line 701 "assstance"
soundness: 2
strengths: - The method proposed is simple and only requires a drop-in replacement of the activation functions, providing a major memory improvement.
- The proposed inversion methods are novel, notably the approximate inverses of the GELU and SiLU functions.
- The computational efficiency of the method is considered on many different frameworks.
- The paper is clear and well-presented.
summary: This article presents a modification to the GELU and SiLU activations to save the output tensor instead of the input, recomputing the input to save approximately a quarter of the memory overhead in practice. They provide several variants of their method, all revolving around the idea of dividing the activation function into monotonous parts which can be inverted. They measure extensively that their approach does not change the computation time and doesn't affect performance.
weaknesses: - The main method proposed has very limited novelty outside of the consideration of GELU and SiLU. Other approaches have already proposed inverting the computations of the activation function such as [1] for similar reasons reducing the memory by half and with further benefits. 
- At no point comparison with activation/gradient checkpointing is considered, despite the method being extremely similar. Indeed, why not simply store the input tensor only and recompute the output of the activation for efficient gradient computation? This is the same method, but it doesn't require an approximate inversion like proposed here. An extensive literature on checkpointing is necessary to consider here to compare the proposed method to, such as [2-4]. Notably, [5] has a particular consideration for the forward and backward of GeLU. It seems very hard to justify this method compared to more standard checkpointing methods, which computes an optimal trade-off between activation and computation time.
- The approximation error of the derivative of the inverse as represented in Figure 2 shows a non-negligible error for values of $x$ close to 0, which seems a bit worrying.
- Other activation functions are not considered. How is ReLU supposed to work in this framework?

[1] In-Place Activated BatchNorm for Memory-Optimized Training of DNNs, S. R. Bulò et al.

[2] Reducing Activation Recomputation in Large Transformer Models, V. Korthikanti et al.

[3] Efficient rematerialization for deep networks, Kumar et al. 

[4] Rockmate: an efficient, fast, automatic and generic tool for re-materialization in pytorch, Zhao et al.

[5] Transcending Runtime-Memory Tradeoffs in Checkpointing by being Fusion Aware, S. Yu et al.

---

confidence: 4
contribution: 4
initial_rating: 5
presentation: 3
questions: + The authors claim that the studied activations are chosen because they are popular choices for transformers, did the authors try to apply their methods to other activations / nonlinearities ?
+ There is a tipo at line 337 ("implementatoins")
soundness: 2
strengths: + The proposed method is derived from a very simple idea, yet it leads to substantial reduction in memory costs, without introducing significant drawbacks or trade-offs (in particular it has no significant impact on either the training time or the model's performance)
+ The method is explained in a very clear way, as well as the research objectives of the paper
+ Thanks to its simplicity, the method is likely implementable within existing frameworks, without changing their core functioning
summary: This paper derives a method that reduces the amount of data that has to be stored in memory during the forward pass for the backpropagation. By doing so, it is possible to save a significant amount of memory during the training of a neural network. The proposed method is easily implementable and can be deployed in a training procedure with almost no impact
+ on the training duration,
+ on the overall performance of the model,

as it is supported by experimental results.
weaknesses: The main weeknesses of the paper lie in its experimental section:
+ Some claims of the paper are not supported by experiments, such as line 60: "effectively reducing the memory footprint by nearly $25\\%$ in practice", it would be nice to see empirical evidence of such a claim
+ More generally, no experiment on the memory footprint achieved by the proposed method is shown, yet since the main objective of the paper is apparently to reduce the memory footprint during training, the paper would benefit greatly from such results
+ In section 3.1, the authors use two different frameworks to assess the computational efficiency of their method, but it seems thus difficult to have a clear view on whether the observed gap in computational efficiency is due to the authors method, or to the fact that two different frameworks are being used
+ Finally, could the authors add a theoretical example of how in practice the proposed method would allow to discard some tensors for back propagation computation would make the contribution of the paper more striking

Meta-review: The paper introduces, a memory-saving method for neural network training that leverages approximate inverse functions for pointwise nonlinearities like GELU and SiLU. By saving the output tensor instead of the input tensor and reconstructing the input during the backward pass, InvAct reduces memory usage in transformer-based architectures without significant impact on model accuracy or computational efficiency. The method is simple, generalizable to popular frameworks like PyTorch. Reviewers pointed out and as acknowledged in the limitations that the technical novelty is limited and narrow in scope as it applies to only a specific set of activations. Moreover, the interaction and comparisons with existing methods for memory saving like activation checkpointing is not fully explored. Overall the work is incremental and limited in scope and may be more fitted for another venue.

Decision: Reject