Compressed Context Memory for Online Language Model Interaction

We thank you for your valuable efforts and time in providing insightful feedback on our work. We would like to address questions below.

Q. Evaluation on more diverse datasets

• Thank you for your question. We would like to first mention that MetaICL consists of 26 test datasets, covering a variety of tasks such as reading comprehension, question answering, and natural language inference. As you mentioned, exploring the applicability of our method beyond ICL and dialog to more general scenarios is indeed practically significant. To investigate the potential for generalization in our approach, we train a compression module using MetaICL train tasks and SODA, a conversation dataset. Then we evaluate the model across various scenarios, including MetaICL test tasks, LaMP, and DailyDialog.
• For detailed experiment results, please refer to the general response above. The results demonstrate that our compression method possesses the ability to generalize to new datasets or scenarios. We will incorporate this analysis into our paper.

Q. Have the comparison with RMT and AutoCompressor in the main text

• Thank you for your comments. We will make sure to incorporate the comparison analysis from the Appendix into the main text.

Q. Eq 1 and Figure 4. All the Mem($\cdot$) are computed after h($\cdot$), for the same layer. Do you condition h(t) on the previous layer's Mem(t-1)?

• We apologize for any confusion caused. Equation 1 signifies that $h(t)$ (keys/values at all layer) is obtained through attention on the keys/values of $\text{Mem}(t-1)$. Specifically, the key/value at each layer is obtained after the attention process from the previous layer. Similarly, in the parallelized operation in Figure 4, the key/value at each layer of $h(t)$ is obtained through the key/value of $\text{Mem}(t-1)$ from the previous layer. We will add this clarification to the main text.

Q. On overfitting situation, could you also report the loss with context, and the test performance with and without context?

• We would like to share the results regarding the points you raised. As mentioned in our paper, training the model using the standard LoRA for compression led to overfitting, evident from the decrease in the training loss without context from an initial 2.69 to 1.84 after training. The training loss with compressed context for this model is 0.58, indicating that the model has indeed learned compression capabilities on the training set.
• However, the test loss with compressed context is 1.01 indicating the existence of generalization gap for compression. On the other hand, the test loss without context is 2.89. Considering the training loss without context of 1.84, the large gap between the train/test losses suggests the presence of overfitting to the input I(t). We will include this analysis in the paper.

Q. The masking doesn't seem to be autoregressive in figure 4b?

• Thank you for bringing this to our attention. It appears that the confusion arose from separating memory and context parts in the visualization of the attention mask. The calculation of memory $\text{Mem}(t-1)$ occurs after $c(t-1)$ and its subsequent <COMP> token. Reordering the top row of the figure 4b to align with this temporal relation will yield an autoregressive mask. We will adjust the figure to enhance clarity.

Q. Table 3 with other two datasets

• We conduct an ablation study for our conditional LoRA on a different dataset, and the results are as follows: | Dataset | Method | Default LoRA | Conditional LoRA (ours) | |---|---|---:|---:| | DailyDialog (Perplexity) | CCM-concat | 6.01 | 5.96 | | DailyDialog (Perplexity) | CCM-merge | 6.42 | 6.33 | | LaMP (Accuracy %) | CCM-concat | 83.9 | 84.7 | | LaMP (Accuracy %) | CCM-merge | 82.6 | 83.9 |

• From the above results, we observe that our conditional LoRA is effective across various datasets. We will include this results in the revision.

Q. Results on T5 without instruction finetuning

• Thank you for your advice. As you suggested, we conduct LaMP experiments on T5-Large without instruction tuning, and the results are as follows: | | No-context | Full-context | Gisting-online | Compressive | CCM-concat | CCM-merge | |---|---|---|---|---|---|---| | Accuracy (%) | 65.9 | 81.7 | 77.6 | 79.0 | 81.7 | 80.6 |

• From these results, we observe that our compression method remains effective in this model. Notably, the performance improvement by the utilization of compressed context, compared to no-context, is 16%, surpassing the 10% improvement observed in the case of Flan-T5 (Table 5). We will incorporate this result into the revised version.

Thank you once again for the valuable feedback. If you have any remained questions, please let us know.

We thank you for your valuable efforts and time in providing insightful feedback on our work. We would like to address questions below.

Q. General tests to prove its compression effect

• Thank you for the pointer. To demonstrate the generalization ability of our method in more general scenarios, we train a single compression model and evaluate its performance across various datasets. Specifically, we leverage MetaICL training tasks and the conversation dataset, SODA, as our training data, and then conduct evaluation on multiple test scenarios: MetaICL test tasks, LaMP, and DailyDialog.
• For detailed experiment results, please refer to the general response above. The results demonstrate that our compression method possesses the ability to generalize to new datasets or scenarios. We will incorporate this analysis into our paper.

Q. For few-shot learning, there doesn't seem to be any online interaction

• Thank you for your insights. We argue that in-context learning scenarios exist where not all demonstrations are provided initially. For instance, in personalization scenarios, examples of user patterns or preferences can be collected through online interactions.
• Moreover, in cases involving a large number of tasks, as discussed in recent work such as "Large-scale Lifelong Learning of In-context Instructions and How to Tackle It, ACL 2023," it is conceivable to consider scenarios where users continuously add demonstrations as data is collected.
• Alternatively, inspired by "Active Example Selection for In-Context Learning", EMNLP 2022, we can imagine a system where users set the goal for a specific task, and the language model learns through interaction by querying the user or internet, allowing for continuous improvement without the user having to prepare multiple demonstrations. We believe that MetaICL and LaMP, evaluating the in-context learning performance for new tasks/users, are effective evaluation datasets for the aforementioned scenarios.

Q. For a specific task/benchmark, there is actually a lot of compression space

• Thank you for your comments. As you pointed out, we believe that for specific tasks, the compression space is larger, making compression more manageable. This means that task-specific compression methods can achieve higher compression efficiency compared to general compression methods. In some practical industrial settings, there are attempts to fine-tune foundation models on individual datasets to obtain specialized models (https://platform.openai.com/docs/guides/fine-tuning). In a similar vein, task-specific compression could be applied to achieve the better compression capability.
• It is noteworthy that our method is task-agnostic. We can apply our approach to a wide range of scenarios, including various forms of data such as multi-tasks in MetaICL, personalization, and dialogues, without requiring any specific modifications. We believe obtaining a task-specific compression module solely based on data without task-specific knowledge or modification is practically meaningful.
• Lastly, we verify that our approach does not overfit to specific datasets or benchmarks as shown in the generality test mentioned above. Notably, the compression model trained on the SODA dataset maintains effective compression performance on the DailyDialog dataset, demonstrating its versatility. We will include these discussions in the paper for further clarity.

Q. A large amount of sampling is required to ensure the model learns well for each position

• Thank you for your feedback. Our compression method operates recurrently using the same embedding and model weights for each compression step, without utilizing different parameters conditioned on time steps or positions. Specifically, compression tokens at different time steps share the same embedding. The gradients computed for the loss during the training process are simultaneously backpropagated to all tokens across all time steps.

• Our approach effectively learns without the need for individual sampling for each position (time step). To illustrate, we conduct a comparison between training CCM-concat by sampling time steps across [1,...,16] and training with a fixed time step of T=16 in MetaICL with LLaMA-7B. The table below shows the test accuracy of two models across a range of time steps. | Time step | 1 | 4 | 16 | |-|-:|-:|-:| | Time sampling | 63.0 | 67.8 | 70.0 | | Only T=16 | 62.5 | 67.4 | 70.1 |

• The results indicate that there is not a significant difference between the two methods. This suggests that substantial sampling for compression at each time step or position is not necessary for our approach. We will include this analysis in the revision.

Q. Regarding performance of RMT in Table 15

• Thank you for your question. The experiments in Table 15 were conducted following the procedure outlined below. Firstly, we obtained the OPT full attention-2.7b-4k model, publicly available from the AutoCompressor GitHub. We fine-tuned this model on the MetaICL training set for ICL meta-training, following the settings in Table 9. Subsequently, we performed compression fine-tuning for RMT, AutoCompressor, and CCM methods, with the same number of steps. The training of RMT and AutoCompressor was carried out using the official AutoCompressor GitHub code. This ensured that all experiments were conducted using the same pretrained model and the same training data. We note that the publicly available RMT and AutoCompressor compression models on GitHub show different ICL performances without compression.
• We acknowledge the approach mentioned above does not assess the original compression capabilities of RMT and AutoCompressor, both trained on a vast corpus comprising billions of tokens. We speculate that the lower performance measured for RMT in Table 15 is due to insufficient training. To address this issue, we obtain the RMT and AutoCompressor compression models from the AutoCompressor GitHub and conduct evaluations on these baseline models. Since these two models have different ICL performances, we conduct separate experiments. For a fair comparison, we also provide finetuned results of the baseline models using the identical training steps to ours. It’s noteworthy that our method only optimizes conditional LoRA for compression and does not affect the ICL performance of the baseline models.

• The first table shows the results by using the RMT model from the GitHub and the table below shows the results by using the AutoCompressor as the full-context model. The results above show a decrease of approximately 5% in both No-context and Full-context performance compared to Table 15, due to the lack of ICL meta-training of baseline models. However, even in this setting, our compression method demonstrates superior efficiency. We acknowledge that the explanation for Table 15 in our paper was insufficient and, to ensure more accurate comparison, we will incorporate the above description and results into the revised version.

Q. Changes in the order of sample arrangement

• Thank you for the pointer. To measure sensitivity on the order of demonstrations, we fix the demonstration set of each test task and evaluate performance by changing only the demonstration order using random seeds. The table below shows the test accuracy on MetaICL with OPT-2.7B over various random seed for the demonstration order. |Model \ Seed|0|1|2|3|4|Stdev.| |-|-|-|-|-|-|-:| | Full-context|59.1|59.8|59.6|59.3|60.5|0.54| | CCM-concat | 53.1 | 52.0 | 52.9 | 52.0 | 52.7 | 0.55 | | CCM-merge | 53.2 | 53.2 | 53.5 | 53.1 | 53.3 | 0.15 |
• In the case of LLaMA-7B, we obtain Standard deviations under 0.25 for all three methods. As evident from the results above, the OPT model shows a larger performance variance based on the order compared to LLaMA-7B. However, it is noteworthy that the variance based on order is not significantly larger compared to the performance difference between methods. Specifically, CCM-Concat outperforms other compression baselines by over 2% in test accuracy (Table 15). We will incorporate these findings into the revision.

Q. In Figure 4 (a), what do the blocks in different colors mean?

• In the figure, each block represents a key/value pair. Blocks in yellow correspond to keys/values of memory, those in red represent keys/values of compression tokens, and those in green correspond to key/value pairs associated with regular text. We will add this explanation in the revision.

Q. More details about the fine-tuning process

• We provide the fine-tuning process in Appendix C. During the fine-tuning process, only the LoRA adapter is additionally trained, following the training settings outlined in Table 9, which is adopted from the Gisting paper. A batch size of 128 are processed on a single GPU through gradient accumulation, with training ranging from 300 to 2000 steps depending on the dataset size. The training time takes around 3 to 24 hours. For a fair comparison with compression baselines considered in the paper, we conducted the comparison after fine-tuning all methods with the same process.

Thank you once again for the valuable feedback. If you have any remained questions, please let us know.

We thank you for your valuable efforts and time in providing insightful feedback on our work. We would like to address questions below.

Q. Missing references

• Thank you for your mention. The papers you mentioned, "Memory Networks" and "Fast Weights to Attend Recent Past," are indeed important seminal works, and we will make sure to add a comparison to related work in our paper. It's worth noting that our approach differs in that we propose a compression adapter to the existing foundation model, whereas they propose entirely new models. As you pointed out, we will also include an explanation of the Compressive Transformer in the preliminary section to enhance the reader's understanding.

Q. Table with number of tokens

• Thank you for the suggestion. As you mentioned, we will add the corresponding row. The results for MetaICL are as follows (Table 4). | Method | No-context | Full-context | Gisting | Gisting-online | Compressive | CCM-concat | CCM-merge | |-|-:|-:|-:|-:|-:|-:|-:| | Max # KV token during inference | 50 | 630 | 588 | 178 | 178 | 178 | 66 |

Q. More sophisticated memory update mechanisms

• It’s a great idea! Based on your advice, we have formulated a more general form of the memory update equation as follows:
  $\text{Mem}(t) = (1 - a_{t-1}) \text{Mem}(t-1) + a_{t-1} h(t), \text{for}\ a_{t-1} \in [0,1].$

• In our paper, we proposed an update method based on the arithmetic average (AA) of the compressed state up to the present time, i.e., $a_{t-1}=1/t$. Another natural design choice can be an exponential moving average (EMA), where $a_{t-1}$ is set to a constant value. This strategy weigh higher importance on recent information, when compared to the arithmetic average. In below, we train a model with EMA and compare it's test performance to the original CCM-merge (AA). | Time step | 1 | 2 | 4 | 8 | 12 | |-|-:|-:|-:|-:|-:| | EMA | 7.49 | 7.06 | 6.79 | 6.49 | 6.38 | | AA | 7.47 | 7.06 | 6.87 | 6.54 | 6.34 |

• The table above provides perplexity results of the arithmetic average and EMA with $a_{t}=0.5$, in DailyDialog with LLaMA-7B. S The results indicate that both methods yield similar performance. When forming the compression state $h(t)$, our method involves referencing the previous memory, $\text{Mem}(t-1)$. We believe this enable preservation of overall context, even with exponentially decreasing coefficient for past states by EMA. We also note that CCM-concat can be interpreted as dynamically inferring coefficients for individual hidden states $h(t)$ through the transformer attention mechanism.

Q. Comparison to retrieval-based approach

• Thank you for your feedback. As you pointed out, memory retrieval-based approaches are indeed important baselines for our work. Existing studies like Memorizing Transformer, and recent studies such as MemoryBank, LongMem, and UnlimitedFormer primarily focus on the retrieval process and training of retrieval-based models, with less emphasis on memory compression. These retrieval studies adopt strategies to cache key-value pairs for individual tokens and assume the existence of extensive storage. However, as seen in Table 4 and 5 of our paper, LLM's keys and values demand a significant amount of storage, even for context size of 1024 tokens, reaching several hundred megabytes. In scenarios like user-level personalization and conversation systems, such high storage requirements can become problematic. We will add these references in the related work section.

• We also note that MemoryBank proposes reducing context history size through summarization. However, this approach comes with additional computational costs for summarization and the overhead of processing the summarized text for subsequent inference. In contrast, our approach allows for more efficient inference without the aforementioned overhead, thanks to the caching of key-value pairs of compression tokens. Following the MemoryBank paper, we conduct experimental comparisons with LLaMA-7B on DailyDialog. Specifically, we use the summarization prompt from the paper to compress context through OpenAI gpt-3.5-turbo API (ChatGPT) and then evaluate models with summarized contexts. The table below compares the test perplexity of methods. | | No-context | Full-context | MemoryBank (summarization) | CCM-concat | CCM-merge | |-|-:|-:|-:|-:|-:| | Perplexity | 10.6 | 5.59 | 7.06 | 5.98 | 6.34 | | Compressed context length | 0 | 222 | 60 | 24 | 2 |

• From the above results, we can confirm that our approach achieves superior performance with less compression overhead and shorter compressed context length. In contrast to text summarization methods that compress context into text form, our approach achieves better compression efficiency by compressing information into key/value parameter space. We will incorporate these findings into the paper.

Thank you once again for the valuable feedback. If you have any remained questions, please let us know.

Q. Explanation for effectiveness

• As explained at the bottom of page 7, in conversation tasks like DailyDialog, where each context history contains a wide range of contextual information, the merging approach exhibits decreased performance when compared to the concatenation approach. Conversely, for context histories consisting of mutually complementary information, such as in MetaICL and LaMP, the merging method demonstrates its effectiveness.

• We further discuss the effectiveness of our approach in specific scenarios. As shown in Table 8, our method excels in the accuracy metric. When evaluating the generation performance with RougeL metric in MetaICL, we obtain the following results: | Metric | No-context | Full-context | Gisting | Gisting-online | Compressive | CCM-merge | CCM-concat | |-|-:|-:|-:|-:|-:|-:|-:| | RougeL | 12.3 | 61.4 | 47.7 | 37.9 | 47.9 | 48.3 | 54.7 | | Accuracy (%) | 51.7 | 70.8 | 66.9 | 57.7 | 67.8 | 69.6 | 70.0 |

• From the table, our method delivers the most accurate generation performance compared to other baselines. However, there is a larger decrease in performance compared to the full context, whereas in the case of accuracy, the performance drop is less than 1%. When we examine the generated outputs with compressed context, we observe instances where synonyms (e.g., "Different" and "Dissimilar" in the medical_questions_pair task) or variations in letter casing (e.g., "Hate" and "hate" in the tweet_eval_hate task) exist, indicating semantic equivalence but differences in expression. Based on these results, we believe our method is particularly effective in scenarios that prioritize preference or nuances, rather than returning precise expressions. We will incorporate these discussions into the revised version.

Q. More sophisticated memory update mechanisms

• Based on your advice, we have formulated a more general form of the memory update equation as follows:
  $\text{Mem}(t) = (1 - a_{t-1}) \text{Mem}(t-1) + a_{t-1} h(t), \text{for}\ a_{t-1} \in [0,1].$
  In our paper, we proposed an update method based on the arithmetic average (AA) of the compressed state up to the present time, i.e., $a_{t-1}=1/t$. Another natural design choice can be an exponential moving average (EMA), where $a_{t-1}$ is set to a constant value. This strategy weigh higher importance on recent information, when compared to the arithmetic average. In below, we train a model with EMA and compare it's test performance to the original CCM-merge (AA). | Time step | 1 | 2 | 4 | 8 | 12 | |-|-:|-:|-:|-:|-:| | EMA | 7.49 | 7.06 | 6.79 | 6.49 | 6.38 | | AA | 7.47 | 7.06 | 6.87 | 6.54 | 6.34 |

• The table above provides perplexity results of the arithmetic average and EMA with $a_{t}=0.5$, in DailyDialog with LLaMA-7B. S The results indicate that both methods yield similar performance. When forming the compression state $h(t)$, our method involves referencing the previous memory, $\text{Mem}(t-1)$. We believe this enable preservation of overall context, even with exponentially decreasing coefficient for past states by EMA. We also note that CCM-concat can be interpreted as dynamically inferring coefficients for individual hidden states $h(t)$ through the transformer attention mechanism.

Q. Lack of comparison to recurrent memory approaches.

• Thank you for your feedback pointing out the significance of recurrent memory approaches as a baseline. In Appendix D, Table 15, we conduct a comparison with recent recurrent memory approaches, RMT and AutoCompressor. Direct comparison with LinearTransformer is non-trivial due to its structural modifications in the Transformer operations, making its application to foundation models like LLaMA challenging. We acknowledge the importance of this previous work and will add discussion in the related work section.
• From comparison to RMT and AutoCompressor, we observed that our parallelized training technique significantly improves training speed by approximately 7x (Appendix, Table 15), while also achieving superior performance at a similar compression level. We recognize the importance of this comparison, and we will incorporate it into the main body of our paper.

Thank you once again for the valuable feedback. If you have any remained questions, please let us know.

We thank you for your valuable efforts and time in providing insightful feedback on our work. We would like to address questions below.

Q: Task-specific limitation in real-world applications with diverse tasks.

• Thank you for your feedback. The MetaICL dataset encompasses evaluations for a wide range of tasks unseen during training, including reading comprehension, question answering, and natural language inference. In our paper, we demonstrate the strong task-generalization performance achieved by our compression method on this multi-task dataset.
• However, as you pointed out, we optimize the model for specific application scenarios in experiments. To demonstrate its utility in more general use cases, we trained a single compression model and evaluated its performance across various scenarios. Specifically, we leverage MetaICL training tasks and the conversation dataset, SODA, as our training data, and then conduct evaluations on MetaICL test tasks, LaMP, and DailyDialog.
  For detailed experiment results, please refer to the general response above. The results presented in the table indicate that our approach are able to generalize across diverse scenarios without the need for additional fine-tuning. We will add this analysis in the revision.

Q. Hybrid approach combining task-specific and task-agnostic components.

• Your idea is indeed insightful. Task-specific compression offers the advantage of having a larger compression space since it focuses on compressing information relevant to a specific task, in contrast to task-agnostic compression, which aims to compress information in a general context. While task-agnostic methods are valuable for compression in generic scenarios, achieving superior compression performance in specific applications may require task-specific adaptations.
• From a technical view, there is potential for significant benefits by integrating both approaches. For instance, one could first apply task-agnostic compression and then follow task-specific compression, combining the strengths of both methods. Alternatively, using a task-agnostic compression module as an initialization and then fine-tuning it in a task-specific context may reduce fine-tuning costs. We will incorporating these discussions into our paper.

Q. Additional resources for training

• Thank you for your question. In our paper, we compared our method with Gisting, Compressive Transformer, and in the appendix, RMT and AutoCompressor. During experiments, all these methods underwent the same process of fine-tuning on each dataset, identical to our method for a fair comparison. We adopted the training protocol from the Gisting paper and adjusted the training steps according to the size of each dataset.

• Our approach is memory-efficient as it only requires training the LoRA adapter for compression. As demonstrated in Appendix Table 15, our parallelized training technique provides a 7x faster training speed compared to existing recurrent memory approaches, RMT and AutoCompressor. Additionally, as indicated in Table 9, we used datasets containing around 100k samples for each application. This dataset size is similar to the recently released instruction finetuning dataset, Alpaca, and is roughly 10% of the size of LmSys-Chat-1M. In contrast, AutoCompressor relies on a massive training dataset containing up to 15 billion tokens for compression, which would demand significant resources and training costs when developing individual models in industry. In this point of view, our approach has advantages in resource-constrained environments by its relatively lower memory and data requirements.

Q. Analysis on performance gap

• Firstly, the existence of a performance gap can be attributed to our compression experiment setup, which demands generalization abilities for new tasks (MetaICL) / users (LaMP) / context histories (DailyDialog). In fact, while the MetaICL train loss for full context LLaMA-7B is 0.59 and CCM-concat's train loss is 0.61, the test set loss diverges to 0.76 (full context) and 0.92 (CCM-concat), indicating the presence of a generalization gap. On the other hand, the Compressive Transformer exhibits a similar level of train loss, but it shows a test loss of 1.21, indicating a larger generalization gap.