Paper: ALAM: Averaged Low-Precision Activation for Memory-Efficient Training of Transformer Models (ICLR-2024)

Reviewer comment (Experiments):
The evaluation part is not sufficient. This paper only compares the result with GACT, which is based on activation compression. However, for model compression, other techniques like pruning and low-rank compression are widely used.

Author response:
We would like to thank the reviewer for bringing up an important point. Following the reviewer’s advice, we conducted detailed comparisons between ALAM and other model compression methods. For comparisons, we selected LoRA [1], which is the state-of-the-art low-rank compression method. We also implemented QLoRA [2], which performs 4-bit quantization for weights and only computes weight gradient for LoRA parameters. QLoRA allows for a 4× reduction in the weights with comparable accuracy, making it widely used for fine-tuning LLMs in constrained environments. The results are presented in Table C.1. Please note that ‘average accuracy’ represents the average accuracy over all tasks in Tables A.1 and A.2 in the global response. Also, ‘param mem’ represents the parameter memory occupied by weights, gradients, and optimizer states during training.

**Table C.1. Comparisons with model compression techniques in fine-tuning LLaMA2-7B and LLaMA2-13B on the LIMA dataset. The micro batch size is set to 8.**

| Model      | Strategy       | Average accuracy | Param mem | Act mem | Total mem |
|------------|----------------|------------------|-----------|---------|-----------|
| LLaMA2-7B  | Baseline       | OOM              |   52 GB   |  21 GB  |   73 GB   |
|            | + LoRA         | 62.5             |   13 GB   |  21 GB  |   34 GB   |
|            | + LoRA + ALAM | 62.2             |   13 GB   |  2.0 GB | **15 GB** |
| LLaMA2-13B | Baseline       | OOM              |   104 GB  |  33 GB  |   137 GB  |
|            | + QLoRA        | 65.1             |   6.6 GB  |  33 GB  |   40 GB   |
|            | + QLoRA + ALAM | 65.0             |   6.6 GB  |  3.2 GB | **10 GB** |

In experiments, LoRA noticeably reduces the memory allocated for gradients and optimizer states, which results in a 4× compression of parameter memory during training. Additionally, QLoRA exhibits a large compression rate of 16× by utilizing 4-bit weights. However, neither LoRA nor QLoRA compresses activation memory. Our ALAM, on the other hand, focuses on compressing activation memory, and hence achieves 10.6× and 10.3× compression rates on LLaMA2-7B and LLaMA2-13B, respectively. As a result, compared to when only LoRA and QLoRA were applied, additionally employing ALAM further saves 56% and 75% of the total memory with comparable accuracy for LLaMA2-7B and LLaMA2-13B, respectively. In summary, our ALAM is orthogonal to other model compression techniques, and they can be applied simultaneously for maximum memory savings.

Please note that only a handful of pruning schemes supporting very large models have been reported in the literature, with limited memory savings. For instance, LLM-Pruner [3] exhibits a significant drop in accuracy even when pruning just 20% of the parameters in LLaMA-7B, translating to similar saving only in the parameter memory.

Reviewer comment (Experiments):
How to choose the parameter group is unclear. Is this parameter related to the sensitivity? What if the adjacent activations have little similarity?

Author response:
We thank the reviewer for raising an important issue, and we apologize for the confusion. As described in Sections 3.1 and 3.2, we group adjacent activations together after simply flattening the activations, without any other complicated processes. As the reviewer insightfully pointed out, varying the group size (g) and the precision of the group average (n) could produce different bit precisions (AQ-n/g bit). However, in all experiments using Average Quantization, we fixed the group size to 4 and performed 2-bit quantization on the group average value to obtain AQ 0.5-bit, which is an optimal configuration identified in experiments, as shown in Table C.2. 

**Table C.2. Performan comparisons of different AQ 0.5-bit configurations when training VGG-11 on CIFAR-100 (n: precision of group average, g: group size).**

| Precision        | FP32     | 1-bit |  |          |    AQ 0.5-bit    |         |          |
|------------------|----------|-------|------------|----------|--------|---------|----------|
| AQ scheme [n,g]  | -        | -     | (1, 2)     | (2, 4)   | (4, 8) | (8, 16) | (32, 64) |
| Accuracy (%)     | **67.7** | 1.0   | 1.0        | **52.3** | 1.87   | 8.07    | 3.08     |

The group size is not related to sensitivity. More specifically, the sensitivity of each layer is first calculated using the GradNormVar algorithm, which determines the layer’s bit precision. In the layers with low sensitivity, we quantize activations into AQ 0.5-bit, where average quantization is performed with a fixed group size of 4 as discussed above.

The reviewer is certainly correct that adjacent activations may show little similarity. This could be mitigated by only grouping activations with similar distributions, but we did not implement this in our algorithm since this process could be time-consuming and require additional memory to save the mask for recovery. Our simple grouping scheme may lead to some loss of information in activations, but the impact on actual accuracy is minimal because ALAM assigns AQ 0.5-bit only to layers with very low sensitivity. Furthermore, the quality of AQ 0.5-bit activations is still higher than that of 1-bit quantized activations, as shown in Table C.2.

Reviewer comment (Experiments):
It would be better to provide the evaluation result with other compression techniques.

Author response:
We would like to thank the reviewer for bringing up an important point. Following the reviewer’s advice, we conducted detailed comparisons between ALAM and other model compression methods. For comparisons, we selected LoRA [1], which is the state-of-the-art low-rank compression method. We also implemented QLoRA [2], which performs 4-bit quantization for weights and only computes weight gradient for LoRA parameters. QLoRA allows for a 4× reduction in the weights with comparable accuracy, making it widely used for fine-tuning LLMs in constrained environments. The results are presented in Table C.1. Please note that ‘average accuracy’ represents the average accuracy over all tasks in Tables A.1 and A.2 in the global response. Also, ‘param mem’ represents the parameter memory occupied by weights, gradients, and optimizer states during training.

**Table C.1. Comparisons with model compression techniques in fine-tuning LLaMA2-7B and LLaMA2-13B on the LIMA dataset. The micro batch size is set to 8.**

| Model      | Strategy       | Average accuracy | Param mem | Act mem | Total mem |
|------------|----------------|------------------|-----------|---------|-----------|
| LLaMA2-7B  | Baseline       | OOM              |   52 GB   |  21 GB  |   73 GB   |
|            | + LoRA         | 62.5             |   13 GB   |  21 GB  |   34 GB   |
|            | + LoRA + ALAM | 62.2             |   13 GB   |  2.0 GB | **15 GB** |
| LLaMA2-13B | Baseline       | OOM              |   104 GB  |  33 GB  |   137 GB  |
|            | + QLoRA        | 65.1             |   6.6 GB  |  33 GB  |   40 GB   |
|            | + QLoRA + ALAM | 65.0             |   6.6 GB  |  3.2 GB | **10 GB** |

In experiments, LoRA noticeably reduces the memory allocated for gradients and optimizer states, which results in a 4× compression of parameter memory during training. Additionally, QLoRA exhibits a large compression rate of 16× by utilizing 4-bit weights. However, neither LoRA nor QLoRA compresses activation memory. Our ALAM, on the other hand, focuses on compressing activation memory, and hence achieves 10.6× and 10.3× compression rates on LLaMA2-7B and LLaMA2-13B, respectively. As a result, compared to when only LoRA and QLoRA were applied, additionally employing ALAM further saves 56% and 75% of the total memory with comparable accuracy for LLaMA2-7B and LLaMA2-13B, respectively. In summary, our ALAM is orthogonal to other model compression techniques, and they can be applied simultaneously for maximum memory savings.

Please note that only a handful of pruning schemes supporting very large models have been reported in the literature, with limited memory savings. For instance, LLM-Pruner [3] exhibits a significant drop in accuracy even when pruning just 20% of the parameters in LLaMA-7B, translating to similar saving only in the parameter memory.

Reviewer comment (Evaluation):
In Table 1, there is an accuracy drop from 98.72 to 96.09 in a small model. However, in the evaluation result, the compressed model retains the accuracy performance. Can you identify the reason for this? This is because the large model has more redundancy or other reasons?

Author response:
We apologize for the confusion. Please note that the experiments in Table 1 did not fully implement our ALAM algorithm; Table 1 simply compares the quality of activations obtained by different compressing methods when the group size is 32 and the group average is not quantized. Also, we assigned identical precision to all layers in this experiment, whereas ALAM assigns different precision to the layers based on their sensitivity. This caused a relatively large accuracy drop in Table 1. 

To fully address the reviewer's concern, we conducted additional experiments of training the same MLP model on MNIST using ALAM, and the experimental results are displayed in Table C.3 below. The results confirm that ALAM successfully trains the MLP model with comparable accuracy to the baseline (FP32). We sincerely appreciate the reviewer’s comments, which have helped us identify unclear aspects of our paper. We will clarify these points in the camera-ready version if the paper is accepted.

**Table C.3. Test accuracy of MLP trained on MNIST with different activation approximation methods.**
|  | FP | Simple group averaging (g=32) | ALAM |
|---|---|---|---|
| Accuracy (%) | 98.72 | 96.09 | 98.60 |

Reviewer comment (Experiments):
It would be better to add QLoRA as a baseline for evaluation.

Author response:
We thank the reviewer for this great suggestion. Following the reviewer’s comment, we applied QLoRA [1] and evaluated our algorithm on larger models: LLaMA2-13B and LLaMA-30B. Please note that we used the recently released LLaMA2 for a 13B model, which has the same structure as LLaMA but is pre-trained on a broader range of datasets, leading to improved accuracy. On the other hand, LLaMA2-30B is not publicly available, and hence LLaMA-30B was employed instead. In this experiment, we fine-tuned the model for 5 epochs on the LIMA dataset. The results are displayed in Table B.1. In the table, 'precision' refers to the target average precision of activation during training as detailed in the manuscript, not the precision of the weight.

**Table B.1. Test accuracy, activation memory (ACT mem) with compression, and memory for calculating sensitivity (sens mem) in fine-tuning LLaMA2-13B and LLaMA-30B on the LIMA dataset. The micro batch size is set to 2.**

| Model       | PEFT  | Scheme   | Precision | Act mem            | MMLU     | Arc-c    | PIQA     | Hellaswag | WinoGrande | BoolQ    | TruthfulQA |
|-------------|-------|----------|-----------|--------------------|----------|----------|----------|-----------|------------|----------|------------|
| LLaMA2-13B  | QLoRA | Baesline | 16-bit    | 8.2 GB             | **55.0** | **50.3** |   80.0   |  **79.2** |    71.5    |   80.9   |  **39.1**  |
|             |       | ALAM     | 1-bit     | **0.8 GB (10.3x)** | 54.8     |   50.0   | **80.2** |    78.8   |  **71.6**  |  **81**  |    38.9    |
| LLaMA-30B   | QLoRA | Baesline | 16-bit    | 15.8 GB            | 56.6     |   52.4   | **81.4** |  **82.1** |    75.0    | **82.7** |  **42.8**  |
|             |       | ALAM     | 1-bit     | **1.6 GB (9.9x)**  | **56.7** | **52.7** |   80.8   |    81.8   |  **75.3**  |   82.5   |    42.5    |

Experimental results demonstrate that our ALAM (1-bit) compresses activation by 10.3× in LLaMA2-13B and 9.9× in LLaMA-30B while closely matching the baseline in accuracy for all tasks. We will include these updated experimental results in the camera-ready version if the paper is accepted.

[1] Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, & Luke Zettlemoyer (2023). QLoRA: Efficient Finetuning of Quantized LLMs. CoRR, abs/2305.14314.

Reviewer comment (Experiments):
It is unclear how the proposed ALAM performs on larger models, such as LLaMa-13B, LLaMa-30B, etc. It would be better to show the results on these larger models.

Author response:
We thank the reviewer for this great suggestion. Following the reviewer’s comment, we applied QLoRA [1] and evaluated our algorithm on larger models: LLaMA2-13B and LLaMA-30B. Please note that we used the recently released LLaMA2 for a 13B model, which has the same structure as LLaMA but is pre-trained on a broader range of datasets, leading to improved accuracy. On the other hand, LLaMA2-30B is not publicly available, and hence LLaMA-30B was employed instead. In this experiment, we fine-tuned the model for 5 epochs on the LIMA dataset. The results are displayed in Table B.1. In the table, 'precision' refers to the target average precision of activation during training as detailed in the manuscript, not the precision of the weight.

**Table B.1. Test accuracy, activation memory (ACT mem) with compression, and memory for calculating sensitivity (sens mem) in fine-tuning LLaMA2-13B and LLaMA-30B on the LIMA dataset. The micro batch size is set to 2.**

| Model       | PEFT  | Scheme   | Precision | Act mem            | MMLU     | Arc-c    | PIQA     | Hellaswag | WinoGrande | BoolQ    | TruthfulQA |
|-------------|-------|----------|-----------|--------------------|----------|----------|----------|-----------|------------|----------|------------|
| LLaMA2-13B  | QLoRA | Baesline | 16-bit    | 8.2 GB             | **55.0** | **50.3** |   80.0   |  **79.2** |    71.5    |   80.9   |  **39.1**  |
|             |       | ALAM     | 1-bit     | **0.8 GB (10.3x)** | 54.8     |   50.0   | **80.2** |    78.8   |  **71.6**  |  **81**  |    38.9    |
| LLaMA-30B   | QLoRA | Baesline | 16-bit    | 15.8 GB            | 56.6     |   52.4   | **81.4** |  **82.1** |    75.0    | **82.7** |  **42.8**  |
|             |       | ALAM     | 1-bit     | **1.6 GB (9.9x)**  | **56.7** | **52.7** |   80.8   |    81.8   |  **75.3**  |   82.5   |    42.5    |

Experimental results demonstrate that our ALAM (1-bit) compresses activation by 10.3× in LLaMA2-13B and 9.9× in LLaMA-30B while closely matching the baseline in accuracy for all tasks. We will include these updated experimental results in the camera-ready version if the paper is accepted.

[1] Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, & Luke Zettlemoyer (2023). QLoRA: Efficient Finetuning of Quantized LLMs. CoRR, abs/2305.14314.