Scalable Language Model with Generalized Continual Learning

Thank you for your review and valuable feedback. We appreciate your recognition of the motivation behind our work, and its simplicity and significance. Based on your feedback, we will make the necessary updates to address the mentioned typos and provide additional clarification.

1: "The number of data points in a task should also be considered..."

We sincerely appreciate your suggestion, which raises an important point regarding the consideration of the scale and size of downstream tasks. Initially, we followed the previous continual learning works, which shared the same task scale and data points. However, our approach provides flexibility and adaptability by allowing the allocation of varying numbers of data points to meet the specific requirements of each task, thereby overcoming this limitation.

2."The computational complexity of the proposed model..."

Thanks for raising this important point. We humbly acknowledge that the proposed method indeed introduces a cost associated with the retrieval process and we are dedicated to discussing and re-claiming it in the discussion section. However, we find the additional cost to be acceptable, and the benefits clearly outweight the additional minimal cost, because:

• Compared to the subsequent inference models, the retrieval stage model used is notably smaller, lighter, and operates at a faster speed. This distinction is particularly significant for the T5 and Llama models.

• In the case of generation models with the decoder architecture, each inference only produces a single token, necessitating multiple inferences to generate a complete sentence. However, the retrieval process is executed only once for yielding each complete answer. Therefor, given $t_r$ as the retrival time, $t_i$ as the inference time, $n$ as the tokens number, the proportion of time consumed is: $\frac{t_r}{t_r+n*t_i}*100\%$

We conducted an experimental comparison to measure the time consumption of different parts of various tasks using Llama on a single A100 GPU.

3."the initialization of (key, value) pairs"

We are sorry for any confusion caused. We will provide more details about the initial process here. Please refer to [1] or the official PyTorch manual for further information. We adopt orthogonal initialization for the keys initialization, specifically,

1. Generate a random matrix $A\in R^{n\times c} \sim \mathcal{N}(0,1)$
2. Perform QR factorization on: $A=QR$
3. Adjust the sign of $Q$ based on the sign of the diagonal elements of $R$
4. Copy the adjusted $Q$ into the original tensor

Regarding the value, we utilize zero and norm initialization.

[1] Exact solutions to the nonlinear dynamics of learning in deep linear neural networks.

4. Typographical errors and improvement

Thanks for your feedback. We will address the typo and consider your suggestions for improving understanding in the final version.

5."As the pretrained model is frozen and is always i...", "Related to the above question..."

This question is considered at the beginning of our method design. While it is inevitable to encounter unexpected results, this issue may not significantly impair the overall performance, attributed to the following facts: 1. The accuracy of the retrieval process is exceptionally high. 2. We have implemented JARe, which takes multiple results into considerations to obtain the agreed one, thereby enhancing the fault tolerance, especially in some cases where the rare outliers may occur. 3. Tasks with similar distributions may share common knowledge and characteristics. As a result, incorporating such tasks does not typically have a detrimental impact. And the similarity analysis can be found in the A.8 of the Appendix. In addition, we provide quantitative statistics about the accuracy of the retrieved related points as follows.

Thank you for your insightful comments and valuable suggestions. We have incorporated more discussions with the references to provide in-depth analysis. Moreover, we will include a main paper paragraph to discuss the weaknesses of our approach. Please find our responses listed below.

1. Comparison and discussion with the replay-based methods.

Thank you for your suggestion, the discussion regarding the replay-based method is indeed crucial. The original version has included replay-method for T5 and Bert benchmarks. We then supplement the Llama experiments of the replay method follow previous work with results, which are listed below in a different order. It is worth noting that the results are better than the naive fine-tuning approach. Additionally, our method has also achieved a significant improvement, with an average accuracy of 82.3%.

Moreover, we would like to share the following interesting findings regarding the replay-based methods.

1. The replay-based methods help maintain the LLMs ability of the format and structure of question answering. After the naive fine-tuning, the LLM can only answer in the format and knowledge related to the last task. Replay-based methods enhance the LLM's ability to maintain diverse answer formats, resulting in improvements compared to the baseline.

2. However, data replay can only preserve basic abilities like Q&A and classification. When faced with more complex tasks that demand advanced logical reasoning and inference abilities, such as MMLU, data replay can certainly give answers but may provide random and incorrect answers. This is a major drawback that contributes to the observed performance degradation of data replay on the complex task.

3. During replay, the quality of the Q&A has also deteriorated, with shorter answers that provide less useful information. The replayed methods tend to provide more generic responses and platitudes, even though the differences may not be obvious through the BertScore metric.

2.Require an additional model and addtitional parameters compared with the baseline

We humbly acknowledge that the proposed method indeed introduces a cost associated with the retrieval process and we are dedicated to discussing and re-claiming such weaknesses in the discussion section as you have suggested. However, we find the additional cost to be acceptable because:

• Compared to the subsequent inference models, the retrieval stage model used is notably smaller, lighter, and operates at a faster speed. This distinction is particularly significant for the T5 and Llama models.

• In the case of generation models with the decoder architecture, each inference only produces a single token, necessitating multiple inferences to generate a complete sentence. However, the retrieval process is executed only once. Therefor, given $t_r$ as the retrivak time, $t_i$ as the inference time, $n$ as the tokens number, the proportion of time consumed is: $\frac{t_r}{t_r+n*t_i}*100\%$

We conducted an experimental comparison to measure the time consumption of different parts of various tasks using Llama on a single A100 GPU.

In terms of storage, the total memory associated with storing different task datasets are list below. The "ZIP" file size represents the size of the original compressed ZIP file, while "with cache" refers to the uncompressed files with additional generated cache, such as '.arrow' files used in the HuggingFace framework. The memory cost for Language Model Models (LLMs) is also significant.

Here we list the memory we used to store the additional parameters:

Moreover, while our method requires additional parameters, these parameters are only used to store the weight increments. They do not incur any computational cost or increase the complexity of the original models.

3. The upper bound performance.

Thank you for your suggestions, and we are sorry for any confusion caused. Our intention was to compare the per-task fine-tuning with our methods. As you mentioned, these results cannot be considered as the upper bound and the models can also benefit from transfer learning, as demonstrated in Table 6. We will refine such statement.

Thank you for taking the time to review our paper and for your thoughtful feedback. We greatly appreciate your careful consideration of our rebuttal and are delighted to hear that you are satisfied with our response. We sincerely appreciate your time and effort in reviewing our work. Your valuable insights have been invaluable in improving the quality of our paper.

Thank you for your valuable feedback. We are encouraged and deeply appreciate your affirmation of the performance and contribution of our work. Our responses to the mentioned weaknesses are listed below.

1."It may be harder to reproduce the method compared to prior ones, as it is quite different from existing methods."

While our approach may differ from prior methods, it is important to note that the paradigm and pipeline we propose are also straightforward and intuitive in nature. Moreover, to address the concern of reproducing our method, we will make our source code and the trained models publicly available for researchers to access and facilitate the understanding regardin the implementation details.

Thank you for your valuable feedback on our paper! We appreciate your recognition of our work's performance and simplicity. We will incorporate your suggestions to refine the method's description and improve its clarity.

Q1: "Is the preparation phase conducted collectively...?", "are the keys and low-rank weights updated utilizing..."

A1: Following previous work, we ensure that each task is treated as an independent entity and executed sequentially. We train the corresponding keys and low-rank weights using only task-specific examples and then keep them frozen without any further fine-tuning. For your convenience, we have included a more detailed definition of the process in Algorithm 1.

Q2: What is the initial count of keys, and is it dependent on the number of future tasks?

A2: As previously mentioned, all tasks are treated as independent and decoupled, without any dependencies on the number of future tasks. The generation of task-specific keys is tailored to the fitting requirements of each task, and it is not influenced by future tasks.

Q3: How is the aggregation across groups achieved

A3: The aggregation across groups is accomplished through a joint weighted sum. In the retrieval process, the query is partitioned into multiple groups. Each group applies a single sub-query to multiple keys, with the keys sharing the same channels as their corresponding sub-query. This operation is similar to the multi-head attention mechanism, and more details can be found in Eq.(6) in the main paper.

1. Basic baseline

We sincerely appreciate your valuable suggestion. In response, we have incorporated the baseline results by utilizing a single key and a single weight, which are showcased in the table presented below. It is worth noting that the proposed method demonstrates a subsequent improvement in the BERT benchmark achieving a score of 79.1%. This is because relying on a single key value can result in incorrect retrieval without fault tolerance.

Furthermore, we have provided a comprehensive analysis of the impact of hyperparameters, which can be found in A.8 and A.9 of the Appendix.

2. Question about Table 6

Thanks for your question! We would like to clarify that:

"The retention of performance on trained tasks seems less beneficial... "

Recently, there has been a surge in fine-tuning pretrained large language models (LLMs) on domain-specific datasets to create customized assistants. However, it has been observed that fine-tuning on small-scale datasets that differ significantly from the training data can have a negative impact on the LLM's generality and adaptability, shown Table 6. Despite the performance of Llama, it exhibits significant degradation in small-scale fine-tuning tasks. And our method remidies this issue, and aims to mitigate the catastrophic forgetting and even brings about improvements in transferability.

Furthermore, in the fundamental setting of continual learning, where the source of inputs is unknown and the objective is to progressively improve the model's capabilities, it is impractical to use multiple trained models and entails resource redundancy.

"zero-shot performance with every newly learned task"

Our experiments are in line with [1], where the training and zero-shot evaluation datasets cover various task types without overlap. The results in the original main paper only include the single result of the last task, following the recent literature. We have provided the evaluation results after each individual task fine-tuning in the table below.

Finance:

MMLU:

Meical:

As discussed in [1], the language models (LMs) acquire transfer learning capability by learning the format and structure of question answering. However, achieving an improvement in general transfer ability necessitates large-scale multi-task training conducted collectively. This deviates from the principles of continual learning and entails significant resource requirements. These two tasks have different focuses. Our primary focus is on continual learning, ensuring fast adaptation and retaining the LLM's general abilities. Additionally, our approach shows potential for delivering fast adaptation ability to various LMs, even with small-scale training data.

[1] Multitask Prompted Training Enables Zero-Shot Task Generalization

Thank you for taking the time to share your valuable feedbacks and comments. We sincerely appreciate the opportunity to provide our explanations and clarifications.

1: Explain how the number of keys needed

The selection of the number of keys is flexible and can be regarded as a hyper-parameter within the framework. However, given the fact that the consumption of each key-value pair is minimal, particularly when compared to the original model, it is acceptable to allocate an adequate number of keys at a low cost. And we found that four pairs are sufficient to ensure the desired outcome, and we found it can well generalize to different datasets and tasks.

2. The number of keys would influence the final rank

In our experiments, we discovered that the performance improvement mostly arises from the fact that multiple retrieval increase the fault tolerance mechanism and sharing knowledge. Continuously increasing parameters directly on a certain basis, such as low-rank, does not yield a substantial impact on the overall performance. More derails are shown in A.8 and A.9 in the Appendix.

3. Is the same number of parameters used in both methods?

Thanks for your kind reminder. In this experiment, we keep the number of parameters same by increasing the low-rank number. As mentioned above, the performance degradation primarily stems from incorrect retrieval and the lack of shared knowledge. It is possible that the scale of learnable parameters is already sufficient, which is also supported by recent parameter-efficient fine-tuning methods.

4. Is it frozen or does it continue to train

The keys are frozen after the preparation stage, indicating that they are not modified thereafter. The preparation stage is responsible for allocating and training the keys, while the fine-tuning stage focuses on adjusting the weights. Once the preparation stage is completed, the keys are frozen. More details are shown in Algorithm 1.

5. After training, are the keys made orthogonal again?

At the beginning of the preparation stage, all the allocated keys are initialized to be orthogonal to each other. Since only the top-K keys are selected and updated in each iteration, this strategy aims to make a concerted effort to retrieve different keys, enhancing the overall diversity. However, with the proposed random-mask strategy, we discovered that the initialization method has limited influence on the overall outcome. After training, all keys are adapted to the task distribution, and the initial orthogonality is no longer a constraint.

6. The source tasks considered do not enable effective zero-shot generalization for the chosen datasets

We fully endorse your point of view. The existence of similarities and shared knowledge between the source tasks and the target tasks is instrumental in significantly improving transfering effectively. Our primary objective is to showcase the exceptional performance of our methods in effectively preserving knowledge, even when applied to tasks with substantial gap in data scale, domains, and even task types.

As for the benefits derived from held-in datasets, the results obtained from fine-tuning on the MMLU task can offer the evidence. Despite encompassing different domains, both the MMLU dataset and Arc datasets are multiple-choice question-answering datasets that cover various reasoning questions. They share the commonality of featuring similar question types. Our method demonstrates a remarkable improvement of up to 5.5 points, which corresponds to a relative improvement of 12.5% compared to the baseline. Notably, this improvement is achieved even without prior exposure to any questions in the ARC evaluation, highlighting the impressive capability of our methods.

Finetune on MMLU:

Moreover, we also follow the previous related works to conduct the transfer learning([1], [2]). We have selected four pairs of in-pair tasks with 20-shots for T5 evaluation. Following [2], the selected pairs are: IMDb and SST2 (2-class sentiment classification), Amazon and Yelp reviews (5-class sentiment classification). And we also don't provide $**the task-ID**$ durning the transfer learning process, which makes the process more challenging, as shown below:

[1] Pretrained transformers improve out-of-distribution robustness [2] Progressive prompts: continual learning for language models

Thank you for providing your valuable feedback; it has truly enriched our discussion and made it highly rewarding. We actively follow the thoughts you have mentioned to enhance both the generality and efficiency of our method. Moreover, the transfer learning capabilities we previously discussed further highlight our method's capacity to share knowledge, thus serving as a robust supplement.

We apologize for any confusion caused by the separate training results. Following previous work, in the table provided, the first set of results (w/o transfer) represents the outcomes of training the tasks separately, without any influence from other tasks. Under the transfer learning setting, the reported results pertain to scenarios where the source tasks and target tasks possess similarities in terms of knowledge and task types. It is evident that our methods demonstrate the ability to share common capabilities among tasks.

In a more general setting of the Bert benchmark, we have conducted such experiments before. Upon reproducing these experiments with the same setting and configs, we have obtained the results of 78.9. And our method can achieve a performance of 80.0 under the same experimental conditions. It is crucial to acknowledge that the improvement achieved in this scenario may be relatively smaller compared to task-pair transfer learning, as we have previously discussed, primarily due to the larger task gap among the tasks involved.

As the discussion deadline is approaching, we kindly request if there are any remaining questions or concerns, such that we can provide further responses before the deadline. Thank you once again for your valuable time.