Network Memory Footprint Compression Through Jointly Learnable Codebooks and Mappings

It is hard to understanding the background part, especially the challenges of conventional approach. It would be better to provide a straightforward illustration with a figure or algorithm.

We introduced a novel figure to highlight these elements. First, we illustrate the challenges of the baseline approach and second, our contributions both in terms of multiple codebooks/scalings and gradient updates.

Some typos. In the beginning of sec3.2, “the” should be “The“; “ram” in sec 3.1 should be “RAM”; Something is missing in Equation 3.

Thank you for pointing these out. We fixed these typos in the revision version and doubled checked the remainder of the article.

The metric in evaluation part is accuracy and compression ratio. Can the proposed idea bring more benefits like accelerating the inference/training/finetuning?

the primary aim of the proposed JLCM method is to preserve the accuracy of the model while significantly reducing its memory footprint. However, as illustrated in Table 2, memory limitations can force the weights to be loaded on-the-fly for inference, leading to severe computational burden: this is especially true on large models (see last row in Table 2 for instance), as the architectures become more and more parameter hungry.

In this context, using JLCM can translate in huge latency improvements at inference time, e.g. on low-power or low-memory hardware where the original model cannot fit entirely, or if we consider very large models (e.g. LLMs).

Secondly, to answer your question, and although not studied in the present paper, one could decide to only optimize the codebooks during a post-hoc fine-tuning phase. In this case ,the proposed method would behave similarly to a foldeable adapter (e.g. LoRa [1], which can allow to finetune or fewer examples or using less memory), which could be an interesting lead for future work.

It would be better to explain why the second term $\log(\Omega(C))\Omega(W)$ dominates. It seems a common sense but I cannot see it in the context.

To answer you concern, we clarified this point in the revised version of the paper. In short, for a dense layer of weights $W \in \mathbb{R}^{n_i \times n_o}$, we have $\Omega(W) = n_i \times n_o$ which is often larger than $100000$ and can grow to several millions in some LLMs. On the other end, $\log(\Omega(C))$ corresponds to the size of our codebooks which is empirically bounded by $\Omega(C) \leq 16$, i.e. we have $\log(\Omega(C)) \leq 4 \ll 100000 \leq \Omega(W)$.

Can you provide some evaluation results in inference or training or finetuning performance in efficiency (throughput or energy efficiency)?

In terms of latency, we refer to Table 2 in our original submission (see also the answer to your previous question, in short JLCM has the potential to increase the latency in the cases where the model is too big to be loaded on device at once, and memory swaps have to be done). This was shown in SQueeze LLM [2]. With respect to throughput, we share these new results

In both cases, we measure the impact of a model not fitting on the device cache (e.g. the V-RAM for a GPU) and thus needing data transfers from disk. This occurs when a model is too large to be fully loaded on the device cache (e.g. V-RAm for a GPU). In such circumstances, we need to offload some paramters on the disk and load them when necessary (this process is optimized with the accelerate library from huggingface). In this experiment, we imitate the situation where the GPU has 2GB of VRAM (simulating the memory on a smartphone). Still, having the model fully loaded leads to massive throughput benefits. This can be acheived through memory footprint compression as in JLCM. We hope that these results further motivate the relevance of the proposed JLCM method.

references

[1] Hu, Edward J., et al. "Lora: Low-rank adaptation of large language models." arXiv preprint arXiv:2106.09685 (2021).

[2] Kim, Sehoon, et al. "SqueezeLLM: Dense-and-Sparse Quantization." arXiv preprint arXiv:2306.07629 (2023).

We thank the reviewer for their interest in the method and we appreciate the comment on readibility, especially provided this is not their exact domain of expertise.

The neuron permutations you use in your method remind me of the channel permutations that are used by N:M sparsification methods [1]. Drawing this connection could be interesting in your related work section.

We agree that our paper would benefit from this and other references on weight permutation. They were added directly in the corresponding section in the methodology.

In a number of places you use “Go” as a unit - is this intentional? Was “Go” supposed to be “GB”?

This was indeed a typo that was corrected in the revised version.

In Table 2, I’m curious to understand how you implemented offloading of parameters to disk. Are transfers from disk to GPU pipelined with computation to hide as much transfer latency as possible? What batch sizes were used for each model?

The implementation of this offloading can be achieved using the accelerate library from huggingface, from our experiments it is well optimized in order to hide as much as possible the cost of data transfers.

Regarding the batch-size, we measure the latency which is defined using a batch size of 1 (table 2). On the flip side, if we want to maximize the throughput (number of predictions per second with optimal batch size = 32), we observe very similar results (especially on larger models). In this experiment, we imitate the situation where the GPU has 2GB of VRAM (simulating the memory on a smartphone).

We hope that these answers address your questions.

We would like to thank you for your interest in the proposed method and for highlighting the empricial benefits offered by the proposed approach. In the following, we share our responses to your concerns.

The settings for the number of codebooks and the number of scaling factors should be further explained. Although the authors mention in the paper that these values are determined by compression goal, the process needs more details and therelationship with the weight distribution needs to be explained as well.

In our original submission, we define the memory footprint of a weight tensor $W \in \mathbb{R}^{n_i \times n_o}$ based on the cost of encoding a single scalar value (in 16 bits floating point) mutliplied by the number of weights (end of introduction of section 3.2), i.e.

$$M = \Omega(W) \times 16 = n_i \times n_o \times 16$$

Now assuming that we know the target hardware capacities (in terms of memory) $M_{\text{capa}}$ and the original memory footprint of the original model $M_{\text{ref}}$, we can derive the target compression goal $\alpha\geq\frac{M_{\text{ref}}}{M_{\text{capa}}}$. Now, we have two cases: multi-codebooks and multi-scalings. In our original submission, we define the memory cost for both as follows

$$\begin{cases} \text{multi-codebooks:} & \Omega(W) \times \log(\Omega(C)) + k \times \Omega(C) \times 16 \\ \text{multi-scalings:} & \Omega(W) \times \log(\Omega(C)) + (\Omega(s) + \Omega(C)) \times 16 \end{cases}$$

where $k$ is the number of codebooks and $\Omega(s)$ the number of scaling terms. If we combine this equation with the constraint of $\alpha$, we can derive the number of codebooks and scaling terms immediately (as in equation 3 of the original paper)

$$\begin{cases} \text{multi-codebooks:} & \Omega(W) \times 16 = \alpha \left(\Omega(W) \times \log(\Omega(C)) + k \times \Omega(C) \times 16\right) \\ \text{multi-scalings:} & \Omega(W) \times 16 = \alpha \left(\Omega(W) \times \log(\Omega(C))+ (\Omega(s) + \Omega(C)) \times 16\right) \end{cases}$$

The multi-codebook encoding is more robust to weight tensor with a lot of variation per-neuron in their distribution. On the flip side, multi-scaling encoding is more robust to weight tensors with signifant shifts in their support range but low variation in their distribution.

We aknowledge that our paper would benefit from these details and we added them to our rebuttal revision.

This paper does not explain how to cluster the neurons. There is also no data to prove that “two neurons far from each other in the weight matrix can be similarly distributed”.

In order to cluster neurons, we use the same clsutering technique $\mathcal{C}$ for the weights and the neurons. This information was missing above equation 4. We added it in the revised article. In order to re-order the output neurons of a layer, one simply needs to re-order the input neurons of the subsequent layers. We added references which detail this process which has been widely adopted in the literature $1,2,3$ . Similarly, the fact that two neurons of the same layer can be similalry distributed and be "far from each other" (i.e. in term of indexes) was illustrated in DFQ [4]. For the sake of having a self-contained paper, we added these elements of clarification in the revised manuscript.

The meaning of the X-axis in Figure 1 is confusing. For different curves, the X-axis seems to have different meanings. These should be marked on the figure.

This figure was indeed confusing, we updated it for a novel one which illustrates both our main contributions at once. Please refer to the revised article.