Paper: In-context Autoencoder for Context Compression in a Large Language Model (ICLR-2024)

Reviewer comment (Experiments):
The paper compresses the contexts into black-box memories. Instead, one can simply apply a summarization model (e.g., GPT-4) to summarize the context and feed it to the model. While Figure 2 claims that black-box memory can be more compressive than words, this claim is not supported by the experiments. Besides, summarization has a merit in that it can be applied to black-box API models, unlike this method, which is only applicable to open-source white-box models.

Author response:
We have experimental results to show that black-box memories can be more compressive than a summary. Firstly, both Figure 4 and 5 demonstrate that using 128-length black-box memories can almost losslessly restore the original 512-length context (with BLEU>99%), which cannot be achieved using a 128-length summary. Secondly, in the last row of Table 6, we compared 128 black-box memory and 128-token summary, which also confirmed that black-box memory is more compact and informative than a summary with the equal length.

Reviewer comment (Novelty):
Though concurrent, discussing other techniques for long context LLM in the related work section would help readers grasp the positioning of this paper. I can see many works on OpenReview and arXiv, such as [4] LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models and [5] Functional Interpolation for Relative Positions improves Long Context Transformers, to name a few.

Author response:
Thank you for your reminder and recommendations. We will include and extend the discussion of these contemporary papers in the revised version.

Reviewer comment (Presentation):
Using a different color for linkcolor would help readers better recognize the link to the tables.

Author response:
We appreciate your suggestions on the presentation of this paper. We will fix these issues in the next version according to your advice.

Reviewer comment (Presentation):
Putting tables/figures at the top instead of within the text would help readers to better skim the experiments.

Author response:
We appreciate your suggestions on the presentation of this paper. We will fix these issues in the next version according to your advice.

Reviewer comment (Experiments):
A comparison with the summarization baseline, which simply compresses the context into words instead of black-box memories, is needed.

Author response:
As we respond to Weakness 1, we have the experiment results in Figure 4/5 and Table 6 to demonstrate that the memory slots are more compact and informative than the summary with the equal length.

Reviewer comment (Reproducibility):
I likely have missed something in the paper. It's unclear to me whether the authors initialized the encoder LLM from the same checkpoint as the decoder LLM? If that's the case, how did the author address the difference in autoregressive attention in the decoder (and the base LLM, such as LLAMA) and bidirectional attention potentially used for the encoder.

Author response:
The encoder in our paper refers to the encoder in the autoencoder terminology, which is used for encoding text into memory slots. It is not the same as the encoder in the encoder-decoder architecture.

As illustrated in Figure 3, the encoder of ICAE is a frozen LLM (specifically Llama in this paper) + a learnable LoRA + a learnable embedding for memory tokens. The decoder of ICAE is the frozen LLM.

Therefore, the encoder does not use bidirectional attention; instead, it uses causal attention, as in the decoder.

We apologize for the confusion. We'll make it clearer in the revised version.

Reviewer comment (Experiments):
Why not use an encoder-decoder language model such as T5 for this?

Author response:
As we mentioned in our response to Q1, our encoder and decoder refer to the encoding and decoding modules of the autoencoder (i.e., ICAE), which are totally different from the encoder/decoder in the encoder-decoder architecture. Our goal is to compress contexts for an existing Large Language Model (LLM), specifically Llama in our experiments. Building the ICAE on the same LLM (Llama) allows us to better align with the LLM to generate a more compressed context.

Reviewer comment (Theory):
As mentioned above, does it matter for the encoder's input embeddings to be in the same space as the decoder's input embeddings?

Author response:
As we previously mentioned in our response to Q1, the encoder's input embeddings and the decoder's input embeddings are indeed in the same space.

Reviewer comment (Reproducibility):
What's the actual size of each of the pretraining tasks, in terms of number of examples or number of tokens consumed.

Author response:
We have included these details in Appendix A. For pretraining, our batch size and training steps are listed in Table 9. The number of training tokens can be calculated as 256 * 512 * 200k = 26.2 billion tokens.

Reviewer comment (Experiments):
I would imagine the autoencoding task to be relatively easy to "learn". How many steps or number of examples does it take for the model to converge?

Author response:
You’re exactly right. Autoencoding is easier to learn than Language Modeling. However, our setting is 4x compression (e.g., using 128 tokens to restore 512 tokens). According to our observations, the validation loss still keeps decreasing even after 150k training steps.

Reviewer comment (Experiments):
For the instruction fine-tuning, have you tried passing the prompt to the encoder instead of the decoder, and what's the effect?

Author response:
Yes, we tried. As long as we mixed such data in our instruction fine-tuning, the ICAE could still perform well because compressing task description prompt is not hard.

Reviewer comment (Reproducibility):
Do you initialize the encoder from the same LLM checkpoint as the decoder?

Author response:
Yes, they are the same — they are both Llama. The only difference is that the ICAE encoder adds a learnable LoRA module and a learnable embedding for memory tokens (see Figure 3).

Reviewer comment (Evaluation):
Could you also provide some explanation on why the first row of table 1 shows a PPLX increase even though there's no compression going on.

Author response:
This is because we replaced the original natural language tokens with compressed memory slots, which inevitably causes some loss (even without compression). This is understandable as our memory slots are pretrained on 26.2 billion tokens, while natural language tokens are pretrained on trillions of tokens by Llama.

Reviewer comment (Experiments):
Regarding the claim that pretrained ICAE suffers less from hallucination than non-pretrained model, I would like to clarify that both variants are trained on the same instruction fine-tuning set right?

Author response:
Yes, the only difference is that one is pretrained with AE+LM before instruction fine-tuning, while the other is not pretrained and is directly trained on the instruction fine-tuning data.

Reviewer comment (Experiments):
Have you tried full model fine-tuning? In other words, what's the impact of lora on the final quality?

Author response:
We have tried it on a small scale, but due to computational resource limitations, we have not done it on a large scale. However, under the same training steps (10k), we found that the results of full model fine-tuning are better than LoRA. But the advantage of LoRA is that it only adds 1% extra memory when used, which is important for practical applications.

Reviewer comment (Experiments):
In table 5, in rows 2, 3, 4 and 5 is column#2 (i.e. "system 2") also fine-tuned on the pwc dataset? I presume so, but I do want to check. Without this, it would be a very slightly unfair comparison if "system 1" got to see the pwc-train set but "system 2" did not.

Author response:
All the llama-2-chat models in the column of system 2 are fine-tuned with the PwC dataset using LoRA.

Reviewer comment (Experiments):
One drawback is that ICAE should be anyway pre-trained, despite few additional parameters and fewer training steps than LLM pre-trainings.

Author response:
When we use the same context (i.e., cached in advance), we can significantly save on latency/memory overhead (over 7x speedup, as we discussed in Section 3.3.2). This scenario is **common and practical**: For example, a customized ChatGPT for a specific user can compress the user's chat history in advance and use the compressed memory during inference to maintain the long-term memory without online latency overhead for compression. Regarding your concern about the need to recompress for different contexts, different contexts indeed require an additional compression step. However, **as we show in Table 8, this additional compression still brings significant end-to-end speedup compared to using the original context when generating a long sequence**, making it still a highly practical method.

Reviewer comment (Evaluation):
Compression will be extremely beneficial if we use the same context multiple times. However, otherwise, it might only incur overhead.

Author response:
When we use the same context (i.e., cached in advance), we can significantly save on latency/memory overhead (over 7x speedup, as we discussed in Section 3.3.2). This scenario is **common and practical**: For example, a customized ChatGPT for a specific user can compress the user's chat history in advance and use the compressed memory during inference to maintain the long-term memory without online latency overhead for compression. Regarding your concern about the need to recompress for different contexts, different contexts indeed require an additional compression step. However, **as we show in Table 8, this additional compression still brings significant end-to-end speedup compared to using the original context when generating a long sequence**, making it still a highly practical method.

Reviewer comment (Reproducibility):
What is the format of memory slots, and how are they generated from the model? I could not find where it is described, but they should be vectors on the embedding space to be fed into the decoder as input. Are they representations before the final softmax layer to get the next word probability distribution? It might be meaningless, but what will we see if we map each memory slot vector to the most likely word?

Author response:
Memory slots are indeed vectors in the embedding space that are fed into the decoder as input. They are the final output of the ICAE encoder (before softmax) in the positions of memory tokens. When we map the memory slot vector to the vocabulary space, we find that memory slots tend to retain some keyword information from the original context and are relatively concentrated in the front part of the original context. This can be explained by the memorization intuition we discussed in our paper, where the model will only remember keywords that provide crucial hints.

Reviewer comment (Theory):
Is it correct that memory slots are order-dependent because they use the same causal attention of language models? If yes, is it optimal because we have access to all information when we create memory slots?

Author response:
Yes, you’re exactly correct. They are generated in parallel. They are order-dependent and use causal attention. We think it is an appropriate way to create memory slots just as people memorize in their working memory.

Reviewer comment (Experiments):
Based on the fact that the amount of information varies depending on the content of the context, not solely depending on its length, is it possible to have an adaptive size of memory slots for more compact compression?

Author response:
As shown in Table 3, texts with lower PPL have greater compression potential. This is because lower PPL mean higher certainty and less difficulty to memorize. Your suggestion is very correct. We are working on adaptive memory slots to dynamically adjust the compression rate according to the difficulty of the text (PPL), in order to achieve the best compression performance without affecting the results.

Reviewer comment (Experiments):
I agree that AE and LM objectives are straightforward for pre-training ICAE. I wonder if you have tried different objectives. Assuming we can pre-train ICAE from scratch instead of further pre-training from a pre-trained LLM, what will be the best way? Still the same objectives?

Author response:
We have tried pretraining an ICAE with the same objective using a GPT-style model from scratch with a similar number of parameters as we used in this paper. However, its results are weaker than the ICAE based on a pretrained LLM. Although we have not found a more efficient and scalable objective than AE/LM, we do not deny that there may be a better training objective to pretrain ICAE, as you suggested.

Reviewer comment (Experiments):
In Section 3.3.1, you mentioned the scalability in terms of context length. Are models trained separately for the sequence lengths? How does ICAE generalize well to different (L, k) combinations that are different from the one used for training?

Author response:
How well ICAE generalizes to different (L, k) combinations depends on how we train the ICAE. If we only use L→k to train the ICAE, it will not generalize well to different length combinations. However, if we mix different lengths during the pretraining phase, the model can generalize well to different lengths.

Reviewer comment (Experiments):
You mentioned that the maximal token length is 512 in the Data paragraph in Section 3.1. However, most samples are longer than 512 tokens, according to Figure 10. Did you truncate them by 512?

Author response:
Yes, in the 512 fixed-length experiment group, we truncate to 512 and then compare the results fairly. As most samples are longer than 512 tokens, it is reliable to say ICAE is performing a 4x compression for the majority of test cases.

Reviewer comment (Reproducibility):
What value of N did you use for the text continuation training?

Author response:
Do you mean Figure 6(Left)? In our experiments for text continuation training, N ranges from 1 to 4.

Reviewer comment (Reproducibility):
Training requires 8 A100 GPUs with 80GB. Can you describe the approximate training cost (time and money)?

Author response:
Training does not require 8 A100 GPUs with 80GB, it can be done with just 1 A100 (because we use LoRA, there is no need to control the GPU memory with parallelism techniques). We used 8 A100 GPUs just for speeding up the pretraining process, which takes about 4 days on 26.2 billion tokens.

Reviewer comment (Presentation):
(minor) the first paragraph in Section 3.3.3: Table 6(Right) -> Figure 6 (Right)

Author response:
Thank you very much for pointing out the error in Section 3.3.3. We will correct it in the next version.