Inverted Activations: Reducing Memory Footprint in Neural Network Training

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

Consider an example of a Transformer MLP block with two linear layers and a nonlinearity $f$ between them. For the case of the standard nonlinearity $f = \text{GELU}$, the forward and backward passes would look like this:

Forward pass:

1. $I$ – input of the MLP block.
2. $X = \text{Linear}_1(I) = I W_1$ – application of the first linear layer. Tensor $I$ is saved for the backward pass.
3. $Y = f(X)$ – application of the GELU nonlinearity. Tensor $X$ is saved for the backward pass.
4. $O = \text{Linear}_2(Y) = Y W_2$ – application of the second linear layer. Tensor $Y$ is saved for the backward pass.

In total, three tensors are saved for the backward pass: $I$, $X$, and $Y$.

Backward pass:

1. $\frac{\partial L}{\partial O}$ – gradient of the loss with respect to tensor $O$, used as input for backward computation.
2. $\frac{\partial L}{\partial Y} = \text{Linear}_2\textrm{-backward}\left(\frac{\partial L}{\partial O}\right) = \frac{\partial L}{\partial O} W_2^T$ and $\frac{\partial L}{\partial W_2} = Y^T \frac{\partial L}{\partial Y}$: computation of gradient of loss w.r.t. tensor $Y$ and gradient of second linear layer matrix.
3. $\frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y} \cdot f'(X)$ – backward pass for nonlinearity. Here is the only place where we use saved activation tensor $X$.
4. $\frac{\partial L}{\partial I} = \frac{\partial L}{\partial X} W_1^T$ and $\frac{\partial L}{\partial W_1} = I^T \frac{\partial L}{\partial X}$

Now consider the situation with InvAct nonlinearity:
Forward pass:

• $I$ – input of the MLP block.
• $X = \text{Linear}_1(I) = I W_1$ – application of the first linear layer. Tensor $I$ is saved for the backward pass.
• $Y = f(X)$ – application of the nonlinearity. Instead of saving tensor $X$, a compressed boolean tensor $S = X < \text{const}$ is saved, where $\text{const}$ is a globally defined constant real value. For Precision-bit InvAct, $S$ is saved inside lower precision bit of $Y$ itself.
• $O = \text{Linear}_2(Y) = Y W_2$ – application of the second linear layer. $Y$ is saved for the backward pass.

In total, we save two tensors ($I$, $Y$), and instead of the full tensor $X$, only the heavily compressed boolean tensor $S$ is saved (if $X$ has data type float 32, then tensor $S$ consumes 32 times less memory). It is roughly 33% less memory in total.

Backward pass:

• $\frac{\partial L}{\partial O}$ – gradient of the loss with respect to tensor $O$, used as input for backward computation.
• $\frac{\partial L}{\partial Y} = \text{Linear}_2\textrm{-backward}\left(\frac{\partial L}{\partial O}\right) = \frac{\partial L}{\partial O} W_2^T$ and $\frac{\partial L}{\partial W_2} = Y^T \frac{\partial L}{\partial Y}$.
• $\frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y} \cdot f'\left(f^{-1}(Y, S)\right)$ – here, instead of tensor $X$, which is not saved,we use tensors $Y$ and $S$.
• $\frac{\partial L}{\partial I} = \frac{\partial L}{\partial X} W_1^T$, and $\frac{\partial L}{\partial W_1} = I^T \frac{\partial L}{\partial X}$.

Comparison: The standard nonlinearity saves three tensors ($I$, $X$, $Y$), whereas InvAct saves two $I$, $Y$, and a compressed boolean tensor $S$. It is roughly 33% less memory.

In section 3.1, the authors use two different frameworks to assess the computational efficiency of their method.

For all our experiments, we use only one framework: PyTorch. The PyTorch implementation of InvAct is done with Triton kernels, while PyTorch GELU is implemented in CUDA. However, this is a standard practice because Triton is simply a more user-friendly version of CUDA.

In our time measurements, we carefully ensure that the only difference in the execution of computational blocks (e.g., MLP and GeGLU) lies in the pointwise nonlinearity. All other components are computed identically. Specifically, we use NVIDIA Nsight Compute to verify which kernels are executed on the GPU.

Some claims of the paper are not supported by experiments.

More generally, no experiment on the memory footprint achieved by the proposed method is shown.

All claims about memory consumption are based on experimental measurements, and specific values can be found in the table in the Introduction section under Listing 1. Memory measurements were conducted as follows:

1. Instantiate the model.
2. Measure the memory consumed by the model parameters.
3. Perform a forward pass.
4. Measure the currently allocated memory and subtract the model memory to get the memory allocated for activations.

This procedure is then performed for both the standard nonlinearity and the InvAct nonlinearity. The decrease in the memory consumed by activations is reported in the Table in Introduction section.

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?**

Most popular computer vision nonlinearities, such as ReLU and Sigmoid, do not require our method because their derivatives can be expressed directly in terms of the output tensors:

These are already implemented optimally in PyTorch.

We did not test our approach on other nonlinearities, but there is no fundamental reason why it would not work:

• Monotonic nonlinearities like SELU require only an appropriate approximation of $f'\left(f^{-1}(y)\right)$ and does not require additional bit per element for inversibility.
• Unimodal functions can be processed analogously to SiLU and GELU.
• Nonlinearities with multiple monotonic parts might require more bits per element, potentially reducing the method's applicability.

If there is a specific nonlinearity you would like us to address, we would be happy to explore it.

Other activation functions are not considered. How is ReLU supposed to work in this framework?

The ReLU nonlinearity, $y = ReLU(x)$, does not require additional memory-reduction methods, as its derivative $f′$ can be directly derived from $y$:

Similarly, for example, for sigmoid nonlinearity, $y =σ(x)$, the derivative can be expressed as $y (1−y)$. For both these nonlinearities PyTorch already does not save input tensor x. That is why we did not conduct experiments on ResNets, as ReLU is the dominant nonlinearity function in this family of models.

In this work, we focus on the most commonly used nonlinearities in transformer-based architectures: GELU and SiLU. However, our method can be extended to other nonlinearities:

• Monotonic nonlinearities like SELU require only an appropriate approximation of $f^′(f^{−1}(y))$ and do not need an additional bit per element.
• Any unimodal function can be processed analogically to SiLU and GELU
• Nonlinearities with multiple monotonic parts may require more bits per element, making the method less applicable.

If there is a specific nonlinearity you would like us to address, we would be happy to explore it.

How were the coefficients obtained?

Coefficients for both parts for SiLU and the right part for GELU were obtained using scipy.minimize to minimize the $L_2$ distance between the approximation $q(y)$ and the target function $f^′(f^{−1}(y))$. The symbolic form of these parts was constructed manually. The left part for GELU was constructed and minimized using the PySR framework (https://github.com/MilesCranmer/PySR).

Why is there an increase in execution time for InvAct only for the plain GELU? What about the execution time for SiLU?

Nonlinearity operations are significantly faster than linear layers, so the minor overhead introduced by InvAct becomes more noticeable when benchmarked in isolation (i.e., for plain GELU) compared to when it is embedded in computational blocks or complete neural network architectures. Additionally, when benchmarking the time for plain nonlinearity (e.g., y = f(x)), both y and x must be calculated and stored in memory, negating the memory-saving benefit.

In contrast, when InvAct is benchmarked inside larger computational blocks, memory savings can lead to more optimal memory access patterns, reducing the relative overhead.

The Sign-bit Inverted activation should be explained more clearly.

Thank you for pointing that out. We will improve the description of the Sign-bit Inverted activation in the revised manuscript. For now, here is more detail: to disambiguate the inverse of the nonlinearity function, we must store an additional boolean value $s = x<T$ along with the nonlinearity output $y = f(x)$. This boolean value requires only one bit of information per element of tensor $x$. The Sign-bit approach uses the lowest precision bit inside $y$ itself. $y$, being a floating-point value, has the following representation:

Image: link

The lowest precision bit is bit number 0 in this representation. We can replace its original value with $s$, introducing slight precision error during the forward pass. Sacrificing one precision bit does not significantly impact results, as both float16 and float32 formats provide sufficient precision. Our experiments show no measurable difference in outcomes.

Figure 3 requires labels for the rows and columns to be clearer.

Thank you for the suggestion. We will improve Figure 3 by adding appropriate labels for the rows and columns.

At no point comparison with activation/gradient checkpointing is considered

Achieving the same results as our method via checkpointing would require applying it to the pair of layers Activation + Following Linear Layer, which has several downsides:

• Performance Overhead: Recomputing GELU during the backward pass adds additional computational cost. Consider, for example, measurements for Transformer MLP Block with standard nonlinearity, InvAct nonlinearity and checkpointed nonlinearity::

Image: link

Benchmark shows that checkpointing introduces an overhead of $0.79$%, compared to $0.06$% for our method.

• Complexity: Checkpointing requires modifications to model architectures and possibly forward computation code, making it less user-friendly. In contrast, our method functions as a drop-in replacement for standard nonlinearities.

Making automatic fusing of such operations and constructing the required frameworks is an intense, ongoing area of research. For example: link. Even advanced approaches like [5] (if we understood it correctly) focus only on specific cases of GELU fusing, and they do not generalize to more complex setups involving subsequent linear layers. This limitation arises because PyTorch’s compile backend (which is used in work [5]) does not yet support such intricate fusions. Situation will be even worse if the output of nonlinearity layer is used in several branches of computation: then either such computations all have to be fused separately, further multiplying number of forward-pass GELU recalculations, or not fused at all, which increase number of memory accesses required to perform backward pass. That is not the case for our InvAct nonlinearity approach.

Instead, more widespread applications of checkpointing is when it is used over the entire network. In these cases, our method can be employed in tandem with gradient checkpointing. By reducing the memory footprint of individual nonlinearities, our approach allows methods like [4, 5] find better checkpointing schedules.

In summary: Specialized "small" checkpointing for individual nonlinearities is slower than our method and requires manual implementation, meaning it cannot function as a drop-in replacement. For traditional network-wide checkpointing, our approach can be applied simultaneously, yielding additional memory savings and making larger models feasible.

We will add discussion dedicated to checkpointing to manuscript, thank you!

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.

Please note that in Figure 2, there is a separate $y$-axis on the right of each plot and is scaled by a common multiplier, indicated in the upper-right corner. For example, for GELU, the absolute error has a magnitude of $1e−2$. Our experiments in Sections 3.2 and 3.3 demonstrate that this level of error does not affect training. We also conducted additional experiments using the GELU benchmark with consistent results. Please see our response to Reviewer jQ9m for further details: link.

Could the authors please answer if the execution time is similar for SiLU to that for GELU?

Execution times for SiLU are the same as for GELU. Here are the exact value:

The GeGLU component uses SiLU as a gating unit.

Thank you for your active feedback. Can you provide some additional comments that we can addresses so that you can consider raising your score?

In our work, we aim to demonstrate that approximation errors have no measurable impact on final performance metrics. However, achieving this goal involves challenges:

1. Statistical significance: Proving the absence of a difference is inherently more difficult than proving its presence. This requires conducting numerous runs for credibility.
2. Comprehensive evaluation: There will always be "one more" benchmark or model architecture to test.

To address these challenges, we focused on demonstrating that approximation errors are so minor that training losses closely follow those of the original training process — both for training and validation losses and for task-specific metrics.

Nevertheless, we performed additional experiments on the GLUE benchmark using the RoBERTa-base model (based on the configuration taken from here). For each model (standard nonlinearity and InvAct), we conducted 16 runs per GLUE dataset. Below, you can see the corresponding metrics for each sub-dataset:

Image: https://ibb.co/pxR4Gdg

Result: InvAct matches the standard nonlinearity perfectly.

Statistical test: A Mann–Whitney U test on average metric values shows no statistically significant difference (p-value = 0.89).

We also conducted additional experiments on the Recognizing Textual Entailment (RTE) dataset: We ran 20 runs of 1-bit FewBit backward and compared it with 20 runs of the standard nonlinearity. The Mann–Whitney U test shows that 1-bit FewBit performs statistically significantly worse than the standard nonlinearity (p-value = 0.027). In contrast, InvAct does not exhibit a statistically significant difference. To confirm this, we performed an additional 128 runs for both the standard nonlinearity and InvAct. Even with 128 samples for each (6 times more than for 1-bit FewBit), there was no statistical significance (p-value = 0.7). The respective boxplots are provided below:

Image: https://ibb.co/MN6zW3d

For completeness, we also provide training losses and task-specific metrics averaged across 16 runs to further validate that InvAct’s minor approximation error do not affect the training process:

Image: https://ibb.co/2WgcHxJ

Image: https://ibb.co/nsRB1hQ

Moreover, as discussed in Section 3.3, other methods report no loss of quality for 4-bit quantization. We demonstrate that our approximation error is smaller than even 8-bit quantization errors. This further supports our claim that InvAct has no noticeable influence on neural network training, which is inherently noisy, stochastic, and robust to larger perturbations.

Sample Size and Standard Deviation in Table 1

Yes, all measurements in Table 1 are averaged across several runs. We ensured enough measurements were taken so that the standard error of the mean is several orders of magnitude lower than the differences between the compared values. To simplify table readability, we decided to omit the error levels. We will add this information to the table caption—thank you for pointing this out! Here are the details:

• Plain GELU, Linear + GELU, and MLP Block: each was averaged across 1000 runs, with standard errors of the mean below 2.5e-4, 1e-3, and 3e-3, respectively.
• GeGLU: Averaged across 100 runs, with a standard errors below 3e-3.
• BERT and Llama: Each averaged across 10 runs, with a standard error below 1e-1.

Precision Bit InvAct

The results for Precision Bit InvAct in our experiments were identical to those of the standard InvAct implementation. We attribute this to the fact that full-precision floating-point numbers provide sufficient precision for pointwise nonlinearities, specifically, and neural network training in general. This sufficiency explains why reducing floating-point precision is viable in the first place. Thus, one less precision bit in Precision Bit InvAct does not lead to performance degradation.

As stated in Section 2.1, this may change under more aggressive training settings involving extremely low-precision data types, like float8. For this reason, we concentrated on the standard InvAct implementation and reserved the Precision Bit variant for future, more rigorous testing.

Clarification on Experimental Consistency

Yes, the comparison with the FewBit backward approach in Section 3.3 was performed using the same Llama v3.1 experimental setup as in Section 3.2.

Sequence Length in BERT Experiment

The sequence length used was 512. The value of 1024 in the paper was an error. Thank you very much for noticing this — it has been corrected in the manuscript.

Model Choice

Transformer-based models are a dominant architecture type, often using GELU and SiLU instead of plain ReLU. For this reason, we selected BERT as a representative small model and Llama v3.1 8B as a representative large model. We believe these two examples adequately reflect the behavior of other transformer-based models, which would likely show very similar results.

Furthermore, we argue that the results presented in the experimental sections are task-independent. Since the minimal approximation error does not influence the training loss, it should not affect other types of losses and, consequently, task-specific metrics either.

Computational Overhead for Other Models

We appreciate your suggestion to provide computational overhead analysis for the models mentioned in the Introduction. We will add this analysis for completeness—thank you for highlighting this!

At present, we include computational overhead analysis for fundamental building blocks, such as the Transformer MLP block and GeGLU, in Section 3.1. These results can be easily extrapolated to other models not explicitly analyzed in our work. The computational overhead is consistently less than 0.1%, a level that will be dominated by other network components. This means that, for any full-scale model, the computational overhead is effectively negligible.

The same holds for ViT, CLIP, Mistral, and all other models we tested. Our method introduces no computational overhead on practice, making it a practical and efficient solution for memory savings in transformer-based architectures.

Thank you very much for the suggested comments. We will correct them in the manuscript. You wrote that your score could be increased if we answer the questions. Is there anything else that we haven't answered or haven't answered fully?