Paper: Scalable Language Model with Generalized Continual Learning (ICLR-2024)

Reviewer comment (Reproducibility):
It may be harder to reproduce the method compared to prior ones, as it is quite different from existing methods.

Author response:
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.

Reviewer comment (Evaluation):
I disagree with the claim regarding the weakness of replay-based methods in the introduction. While I do agree that replay-based methods require additional storage to save the previous datasets, this paper also requires saving the parameters of the previous task's low-rank parameters and saving the keys. Since text dataset saving does not require much storage, I believe comparing the storage is needed for a fair comparison.

Author response:
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%.

| Task | Finance | MMLU | Medical | Avg |
|--------|--------|--------|--------|--------|
|  Replay | 71.5  | 23.3 | 85.0 | 59.9 |

| Task | Medical | MMLU | Finance | Avg |
|--------|--------|--------|--------|--------|
|  Replay | 83.7  | 23.6 | 86.9 | 64.7 |

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.

Reviewer comment (Evaluation):
Comparing with the upper bound performance will be interesting by assuming the ideal update is made (i.e., D(p, p_i)=0 for non-relevant tasks and D(p, p_i)=1 for the current task). Maybe such results are not an upper bound, as the method has benefits on forward transfer.

Author response:
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.

Reviewer comment (Experiments):
Due to the fact that data replay is not a big issue (as mentioned in Point 1), I think adding more baselines related to data replay is needed, e.g., online SGD or other replay-based approaches [1].

Author response:
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%.

| Task | Finance | MMLU | Medical | Avg |
|--------|--------|--------|--------|--------|
|  Replay | 71.5  | 23.3 | 85.0 | 59.9 |

| Task | Medical | MMLU | Finance | Avg |
|--------|--------|--------|--------|--------|
|  Replay | 83.7  | 23.6 | 86.9 | 64.7 |

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.

Reviewer comment (Experiments):
The number of data points in a task should also be considered while using the weight increments from that task during the retrieval step, otherwise, weight increments of the task with 100 data points will have the same importance/impact as the task with 10000 data points. Just like it happens in the weighted averaging strategy used in the FedAvg algorithm in federated learning settings.

Author response:
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.

Reviewer comment (Evaluation):
The computational complexity of the proposed model seems high as the training is done for one data instance at a time and that too involves vector space retrieval of top-K relevant weight increments for each training example. Therefore, it would be fair if the authors could provide a comparison of time spent in training the proposed SLM method as compared to the baseline methods.

Author response:
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.

| Task | Retrieval | Generation | All | Per |
|--------|--------|--------|--------|--------|
|  Finance | 7.2ms  | 147.6ms | 154.8ms | 4.7% |
|  MMLU | 9.7ms  | 161.2ms | 170.9ms | 5.7% |
| Medical  | 8ms  | 2600ms | 2608ms | 0.3% |

Reviewer comment (Reproducibility):
It is not clear how authors have performed the initialization of (key, value) pairs for each task as mentioned in algorithm 1. It would be better if the authors could provide the mathematical formulation of this initialization strategy.

Author response:
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.

Reviewer comment (Writing):
[Section-3.1, Line-1] "researchs have" -> "research has" or "works have".

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

Reviewer comment (Experiments):
As the pretrained model $f_s$ is frozen and $f_{\theta}$ is always initialized with the same pre-trained initial weights, the keys of tasks with similar distributions can overlap with each other. In such cases, their values i.e., weight increments, should also be similar to each other. To understand this numerically, have authors done any such analysis to study the similarity between (key, value) pairs across tasks?

Author response:
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.  

| Task | ag_news | amazon | abpedia | yahoo | yelp |
|--------|--------|--------|--------|--------|--------|
|  Acc | 99.1%  | 98.1% | 99.5% | 99.2% | 97.2% |

Reviewer comment (Reproducibility):
Related to the above question, during the inference phase if DTKR provide similar (key, value) pairs from different tasks, then is there any mechanism to avoid this as they will provide redundant information and can be ignored for other relevant but non-overlapping knowledge?

Author response:
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.  

| Task | ag_news | amazon | abpedia | yahoo | yelp |
|--------|--------|--------|--------|--------|--------|
|  Acc | 99.1%  | 98.1% | 99.5% | 99.2% | 97.2% |

Reviewer comment (Writing):
The method's description appears somewhat convoluted, making it challenging to grasp the benefits fully. Relevant questions highlighting each part are in the Questions section.

Author response:
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.

Reviewer comment (Experiments):
A basic baseline employing an average feature extractor across examples of a task as a key and training a singular low-rank weight for each task as value is absent. In this method, each task would have just one key and one low-rank weight, and retrieval for examples from a new task or an existing task could be executed using either Top-1 or Top-k with existing keys.

Author response:
**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.

| Order | 4 | 5 | 6 | 7 | Avg |
|--------|--------|--------|--------|--------|--------|
| Acc  | 77.7  | 77.9 | 77.8 | 77.5 | 77.7 |

Reviewer comment (Experiments):
The retention of performance on trained tasks seems less beneficial since for inference on an example originating from one of the trained tasks, the already trained model could simply be employed. The principal advantage of these continual systems lies in their capability to be deployed for inference as zero-shot evaluation. However, the results from Table 6 appear unconvincing as the baseline LLAMA model already exhibits decent performance. A more robust experiment could entail learning a sequence of tasks continually and evaluating the zero-shot performance with every newly learned task, please see training datasets from T0 Held-in as depicted in Figure 2 from https://arxiv.org/abs/2110.08207 for experimental setup.

Author response:
**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:
| Task | Arc_c | Arc_e | Piqa | Wino |
|--------|--------|--------|--------|--------|
| FineTune  | 31.8 | 42.6 | 67.9 | 64.3 |
|  SLM | 44.7 | 76.0 | 76.3 | 67.7 |

MMLU:
| Task | Arc_c | Arc_e | Piqa | Wino |
|--------|--------|--------|--------|--------|
| FineTune  | 30.0  | 39.7 | 63.6 | 66.3 |
|  SLM | 49.4  | 76.7 | 76.6 | 66.2 |

Meical:
| Task | Arc_c | Arc_e | Piqa | Wino |
|--------|--------|--------|--------|--------|
| FineTune  | 38.9  | 67.1 | 76.5 | 66.9 |
|  SLM | 44.3  | 75.0 | 77.8 | 67.8 |

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